Fix residue naming

news/fix_unk_resnames.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* Small molecules in RelativeHybridTopology topologies (including the output PDB) are now named LIG (alchemical ligand) and CF1, CF2… (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,7 @@
 from openmmtools import multistate
 
 import openfe
+from openfe.protocols.openmm_utils.offmolecule_utils import _get_offmol_resname, _set_offmol_resname
 from openfe.protocols.openmm_utils.omm_settings import (
     BasePartialChargeSettings,
 )
@@ -753,6 +754,20 @@ 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"])
+        cofactors = [smc for smc in small_mols if smc not in alchemical]
+
+        for smc, offmol in small_mols.items():
+            if _get_offmol_resname(offmol) is not None:
+                continue
+            if smc in alchemical:
+                _set_offmol_resname(offmol, "LIG")
+            else:
+                _set_offmol_resname(offmol, f"CF{cofactors.index(smc) + 1}")
+
+        # Record the resname the alchemical ligand actually ended up with
+        alchem_resname = _get_offmol_resname(small_mols[mapping.componentA])
+
         # Assign partial charges now to avoid any discrepancies later
         self._assign_partial_charges(settings["charge_settings"], small_mols)
 
@@ -810,6 +825,7 @@ def run(
             "positions": positions_outfile,
             "pdb_structure": self.shared_basepath / settings["output_settings"].output_structure,
             "selection_indices": selection_indices,
+            "alchemical_resname": alchem_resname,
         }
 
         if dry:
@@ -1505,6 +1521,7 @@ def _structural_analysis(
         trj_file: pathlib.Path,
         output_directory: pathlib.Path,
         dry: bool,
+        ligand_resname: str = "LIG",
     ) -> dict[str, str | pathlib.Path]:
         """
         Run structural analysis using ``openfe-analysis``.
@@ -1520,6 +1537,8 @@ def _structural_analysis(
           will be stored.
         dry : bool
           Whether or not we are running a dry run.
+        ligand_resname: str
+          The residue name of the hybrid topology ligand. Default: LIG
 
         Returns
         -------
@@ -1535,7 +1554,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=f"resname {ligand_resname}"
+            )
         # TODO: eventually change this to more specific exception types
         except Exception as e:
             return {"structural_analysis_error": str(e)}
@@ -1573,6 +1594,7 @@ def run(
         pdb_file: pathlib.Path,
         trajectory: pathlib.Path,
         checkpoint: pathlib.Path,
+        ligand_resname: str = "LIG",
         dry: bool = False,
         verbose: bool = True,
         scratch_basepath: pathlib.Path | None = None,
@@ -1588,6 +1610,8 @@ def run(
           Path to the MultiStateReporter generated NetCDF file.
         checkpoint : pathlib.Path
           Path to the checkpoint file generated by MultiStateReporter.
+        ligand_resname: str
+          The residue name of the hybrid topology ligand. Default: LIG
         dry : bool
           Do a dry run of the calculation, creating all necessary hybrid
           system components (topology, system, sampler, etc...) but without
@@ -1641,6 +1665,7 @@ def run(
             trj_file=trajectory,
             output_directory=self.shared_basepath,
             dry=dry,
+            ligand_resname=ligand_resname,
         )
 
         # Return relevant things
@@ -1669,6 +1694,9 @@ def _execute(
             pdb_file=pdb_file,
             trajectory=trajectory,
             checkpoint=checkpoint,
+            ligand_resname=setup_results.outputs.get(
+                "alchemical_resname", "LIG"
+            ),  # Adding this for backward compatibility
             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,,