lammps/tools/eff/lmp2data.py

145 lines
4.0 KiB
Python

#!/usr/local/bin/python-2.5/bin/python
Info="""
Module name: lmp2data.py
Author: (c) Andres Jaramillo-Botero
California Institute of Technology
ajaramil@wag.caltech.edu
Project: pEFF
Version: August 2009
Extracts the electron radii from a lammps trajectory dump of style custom:
dump 1 all custom period dump_file id type x y z spin radius ...
NOTE: The radius must be the i'th column per trajectory entry in the dump file
"""
# import essentials:
import sys, os
from math import log10
from shutil import rmtree
from getopt import gnu_getopt as getopt
import numpy
def printHelp():
print Info
print "Usage: python lmp2data.py test.lammpstrj\n"
return
def makeradii(infile,outfile,column,flag_all):
print "Reading %s ... [WAIT]"%infile,
fin = open(infile,'r')
lines = fin.xreadlines()
print 7*"\b"+"[DONE]"
frame=0
radii=[]
# grep the number of frames and atoms/frame
os.system("grep TIMESTEP %s | wc -l > frames; grep -m 1 -A 1 ATOMS %s > atoms; grep -m 1 \"ITEM: ATOMS\" %s > params"%(infile,infile,infile))
tmp=open("frames",'r')
frames=int(tmp.readline().split()[0])
tmp.close()
tmp=open("atoms",'r')
atoms=int(tmp.readlines()[1].split()[0])
tmp.close()
tmp=open("params",'r')
ids=tmp.readline().split()[2:]
os.system("rm -rf frames atoms params")
arry=numpy.zeros((atoms,frames),dtype=str)
framecnt=0
header=9
ecount=0
if flag_all==True: atom_type="nuclei and electron"
else: atom_type="electron"
print "Extracting %s %s per frame from %s ... "%(atom_type,ids[column],infile),
for i,line in enumerate(lines):
lo=(atoms+header)*framecnt+header
hi=lo+atoms
if (i<lo):
continue
elif (i >= lo) and (i < hi):
lparse=line.split()
id=int(lparse[0])
# r=float(lparse[column-1])
r=lparse[column]
# if (float(r)!=0):
arry[id-1][framecnt]=r
print arry[id-1][framecnt],r,raw_input()
if (float(r)!=0) and (framecnt==0): ecount+=1
# else: arry[id-1][framecnt]=r
if (i==lo+1):
sys.stdout.write("%d/%d%s"%(framecnt+1,frames,(int(log10(framecnt+1))+3+int(log10(frames)))*"\b"))
sys.stdout.flush()
if (i == hi+1):
framecnt+=1
print
if outfile=="":
outfile=infile+'.%s'%(ids[column])
fout=open(outfile,'w')
else: fout=open(outfile,'w')
print "Writing %s/frame table to %s ... "%(ids[column],outfile),
sys.stdout.flush()
for i in range(frames):
fout.writelines('\tF'+str(i))
fout.writelines("\n")
e=1
for a in range(atoms):
if flag_all==True:
sys.stdout.write("%d/%d%s"%(a+1,atoms,(int(log10(a+1))+int(log10(atoms))+3)*"\b"))
sys.stdout.flush()
fout.writelines("%d\t"%(a+1))
for f in range(frames):
fout.writelines("%s\t"%(arry[a][f]))
fout.writelines("\n")
else:
if arry[a][0] == 0.0:
continue
else:
sys.stdout.write("%d/%d%s"%(e,ecount,(int(log10(e))+int(log10(ecount))+3)*"\b"))
sys.stdout.flush()
e+=1
fout.writelines("%d\t"%(a+1))
for f in range(frames):
fout.writelines("%s\t"%(arry[a][f]))
fout.writelines("\n")
print
print "DONE .... GOODBYE !!"
fout.close()
fin.close()
if __name__ == '__main__':
# set defaults
outfile = ""
flag_all = False
column=6 # default = radius
# check for input:
opts, argv = getopt(sys.argv[1:], 'c:o:ha')
# if no input, print help and exit
if len(argv) != 1:
printHelp()
sys.exit(1)
else:
infile=argv[0]
# read options
for opt, arg in opts:
if opt == '-h': # -h: print help
printHelp()
if opt == '-o': # output file name
outfile=arg
if opt == '-a': # all nuclii+electrons
flag_all=True
if opt == '-c': # select column from lammpstrj file to tabulate
column=int(arg)
makeradii(infile,outfile,column,flag_all)