diff --git a/.gitignore b/.gitignore index 0d20b64..8d35cb3 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ +__pycache__ *.pyc diff --git a/pizza/__init__.py b/pizza/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/log.py b/pizza/log.py similarity index 74% rename from src/log.py rename to pizza/log.py index aeca1d8..f0735e9 100644 --- a/src/log.py +++ b/pizza/log.py @@ -5,6 +5,7 @@ # 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. +from __future__ import print_function # log tool @@ -28,7 +29,7 @@ nlen = l.nlen length of each vectors names = l.names list of vector names t,pe,... = l.get("Time","KE",...) return one or more vectors of values -l.write("file.txt") write all vectors to a file +l.write("file.txt") write all vectors to a file l.write("file.txt","Time","PE",...) write listed vectors to a file get and write allow abbreviated (uniquely) vector names @@ -53,10 +54,7 @@ # Imports and external programs import sys, re, glob -from os import popen - -try: tmp = PIZZA_GUNZIP -except: PIZZA_GUNZIP = "gunzip" +import gzip # Class definition @@ -76,14 +74,14 @@ def __init__(self,*list): self.flist = [] for word in words: self.flist += glob.glob(word) if len(self.flist) == 0 and len(list) == 1: - raise StandardError,"no log file specified" + raise Exception("no log file specified") if len(list) == 1: self.increment = 0 self.read_all() else: if len(self.flist) > 1: - raise StandardError,"can only incrementally read one log file" + raise Exception("can only incrementally read one log file") self.increment = 1 self.eof = 0 @@ -92,24 +90,24 @@ def __init__(self,*list): def read_all(self): self.read_header(self.flist[0]) - if self.nvec == 0: raise StandardError,"log file has no values" + if self.nvec == 0: raise Exception("log file has no values") # read all files for file in self.flist: self.read_one(file) - print + print() # sort entries by timestep, cull duplicates - self.data.sort(self.compare) + self.data.sort(key=lambda x: x[0]) self.cull() self.nlen = len(self.data) - print "read %d log entries" % self.nlen + print("read %d log entries" % self.nlen) # -------------------------------------------------------------------- def next(self): - if not self.increment: raise StandardError,"cannot read incrementally" + if not self.increment: raise Exception("cannot read incrementally") if self.nvec == 0: try: open(self.flist[0],'r') @@ -124,28 +122,28 @@ def next(self): def get(self,*keys): if len(keys) == 0: - raise StandardError, "no log vectors specified" + raise Exception("no log vectors specified") - map = [] + m = [] for key in keys: - if self.ptr.has_key(key): - map.append(self.ptr[key]) + if key in self.ptr: + m.append(self.ptr[key]) else: count = 0 for i in range(self.nvec): - if self.names[i].find(key) == 0: - count += 1 - index = i + if self.names[i].find(key) == 0: + count += 1 + index = i if count == 1: - map.append(index) + m.append(index) else: - raise StandardError, "unique log vector %s not found" % key + raise Exception("unique log vector %s not found" % key) vecs = [] for i in range(len(keys)): vecs.append(self.nlen * [0]) - for j in xrange(self.nlen): - vecs[i][j] = self.data[j][map[i]] + for j in range(self.nlen): + vecs[i][j] = self.data[j][m[i]] if len(keys) == 1: return vecs[0] else: return vecs @@ -154,39 +152,26 @@ def get(self,*keys): def write(self,filename,*keys): if len(keys): - map = [] + m = [] for key in keys: - if self.ptr.has_key(key): - map.append(self.ptr[key]) + if key in self.ptr: + m.append(self.ptr[key]) else: count = 0 for i in range(self.nvec): - if self.names[i].find(key) == 0: - count += 1 - index = i + if self.names[i].find(key) == 0: + count += 1 + index = i if count == 1: - map.append(index) + m.append(index) else: - raise StandardError, "unique log vector %s not found" % key + raise Exception("unique log vector %s not found" % key) else: - map = range(self.nvec) + m = list(range(self.nvec)) - f = open(filename,"w") - for i in xrange(self.nlen): - for j in xrange(len(map)): - print >>f,self.data[i][map[j]], - print >>f - f.close() - - # -------------------------------------------------------------------- - - def compare(self,a,b): - if a[0] < b[0]: - return -1 - elif a[0] > b[0]: - return 1 - else: - return 0 + with open(filename,"w") as f: + for i in range(self.nlen): + print(*[self.data[i][x] for x in m], file=f) # -------------------------------------------------------------------- @@ -198,14 +183,16 @@ def cull(self): # -------------------------------------------------------------------- - def read_header(self,file): + def read_header(self,filename): str_multi = "----- Step" str_one = "Step " - if file[-3:] == ".gz": - txt = popen("%s -c %s" % (PIZZA_GUNZIP,file),'r').read() + if filename[-3:] == ".gz": + with gzip.open(filename, 'rt') as f: + txt = f.read() else: - txt = open(file).read() + with open(filename, 'rt') as f: + txt = f.read() if txt.find(str_multi) >= 0: self.firststr = str_multi @@ -226,37 +213,36 @@ def read_header(self,file): keywords.insert(0,"Step") i = 0 for keyword in keywords: - self.names.append(keyword) + self.names.append(keyword) self.ptr[keyword] = i i += 1 - else: s1 = txt.find(self.firststr) s2 = txt.find("\n",s1) line = txt[s1:s2] words = line.split() for i in range(len(words)): - self.names.append(words[i]) + self.names.append(words[i]) self.ptr[words[i]] = i self.nvec = len(self.names) # -------------------------------------------------------------------- - def read_one(self,*list): + def read_one(self,*args): # if 2nd arg exists set file ptr to that value # read entire (rest of) file into txt - file = list[0] - if file[-3:] == ".gz": - f = popen("%s -c %s" % (PIZZA_GUNZIP,file),'rb') + filename = args[0] + if filename[-3:] == ".gz": + f = gzip.open(filename, 'rt') else: - f = open(file,'rb') + f = open(filename,'rt') - if len(list) == 2: f.seek(list[1]) + if len(args) == 2: f.seek(args[1]) txt = f.read() - if file[-3:] == ".gz": eof = 0 + if filename[-3:] == ".gz": eof = 0 else: eof = f.tell() f.close() @@ -275,36 +261,36 @@ def read_one(self,*list): if s1 >= 0 and s2 >= 0 and s1 < s2: # found s1,s2 with s1 before s2 if self.style == 2: - s1 = txt.find("\n",s1) + 1 + s1 = txt.find("\n",s1) + 1 elif s1 >= 0 and s2 >= 0 and s2 < s1: # found s1,s2 with s2 before s1 s1 = 0 elif s1 == -1 and s2 >= 0: # found s2, but no s1 - last = 1 + last = 1 s1 = 0 elif s1 >= 0 and s2 == -1: # found s1, but no s2 last = 1 if self.style == 1: s2 = txt.rfind("\n--",s1) + 1 else: - s1 = txt.find("\n",s1) + 1 + s1 = txt.find("\n",s1) + 1 s2 = txt.rfind("\n",s1) + 1 - eof -= len(txt) - s2 + eof -= len(txt) - s2 elif s1 == -1 and s2 == -1: # found neither # could be end-of-file section - # or entire read was one chunk + # or entire read was one chunk if txt.find("Loop time of",start) == start: # end of file, so exit - eof -= len(txt) - start # reset eof to "Loop" - break + eof -= len(txt) - start # reset eof to "Loop" + break - last = 1 # entire read is a chunk + last = 1 # entire read is a chunk s1 = 0 if self.style == 1: s2 = txt.rfind("\n--",s1) + 1 else: s2 = txt.rfind("\n",s1) + 1 - eof -= len(txt) - s2 - if s1 == s2: break + eof -= len(txt) - s2 + if s1 == s2: break chunk = txt[s1:s2-1] start = s2 @@ -320,17 +306,18 @@ def read_one(self,*list): word1 = [re.search(pat1,section).group(1)] word2 = re.findall(pat2,section) words = word1 + word2 - self.data.append(map(float,words)) - + self.data.append(list(map(float,words))) else: lines = chunk.split("\n") + alpha = re.compile('[a-df-zA-DF-Z]') # except e or E for floating-point numbers for line in lines: + if alpha.search(line): continue words = line.split() - self.data.append(map(float,words)) + self.data.append(list(map(float,words))) # print last timestep of chunk - print int(self.data[len(self.data)-1][0]), + print(int(self.data[len(self.data)-1][0]), end='') sys.stdout.flush() return eof diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..83dc87f --- /dev/null +++ b/setup.py @@ -0,0 +1,15 @@ +#!/usr/bin/env python + +from distutils.core import setup + +setup(name='Pizza.py', + version='1.0', + description='Pizza.py is a loosely integrated collection of tools written in Python, many of which provide pre- ' + 'and post-processing capability for the LAMMPS molecular dynamics package. There are tools to create ' + 'input files, convert between file formats, process log and dump files, create plots, and visualize ' + 'and animate simulation snapshots.', + author='Steve Plimpton', + author_email='sjplimp@sandia.gov', + url='http://pizza.sandia.gov/', + packages=['pizza'], + ) diff --git a/tests/test_log.py b/tests/test_log.py new file mode 100644 index 0000000..43a2c3a --- /dev/null +++ b/tests/test_log.py @@ -0,0 +1,44 @@ +import unittest +import csv +import os +from pizza.log import log + +TESTS_DIR=os.path.dirname(os.path.abspath(__file__)) + +class TestLog(unittest.TestCase): + EXAMPLE_FILE = os.path.join(TESTS_DIR, "../examples/files/log.obstacle") + + def test_read_log(self): + lg = log(TestLog.EXAMPLE_FILE) + self.assertEqual(lg.nvec, 7) + self.assertEqual(lg.nlen, 26) + self.assertListEqual(lg.names, ["Step", "Temperature", "E_pair", "E_bond", "E_total", "Pressure", "Volume"]) + + def test_log_get(self): + lg = log(TestLog.EXAMPLE_FILE) + time, temp, press = lg.get("Step", "Temp", "Press") + self.assertEqual(len(time), 26) + self.assertEqual(len(temp), 26) + self.assertEqual(len(press), 26) + + def test_write_all(self): + tmp_file = "/tmp/tmp.log" + lg = log(TestLog.EXAMPLE_FILE) + lg.write(tmp_file) + self.assertTrue(os.path.exists(tmp_file)) + + with open(tmp_file, 'rt') as csvfile: + lines = list(csv.reader(csvfile, delimiter=' ')) + self.assertEqual(len(lines), 26) + self.assertEqual(len(lines[0]), 7) + + def test_write_columns(self): + tmp_file = "/tmp/tmp.log.two" + lg = log(TestLog.EXAMPLE_FILE) + lg.write(tmp_file, "Step", "E_pair") + self.assertTrue(os.path.exists(tmp_file)) + + with open(tmp_file, 'rt') as csvfile: + lines = list(csv.reader(csvfile, delimiter=' ')) + self.assertEqual(len(lines), 26) + self.assertEqual(len(lines[0]), 2)