diff --git a/news/fix_unk_resnames_plainmd.rst b/news/fix_unk_resnames_plainmd.rst new file mode 100644 index 000000000..00c5aec45 --- /dev/null +++ b/news/fix_unk_resnames_plainmd.rst @@ -0,0 +1,27 @@ +**Added:** + +* + +**Changed:** + +* The PlainMDProtocol now assigns the residue name "LIG" (or "LG?" where + ? is a number between 1 and 9 to yield a unique residue name) and + a unique residue number to all SmallMoleculeComponents. Previously + these would have been assigned the residue name "UNK" + (`PR #2178 `_). + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* diff --git a/src/openfe/protocols/openmm_md/plain_md_methods.py b/src/openfe/protocols/openmm_md/plain_md_methods.py index c16c3d3cb..48b3fce14 100644 --- a/src/openfe/protocols/openmm_md/plain_md_methods.py +++ b/src/openfe/protocols/openmm_md/plain_md_methods.py @@ -57,6 +57,13 @@ system_creation, system_validation, ) +from openfe.protocols.openmm_utils.offmolecule_utils import ( + _get_offmol_metadata, + _get_unique_name, + _get_used_offmol_property, + _next_available_number, + _set_offmol_metadata, +) from openfe.protocols.openmm_utils.omm_settings import ( BasePartialChargeSettings, FemtosecondQuantity, @@ -441,6 +448,9 @@ def run( ) solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) + + # Set the solvent component to be the protein component if the latter + # is already solvated. if isinstance(protein_comp, SolvatedPDBComponent): solvent_comp = protein_comp @@ -450,10 +460,32 @@ def run( i: i.to_openff() for i in small_mols } - # a. assign partial charges to smcs + # a. assign residue names / numbers to the small molecules. + # First gather all the unique names & resnums + used_names = _get_used_offmol_property(list(smc_components.values()), "residue_name") + used_resnums = { + int(i) + for i in _get_used_offmol_property(list(smc_components.values()), "residue_number") + } + + # We assign everything LIG unless LIG is already a name + resname = _get_unique_name("LIG", "LG", used_names) + + # Now we assign all the residue names & numbers + for offmol in smc_components.values(): + name = _get_offmol_metadata(offmol, "residue_name") + num = _get_offmol_metadata(offmol, "residue_number") + if name is None: + _set_offmol_metadata(offmol, "residue_name", resname) + if num is None: + resnum = _next_available_number(used_resnums) + _set_offmol_metadata(offmol, "residue_number", str(resnum)) + used_resnums.add(resnum) + + # b. assign partial charges to smcs self._assign_partial_charges(charge_settings, smc_components) - # b. get a system generator + # c. get a system generator if output_settings.forcefield_cache is not None: ffcache = self.shared_basepath / output_settings.forcefield_cache else: @@ -475,7 +507,7 @@ def run( for mol in smc_components.values(): system_generator.create_system(mol.to_topology().to_openmm(), molecules=[mol]) - # c. get OpenMM Modeller + a resids dictionary for each component + # d. get OpenMM Modeller + a resids dictionary for each component stateA_modeller, comp_resids = system_creation.get_omm_modeller( protein_comp=protein_comp, solvent_comp=solvent_comp, @@ -484,22 +516,22 @@ def run( solvent_settings=solvation_settings, ) - # d. get topology & positions + # e. get topology & positions # Note: roundtrip positions to remove vec3 issues stateA_topology = stateA_modeller.getTopology() stateA_positions = to_openmm(from_openmm(stateA_modeller.getPositions())) - # e. create the stateA System + # f. create the stateA System stateA_system = system_generator.create_system( stateA_topology, molecules=[s.to_openff() for s in small_mols], ) - # f. Save pdb of entire system topology to file, this is always needed for restarts + # g. Save pdb of entire system topology to file, this is always needed for restarts with open(self.shared_basepath / output_settings.preminimized_structure, "w") as f: openmm.app.PDBFile.writeFile(stateA_topology, stateA_positions, file=f, keepIds=True) - # g. Save the system and positions to file + # h. Save the system and positions to file system_outfile = self.shared_basepath / "system.xml.bz2" serialization.serialize(stateA_system, system_outfile) positions_outfile = self.shared_basepath / "input_positions.npy" diff --git a/src/openfe/protocols/openmm_utils/offmolecule_utils.py b/src/openfe/protocols/openmm_utils/offmolecule_utils.py index 12805e087..50a69b1ae 100644 --- a/src/openfe/protocols/openmm_utils/offmolecule_utils.py +++ b/src/openfe/protocols/openmm_utils/offmolecule_utils.py @@ -70,45 +70,30 @@ def _get_offmol_metadata(offmol: OFFMolecule, key: Any) -> Any | None: return value -def _set_offmol_resname( - offmol: OFFMolecule, - resname: str | None, -) -> None: - """ - Helper method to set offmol residue names +def _get_used_offmol_property(offmols: list[OFFMolecule], offmol_property: str) -> set[str]: + used_property: set[str] = set() + for mol in offmols: + prop = _get_offmol_metadata(mol, offmol_property) + if prop is not None: + used_property.add(prop) + return used_property - Parameters - ---------- - offmol : openff.toolkit.Molecule - Molecule to assign a residue name to. - resname : str | None - Residue name to be set. Set to None to clear it. - Returns - ------- - None - """ - _set_offmol_metadata(offmol, "residue_name", resname) +def _get_unique_name(default: str, stem: str, used_names: set[str]) -> str: + if default not in used_names: + return default + for i in range(1, 10): + if (candidate := f"{stem}{i}") not in used_names: + return candidate + raise ValueError(f"Could not assign a unique residue name with stem {stem!r}.") -def _get_offmol_resname(offmol: OFFMolecule) -> str | None: - """ - Helper method to get an offmol's residue name and make sure it is - consistent across all atoms in the Molecule. - - Parameters - ---------- - offmol : openff.toolkit.Molecule - Molecule to get the residue name from. - - Returns - ------- - resname : Optional[str] - Residue name of the molecule. ``None`` if the Molecule - does not have a residue name, or if the residue name is - inconsistent across all the atoms. - """ - return _get_offmol_metadata(offmol, "residue_name") +def _next_available_number(numbers: set[int]) -> int: + """Return the lowest residue number not already in use.""" + number = 1 + while number in numbers: + number += 1 + return number def assign_offmol_residue_metadata( @@ -135,43 +120,24 @@ def assign_offmol_residue_metadata( alchemical = set(alchemical_components) - used_names: set[str] = set() - used_resnums: set[int] = set() - for offmol in small_mols.values(): - name = _get_offmol_resname(offmol) - resnum = _get_offmol_metadata(offmol, "residue_number") - if name is not None: - used_names.add(name) - if resnum is not None: - used_resnums.add(int(resnum)) - - def _unique(default: str, stem: str) -> str: - if default not in used_names: - return default - for i in range(1, 10): - if (candidate := f"{stem}{i}") not in used_names: - return candidate - raise ValueError(f"Could not assign a unique residue name with stem {stem!r}.") - - lig_name = _unique(ligand_resname, ligand_stem) - used_names.add(lig_name) - cof_name = _unique(cofactor_resname, cofactor_stem) + used_names = _get_used_offmol_property(list(small_mols.values()), "residue_name") + used_resnums = { + int(i) for i in _get_used_offmol_property(list(small_mols.values()), "residue_number") + } - def _next_resnum() -> int: - """Return the lowest residue number not already in use.""" - resnum = 1 - while resnum in used_resnums: - resnum += 1 - used_resnums.add(resnum) - return resnum + lig_name = _get_unique_name(ligand_resname, ligand_stem, used_names) + used_names.add(lig_name) + cof_name = _get_unique_name(cofactor_resname, cofactor_stem, used_names) assigned: dict[SmallMoleculeComponent, str] = {} for smc, offmol in small_mols.items(): - name = _get_offmol_resname(offmol) + name = _get_offmol_metadata(offmol, "residue_name") if name is None: name = lig_name if smc in alchemical else cof_name - _set_offmol_resname(offmol, name) + _set_offmol_metadata(offmol, "residue_name", name) if _get_offmol_metadata(offmol, "residue_number") is None: - _set_offmol_metadata(offmol, "residue_number", _next_resnum()) + resnum = _next_available_number(used_resnums) + _set_offmol_metadata(offmol, "residue_number", str(resnum)) + used_resnums.add(resnum) assigned[smc] = name return assigned diff --git a/src/openfe/protocols/openmm_utils/system_validation.py b/src/openfe/protocols/openmm_utils/system_validation.py index 906999910..5d5a6b0b7 100644 --- a/src/openfe/protocols/openmm_utils/system_validation.py +++ b/src/openfe/protocols/openmm_utils/system_validation.py @@ -205,14 +205,13 @@ def validate_barostat(state: ChemicalSystem, barostat: str): logger.warning(wmsg) -ParseCompRet = Tuple[ - Optional[SolventComponent], - Optional[ProteinComponent], +def get_components( + state: ChemicalSystem, +) -> tuple[ + SolventComponent | None, + ProteinComponent | None, list[SmallMoleculeComponent], -] - - -def get_components(state: ChemicalSystem) -> ParseCompRet: +]: """ Establish all necessary Components for the transformation. @@ -223,11 +222,12 @@ def get_components(state: ChemicalSystem) -> ParseCompRet: Returns ------- - solvent_comp : Optional[SolventComponent] + solvent_comp : SolventComponent | None If it exists, the SolventComponent for the state, otherwise None. - protein_comp : Optional[ProteinComponent] + protein_comp : ProteinComponent | None If it exists, the ProteinComponent for the state, otherwise None. small_mols : list[SmallMoleculeComponent] + List of SmallMoleculeComponent, if none exists, returns an empty list. """ def _get_single_comps(state, comptype): diff --git a/src/openfe/tests/protocols/openmm_md/test_plain_md_slow.py b/src/openfe/tests/protocols/openmm_md/test_plain_md_slow.py index edcc5944f..33e976bfd 100644 --- a/src/openfe/tests/protocols/openmm_md/test_plain_md_slow.py +++ b/src/openfe/tests/protocols/openmm_md/test_plain_md_slow.py @@ -2,6 +2,7 @@ # For details, see https://github.com/OpenFreeEnergy/openfe import pathlib +import MDAnalysis as mda import pytest from gufe.protocols import execute_DAG from openff.units import unit @@ -75,6 +76,11 @@ def test_vacuum_sim( assert pur.outputs["npt_equil_pdb"] == unit_shared / "equil_npt.pdb" assert pur.outputs["nvt_equil_pdb"] is None + # Check that the residue name is correct + univ = mda.Universe(pur.outputs["npt_equil_pdb"]) + assert univ.atoms.resnames[0] == "LIG" + assert univ.atoms.resnums[0] == 1 + @pytest.mark.integration @pytest.mark.parametrize("platform", ["CUDA"]) diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py index 611e6d0f9..9b6b3872f 100644 --- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py +++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py @@ -52,8 +52,7 @@ HAS_OPENEYE, ) from openfe.protocols.openmm_utils.offmolecule_utils import ( - _get_offmol_resname, - _set_offmol_resname, + _set_offmol_metadata, ) @@ -2429,7 +2428,7 @@ def _run_setup_dry(stateA, stateB, mapping, settings, tmp_path): def _named_smc(smc, resname): off = Molecule.from_rdkit(copy.deepcopy(smc.to_rdkit())) - _set_offmol_resname(off, resname) + _set_offmol_metadata(off, "residue_name", resname) return openfe.SmallMoleculeComponent.from_openff(off) diff --git a/src/openfe/tests/protocols/test_openmmutils.py b/src/openfe/tests/protocols/test_openmmutils.py index f1672b6ba..f4dcf6bff 100644 --- a/src/openfe/tests/protocols/test_openmmutils.py +++ b/src/openfe/tests/protocols/test_openmmutils.py @@ -45,9 +45,7 @@ ) from openfe.protocols.openmm_utils.offmolecule_utils import ( _get_offmol_metadata, - _get_offmol_resname, _set_offmol_metadata, - _set_offmol_resname, assign_offmol_residue_metadata, ) @@ -1338,7 +1336,7 @@ class TestAssignOffmolResidueMetadata: _RESIDUE_METADATA_CASES = { "single ligand": ( [("benzene", True, None, None)], - [("LIG", 1)], + [("LIG", "1")], ), "ligand and cofactors": ( [ @@ -1346,7 +1344,7 @@ class TestAssignOffmolResidueMetadata: ("toluene", False, None, None), ("phenol", False, None, None), ], - [("LIG", 1), ("COF", 2), ("COF", 3)], + [("LIG", "1"), ("COF", "2"), ("COF", "3")], ), "two ligands pool to LIG": ( [ @@ -1354,15 +1352,15 @@ class TestAssignOffmolResidueMetadata: ("toluene", True, None, None), ("phenol", False, None, None), ], - [("LIG", 1), ("LIG", 2), ("COF", 3)], + [("LIG", "1"), ("LIG", "2"), ("COF", "3")], ), "custom cofactor name kept": ( [("benzene", True, None, None), ("toluene", False, "NAD", None)], - [("LIG", 1), ("NAD", 2)], + [("LIG", "1"), ("NAD", "2")], ), "custom ligand name kept": ( [("benzene", True, "BNZ", None), ("toluene", False, None, None)], - [("BNZ", 1), ("COF", 2)], + [("BNZ", "1"), ("COF", "2")], ), "cofactor named LIG, ligands fall back to LG1": ( [ @@ -1371,7 +1369,7 @@ class TestAssignOffmolResidueMetadata: ("phenol", False, "LIG", None), ("benzonitrile", False, None, None), ], - [("LG1", 1), ("LG1", 2), ("LIG", 3), ("COF", 4)], + [("LG1", "1"), ("LG1", "2"), ("LIG", "3"), ("COF", "4")], ), "ligand named COF, cofactors fall back to CF1": ( [ @@ -1379,19 +1377,19 @@ class TestAssignOffmolResidueMetadata: ("toluene", False, None, None), ("phenol", False, None, None), ], - [("COF", 1), ("CF1", 2), ("CF1", 3)], + [("COF", "1"), ("CF1", "2"), ("CF1", "3")], ), "pinned residue number respected": ( - [("benzene", True, None, None), ("toluene", False, None, 5)], - [("LIG", 1), ("COF", 5)], + [("benzene", True, None, None), ("toluene", False, None, "5")], + [("LIG", "1"), ("COF", "5")], ), "auto residue number steps past a pinned one": ( [ ("benzene", True, None, None), - ("toluene", False, None, 2), + ("toluene", False, None, "2"), ("phenol", False, None, None), ], - [("LIG", 1), ("COF", 2), ("COF", 3)], + [("LIG", "1"), ("COF", "2"), ("COF", "3")], ), } @@ -1405,7 +1403,7 @@ def _build_small_molecules(benzene_modifications, specs): off = copy.deepcopy(smc.to_openff()) if resname is not None: - _set_offmol_resname(off, resname) + _set_offmol_metadata(off, "residue_name", resname) if residue_number is not None: _set_offmol_metadata(off, "residue_number", residue_number) @@ -1426,6 +1424,6 @@ def test_assign_offmol_residue_metadata(self, benzene_modifications, specs, expe assigned = assign_offmol_residue_metadata(small, alchemical) for smc, off, (resname, residue_number) in zip(small.keys(), small.values(), expected): - assert _get_offmol_resname(off) == resname + assert _get_offmol_metadata(off, "residue_name") == resname assert _get_offmol_metadata(off, "residue_number") == residue_number assert assigned[smc] == resname