Source Code for Module MolKit.trajParser

  1  # 
  2  # $Header: /opt/cvs/python/packages/share1.5/MolKit/trajParser.py,v 1.2 2006/10/17 18:22:13 annao Exp $ 
  3  # 
  4  # $Id: trajParser.py,v 1.2 2006/10/17 18:22:13 annao Exp $ 
  5  # 
  6   
  7  ## This module implements classes for parsing GROMACS trajectory files 
  8   
  9  from struct import calcsize 
 10  from xdrlib import Unpacker 
 11  DIM = 3 
 12  import os 
 13   
 14 -class trrParser: 
 15      """ Parses .trr Gromacs trajectory file """ 
 16       
 17 -    def __init__(self, file): 
 18   
 19          self.nframes = 0 
 20          self.headers=[] 
 21  ##         header = { 
 22  ##             "ir_size":None, 
 23  ##             "e_size": None, 
 24  ##             "box_size" :None, 
 25  ##             "vir_size": None, 
 26  ##             "pres_size": None, 
 27  ##             "top_size ": None, 
 28  ##             "sym_size ": None, 
 29  ##             "x_size ": None, 
 30  ##             "v_size ": None, 
 31  ##             "f_size ": None, 
 32  ##             "natoms ": None, 
 33  ##             "step ": None, 
 34  ##             "nre ": None, 
 35  ##             "version":'GMX_trn_file', 
 36  ##             "magicnum" :1993, 
 37  ##             "bDouble ": None, 
 38  ##             } 
 39          self.file = file 
 40          self.coords = [] 
 41          self.velocities = {} 
 42          self.forces = {} 
 43   
 44   
 45 -    def nFloatSize(self, h): 
 46   
 47          nflsize=0; 
 48          if h["box_size"]: 
 49              nflsize = h["box_size"]/(DIM*DIM); 
 50          elif h["x_size"]: 
 51              nflsize = h["x_size"]/(h["natoms"]*DIM); 
 52          elif h["v_size"]: 
 53              nflsize = h["v_size"]/(h["natoms"]*DIM); 
 54          elif h["f_size"]: 
 55              nflsize = h["f_size"]/(h["natoms"]*DIM); 
 56          else:  
 57              print "Can not determine precision of trr file" 
 58     
 59          if (nflsize != calcsize("f")) and (nflsize != calcsize("d")): 
 60              print "Float size %d. Maybe different CPU?"%nflsize 
 61         
 62          return nflsize 
 63   
 64 -    def read(self): 
 65          fext = os.path.splitext(self.file)[-1] 
 66          assert fext == ".trr" 
 67          fp = open(self.file, "rb") 
 68          self.data = data = fp.read() 
 69          self.coords = [] 
 70          self.v = {} 
 71          self.f = {} 
 72          self.up = Unpacker(data) 
 73          curpos = self.up.get_position() 
 74          datasize = len(data) 
 75          nframe = 0 
 76          #each frame begins with a header 
 77          while curpos < datasize: 
 78              #print "current position:", curpos 
 79              h = self.readHeader(nframe) 
 80              self.headers.append(h) 
 81              self.readData(nframe) 
 82              nframe = nframe + 1 
 83              curpos = self.up.get_position() 
 84          #print "end of readTraj, cur position : %d, datazize: %d" %(self.up.get_position(), datasize) 
 85          self.nframes = nframe 
 86          if self.nframes: 
 87              return 1 
 88          else: 
 89              return 0 
 90   
 91 -    def readHeader(self, nframe): 
 92          #print "reading header, frame %d" %nframe 
 93          up = self.up 
 94          header = {} 
 95          header["magicnum"] = up.unpack_int() 
 96          #print "magicnum: ", header["magicnum"] 
 97          #1993 
 98          i1 = up.unpack_int() 
 99          #13 
