Fix residue names & numbers for the PlainMDProtocol

news/fix_unk_resnames_plainmd.rst

@@ -0,0 +1,27 @@
+**Added:**
+
+* <news item>
+
+**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 <https://github.com/OpenFreeEnergy/openfe/pull/2178>`_).
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

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"

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

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):

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"])

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)