Adding integrator settings model

feflow/protocols/nonequilibrium_cycling.py

@@ -301,16 +301,13 @@ def _execute(self, ctx, *, state_a, state_b, mapping, settings, **inputs):
 
  # Set up integrator
  temperature = to_openmm(thermodynamic_settings.temperature)
- neq_steps = settings.eq_steps
- eq_steps = settings.neq_steps
- timestep = to_openmm(settings.timestep)
- splitting = settings.neq_splitting
+ integrator_settings = settings.integrator_settings
  integrator = PeriodicNonequilibriumIntegrator(
  alchemical_functions=settings.lambda_functions,
- nsteps_neq=neq_steps,
- nsteps_eq=eq_steps,
- splitting=splitting,
- timestep=timestep,
+ nsteps_neq=integrator_settings.nonequilibrium_steps,
+ nsteps_eq=integrator_settings.equilibrium_steps,
+ splitting=integrator_settings.splitting,
+ timestep=to_openmm(integrator_settings.timestep), # needs openmm Quantity
  temperature=temperature,
  )
 
@@ -454,8 +451,8 @@ def _execute(self, ctx, *, setup, settings, **inputs):
  context.setVelocitiesToTemperature(temperature)
 
  # Extract settings used below
- neq_steps = settings.eq_steps
- eq_steps = settings.neq_steps
+ neq_steps = settings.integrator_settings.nonequilibrium_steps
+ eq_steps = settings.integrator_settings.equilibrium_steps
  traj_save_frequency = settings.traj_save_frequency
  work_save_frequency = (
  settings.work_save_frequency
@@ -609,7 +606,7 @@ def _execute(self, ctx, *, setup, settings, **inputs):
  )
 
  # Computing performance in ns/day
- timestep = to_openmm(settings.timestep)
+ timestep = to_openmm(settings.integrator_settings.timestep)
  simulation_time = 2 * (eq_steps + neq_steps) * timestep
  walltime_in_seconds = cycle_walltime.total_seconds() * openmm_unit.seconds
  estimated_performance = simulation_time.value_in_unit(
@@ -863,7 +860,10 @@ def __init__(self, settings: Settings):
 
  @classmethod
  def _default_settings(cls):
- from feflow.settings.nonequilibrium_cycling import NonEquilibriumCyclingSettings
+ from feflow.settings import (
+ NonEquilibriumCyclingSettings,
+ PeriodicNonequilibriumIntegratorSettings,
+ )
  from gufe.settings import OpenMMSystemGeneratorFFSettings, ThermoSettings
  from openfe.protocols.openmm_utils.omm_settings import (
  SystemSettings,
@@ -877,6 +877,7 @@ def _default_settings(cls):
  system_settings=SystemSettings(),
  solvation_settings=SolvationSettings(),
  alchemical_settings=AlchemicalSettings(),
+ integrator_settings=PeriodicNonequilibriumIntegratorSettings(),
  )
 
  # NOTE: create method should be really fast, since it would be running in the work units not the clients!!

feflow/settings/__init__.py

@@ -0,0 +1,2 @@
+from .integrators import PeriodicNonequilibriumIntegratorSettings
+from .nonequilibrium_cycling import NonEquilibriumCyclingSettings

feflow/settings/integrators.py

@@ -0,0 +1,58 @@
+"""
+Settings objects for the different integrators used in the protocols.
+
+The reasoning is to have a base class that contains the settings that
+would be shared between different integrators, and subclasses of it
+for the specific integrator settings.
+"""
+
+try:
+ from pydantic.v1 import validator
+except ImportError:
+ from pydantic import validator # type: ignore[assignment]
+
+from openff.units import unit
+from openff.models.types import FloatQuantity
+from gufe.settings import SettingsBaseModel
+
+
+class PeriodicNonequilibriumIntegratorSettings(SettingsBaseModel):
+ """Settings for the PeriodicNonequilibriumIntegrator"""
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ timestep: FloatQuantity["femtosecond"] = 4 * unit.femtoseconds
+ """Size of the simulation timestep. Default 4 fs."""
+ splitting: str = "V R H O R V"
+ """Operator splitting"""
+ equilibrium_steps: int = 250000
+ """Number of steps for the equilibrium parts of the cycle. Default 250000"""
+ nonequilibrium_steps: int = 250000
+ """Number of steps for the non-equilibrium parts of the cycle. Default 250000"""
+
+ # TODO: This validator is used in other settings, better create a new Type
+ @validator("timestep")
+ def must_be_positive(cls, v):
+ if v <= 0:
+ errmsg = (f"timestep must be positive, received {v}.")
+ raise ValueError(errmsg)
+ return v
+
+ # TODO: This validator is used in other settings, better create a new Type
+ @validator('timestep')
+ def is_time(cls, v):
+ # these are time units, not simulation steps
+ if not v.is_compatible_with(unit.picosecond):
+ raise ValueError("timestep must be in time units "
+ "(i.e. picoseconds)")
+ return v
+
+ # TODO: This validator is used in other settings, better create a new Type
+ @validator("equilibrium_steps", "nonequilibrium_steps")
+ def must_be_positive_or_zero(cls, v):
+ if v < 0:
+ errmsg = ("langevin_collision_rate, and n_restart_attempts must be"
+ f" zero or positive values, got {v}.")
+ raise ValueError(errmsg)
+ return v

feflow/settings/nonequilibrium_cycling.py

@@ -6,12 +6,15 @@
 """
 
 from typing import Optional
+
+from feflow.settings import PeriodicNonequilibriumIntegratorSettings
+
 from gufe.settings import Settings
-from openff.units import unit
 from pydantic import root_validator
 from openfe.protocols.openmm_utils.omm_settings import SystemSettings, SolvationSettings
 from openfe.protocols.openmm_rfe.equil_rfe_settings import AlchemicalSettings
 
+
 # Default settings for the lambda functions
 x = "lambda"
 DEFAULT_ALCHEMICAL_FUNCTIONS = {
@@ -64,11 +67,8 @@ class Config:
  # alchemical settings
  alchemical_settings: AlchemicalSettings
 
- # NEQ integration settings
- timestep = 4.0 * unit.femtoseconds
- neq_splitting = "V R H O R V"
- eq_steps = 250000
- neq_steps = 250000
+ # integrator settings
+ integrator_settings: PeriodicNonequilibriumIntegratorSettings
 
  # platform and serialization
  platform = "CUDA"

feflow/tests/conftest.py

@@ -74,8 +74,8 @@ def short_settings():
  settings = NonEquilibriumCyclingProtocol.default_settings()
 
  settings.thermo_settings.temperature = 300 * unit.kelvin
- settings.eq_steps = 25000
- settings.neq_steps = 25000
+ settings.integrator_settings.equilibrium_steps = 25000
+ settings.integrator_settings.nonequilibrium_steps = 25000
  settings.work_save_frequency = 50
  settings.traj_save_frequency = 250
  settings.platform = "CPU"