space_group.py

import os
import inspect
import re
import pickle
import itertools
newpath = os.path.dirname(inspect.getfile(inspect.currentframe()))  # script directory
from xyzMath import Vec, Mat, Xform, Ux, Uy, Uz
from pymol_util import pymol, cgo_sphere, cube

class SpaceGroup(object):
   """docstring for SpaceGroup"""
   def __init__(self, names, celltype, ops):
      super(SpaceGroup, self).__init__()
      self.name = names[0]
      self.names = names
      self.celltype = celltype
      self.ops = ops
      self.frames = list()
      for op in self.ops:
         self.frames.append(self.op2xform(op))

   def op2vec(self, op):
      op = op.replace(" ", "")
      v, t = Vec(), 0
      for o in op.replace("-", "+-").strip("+").split("+"):
         if o == 'X':
            v += Ux
         elif o == '-X':
            v -= Ux
         elif o == 'Y':
            v += Uy
         elif o == '-Y':
            v -= Uy
         elif o == 'Z':
            v += Uz
         elif o == '-Z':
            v -= Uz
         elif o.count("/"):
            n, d = o.split("/")
            t += float(n) / float(d)
         else:
            raise ValueError
      return v, t

   def op2xform(self, op):
      Xop, Yop, Zop = op.split(",")
      # print op
      X, Tx = self.op2vec(Xop)
      Y, Ty = self.op2vec(Yop)
      Z, Tz = self.op2vec(Zop)
      # print X
      # print Y
      # print Z
      # print Tx,Ty,Tz
      # print
      return Xform(Mat(X.x, X.y, X.z, Y.x, Y.y, Y.z, Z.x, Z.y, Z.z), Vec(Tx, Ty, Tz))
      # return Xform( Mat( X.x, Y.x, Z.x,
      # 	               X.y, Y.y, Z.y,
      # 	               X.z, Y.z, Z.z,
      #                    ), Vec(Tx,Ty,Tz) )
      # raise NotImplementedError

   def __str__(self):
      return "SpaceGroup %s, unit cell size: %i" % (self.name, len(self.frames))

   def show(self, points=[Vec(7, 12, 13), Vec(33, 34, 35)]):
      v = pymol.cmd.get_view()
      CGO = list()
      for x in self.frames:
         x.t = 100.0 * x.t
         assert self.celltype == "CUBIC"
         for p in points:
            for i, j, k in itertools.product(list(range(-1, 2)), list(range(-1, 2)),
                                             list(range(-1, 2))):
               CGO.extend(cgo_sphere(x * p + Vec(i, j, k) * 100.0))
      pymol.cmd.load_cgo(CGO, self.name)
      cube(Vec(0, 0, 0), Vec(100, 100, 100))
      pymol.cmd.set_view(v)

path = os.path.dirname(inspect.getfile(inspect.currentframe()))
picklefile = "data/spacegroups.pkl"
if False:  # os.path.exists(picklefile):
   with open(picklefile) as infile:
      spacegroups = pickle.load(infile)
else:
   spacegroups = dict()
   with open(path + '/data/symop.lib') as f:
      while f:
         header = f.readline().strip()
         if not header:
            break
         header = re.sub("!+", "!", header)
         # print header
         comment = ''
         if header.find("!") >= 0:
            splt = header.split("!")
            header, comment = splt[0], "".join(splt[1:])
            header = header.strip()
            comment = comment.strip()
         header = header.split()
         idx, nop, dunno, name1, name2, celltype = header[:6]
         idx, nop, dunno = int(idx), int(nop), int(dunno)
         names = " ".join(header[6:]).strip("'").split("' '")
         ops = list()
         for iop in range(nop):
            ops.append(f.readline().strip())
         # print idx,names,len(ops),comment
         sg = SpaceGroup(names, celltype, ops)
         spacegroups[sg.name] = sg
   with open(picklefile, "w") as out:
      pickle.dump(spacegroups, out)

print("loaded", len(spacegroups), "spacegroups")