From a6ca05e2b93ba1231f3dec6b75ac5cf769fedac6 Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Tue, 28 Apr 2026 16:05:39 +0100 Subject: [PATCH 1/5] include functionality to read LAMMPSDUMP trajectory file --- CodeEntropy/config/runtime.py | 10 ++++- CodeEntropy/levels/mda.py | 69 +++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 2 deletions(-) diff --git a/CodeEntropy/config/runtime.py b/CodeEntropy/config/runtime.py index b55bcbe..6ef626b 100644 --- a/CodeEntropy/config/runtime.py +++ b/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 diff --git a/CodeEntropy/levels/mda.py b/CodeEntropy/levels/mda.py index f8cdec4..d5fa703 100644 --- a/CodeEntropy/levels/mda.py +++ b/CodeEntropy/levels/mda.py @@ -121,6 +121,75 @@ 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 + + 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: + raise + else: + raise ValueError( + f"Incorrect file format: {fileformat}, LAMMPSDUMP expected" + ) + def merge_forces( self, tprfile: str, From 5f6b298f94d2966dc6431da3c6a78618419510a5 Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Tue, 28 Apr 2026 16:06:35 +0100 Subject: [PATCH 2/5] handle extra restraint bonds to hydrogens for symmetry calcs --- CodeEntropy/levels/neighbors.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/CodeEntropy/levels/neighbors.py b/CodeEntropy/levels/neighbors.py index 7f77fc9..5e0ef4b 100644 --- a/CodeEntropy/levels/neighbors.py +++ b/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: Constriant 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 From c0f7a6a7881e6d35c29f360efb720431d738cfd4 Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Tue, 28 Apr 2026 16:19:37 +0100 Subject: [PATCH 3/5] handle LAMMPSDUMP with no atom energy information --- CodeEntropy/levels/mda.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/CodeEntropy/levels/mda.py b/CodeEntropy/levels/mda.py index d5fa703..d9cfc37 100644 --- a/CodeEntropy/levels/mda.py +++ b/CodeEntropy/levels/mda.py @@ -174,6 +174,19 @@ def _convert_lammps_forces_energies(ts): 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( @@ -184,7 +197,16 @@ def _convert_lammps_forces_energies(ts): transformations=[_convert_lammps_forces_energies], ) except KeyError: - raise + 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" From fffeb8367eb4fb32783847f20141531ee24a429e Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Wed, 29 Apr 2026 15:04:54 +0100 Subject: [PATCH 4/5] add unit tests for lammps conversion --- .../levels/test_mda_universe_operations.py | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py index 98e3a7d..baa89f3 100644 --- a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py +++ b/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() From b2efc2658cea592a73ae5587697af92644c3f4c4 Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Wed, 29 Apr 2026 16:14:28 +0100 Subject: [PATCH 5/5] fix typos and c_7 shape in tests --- CodeEntropy/levels/neighbors.py | 2 +- tests/unit/CodeEntropy/levels/test_mda_universe_operations.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CodeEntropy/levels/neighbors.py b/CodeEntropy/levels/neighbors.py index 5e0ef4b..a385d21 100644 --- a/CodeEntropy/levels/neighbors.py +++ b/CodeEntropy/levels/neighbors.py @@ -170,7 +170,7 @@ def _get_rdkit_mol(self, universe, mol_id): try: rdkit_mol = molecule.convert_to("RDKIT", force=True) except Exception: - logger.debug("Warning: Constriant bonds to H atoms found") + logger.debug("Warning: Constraint bonds to H atoms found") rdkit_mol = molecule.convert_to("RDKIT", force=True, inferrer=None) number_heavy = rdkit_mol.GetNumHeavyAtoms() diff --git a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py index baa89f3..5fb1609 100644 --- a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py +++ b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py @@ -141,7 +141,7 @@ def capture_universe(*args, **kwargs): 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) + assert np.allclose(ts.data["c_7"], np.array([2.0], dtype=float) * 4.184) def test_convert_lammps_fallback_on_keyerror(monkeypatch):