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..d9cfc37 100644 --- a/CodeEntropy/levels/mda.py +++ b/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, diff --git a/CodeEntropy/levels/neighbors.py b/CodeEntropy/levels/neighbors.py index 7f77fc9..a385d21 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: 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 diff --git a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py index 98e3a7d..5fb1609 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()