From f92089298d869a7892beed62162685810c03284d Mon Sep 17 00:00:00 2001
From: Richard Berger <richard.berger@temple.edu>
Date: Thu, 25 Feb 2021 16:56:19 -0500
Subject: [PATCH] Fix bug in AvgChunkReader and add docs

---
 doc/src/Python_formats.rst        | 11 +++++++++++
 doc/src/Python_head.rst           |  1 +
 python/lammps/formats.py          | 25 ++++++++++++++++++++++++-
 unittest/python/python-formats.py |  2 ++
 4 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 doc/src/Python_formats.rst

diff --git a/doc/src/Python_formats.rst b/doc/src/Python_formats.rst
new file mode 100644
index 0000000000..b9ffdba632
--- /dev/null
+++ b/doc/src/Python_formats.rst
@@ -0,0 +1,11 @@
+Output Readers
+==============
+
+.. py:module:: lammps.formats
+
+The Python package contains the :py:mod:`lammps.formats` module, which
+provides classes to post-process some of the output files generated by LAMMPS.
+
+.. automodule:: lammps.formats
+   :members:
+   :noindex:
diff --git a/doc/src/Python_head.rst b/doc/src/Python_head.rst
index 09071171ad..3e84ed4506 100644
--- a/doc/src/Python_head.rst
+++ b/doc/src/Python_head.rst
@@ -13,6 +13,7 @@ together.
    Python_module
    Python_ext
    Python_call
+   Python_formats
    Python_examples
    Python_error
    Python_trouble
diff --git a/python/lammps/formats.py b/python/lammps/formats.py
index 6cbff321a2..3ebc2e1bac 100644
--- a/python/lammps/formats.py
+++ b/python/lammps/formats.py
@@ -12,13 +12,25 @@
 # -------------------------------------------------------------------------
 
 ################################################################################
-# LAMMPS data formats
+# LAMMPS output formats
 # Written by Richard Berger <richard.berger@temple.edu>
 ################################################################################
 
 import re
 
 class LogFile:
+  """Reads LAMMPS log files and extracts the thermo information
+
+  It supports both the default thermo output style (including custom) and multi.
+
+  :param filename: path to log file
+  :type  filename: str
+
+  :ivar runs: List of LAMMPS runs in log file. Each run is a dictionary with
+              thermo fields as keys, storing the values over time
+  :ivar errors: List of error lines in log file
+  """
+
   STYLE_DEFAULT = 0
   STYLE_MULTI   = 1
 
@@ -73,7 +85,17 @@ class LogFile:
                         else:
                             current_run[k].append(float(v))
 
+
 class AvgChunkFile:
+  """Reads files generated by fix ave/chunk
+
+  :param filename: path to ave/chunk file
+  :type  filename: str
+
+  :ivar timesteps: List of timesteps stored in file
+  :ivar total_count: total count over time
+  :ivar chunks: List of chunks. Each chunk is a dictionary containing its ID, the coordinates, and the averaged quantities
+  """
   def __init__(self, filename):
     with open(filename, 'rt') as f:
       timestep = None
@@ -121,6 +143,7 @@ class AvgChunkFile:
           total_count = float(parts[2])
 
           self.timesteps.append(timestep)
+          self.total_count.append(total_count)
 
           for i in range(num_chunks):
             self.chunks.append({
diff --git a/unittest/python/python-formats.py b/unittest/python/python-formats.py
index a5d4e84cfa..ba68e50f4b 100644
--- a/unittest/python/python-formats.py
+++ b/unittest/python/python-formats.py
@@ -73,7 +73,9 @@ class AvgChunkFiles(unittest.TestCase):
         self.assertEqual(cfile.timesteps, list(range(10000, 110000, 5000)))
 
         ntimesteps = len(cfile.timesteps)
+        ntotal_count = len(cfile.total_count)
         nchunks = len(cfile.chunks)
+        self.assertEqual(ntimesteps, ntotal_count)
         self.assertEqual(nchunks, 20)
 
         for i in range(1, nchunks+1):