Merge branch 'main' into maint/pydantic_3.13

Author: atravitz

.github/workflows/mypy.yaml

@@ -1,8 +1,6 @@
 name: "PR: mypy static type checking"
 on:
   pull_request:
-    branches:
-      - main
   push:
     branches:
       - main

docs/conf.py

@@ -148,6 +148,7 @@
     ],
     "accent_color": "cantina-purple",
     "navigation_with_keys": False,
+    "announcement": "The OpenFE team is conducting user interviews until August 15th! <a href=https://app.reclaim.ai/m/james-omsf/openfe-user-interview>Click here</a> to schedule an interview.",
 }
 html_logo = "_static/OFE-color-icon.svg"
 html_favicon = "_static/OFE-color-icon.svg"

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/conftest.py

@@ -223,7 +223,7 @@ def benzene_toluene_topology():
 
 
 @pytest.fixture(scope="session")
-def benzene_modifications():
+def benzene_modifications_uncharged():
     files = {}
     with resources.as_file(resources.files("openfe.tests.data")) as d:
         fn = str(d / "benzene_modifications.sdf")
@@ -233,6 +233,17 @@ def benzene_modifications():
     return files
 
 
+@pytest.fixture(scope="session")
+def benzene_modifications():
+    files = {}
+    with resources.as_file(resources.files("openfe.tests.data")) as d:
+        fn = str(d / "benzene_modifications_am1bcc.sdf")
+        supp = Chem.SDMolSupplier(str(fn), removeHs=False)
+        for rdmol in supp:
+            files[rdmol.GetProp("_Name")] = SmallMoleculeComponent(rdmol)
+    return files
+
+
 @pytest.fixture(scope="session")
 def charged_benzene_modifications():
     files = {}
@@ -266,7 +277,7 @@ def benzene_transforms():
     # a dict of Molecules for benzene transformations
     mols = {}
     with resources.as_file(resources.files("openfe.tests.data")) as d:
-        fn = str(d / "benzene_modifications.sdf")
+        fn = str(d / "benzene_modifications_am1bcc.sdf")
         supplier = Chem.SDMolSupplier(fn, removeHs=False)
         for mol in supplier:
             mols[mol.GetProp("_Name")] = SmallMoleculeComponent(mol)
@@ -302,22 +313,36 @@ def eg5_protein_pdb():
 
 
 @pytest.fixture()
-def eg5_ligands_sdf():
+def eg5_ligands_uncharged_sdf():
     with resources.as_file(resources.files("openfe.tests.data.eg5")) as d:
         yield str(d / "eg5_ligands.sdf")
 
 
+@pytest.fixture()
+def eg5_ligands_sdf():
+    with resources.as_file(resources.files("openfe.tests.data.eg5")) as d:
+        yield str(d / "eg5_ligands_am1bcc.sdf")
+
+
 @pytest.fixture()
 def eg5_cofactor_sdf():
     with resources.as_file(resources.files("openfe.tests.data.eg5")) as d:
-        yield str(d / "eg5_cofactor.sdf")
+        yield str(d / "eg5_cofactor_am1bcc.sdf")
 
 
 @pytest.fixture()
 def eg5_protein(eg5_protein_pdb) -> openfe.ProteinComponent:
     return openfe.ProteinComponent.from_pdb_file(eg5_protein_pdb)
 
 
+@pytest.fixture()
+def eg5_ligands_uncharged(eg5_ligands_uncharged_sdf) -> list[SmallMoleculeComponent]:
+    return [
+        SmallMoleculeComponent(m)
+        for m in Chem.SDMolSupplier(eg5_ligands_uncharged_sdf, removeHs=False)
+    ]
+
+
 @pytest.fixture()
 def eg5_ligands(eg5_ligands_sdf) -> list[SmallMoleculeComponent]:
     return [SmallMoleculeComponent(m) for m in Chem.SDMolSupplier(eg5_ligands_sdf, removeHs=False)]