add unit test for yaml style thermo output and updated logfile class

python/lammps/formats.py

@@ -14,14 +14,19 @@
 ################################################################################
 # LAMMPS output formats
 # Written by Richard Berger <richard.berger@temple.edu>
+# and Axel Kohlmeyer <akohlmey@gmail.com>
 ################################################################################
 
-import re
+import re, yaml
+try:
+  from yaml import CSafeLoader as Loader, CSafeDumper as Dumper
+except ImportError:
+    from yaml import SafeLoader as Loader, SafeDumper as Dumper
 
 class LogFile:
   """Reads LAMMPS log files and extracts the thermo information
 
-  It supports both the default thermo output style (including custom) and multi.
+  It supports the line, multi, and yaml thermo output styles.
 
   :param filename: path to log file
   :type  filename: str
@@ -33,11 +38,13 @@ class LogFile:
 
   STYLE_DEFAULT = 0
   STYLE_MULTI   = 1
+  STYLE_YAML    = 2
 
   def __init__(self, filename):
     alpha = re.compile(r'[a-df-zA-DF-Z]') # except e or E for floating-point numbers
     kvpairs = re.compile(r'([a-zA-Z_0-9]+)\s+=\s*([0-9\.eE\-]+)')
     style = LogFile.STYLE_DEFAULT
+    yamllog = ""
     self.runs = []
     self.errors = []
     with open(filename, 'rt') as f:
@@ -54,6 +61,24 @@ class LogFile:
                 current_run = {}
                 for k in keys:
                     current_run[k] = []
+
+            elif re.match(r'^(keywords:.*$|data:$|---$|  - \[.*\]$)', line):
+                style = LogFile.STYLE_YAML
+                yamllog += line;
+                current_run = {}
+
+            elif re.match(r'^\.\.\.$', line):
+                thermo = yaml.load(yamllog, Loader=Loader)
+                for k in thermo['keywords']:
+                    current_run[k] = []
+                for step in thermo['data']:
+                    icol = 0
+                    for k in thermo['keywords']:
+                        current_run[k].append(step[icol])
+                        icol += 1
+                self.runs.append(current_run)
+                yamllog = ""
+
             elif re.match(r'^------* Step ', line):
                 if not in_thermo:
                    current_run = {'Step': [], 'CPU': []}
@@ -65,28 +90,29 @@ class LogFile:
                 cpu  = float(str_cpu.split('=')[1].split()[0])
                 current_run["Step"].append(step)
                 current_run["CPU"].append(cpu)
+
             elif line.startswith('Loop time of'):
                 in_thermo = False
-                self.runs.append(current_run)
+                if style != LogFile.STYLE_YAML:
+                    self.runs.append(current_run)
+
             elif in_thermo and in_data_section:
                 if style == LogFile.STYLE_DEFAULT:
                     if alpha.search(line):
                         continue
-
                     for k, v in zip(keys, map(float, line.split())):
                         current_run[k].append(v)
+
                 elif style == LogFile.STYLE_MULTI:
                     if '=' not in line:
                         in_data_section = False
                         continue
-
                     for k,v in kvpairs.findall(line):
                         if k not in current_run:
                             current_run[k] = [float(v)]
                         else:
                             current_run[k].append(float(v))
 
-
 class AvgChunkFile:
   """Reads files generated by fix ave/chunk