Fix residue naming

news/fix_unk_resnames.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* Small molecules in RelativeHybridTopologyProtocol topologies (including the output PDB) are now named LIG (alchemical ligand) and COF (cofactors) instead of UNK. If a residue name was already assigned, the assigned one is kept.
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* Fixed inflated ligand RMSD in the RelativeHybridTopology protocol's structural analysis for systems containing cofactors; the ligand RMSD is now computed for the alchemical ligand alone rather than conflating it with cofactors that shared the UNK residue name.
+
+**Security:**
+
+* <news item>

src/openfe/protocols/openmm_rfe/hybridtop_units.py

@@ -46,6 +46,11 @@
 from openmmtools import multistate
 
 import openfe
+from openfe.protocols.openmm_utils.offmolecule_utils import (
+    _get_offmol_resname,
+    _set_offmol_metadata,
+    _set_offmol_resname,
+)
 from openfe.protocols.openmm_utils.omm_settings import (
     BasePartialChargeSettings,
 )
@@ -753,6 +758,46 @@ def run(
         alchem_comps = self._inputs["alchemical_components"]
         solvent_comp, protein_comp, small_mols = self._get_components(stateA, stateB)
 
+        alchemical = set(alchem_comps["stateA"]) | set(alchem_comps["stateB"])
+
+        def _unique(stem: str, used: set[str]) -> str:
+            for i in range(1, 10):
+                candidate = f"{stem}{i}"
+                if candidate not in used:
+                    return candidate
+            raise ValueError(
+                f"Could not assign a unique residue name with stem {stem!r}; "
+                "too many colliding names."
+            )
+
+        # Seed with user-provided resnames so auto-assigned names avoid them.
+        used: set[str] = set()
+        for offmol in small_mols.values():
+            name = _get_offmol_resname(offmol)
+            if name is not None:
+                used.add(name)
+
+        lig_name = "LIG" if "LIG" not in used else _unique("LG", used)
+
+        resnum = 1
+        for smc, offmol in small_mols.items():
+            if _get_offmol_resname(offmol) is not None:
+                continue
+            if smc in alchemical:
+                name = lig_name
+            else:
+                name = "COF" if "COF" not in used else _unique("CO", used)
+            _set_offmol_resname(offmol, name)
+            _set_offmol_metadata(offmol, "residue_number", resnum)
+            resnum += 1
+
+        names: set[str] = set()
+        for comp in alchemical:
+            name = _get_offmol_resname(small_mols[comp])
+            assert name is not None
+            names.add(name)
+        alchem_resnames = sorted(names)
+
         # Assign partial charges now to avoid any discrepancies later
         self._assign_partial_charges(settings["charge_settings"], small_mols)
 
@@ -810,6 +855,7 @@ def run(
             "positions": positions_outfile,
             "pdb_structure": self.shared_basepath / settings["output_settings"].output_structure,
             "selection_indices": selection_indices,
+            "alchemical_resnames": alchem_resnames,
         }
 
         if dry:
@@ -1505,6 +1551,7 @@ def _structural_analysis(
         trj_file: pathlib.Path,
         output_directory: pathlib.Path,
         dry: bool,
+        ligand_resnames: list[str],
     ) -> dict[str, str | pathlib.Path]:
         """
         Run structural analysis using ``openfe-analysis``.
@@ -1520,6 +1567,8 @@ def _structural_analysis(
           will be stored.
         dry : bool
           Whether or not we are running a dry run.
+        ligand_resnames: list[str]
+          The residue names of the ligands.
 
         Returns
         -------
@@ -1535,7 +1584,9 @@ def _structural_analysis(
         from openfe_analysis import rmsd
 
         try:
-            data = rmsd.gather_rms_data(pdb_file, trj_file)
+            data = rmsd.gather_rms_data(
+                pdb_file, trj_file, ligand_selection="resname " + " ".join(ligand_resnames)
+            )
         # TODO: eventually change this to more specific exception types
         except Exception as e:
             return {"structural_analysis_error": str(e)}
@@ -1573,6 +1624,7 @@ def run(
         pdb_file: pathlib.Path,
         trajectory: pathlib.Path,
         checkpoint: pathlib.Path,
+        ligand_resnames: list[str],
         dry: bool = False,
         verbose: bool = True,
         scratch_basepath: pathlib.Path | None = None,
@@ -1588,6 +1640,8 @@ def run(
           Path to the MultiStateReporter generated NetCDF file.
         checkpoint : pathlib.Path
           Path to the checkpoint file generated by MultiStateReporter.
+        ligand_resnames: list[str]
+          The residue names of the ligands.
         dry : bool
           Do a dry run of the calculation, creating all necessary hybrid
           system components (topology, system, sampler, etc...) but without
@@ -1641,6 +1695,7 @@ def run(
             trj_file=trajectory,
             output_directory=self.shared_basepath,
             dry=dry,
+            ligand_resnames=ligand_resnames,
         )
 
         # Return relevant things
@@ -1664,11 +1719,13 @@ def _execute(
         selection_indices = setup_results.outputs["selection_indices"]
         trajectory = simulation_results.outputs["nc"]
         checkpoint = simulation_results.outputs["checkpoint"]
+        ligand_resnames = setup_results.outputs["alchemical_resnames"]
 
         outputs = self.run(
             pdb_file=pdb_file,
             trajectory=trajectory,
             checkpoint=checkpoint,
+            ligand_resnames=ligand_resnames,
             scratch_basepath=ctx.scratch,
             shared_basepath=ctx.shared,
         )

src/openfe/protocols/openmm_utils/offmolecule_utils.py

@@ -0,0 +1,109 @@
+# This code is part of OpenFE and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/openfe
+import logging
+from typing import Any
+
+from openff.toolkit import Molecule as OFFMolecule
+
+logger = logging.getLogger(__name__)
+
+
+def _set_offmol_metadata(
+    offmol: OFFMolecule,
+    key: Any,
+    val: Any | None,
+) -> None:
+    """
+    Set a given metadata entry for a whole Molecule.
+
+    Parameters
+    ----------
+    offmol : openff.toolkit.Molecule
+      The Molecule to set the metadata for.
+    key : Any
+      The metadata key.
+    val : Any
+      The value to set the metadata entry to.
+    """
+    if val is None:
+        for a in offmol.atoms:
+            a.metadata.pop(key, None)
+    else:
+        for a in offmol.atoms:
+            a.metadata[key] = val
+
+
+def _get_offmol_metadata(offmol: OFFMolecule, key: Any) -> Any | None:
+    """
+    Get an offmol's given metadata entry and make sure it is
+    consistent across all atoms in the Molecule.
+
+    Parameters
+    ----------
+    offmol : openff.toolkit.Molecule
+      Molecule to get the metadata value from.
+    key: Any
+      The metadata entry key.
+
+    Returns
+    -------
+    value : Any | None
+      Metadata for the given key in the molecule. ``None`` if the
+      Molecule does not have that metadata entry set, or if
+      the value is inconsistent across all the atoms.
+    """
+    value: Any | None = None
+    for a in offmol.atoms:
+        if value is None:
+            try:
+                value = a.metadata[key]
+            except KeyError:
+                return None
+
+        if value != a.metadata[key]:
+            wmsg = f"Inconsistent metadata {key} in OFFMol: {offmol}"
+            logger.warning(wmsg)
+            return 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")

src/openfe/tests/data/openmm_rfe/benzene_toluene_hybrid_top/hybrid_topology_atoms.csv

@@ -1,17 +1,17 @@
 ,serial,name,element,resSeq,resName,chainID,segmentID,formal_charge
-0,,C1x,C,0,UNK,0,,
-1,,C2x,C,0,UNK,0,,
-2,,C3x,C,0,UNK,0,,
-3,,C4x,C,0,UNK,0,,
-4,,C5x,C,0,UNK,0,,
-5,,C6x,C,0,UNK,0,,
-6,,H1x,H,0,UNK,0,,
-7,,H2x,H,0,UNK,0,,
-8,,H3x,H,0,UNK,0,,
-9,,H4x,H,0,UNK,0,,
-10,,H5x,H,0,UNK,0,,
-11,,H6x,H,0,UNK,0,,
-12,,H1x,H,0,UNK,0,,
-13,,H2x,H,0,UNK,0,,
-14,,C1x,C,0,UNK,0,,
-15,,H3x,H,0,UNK,0,,
+0,,C1x,C,0,LIG,0,,
+1,,C2x,C,0,LIG,0,,
+2,,C3x,C,0,LIG,0,,
+3,,C4x,C,0,LIG,0,,
+4,,C5x,C,0,LIG,0,,
+5,,C6x,C,0,LIG,0,,
+6,,H1x,H,0,LIG,0,,
+7,,H2x,H,0,LIG,0,,
+8,,H3x,H,0,LIG,0,,
+9,,H4x,H,0,LIG,0,,
+10,,H5x,H,0,LIG,0,,
+11,,H6x,H,0,LIG,0,,
+12,,H1x,H,0,LIG,0,,
+13,,H2x,H,0,LIG,0,,
+14,,C1x,C,0,LIG,0,,
+15,,H3x,H,0,LIG,0,,