Add AbsoluteFactory options to AlchemicalSettings (#1742)

* Add absolute alchemical settings

* Expose settings in AFE protocols

* expose settings in septop

* Fix up settings

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add some tests

* Add septop tests for sterics force settings

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add news item

* add an extra parameterize option

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: IAlibay

news/absolute_settings.rst

@@ -0,0 +1,27 @@
+**Added:**
+
+* New options have been added to the ``AlchemicalSettings``
+  of the ``SepTopProtocol``, ``AbsoluteSolvationProtocol``
+  and ``AbsoluteBindingProtocol``. Notably, these options allow users to
+  control the softcore parameters as well as the use of
+  long range dispersion corrections.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

openfe/protocols/openmm_afe/base.py

@@ -36,7 +36,7 @@
 from openff.units import Quantity, unit
 from openff.units.openmm import ensure_quantity, from_openmm, to_openmm
 from openmm import app
-from openmm import unit as omm_unit
+from openmm import unit as ommunit
 from openmmforcefields.generators import SystemGenerator
 from openmmtools import multistate
 from openmmtools.alchemy import (
@@ -52,6 +52,7 @@
 )
 
 from openfe.protocols.openmm_afe.equil_afe_settings import (
+    AlchemicalSettings,
     BaseSolvationSettings,
     IntegratorSettings,
     MultiStateOutputSettings,
@@ -168,10 +169,10 @@ def _pre_equilibrate(
         self,
         system: openmm.System,
         topology: openmm.app.Topology,
-        positions: omm_unit.Quantity,
+        positions: ommunit.Quantity,
         settings: dict[str, SettingsBaseModel],
         dry: bool,
-    ) -> tuple[omm_unit.Quantity, omm_unit.Quantity]:
+    ) -> tuple[ommunit.Quantity, ommunit.Quantity]:
         """
         Run a non-alchemical equilibration to get a stable system.
 
@@ -619,6 +620,7 @@ def _get_alchemical_system(
         system: openmm.System,
         comp_resids: dict[Component, npt.NDArray],
         alchem_comps: dict[str, list[Component]],
+        alchemical_settings: AlchemicalSettings,
     ) -> tuple[AbsoluteAlchemicalFactory, openmm.System, list[int]]:
         """
         Get an alchemically modified system and its associated factory
@@ -633,6 +635,8 @@ def _get_alchemical_system(
           A dictionary of residues for each component in the System.
         alchem_comps : dict[str, list[Component]]
           A dictionary of alchemical components for each end state.
+        alchemical_settings : AlchemicalSettings
+          Settings controlling how the alchemical system is built.
 
         Returns
         -------
@@ -652,9 +656,26 @@ def _get_alchemical_system(
 
         alchemical_region = AlchemicalRegion(
             alchemical_atoms=alchemical_indices,
+            softcore_alpha=alchemical_settings.softcore_alpha,
+            annihilate_electrostatics=True,
+            annihilate_sterics=alchemical_settings.annihilate_sterics,
+            softcore_a=alchemical_settings.softcore_a,
+            softcore_b=alchemical_settings.softcore_b,
+            softcore_c=alchemical_settings.softcore_c,
+            softcore_beta=0.0,
+            softcore_d=1.0,
+            softcore_e=1.0,
+            softcore_f=2.0,
         )
 
-        alchemical_factory = AbsoluteAlchemicalFactory()
+        alchemical_factory = AbsoluteAlchemicalFactory(
+            consistent_exceptions=False,
+            switch_width=1.0 * ommunit.angstroms,
+            alchemical_pme_treatment="exact",
+            alchemical_rf_treatment="switched",
+            disable_alchemical_dispersion_correction=alchemical_settings.disable_alchemical_dispersion_correction,
+            split_alchemical_forces=True,
+        )
         alchemical_system = alchemical_factory.create_alchemical_system(system, alchemical_region)
 
         return alchemical_factory, alchemical_system, alchemical_indices
@@ -1141,7 +1162,11 @@ def run(
 
         # 7. Get alchemical system
         alchem_factory, alchem_system, alchem_indices = self._get_alchemical_system(
-            omm_topology, restrained_omm_system, comp_resids, alchem_comps
+            topology=omm_topology,
+            system=restrained_omm_system,
+            comp_resids=comp_resids,
+            alchem_comps=alchem_comps,
+            alchemical_settings=settings["alchemical_settings"],
         )
 
         # 8. Get compound and sampler states

openfe/protocols/openmm_afe/equil_afe_settings.py

@@ -42,9 +42,48 @@
 
 
 class AlchemicalSettings(SettingsBaseModel):
-    """Settings for the alchemical protocol
+    """
+    Alchemical settings for Protocols which use the
+    AbsoluteAlchemicalFactory.
+    """
+
+    disable_alchemical_dispersion_correction: bool = False
+    """
+    If True, the long-range dispersion correction will not
+    be included for the alchemical region, avoiding the need
+    to recompute the correction. This can improve performance,
+    at the cost of accuracy. Default is False.
+    """
+    annihilate_sterics: bool = False
+    """
+    If True, sterics (Lennard-Jones) will be annhilated instead
+    of decoupled. Default is False.
+    """
+    softcore_alpha: float = 0.5
+    """
+    Alchemical softcore parameter for the Lennard-Jones interactions
+    (default is 0.5).
 
-    Empty place holder for right now.
+    The generalized softcore potential formalism introduced by
+    Pham and Shirts, J. Chem. Phys. 135, 034114 (2011), equation 13,
+    is used here. The ``softcore_a``, ``softcore_b``, and
+    ``softcore_c`` parameters are used alongside ``softcore_alpha``
+    to control how the potential is scaled.
+    """
+    softcore_a: float = 1.0
+    """
+    Scaling constant ``a`` in
+    Eq. 13 from Pham and Shirts, J. Chem. Phys. 135, 034114 (2011).
+    """
+    softcore_b: float = 1.0
+    """
+    Scaling constant ``b`` in
+    Eq. 13 from Pham and Shirts, J. Chem. Phys. 135, 034114 (2011).
+    """
+    softcore_c: float = 6.0
+    """
+    Scaling constant ``c`` in
+    Eq. 13 from Pham and Shirts, J. Chem. Phys. 135, 034114 (2011).
     """
 
 

openfe/protocols/openmm_septop/base.py

@@ -45,6 +45,7 @@
 )
 
 from openfe.protocols.openmm_afe.equil_afe_settings import (
+    AlchemicalSettings,
     BaseSolvationSettings,
     IntegratorSettings,
     MultiStateOutputSettings,
@@ -285,6 +286,7 @@ def _get_alchemical_system(
         system: openmm.System,
         alchem_indices_A: list[int],
         alchem_indices_B: list[int],
+        alchemical_settings: AlchemicalSettings,
     ) -> tuple[AbsoluteAlchemicalFactory, openmm.System]:
         """
         Get an alchemically modified system and its associated factory
@@ -293,12 +295,14 @@ def _get_alchemical_system(
         ----------
         system : openmm.System
           System to alchemically modify.
-        alchem_indices_A: list[int]
+        alchem_indices_A : list[int]
           A list of atom indices for the alchemically modified
           ligand A in the system.
-        alchem_indices_B: list[int]
+        alchem_indices_B : list[int]
           A list of atom indices for the alchemically modified
           ligand B in the system.
+        alchemical_settings : AlchemicalSettings
+          Settings controlling how the alchemical system will be built.
 
         Returns
         -------
@@ -313,13 +317,42 @@ def _get_alchemical_system(
             switch_width=1.0 * unit.angstroms,
             alchemical_pme_treatment="exact",
             alchemical_rf_treatment="switched",
-            disable_alchemical_dispersion_correction=False,
+            disable_alchemical_dispersion_correction=alchemical_settings.disable_alchemical_dispersion_correction,
             split_alchemical_forces=True,
         )
+
         # Alchemical Region for ligand A
-        alchemical_region_A = AlchemicalRegion(alchemical_atoms=alchem_indices_A, name="A")
+        alchemical_region_A = AlchemicalRegion(
+            alchemical_atoms=alchem_indices_A,
+            name="A",
+            softcore_alpha=alchemical_settings.softcore_alpha,
+            annihilate_electrostatics=True,
+            annihilate_sterics=alchemical_settings.annihilate_sterics,
+            softcore_a=alchemical_settings.softcore_a,
+            softcore_b=alchemical_settings.softcore_b,
+            softcore_c=alchemical_settings.softcore_c,
+            softcore_beta=0.0,
+            softcore_d=1.0,
+            softcore_e=1.0,
+            softcore_f=2.0,
+        )
+
         # Alchemical Region for ligand B
-        alchemical_region_B = AlchemicalRegion(alchemical_atoms=alchem_indices_B, name="B")
+        alchemical_region_B = AlchemicalRegion(
+            alchemical_atoms=alchem_indices_B,
+            name="B",
+            softcore_alpha=alchemical_settings.softcore_alpha,
+            annihilate_electrostatics=True,
+            annihilate_sterics=alchemical_settings.annihilate_sterics,
+            softcore_a=alchemical_settings.softcore_a,
+            softcore_b=alchemical_settings.softcore_b,
+            softcore_c=alchemical_settings.softcore_c,
+            softcore_beta=0.0,
+            softcore_d=1.0,
+            softcore_e=1.0,
+            softcore_f=2.0,
+        )
+
         alchemical_system = alchemical_factory.create_alchemical_system(
             system, [alchemical_region_A, alchemical_region_B]
         )

openfe/protocols/openmm_septop/equil_septop_method.py

@@ -1971,6 +1971,7 @@ def run(
             omm_system_AB,
             atom_indices_AB_A,
             atom_indices_AB_B,
+            settings["alchemical_settings"],
         )
 
         # 10. Apply Restraints
@@ -2254,6 +2255,7 @@ def run(
             omm_system_AB,
             atom_indices_AB_A,
             atom_indices_AB_B,
+            settings["alchemical_settings"],
         )
 
         # 8. Apply Restraints

openfe/protocols/openmm_septop/equil_septop_settings.py

@@ -23,6 +23,7 @@
 from openff.units import unit as offunit
 from pydantic import field_validator
 
+from openfe.protocols.openmm_afe.equil_afe_settings import AlchemicalSettings
 from openfe.protocols.openmm_utils.omm_settings import (
     IntegratorSettings,
     MDOutputSettings,
@@ -36,13 +37,6 @@
 from openfe.protocols.restraint_utils.settings import BaseRestraintSettings
 
 
-class AlchemicalSettings(SettingsBaseModel):
-    """Settings for the alchemical protocol
-
-    Empty place holder for right now.
-    """
-
-
 class LambdaSettings(SettingsBaseModel):
     """Lambda schedule settings.