diff --git a/python/lammps.py b/python/lammps.py index daf526f0e903f29cf1de4ab299dceded40e694d4..bc940c004c699c64f387f957b6aabe6ecd39abee 100644 --- a/python/lammps.py +++ b/python/lammps.py @@ -436,6 +436,42 @@ class Atom2D(Atom): self.lmp.eval("fy[%d]" % self.index)) +def get_thermo_data(output): + """ traverse output of runs and extract thermo data columns """ + if isinstance(output, str): + lines = output.splitlines() + else: + lines = output + + runs = [] + columns = [] + in_run = False + + for line in lines: + if line.startswith("Memory usage per processor"): + in_run = True + elif in_run and len(columns) == 0: + # first line after memory usage are column names + columns = line.split() + + current_run = {} + + for col in columns: + current_run[col] = [] + + elif line.startswith("Loop time of "): + in_run = False + columns = None + thermo_data = namedtuple('ThermoData', list(current_run.keys()))(*list(current_run.values())) + r = {'thermo' : thermo_data } + runs.append(namedtuple('Run', list(r.keys()))(*list(r.values()))) + elif in_run and len(columns) > 0: + values = [float(x) for x in line.split()] + + for i, col in enumerate(columns): + current_run[col].append(values[i]) + return runs + class PyLammps(object): """ More Python-like wrapper for LAMMPS (e.g., for iPython) @@ -455,6 +491,7 @@ class PyLammps(object): self.lmp = lammps(name=name,cmdargs=cmdargs,ptr=None,comm=comm) print("LAMMPS output is captured by PyLammps wrapper") self._cmd_history = [] + self.runs = [] def __del__(self): if self.lmp: self.lmp.close() @@ -480,6 +517,17 @@ class PyLammps(object): self.lmp.command(cmd) self._cmd_history.append(cmd) + def run(self, *args, **kwargs): + output = self.__getattr__('run')(*args, **kwargs) + self.runs += get_thermo_data(output) + return output + + @property + def last_run(self): + if len(self.runs) > 0: + return self.runs[-1] + return None + @property def atoms(self): return AtomList(self)