From bf10d71d2763550911b5a4d6fe82ba68961876f4 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Sat, 29 Aug 2026 18:11:19 -0500 Subject: [PATCH 1/7] Add residue naming for the PlainMDProtocol --- .../protocols/openmm_md/plain_md_methods.py | 46 +++++++++++-- .../openmm_utils/offmolecule_utils.py | 66 +++++++++++-------- .../openmm_utils/system_validation.py | 19 +++--- .../protocols/openmm_md/test_plain_md_slow.py | 6 ++ 4 files changed, 90 insertions(+), 47 deletions(-) diff --git a/src/openfe/protocols/openmm_md/plain_md_methods.py b/src/openfe/protocols/openmm_md/plain_md_methods.py index c16c3d3cb..145bacc06 100644 --- a/src/openfe/protocols/openmm_md/plain_md_methods.py +++ b/src/openfe/protocols/openmm_md/plain_md_methods.py @@ -61,6 +61,13 @@ BasePartialChargeSettings, FemtosecondQuantity, ) +from openfe.protocols.openmm_utils.offmolecule_utils import ( + _get_used_offmol_property, + _get_unique_name, + _get_offmol_metadata, + _set_offmol_metadata, + _next_available_number, +) from openfe.utils import log_system_probe, without_oechem_backend logger = logging.getLogger(__name__) @@ -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..9edd242b1 100644 --- a/src/openfe/protocols/openmm_utils/offmolecule_utils.py +++ b/src/openfe/protocols/openmm_utils/offmolecule_utils.py @@ -111,6 +111,32 @@ def _get_offmol_resname(offmol: OFFMolecule) -> str | None: return _get_offmol_metadata(offmol, "residue_name") +def _get_used_offmol_property(offmols: list[OFFMolecule], offmol_property: str) -> list[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 + + +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 _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( small_mols: dict[SmallMoleculeComponent, OFFMolecule], alchemical_components: Collection[Component], @@ -135,35 +161,15 @@ 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(): @@ -172,6 +178,8 @@ def _next_resnum() -> int: name = lig_name if smc in alchemical else cof_name _set_offmol_resname(offmol, 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..192b055c2 100644 --- a/src/openfe/protocols/openmm_utils/system_validation.py +++ b/src/openfe/protocols/openmm_utils/system_validation.py @@ -205,14 +205,11 @@ def validate_barostat(state: ChemicalSystem, barostat: str): logger.warning(wmsg) -ParseCompRet = Tuple[ - Optional[SolventComponent], - Optional[ProteinComponent], - list[SmallMoleculeComponent], -] - - -def get_components(state: ChemicalSystem) -> ParseCompRet: +def get_components(state: ChemicalSystem) -> tuple[ + SolventComponent | None, + ProteinComponent | None, + list[SmallMoleculeComponent] | None, +]: """ Establish all necessary Components for the transformation. @@ -223,11 +220,11 @@ 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] + small_mols : list[SmallMoleculeComponent] | None """ 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..1a0cbcda7 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 @@ -5,6 +5,7 @@ import pytest from gufe.protocols import execute_DAG from openff.units import unit +import MDAnalysis as mda from openfe.protocols import openmm_md @@ -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"]) From 26fe6214bd6af2d9758f37f6c4c2c1521f90a7c5 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Sat, 29 Aug 2026 18:19:35 -0500 Subject: [PATCH 2/7] Add news item --- news/fix_unk_resnames_plainmd.rst | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 news/fix_unk_resnames_plainmd.rst diff --git a/news/fix_unk_resnames_plainmd.rst b/news/fix_unk_resnames_plainmd.rst new file mode 100644 index 000000000..f143ff1b1 --- /dev/null +++ b/news/fix_unk_resnames_plainmd.rst @@ -0,0 +1,24 @@ +**Added:** + +* + +**Changed:** + +* The PlainMDProtocol now assigns the residue name LIG (or LG?) and + a unique residue number to all SmallMoleculeComponents. + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* From b98120ab669b049c2abe1bce1972580287732863 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 23:24:02 +0000 Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../protocols/openmm_md/plain_md_methods.py | 16 ++++++++-------- .../protocols/openmm_utils/offmolecule_utils.py | 3 +-- .../protocols/openmm_utils/system_validation.py | 4 +++- .../protocols/openmm_md/test_plain_md_slow.py | 2 +- 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/src/openfe/protocols/openmm_md/plain_md_methods.py b/src/openfe/protocols/openmm_md/plain_md_methods.py index 145bacc06..48b3fce14 100644 --- a/src/openfe/protocols/openmm_md/plain_md_methods.py +++ b/src/openfe/protocols/openmm_md/plain_md_methods.py @@ -57,16 +57,16 @@ system_creation, system_validation, ) -from openfe.protocols.openmm_utils.omm_settings import ( - BasePartialChargeSettings, - FemtosecondQuantity, -) from openfe.protocols.openmm_utils.offmolecule_utils import ( - _get_used_offmol_property, - _get_unique_name, _get_offmol_metadata, - _set_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, ) from openfe.utils import log_system_probe, without_oechem_backend @@ -469,7 +469,7 @@ def run( } # We assign everything LIG unless LIG is already a name - resname = _get_unique_name("LIG", "LG", used_names) + resname = _get_unique_name("LIG", "LG", used_names) # Now we assign all the residue names & numbers for offmol in smc_components.values(): diff --git a/src/openfe/protocols/openmm_utils/offmolecule_utils.py b/src/openfe/protocols/openmm_utils/offmolecule_utils.py index 9edd242b1..0d1956475 100644 --- a/src/openfe/protocols/openmm_utils/offmolecule_utils.py +++ b/src/openfe/protocols/openmm_utils/offmolecule_utils.py @@ -163,8 +163,7 @@ def assign_offmol_residue_metadata( 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") + int(i) for i in _get_used_offmol_property(list(small_mols.values()), "residue_number") } lig_name = _get_unique_name(ligand_resname, ligand_stem, used_names) diff --git a/src/openfe/protocols/openmm_utils/system_validation.py b/src/openfe/protocols/openmm_utils/system_validation.py index 192b055c2..9e20fdef1 100644 --- a/src/openfe/protocols/openmm_utils/system_validation.py +++ b/src/openfe/protocols/openmm_utils/system_validation.py @@ -205,7 +205,9 @@ def validate_barostat(state: ChemicalSystem, barostat: str): logger.warning(wmsg) -def get_components(state: ChemicalSystem) -> tuple[ +def get_components( + state: ChemicalSystem, +) -> tuple[ SolventComponent | None, ProteinComponent | None, list[SmallMoleculeComponent] | None, 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 1a0cbcda7..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,10 +2,10 @@ # 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 -import MDAnalysis as mda from openfe.protocols import openmm_md From fc4dd6e9ee44907816f21e3e69e67aa36ad185c4 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Sat, 29 Aug 2026 18:32:46 -0500 Subject: [PATCH 4/7] fix mypy --- src/openfe/protocols/openmm_utils/offmolecule_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/openfe/protocols/openmm_utils/offmolecule_utils.py b/src/openfe/protocols/openmm_utils/offmolecule_utils.py index 9edd242b1..6f18a6426 100644 --- a/src/openfe/protocols/openmm_utils/offmolecule_utils.py +++ b/src/openfe/protocols/openmm_utils/offmolecule_utils.py @@ -111,7 +111,7 @@ def _get_offmol_resname(offmol: OFFMolecule) -> str | None: return _get_offmol_metadata(offmol, "residue_name") -def _get_used_offmol_property(offmols: list[OFFMolecule], offmol_property: str) -> list[str]: +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) From 2dd38b62932f2a1509b8955247d5ccc72fe06c22 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Tue, 1 Sep 2026 12:13:59 -0400 Subject: [PATCH 5/7] fix up tests to check for strings instead --- .../tests/protocols/test_openmmutils.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/openfe/tests/protocols/test_openmmutils.py b/src/openfe/tests/protocols/test_openmmutils.py index f1672b6ba..b05512f96 100644 --- a/src/openfe/tests/protocols/test_openmmutils.py +++ b/src/openfe/tests/protocols/test_openmmutils.py @@ -1338,7 +1338,7 @@ class TestAssignOffmolResidueMetadata: _RESIDUE_METADATA_CASES = { "single ligand": ( [("benzene", True, None, None)], - [("LIG", 1)], + [("LIG", "1")], ), "ligand and cofactors": ( [ @@ -1346,7 +1346,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 +1354,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 +1371,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 +1379,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")], ), } From 724d5b219eb286588300820b2b3c7bc50568667e Mon Sep 17 00:00:00 2001 From: IAlibay Date: Tue, 1 Sep 2026 22:24:36 -0400 Subject: [PATCH 6/7] fix mypy --- src/openfe/protocols/openmm_utils/system_validation.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/openfe/protocols/openmm_utils/system_validation.py b/src/openfe/protocols/openmm_utils/system_validation.py index 9e20fdef1..5d5a6b0b7 100644 --- a/src/openfe/protocols/openmm_utils/system_validation.py +++ b/src/openfe/protocols/openmm_utils/system_validation.py @@ -210,7 +210,7 @@ def get_components( ) -> tuple[ SolventComponent | None, ProteinComponent | None, - list[SmallMoleculeComponent] | None, + list[SmallMoleculeComponent], ]: """ Establish all necessary Components for the transformation. @@ -226,7 +226,8 @@ def get_components( If it exists, the SolventComponent for the state, otherwise None. protein_comp : ProteinComponent | None If it exists, the ProteinComponent for the state, otherwise None. - small_mols : list[SmallMoleculeComponent] | None + small_mols : list[SmallMoleculeComponent] + List of SmallMoleculeComponent, if none exists, returns an empty list. """ def _get_single_comps(state, comptype): From 0e9586ef723be5f50cd198199d1474ec41de8cf4 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Wed, 2 Sep 2026 21:22:55 -0400 Subject: [PATCH 7/7] include fixes from code review --- news/fix_unk_resnames_plainmd.rst | 7 ++- .../openmm_utils/offmolecule_utils.py | 45 +------------------ .../openmm_rfe/test_hybrid_top_protocol.py | 5 +-- .../tests/protocols/test_openmmutils.py | 6 +-- 4 files changed, 11 insertions(+), 52 deletions(-) diff --git a/news/fix_unk_resnames_plainmd.rst b/news/fix_unk_resnames_plainmd.rst index f143ff1b1..00c5aec45 100644 --- a/news/fix_unk_resnames_plainmd.rst +++ b/news/fix_unk_resnames_plainmd.rst @@ -4,8 +4,11 @@ **Changed:** -* The PlainMDProtocol now assigns the residue name LIG (or LG?) and - a unique residue number to all SmallMoleculeComponents. +* 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:** diff --git a/src/openfe/protocols/openmm_utils/offmolecule_utils.py b/src/openfe/protocols/openmm_utils/offmolecule_utils.py index 3622984bb..50a69b1ae 100644 --- a/src/openfe/protocols/openmm_utils/offmolecule_utils.py +++ b/src/openfe/protocols/openmm_utils/offmolecule_utils.py @@ -70,47 +70,6 @@ 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 - - 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_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 _get_used_offmol_property(offmols: list[OFFMolecule], offmol_property: str) -> set[str]: used_property: set[str] = set() for mol in offmols: @@ -172,10 +131,10 @@ def assign_offmol_residue_metadata( 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: resnum = _next_available_number(used_resnums) _set_offmol_metadata(offmol, "residue_number", str(resnum)) 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 b05512f96..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, ) @@ -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