refine tests

package/CHANGELOG

@@ -29,7 +29,7 @@ Fixes
   * Fixes hydrogenbonds tutorial path to point to hbonds (Issue #4285, PR #4286)
 
 Enhancements
-  * Added parsing of arbitrary columns of the LAMMPS dump parser.
+  * Added parsing of arbitrary columns of the LAMMPS dump parser. (Issue #3504)
   * Adds external sidebar links (Issue #4296)
   * Updated lib.qcprot.CalcRMSDRotationalMatrix to accept either float32 or float64
     inputs (PR #4273, part of #3927)

package/MDAnalysis/coordinates/LAMMPS.py

@@ -477,28 +477,24 @@ class DumpReader(base.ReaderBase):
     length unit (Å), and angles are in degrees.
 
     By using the keyword `additional_columns`, you can specify arbitrary data
-    to be read alongside the coordinates. If specified, the keyword expects a
-    list of the names of the columns that you want to have read. The results
-    of the parsing are saved to the time step :attr:`Timestep.data` dictionary
-    alongside the name of the data column. For instance, if you have time-dependent
-    charges saved in a LAMMPS dump such as
+    to be read. The keyword expects a list of the names of the columns or `True` to read
+    all additional columns. The results are saved to :attr:`Timestep.data`. 
+    For example, if your LAMMPS dump looks like this
 
     .. code-block::
 
         ITEM: ATOMS id x y z q l
         1 2.84 8.17 -25 0.00258855 1.1
         2 7.1 8.17 -25 6.91952e-05 1.2
 
-    Then you may parse the additional columns `q` and `l` via.
+    Then you may parse the additional columns `q` and `l` via:
 
     .. code-block:: python
 
         u = mda.Universe('path_to_data', 'path_to_lammpsdump',
                          additional_columns=['q', 'l'])
 
-    The additional data is then available for each time step via
-    (as the value of the :attr:`Timestep.data` dictionary, sorted by the ids of the
-    atoms).
+    The additional data is then available for each time step via:
 
     .. code-block:: python
 
@@ -587,7 +583,7 @@ def __init__(self, filename,
         self._unwrap = unwrap_images
 
         if (util.iterable(additional_columns)
-                or additional_columns is None 
+                or additional_columns is None
                 or additional_columns is True):
             self._additional_columns = additional_columns
         else:
@@ -745,10 +741,12 @@ def _read_next_timestep(self):
         # Create the data arrays for additional attributes which will be saved
         # under ts.data
         if self._additional_columns is True:
-            # Parse every column that is not already parsed elsewhere (total \ parsable)
+            # Parse every column that is not already parsed
+            # elsewhere (total \ parsable)
             additional_keys = set(attrs).difference(self._parsable_columns)
         elif self._additional_columns:
-            additional_keys = [key for key in self._additional_columns if key in attrs]
+            additional_keys = \
+                [key for key in self._additional_columns if key in attrs]
         else:
             additional_keys = []
         for key in additional_keys:

testsuite/MDAnalysisTests/coordinates/reference.py

@@ -230,7 +230,7 @@ class RefLAMMPSDataMini(object):
 
 
 class RefLAMMPSDataAdditionalColumns(object):
-    charges = np.array([2.58855e-03, 6.91952e-05, 1.05548e-02, 4.20319e-03,
-                        9.19172e-03, 4.79777e-03, 6.36864e-04, 5.87125e-03,
-                        -2.18125e-03, 6.88910e-03])
-    additional_data = np.array(5 * [1.1, 1.2])
+    q = np.array([2.58855e-03, 6.91952e-05, 1.05548e-02, 4.20319e-03,
+                  9.19172e-03, 4.79777e-03, 6.36864e-04, 5.87125e-03,
+                  -2.18125e-03, 6.88910e-03])
+    l = np.array(5 * [1.1, 1.2])

testsuite/MDAnalysisTests/coordinates/test_lammps.py

@@ -29,7 +29,7 @@
 import MDAnalysis as mda
 from MDAnalysis import NoDataError
 
-from numpy.testing import (assert_equal, assert_allclose, assert_allclose)
+from numpy.testing import (assert_equal, assert_allclose)
 
 from MDAnalysisTests import make_Universe
 from MDAnalysisTests.coordinates.reference import (
@@ -515,15 +515,28 @@ def u(self, tmpdir, request):
                            lammps_coordinate_convention="auto")
 
     @pytest.fixture()
-    def u_additional_columns(self):
+    def u_additional_columns_true(self):
         f = LAMMPSDUMP_additional_columns
         top = LAMMPSdata_additional_columns
-        yield (mda.Universe(top, f, format='LAMMPSDUMP',
+        return mda.Universe(top, f, format='LAMMPSDUMP',
                             lammps_coordinate_convention="auto",
-                            additional_columns=['q']),
-               mda.Universe(top, f, format='LAMMPSDUMP',
+                            additional_columns=True)
+
+    @pytest.fixture()
+    def u_additional_columns_single(self):
+        f = LAMMPSDUMP_additional_columns
+        top = LAMMPSdata_additional_columns
+        return mda.Universe(top, f, format='LAMMPSDUMP',
+                            lammps_coordinate_convention="auto",
+                            additional_columns=['q'])
+
+    @pytest.fixture()
+    def u_additional_columns_multiple(self):
+        f = LAMMPSDUMP_additional_columns
+        top = LAMMPSdata_additional_columns
+        return mda.Universe(top, f, format='LAMMPSDUMP',
                             lammps_coordinate_convention="auto",
-                            additional_columns=['q', 'l']))
+                            additional_columns=['q', 'l'])
 
     @pytest.fixture()
     def reference_positions(self):
@@ -606,15 +619,17 @@ def test_atom_reordering(self, u, reference_positions):
             assert_allclose(atom1.position, atom1_pos-bmin, atol=1e-5)
             assert_allclose(atom13.position, atom13_pos-bmin, atol=1e-5)
 
-    def test_additional_columns(self, u_additional_columns):
-        # this is the universe with just q
-        charges = u_additional_columns[0].trajectory[0].data['q']
-        # this is the universe with both
-        second = u_additional_columns[1].trajectory[0].data['l']
-        assert_almost_equal(charges,
-                            RefLAMMPSDataAdditionalColumns.charges)
-        assert_almost_equal(second,
-                            RefLAMMPSDataAdditionalColumns.additional_data)
+    @pytest.mark.parametrize("system, fields", [
+        ('u_additional_columns_true', ['q', 'l']),
+        ('u_additional_columns_single', ['q']),
+        ('u_additional_columns_multiple', ['q', 'l']),
+    ])
+    def test_additional_columns(self, system, fields, request):
+        u = request.getfixturevalue(system)
+        for field in fields:
+            data = u.trajectory[0].data[field]
+            assert_allclose(data,
+                            getattr(RefLAMMPSDataAdditionalColumns, field))
 
 
 @pytest.mark.parametrize("convention",