Lammps Dump Parser

package/CHANGELOG

@@ -13,22 +13,26 @@ The rules for this file:
   * release numbers follow "Semantic Versioning" http://semver.org
 
 ------------------------------------------------------------------------------
-??/??/?? IAlibay, Luthaf, hmacdope, rafaelpap, jbarnoud, BFedder, aya9aladdin
+??/??/?? IAlibay, Luthaf, hmacdope, rafaelpap, jbarnoud, BFedder, aya9aladdin,
+         jaclark5
 
  * 2.4.0
 
 Fixes
+  * LAMMPS Dump Reader translates the box to the origin (#3741)
   * Auxiliary; determination of representative frames: Removed undefined
     behaviour for cutoff values < -1 (PR # 3749)
   * Fixes distances.between not always returning AtomGroup (Issue #3794)
   * Upgrade to chemfiles 0.10.3 (Issue #3798)
   * fixes guessed_attributes and read_attributes methods of the Topology class, as 
     those methods were giving unexpected results for some topology attributes
     (e.g. bonds, angles) (PR #3779).
-   
+  * LAMMPS Dump Readerr translates the box to the origin (#3741)
 
 
 Enhancements
+  * Optionally unwraps trajectories with image flags upon loading (#3843)
+  * LAMMPS Dump Reader now imports velocities and forces (#3843)
   * AuxReaders now have a memory_limit parameter to control when memory usage
     warnings are issued. (PR # 3749)
   * MDAnalysis.units now has a lookup dictionary called MDANALYSIS_BASE_UNITS

package/MDAnalysis/coordinates/LAMMPS.py

@@ -455,7 +455,7 @@ class DumpReader(base.ReaderBase):
     """Reads the default `LAMMPS dump format`_
 
     Supports coordinates in the LAMMPS "unscaled" (x,y,z), "scaled" (xs,ys,zs),
-    "unwrapped" (xu,yu,zu) and "scaled_unwrapped" (xsu,ysu,zsu) coordinate
+    "unwrapped" (xu,yu,zu), and "scaled_unwrapped" (xsu,ysu,zsu) coordinate
     conventions (see https://docs.lammps.org/dump.html for more details).
     If `lammps_coordinate_convention='auto'` (default),
     one will be guessed. Guessing checks whether the coordinates fit each convention in the order "unscaled",
@@ -465,12 +465,18 @@ class DumpReader(base.ReaderBase):
     (xsu,ysu,zsu) they will automatically be converted from their
     scaled/fractional representation to their real values.
 
+    If a dump file has image flags it can be unwrapped upon loading with 
+    the keyword `unwrap_images=True`.
+
     Supports both orthogonal and triclinic simulation box dimensions (for more
     details see https://docs.lammps.org/Howto_triclinic.html). In either case,
     MDAnalysis will always use ``(*A*, *B*, *C*, *alpha*, *beta*, *gamma*)``
     to represent the unit cell. Lengths *A*, *B*, *C* are in the MDAnalysis
     length unit (Å), and angles are in degrees.
 
+    .. versionchanged:: 2.4.0
+       Now imports velocities and forces, translates the box to the origin,
+       and optionally unwraps trajectories with image flags upon loading.
     .. versionchanged:: 2.2.0
        Triclinic simulation boxes are supported.
        (Issue `#3383 <https://github.com/MDAnalysis/mdanalysis/issues/3383>`__)
@@ -479,8 +485,7 @@ class DumpReader(base.ReaderBase):
     .. versionadded:: 0.19.0
     """
     format = 'LAMMPSDUMP'
-    _conventions = ["auto", "unscaled", "scaled", "unwrapped",
-                    "scaled_unwrapped"]
+    _conventions = ["auto", "unscaled", "scaled", "unwrapped", "scaled_unwrapped"]
     _coordtype_column_names = {
         "unscaled": ["x", "y", "z"],
         "scaled": ["xs", "ys", "zs"],
@@ -489,7 +494,9 @@ class DumpReader(base.ReaderBase):
     }
 
     @store_init_arguments
-    def __init__(self, filename, lammps_coordinate_convention="auto",
+    def __init__(self, filename, 
+                 lammps_coordinate_convention="auto",
+                 unwrap_images=False,
                  **kwargs):
         super(DumpReader, self).__init__(filename, **kwargs)
 
@@ -503,6 +510,13 @@ def __init__(self, filename, lammps_coordinate_convention="auto",
                              " is not a valid option. "
                              f"Please choose one of {option_string}")
 
+        self._has_vels, self._has_forces = self._check_vels_forces()
+
+        if unwrap_images:
+            self._unwrap = unwrap_images
+        else:
+            self._unwrap = False
+
         self._cache = {}
 
         self._reopen()
@@ -512,9 +526,31 @@ def __init__(self, filename, lammps_coordinate_convention="auto",
     def _reopen(self):
         self.close()
         self._file = util.anyopen(self.filename)
-        self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs)
+        self.ts = self._Timestep(self.n_atoms, # NoteHere
+                                 velocities=self._has_vels,
+                                 forces=self._has_forces,
+                                 **self._ts_kwargs)
         self.ts.frame = -1
 
+    def _check_vels_forces(self):
+        with util.anyopen(self.filename) as f:
+            for i in range(8):
+                tmp = f.readline()
+            tmp = f.readline().split() 
+            _has_vels = "vx" in tmp
+            if _has_vels:
+                self._vel_ind = []
+                for i in range(len(tmp)):
+                    if tmp[i][0] == "v":
+                        self._vel_ind.append(i-2)
+            _has_forces = "fx" in tmp
+            if _has_forces:
+                self._forces_ind = []
+                for i in range(len(tmp)):
+                    if tmp[i][0] == "f":
+                        self._forces_ind.append(i-2)
+        return _has_vels, _has_forces
+
     @property
     @cached('n_atoms')
     def n_atoms(self):
@@ -609,6 +645,13 @@ def _read_next_timestep(self):
                 convention_to_col_ix[cv_name] = [attr_to_col_ix[x] for x in cv_col_names]
             except KeyError:
                 pass
+
+        if self._unwrap:
+            try:
+                image_cols = [attr_to_col_ix[x] for x in ["ix", "iy", "iz"]]
+            except:
+                raise ValueError("Trajectory must have image flag in order to unwrap.")
+
         # this should only trigger on first read of "ATOM" card, after which it
         # is fixed to the guessed value. Auto proceeds unscaled -> scaled
         # -> unwrapped -> scaled_unwrapped
@@ -620,22 +663,47 @@ def _read_next_timestep(self):
             except IndexError:
                 raise ValueError("No coordinate information detected")
         elif not self.lammps_coordinate_convention in convention_to_col_ix:
-            raise ValueError(f"No coordinates following convention {self.lammps_coordinate_convention} found in timestep")
+            raise ValueError("No coordinates following convention {} found in "
+                "timestep".format(self.lammps_coordinate_convention))
 
         coord_cols = convention_to_col_ix[self.lammps_coordinate_convention]
+        if self._unwrap:
+            coord_cols.extend(image_cols)
 
         ids = "id" in attr_to_col_ix
         for i in range(self.n_atoms):
             fields = f.readline().split()
             if ids:
                 indices[i] = fields[attr_to_col_ix["id"]]
-            ts.positions[i] = [fields[dim] for dim in coord_cols]
+            coords = np.array([fields[dim] for dim in coord_cols], np.float32)
+
+            if self._unwrap:
+                if len(coords) >= 6:
+                    images = coords[3:]
+                    coords = coords[:3]
+                    coords += images * ts.dimensions[:3]
+                else:
+                    raise ValueError("Cannot unwrap coordinates without image flags.")
+            else:
+                coords = coords[:3]
+            ts.positions[i] = coords
+
+            if self._has_vels:
+                ts.velocities[i] = [fields[dim] for dim in self._vel_ind]
+            if self._has_forces:
+                ts.forces[i] = [fields[dim] for dim in self._forces_ind]
 
         order = np.argsort(indices)
         ts.positions = ts.positions[order]
+        if self._has_vels:
+            ts.velocities = ts.velocities[order]
+        if self._has_forces:
+            ts.forces = ts.forces[order]
         if (self.lammps_coordinate_convention.startswith("scaled")):
             # if coordinates are given in scaled format, undo that
             ts.positions = distances.transform_StoR(ts.positions,
                                                     ts.dimensions)
+        # Transform to origin after transformation of scaled variables
+        ts.positions -= np.array([xlo, ylo, zlo])[None,:]
 
         return ts

testsuite/MDAnalysisTests/coordinates/test_lammps.py

@@ -37,7 +37,8 @@
 )
 from MDAnalysisTests.datafiles import (
     LAMMPScnt, LAMMPShyd, LAMMPSdata, LAMMPSdata_mini, LAMMPSdata_triclinic,
-    LAMMPSDUMP, LAMMPSDUMP_allcoords, LAMMPSDUMP_nocoords, LAMMPSDUMP_triclinic
+    LAMMPSDUMP, LAMMPSDUMP_allcoords, LAMMPSDUMP_nocoords, LAMMPSDUMP_triclinic,
+    LAMMPSDUMP_image_vf, LAMMPS_image_vf
 )
 
 
@@ -113,6 +114,37 @@ def LAMMPSDATAWriter(request, tmpdir_factory):
 
     return u, u_new
 
+def test_unwrap_vel_force():
+
+    u_wrapped = mda.Universe(LAMMPS_image_vf, [LAMMPSDUMP_image_vf], 
+                             format="LAMMPSDUMP")
+    u_wrapped.trajectory[-1]
+
+    assert_almost_equal(u_wrapped.atoms.positions[0], 
+                        np.array([2.56616, 6.11565, 7.37956]),
+                        decimal=5)
+    assert hasattr(u_wrapped.atoms, "velocities")
+    assert hasattr(u_wrapped.atoms, "forces")
+
+def test_unwrap_image_wrap():
+    u_unwrapped = mda.Universe(LAMMPS_image_vf, LAMMPSDUMP_image_vf,
+                               format="LAMMPSDUMP", unwrap_images=True)
+    u_unwrapped.trajectory[-1]
+
+    pos = (np.array([2.56616, 6.11565, 7.37956]) + 
+        np.array([3, 1, -4])*u_unwrapped.dimensions[:3])
+    assert_almost_equal(u_unwrapped.atoms.positions[0], 
+                        pos,
+                        decimal=5,
+                       )
+
+def test_unwrap_no_image():
+    with pytest.raises(ValueError, match="must have image flag"):
+        u_unwrapped = mda.Universe( 
+                                   LAMMPSDUMP_allcoords,
+                                   format="LAMMPSDUMP", 
+                                   unwrap_images=True)
+
 
 class TestLAMMPSDATAWriter(object):
     def test_Writer_dimensions(self, LAMMPSDATAWriter):
@@ -455,6 +487,12 @@ def reference_positions(self):
             np.array([length2, length2, length2, 90., 90., 90.]),
         ]
         data['box'] = boxes
+        box_mins = [
+            np.array([0., 0., 0.]),
+            np.array([0.21427867, 0.21427867, 0.21427867]),
+            np.array([-0.00544581, -0.00544581, -0.00544581]),
+        ]
+        data["mins"] = box_mins
 
         # data for atom id 1 in traj (ie first in frame)
         # isn't sensitive to required sorting
@@ -489,11 +527,13 @@ def test_seeking(self, u, reference_positions):
 
         assert_almost_equal(u.dimensions, reference_positions['box'][1],
                             decimal=5)
-        assert_almost_equal(u.atoms[0].position,
-                            reference_positions['atom1_pos'][1],
+        pos = (reference_positions['atom1_pos'][1] - 
+            reference_positions['mins'][1])
+        assert_almost_equal(u.atoms[0].position, pos,
                             decimal=5)
-        assert_almost_equal(u.atoms[12].position,
-                            reference_positions['atom13_pos'][1],
+        pos = (reference_positions['atom13_pos'][1] -
+            reference_positions['mins'][1])
+        assert_almost_equal(u.atoms[12].position, pos,
                             decimal=5)
 
     def test_boxsize(self, u, reference_positions):
@@ -504,11 +544,12 @@ def test_boxsize(self, u, reference_positions):
     def test_atom_reordering(self, u, reference_positions):
         atom1 = u.atoms[0]
         atom13 = u.atoms[12]
-        for ts, atom1_pos, atom13_pos in zip(u.trajectory,
+        for ts, atom1_pos, atom13_pos, bmin in zip(u.trajectory,
                                              reference_positions['atom1_pos'],
-                                             reference_positions['atom13_pos']):
-            assert_almost_equal(atom1.position, atom1_pos, decimal=5)
-            assert_almost_equal(atom13.position, atom13_pos, decimal=5)
+                                             reference_positions['atom13_pos'],
+                                             reference_positions['mins']):
+            assert_almost_equal(atom1.position, atom1_pos-bmin, decimal=5)
+            assert_almost_equal(atom13.position, atom13_pos-bmin, decimal=5)
 
 
 @pytest.mark.parametrize("convention",
@@ -556,9 +597,10 @@ def universes(self):
     def reference_unscaled_positions(self):
         # copied from trajectory file
         # atom 340 is the first one in the trajectory so we use that
-        atom340_pos1_unscaled = [4.48355, 0.331422, 1.59231]
-        atom340_pos2_unscaled = [4.41947, 35.4403, 2.25115]
-        atom340_pos3_unscaled = [4.48989, 0.360633, 2.63623]
+        bmin = np.array([0.02645, 0.02645, 0.02641])
+        atom340_pos1_unscaled = np.array([4.48355, 0.331422, 1.59231]) - bmin
+        atom340_pos2_unscaled = np.array([4.41947, 35.4403, 2.25115]) - bmin
+        atom340_pos3_unscaled = np.array([4.48989, 0.360633, 2.63623]) - bmin
         return np.asarray([atom340_pos1_unscaled, atom340_pos2_unscaled,
                           atom340_pos3_unscaled])
 

testsuite/MDAnalysisTests/data/lammps/image_vf.data

@@ -0,0 +1,48 @@
+LAMMPS data file via write_data, version 30 Jul 2021, timestep = 0
+
+7 atoms
+2 atom types
+1 bonds
+1 bond types
+
+0 10 xlo xhi
+0 10 ylo yhi
+0 10 zlo zhi
+
+Masses
+
+1 1
+2 1
+
+Pair Coeffs # lj/cut
+
+1 1 1
+2 1 1
+
+Bond Coeffs # harmonic
+
+1 1000 1
+
+Atoms # full
+
+4 0 2 0 5.891131260960588 3.398519062578611 0.23689615365476138 0 0 0
+1 0 1 0 4.999443228802319 5.0001459354508775 5.5008776144874 0 0 0
+2 0 1 0 5.000001730605194 4.999999546244266 4.500875455656523 0 0 0
+6 0 2 0 8.27919675422795 8.459848309149896 4.670531882285388 0 0 0
+3 0 2 0 3.362012829847722 5.522388563244657 8.669485965475673 0 0 0
+5 0 2 0 6.586761886625301 3.97905466100164 9.576146361865367 0 0 0
+7 0 2 0 7.621069760835089 6.390558075605001 9.73656065860773 0 0 0
+
+Velocities
+
+4 -0.07044405565641114 0.22797649438575432 0.9964537327696037
+1 1.6773916431557685 0.920692478778414 -2.57312540408295
+2 0.6247600152168848 1.0635940922731792 -0.09111771815842719
+6 0.28742689473258004 -0.5462919186206593 -1.8993269764408356
+3 -0.9529069979679257 -0.9763723022827592 3.1367559503110862
+5 -0.7649018752874791 -2.004540572487937 0.652141765065628
+7 -0.8013256241934161 1.3149417279540099 -0.22178134946410677
+
+Bonds
+
+1 1 1 2

testsuite/MDAnalysisTests/data/lammps/image_vf.lammpstrj

@@ -0,0 +1,48 @@
+ITEM: TIMESTEP
+0
+ITEM: NUMBER OF ATOMS
+7
+ITEM: BOX BOUNDS pp pp pp
+0.0000000000000000e+00 1.0000000000000000e+01
+0.0000000000000000e+00 1.0000000000000000e+01
+0.0000000000000000e+00 1.0000000000000000e+01
+ITEM: ATOMS id mol type q x y z ix iy iz vx vy vz fx fy fz
+4 0 2 0 5.89113 3.39852 0.236896 0 0 0 -0.0704441 0.227976 0.996454 -0.0384407 -0.027219 0.0347675
+1 0 1 0 4.99944 5.00015 5.50088 0 0 0 1.67739 0.920692 -2.57313 -0.00114452 0.000290546 -0.00139512
+2 0 1 0 5 5 4.50088 0 0 0 0.62476 1.06359 -0.0911177 -2.59759e-06 6.80857e-07 0.00465101
+6 0 2 0 8.2792 8.45985 4.67053 0 0 0 0.287427 -0.546292 -1.89933 0 0 0
+3 0 2 0 3.36201 5.52239 8.66949 0 0 0 -0.952907 -0.976372 3.13676 0.00617181 -0.00298752 -0.000692718
+5 0 2 0 6.58676 3.97905 9.57615 0 0 0 -0.764902 -2.00454 0.652142 0.0467422 0.0585064 -0.0360363
+7 0 2 0 7.62107 6.39056 9.73656 0 0 0 -0.801326 1.31494 -0.221781 -0.0133262 -0.0285912 -0.00129431
+ITEM: TIMESTEP
+1000
+ITEM: NUMBER OF ATOMS
+7
+ITEM: BOX BOUNDS pp pp pp
+0.0000000000000000e+00 1.0000000000000000e+01
+0.0000000000000000e+00 1.0000000000000000e+01
+0.0000000000000000e+00 1.0000000000000000e+01
+ITEM: ATOMS id mol type q x y z ix iy iz vx vy vz fx fy fz
+4 0 2 0 3.94665 -0.0323028 8.68651 0 0 0 -0.161974 -0.382081 0.798421 0.00521574 0.00769953 0.00166801
+1 0 1 0 8.86026 1.45707 6.49955 1 1 -2 2.56985 0.999077 -2.32084 9.89117 3.11795 4.16786
+2 0 1 0 7.97925 1.17482 6.13483 1 1 -2 0.143288 -0.0643266 -1.75995 -9.88081 -3.17777 -4.06472
+6 0 2 0 5.99448 3.88197 6.60468 0 1 -1 -0.694542 1.53609 0.18424 0.0038944 -0.00933068 0.00627887
+3 0 2 0 5.77413 2.29184 9.29137 -1 -1 3 -0.279753 -2.05124 2.46002 -0.00266798 -0.00561799 -0.0113445
+5 0 2 0 0.31665 8.73726 7.2999 0 -2 1 0.0467709 -1.59613 0.0234574 -0.0414919 0.047346 0.0252183
+7 0 2 0 8.86819 10.2398 8.37454 -1 1 0 -1.62364 1.5586 0.614649 0.024686 0.019719 -0.124957
+ITEM: TIMESTEP
+2000
+ITEM: NUMBER OF ATOMS
+7
+ITEM: BOX BOUNDS pp pp pp
+0.0000000000000000e+00 1.0000000000000000e+01
+0.0000000000000000e+00 1.0000000000000000e+01
+0.0000000000000000e+00 1.0000000000000000e+01
+ITEM: ATOMS id mol type q x y z ix iy iz vx vy vz fx fy fz
+1 0 1 0 2.56616 6.11565 7.37956 3 1 -4 0.0243353 -0.367818 -0.102071 2.48926 1.21693 -1.15048
+2 0 1 0 1.71413 6.13509 7.90135 3 1 -4 2.11307 0.0287745 -1.49691 -1.11544 0.0722323 0.689288
+6 0 2 0 8.68886 4.14263 8.29898 -1 2 -1 -0.646657 -0.404476 0.292564 0.00294106 0.00184583 -0.0004509
+3 0 2 0 0.709966 -0.100163 1.44498 -1 -3 6 -0.61121 -1.74802 1.39723 0 0 0
+7 0 2 0 2.98365 5.50022 4.32967 -2 3 1 -1.14938 1.11605 0.443808 -0.000791519 0.00601418 0.0160683
+5 0 2 0 1.54344 8.91853 6.46247 0 -3 1 -0.130583 1.71424 -1.00137 0.015867 -0.0269899 0.00946186
+4 0 2 0 3.53342 7.03855 7.07437 0 -1 1 0.400428 -0.338748 0.466751 -1.39184 -1.27004 0.436115