Refactoring

.gitignore

@@ -1 +1,2 @@
+__pycache__
 *.pyc

src/log.py → 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

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'],
+     )

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)