Add LAMMPS trajectory reader (#327)

* Implement lammps reader

* Add tests

* Rename box_file to data_file

* Update readme

README.md

@@ -15,7 +15,7 @@ Gemdat is a Python library for the analysis of diffusion in solid-state electrol
 With Gemdat, you can:
 
 - Explore your MD simulation via an easy-to-use Python API
-- Load and analyze trajectories from VASP simulation data
+- Load and analyze trajectories from VASP and LAMMPS simulation data
 - Find jumps and transitions between sites
 - Effortlessly calculate tracer and jump diffusivity
 - Characterize and visualize diffusion pathways

src/gemdat/trajectory.py

@@ -12,7 +12,7 @@
 from typing import TYPE_CHECKING, Collection, Optional
 
 import numpy as np
-from pymatgen.core import Lattice
+from pymatgen.core import Element, Lattice
 from pymatgen.core.trajectory import Trajectory as PymatgenTrajectory
 from pymatgen.io import vasp
 
@@ -249,6 +249,97 @@ def from_vasprun(
 
         return obj
 
+    @classmethod
+    def from_lammps(
+        cls,
+        *,
+        coords_file: Path | str,
+        data_file: Path | str,
+        temperature: float,
+        time_step: float,
+        coords_format: str = 'xyz',
+        atom_style: str = 'atomic',
+        cache: Optional[str | Path] = None,
+        constant_lattice: bool = True,
+    ) -> Trajectory:
+        """Load data from LAMMPS.
+
+        Parameters
+        ----------
+        coords_file : ...
+            LAMMPS coords file with trajectory data
+        data_file : ...
+            LAMMPS data file with the lattice
+        temperature : float
+            Temperature of simulation in K
+        time_step : float
+            Time step of the simulation in fs
+        coords_format: str
+            Format of the coords file
+        atom_style : str
+            Atom style for box file
+        cache : Optional[Path], optional
+            Path to cache data for vasprun.xml
+        constant_lattice : bool
+            Whether the lattice changes during the simulation,
+            such as in an NPT MD simulation.
+
+        Returns
+        -------
+        trajectory : Trajectory
+            Output trajectory
+        """
+        from MDAnalysis import Universe
+        from pymatgen.io.lammps.data import LammpsData
+
+        coords_file = str(coords_file)
+        data_file = str(data_file)
+
+        if not cache:
+            kwargs = {
+                'coords_file': coords_file,
+                'data_file': data_file,
+                'temperature': temperature,
+                'time_step': time_step,
+            }
+            serialized = json.dumps(kwargs, sort_keys=True).encode()
+            hashid = hashlib.sha1(serialized).hexdigest()[:8]
+            cache = Path(coords_file).with_suffix(f'.{coords_format}.{hashid}.cache')
+
+        if Path(cache).exists():
+            try:
+                return cls.from_cache(cache)
+            except Exception as e:
+                print(e)
+                print(f'Error reading from cache, reading {coords_file!r}')
+
+        if not constant_lattice:
+            raise NotImplementedError('Lammps reader does not support NPT simulations')
+
+        lammps_data = LammpsData.from_file(filename=data_file, atom_style=atom_style)
+        lattice = lammps_data.structure.lattice
+
+        utraj = Universe(coords_file, format=coords_format)
+        coords = utraj.trajectory.timeseries()
+        coords = lattice.get_fractional_coords(coords)
+
+        species = [Element(sp) for sp in utraj.atoms.elements]
+
+        obj = cls(
+            species=species,
+            coords=coords,
+            lattice=lattice,
+            time_step=time_step,
+            constant_lattice=constant_lattice,
+            metadata={'temperature': temperature},
+        )
+        obj.to_positions()
+
+        if cache:
+            obj.to_cache(cache)
+
+        return obj
+
     def get_lattice(self, idx: int | None = None) -> Lattice:
         """Get lattice.
 

tests/trajectory_test.py

@@ -158,3 +158,19 @@ def test_mean_squared_displacement(trajectory):
     assert_allclose(msd[0], [0.0, 0.0525, 0.19, 0.425, 0.81])
     assert isinstance(msd, np.ndarray)
     assert_allclose(msd.mean(), 0.073875)
+
+
+def test_from_lammps():
+    from pathlib import Path
+
+    data_dir = Path(__file__).parent / 'data' / 'lammps'
+
+    traj = Trajectory.from_lammps(
+        coords_file=data_dir / 'lammps_coords.xyz',
+        data_file=data_dir / 'lammps_data.txt',
+        temperature=700,
+        time_step=2,
+    )
+
+    assert traj.positions.shape == (4, 80, 3)
+    assert len(traj.species) == 80