Merge pull request #322 from CCPBioSim/lammps-compatibility

Lammps compatibility

Author: jkalayan

CodeEntropy/config/runtime.py

@@ -335,8 +335,14 @@ def _build_universe(
         kcal_units = args.kcal_force_units
 
         if forcefile is None:
-            logger.debug(f"Loading Universe with {tprfile} and {trrfile}")
-            return mda.Universe(tprfile, trrfile, format=fileformat)
+            if fileformat == "LAMMPSDUMP":
+                logger.debug(
+                    f"Loading Universe with {tprfile} and {trrfile} (LAMMPSDUMP)"
+                )
+                return universe_operations.convert_lammps(tprfile, trrfile, fileformat)
+            else:
+                logger.debug(f"Loading Universe with {tprfile} and {trrfile}")
+                return mda.Universe(tprfile, trrfile, format=fileformat)
 
         return universe_operations.merge_forces(
             tprfile, trrfile, forcefile, fileformat, kcal_units

CodeEntropy/levels/mda.py

@@ -121,6 +121,97 @@ def extract_fragment(
         selection_string = f"index {frag.indices[0]}:{frag.indices[-1]}"
         return self.select_atoms(universe, selection_string)
 
+    def convert_lammps(
+        self,
+        tprfile: str,
+        trrfile,
+        fileformat: str | None = None,
+    ) -> mda.Universe:
+        """Update the units produced from the universe produced from LAMMPS
+        format topology and trajectory files. MDA currently has a bug that
+        results in forces not being converted to the correct units
+        (see issue for more details:
+        https://github.com/MDAnalysis/mdanalysis/issues/5115
+        )
+        The method currently expects the following additional columns in the
+        lammps dump file: fx fy fz c_5 c_7
+        where c_5 and c_7 are the atom potential and kinetic energies respectively.
+
+        This method loads:
+
+        - Coordinates and dimensions from the coordinate trajectory
+          (``tprfile`` + ``trrfile``).
+
+        Args:
+            tprfile: Topology input file.
+            trrfile: Coordinate trajectory file(s). This can be a single path or a
+                list, as accepted by MDAnalysis.
+            fileformat: Optional file format for the coordinate trajectory, as
+                recognised by MDAnalysis.
+
+        Returns:
+            MDAnalysis.Universe: A new Universe containing coordinates, forces and
+            dimensions loaded into memory.
+
+        Raises:
+            ValueError: If fileformat is not one of the supported values.
+        """
+
+        def _convert_lammps_forces_energies(ts):
+            """
+            Convert lammps forces from kcal/mol/Ang to kJ/mol/Ang.
+            Assumes columns for per-atom potential (c_5) and kinetic energies (c_7)
+            are provided and converts these too.
+
+            Args:
+                ts: MDAnalysis timeseries from the trajectory.
+
+            Returns:
+                A converted time series.
+            """
+            ts.forces *= 4.184
+            ts.data["c_5"] *= 4.184
+            ts.data["c_7"] *= 4.184
+            return ts
+
+        def _convert_lammps_forces(ts):
+            """
+            Convert lammps forces from kcal/mol/Ang to kJ/mol/Ang.
+
+            Args:
+                ts: MDAnalysis timeseries from the trajectory.
+
+            Returns:
+                A converted time series.
+            """
+            ts.forces *= 4.184
+            return ts
+
+        if fileformat == "LAMMPSDUMP":
+            try:
+                return mda.Universe(
+                    tprfile,
+                    trrfile,
+                    format=fileformat,
+                    additional_columns=["fx", "fy", "fz", "c_5", "c_7"],
+                    transformations=[_convert_lammps_forces_energies],
+                )
+            except KeyError:
+                logger.debug(
+                    f"Warning: Energy columns not found in LAMMPSDUMP: {trrfile}"
+                )
+                return mda.Universe(
+                    tprfile,
+                    trrfile,
+                    format=fileformat,
+                    additional_columns=["fx", "fy", "fz"],
+                    transformations=[_convert_lammps_forces],
+                )
+        else:
+            raise ValueError(
+                f"Incorrect file format: {fileformat}, LAMMPSDUMP expected"
+            )
+
     def merge_forces(
         self,
         tprfile: str,

CodeEntropy/levels/neighbors.py

@@ -167,7 +167,11 @@ def _get_rdkit_mol(self, universe, mol_id):
             rdkit_mol = frag.convert_to("RDKIT", force=True, inferrer=None)
             logger.debug("Warning: Dummy atoms found")
         else:
-            rdkit_mol = molecule.convert_to("RDKIT", force=True)
+            try:
+                rdkit_mol = molecule.convert_to("RDKIT", force=True)
+            except Exception:
+                logger.debug("Warning: Constraint bonds to H atoms found")
+                rdkit_mol = molecule.convert_to("RDKIT", force=True, inferrer=None)
 
         number_heavy = rdkit_mol.GetNumHeavyAtoms()
         number_hydrogen = rdkit_mol.GetNumAtoms() - number_heavy

tests/unit/CodeEntropy/levels/test_mda_universe_operations.py

@@ -118,6 +118,59 @@ def test_merge_forces_scales_kcal(monkeypatch):
     assert np.allclose(forces, np.ones((2, 2, 3)) * 4.184)
 
 
+def test_convert_lammps_transforms_forces_and_energies(monkeypatch):
+    ops = UniverseOperations()
+
+    mock_universe = MagicMock()
+    transformations_captured = []
+
+    def capture_universe(*args, **kwargs):
+        if "transformations" in kwargs:
+            transformations_captured.extend(kwargs["transformations"])
+        return mock_universe
+
+    monkeypatch.setattr("CodeEntropy.levels.mda.mda.Universe", capture_universe)
+
+    ops.convert_lammps("tpr", "trr", "LAMMPSDUMP")
+
+    ts = MagicMock()
+    ts.forces = np.array([[1.0, 2.0, 3.0]], dtype=float)
+    ts.data = {"c_5": np.array([1.0]), "c_7": np.array([2.0])}
+
+    transformations_captured[0](ts)
+
+    assert np.allclose(ts.forces, np.array([[1.0, 2.0, 3.0]], dtype=float) * 4.184)
+    assert np.allclose(ts.data["c_5"], np.array([1.0], dtype=float) * 4.184)
+    assert np.allclose(ts.data["c_7"], np.array([2.0], dtype=float) * 4.184)
+
+
+def test_convert_lammps_fallback_on_keyerror(monkeypatch):
+    ops = UniverseOperations()
+
+    transformations_captured = []
+    call_count = [0]
+
+    def mock_universe(*args, **kwargs):
+        call_count[0] += 1
+        if call_count[0] == 1:
+            raise KeyError("c_5")
+        if "transformations" in kwargs:
+            transformations_captured.extend(kwargs["transformations"])
+        return MagicMock()
+
+    monkeypatch.setattr("CodeEntropy.levels.mda.mda.Universe", mock_universe)
+
+    ops.convert_lammps("tpr", "trr", "LAMMPSDUMP")
+
+    ts = MagicMock()
+    ts.forces = np.array([[1.0, 2.0, 3.0]], dtype=float)
+
+    transformations_captured[0](ts)
+
+    assert np.allclose(ts.forces, np.array([[1.0, 2.0, 3.0]], dtype=float) * 4.184)
+    assert call_count[0] == 2
+
+
 def test_select_atoms_builds_merged_universe_and_loads_timeseries(monkeypatch):
     ops = UniverseOperations()