Enable user defn boresch restraints in ABFE Protocol

news/abfe_user_dfn_restraints.rst

@@ -0,0 +1,24 @@
+**Added:**
+
+* Added support for user-defined Boresch restraints in the
+  ABFE Protocol (`PR #2019 <https://github.com/OpenFreeEnergy/openfe/pull/2019>`_).
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/openfe/protocols/openmm_afe/abfe_units.py

@@ -26,7 +26,7 @@
 from rdkit import Chem
 
 from openfe.protocols.openmm_afe.equil_afe_settings import (
-    BoreschRestraintSettings,
+    ABFEBoreschRestraintSettings,
     SettingsBaseModel,
 )
 from openfe.protocols.openmm_utils import system_validation
@@ -218,7 +218,7 @@ def _get_boresch_restraint(
         guest_atom_ids: list[int],
         host_atom_ids: list[int],
         temperature: Quantity,
-        settings: BoreschRestraintSettings,
+        settings: ABFEBoreschRestraintSettings,
     ) -> tuple[BoreschRestraintGeometry, BoreschRestraint]:
         """
         Get a Boresch-like restraint Geometry and OpenMM restraint force
@@ -236,7 +236,7 @@ def _get_boresch_restraint(
           A list of atom indices defining the host molecules in the universe.
         temperature : openff.units.Quantity
           The temperature of the simulation where the restraint will be added.
-        settings : BoreschRestraintSettings
+        settings : ABFEBoreschRestraintSettings
           Settings on how the Boresch-like restraint should be defined.
 
         Returns
@@ -254,6 +254,12 @@ def _get_boresch_restraint(
             guest_rdmol=guest_rdmol,
             guest_idxs=guest_atom_ids,
             host_idxs=host_atom_ids,
+            guest_restraint_atoms_idxs=list(settings.guest_restraint_ids)
+            if settings.guest_restraint_ids is not None
+            else None,
+            host_restraint_atoms_idxs=list(settings.host_restraint_ids)
+            if settings.host_restraint_ids is not None
+            else None,
             host_selection=settings.host_selection,
             anchor_finding_strategy=settings.anchor_finding_strategy,
             dssp_filter=settings.dssp_filter,
@@ -355,7 +361,7 @@ def _add_restraints(
             self.shared_basepath / settings["equil_output_settings"].production_trajectory_filename,
         )
 
-        if isinstance(settings["restraint_settings"], BoreschRestraintSettings):
+        if isinstance(settings["restraint_settings"], ABFEBoreschRestraintSettings):
             rest_geom, restraint = self._get_boresch_restraint(
                 univ,
                 guest_rdmol,

src/openfe/protocols/openmm_afe/equil_afe_settings.py

@@ -213,6 +213,29 @@ def must_be_all(cls, v):
         return v
 
 
+class ABFEBoreschRestraintSettings(BoreschRestraintSettings):
+    host_restraint_ids: tuple[int, int, int] | None = None
+    """
+    The indices of the host component atoms to restrain.
+    The entries define the H0, H1, and H2 atoms in order.
+    If defined, these will override any automatic selection.
+    """
+    guest_restraint_ids: tuple[int, int, int] | None = None
+    """
+    The indices of the guest component atoms to restraint.
+    The entries define the G0, G1, and G2 atoms in order.
+    If defined, these will override any automatic selection.
+    """
+
+    @field_validator("guest_restraint_ids", "host_restraint_ids")
+    def positive_idxs_three_tuple(cls, v):
+        if v is not None:
+            if any([i < 0 for i in v]):
+                errmsg = "``guest_atoms`` and ``host_atoms`` cannot have negative indices."
+                raise ValueError(errmsg)
+        return v
+
+
 # This subclasses from SettingsBaseModel as it has vacuum_forcefield and
 # solvent_forcefield fields, not just a single forcefield_settings field
 class AbsoluteSolvationSettings(SettingsBaseModel):

src/openfe/protocols/openmm_afe/equil_binding_afe_method.py

@@ -43,10 +43,10 @@
 
 from openfe.due import Doi, due
 from openfe.protocols.openmm_afe.equil_afe_settings import (
+    ABFEBoreschRestraintSettings,
     ABFEPreEquilOutputSettings,
     AbsoluteBindingSettings,
     AlchemicalSettings,
-    BoreschRestraintSettings,
     IntegratorSettings,
     LambdaSettings,
     MDSimulationSettings,
@@ -174,7 +174,7 @@ def _default_settings(cls):
             engine_settings=OpenMMEngineSettings(),
             solvent_integrator_settings=IntegratorSettings(),
             complex_integrator_settings=IntegratorSettings(),
-            restraint_settings=BoreschRestraintSettings(),
+            restraint_settings=ABFEBoreschRestraintSettings(),
             solvent_equil_simulation_settings=MDSimulationSettings(
                 equilibration_length_nvt=0.1 * offunit.nanosecond,
                 equilibration_length=0.2 * offunit.nanosecond,

src/openfe/protocols/restraint_utils/settings.py

@@ -174,17 +174,6 @@ class BoreschRestraintSettings(BaseRestraintSettings):
     Boresch-like restraint search parameter.
     The maximum distance between any host atom and the guest G0 atom. Must be in units compatible with nanometer.
     """
-    # TODO: re-enable this (Issue #1555)
-    # host_atoms: Optional[list[int]] = None
-    # """
-    # The indices of the host component atoms to restrain.
-    # If defined, these will override any automatic selection.
-    # """
-    # guest_atoms: Optional[list[int]] = None
-    # """
-    # The indices of the guest component atoms to restraint.
-    # If defined, these will override any automatic selection.
-    # """
     anchor_finding_strategy: Literal["multi-residue", "bonded"] = "bonded"
     """
     The Boresch atom picking strategy to use.
@@ -193,11 +182,3 @@ class BoreschRestraintSettings(BaseRestraintSettings):
       * `bonded`: pick host atoms that are bonded to each other.
       * `multi-residue`: pick host atoms which can span multiple residues.
     """
-
-
-#     @field_validator("guest_atoms", "host_atoms")
-#     def positive_idxs_list(cls, v):
-#         if v is not None and any([i < 0 for i in v]):
-#             errmsg = "negative indices passed"
-#             raise ValueError(errmsg)
-#         return v

src/openfe/tests/conftest.py

@@ -233,6 +233,17 @@ def benzene_modifications():
     return files
 
 
+@pytest.fixture(scope="session")
+def benzene_modifications_am1bcc():
+    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 = {}