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 LG1 amd LG2 (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,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,
 )
@@ -692,9 +693,14 @@ def _subsample_topology(
         bfactors[np.isin(selection_indices, list(atom_classes["unique_new_atoms"]))] = 0.75
 
         if len(selection_indices) > 0:
+            sub_top = hybrid_topology.subset(selection_indices)
+            # Renumber sequentially so same-named residues (e.g. two COF cofactors,
+            # both resSeq=0) don't merge into one residue on PDB write/read.
+            for i, res in enumerate(sub_top.residues, start=1):
+                res.resSeq = i
             traj = mdt.Trajectory(
                 hybrid_positions[selection_indices, :],
-                hybrid_topology.subset(selection_indices),
+                sub_top,
             ).save_pdb(
                 self.shared_basepath / output_filename,
                 bfactors=bfactors,
@@ -753,6 +759,47 @@ 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"])
+        stateA_comps = set(alchem_comps["stateA"])
+
+        def _unique(used: set[str]) -> str:
+            """Next free 'LG{n}' name (n = 1..9) for the alchemical ligands."""
+            for i in range(1, 10):
+                candidate = f"LG{i}"
+                if candidate not in used:
+                    return candidate
+            raise ValueError(
+                "Could not assign a unique ligand residue name; too many colliding 'LG#' 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)
+
+        # ligands take LG1 (endstate A) and LG2 (endstate B) by default, only
+        # bumping to LG3+ if a user pre-named something LG1/LG2. All cofactors
+        # share "COF" and are distinguished by residue index, not name.
+        for smc, offmol in small_mols.items():
+            if _get_offmol_resname(offmol) is not None:
+                continue
+            if smc in alchemical:
+                base = "LG1" if smc in stateA_comps else "LG2"
+                name = base if base not in used else _unique(used)
+                used.add(name)
+            else:
+                name = "COF"
+            _set_offmol_resname(offmol, name)
+
+        names: set[str] = set()
+        for comp in (mapping.componentA, mapping.componentB):
+            name = _get_offmol_resname(small_mols[comp])
+            assert name is not None  # for typing reasons, cannot be 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 +857,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 +1553,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 +1569,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 +1586,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 +1626,7 @@ def run(
         pdb_file: pathlib.Path,
         trajectory: pathlib.Path,
         checkpoint: pathlib.Path,
+        ligand_resnames: list[str] = ["LG1", "LG2"],
         dry: bool = False,
         verbose: bool = True,
         scratch_basepath: pathlib.Path | None = None,
@@ -1588,6 +1642,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. Default: ["LG1", "LG2"]
         dry : bool
           Do a dry run of the calculation, creating all necessary hybrid
           system components (topology, system, sampler, etc...) but without
@@ -1641,6 +1697,7 @@ def run(
             trj_file=trajectory,
             output_directory=self.shared_basepath,
             dry=dry,
+            ligand_resnames=ligand_resnames,
         )
 
         # Return relevant things
@@ -1669,6 +1726,7 @@ def _execute(
             pdb_file=pdb_file,
             trajectory=trajectory,
             checkpoint=checkpoint,
+            ligand_resnames=setup_results.outputs.get("alchemical_resnames", ["LG1", "LG2"]),
             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,LG1,0,,
+1,,C2x,C,0,LG1,0,,
+2,,C3x,C,0,LG1,0,,
+3,,C4x,C,0,LG1,0,,
+4,,C5x,C,0,LG1,0,,
+5,,C6x,C,0,LG1,0,,
+6,,H1x,H,0,LG1,0,,
+7,,H2x,H,0,LG1,0,,
+8,,H3x,H,0,LG1,0,,
+9,,H4x,H,0,LG1,0,,
+10,,H5x,H,0,LG1,0,,
+11,,H6x,H,0,LG1,0,,
+12,,H1x,H,0,LG1,0,,
+13,,H2x,H,0,LG1,0,,
+14,,C1x,C,0,LG1,0,,
+15,,H3x,H,0,LG1,0,,