Adding LambdaSettings object for AFE

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
@@ -406,14 +407,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:]]
-        )
+
+        lambda_elec = settings['lambda_settings'].lambda_elec
+        lambda_vdw = settings['lambda_settings'].lambda_vdw
+
+        # 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
 
         n_replicas = settings['sampler_settings'].n_replicas
 
@@ -761,11 +764,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")
@@ -775,7 +776,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

@@ -40,20 +40,52 @@
 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 ")
-            raise ValueError(errmsg)
+class LambdaSettings(SettingsBaseModel):
+    """Settings for lambda schedule
+    """
+    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')
+    def must_have_equal_number_elec_vdw_restraints_windows(cls, v, values):
+        if 'lambda_vdw' in values and 'lambda_restraints' in values:
+            lambdas = [v, values['lambda_vdw'], values['lambda_restraints']]
+            it = iter(lambdas)
+            the_len = len(next(it))
+            if not all(len(l) == the_len for l in it):
+                raise ValueError('Components elec, vdw and restraints must have '
+                                 'equal amount of lambda windows')
         return v
 
 
@@ -91,11 +123,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.
     """
 
@@ -129,3 +166,11 @@ class Config:
     Simulation control settings, including simulation lengths and
     record-keeping for the solvent transformation.
     """
+
+    def must_have_same_number_windows(cls, values):
+        n_elec = len(values['lambda_settings']['lambda_elec'])
+        n_replicas = values['alchemsampler_settings']['n_replicas']
+        if n_elec != n_replicas:
+            errmsg = "Number of lambda windows must equal number of replicas."
+            raise ValueError(errmsg)
+        return values

openfe/protocols/openmm_afe/equil_solvation_afe_method.py

