Adding LambdaSettings object for AFE

devtools/data/gen-serialized-results.py

@@ -74,8 +74,12 @@ def generate_ahfe_json(smc):
     settings.solvent_simulation_settings.production_length = 500 * unit.picosecond
     settings.vacuum_simulation_settings.equilibration_length = 10 * unit.picosecond
     settings.vacuum_simulation_settings.production_length = 1000 * unit.picosecond
-    settings.alchemical_settings.lambda_elec_windows = 5
-    settings.alchemical_settings.lambda_vdw_windows = 9
+    settings.lambda_settings.lambda_elec = [0.0, 0.25, 0.5, 0.75, 1.0, 1.0,
+                                            1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
+                                            1.0]
+    settings.lambda_settings.lambda_vdw = [0.0, 0.0, 0.0, 0.0, 0.0, 0.12, 0.24,
+                                           0.36, 0.48, 0.6, 0.7, 0.77, 0.85,
+                                           1.0]
     settings.alchemsampler_settings.n_repeats = 3
     settings.alchemsampler_settings.n_replicas = 14
     settings.alchemsampler_settings.online_analysis_target_error = 0.2 * unit.boltzmann_constant * unit.kelvin

docs/reference/api/openmm_rfe.rst

@@ -80,6 +80,17 @@ Below are the settings which can be tweaked in the protocol. The default setting
    :inherited-members: SettingsBaseModel
    :member-order: bysource
 
+.. autopydantic_model:: LambdaSettings
+   :model-show-json: False
+   :model-show-field-summary: False
+   :model-show-config-member: False
+   :model-show-config-summary: False
+   :model-show-validator-members: False
+   :model-show-validator-summary: False
+   :field-list-validators: False
+   :inherited-members: SettingsBaseModel
+   :member-order: bysource
+
 .. autopydantic_model:: OpenMMEngineSettings
    :model-show-json: False
    :model-show-field-summary: False

docs/reference/api/openmm_solvation_afe.rst

@@ -85,6 +85,18 @@ Below are the settings which can be tweaked in the protocol. The default setting
    :member-order: bysource
    :noindex:
 
+.. autopydantic_model:: LambdaSettings
+   :model-show-json: False
+   :model-show-field-summary: False
+   :model-show-config-member: False
+   :model-show-config-summary: False
+   :model-show-validator-members: False
+   :model-show-validator-summary: False
+   :field-list-validators: False
+   :inherited-members: SettingsBaseModel
+   :member-order: bysource
+   :noindex:
+
 .. autopydantic_model:: OpenMMEngineSettings
    :model-show-json: False
    :model-show-field-summary: False

openfe/protocols/openmm_afe/base.py

