Lammps Dump Parser

package/CHANGELOG

@@ -13,7 +13,8 @@ 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
 
@@ -25,6 +26,9 @@ Fixes
   * 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 Reader now imports velocities and forces (#3843)
+  * LAMMPS Dump Readerr translates the box to the origin (#3741)
+  * Optionally unwraps trajectories with image flags upon loading (#3843)
 
 
 

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.
 
+    .. versionchanges:: 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"],
@@ -503,6 +508,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" in kwargs:
+            self._unwrap = kwargs["unwrap_images"]
+        else:
+            self._unwrap = False
+
         self._cache = {}
 
         self._reopen()
@@ -512,9 +524,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 +643,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,19 +661,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(
+                f"No coordinates following convention {self.lammps_coordinate_convention}"\
+                " found in timestep")
 
         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 len(coords) > 3:
+                if self._unwrap:
+                    if coords.shape[1] == 6:
+                        images = coords[:, 3:]
+                        coords = coords[:, :3]
+                        coords += images * ts.dimensions[:3]
+                    else:
+                        raise ValueError("Six coord elements are needed to unwrap with image flags.") 
+                else:
+                    coords = coords[:3]
+            else:
+                if self._unwrap:
+                    raise ValueError("Cannot unwrap coordinates without image flags.")
+            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]
+        ts.positions = ts.positions[order] - [xlo, ylo, zlo]
+        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,