382 lines
12 KiB
Python
382 lines
12 KiB
Python
# Pizza.py toolkit, www.cs.sandia.gov/~sjplimp/pizza.html
|
|
# Steve Plimpton, sjplimp@sandia.gov, Sandia National Laboratories
|
|
#
|
|
# Copyright (2005) Sandia Corporation. Under the terms of Contract
|
|
# DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
|
|
# certain rights in this software. This software is distributed under
|
|
# the GNU General Public License.
|
|
|
|
# data tool
|
|
|
|
oneline = "Read, write, manipulate LAMMPS data files"
|
|
|
|
docstr = """
|
|
d = data("data.poly") read a LAMMPS data file, can be gzipped
|
|
d = data() create an empty data file
|
|
|
|
d.map(1,"id",3,"x") assign names to atom columns (1-N)
|
|
|
|
coeffs = d.get("Pair Coeffs") extract info from data file section
|
|
q = d.get("Atoms",4)
|
|
|
|
1 arg = all columns returned as 2d array of floats
|
|
2 args = Nth column returned as vector of floats
|
|
|
|
d.reorder("Atoms",1,3,2,4,5) reorder columns (1-N) in a data file section
|
|
|
|
1,3,2,4,5 = new order of previous columns, can delete columns this way
|
|
|
|
d.title = "My LAMMPS data file" set title of the data file
|
|
d.headers["atoms"] = 1500 set a header value
|
|
d.sections["Bonds"] = lines set a section to list of lines (with newlines)
|
|
d.delete("bonds") delete a keyword or section of data file
|
|
d.delete("Bonds")
|
|
d.replace("Atoms",5,vec) replace Nth column of section with vector
|
|
d.newxyz(dmp,1000) replace xyz in Atoms with xyz of snapshot N
|
|
|
|
newxyz assumes id,x,y,z are defined in both data and dump files
|
|
also replaces ix,iy,iz if they are defined
|
|
|
|
index,time,flag = d.iterator(0/1) loop over single data file snapshot
|
|
time,box,atoms,bonds,tris,lines = d.viz(index) return list of viz objects
|
|
|
|
iterator() and viz() are compatible with equivalent dump calls
|
|
iterator() called with arg = 0 first time, with arg = 1 on subsequent calls
|
|
index = timestep index within dump object (only 0 for data file)
|
|
time = timestep value (only 0 for data file)
|
|
flag = -1 when iteration is done, 1 otherwise
|
|
viz() returns info for specified timestep index (must be 0)
|
|
time = 0
|
|
box = [xlo,ylo,zlo,xhi,yhi,zhi]
|
|
atoms = id,type,x,y,z for each atom as 2d array
|
|
bonds = id,type,x1,y1,z1,x2,y2,z2,t1,t2 for each bond as 2d array
|
|
NULL if bonds do not exist
|
|
tris = NULL
|
|
lines = NULL
|
|
|
|
d.write("data.new") write a LAMMPS data file
|
|
"""
|
|
|
|
# History
|
|
# 8/05, Steve Plimpton (SNL): original version
|
|
# 11/07, added triclinic box support
|
|
|
|
# ToDo list
|
|
|
|
# Variables
|
|
# title = 1st line of data file
|
|
# names = dictionary with atom attributes as keys, col #s as values
|
|
# headers = dictionary with header name as key, value or tuple as values
|
|
# sections = dictionary with section name as key, array of lines as values
|
|
# nselect = 1 = # of snapshots
|
|
|
|
# Imports and external programs
|
|
|
|
from os import popen
|
|
|
|
try: tmp = PIZZA_GUNZIP
|
|
except: PIZZA_GUNZIP = "gunzip"
|
|
|
|
# Class definition
|
|
|
|
class data:
|
|
|
|
# --------------------------------------------------------------------
|
|
|
|
def __init__(self,*list):
|
|
self.nselect = 1
|
|
|
|
if len(list) == 0:
|
|
self.title = "LAMMPS data file"
|
|
self.names = {}
|
|
self.headers = {}
|
|
self.sections = {}
|
|
return
|
|
|
|
file = list[0]
|
|
if file[-3:] == ".gz": f = popen("%s -c %s" % (PIZZA_GUNZIP,file),'r')
|
|
else: f = open(file)
|
|
|
|
self.title = f.readline()
|
|
self.names = {}
|
|
|
|
headers = {}
|
|
while 1:
|
|
line = f.readline()
|
|
line = line.strip()
|
|
if len(line) == 0:
|
|
continue
|
|
found = 0
|
|
for keyword in hkeywords:
|
|
if line.find(keyword) >= 0:
|
|
found = 1
|
|
words = line.split()
|
|
if keyword == "xlo xhi" or keyword == "ylo yhi" or \
|
|
keyword == "zlo zhi":
|
|
headers[keyword] = (float(words[0]),float(words[1]))
|
|
elif keyword == "xy xz yz":
|
|
headers[keyword] = \
|
|
(float(words[0]),float(words[1]),float(words[2]))
|
|
else:
|
|
headers[keyword] = int(words[0])
|
|
if not found:
|
|
break
|
|
|
|
sections = {}
|
|
while 1:
|
|
found = 0
|
|
for pair in skeywords:
|
|
keyword,length = pair[0],pair[1]
|
|
if keyword == line:
|
|
found = 1
|
|
if not headers.has_key(length):
|
|
raise StandardError, \
|
|
"data section %s has no matching header value" % line
|
|
f.readline()
|
|
list = []
|
|
for i in xrange(headers[length]): list.append(f.readline())
|
|
sections[keyword] = list
|
|
if not found:
|
|
raise StandardError,"invalid section %s in data file" % line
|
|
f.readline()
|
|
line = f.readline()
|
|
if not line:
|
|
break
|
|
line = line.strip()
|
|
|
|
f.close()
|
|
self.headers = headers
|
|
self.sections = sections
|
|
|
|
# --------------------------------------------------------------------
|
|
# assign names to atom columns
|
|
|
|
def map(self,*pairs):
|
|
if len(pairs) % 2 != 0:
|
|
raise StandardError, "data map() requires pairs of mappings"
|
|
for i in range(0,len(pairs),2):
|
|
j = i + 1
|
|
self.names[pairs[j]] = pairs[i]-1
|
|
|
|
# --------------------------------------------------------------------
|
|
# extract info from data file fields
|
|
|
|
def get(self,*list):
|
|
if len(list) == 1:
|
|
field = list[0]
|
|
array = []
|
|
lines = self.sections[field]
|
|
for line in lines:
|
|
words = line.split()
|
|
values = map(float,words)
|
|
array.append(values)
|
|
return array
|
|
elif len(list) == 2:
|
|
field = list[0]
|
|
n = list[1] - 1
|
|
vec = []
|
|
lines = self.sections[field]
|
|
for line in lines:
|
|
words = line.split()
|
|
vec.append(float(words[n]))
|
|
return vec
|
|
else:
|
|
raise StandardError, "invalid arguments for data.get()"
|
|
|
|
# --------------------------------------------------------------------
|
|
# reorder columns in a data file field
|
|
|
|
def reorder(self,name,*order):
|
|
n = len(order)
|
|
natoms = len(self.sections[name])
|
|
oldlines = self.sections[name]
|
|
newlines = natoms*[""]
|
|
for index in order:
|
|
for i in xrange(len(newlines)):
|
|
words = oldlines[i].split()
|
|
newlines[i] += words[index-1] + " "
|
|
for i in xrange(len(newlines)):
|
|
newlines[i] += "\n"
|
|
self.sections[name] = newlines
|
|
|
|
# --------------------------------------------------------------------
|
|
# replace a column of named section with vector of values
|
|
|
|
def replace(self,name,icol,vector):
|
|
lines = self.sections[name]
|
|
newlines = []
|
|
j = icol - 1
|
|
for i in xrange(len(lines)):
|
|
line = lines[i]
|
|
words = line.split()
|
|
words[j] = str(vector[i])
|
|
newline = ' '.join(words) + '\n'
|
|
newlines.append(newline)
|
|
self.sections[name] = newlines
|
|
|
|
# --------------------------------------------------------------------
|
|
# replace x,y,z in Atoms with x,y,z values from snapshot ntime of dump object
|
|
# assumes id,x,y,z are defined in both data and dump files
|
|
# also replaces ix,iy,iz if they are defined
|
|
|
|
def newxyz(self,dm,ntime):
|
|
nsnap = dm.findtime(ntime)
|
|
|
|
dm.sort(ntime)
|
|
x,y,z = dm.vecs(ntime,"x","y","z")
|
|
|
|
self.replace("Atoms",self.names['x']+1,x)
|
|
self.replace("Atoms",self.names['y']+1,y)
|
|
self.replace("Atoms",self.names['z']+1,z)
|
|
|
|
if dm.names.has_key("ix") and self.names.has_key("ix"):
|
|
ix,iy,iz = dm.vecs(ntime,"ix","iy","iz")
|
|
self.replace("Atoms",self.names['ix']+1,ix)
|
|
self.replace("Atoms",self.names['iy']+1,iy)
|
|
self.replace("Atoms",self.names['iz']+1,iz)
|
|
|
|
# --------------------------------------------------------------------
|
|
# delete header value or section from data file
|
|
|
|
def delete(self,keyword):
|
|
|
|
if self.headers.has_key(keyword): del self.headers[keyword]
|
|
elif self.sections.has_key(keyword): del self.sections[keyword]
|
|
else: raise StandardError, "keyword not found in data object"
|
|
|
|
# --------------------------------------------------------------------
|
|
# write out a LAMMPS data file
|
|
|
|
def write(self,file):
|
|
f = open(file,"w")
|
|
print >>f,self.title
|
|
for keyword in hkeywords:
|
|
if self.headers.has_key(keyword):
|
|
if keyword == "xlo xhi" or keyword == "ylo yhi" or \
|
|
keyword == "zlo zhi":
|
|
pair = self.headers[keyword]
|
|
print >>f,pair[0],pair[1],keyword
|
|
elif keyword == "xy xz yz":
|
|
triple = self.headers[keyword]
|
|
print >>f,triple[0],triple[1],triple[2],keyword
|
|
else:
|
|
print >>f,self.headers[keyword],keyword
|
|
for pair in skeywords:
|
|
keyword = pair[0]
|
|
if self.sections.has_key(keyword):
|
|
print >>f,"\n%s\n" % keyword
|
|
for line in self.sections[keyword]:
|
|
print >>f,line,
|
|
f.close()
|
|
|
|
# --------------------------------------------------------------------
|
|
# iterator called from other tools
|
|
|
|
def iterator(self,flag):
|
|
if flag == 0: return 0,0,1
|
|
return 0,0,-1
|
|
|
|
# --------------------------------------------------------------------
|
|
# time query from other tools
|
|
|
|
def findtime(self,n):
|
|
if n == 0: return 0
|
|
raise StandardError, "no step %d exists" % (n)
|
|
|
|
# --------------------------------------------------------------------
|
|
# return list of atoms and bonds to viz for data object
|
|
|
|
def viz(self,isnap):
|
|
if isnap: raise StandardError, "cannot call data.viz() with isnap != 0"
|
|
|
|
id = self.names["id"]
|
|
type = self.names["type"]
|
|
x = self.names["x"]
|
|
y = self.names["y"]
|
|
z = self.names["z"]
|
|
|
|
xlohi = self.headers["xlo xhi"]
|
|
ylohi = self.headers["ylo yhi"]
|
|
zlohi = self.headers["zlo zhi"]
|
|
box = [xlohi[0],ylohi[0],zlohi[0],xlohi[1],ylohi[1],zlohi[1]]
|
|
|
|
# create atom list needed by viz from id,type,x,y,z
|
|
|
|
atoms = []
|
|
atomlines = self.sections["Atoms"]
|
|
for line in atomlines:
|
|
words = line.split()
|
|
atoms.append([int(words[id]),int(words[type]),
|
|
float(words[x]),float(words[y]),float(words[z])])
|
|
|
|
# create list of current bond coords from list of bonds
|
|
# assumes atoms are sorted so can lookup up the 2 atoms in each bond
|
|
|
|
bonds = []
|
|
if self.sections.has_key("Bonds"):
|
|
bondlines = self.sections["Bonds"]
|
|
for line in bondlines:
|
|
words = line.split()
|
|
bid,btype = int(words[0]),int(words[1])
|
|
atom1,atom2 = int(words[2]),int(words[3])
|
|
atom1words = atomlines[atom1-1].split()
|
|
atom2words = atomlines[atom2-1].split()
|
|
bonds.append([bid,btype,
|
|
float(atom1words[x]),float(atom1words[y]),
|
|
float(atom1words[z]),
|
|
float(atom2words[x]),float(atom2words[y]),
|
|
float(atom2words[z]),
|
|
float(atom1words[type]),float(atom2words[type])])
|
|
|
|
tris = []
|
|
lines = []
|
|
return 0,box,atoms,bonds,tris,lines
|
|
|
|
# --------------------------------------------------------------------
|
|
# return box size
|
|
|
|
def maxbox(self):
|
|
xlohi = self.headers["xlo xhi"]
|
|
ylohi = self.headers["ylo yhi"]
|
|
zlohi = self.headers["zlo zhi"]
|
|
return [xlohi[0],ylohi[0],zlohi[0],xlohi[1],ylohi[1],zlohi[1]]
|
|
|
|
# --------------------------------------------------------------------
|
|
# return number of atom types
|
|
|
|
def maxtype(self):
|
|
return self.headers["atom types"]
|
|
|
|
# --------------------------------------------------------------------
|
|
# data file keywords, both header and main sections
|
|
|
|
hkeywords = ["atoms","ellipsoids","lines","triangles","bodies",
|
|
"bonds","angles","dihedrals","impropers","pitorsions",
|
|
"atom types","bond types","angle types","dihedral types",
|
|
"improper types","pitorsion types",
|
|
"xlo xhi","ylo yhi","zlo zhi","xy xz yz"]
|
|
|
|
skeywords = [["Masses","atom types"],
|
|
["Atoms","atoms"],["Ellipsoids","ellipsoids"],
|
|
["Lines","lines"],["Triangles","triangles"],["Bodies","bodies"],
|
|
["Bonds","bonds"],
|
|
["Angles","angles"],
|
|
["Dihedrals","dihedrals"],
|
|
["Impropers","impropers"],
|
|
["PiTorsions","pitorsions"],
|
|
["Velocities","atoms"],
|
|
["Pair Coeffs","atom types"],
|
|
["Bond Coeffs","bond types"],["Angle Coeffs","angle types"],
|
|
["Dihedral Coeffs","dihedral types"],
|
|
["Improper Coeffs","improper types"],
|
|
["PiTorsion Coeffs","pitorsion types"],
|
|
["BondBond Coeffs","angle types"],
|
|
["BondAngle Coeffs","angle types"],
|
|
["MiddleBondTorsion Coeffs","dihedral types"],
|
|
["EndBondTorsion Coeffs","dihedral types"],
|
|
["AngleTorsion Coeffs","dihedral types"],
|
|
["AngleAngleTorsion Coeffs","dihedral types"],
|
|
["BondBond13 Coeffs","dihedral types"],
|
|
["AngleAngle Coeffs","improper types"],
|
|
["Molecules","atoms"]]
|