@@ -53,7 +53,7 @@
 from openfe.protocols.openmm_afe.equil_afe_settings import (
     SolvationSettings,
     AlchemicalSamplerSettings, OpenMMEngineSettings,
-    IntegratorSettings, SimulationSettings,
+    IntegratorSettings, SimulationSettings, LambdaSettings,
 )
 from openfe.protocols.openmm_rfe._rfe_utils import compute
 from ..openmm_utils import (
@@ -233,6 +233,7 @@ def _handle_settings(self):
           * system_settings : SystemSettings
           * solvation_settings : SolvationSettings
           * alchemical_settings : AlchemicalSettings
+          * lambda_settings : LambdaSettings
           * sampler_settings : AlchemicalSamplerSettings
           * engine_settings : OpenMMEngineSettings
           * integrator_settings : IntegratorSettings
@@ -407,21 +408,16 @@ def _get_lambda_schedule(
           LambdaProtocol
         """
         lambdas = dict()
-        n_elec = settings['alchemical_settings'].lambda_elec_windows
-        n_vdw = settings['alchemical_settings'].lambda_vdw_windows + 1
-        lambdas['lambda_electrostatics'] = np.concatenate(
-                [np.linspace(1, 0, n_elec), np.linspace(0, 0, n_vdw)[1:]]
-        )
-        lambdas['lambda_sterics'] = np.concatenate(
-                [np.linspace(1, 1, n_elec), np.linspace(1, 0, n_vdw)[1:]]
-        )
 
-        n_replicas = settings['sampler_settings'].n_replicas
+        lambda_elec = settings['lambda_settings'].lambda_elec
+        lambda_vdw = settings['lambda_settings'].lambda_vdw
 
-        if n_replicas != (len(lambdas['lambda_sterics'])):
-            errmsg = (f"Number of replicas {n_replicas} "
-                      "does not equal the number of lambda windows ")
-            raise ValueError(errmsg)
+        # Reverse lambda schedule since in AbsoluteAlchemicalFactory 1
+        # means fully interacting, not stateB
+        lambda_elec = [1-x for x in lambda_elec]
+        lambda_vdw = [1-x for x in lambda_vdw]
+        lambdas['lambda_electrostatics'] = lambda_elec
+        lambdas['lambda_sterics'] = lambda_vdw
 
         return lambdas
 
@@ -762,11 +758,9 @@ def _run_simulation(
             # minimize
             if self.verbose:
                 self.logger.info("minimizing systems")
-
             sampler.minimize(
                 max_iterations=settings['simulation_settings'].minimization_steps
             )
-
             # equilibrate
             if self.verbose:
                 self.logger.info("equilibrating systems")
@@ -776,7 +770,6 @@ def _run_simulation(
             # production
             if self.verbose:
                 self.logger.info("running production phase")
-
             sampler.extend(int(prod_steps / mc_steps))  # type: ignore
 
             if self.verbose:

openfe/protocols/openmm_afe/equil_afe_settings.py

@@ -29,6 +29,7 @@
     IntegratorSettings,
     SimulationSettings
 )
+import numpy as np
 
 
 try:
@@ -40,20 +41,69 @@
 class AlchemicalSettings(SettingsBaseModel):
     """Settings for the alchemical protocol
 
-    These settings describe the lambda schedule and the creation of the
-    hybrid system.
+    Empty place holder for right now.
     """
 
-    lambda_elec_windows = 12
-    """Number of lambda electrostatic alchemical steps, default 12"""
-    lambda_vdw_windows = 12
-    """Number of lambda vdw alchemical steps, default 12"""
 
-    @validator('lambda_elec_windows', 'lambda_vdw_windows')
-    def must_be_positive(cls, v):
-        if v <= 0:
-            errmsg = ("Number of lambda steps must be positive ")
+class LambdaSettings(SettingsBaseModel):
+    """Lambda schedule settings.
+
+    Defines lists of floats to control various aspects of the alchemical
+    transformation.
+
+    Notes
+    --------
+    * In all cases a lambda value of 0 defines a fully interacting state A and
+    a non-interacting state B, whilst a value of 1 defines a fully interacting
+     state B and a non-interacting state A.
+    * ``lambda_elec``, `lambda_vdw``, and ``lambda_restraints`` must all be of
+    the same length, defining all the windows of the transformation.
+    """
+    lambda_elec: list[float] = [
+        0.0, 0.25, 0.5, 0.75, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
+        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
+    ]
+    """
+    List of floats of lambda values for the electrostatics. 
+    Zero means state A and 1 means state B.
+    Length of this list needs to match length of lambda_vdw and lambda_restraints.
+    """
+    lambda_vdw: list[float] = [
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.05, 0.1, 0.2, 0.3, 0.4,
+        0.5, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0,
+    ]
+    """
+    List of floats of lambda values for the van der Waals.
+    Zero means state A and 1 means state B.
+    Length of this list needs to match length of lambda_elec and lambda_restraints.
+    """
+    lambda_restraints: list[float] = [
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+    ]
+    """
+    List of floats of lambda values for the restraints.
+    Zero means state A and 1 means state B.
+    Length of this list needs to match length of lambda_vdw and lambda_elec.
+    """
+
+    @validator('lambda_elec', 'lambda_vdw', 'lambda_restraints')
+    def must_be_between_0_and_1(cls, v):
+        for window in v:
+            if not 0 <= window <= 1:
+                errmsg = "Lambda windows must be between 0 and 1."
+                raise ValueError(errmsg)
+        return v
+
+    @validator('lambda_elec', 'lambda_vdw', 'lambda_restraints')
+    def must_be_monotonic(cls, v):
+
+        difference = np.diff(v)
+
+        if not all(i >= 0. for i in difference):
+            errmsg = "The lambda schedule is not monotonic."
             raise ValueError(errmsg)
+
         return v
 
 
@@ -91,11 +141,16 @@ class Config:
     # Alchemical settings
     alchemical_settings: AlchemicalSettings
     """
-    Alchemical protocol settings including lambda windows.
+    Alchemical protocol settings.
+    """
+    lambda_settings: LambdaSettings
+    """
+    Settings for controlling the lambda schedule for the different components 
+    (vdw, elec, restraints).
     """
     alchemsampler_settings: AlchemicalSamplerSettings
     """
-    Settings for controling how we sample alchemical space, including the
+    Settings for controlling how we sample alchemical space, including the
     number of repeats.
     """
 

openfe/protocols/openmm_afe/equil_solvation_afe_method.py

@@ -30,6 +30,7 @@
 
 import pathlib
 import logging
+import warnings
 from collections import defaultdict
 import gufe
 from gufe.components import Component
@@ -50,7 +51,7 @@
 )
 from openfe.protocols.openmm_afe.equil_afe_settings import (
     AbsoluteSolvationSettings, SystemSettings,
-    SolvationSettings, AlchemicalSettings,
+    SolvationSettings, AlchemicalSettings, LambdaSettings,
     AlchemicalSamplerSettings, OpenMMEngineSettings,
     IntegratorSettings, SimulationSettings,
     SettingsBaseModel,
@@ -404,8 +405,17 @@ def _default_settings(cls):
             solvent_system_settings=SystemSettings(),
             vacuum_system_settings=SystemSettings(nonbonded_method='nocutoff'),
             alchemical_settings=AlchemicalSettings(),
+            lambda_settings=LambdaSettings(
+                lambda_elec=[
+                    0.0, 0.25, 0.5, 0.75, 1.0, 1.0, 1.0,
+                    1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0
+                ],
+                lambda_vdw=[
+                    0.0, 0.0, 0.0, 0.0, 0.0, 0.12, 0.24,
+                    0.36, 0.48, 0.6, 0.7, 0.77, 0.85, 1.0],
+            ),
             alchemsampler_settings=AlchemicalSamplerSettings(
-                n_replicas=24,
+                n_replicas=14,
             ),
             solvation_settings=SolvationSettings(),
             vacuum_engine_settings=OpenMMEngineSettings(),
@@ -517,6 +527,71 @@ def _validate_alchemical_components(
                           "are not currently supported")
                 raise ValueError(errmsg)
 
+    @staticmethod
+    def _validate_lambda_schedule(
+            lambda_settings: LambdaSettings,
+            alchemsampler_settings: AlchemicalSamplerSettings,
+    ) -> None:
+        """
+        Checks that the lambda schedule is set up correctly.
+
+        Parameters
+        ----------
+        settings : AbsoluteSolvationSettings
+          Settings object.
+
+        Raises
+        ------
+        ValueError
+          If the number of lambda windows differs for electrostatics and sterics.
+          If the number of replicas does not match the number of lambda windows.
+          If there are states with naked charges.
+        Warnings
+          If there are non-zero values for restraints (lambda_restraints).
+        """
+
+        lambda_elec = lambda_settings.lambda_elec
+        lambda_vdw = lambda_settings.lambda_vdw
+        lambda_restraints = lambda_settings.lambda_restraints
+        n_replicas = alchemsampler_settings.n_replicas
+
+        # Ensure that all lambda components have equal amount of windows
+        lambda_components = [lambda_vdw, lambda_elec]
+        it = iter(lambda_components)
+        the_len = len(next(it))
+        if not all(len(l) == the_len for l in it):
+            errmsg = (
+                "Components elec and vdw must have equal amount"
+                f" of lambda windows. Got {len(lambda_elec)} elec lambda"
+                f" windows and {len(lambda_vdw)} vdw lambda windows.")
+            raise ValueError(errmsg)
+
+        # Ensure that number of overall lambda windows matches number of lambda
+        # windows for individual components
+        if n_replicas != len(lambda_vdw):
+            errmsg = (f"Number of replicas {n_replicas} does not equal the"
+                      f" number of lambda windows {len(lambda_vdw)}")
+            raise ValueError(errmsg)
+
+        # Check if there are lambda windows with naked charges
+        for inx, lam in enumerate(lambda_elec):
+            if lam < 1 and lambda_vdw[inx] == 1:
+                errmsg = (
+                    "There are states along this lambda schedule "
+                    "where there are atoms with charges but no LJ "
+                    f"interactions: lambda {inx}: "
+                    f"elec {lam} vdW {lambda_vdw[inx]}")
+                raise ValueError(errmsg)
+
+        # Check if there are lambda windows with non-zero restraints
+        if len([r for r in lambda_restraints if r != 0]) > 0:
+            wmsg = ("Non-zero restraint lambdas applied. The absolute "
+                    "solvation protocol doesn't apply restraints, "
+                    "therefore restraints won't be applied. "
+                    f"Given lambda_restraints: {lambda_restraints}")
+            logger.warning(wmsg)
+            warnings.warn(wmsg)
+
     def _create(
         self,
         stateA: ChemicalSystem,
@@ -535,6 +610,10 @@ def _create(
         )
         self._validate_alchemical_components(alchem_comps)
 
+        # Validate the lambda schedule
+        self._validate_lambda_schedule(self.settings.lambda_settings,
+                                       self.settings.alchemsampler_settings)
+
         # Check nonbond & solvent compatibility
         solv_nonbonded_method = self.settings.solvent_system_settings.nonbonded_method
         vac_nonbonded_method = self.settings.vacuum_system_settings.nonbonded_method
@@ -655,6 +734,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
             * system_settings : SystemSettings
             * solvation_settings : SolvationSettings
             * alchemical_settings : AlchemicalSettings
+            * lambda_settings : LambdaSettings
             * sampler_settings : AlchemicalSamplerSettings
             * engine_settings : OpenMMEngineSettings
             * integrator_settings : IntegratorSettings
@@ -668,6 +748,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
         settings['system_settings'] = prot_settings.vacuum_system_settings
         settings['solvation_settings'] = prot_settings.solvation_settings
         settings['alchemical_settings'] = prot_settings.alchemical_settings
+        settings['lambda_settings'] = prot_settings.lambda_settings
         settings['sampler_settings'] = prot_settings.alchemsampler_settings
         settings['engine_settings'] = prot_settings.vacuum_engine_settings
         settings['integrator_settings'] = prot_settings.integrator_settings
@@ -739,6 +820,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
             * system_settings : SystemSettings
             * solvation_settings : SolvationSettings
             * alchemical_settings : AlchemicalSettings
+            * lambda_settings : LambdaSettings
             * sampler_settings : AlchemicalSamplerSettings
             * engine_settings : OpenMMEngineSettings
             * integrator_settings : IntegratorSettings
@@ -752,6 +834,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
         settings['system_settings'] = prot_settings.solvent_system_settings
         settings['solvation_settings'] = prot_settings.solvation_settings
         settings['alchemical_settings'] = prot_settings.alchemical_settings
+        settings['lambda_settings'] = prot_settings.lambda_settings
         settings['sampler_settings'] = prot_settings.alchemsampler_settings
         settings['engine_settings'] = prot_settings.solvent_engine_settings
         settings['integrator_settings'] = prot_settings.integrator_settings