qso_plot_species.py 4.58 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170
#!/usr/bin/python
# Copyright 2018 The Regents of the University of California
# This file is part of Qbox
#
# qso_plot_species.py: plot QSO potential in Gnuplot format
# using SAX incremental parsing
#
# use: qso_plot_species.py {file|URL}
import os.path
import xml.sax
import sys
import urllib2

def usage():
  print "use: ",sys.argv[0]," {file|URL}"
  sys.exit()

argc=len(sys.argv)
if ( argc != 2 ):
  usage()

input_source = sys.argv[1]

# Qbox output handler to extract and process data
class QSOSpeciesHandler(xml.sax.handler.ContentHandler):
  def __init__(self):

    self.inDescription = False
    self.inSymbol = False
    self.inAtomicNumber = False
    self.inMass = False
    self.inValenceCharge = False
    self.inMeshSpacing = False
    self.inProjector = False

    self.inNCP = False
    self.inLmax = False
    self.inLLocal = False
    self.inNquad = False
    self.inRquad = False
    self.inRadialPotential = False
    self.inRadialFunction = False

    self.inNCSLP = False
    self.inLocalPotential = False
    self.inDij = False

    self.buffer=""

  def startElement(self, name, attributes):
    if name == "symbol":
      self.symbol = ""
      self.buffer = ""
    elif name == "atomic_number":
      self.atomic_number = 0
      self.buffer = ""
    elif name == "valence_charge":
      self.valence_charge= 0
      self.buffer = ""
    elif (name == "mass"):
      self.mass = 0.0
      self.buffer = ""
    elif (name == "mesh_spacing"):
      self.mesh_spacing = 0.0
      self.buffer = ""
    elif (name == "norm_conserving_semilocal_pseudopotential"):
      self.inNCSLP = True
      print "# Norm-conserving semilocal pseudopotential"
    elif (name == "norm_conserving_pseudopotential"):
      self.inNCP = True
      print "# Norm-conserving pseudopotential"
    elif (name == "projector"):
      if self.inNCP:
        self.l = int(attributes["l"])
        self.size = int(attributes["size"])
        print "# projector l=",self.l," size=",self.size
      if self.inNCSLP:
        self.l = int(attributes["l"])
        self.i = int(attributes["i"])
        self.size = int(attributes["size"])
        print "# projector l=",self.l," i=",self.i," size=",self.size
      self.buffer = ""
    elif (name == "local_potential"):
      self.size = int(attributes["size"])
      print "# local potential, size=",self.size
      self.buffer = ""
    elif (name == "radial_potential"):
      self.buffer = ""

  def characters(self, data):
      self.buffer += data

  def endElement(self, name):
    if (name == "symbol"):
      print "# symbol:",self.buffer
      self.buffer = ""
    elif name == "atomic_number":
      print "# Z:",self.buffer
      self.buffer = ""
    elif name == "valence_charge":
      self.valence_charge=int(self.buffer)
      print "# valence charge:",self.valence_charge
      self.buffer = ""
    elif (name == "mass"):
      print "# mass:",self.buffer
      self.buffer = ""
    elif (name == "mesh_spacing"):
      self.mesh_spacing = float(self.buffer)
      print "# mesh spacing:",self.mesh_spacing
      self.buffer = ""
    elif (name == "norm_conserving_semilocal_pseudopotential"):
      self.inNCSLP = False
    elif (name == "norm_conserving_pseudopotential"):
      self.inNCP = False
    elif (name == "local_potential"):
      self.p = self.buffer.split()
      for i in range(len(self.p)):
        r = i * self.mesh_spacing
        val = float(self.p[i])
        print '%.6f'%r, '%.10e'%val
      print
      print
      self.buffer = ""
    elif (name == "projector"):
      if self.inNCSLP:
        self.p = self.buffer.split()
        for i in range(len(self.p)):
          r = i * self.mesh_spacing
          val = float(self.p[i])
          print '%.6f'%r, '%.10e'%val
        print
        print
        self.buffer = ""
    elif (name == "radial_potential"):
      print "# radial potential"
      self.p = self.buffer.split()
      for i in range(len(self.p)):
        r = i * self.mesh_spacing
        val = float(self.p[i])
        print '%.6f'%r, '%.10e'%val
      print
      print
      self.buffer = ""

parser = xml.sax.make_parser()
handler = QSOSpeciesHandler()
parser.setContentHandler(handler)
# test if input_source is a local file
# if not, process as a URL
if ( os.path.isfile(input_source) ):
  file = open(input_source)
  s = file.read(8192)
  while ( s !="" ):
    parser.feed(s)
    s = file.read(8192)
  file.close()
else:
  # attempt to open as a URL
  try:
    f = urllib2.urlopen(input_source)
    s = f.read(8192)
    while ( s !="" ):
      parser.feed(s)
      s = f.read(8192)
    f.close()
  except (ValueError,urllib2.HTTPError) as e:
    print e
    sys.exit()

parser.reset()