100          version = up.unpack_string() 
101          #'GMX_trn_file' 
102          header["version"] = version 
103          #print "version:", version 
104          header["ir_size"] = up.unpack_int() 
105          #print "ir_size=", header["ir_size"] 
106          header["e_size"]=up.unpack_int() 
107          #print "e_size=", header["e_size"] 
108          header["box_size"] =up.unpack_int() 
109          #print "box_size=", header["box_size"] 
110          header["vir_size"] = up.unpack_int() 
111          #print "vir_size=", header["vir_size"] 
112          header["pres_size"] = up.unpack_int() 
113          #print "pres_size=", header["pres_size"] 
114          header["top_size"]=up.unpack_int() 
115          #print "top_size=", header["top_size"] 
116   
117          header["sym_size"]=up.unpack_int() 
118          #print "sym_size=", header["sym_size"] 
119   
120          header["x_size"]=up.unpack_int() 
121          #print "x_size=", header["x_size"] 
122   
123          header["v_size"]=up.unpack_int() 
124          #print "v_size=", header["v_size"] 
125   
126          header["f_size"]=up.unpack_int() 
127          #print "f_size=", header["f_size"] 
128   
129          header["natoms"]=up.unpack_int() 
130          #print "natoms=", header["natoms"] 
131          header["step"]=up.unpack_int() 
132          #print "step=", header["step"] 
133          header["nre"]=up.unpack_int() 
134          #print "nre=", header["nre"] 
135   
136           
137          if self.nFloatSize(header) == calcsize("d"): 
138              header["bDouble"] = True 
139          else: 
140              header["bDouble"] = False 
141   
142          if header["bDouble"]: 
143              header["time"] = up.unpack_double() 
144              header["lam"] = up.umpack_double() 
145          else: 
146              header["time"] = up.unpack_float() 
147              header["lam"] = up.unpack_float() 
148   
149          #print "time = ", header["time"] 
150          #print "lambda = ", header["lam"] 
151          #print "natoms=%10d  step=%10d  time=%10g  lambda=%10g"% (header["natoms"],header["step"],header["time"],header["lam"]) 
152          #print "current position:",  up.get_position() 
153          return header 
154   
155   
156 -    def readData(self, nframe): 
157          up = self.up 
158          h = self.headers[nframe] 
159          box = [] 
160          if h["box_size"] != 0 : 
161              for i in range(3): 
162                  box.append(up.unpack_farray(3, up.unpack_float)) 
163              #print " box (3x3):" 
164              #print box 
165              self.headers[nframe]["box"] = box 
166          pv = [] 
167          if h["vir_size"] != 0: 
168              for i in range(3): 
169                  pv.append(up.unpack_farray(3, up.unpack_float)) 
170              #print "pv:" 
171              #print pv 
172              self.headers[nframe]["pv"] = pv 
173          if h["pres_size"]!= 0: 
174              pv.append(up.unpack_farray(3, up.unpack_float)) 
175              #print "pv:" 
176              #print pv 
177              self.headers[nframe]["pv"] = pv 
178          natoms = h["natoms"] 
179          if h["x_size"] != 0: 
180              x= [] 
181              for i in range (natoms):  
182                  x.append(up.unpack_farray(3, up.unpack_float)) 
183              #self.coords["frame%d"%nframe] = x 
184              self.coords.append(x) 
185   
186          if h["v_size"] != 0: 
187              v = [] 
188              for i in range (natoms): 
189                  v.append(up.unpack_farray(3, up.unpack_float)) 
190              self.velocities["frame%d"%nframe] = v 
191          if h["f_size"] != 0: 
192              f = [] 
193              for i in range (natoms): 
194                  f.append(up.unpack_farray(3, up.unpack_float)) 
195              self.forces["frame%d"%nframe] = f 
196          #print up.get_position() 
197   
198   
199   
200   
201   
202 -class xtcParser: 
203   
204      """ Parses .xtc Gromacs trajectory file """ 
205   
206 -    def __init__(self, file): 
207           
208          self.nframes = 0 
209          self.headers=[] 
210  ##         header = { 
211  ##             'step ': None, 
212  ##             'frame':None, 
213  ##             'time':None, 
214  ##             'prec':None 
215  ##             } 
216          self.file = file 
217          self.coords = [] 
218          self.headers = [] 
219          self.nframes = 0 
220          self.velocities = None #not available inthis file format 
221          self.forces = None 
222          # try to import a platform-dependent module xtcparcer 
223          # containing C functions to read this file format 
224          try: 
225              from cMolKit import xtcparser 
226          except: 
227              print "WARNING: could not import cMolKit.xtcparser - No parser is available for xtc files." 
228              self.file=None 
229   
230   
231 -    def read(self): 
232          if not self.file: 
233              return 0 
234          assert os.path.exists(self.file) 
235          fext = os.path.splitext(self.file)[-1] 
236          assert fext == ".xtc" 
237           
238          from cMolKit import xtcparser 
239          self.coords, self.headers = xtcparser.read_xtc(self.file) 
240          if self.coords: 
241              self.nframes = len(self.coords) 
242              return 1 
243          else: 
244              return 0 
245