qbox_xyz.py 3.33 KB
Newer Older
1
#!/usr/bin/python
2 3 4 5 6 7 8 9 10
# Copyright 2016 The Regents of the University of California
# This file is part of Qbox
#
# qbox_xyz.py: extract first (or all) set(s) of atomic positions in xyz format
# from a Qbox output file or from a Qbox sample file using SAX 
# incremental parsing
#
# use: qbox_xyz.py [-all] {file|URL}
import os.path
11 12
import xml.sax
import sys
13
import urllib2
14

15 16
def usage():
  print "use: ",sys.argv[0]," [-all] {file|URL}"
17 18
  sys.exit()

19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34
argc=len(sys.argv)
if ( argc < 2 or argc > 3 ):
  usage()

# check if option "-all" is used
# "-all" option: extract all atomsets
# default: extract first atomset only
first_only = True
input_source = sys.argv[1]
if ( sys.argv[1] == "-all" ):
  if ( argc != 3 ):
    usage()
  first_only = False
  input_source = sys.argv[2]

# conversion from Bohr to Angstrom
35 36 37 38 39 40 41 42 43
a0=0.529177

# Qbox output handler to extract and process data
class QboxOutputHandler(xml.sax.handler.ContentHandler):
  def __init__(self):
    self.step = 0
    self.inAtomset = 0
    self.inAtom = 0
    self.inPosition = 0
44
    self.done_first = False
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

  def startElement(self, name, attributes):
    if name == "atomset":
      self.tau=[]
      self.atomname=[]
      self.inAtomset = 1
    elif (name == "unit_cell") & self.inAtomset:
      self.cell_a = attributes["a"]
      self.cell_b = attributes["b"]
      self.cell_c = attributes["c"]
    elif (name == "atom") & self.inAtomset:
      self.atomname.append(attributes["name"])
      self.inAtom = 1
    elif (name == "position") & self.inAtom:
        self.buffer = ""
        self.inPosition = 1

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

  def endElement(self, name):
    if (name == "atom") and self.inAtomset:
      self.inAtom = 0
    if (name == "position") & self.inAtom:
      pos = self.buffer.split()
      x = a0*float(pos[0])
      y = a0*float(pos[1])
      z = a0*float(pos[2])
      self.tau.append([x,y,z])
      self.inPosition = 0
    elif name == "atomset":
      self.step += 1
      print len(self.tau)
      avec = self.cell_a.split()
      bvec = self.cell_b.split()
      cvec = self.cell_c.split()
      print self.step,\
      '%.6f'%(a0*float(avec[0])),\
      '%.6f'%(a0*float(avec[1])),\
      '%.6f'%(a0*float(avec[2])),\
      '%.6f'%(a0*float(bvec[0])),\
      '%.6f'%(a0*float(bvec[1])),\
      '%.6f'%(a0*float(bvec[2])),\
      '%.6f'%(a0*float(cvec[0])),\
      '%.6f'%(a0*float(cvec[1])),\
      '%.6f'%(a0*float(cvec[2]))
      for i in range(len(self.tau)):
        print self.atomname[i],'%.6f'%self.tau[i][0],\
                               '%.6f'%self.tau[i][1],\
                               '%.6f'%self.tau[i][2]
      self.inAtomset = 0
97
      self.done_first = True
98 99 100 101

parser = xml.sax.make_parser()
handler = QboxOutputHandler()
parser.setContentHandler(handler)
102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124
# 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 !="" and not (first_only and handler.done_first) ):
    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 !="" and not (first_only and handler.done_first) ):
      parser.feed(s)
      s = f.read(8192)
    f.close()
  except (ValueError,urllib2.HTTPError) as e:
    print e
    sys.exit()

parser.reset()