@@ -50,7 +50,7 @@
 )
 from openfe.protocols.openmm_afe.equil_afe_settings import (
     AbsoluteSolvationSettings, SystemSettings,
-    SolvationSettings, AlchemicalSettings,
+    SolvationSettings, AlchemicalSettings, LambdaSettings,
     AlchemicalSamplerSettings, OpenMMEngineSettings,
     IntegratorSettings, SimulationSettings,
     SettingsBaseModel,
@@ -404,8 +404,9 @@ def _default_settings(cls):
             solvent_system_settings=SystemSettings(),
             vacuum_system_settings=SystemSettings(nonbonded_method='nocutoff'),
             alchemical_settings=AlchemicalSettings(),
+            lambda_settings=LambdaSettings(),
             alchemsampler_settings=AlchemicalSamplerSettings(
-                n_replicas=24,
+                n_replicas=20,
             ),
             solvation_settings=SolvationSettings(),
             vacuum_engine_settings=OpenMMEngineSettings(),
@@ -655,6 +656,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 +670,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 +742,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 +756,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

openfe/protocols/openmm_rfe/equil_rfe_methods.py

@@ -53,7 +53,7 @@
 
 from .equil_rfe_settings import (
     RelativeHybridTopologyProtocolSettings, SystemSettings,
-    SolvationSettings, AlchemicalSettings,
+    SolvationSettings, AlchemicalSettings, LambdaSettings,
     AlchemicalSamplerSettings, OpenMMEngineSettings,
     IntegratorSettings, SimulationSettings
 )
@@ -456,6 +456,7 @@ def _default_settings(cls):
             system_settings=SystemSettings(),
             solvation_settings=SolvationSettings(),
             alchemical_settings=AlchemicalSettings(),
+            lambda_settings=LambdaSettings(),
             alchemical_sampler_settings=AlchemicalSamplerSettings(),
             engine_settings=OpenMMEngineSettings(),
             integrator_settings=IntegratorSettings(),
@@ -626,6 +627,7 @@ def run(self, *, dry=False, verbose=True,
         forcefield_settings: settings.OpenMMSystemGeneratorFFSettings = protocol_settings.forcefield_settings
         thermo_settings: settings.ThermoSettings = protocol_settings.thermo_settings
         alchem_settings: AlchemicalSettings = protocol_settings.alchemical_settings
+        lambda_settings: LambdaSettings = protocol_settings.lambda_settings
         system_settings: SystemSettings = protocol_settings.system_settings
         solvation_settings: SolvationSettings = protocol_settings.solvation_settings
         sampler_settings: AlchemicalSamplerSettings = protocol_settings.alchemical_sampler_settings
@@ -793,8 +795,8 @@ def run(self, *, dry=False, verbose=True,
         # TODO - this should be exposed to users, maybe we should offer the
         # ability to print the schedule directly in settings?
         lambdas = _rfe_utils.lambdaprotocol.LambdaProtocol(
-            functions=alchem_settings.lambda_functions,
-            windows=alchem_settings.lambda_windows
+            functions=lambda_settings.lambda_functions,
+            windows=lambda_settings.lambda_windows
         )
 
         # PR #125 temporarily pin lambda schedule spacing to n_replicas

openfe/protocols/openmm_rfe/equil_rfe_settings.py

@@ -30,24 +30,33 @@
     from pydantic import validator  # type: ignore[assignment]
 
 
-class AlchemicalSettings(SettingsBaseModel):
+class LambdaSettings(SettingsBaseModel):
     class Config:
         extra = 'ignore'
         arbitrary_types_allowed = True
 
-    """Settings for the alchemical protocol
-
-    This describes the lambda schedule and the creation of the
-    hybrid system.
+    """Settings for the lambda protocol
+    This describes the lambda schedule.
     """
-    # Lambda settings
     lambda_functions = 'default'
     """
     Key of which switching functions to use for alchemical mutation.
     Default 'default'.
     """
     lambda_windows = 11
     """Number of lambda windows to calculate. Default 11."""
+
+
+class AlchemicalSettings(SettingsBaseModel):
+    class Config:
+        extra = 'ignore'
+        arbitrary_types_allowed = True
+
+    """Settings for the alchemical protocol
+
+    This describes the creation of the hybrid system.
+    """
+
     unsampled_endstates = False
     """
     Whether to have extra unsampled endstate windows for long range
@@ -117,9 +126,13 @@ class Config:
     """Settings for solvating the system."""
 
     # Alchemical settings
+    lambda_settings: LambdaSettings
+    """
+    Lambda protocol settings including lambda windows and lambda functions.
+    """
     alchemical_settings: AlchemicalSettings
     """
-    Alchemical protocol settings including lambda windows and soft core scaling.
+    Alchemical protocol settings including soft core scaling.
     """
     alchemical_sampler_settings: AlchemicalSamplerSettings
     """

openfe/tests/protocols/test_openmm_afe_slow.py

@@ -56,9 +56,14 @@ def test_openmm_run_engine(platform,
     s.integrator_settings.n_steps = 5 * unit.timestep
     s.vacuum_simulation_settings.checkpoint_interval = 5 * unit.timestep
     s.solvent_simulation_settings.checkpoint_interval = 5 * unit.timestep
-    s.alchemsampler_settings.n_replicas = 14
-    s.alchemical_settings.lambda_elec_windows = 5
-    s.alchemical_settings.lambda_vdw_windows = 9
+    s.alchemsampler_settings.n_replicas = 20
+    s.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, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
+    s.lambda_settings.lambda_vdw = \
+        [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]
+
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(
             settings=s,
@@ -103,4 +108,4 @@ def test_openmm_run_engine(platform,
     states = results.get_replica_states()
     assert len(states.items()) == 2
     assert len(states['solvent']) == 1
-    assert states['solvent'][0].shape[1] == 14
+    assert states['solvent'][0].shape[1] == 20

openfe/tests/protocols/test_openmm_afe_solvation_protocol.py

@@ -33,17 +33,28 @@ def test_create_default_settings():
 
 
 @pytest.mark.parametrize('val', [
-    {'elec': 0, 'vdw': 5},
-    {'elec': -2, 'vdw': 5},
-    {'elec': 5, 'vdw': -2},
-    {'elec': 5, 'vdw': 0},
+    {'elec': [0.0, -1], 'vdw': [0.0, 1.0], 'restraints': [0.0, 1.0]},
+    {'elec': [0.0, 1.5], 'vdw': [0.0, 1.5], 'restraints': [-0.1, 1.0]}
 ])
 def test_incorrect_window_settings(val, default_settings):
-    errmsg = "lambda steps must be positive"
-    alchem_settings = default_settings.alchemical_settings
+    errmsg = "Lambda windows must be between 0 and 1."
+    lambda_settings = default_settings.lambda_settings
     with pytest.raises(ValueError, match=errmsg):
-        alchem_settings.lambda_elec_windows = val['elec']
-        alchem_settings.lambda_vdw_windows = val['vdw']
+        lambda_settings.lambda_elec = val['elec']
+        lambda_settings.lambda_vdw = val['vdw']
+        lambda_settings.lambda_restraints = val['restraints']
+
+@pytest.mark.parametrize('val', [
+    {'elec': [0.0, 0.1, 1.0], 'vdw': [0.0, 1.0], 'restraints': [0.0, 1.0]},
+])
+def test_incorrect_number_windows(val, default_settings):
+    errmsg = "Components elec, vdw and restraints must have equal amount " \
+             "of lambda windows"
+    lambda_settings = default_settings.lambda_settings
+    with pytest.raises(ValueError, match=errmsg):
+        lambda_settings.lambda_elec = val['elec']
+        lambda_settings.lambda_vdw = val['vdw']
+        lambda_settings.lambda_restraints = val['restraints']
 
 
 def test_create_default_protocol(default_settings):