1
+ # Pizza.py toolkit, www.cs.sandia.gov/~sjplimp/pizza.html
2
+ # Steve Plimpton, sjplimp@sandia.gov, Sandia National Laboratories
3
+ #
4
+ # Copyright (2005) Sandia Corporation. Under the terms of Contract
5
+ # DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
6
+ # certain rights in this software. This software is distributed under
7
+ # the GNU General Public License.
8
+
9
+ # data tool
10
+
11
+ oneline = "Read, write, manipulate LAMMPS data files"
12
+
13
+ docstr = """
14
+ d = data("data.poly") read a LAMMPS data file, can be gzipped
15
+ d = data() create an empty data file
16
+
17
+ d.map(1,"id",3,"x") assign names to atom columns (1-N)
18
+
19
+ coeffs = d.get("Pair Coeffs") extract info from data file section
20
+ q = d.get("Atoms",4)
21
+
22
+ 1 arg = all columns returned as 2d array of floats
23
+ 2 args = Nth column returned as vector of floats
24
+
25
+ d.reorder("Atoms",1,3,2,4,5) reorder columns (1-N) in a data file section
26
+
27
+ 1,3,2,4,5 = new order of previous columns, can delete columns this way
28
+
29
+ d.title = "My LAMMPS data file" set title of the data file
30
+ d.headers["atoms"] = 1500 set a header value
31
+ d.sections["Bonds"] = lines set a section to list of lines (with newlines)
32
+ d.delete("bonds") delete a keyword or section of data file
33
+ d.delete("Bonds")
34
+ d.replace("Atoms",5,vec) replace Nth column of section with vector
35
+ d.newxyz(dmp,1000) replace xyz in Atoms with xyz of snapshot N
36
+
37
+ newxyz assumes id,x,y,z are defined in both data and dump files
38
+ also replaces ix,iy,iz if they are defined
39
+
40
+ index,time,flag = d.iterator(0/1) loop over single data file snapshot
41
+ time,box,atoms,bonds,tris,lines = d.viz(index) return list of viz objects
42
+
43
+ iterator() and viz() are compatible with equivalent dump calls
44
+ iterator() called with arg = 0 first time, with arg = 1 on subsequent calls
45
+ index = timestep index within dump object (only 0 for data file)
46
+ time = timestep value (only 0 for data file)
47
+ flag = -1 when iteration is done, 1 otherwise
48
+ viz() returns info for specified timestep index (must be 0)
49
+ time = 0
50
+ box = [xlo,ylo,zlo,xhi,yhi,zhi]
51
+ atoms = id,type,x,y,z for each atom as 2d array
52
+ bonds = id,type,x1,y1,z1,x2,y2,z2,t1,t2 for each bond as 2d array
53
+ NULL if bonds do not exist
54
+ tris = NULL
55
+ lines = NULL
56
+
57
+ d.write("data.new") write a LAMMPS data file
58
+ """
59
+
60
+ # History
61
+ # 8/05, Steve Plimpton (SNL): original version
62
+ # 11/07, added triclinic box support
63
+
64
+ # ToDo list
65
+
66
+ # Variables
67
+ # title = 1st line of data file
68
+ # names = dictionary with atom attributes as keys, col #s as values
69
+ # headers = dictionary with header name as key, value or tuple as values
70
+ # sections = dictionary with section name as key, array of lines as values
71
+ # nselect = 1 = # of snapshots
72
+
73
+ # Imports and external programs
74
+
75
+ from os import popen
76
+
77
+ try : tmp = PIZZA_GUNZIP
78
+ except : PIZZA_GUNZIP = "gunzip"
79
+
80
+ # Class definition
81
+
82
+ class data :
83
+
84
+ # --------------------------------------------------------------------
85
+
86
+ def __init__ (self ,* list ):
87
+ self .nselect = 1
88
+
89
+ if len (list ) == 0 :
90
+ self .title = "LAMMPS data file"
91
+ self .names = {}
92
+ self .headers = {}
93
+ self .sections = {}
94
+ return
95
+
96
+ file = list [0 ]
97
+ if file [- 3 :] == ".gz" : f = popen ("%s -c %s" % (PIZZA_GUNZIP ,file ),'r' )
98
+ else : f = open (file )
99
+
100
+ self .title = f .readline ()
101
+ self .names = {}
102
+
103
+ headers = {}
104
+ while 1 :
105
+ line = f .readline ()
106
+ line = line .strip ()
107
+ if len (line ) == 0 :
108
+ continue
109
+ found = 0
110
+ for keyword in hkeywords :
111
+ if line .find (keyword ) >= 0 :
112
+ found = 1
113
+ words = line .split ()
114
+ if keyword == "xlo xhi" or keyword == "ylo yhi" or \
115
+ keyword == "zlo zhi" :
116
+ headers [keyword ] = (float (words [0 ]),float (words [1 ]))
117
+ elif keyword == "xy xz yz" :
118
+ headers [keyword ] = \
119
+ (float (words [0 ]),float (words [1 ]),float (words [2 ]))
120
+ else :
121
+ headers [keyword ] = int (words [0 ])
122
+ if not found :
123
+ break
124
+
125
+ sections = {}
126
+ while 1 :
127
+ found = 0
128
+ for pair in skeywords :
129
+ keyword ,length = pair [0 ],pair [1 ]
130
+ if keyword == line :
131
+ found = 1
132
+ if not headers .has_key (length ):
133
+ raise StandardError , \
134
+ "data section %s has no matching header value" % line
135
+ f .readline ()
136
+ list = []
137
+ for i in xrange (headers [length ]): list .append (f .readline ())
138
+ sections [keyword ] = list
139
+ if not found :
140
+ raise StandardError ,"invalid section %s in data file" % line
141
+ f .readline ()
142
+ line = f .readline ()
143
+ if not line :
144
+ break
145
+ line = line .strip ()
146
+
147
+ f .close ()
148
+ self .headers = headers
149
+ self .sections = sections
150
+
151
+ # --------------------------------------------------------------------
152
+ # assign names to atom columns
153
+
154
+ def map (self ,* pairs ):
155
+ if len (pairs ) % 2 != 0 :
156
+ raise StandardError , "data map() requires pairs of mappings"
157
+ for i in range (0 ,len (pairs ),2 ):
158
+ j = i + 1
159
+ self .names [pairs [j ]] = pairs [i ]- 1
160
+
161
+ # --------------------------------------------------------------------
162
+ # extract info from data file fields
163
+
164
+ def get (self ,* list ):
165
+ if len (list ) == 1 :
166
+ field = list [0 ]
167
+ array = []
168
+ lines = self .sections [field ]
169
+ for line in lines :
170
+ words = line .split ()
171
+ values = map (float ,words )
172
+ array .append (values )
173
+ return array
174
+ elif len (list ) == 2 :
175
+ field = list [0 ]
176
+ n = list [1 ] - 1
177
+ vec = []
178
+ lines = self .sections [field ]
179
+ for line in lines :
180
+ words = line .split ()
181
+ vec .append (float (words [n ]))
182
+ return vec
183
+ else :
184
+ raise StandardError , "invalid arguments for data.get()"
185
+
186
+ # --------------------------------------------------------------------
187
+ # reorder columns in a data file field
188
+
189
+ def reorder (self ,name ,* order ):
190
+ n = len (order )
191
+ natoms = len (self .sections [name ])
192
+ oldlines = self .sections [name ]
193
+ newlines = natoms * ["" ]
194
+ for index in order :
195
+ for i in xrange (len (newlines )):
196
+ words = oldlines [i ].split ()
197
+ newlines [i ] += words [index - 1 ] + " "
198
+ for i in xrange (len (newlines )):
199
+ newlines [i ] += "\n "
200
+ self .sections [name ] = newlines
201
+
202
+ # --------------------------------------------------------------------
203
+ # replace a column of named section with vector of values
204
+
205
+ def replace (self ,name ,icol ,vector ):
206
+ lines = self .sections [name ]
207
+ newlines = []
208
+ j = icol - 1
209
+ for i in xrange (len (lines )):
210
+ line = lines [i ]
211
+ words = line .split ()
212
+ words [j ] = str (vector [i ])
213
+ newline = ' ' .join (words ) + '\n '
214
+ newlines .append (newline )
215
+ self .sections [name ] = newlines
216
+
217
+ # --------------------------------------------------------------------
218
+ # replace x,y,z in Atoms with x,y,z values from snapshot ntime of dump object
219
+ # assumes id,x,y,z are defined in both data and dump files
220
+ # also replaces ix,iy,iz if they are defined
221
+
222
+ def newxyz (self ,dm ,ntime ):
223
+ nsnap = dm .findtime (ntime )
224
+
225
+ dm .sort (ntime )
226
+ x ,y ,z = dm .vecs (ntime ,"x" ,"y" ,"z" )
227
+
228
+ self .replace ("Atoms" ,self .names ['x' ]+ 1 ,x )
229
+ self .replace ("Atoms" ,self .names ['y' ]+ 1 ,y )
230
+ self .replace ("Atoms" ,self .names ['z' ]+ 1 ,z )
231
+
232
+ if dm .names .has_key ("ix" ) and self .names .has_key ("ix" ):
233
+ ix ,iy ,iz = dm .vecs (ntime ,"ix" ,"iy" ,"iz" )
234
+ self .replace ("Atoms" ,self .names ['ix' ]+ 1 ,ix )
235
+ self .replace ("Atoms" ,self .names ['iy' ]+ 1 ,iy )
236
+ self .replace ("Atoms" ,self .names ['iz' ]+ 1 ,iz )
237
+
238
+ # --------------------------------------------------------------------
239
+ # delete header value or section from data file
240
+
241
+ def delete (self ,keyword ):
242
+
243
+ if self .headers .has_key (keyword ): del self .headers [keyword ]
244
+ elif self .sections .has_key (keyword ): del self .sections [keyword ]
245
+ else : raise StandardError , "keyword not found in data object"
246
+
247
+ # --------------------------------------------------------------------
248
+ # write out a LAMMPS data file
249
+
250
+ def write (self ,file ):
251
+ f = open (file ,"w" )
252
+ print >> f ,self .title
253
+ for keyword in hkeywords :
254
+ if self .headers .has_key (keyword ):
255
+ if keyword == "xlo xhi" or keyword == "ylo yhi" or \
256
+ keyword == "zlo zhi" :
257
+ pair = self .headers [keyword ]
258
+ print >> f ,pair [0 ],pair [1 ],keyword
259
+ elif keyword == "xy xz yz" :
260
+ triple = self .headers [keyword ]
261
+ print >> f ,triple [0 ],triple [1 ],triple [2 ],keyword
262
+ else :
263
+ print >> f ,self .headers [keyword ],keyword
264
+ for pair in skeywords :
265
+ keyword = pair [0 ]
266
+ if self .sections .has_key (keyword ):
267
+ print >> f ,"\n %s\n " % keyword
268
+ for line in self .sections [keyword ]:
269
+ print >> f ,line ,
270
+ f .close ()
271
+
272
+ # --------------------------------------------------------------------
273
+ # iterator called from other tools
274
+
275
+ def iterator (self ,flag ):
276
+ if flag == 0 : return 0 ,0 ,1
277
+ return 0 ,0 ,- 1
278
+
279
+ # --------------------------------------------------------------------
280
+ # time query from other tools
281
+
282
+ def findtime (self ,n ):
283
+ if n == 0 : return 0
284
+ raise StandardError , "no step %d exists" % (n )
285
+
286
+ # --------------------------------------------------------------------
287
+ # return list of atoms and bonds to viz for data object
288
+
289
+ def viz (self ,isnap ):
290
+ if isnap : raise StandardError , "cannot call data.viz() with isnap != 0"
291
+
292
+ id = self .names ["id" ]
293
+ type = self .names ["type" ]
294
+ x = self .names ["x" ]
295
+ y = self .names ["y" ]
296
+ z = self .names ["z" ]
297
+
298
+ xlohi = self .headers ["xlo xhi" ]
299
+ ylohi = self .headers ["ylo yhi" ]
300
+ zlohi = self .headers ["zlo zhi" ]
301
+ box = [xlohi [0 ],ylohi [0 ],zlohi [0 ],xlohi [1 ],ylohi [1 ],zlohi [1 ]]
302
+
303
+ # create atom list needed by viz from id,type,x,y,z
304
+
305
+ atoms = []
306
+ atomlines = self .sections ["Atoms" ]
307
+ for line in atomlines :
308
+ words = line .split ()
309
+ atoms .append ([int (words [id ]),int (words [type ]),
310
+ float (words [x ]),float (words [y ]),float (words [z ])])
311
+
312
+ # create list of current bond coords from list of bonds
313
+ # assumes atoms are sorted so can lookup up the 2 atoms in each bond
314
+
315
+ bonds = []
316
+ if self .sections .has_key ("Bonds" ):
317
+ bondlines = self .sections ["Bonds" ]
318
+ for line in bondlines :
319
+ words = line .split ()
320
+ bid ,btype = int (words [0 ]),int (words [1 ])
321
+ atom1 ,atom2 = int (words [2 ]),int (words [3 ])
322
+ atom1words = atomlines [atom1 - 1 ].split ()
323
+ atom2words = atomlines [atom2 - 1 ].split ()
324
+ bonds .append ([bid ,btype ,
325
+ float (atom1words [x ]),float (atom1words [y ]),
326
+ float (atom1words [z ]),
327
+ float (atom2words [x ]),float (atom2words [y ]),
328
+ float (atom2words [z ]),
329
+ float (atom1words [type ]),float (atom2words [type ])])
330
+
331
+ tris = []
332
+ lines = []
333
+ return 0 ,box ,atoms ,bonds ,tris ,lines
334
+
335
+ # --------------------------------------------------------------------
336
+ # return box size
337
+
338
+ def maxbox (self ):
339
+ xlohi = self .headers ["xlo xhi" ]
340
+ ylohi = self .headers ["ylo yhi" ]
341
+ zlohi = self .headers ["zlo zhi" ]
342
+ return [xlohi [0 ],ylohi [0 ],zlohi [0 ],xlohi [1 ],ylohi [1 ],zlohi [1 ]]
343
+
344
+ # --------------------------------------------------------------------
345
+ # return number of atom types
346
+
347
+ def maxtype (self ):
348
+ return self .headers ["atom types" ]
349
+
350
+ # --------------------------------------------------------------------
351
+ # data file keywords, both header and main sections
352
+
353
+ hkeywords = ["atoms" ,"ellipsoids" ,"lines" ,"triangles" ,"bodies" ,
354
+ "bonds" ,"angles" ,"dihedrals" ,"impropers" ,
355
+ "atom types" ,"bond types" ,"angle types" ,"dihedral types" ,
356
+ "improper types" ,"xlo xhi" ,"ylo yhi" ,"zlo zhi" ,"xy xz yz" ]
357
+
358
+ skeywords = [["Masses" ,"atom types" ],
359
+ ["Atoms" ,"atoms" ],["Ellipsoids" ,"ellipsoids" ],
360
+ ["Lines" ,"lines" ],["Triangles" ,"triangles" ],["Bodies" ,"bodies" ],
361
+ ["Bonds" ,"bonds" ],
362
+ ["Angles" ,"angles" ],["Dihedrals" ,"dihedrals" ],
363
+ ["Impropers" ,"impropers" ],["Velocities" ,"atoms" ],
364
+ ["Pair Coeffs" ,"atom types" ],
365
+ ["Bond Coeffs" ,"bond types" ],["Angle Coeffs" ,"angle types" ],
366
+ ["Dihedral Coeffs" ,"dihedral types" ],
367
+ ["Improper Coeffs" ,"improper types" ],
368
+ ["BondBond Coeffs" ,"angle types" ],
369
+ ["BondAngle Coeffs" ,"angle types" ],
370
+ ["MiddleBondTorsion Coeffs" ,"dihedral types" ],
371
+ ["EndBondTorsion Coeffs" ,"dihedral types" ],
372
+ ["AngleTorsion Coeffs" ,"dihedral types" ],
373
+ ["AngleAngleTorsion Coeffs" ,"dihedral types" ],
374
+ ["BondBond13 Coeffs" ,"dihedral types" ],
375
+ ["AngleAngle Coeffs" ,"improper types" ],
376
+ ["Molecules" ,"atoms" ],
377
+ ["Tinker Types" ,"atoms" ]]
378
+
0 commit comments