From 4d9cdc47d57eb4e338d8e1127c91410b35168e9e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Fri, 13 Sep 2024 16:42:48 -0400 Subject: [PATCH 01/21] WIP -- Initial skeleton for NEq switching protocol --- feflow/protocols/nonequilibrium_switching.py | 169 +++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 feflow/protocols/nonequilibrium_switching.py diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py new file mode 100644 index 00000000..30fcace5 --- /dev/null +++ b/feflow/protocols/nonequilibrium_switching.py @@ -0,0 +1,169 @@ +import mdtraj + + +class BaseSwitchingUnit(ProtocolUnit): + """ + Monolithic unit for the cycle part of the simulation. + It runs a number of NEq cycles from the outputs of a setup unit and stores the work computed in + numpy-formatted files, to be analyzed by a result unit. + """ + + @staticmethod + def extract_positions(context, initial_atom_indices, final_atom_indices): + """ + Extract positions from initial and final systems based from the hybrid topology. + + Parameters + ---------- + context: openmm.Context + Current simulation context where from extract positions. + hybrid_topology_factory: HybridTopologyFactory + Hybrid topology factory where to extract positions and mapping information + + Returns + ------- + + Notes + ----- + It achieves this by taking the positions and indices from the initial and final states of + the transformation, and computing the overlap of these with the indices of the complete + hybrid topology, filtered by some mdtraj selection expression. + + 1. Get positions from context + 2. Get topology from HTF (already mdtraj topology) + 3. Merge that information into mdtraj.Trajectory + 4. Filter positions for initial/final according to selection string + """ + import numpy as np + + # Get positions from current openmm context + positions = context.getState(getPositions=True).getPositions(asNumpy=True) + + # Get indices for initial and final topologies in hybrid topology + initial_indices = np.asarray(initial_atom_indices) + final_indices = np.asarray(final_atom_indices) + + initial_positions = positions[initial_indices, :] + final_positions = positions[final_indices, :] + + return initial_positions, final_positions + + def _execute(self, ctx, *, protocol, md_unit, index, **inputs): + """ + Execute the simulation part of the Nonequilibrium switching protocol using GUFE objects. + + Parameters + ---------- + ctx : gufe.protocols.protocolunit.Context + The gufe context for the unit. + protocol : gufe.protocols.Protocol + The Protocol used to create this Unit. Contains key information + such as the settings. + md_unit : gufe.protocols.ProtocolUnit + The SetupUnit + index: int + TODO: Index for the snapshot to use as input + + Returns + ------- + dict : dict[str, str] + Dictionary with paths to work arrays, both forward and reverse, and trajectory coordinates for systems + A and B. + """ + import openmm + from openmmtools.integrators import PeriodicNonequilibriumIntegrator + + # Setting up logging to file in shared filesystem + file_logger = logging.getLogger("neq-cycling") + output_log_path = ctx.shared / "feflow-neq-cycling.log" + file_handler = logging.FileHandler(output_log_path, mode="w") + file_handler.setLevel(logging.DEBUG) # TODO: Set to INFO in production + log_formatter = logging.Formatter( + fmt="%(asctime)s %(levelname)-8s %(message)s", datefmt="%Y-%m-%d %H:%M:%S" + ) + file_handler.setFormatter(log_formatter) + file_logger.addHandler(file_handler) + + system = deserialize(md_unit.inputs["setup"].outputs["system"]) + state = deserialize(md_unit.inputs["setup"].outputs["state"]) + integrator = deserialize(md_unit.inputs["setup"].outputs["integrator"]) + + PeriodicNonequilibriumIntegrator.restore_interface(integrator) + + # Get atom indices for either end of the hybrid topology + initial_atom_indices = setup.outputs["initial_atom_indices"] + final_atom_indices = setup.outputs["final_atom_indices"] + + # Extract settings from the Protocol + settings = protocol.settings + + # Load positions from snapshots + xtc_file = md_unit.outputs["production_trajectory"] + md_traj_ob = mdtraj.load_frame(xtc_file, index=index) + input_positions = md_traj_ob.openmm_positions(0) + # Set up context + platform = get_openmm_platform(settings.engine_settings.compute_platform) + context = openmm.Context(system, integrator, platform) + context.setState(state) + # TODO: This is kinda ugly, is there a better way to set positions? + context.setPositions(input_positions) + + # Setting velocities to temperatures + thermodynamic_settings = settings.thermo_settings + temperature = to_openmm(thermodynamic_settings.temperature) + context.setVelocitiesToTemperature(temperature) + + # Extract settings used below + neq_steps = settings.integrator_settings.nonequilibrium_steps + traj_save_frequency = settings.traj_save_frequency + work_save_frequency = ( + settings.work_save_frequency + ) # Note: this is divisor of traj save freq. + selection_expression = settings.atom_selection_expression + + try: + # Coarse number of steps -- each coarse consists of work_save_frequency steps + coarse_neq_steps = int( + neq_steps / work_save_frequency + ) # Note: neq_steps is multiple of work save steps + + # TODO: Also get the GPU information (plain try-except with nvidia-smi) + + + integrator.step(NSTEPS) + + + + + # Equilibrium (lambda = 0) + # start timer + start_time = time.perf_counter() + # Run neq + # Forward (0 -> 1) + # Initialize works with current value + forward_works = [] + for fwd_step in range(coarse_neq_steps): + integrator.step(work_save_frequency) + forward_works.append(integrator.get_protocol_work(dimensionless=True)) + if fwd_step % traj_save_frequency == 0: + initial_positions, final_positions = self.extract_positions( + context, initial_atom_indices, final_atom_indices + ) + forward_neq_initial.append(initial_positions) + forward_neq_final.append(final_positions) + # Make sure trajectories are stored at the end of the neq loop + initial_positions, final_positions = self.extract_positions( + context, initial_atom_indices, final_atom_indices + ) + forward_neq_initial.append(initial_positions) + forward_neq_final.append(final_positions) + + neq_forward_time = time.perf_counter() + neq_forward_walltime = datetime.timedelta( + seconds=neq_forward_time - eq_forward_time + ) + file_logger.info( + f"replicate_{self.name} Forward nonequilibrium time (lambda 0 -> 1): {neq_forward_walltime}" + ) + + # TODO: We should return the work in one direction From e2832c52d5c8e5077cf788d4348b8e6798d8aea3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Fri, 16 Jan 2026 17:05:34 -0500 Subject: [PATCH 02/21] migrating integrator settings pydantic v2 --- feflow/settings/integrators.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/feflow/settings/integrators.py b/feflow/settings/integrators.py index 4d9e73bf..f4b99cc5 100644 --- a/feflow/settings/integrators.py +++ b/feflow/settings/integrators.py @@ -8,11 +8,10 @@ from typing import Annotated, TypeAlias -from pydantic.v1 import validator - from openff.units import unit from gufe.settings import SettingsBaseModel from gufe.settings.typing import GufeQuantity, specify_quantity_units +from pydantic import field_validator, ConfigDict FemtosecondQuantity: TypeAlias = Annotated[ GufeQuantity, specify_quantity_units("femtoseconds") @@ -24,9 +23,7 @@ class PeriodicNonequilibriumIntegratorSettings(SettingsBaseModel): """Settings for the PeriodicNonequilibriumIntegrator""" - - class Config: - arbitrary_types_allowed = True + model_config = ConfigDict(arbitrary_types_allowed=True) timestep: FemtosecondQuantity = 4 * unit.femtoseconds """Size of the simulation timestep. Default 4 fs.""" @@ -48,7 +45,8 @@ class Config: """ # TODO: This validator is used in other settings, better create a new Type - @validator("timestep") + @field_validator("timestep") + @classmethod def must_be_positive(cls, v): if v <= 0: errmsg = f"timestep must be positive, received {v}." @@ -56,7 +54,8 @@ def must_be_positive(cls, v): return v # TODO: This validator is used in other settings, better create a new Type - @validator("timestep") + @field_validator("timestep") + @classmethod def is_time(cls, v): # these are time units, not simulation steps if not v.is_compatible_with(unit.picosecond): @@ -64,7 +63,8 @@ def is_time(cls, v): return v # TODO: This validator is used in other settings, better create a new Type - @validator("equilibrium_steps", "nonequilibrium_steps") + @field_validator("equilibrium_steps", "nonequilibrium_steps") + @classmethod def must_be_positive_or_zero(cls, v): if v < 0: errmsg = ( From 64dc3e33afdd49de8c51f605a476d92f45c7260b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Fri, 16 Jan 2026 18:18:53 -0500 Subject: [PATCH 03/21] Base settings object for integrators for cycling and switching. --- feflow/settings/integrators.py | 45 +++++++++++++++++++++------------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/feflow/settings/integrators.py b/feflow/settings/integrators.py index f4b99cc5..6d4ee5da 100644 --- a/feflow/settings/integrators.py +++ b/feflow/settings/integrators.py @@ -21,28 +21,13 @@ ] -class PeriodicNonequilibriumIntegratorSettings(SettingsBaseModel): - """Settings for the PeriodicNonequilibriumIntegrator""" +class BaseNonequilibriumIntegrator(SettingsBaseModel): + """Base class for nonequilibrium integrator settings""" model_config = ConfigDict(arbitrary_types_allowed=True) timestep: FemtosecondQuantity = 4 * unit.femtoseconds """Size of the simulation timestep. Default 4 fs.""" splitting: str = "V R H O R V" - """Operator splitting""" - equilibrium_steps: int = 12500 - """Number of steps for the equilibrium parts of the cycle. Default 12500""" - nonequilibrium_steps: int = 12500 - """Number of steps for the non-equilibrium parts of the cycle. Default 12500""" - barostat_frequency: TimestepQuantity = 25 * unit.timestep - """ - Frequency at which volume scaling changes should be attempted. - Note: The barostat frequency is ignored for gas-phase simulations. - Default 25 * unit.timestep. - """ - remove_com: bool = False - """ - Whether or not to remove the center of mass motion. Default False. - """ # TODO: This validator is used in other settings, better create a new Type @field_validator("timestep") @@ -62,6 +47,32 @@ def is_time(cls, v): raise ValueError("timestep must be in time units " "(i.e. picoseconds)") return v + +class AlchemicalNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): + """Settings for the AlchemicalNonequilibriumIntegrator used for switching""" + ... + + +class PeriodicNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): + """Settings for the PeriodicNonequilibriumIntegrator""" + model_config = ConfigDict(arbitrary_types_allowed=True) + + """Operator splitting""" + equilibrium_steps: int = 12500 + """Number of steps for the equilibrium parts of the cycle. Default 12500""" + nonequilibrium_steps: int = 12500 + """Number of steps for the non-equilibrium parts of the cycle. Default 12500""" + barostat_frequency: TimestepQuantity = 25 * unit.timestep + """ + Frequency at which volume scaling changes should be attempted. + Note: The barostat frequency is ignored for gas-phase simulations. + Default 25 * unit.timestep. + """ + remove_com: bool = False + """ + Whether or not to remove the center of mass motion. Default False. + """ + # TODO: This validator is used in other settings, better create a new Type @field_validator("equilibrium_steps", "nonequilibrium_steps") @classmethod From 049cbd81ed1f2b127d927ab56f20cafb4d47b477 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 16 Jan 2026 23:19:14 +0000 Subject: [PATCH 04/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- feflow/settings/integrators.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/feflow/settings/integrators.py b/feflow/settings/integrators.py index 6d4ee5da..e970404f 100644 --- a/feflow/settings/integrators.py +++ b/feflow/settings/integrators.py @@ -23,6 +23,7 @@ class BaseNonequilibriumIntegrator(SettingsBaseModel): """Base class for nonequilibrium integrator settings""" + model_config = ConfigDict(arbitrary_types_allowed=True) timestep: FemtosecondQuantity = 4 * unit.femtoseconds @@ -50,11 +51,13 @@ def is_time(cls, v): class AlchemicalNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): """Settings for the AlchemicalNonequilibriumIntegrator used for switching""" + ... class PeriodicNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): """Settings for the PeriodicNonequilibriumIntegrator""" + model_config = ConfigDict(arbitrary_types_allowed=True) """Operator splitting""" From 69cf36e59c9db6e3c54b4d6884006cfd44c3889f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 31 Mar 2026 16:49:15 -0400 Subject: [PATCH 05/21] class config with pydantic v2 support --- feflow/settings/nonequilibrium_cycling.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/feflow/settings/nonequilibrium_cycling.py b/feflow/settings/nonequilibrium_cycling.py index a765b73b..f0e31fd3 100644 --- a/feflow/settings/nonequilibrium_cycling.py +++ b/feflow/settings/nonequilibrium_cycling.py @@ -17,6 +17,7 @@ ThermoSettings, ) from openfe.protocols.openmm_rfe.equil_rfe_settings import AlchemicalSettings +from pydantic import ConfigDict # Default settings for the lambda functions @@ -51,10 +52,7 @@ class NonEquilibriumCyclingSettings(Settings): alchemical : AlchemicalSettings The alchemical settings to use. """ - - # TODO: Add type hints - class Config: - arbitrary_types_allowed = True + model_config = ConfigDict(arbitrary_types_allowed=True) forcefield_cache: Optional[str] = ( "db.json" # TODO: Remove once it has been integrated with openfe settings From adf963056be0aed351fa3549e6431438efdc2ebe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 31 Mar 2026 16:49:31 -0400 Subject: [PATCH 06/21] base imports and names --- feflow/protocols/__init__.py | 4 ++++ feflow/protocols/nonequilibrium_switching.py | 9 +++++++++ feflow/settings/__init__.py | 1 + 3 files changed, 14 insertions(+) diff --git a/feflow/protocols/__init__.py b/feflow/protocols/__init__.py index b7f78071..7ed4813d 100644 --- a/feflow/protocols/__init__.py +++ b/feflow/protocols/__init__.py @@ -7,3 +7,7 @@ NonEquilibriumCyclingProtocol, NonEquilibriumCyclingProtocolResult, ) +from .nonequilibrium_switching import ( + NonEquilibriumSwitchingProtocol, + NonEquilibriumSwitchingProtocolResult, +) diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py index 30fcace5..92589302 100644 --- a/feflow/protocols/nonequilibrium_switching.py +++ b/feflow/protocols/nonequilibrium_switching.py @@ -1,4 +1,5 @@ import mdtraj +from gufe import Protocol, ProtocolUnit, ProtocolResult class BaseSwitchingUnit(ProtocolUnit): @@ -167,3 +168,11 @@ def _execute(self, ctx, *, protocol, md_unit, index, **inputs): ) # TODO: We should return the work in one direction + + +class NonEquilibriumSwitchingProtocol(Protocol): + ... + + +class NonEquilibriumSwitchingProtocolResult(ProtocolResult): + ... \ No newline at end of file diff --git a/feflow/settings/__init__.py b/feflow/settings/__init__.py index 06fd6a44..cf15b085 100644 --- a/feflow/settings/__init__.py +++ b/feflow/settings/__init__.py @@ -1,3 +1,4 @@ from .integrators import PeriodicNonequilibriumIntegratorSettings from .small_molecules import OpenFFPartialChargeSettings from .nonequilibrium_cycling import NonEquilibriumCyclingSettings +from .nonequilibrium_switching import NonEquilibriumSwitchingSettings From d38d63fcd049b52096b47fe520b39109ad201198 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 31 Mar 2026 20:50:15 +0000 Subject: [PATCH 07/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- feflow/settings/nonequilibrium_cycling.py | 1 + 1 file changed, 1 insertion(+) diff --git a/feflow/settings/nonequilibrium_cycling.py b/feflow/settings/nonequilibrium_cycling.py index f0e31fd3..a95b5038 100644 --- a/feflow/settings/nonequilibrium_cycling.py +++ b/feflow/settings/nonequilibrium_cycling.py @@ -52,6 +52,7 @@ class NonEquilibriumCyclingSettings(Settings): alchemical : AlchemicalSettings The alchemical settings to use. """ + model_config = ConfigDict(arbitrary_types_allowed=True) forcefield_cache: Optional[str] = ( From 6200b3246bd44da67558e4501f49f3a93e84e30a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 7 Apr 2026 19:27:22 -0400 Subject: [PATCH 08/21] First iteration on Settings for NES --- feflow/settings/__init__.py | 7 +- feflow/settings/integrators.py | 31 ++- feflow/settings/nonequilibrium_switching.py | 223 ++++++++++++++++++++ 3 files changed, 257 insertions(+), 4 deletions(-) create mode 100644 feflow/settings/nonequilibrium_switching.py diff --git a/feflow/settings/__init__.py b/feflow/settings/__init__.py index cf15b085..d810455f 100644 --- a/feflow/settings/__init__.py +++ b/feflow/settings/__init__.py @@ -1,4 +1,7 @@ -from .integrators import PeriodicNonequilibriumIntegratorSettings +from .integrators import ( + PeriodicNonequilibriumIntegratorSettings, + AlchemicalNonequilibriumIntegratorSettings, +) from .small_molecules import OpenFFPartialChargeSettings from .nonequilibrium_cycling import NonEquilibriumCyclingSettings -from .nonequilibrium_switching import NonEquilibriumSwitchingSettings +from .nonequilibrium_switching import NonEquilibriumSwitchingSettings, SnapshotSettings diff --git a/feflow/settings/integrators.py b/feflow/settings/integrators.py index e970404f..40bc37dd 100644 --- a/feflow/settings/integrators.py +++ b/feflow/settings/integrators.py @@ -19,6 +19,9 @@ TimestepQuantity: TypeAlias = Annotated[ GufeQuantity, specify_quantity_units("timestep") ] +InversePicosecondQuantity: TypeAlias = Annotated[ + GufeQuantity, specify_quantity_units("1/picoseconds") +] class BaseNonequilibriumIntegrator(SettingsBaseModel): @@ -50,9 +53,33 @@ def is_time(cls, v): class AlchemicalNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): - """Settings for the AlchemicalNonequilibriumIntegrator used for switching""" + """Settings for the AlchemicalNonequilibriumLangevinIntegrator used for one-way NEQ switching""" - ... + timestep: FemtosecondQuantity = 4 * unit.femtoseconds + """Size of the simulation timestep. Default 4 fs.""" + splitting: str = "V R H O R V" + """Operator splitting for the Langevin integrator.""" + collision_rate: InversePicosecondQuantity = 1.0 / unit.picoseconds + """Langevin collision rate (friction coefficient). Default 1/ps.""" + nonequilibrium_steps: int = 2500 + """Number of steps for the non-equilibrium switching (lambda 0->1 or 1->0). Default 2500 (10 ps at 4 fs).""" + equilibrium_steps: int = 1000 + """Number of equilibration steps at the endpoint before each switch. Default 1000.""" + + @field_validator("nonequilibrium_steps", "equilibrium_steps") + @classmethod + def must_be_positive_or_zero(cls, v): + if v < 0: + errmsg = f"nonequilibrium_steps and equilibrium_steps must be zero or positive, got {v}." + raise ValueError(errmsg) + return v + + @field_validator("collision_rate") + @classmethod + def collision_rate_must_be_positive(cls, v): + if v <= 0: + raise ValueError(f"collision_rate must be positive, received {v}.") + return v class PeriodicNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): diff --git a/feflow/settings/nonequilibrium_switching.py b/feflow/settings/nonequilibrium_switching.py new file mode 100644 index 00000000..461b9e56 --- /dev/null +++ b/feflow/settings/nonequilibrium_switching.py @@ -0,0 +1,223 @@ +""" +Settings objects for the nonequilibrium switching protocol. +""" + +from typing import Optional + +from feflow.settings import ( + AlchemicalNonequilibriumIntegratorSettings, + OpenFFPartialChargeSettings, +) + +from gufe.settings import Settings, OpenMMSystemGeneratorFFSettings, SettingsBaseModel +from pydantic.v1 import root_validator +from openfe.protocols.openmm_utils.omm_settings import ( + OpenMMSolvationSettings, + OpenMMEngineSettings, + ThermoSettings, +) +from openfe.protocols.openmm_rfe.equil_rfe_settings import AlchemicalSettings +from pydantic import ConfigDict + + +# Default settings for the lambda functions +x = "lambda" +DEFAULT_ALCHEMICAL_FUNCTIONS = { + "lambda_sterics_core": x, + "lambda_electrostatics_core": x, + "lambda_sterics_insert": f"select(step({x} - 0.5), 1.0, 2.0 * {x})", + "lambda_sterics_delete": f"select(step({x} - 0.5), 2.0 * ({x} - 0.5), 0.0)", + "lambda_electrostatics_insert": f"select(step({x} - 0.5), 2.0 * ({x} - 0.5), 0.0)", + "lambda_electrostatics_delete": f"select(step({x} - 0.5), 1.0, 2.0 * {x})", + "lambda_bonds": x, + "lambda_angles": x, + "lambda_torsions": x, +} + + +class SnapshotSettings(SettingsBaseModel): + """ + Settings for loading pre-equilibrated snapshots from a trajectory file via + MDAnalysis, instead of running internal equilibration. + + The trajectory must have been generated from the **hybrid topology** produced + by the SetupUnit (i.e. the same system that will be used for NEQ switching). + Replicate ``i`` uses frame ``i * stride`` from the trajectory. + """ + + topology_file: str + """Path to a topology file (PDB, GRO, …) compatible with the hybrid topology.""" + trajectory_file: str + """Path to a trajectory file (XTC, DCD, …) of equilibrated configurations.""" + stride: int = 1 + """ + Frame stride for snapshot selection. + Replicate i loads frame i * stride. Default 1 (consecutive frames). + """ + + +class NonEquilibriumSwitchingSettings(Settings): + """ + Settings for the NEQ switching protocol. + + The protocol drives the hybrid system from lambda=0 to lambda=1 (forward + switches) and from lambda=1 to lambda=0 (reverse switches) using the + AlchemicalNonequilibriumLangevinIntegrator from openmmtools. Free energy + estimates are obtained via BAR over the replicate work values. + + Starting snapshots for each direction can either be generated internally + via ``integrator_settings.equilibrium_steps`` (set ``lambda0_snapshots`` / + ``lambda1_snapshots`` to ``None``) or loaded from an existing trajectory with + MDAnalysis. When snapshot settings are provided, ``equilibrium_steps`` must be + set to 0 to avoid ambiguity. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + forcefield_cache: Optional[str] = "db.json" + + forcefield_settings: OpenMMSystemGeneratorFFSettings + solvation_settings: OpenMMSolvationSettings + partial_charge_settings: OpenFFPartialChargeSettings + thermo_settings: ThermoSettings + + # Lambda schedule + lambda_functions: dict[str, str] = DEFAULT_ALCHEMICAL_FUNCTIONS + + # Alchemical settings (softcore potentials, charge corrections, …) + alchemical_settings: AlchemicalSettings = AlchemicalSettings(softcore_LJ="gapsys") + + # Integrator (AlchemicalNonequilibriumLangevinIntegrator) + integrator_settings: AlchemicalNonequilibriumIntegratorSettings + + # Platform + engine_settings: OpenMMEngineSettings + + # Optional pre-equilibrated snapshot sources + lambda0_snapshots: Optional[SnapshotSettings] = None + """ + If provided, forward switching replicates load their starting snapshot at + lambda=0 from this trajectory instead of running internal equilibration. + Requires ``integrator_settings.equilibrium_steps == 0``. + """ + lambda1_snapshots: Optional[SnapshotSettings] = None + """ + If provided, reverse switching replicates load their starting snapshot at + lambda=1 from this trajectory instead of running internal equilibration. + Requires ``integrator_settings.equilibrium_steps == 0``. + """ + + work_save_frequency: Optional[int] = None + """ + How often (in NEQ steps) to record the protocol work. + Defaults to nonequilibrium_steps // 50, giving ~50 work samples per switch. + """ + traj_save_frequency: Optional[int] = None + """ + How often (in NEQ steps) to save trajectory frames during switching. + Defaults to 5 * work_save_frequency (~10 frames per switch). + Must be a multiple of work_save_frequency. + """ + + num_switches: int = 100 + """ + Number of independent NEQ switch trajectories to run per direction. + The protocol creates this many forward (lambda 0->1) switches and this many + reverse (lambda 1->0) switches, for a total of 2 * num_switches trajectory + runs. Each switch produces one work value used in the BAR free energy estimate. + """ + + @root_validator + def set_and_validate_save_frequencies(cls, values): + """ + Derive save-frequency defaults from nonequilibrium_steps when not set, + then check consistency. + """ + integrator_settings = values.get("integrator_settings") + neq_steps = ( + integrator_settings.nonequilibrium_steps if integrator_settings else 2500 + ) + + work_freq = values.get("work_save_frequency") + traj_freq = values.get("traj_save_frequency") + + if work_freq is None: + work_freq = max(1, neq_steps // 50) + values["work_save_frequency"] = work_freq + if traj_freq is None: + values["traj_save_frequency"] = work_freq * 5 + traj_freq = values["traj_save_frequency"] + + if traj_freq % work_freq != 0: + raise ValueError( + "traj_save_frequency must be a multiple of work_save_frequency. " + "Please specify consistent values." + ) + return values + + @root_validator + def snapshots_require_no_internal_equilibration(cls, values): + """ + When snapshot settings are provided equilibrium_steps must be 0 — + the snapshots are already equilibrated. + """ + integrator_settings = values.get("integrator_settings") + if integrator_settings is None: + return values + + has_snapshots = values.get("lambda0_snapshots") or values.get("lambda1_snapshots") + if has_snapshots and integrator_settings.equilibrium_steps != 0: + raise ValueError( + "When lambda0_snapshots or lambda1_snapshots are provided the " + "snapshots are assumed to be pre-equilibrated. " + "Set integrator_settings.equilibrium_steps = 0 to avoid " + "running redundant equilibration on top of them." + ) + return values + + @root_validator + def snapshot_trajectories_have_enough_frames(cls, values): + """ + When both lambda0_snapshots and lambda1_snapshots are provided, check + that each trajectory contains enough frames to cover all replicates + (i.e. at least num_switches * stride frames), and that both + trajectories expose the same number of usable snapshots. + """ + snap0: Optional[SnapshotSettings] = values.get("lambda0_snapshots") + snap1: Optional[SnapshotSettings] = values.get("lambda1_snapshots") + + if snap0 is None or snap1 is None: + return values + + num_switches: int = values.get("num_switches", 100) + + try: + import MDAnalysis as mda + + n0 = len(mda.Universe(snap0.topology_file, snap0.trajectory_file).trajectory) + n1 = len(mda.Universe(snap1.topology_file, snap1.trajectory_file).trajectory) + except Exception as exc: + raise ValueError( + f"Could not read snapshot trajectories to validate frame counts: {exc}" + ) from exc + + usable0 = n0 // snap0.stride + usable1 = n1 // snap1.stride + + if usable0 < num_switches: + raise ValueError( + f"lambda0 trajectory has only {n0} frames (stride={snap0.stride} → " + f"{usable0} usable), but num_switches={num_switches}." + ) + if usable1 < num_switches: + raise ValueError( + f"lambda1 trajectory has only {n1} frames (stride={snap1.stride} → " + f"{usable1} usable), but num_switches={num_switches}." + ) + if usable0 != usable1: + raise ValueError( + f"lambda0 and lambda1 trajectories expose different numbers of " + f"usable snapshots ({usable0} vs {usable1}). " + "Provide trajectories of equal length or adjust the stride values." + ) + return values From f62837fc2ca4daa1136cdb716c8dd384e6d8a9c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 7 Apr 2026 19:47:21 -0400 Subject: [PATCH 09/21] Writing basic tests for NES --- feflow/tests/conftest.py | 33 +++ feflow/tests/test_nonequilibrium_switching.py | 209 ++++++++++++++++++ 2 files changed, 242 insertions(+) create mode 100644 feflow/tests/test_nonequilibrium_switching.py diff --git a/feflow/tests/conftest.py b/feflow/tests/conftest.py index 324e1592..611eea25 100644 --- a/feflow/tests/conftest.py +++ b/feflow/tests/conftest.py @@ -185,6 +185,39 @@ def short_settings_multiple_cycles_gpu(short_settings_multiple_cycles): return settings +@pytest.fixture +def short_switching_settings(): + from openff.units import unit + from feflow.protocols import NonEquilibriumSwitchingProtocol + + settings = NonEquilibriumSwitchingProtocol.default_settings() + settings.engine_settings.compute_platform = "CPU" + settings.thermo_settings.temperature = 300 * unit.kelvin + settings.integrator_settings.equilibrium_steps = 50 + settings.integrator_settings.nonequilibrium_steps = 100 + # explicit save frequencies consistent with 100 neq steps + settings.work_save_frequency = 10 + settings.traj_save_frequency = 50 + settings.num_switches = 1 + return settings + + +@pytest.fixture +def short_switching_settings_multiple_switches(): + from openff.units import unit + from feflow.protocols import NonEquilibriumSwitchingProtocol + + settings = NonEquilibriumSwitchingProtocol.default_settings() + settings.engine_settings.compute_platform = "CPU" + settings.thermo_settings.temperature = 300 * unit.kelvin + settings.integrator_settings.equilibrium_steps = 50 + settings.integrator_settings.nonequilibrium_steps = 100 + settings.work_save_frequency = 10 + settings.traj_save_frequency = 50 + settings.num_switches = 5 + return settings + + @pytest.fixture def production_settings(short_settings): settings = short_settings.copy(deep=True) diff --git a/feflow/tests/test_nonequilibrium_switching.py b/feflow/tests/test_nonequilibrium_switching.py new file mode 100644 index 00000000..cc2a1584 --- /dev/null +++ b/feflow/tests/test_nonequilibrium_switching.py @@ -0,0 +1,209 @@ +import json +from pathlib import Path + +import pytest + +from feflow.protocols import ( + NonEquilibriumSwitchingProtocol, + ForwardSwitchingUnit, + ReverseSwitchingUnit, +) +from feflow.settings import NonEquilibriumSwitchingSettings +from gufe.protocols.protocoldag import ProtocolDAGResult, execute_DAG +from gufe.protocols.protocolunit import ProtocolUnitResult +from gufe.tokenization import JSON_HANDLER + +# required plugins/fixtures +pytest_plugins = ["feflow.tests.fixtures.tyk2_fixtures"] + + +class TestNonEquilibriumSwitching: + @pytest.fixture + def protocol_short(self, short_switching_settings): + return NonEquilibriumSwitchingProtocol(settings=short_switching_settings) + + @pytest.fixture + def protocol_dag_result( + self, + protocol_short, + benzene_vacuum_system, + toluene_vacuum_system, + mapping_benzene_toluene, + tmpdir, + ): + dag = protocol_short.create( + stateA=benzene_vacuum_system, + stateB=toluene_vacuum_system, + name="Short vacuum switching", + mapping=mapping_benzene_toluene, + ) + with tmpdir.as_cwd(): + shared = Path("shared") + shared.mkdir() + scratch = Path("scratch") + scratch.mkdir() + dagresult: ProtocolDAGResult = execute_DAG( + dag, shared_basedir=shared, scratch_basedir=scratch + ) + return protocol_short, dag, dagresult + + def test_dag_execute(self, protocol_dag_result): + _, _, dagresult = protocol_dag_result + assert dagresult.ok() + assert dagresult.protocol_unit_results[-1].name == "result" + + def test_terminal_units(self, protocol_dag_result): + _, _, res = protocol_dag_result + finals = res.terminal_protocol_unit_results + assert len(finals) == 1 + assert isinstance(finals[0], ProtocolUnitResult) + assert finals[0].name == "result" + + def test_dag_has_forward_and_reverse_units( + self, + protocol_short, + benzene_vacuum_system, + toluene_vacuum_system, + mapping_benzene_toluene, + ): + """DAG must contain both ForwardSwitchingUnit and ReverseSwitchingUnit instances.""" + dag = protocol_short.create( + stateA=benzene_vacuum_system, + stateB=toluene_vacuum_system, + name="unit type check", + mapping=mapping_benzene_toluene, + ) + unit_types = [type(u) for u in dag.protocol_units] + assert ForwardSwitchingUnit in unit_types + assert ReverseSwitchingUnit in unit_types + + def test_num_switches_creates_correct_unit_counts( + self, + short_switching_settings, + benzene_vacuum_system, + toluene_vacuum_system, + mapping_benzene_toluene, + ): + """num_switches forward + num_switches reverse + 1 setup + 1 result units.""" + protocol = NonEquilibriumSwitchingProtocol(settings=short_switching_settings) + num_switches = short_switching_settings.num_switches + dag = protocol.create( + stateA=benzene_vacuum_system, + stateB=toluene_vacuum_system, + name="unit count check", + mapping=mapping_benzene_toluene, + ) + units = dag.protocol_units + n_forward = sum(1 for u in units if isinstance(u, ForwardSwitchingUnit)) + n_reverse = sum(1 for u in units if isinstance(u, ReverseSwitchingUnit)) + assert n_forward == num_switches + assert n_reverse == num_switches + # +2 for SetupUnit and ResultUnit + assert len(units) == 2 * num_switches + 2 + + def test_create_with_invalid_mapping( + self, + protocol_short, + benzene_solvent_system, + toluene_solvent_system, + mapping_benzonitrile_styrene, + ): + """Mapping whose components don't match the states should raise AssertionError.""" + with pytest.raises(AssertionError): + _ = protocol_short.create( + stateA=benzene_solvent_system, + stateB=toluene_solvent_system, + name="bad mapping", + mapping=mapping_benzonitrile_styrene, + ) + + def test_error_with_multiple_mappings( + self, + protocol_short, + benzene_vacuum_system, + toluene_vacuum_system, + mapping_benzene_toluene, + ): + """Passing a list of more than one mapping should raise ValueError.""" + with pytest.raises(ValueError): + _ = protocol_short.create( + stateA=benzene_vacuum_system, + stateB=toluene_vacuum_system, + name="multiple mappings", + mapping=[mapping_benzene_toluene, mapping_benzene_toluene], + ) + + def test_fail_with_multiple_solvent_comps( + self, + protocol_short, + benzene_solvent_system, + toluene_double_solvent_system, + mapping_benzene_toluene, + ): + """A state with more than one solvent component should raise AssertionError.""" + with pytest.raises(AssertionError): + _ = protocol_short.create( + stateA=benzene_solvent_system, + stateB=toluene_double_solvent_system, + name="double solvent", + mapping=mapping_benzene_toluene, + ) + + def test_create_execute_gather( + self, + short_switching_settings_multiple_switches, + benzene_vacuum_system, + toluene_vacuum_system, + mapping_benzene_toluene, + tmpdir, + ): + """ + Run the switching protocol with multiple switches and gather results. + Checks that the execution is successful and that the free energy estimate + and its uncertainty are not NaN. + """ + import numpy as np + + protocol = NonEquilibriumSwitchingProtocol( + settings=short_switching_settings_multiple_switches + ) + dag = protocol.create( + stateA=benzene_vacuum_system, + stateB=toluene_vacuum_system, + name="Multiple switches vacuum", + mapping=mapping_benzene_toluene, + ) + + results = [] + n_repeats = 4 + for i in range(n_repeats): + with tmpdir.as_cwd(): + shared = Path(f"shared_{i}") + shared.mkdir() + scratch = Path(f"scratch_{i}") + scratch.mkdir() + dagresult = execute_DAG( + dag, shared_basedir=shared, scratch_basedir=scratch + ) + results.append(dagresult) + + for dag_result in results: + assert len(dag_result.protocol_unit_failures) == 0, ( + "Unit failure in protocol dag result." + ) + + protocolresult = protocol.gather(results) + fe_estimate = protocolresult.get_estimate() + fe_error = protocolresult.get_uncertainty() + assert not np.isnan(fe_estimate.magnitude), "Free energy estimate is NaN." + assert not np.isnan(fe_error.magnitude), "Free energy uncertainty is NaN." + + +def test_settings_round_trip(): + """Settings must survive a JSON round-trip unchanged.""" + neq_settings = NonEquilibriumSwitchingProtocol.default_settings() + neq_json = json.dumps(neq_settings.model_dump(), cls=JSON_HANDLER.encoder) + neq_settings_2 = NonEquilibriumSwitchingSettings.model_validate( + json.loads(neq_json, cls=JSON_HANDLER.decoder) + ) + assert neq_settings == neq_settings_2 From df6d66429e6d6c3ffb95de2bf137c45918785d99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 7 Apr 2026 19:52:47 -0400 Subject: [PATCH 10/21] NES protocol firs try. Parallel forward/reverse switch units. accepts eq trajectory. --- feflow/protocols/__init__.py | 2 + feflow/protocols/nonequilibrium_switching.py | 629 +++++++++++++++---- 2 files changed, 498 insertions(+), 133 deletions(-) diff --git a/feflow/protocols/__init__.py b/feflow/protocols/__init__.py index 7ed4813d..86e91be4 100644 --- a/feflow/protocols/__init__.py +++ b/feflow/protocols/__init__.py @@ -10,4 +10,6 @@ from .nonequilibrium_switching import ( NonEquilibriumSwitchingProtocol, NonEquilibriumSwitchingProtocolResult, + ForwardSwitchingUnit, + ReverseSwitchingUnit, ) diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py index 92589302..f2cf0304 100644 --- a/feflow/protocols/nonequilibrium_switching.py +++ b/feflow/protocols/nonequilibrium_switching.py @@ -1,178 +1,541 @@ -import mdtraj -from gufe import Protocol, ProtocolUnit, ProtocolResult +# Nonequilibrium switching protocol using AlchemicalNonequilibriumLangevinIntegrator. +# Reuses SetupUnit from nonequilibrium_cycling for hybrid topology construction. + +import datetime +import logging +import re +import time +from typing import Optional, Any +from collections.abc import Iterable + +from gufe.chemicalsystem import ChemicalSystem +from gufe.mapping import ComponentMapping +from gufe.protocols import ( + Protocol, + ProtocolUnit, + ProtocolResult, + ProtocolDAGResult, +) + +from openfe.protocols.openmm_utils.omm_compute import get_openmm_platform +from openff.units import unit +from openff.units.openmm import to_openmm + +from ..settings import NonEquilibriumSwitchingSettings +from ..settings.nonequilibrium_switching import SnapshotSettings +from ..settings.small_molecules import OpenFFPartialChargeSettings +from ..utils.data import deserialize +from .nonequilibrium_cycling import SetupUnit # reuse hybrid topology setup + +logger = logging.getLogger(__name__) + + +def _reversed_lambda_functions(lambda_functions: dict[str, str]) -> dict[str, str]: + """ + Derive alchemical functions for the reverse switch (lambda 1->0) by replacing + the standalone 'lambda' variable with '(1.0 - lambda)' in all expressions. + Uses a word-boundary regex to avoid partial matches. + """ + return { + k: re.sub(r"\blambda\b", "(1.0 - lambda)", v) + for k, v in lambda_functions.items() + } -class BaseSwitchingUnit(ProtocolUnit): +def _load_snapshot(snapshot_settings: SnapshotSettings, index: int): """ - Monolithic unit for the cycle part of the simulation. - It runs a number of NEq cycles from the outputs of a setup unit and stores the work computed in - numpy-formatted files, to be analyzed by a result unit. + Load positions and box vectors for replicate *index* from a trajectory + using MDAnalysis. + + Parameters + ---------- + snapshot_settings : SnapshotSettings + index : int + Replicate index; selects frame ``index * stride`` from the trajectory. + + Returns + ------- + positions_nm : np.ndarray, shape (n_atoms, 3) + Positions in nanometres. + box_vectors_nm : np.ndarray, shape (3, 3) or None + Triclinic box vectors in nanometres, or None for vacuum systems. + """ + import numpy as np + import MDAnalysis as mda + + u = mda.Universe( + snapshot_settings.topology_file, snapshot_settings.trajectory_file + ) + frame_idx = index * snapshot_settings.stride + u.trajectory[frame_idx] + + positions_nm = u.atoms.positions * 0.1 # Angstrom -> nm + + ts = u.trajectory.ts + if ts.dimensions is not None: + from MDAnalysis.lib.mdamath import triclinic_vectors + box_vectors_nm = triclinic_vectors(ts.dimensions) * 0.1 # Angstrom -> nm + else: + box_vectors_nm = None + + return positions_nm, box_vectors_nm + + +class _BaseSwitchingUnit(ProtocolUnit): + """ + Shared machinery for ForwardSwitchingUnit and ReverseSwitchingUnit. + Subclasses provide ``_direction``, ``_snapshot_settings_key``, and + ``_lambda_functions``. """ @staticmethod def extract_positions(context, initial_atom_indices, final_atom_indices): - """ - Extract positions from initial and final systems based from the hybrid topology. - - Parameters - ---------- - context: openmm.Context - Current simulation context where from extract positions. - hybrid_topology_factory: HybridTopologyFactory - Hybrid topology factory where to extract positions and mapping information - - Returns - ------- - - Notes - ----- - It achieves this by taking the positions and indices from the initial and final states of - the transformation, and computing the overlap of these with the indices of the complete - hybrid topology, filtered by some mdtraj selection expression. - - 1. Get positions from context - 2. Get topology from HTF (already mdtraj topology) - 3. Merge that information into mdtraj.Trajectory - 4. Filter positions for initial/final according to selection string - """ import numpy as np - # Get positions from current openmm context positions = context.getState(getPositions=True).getPositions(asNumpy=True) + return ( + positions[list(initial_atom_indices), :], + positions[list(final_atom_indices), :], + ) - # Get indices for initial and final topologies in hybrid topology - initial_indices = np.asarray(initial_atom_indices) - final_indices = np.asarray(final_atom_indices) + # --- to be overridden by subclasses --- + _direction: str = "" # "forward" or "reverse" + _snapshot_settings_key: str = "" # "lambda0_snapshots" or "lambda1_snapshots" - initial_positions = positions[initial_indices, :] - final_positions = positions[final_indices, :] + def _get_lambda_functions(self, settings: NonEquilibriumSwitchingSettings) -> dict: + raise NotImplementedError - return initial_positions, final_positions + # -------------------------------------------------------------- - def _execute(self, ctx, *, protocol, md_unit, index, **inputs): + def _execute(self, ctx, *, protocol, setup, index, **inputs): """ - Execute the simulation part of the Nonequilibrium switching protocol using GUFE objects. + Execute one NEQ switch: + + 1. Obtain starting snapshot (from trajectory or internal equilibration). + 2. Run the NEQ switch, collecting work and trajectory frames. Parameters ---------- ctx : gufe.protocols.protocolunit.Context - The gufe context for the unit. - protocol : gufe.protocols.Protocol - The Protocol used to create this Unit. Contains key information - such as the settings. - md_unit : gufe.protocols.ProtocolUnit - The SetupUnit - index: int - TODO: Index for the snapshot to use as input - - Returns - ------- - dict : dict[str, str] - Dictionary with paths to work arrays, both forward and reverse, and trajectory coordinates for systems - A and B. + protocol : NonEquilibriumSwitchingProtocol + setup : SetupUnit result with hybrid system and initial state. + index : int + Replicate index; used for frame selection and RNG seeding. """ + import numpy as np import openmm - from openmmtools.integrators import PeriodicNonequilibriumIntegrator - - # Setting up logging to file in shared filesystem - file_logger = logging.getLogger("neq-cycling") - output_log_path = ctx.shared / "feflow-neq-cycling.log" - file_handler = logging.FileHandler(output_log_path, mode="w") - file_handler.setLevel(logging.DEBUG) # TODO: Set to INFO in production - log_formatter = logging.Formatter( - fmt="%(asctime)s %(levelname)-8s %(message)s", datefmt="%Y-%m-%d %H:%M:%S" + import openmm.unit as openmm_unit + from openmmtools.integrators import AlchemicalNonequilibriumLangevinIntegrator + + # Logging + file_logger = logging.getLogger(f"neq-switching-{self._direction}") + log_path = ctx.shared / f"feflow-neq-{self._direction}-{self.name}.log" + file_handler = logging.FileHandler(log_path, mode="w") + file_handler.setLevel(logging.DEBUG) + file_handler.setFormatter( + logging.Formatter( + fmt="%(asctime)s %(levelname)-8s %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) ) - file_handler.setFormatter(log_formatter) file_logger.addHandler(file_handler) - system = deserialize(md_unit.inputs["setup"].outputs["system"]) - state = deserialize(md_unit.inputs["setup"].outputs["state"]) - integrator = deserialize(md_unit.inputs["setup"].outputs["integrator"]) - - PeriodicNonequilibriumIntegrator.restore_interface(integrator) - - # Get atom indices for either end of the hybrid topology + settings: NonEquilibriumSwitchingSettings = protocol.settings + int_settings = settings.integrator_settings + + temperature = to_openmm(settings.thermo_settings.temperature) + timestep = to_openmm(int_settings.timestep) + collision_rate = to_openmm(int_settings.collision_rate) + neq_steps = int_settings.nonequilibrium_steps + eq_steps = int_settings.equilibrium_steps + work_save_freq = settings.work_save_frequency + traj_save_freq = settings.traj_save_frequency + coarse_steps = neq_steps // work_save_freq + traj_coarse_freq = traj_save_freq // work_save_freq + + system = deserialize(setup.outputs["system"]) + initial_state = deserialize(setup.outputs["state"]) initial_atom_indices = setup.outputs["initial_atom_indices"] final_atom_indices = setup.outputs["final_atom_indices"] - # Extract settings from the Protocol - settings = protocol.settings - - # Load positions from snapshots - xtc_file = md_unit.outputs["production_trajectory"] - md_traj_ob = mdtraj.load_frame(xtc_file, index=index) - input_positions = md_traj_ob.openmm_positions(0) - # Set up context platform = get_openmm_platform(settings.engine_settings.compute_platform) - context = openmm.Context(system, integrator, platform) - context.setState(state) - # TODO: This is kinda ugly, is there a better way to set positions? - context.setPositions(input_positions) - - # Setting velocities to temperatures - thermodynamic_settings = settings.thermo_settings - temperature = to_openmm(thermodynamic_settings.temperature) - context.setVelocitiesToTemperature(temperature) - - # Extract settings used below - neq_steps = settings.integrator_settings.nonequilibrium_steps - traj_save_frequency = settings.traj_save_frequency - work_save_frequency = ( - settings.work_save_frequency - ) # Note: this is divisor of traj save freq. - selection_expression = settings.atom_selection_expression - - try: - # Coarse number of steps -- each coarse consists of work_save_frequency steps - coarse_neq_steps = int( - neq_steps / work_save_frequency - ) # Note: neq_steps is multiple of work save steps + timing_info = {} - # TODO: Also get the GPU information (plain try-except with nvidia-smi) + # ------------------------------------------------------------------ + # Step 1: obtain the starting snapshot + # ------------------------------------------------------------------ + snapshot_settings: Optional[SnapshotSettings] = getattr( + settings, self._snapshot_settings_key + ) + if snapshot_settings is not None: + file_logger.info( + f"{self.name}: loading snapshot {index} from " + f"{snapshot_settings.trajectory_file} (frame {index * snapshot_settings.stride})" + ) + positions_nm, box_nm = _load_snapshot(snapshot_settings, index) - integrator.step(NSTEPS) + start_integrator = openmm.LangevinMiddleIntegrator( + temperature, collision_rate, timestep + ) + start_ctx = openmm.Context(system, start_integrator, platform) + start_ctx.setState(initial_state) + start_ctx.setPositions(positions_nm * openmm_unit.nanometers) + if box_nm is not None: + start_ctx.setPeriodicBoxVectors(*box_nm * openmm_unit.nanometers) + start_ctx.setVelocitiesToTemperature(temperature) + snap_state = start_ctx.getState( + getPositions=True, + getVelocities=True, + getEnergy=True, + getForces=True, + enforcePeriodicBox=True, + ) + del start_ctx, start_integrator + else: + file_logger.info( + f"{self.name}: equilibrating for {eq_steps} steps " + f"(replicate {index})" + ) + eq_integrator = openmm.LangevinMiddleIntegrator( + temperature, collision_rate, timestep + ) + eq_ctx = openmm.Context(system, eq_integrator, platform) + eq_ctx.setState(initial_state) + # Use index as seed so replicates decorrelate from each other + eq_ctx.setVelocitiesToTemperature(temperature, index) + + t0 = time.perf_counter() + eq_integrator.step(eq_steps) + timing_info["eq_time_in_s"] = datetime.timedelta( + seconds=time.perf_counter() - t0 + ).total_seconds() + + snap_state = eq_ctx.getState( + getPositions=True, + getVelocities=True, + getEnergy=True, + getForces=True, + enforcePeriodicBox=True, + ) + del eq_ctx, eq_integrator + file_logger.info( + f"{self.name}: equilibration done ({timing_info['eq_time_in_s']:.1f} s)" + ) + # ------------------------------------------------------------------ + # Step 2: NEQ switch + # ------------------------------------------------------------------ + lambda_functions = self._get_lambda_functions(settings) + + neq_integrator = AlchemicalNonequilibriumLangevinIntegrator( + alchemical_functions=lambda_functions, + splitting=int_settings.splitting, + temperature=temperature, + collision_rate=collision_rate, + timestep=timestep, + nsteps_neq=neq_steps, + ) + neq_ctx = openmm.Context(system, neq_integrator, platform) + neq_ctx.setState(snap_state) + works = [neq_integrator.get_protocol_work(dimensionless=True)] + initial_traj, final_traj = [], [] - # Equilibrium (lambda = 0) - # start timer - start_time = time.perf_counter() - # Run neq - # Forward (0 -> 1) - # Initialize works with current value - forward_works = [] - for fwd_step in range(coarse_neq_steps): - integrator.step(work_save_frequency) - forward_works.append(integrator.get_protocol_work(dimensionless=True)) - if fwd_step % traj_save_frequency == 0: - initial_positions, final_positions = self.extract_positions( - context, initial_atom_indices, final_atom_indices + file_logger.info( + f"{self.name}: {self._direction} NEQ switch ({neq_steps} steps)" + ) + t0 = time.perf_counter() + try: + for step in range(coarse_steps): + neq_integrator.step(work_save_freq) + works.append(neq_integrator.get_protocol_work(dimensionless=True)) + if step % traj_coarse_freq == 0: + i_pos, f_pos = self.extract_positions( + neq_ctx, initial_atom_indices, final_atom_indices ) - forward_neq_initial.append(initial_positions) - forward_neq_final.append(final_positions) - # Make sure trajectories are stored at the end of the neq loop - initial_positions, final_positions = self.extract_positions( - context, initial_atom_indices, final_atom_indices + initial_traj.append(i_pos) + final_traj.append(f_pos) + # Always capture the final frame + i_pos, f_pos = self.extract_positions( + neq_ctx, initial_atom_indices, final_atom_indices ) - forward_neq_initial.append(initial_positions) - forward_neq_final.append(final_positions) + initial_traj.append(i_pos) + final_traj.append(f_pos) + finally: + del neq_ctx, neq_integrator + + neq_elapsed = datetime.timedelta(seconds=time.perf_counter() - t0) + timing_info[f"neq_{self._direction}_time_in_s"] = neq_elapsed.total_seconds() + file_logger.info(f"{self.name}: NEQ switch time: {neq_elapsed}") + + # ------------------------------------------------------------------ + # Serialize outputs + # ------------------------------------------------------------------ + work_path = ctx.shared / f"{self._direction}_{self.name}.npy" + initial_traj_path = ctx.shared / f"{self._direction}_initial_{self.name}.npy" + final_traj_path = ctx.shared / f"{self._direction}_final_{self.name}.npy" + + with open(work_path, "wb") as f: + np.save(f, works) + with open(initial_traj_path, "wb") as f: + np.save(f, np.array(initial_traj)) + with open(final_traj_path, "wb") as f: + np.save(f, np.array(final_traj)) + + return { + "work": work_path, + "initial_traj": initial_traj_path, + "final_traj": final_traj_path, + "timing_info": timing_info, + "log": log_path, + } + + +class ForwardSwitchingUnit(_BaseSwitchingUnit): + """Runs one forward NEQ switch (lambda 0->1).""" + + _direction = "forward" + _snapshot_settings_key = "lambda0_snapshots" + + def _get_lambda_functions(self, settings): + return settings.lambda_functions + + +class ReverseSwitchingUnit(_BaseSwitchingUnit): + """Runs one reverse NEQ switch (lambda 1->0).""" + + _direction = "reverse" + _snapshot_settings_key = "lambda1_snapshots" - neq_forward_time = time.perf_counter() - neq_forward_walltime = datetime.timedelta( - seconds=neq_forward_time - eq_forward_time + def _get_lambda_functions(self, settings): + return _reversed_lambda_functions(settings.lambda_functions) + + +class ResultUnit(ProtocolUnit): + """Collects per-switch work arrays from all ForwardSwitchingUnits and ReverseSwitchingUnits.""" + + @staticmethod + def _execute(ctx, *, forward_switches, reverse_switches, **inputs): + import numpy as np + + forward_works, reverse_works = [], [] + for unit in forward_switches: + w = np.load(unit.outputs["work"]) + forward_works.append(w - w[0]) + for unit in reverse_switches: + w = np.load(unit.outputs["work"]) + reverse_works.append(w - w[0]) + + return { + "forward_work": forward_works, + "reverse_work": reverse_works, + "forward_work_paths": [u.outputs["work"] for u in forward_switches], + "reverse_work_paths": [u.outputs["work"] for u in reverse_switches], + } + + +class NonEquilibriumSwitchingProtocolResult(ProtocolResult): + """ + Collects results and computes free energy estimates via BAR from the + forward (lambda 0->1) and reverse (lambda 1->0) work values. + """ + + def get_estimate(self): + """Free energy estimate via BAR in kcal/mol.""" + import numpy as np + import numpy.typing as npt + import pymbar + + forward: npt.NDArray = np.array([w[-1] for w in self.data["forward_work"]]) + reverse: npt.NDArray = np.array([w[-1] for w in self.data["reverse_work"]]) + bar_data = pymbar.bar(forward, reverse) + return ( + bar_data["Delta_f"] + * unit.k + * self.data["temperature"] + * unit.avogadro_constant + ).to("kcal/mol") + + def get_uncertainty(self, n_bootstraps: int = 1000): + """Uncertainty via bootstrapped BAR standard deviation in kcal/mol.""" + import numpy as np + import numpy.typing as npt + + forward: npt.NDArray = np.array([w[-1] for w in self.data["forward_work"]]) + reverse: npt.NDArray = np.array([w[-1] for w in self.data["reverse_work"]]) + all_dgs = self._do_bootstrap(forward, reverse, n_bootstraps) + return ( + np.std(all_dgs) + * unit.k + * self.data["temperature"] + * unit.avogadro_constant + ).to("kcal/mol") + + def get_rate_of_convergence(self): ... + + def _do_bootstrap(self, forward, reverse, n_bootstraps: int = 1000): + import numpy as np + import pymbar + + assert len(forward) == len(reverse), ( + "Forward and reverse work arrays must have the same length." + ) + n = len(forward) + all_dgs = np.zeros(n_bootstraps) + for i in range(n_bootstraps): + idx = np.random.choice(np.arange(n), size=n, replace=True) + all_dgs[i] = pymbar.bar(forward[idx], reverse[idx])["Delta_f"] + return all_dgs + + +class NonEquilibriumSwitchingProtocol(Protocol): + """ + RBFE protocol using non-equilibrium switching with the + AlchemicalNonequilibriumLangevinIntegrator from openmmtools. + + For each of ``num_switches`` replicates the protocol creates: + - one ForwardSwitchingUnit (lambda 0->1) + - one ReverseSwitchingUnit (lambda 1->0) + + Both sets of units depend only on the shared SetupUnit and can therefore + run in parallel. Free energy is estimated via BAR over all work values. + + Starting snapshots are either generated by internal equilibration + (``integrator_settings.equilibrium_steps > 0``) or loaded from + pre-equilibrated trajectories via ``lambda0_snapshots`` / ``lambda1_snapshots``. + """ + + _settings_cls = NonEquilibriumSwitchingSettings + result_cls = NonEquilibriumSwitchingProtocolResult + + @classmethod + def _default_settings(cls): + from feflow.settings import ( + NonEquilibriumSwitchingSettings, + AlchemicalNonequilibriumIntegratorSettings, + ) + from gufe.settings import OpenMMSystemGeneratorFFSettings + from openfe.protocols.openmm_utils.omm_settings import ( + OpenMMSolvationSettings, + OpenMMEngineSettings, + ThermoSettings, + ) + from openfe.protocols.openmm_rfe.equil_rfe_settings import AlchemicalSettings + + return NonEquilibriumSwitchingSettings( + forcefield_settings=OpenMMSystemGeneratorFFSettings(), + thermo_settings=ThermoSettings( + temperature=300 * unit.kelvin, pressure=1 * unit.bar + ), + solvation_settings=OpenMMSolvationSettings(), + partial_charge_settings=OpenFFPartialChargeSettings(), + alchemical_settings=AlchemicalSettings(softcore_LJ="gapsys"), + integrator_settings=AlchemicalNonequilibriumIntegratorSettings(), + engine_settings=OpenMMEngineSettings(), + ) + + def _create( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: Optional[ComponentMapping | list[ComponentMapping]], + extends: Optional[ProtocolDAGResult] = None, + ) -> list[ProtocolUnit]: + if isinstance(mapping, list): + if len(mapping) != 1: + raise ValueError( + "Exactly one mapping must be provided for this protocol." + ) + mapping = mapping[0] + self.validate(stateA=stateA, stateB=stateB, mapping=mapping, extends=extends) + + num_switches = self.settings.num_switches + + setup = SetupUnit( + protocol=self, + state_a=stateA, + state_b=stateB, + mapping=mapping, + name="setup", + ) + + # Forward and reverse units both depend only on setup — they can run + # in parallel with each other and across replicates. + forward_switches = [ + ForwardSwitchingUnit( + protocol=self, setup=setup, index=i, name=f"forward_{i}" ) - file_logger.info( - f"replicate_{self.name} Forward nonequilibrium time (lambda 0 -> 1): {neq_forward_walltime}" + for i in range(num_switches) + ] + reverse_switches = [ + ReverseSwitchingUnit( + protocol=self, setup=setup, index=i, name=f"reverse_{i}" ) + for i in range(num_switches) + ] - # TODO: We should return the work in one direction + end = ResultUnit( + name="result", + forward_switches=forward_switches, + reverse_switches=reverse_switches, + ) + return [setup, *forward_switches, *reverse_switches, end] -class NonEquilibriumSwitchingProtocol(Protocol): - ... + @staticmethod + def _check_mappings_consistency(mapping, chemical_system_a, chemical_system_b): + mapping_comp_a = mapping.componentA + mapping_comp_b = mapping.componentB + chem_sys_a_keys = [c.key for _, c in chemical_system_a.components.items()] + chem_sys_b_keys = [c.key for _, c in chemical_system_b.components.items()] + assert mapping_comp_a.key in chem_sys_a_keys, ( + "Component A in mapping not found in chemical system A." + ) + assert mapping_comp_b.key in chem_sys_b_keys, ( + "Component B in mapping not found in chemical system B." + ) + + def _validate( + self, + *, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: ComponentMapping | list[ComponentMapping] | None, + extends: Optional[ProtocolDAGResult] = None, + ): + from gufe import SolventComponent + + if mapping is None: + raise ValueError("`mapping` is required for this Protocol") + if extends: + raise NotImplementedError("Can't extend simulations yet") + + self._check_mappings_consistency( + mapping=mapping, chemical_system_a=stateA, chemical_system_b=stateB + ) + + state_a_solv = stateA.get_components_of_type(SolventComponent) + state_b_solv = stateB.get_components_of_type(SolventComponent) + assert len(state_a_solv) <= 1, ( + f"State A has {len(state_a_solv)} solvent components. Only 0 or 1 allowed." + ) + assert len(state_b_solv) <= 1, ( + f"State B has {len(state_b_solv)} solvent components. Only 0 or 1 allowed." + ) + def _gather( + self, protocol_dag_results: Iterable[ProtocolDAGResult] + ) -> dict[str, Any]: + from collections import defaultdict -class NonEquilibriumSwitchingProtocolResult(ProtocolResult): - ... \ No newline at end of file + outputs = defaultdict(list) + for pdr in protocol_dag_results: + for pur in pdr.protocol_unit_results: + if pur.name == "result": + outputs["forward_work"].extend(pur.outputs["forward_work"]) + outputs["reverse_work"].extend(pur.outputs["reverse_work"]) + + outputs["temperature"] = self.settings.thermo_settings.temperature + return outputs From 45f4bf43974c00de055c7e3201e768c6cb331150 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 7 Apr 2026 20:10:35 -0400 Subject: [PATCH 11/21] adding remove_com attribute to base settings --- feflow/settings/integrators.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/feflow/settings/integrators.py b/feflow/settings/integrators.py index 40bc37dd..54436fc4 100644 --- a/feflow/settings/integrators.py +++ b/feflow/settings/integrators.py @@ -32,6 +32,10 @@ class BaseNonequilibriumIntegrator(SettingsBaseModel): timestep: FemtosecondQuantity = 4 * unit.femtoseconds """Size of the simulation timestep. Default 4 fs.""" splitting: str = "V R H O R V" + remove_com: bool = False + """ + Whether or not to remove the center of mass motion. Default False. + """ # TODO: This validator is used in other settings, better create a new Type @field_validator("timestep") @@ -98,10 +102,6 @@ class PeriodicNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): Note: The barostat frequency is ignored for gas-phase simulations. Default 25 * unit.timestep. """ - remove_com: bool = False - """ - Whether or not to remove the center of mass motion. Default False. - """ # TODO: This validator is used in other settings, better create a new Type @field_validator("equilibrium_steps", "nonequilibrium_steps") From e9020c0b57bde4ea3c3520aafa88fe1746353522 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 7 Apr 2026 20:10:51 -0400 Subject: [PATCH 12/21] adding store minimized to switching settings --- feflow/settings/nonequilibrium_switching.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/feflow/settings/nonequilibrium_switching.py b/feflow/settings/nonequilibrium_switching.py index 461b9e56..45826f33 100644 --- a/feflow/settings/nonequilibrium_switching.py +++ b/feflow/settings/nonequilibrium_switching.py @@ -127,6 +127,10 @@ class NonEquilibriumSwitchingSettings(Settings): runs. Each switch produces one work value used in the BAR free energy estimate. """ + # Debugging settings + store_minimized_pdb: bool = True + """Setting for storing pdb right after minimization (right before neq cycle)""" + @root_validator def set_and_validate_save_frequencies(cls, values): """ From f0cbbdf89e04c3706fd186f45e794e0302d4d8ae Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Apr 2026 00:13:09 +0000 Subject: [PATCH 13/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- feflow/protocols/nonequilibrium_switching.py | 40 +++++++++---------- feflow/settings/nonequilibrium_switching.py | 12 ++++-- feflow/tests/test_nonequilibrium_switching.py | 6 +-- 3 files changed, 30 insertions(+), 28 deletions(-) diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py index f2cf0304..ffe0548a 100644 --- a/feflow/protocols/nonequilibrium_switching.py +++ b/feflow/protocols/nonequilibrium_switching.py @@ -63,9 +63,7 @@ def _load_snapshot(snapshot_settings: SnapshotSettings, index: int): import numpy as np import MDAnalysis as mda - u = mda.Universe( - snapshot_settings.topology_file, snapshot_settings.trajectory_file - ) + u = mda.Universe(snapshot_settings.topology_file, snapshot_settings.trajectory_file) frame_idx = index * snapshot_settings.stride u.trajectory[frame_idx] @@ -74,6 +72,7 @@ def _load_snapshot(snapshot_settings: SnapshotSettings, index: int): ts = u.trajectory.ts if ts.dimensions is not None: from MDAnalysis.lib.mdamath import triclinic_vectors + box_vectors_nm = triclinic_vectors(ts.dimensions) * 0.1 # Angstrom -> nm else: box_vectors_nm = None @@ -367,10 +366,7 @@ def get_uncertainty(self, n_bootstraps: int = 1000): reverse: npt.NDArray = np.array([w[-1] for w in self.data["reverse_work"]]) all_dgs = self._do_bootstrap(forward, reverse, n_bootstraps) return ( - np.std(all_dgs) - * unit.k - * self.data["temperature"] - * unit.avogadro_constant + np.std(all_dgs) * unit.k * self.data["temperature"] * unit.avogadro_constant ).to("kcal/mol") def get_rate_of_convergence(self): ... @@ -379,9 +375,9 @@ def _do_bootstrap(self, forward, reverse, n_bootstraps: int = 1000): import numpy as np import pymbar - assert len(forward) == len(reverse), ( - "Forward and reverse work arrays must have the same length." - ) + assert len(forward) == len( + reverse + ), "Forward and reverse work arrays must have the same length." n = len(forward) all_dgs = np.zeros(n_bootstraps) for i in range(n_bootstraps): @@ -490,12 +486,12 @@ def _check_mappings_consistency(mapping, chemical_system_a, chemical_system_b): mapping_comp_b = mapping.componentB chem_sys_a_keys = [c.key for _, c in chemical_system_a.components.items()] chem_sys_b_keys = [c.key for _, c in chemical_system_b.components.items()] - assert mapping_comp_a.key in chem_sys_a_keys, ( - "Component A in mapping not found in chemical system A." - ) - assert mapping_comp_b.key in chem_sys_b_keys, ( - "Component B in mapping not found in chemical system B." - ) + assert ( + mapping_comp_a.key in chem_sys_a_keys + ), "Component A in mapping not found in chemical system A." + assert ( + mapping_comp_b.key in chem_sys_b_keys + ), "Component B in mapping not found in chemical system B." def _validate( self, @@ -518,12 +514,12 @@ def _validate( state_a_solv = stateA.get_components_of_type(SolventComponent) state_b_solv = stateB.get_components_of_type(SolventComponent) - assert len(state_a_solv) <= 1, ( - f"State A has {len(state_a_solv)} solvent components. Only 0 or 1 allowed." - ) - assert len(state_b_solv) <= 1, ( - f"State B has {len(state_b_solv)} solvent components. Only 0 or 1 allowed." - ) + assert ( + len(state_a_solv) <= 1 + ), f"State A has {len(state_a_solv)} solvent components. Only 0 or 1 allowed." + assert ( + len(state_b_solv) <= 1 + ), f"State B has {len(state_b_solv)} solvent components. Only 0 or 1 allowed." def _gather( self, protocol_dag_results: Iterable[ProtocolDAGResult] diff --git a/feflow/settings/nonequilibrium_switching.py b/feflow/settings/nonequilibrium_switching.py index 45826f33..f75ed62e 100644 --- a/feflow/settings/nonequilibrium_switching.py +++ b/feflow/settings/nonequilibrium_switching.py @@ -169,7 +169,9 @@ def snapshots_require_no_internal_equilibration(cls, values): if integrator_settings is None: return values - has_snapshots = values.get("lambda0_snapshots") or values.get("lambda1_snapshots") + has_snapshots = values.get("lambda0_snapshots") or values.get( + "lambda1_snapshots" + ) if has_snapshots and integrator_settings.equilibrium_steps != 0: raise ValueError( "When lambda0_snapshots or lambda1_snapshots are provided the " @@ -198,8 +200,12 @@ def snapshot_trajectories_have_enough_frames(cls, values): try: import MDAnalysis as mda - n0 = len(mda.Universe(snap0.topology_file, snap0.trajectory_file).trajectory) - n1 = len(mda.Universe(snap1.topology_file, snap1.trajectory_file).trajectory) + n0 = len( + mda.Universe(snap0.topology_file, snap0.trajectory_file).trajectory + ) + n1 = len( + mda.Universe(snap1.topology_file, snap1.trajectory_file).trajectory + ) except Exception as exc: raise ValueError( f"Could not read snapshot trajectories to validate frame counts: {exc}" diff --git a/feflow/tests/test_nonequilibrium_switching.py b/feflow/tests/test_nonequilibrium_switching.py index cc2a1584..21f2b04d 100644 --- a/feflow/tests/test_nonequilibrium_switching.py +++ b/feflow/tests/test_nonequilibrium_switching.py @@ -188,9 +188,9 @@ def test_create_execute_gather( results.append(dagresult) for dag_result in results: - assert len(dag_result.protocol_unit_failures) == 0, ( - "Unit failure in protocol dag result." - ) + assert ( + len(dag_result.protocol_unit_failures) == 0 + ), "Unit failure in protocol dag result." protocolresult = protocol.gather(results) fe_estimate = protocolresult.get_estimate() From 7a07ae65a6a89815a411fe1a34c1fb3f78e59799 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 7 Apr 2026 20:17:53 -0400 Subject: [PATCH 14/21] bump CI From 596f2c952fe54c8588880354303bcb5dc76ff473 Mon Sep 17 00:00:00 2001 From: ijpulidos <2949729+ijpulidos@users.noreply.github.com> Date: Wed, 3 Jun 2026 19:23:06 -0400 Subject: [PATCH 15/21] Separating equilibration in own units. Avoiding rerunning every time --- feflow/protocols/nonequilibrium_switching.py | 299 +++++++++++++------ 1 file changed, 214 insertions(+), 85 deletions(-) diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py index ffe0548a..c6f7cfff 100644 --- a/feflow/protocols/nonequilibrium_switching.py +++ b/feflow/protocols/nonequilibrium_switching.py @@ -24,7 +24,7 @@ from ..settings import NonEquilibriumSwitchingSettings from ..settings.nonequilibrium_switching import SnapshotSettings from ..settings.small_molecules import OpenFFPartialChargeSettings -from ..utils.data import deserialize +from ..utils.data import deserialize, serialize from .nonequilibrium_cycling import SetupUnit # reuse hybrid topology setup logger = logging.getLogger(__name__) @@ -80,11 +80,191 @@ def _load_snapshot(snapshot_settings: SnapshotSettings, index: int): return positions_nm, box_vectors_nm +class _BaseEquilibrationUnit(ProtocolUnit): + """ + Produces ``num_switches`` equilibrated starting snapshots for one lambda + endpoint. Subclasses set ``_snapshot_settings_key`` and ``_endpoint``. + + Two modes: + - **Internal equilibration**: runs a single continuous Langevin trajectory + for ``equilibrium_steps`` total steps, saving ``num_switches`` snapshots + at uniform intervals (every ``equilibrium_steps // num_switches`` steps). + Raises if ``equilibrium_steps < num_switches``. + - **XTC trajectory**: loads ``num_switches`` frames from a pre-equilibrated + trajectory via MDAnalysis (frame ``i * stride``). Raises if the + trajectory does not contain enough frames. + + Outputs + ------- + snap_states : list[pathlib.Path] + Serialized OpenMM State XML files, one per switch replicate. + timing_info : dict + log : pathlib.Path + """ + + _snapshot_settings_key: str = "" # "lambda0_snapshots" or "lambda1_snapshots" + _endpoint: str = "" # "lambda0" or "lambda1" + + def _execute(self, ctx, *, protocol, setup, **inputs): + import openmm + import openmm.unit as openmm_unit + + settings: NonEquilibriumSwitchingSettings = protocol.settings + int_settings = settings.integrator_settings + + temperature = to_openmm(settings.thermo_settings.temperature) + timestep = to_openmm(int_settings.timestep) + collision_rate = to_openmm(int_settings.collision_rate) + eq_steps = int_settings.equilibrium_steps + num_switches = settings.num_switches + + file_logger = logging.getLogger(f"neq-eq-{self._endpoint}") + log_path = ctx.shared / f"feflow-eq-{self._endpoint}-{self.name}.log" + file_handler = logging.FileHandler(log_path, mode="w") + file_handler.setLevel(logging.DEBUG) + file_handler.setFormatter( + logging.Formatter( + fmt="%(asctime)s %(levelname)-8s %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + ) + file_logger.addHandler(file_handler) + + system = deserialize(setup.outputs["system"]) + initial_state = deserialize(setup.outputs["state"]) + platform = get_openmm_platform(settings.engine_settings.compute_platform) + timing_info = {} + + snapshot_settings: Optional[SnapshotSettings] = getattr( + settings, self._snapshot_settings_key + ) + + snap_states = [] + + if snapshot_settings is not None: + # Load num_switches frames from a pre-equilibrated XTC trajectory + import MDAnalysis as mda + + u = mda.Universe( + snapshot_settings.topology_file, snapshot_settings.trajectory_file + ) + available = len(u.trajectory) // snapshot_settings.stride + if available < num_switches: + raise ValueError( + f"{self._endpoint} trajectory has only {len(u.trajectory)} frames " + f"(stride={snapshot_settings.stride} → {available} usable), " + f"but num_switches={num_switches}." + ) + + file_logger.info( + f"{self.name}: loading {num_switches} snapshots from " + f"{snapshot_settings.trajectory_file} (stride={snapshot_settings.stride})" + ) + integrator = openmm.LangevinMiddleIntegrator( + temperature, collision_rate, timestep + ) + ctx_snap = openmm.Context(system, integrator, platform) + ctx_snap.setState(initial_state) + + t0 = time.perf_counter() + for i in range(num_switches): + positions_nm, box_nm = _load_snapshot(snapshot_settings, i) + ctx_snap.setPositions(positions_nm * openmm_unit.nanometers) + if box_nm is not None: + ctx_snap.setPeriodicBoxVectors(*box_nm * openmm_unit.nanometers) + ctx_snap.setVelocitiesToTemperature(temperature) + snap_states.append( + ctx_snap.getState( + getPositions=True, + getVelocities=True, + getEnergy=True, + getForces=True, + enforcePeriodicBox=True, + ) + ) + del ctx_snap, integrator + timing_info["snapshot_load_time_in_s"] = datetime.timedelta( + seconds=time.perf_counter() - t0 + ).total_seconds() + file_logger.info( + f"{self.name}: loaded {num_switches} snapshots " + f"({timing_info['snapshot_load_time_in_s']:.1f} s)" + ) + + else: + # Internal equilibration: single continuous trajectory of eq_steps + # total steps; snapshots taken at uniform intervals. + if eq_steps < num_switches: + raise ValueError( + f"equilibrium_steps ({eq_steps}) must be >= num_switches " + f"({num_switches}) to produce uniformly-spaced snapshots." + ) + + save_interval = eq_steps // num_switches + file_logger.info( + f"{self.name}: running {eq_steps} equilibration steps, " + f"saving {num_switches} snapshots every {save_interval} steps" + ) + eq_integrator = openmm.LangevinMiddleIntegrator( + temperature, collision_rate, timestep + ) + eq_ctx = openmm.Context(system, eq_integrator, platform) + eq_ctx.setState(initial_state) + eq_ctx.setVelocitiesToTemperature(temperature) + + t0 = time.perf_counter() + for i in range(num_switches): + # We run eq simulations and save snapshots where each switch should start + eq_integrator.step(save_interval) + snap_states.append( + eq_ctx.getState( + getPositions=True, + getVelocities=True, + getEnergy=True, + getForces=True, + enforcePeriodicBox=True, + ) + ) + del eq_ctx, eq_integrator + timing_info["eq_time_in_s"] = datetime.timedelta( + seconds=time.perf_counter() - t0 + ).total_seconds() + file_logger.info( + f"{self.name}: equilibration done ({timing_info['eq_time_in_s']:.1f} s)" + ) + + # Serialize each snapshot to its own XML file + snap_paths = [] + for i, state in enumerate(snap_states): + path = ctx.shared / f"{self._endpoint}_snapshot_state_{self.name}_{i}.xml" + serialize(state, path) + snap_paths.append(path) + + return { + "snap_states": snap_paths, + "timing_info": timing_info, + "log": log_path, + } + + +class Lambda0EquilibrationUnit(_BaseEquilibrationUnit): + """Equilibration / snapshot loading at lambda=0 (starting point for forward switches).""" + + _snapshot_settings_key = "lambda0_snapshots" + _endpoint = "lambda0" + + +class Lambda1EquilibrationUnit(_BaseEquilibrationUnit): + """Equilibration / snapshot loading at lambda=1 (starting point for reverse switches).""" + + _snapshot_settings_key = "lambda1_snapshots" + _endpoint = "lambda1" + + class _BaseSwitchingUnit(ProtocolUnit): """ Shared machinery for ForwardSwitchingUnit and ReverseSwitchingUnit. - Subclasses provide ``_direction``, ``_snapshot_settings_key``, and - ``_lambda_functions``. + Subclasses provide ``_direction`` and ``_get_lambda_functions``. """ @staticmethod @@ -99,31 +279,28 @@ def extract_positions(context, initial_atom_indices, final_atom_indices): # --- to be overridden by subclasses --- _direction: str = "" # "forward" or "reverse" - _snapshot_settings_key: str = "" # "lambda0_snapshots" or "lambda1_snapshots" def _get_lambda_functions(self, settings: NonEquilibriumSwitchingSettings) -> dict: raise NotImplementedError # -------------------------------------------------------------- - def _execute(self, ctx, *, protocol, setup, index, **inputs): + def _execute(self, ctx, *, protocol, setup, equilibration, index, **inputs): """ - Execute one NEQ switch: - - 1. Obtain starting snapshot (from trajectory or internal equilibration). - 2. Run the NEQ switch, collecting work and trajectory frames. + Execute one NEQ switch starting from the pre-equilibrated snapshot + produced by an upstream EquilibrationUnit. Parameters ---------- ctx : gufe.protocols.protocolunit.Context protocol : NonEquilibriumSwitchingProtocol - setup : SetupUnit result with hybrid system and initial state. + setup : SetupUnit result with hybrid system and atom indices. + equilibration : EquilibrationUnit result with the starting snapshot. index : int - Replicate index; used for frame selection and RNG seeding. + Replicate index; used for output file naming. """ import numpy as np import openmm - import openmm.unit as openmm_unit from openmmtools.integrators import AlchemicalNonequilibriumLangevinIntegrator # Logging @@ -146,85 +323,22 @@ def _execute(self, ctx, *, protocol, setup, index, **inputs): timestep = to_openmm(int_settings.timestep) collision_rate = to_openmm(int_settings.collision_rate) neq_steps = int_settings.nonequilibrium_steps - eq_steps = int_settings.equilibrium_steps work_save_freq = settings.work_save_frequency traj_save_freq = settings.traj_save_frequency coarse_steps = neq_steps // work_save_freq traj_coarse_freq = traj_save_freq // work_save_freq system = deserialize(setup.outputs["system"]) - initial_state = deserialize(setup.outputs["state"]) initial_atom_indices = setup.outputs["initial_atom_indices"] final_atom_indices = setup.outputs["final_atom_indices"] + snap_state = deserialize(equilibration.outputs["snap_states"][index]) + platform = get_openmm_platform(settings.engine_settings.compute_platform) timing_info = {} # ------------------------------------------------------------------ - # Step 1: obtain the starting snapshot - # ------------------------------------------------------------------ - snapshot_settings: Optional[SnapshotSettings] = getattr( - settings, self._snapshot_settings_key - ) - - if snapshot_settings is not None: - file_logger.info( - f"{self.name}: loading snapshot {index} from " - f"{snapshot_settings.trajectory_file} (frame {index * snapshot_settings.stride})" - ) - positions_nm, box_nm = _load_snapshot(snapshot_settings, index) - - start_integrator = openmm.LangevinMiddleIntegrator( - temperature, collision_rate, timestep - ) - start_ctx = openmm.Context(system, start_integrator, platform) - start_ctx.setState(initial_state) - start_ctx.setPositions(positions_nm * openmm_unit.nanometers) - if box_nm is not None: - start_ctx.setPeriodicBoxVectors(*box_nm * openmm_unit.nanometers) - start_ctx.setVelocitiesToTemperature(temperature) - snap_state = start_ctx.getState( - getPositions=True, - getVelocities=True, - getEnergy=True, - getForces=True, - enforcePeriodicBox=True, - ) - del start_ctx, start_integrator - - else: - file_logger.info( - f"{self.name}: equilibrating for {eq_steps} steps " - f"(replicate {index})" - ) - eq_integrator = openmm.LangevinMiddleIntegrator( - temperature, collision_rate, timestep - ) - eq_ctx = openmm.Context(system, eq_integrator, platform) - eq_ctx.setState(initial_state) - # Use index as seed so replicates decorrelate from each other - eq_ctx.setVelocitiesToTemperature(temperature, index) - - t0 = time.perf_counter() - eq_integrator.step(eq_steps) - timing_info["eq_time_in_s"] = datetime.timedelta( - seconds=time.perf_counter() - t0 - ).total_seconds() - - snap_state = eq_ctx.getState( - getPositions=True, - getVelocities=True, - getEnergy=True, - getForces=True, - enforcePeriodicBox=True, - ) - del eq_ctx, eq_integrator - file_logger.info( - f"{self.name}: equilibration done ({timing_info['eq_time_in_s']:.1f} s)" - ) - - # ------------------------------------------------------------------ - # Step 2: NEQ switch + # NEQ switch # ------------------------------------------------------------------ lambda_functions = self._get_lambda_functions(settings) @@ -296,7 +410,6 @@ class ForwardSwitchingUnit(_BaseSwitchingUnit): """Runs one forward NEQ switch (lambda 0->1).""" _direction = "forward" - _snapshot_settings_key = "lambda0_snapshots" def _get_lambda_functions(self, settings): return settings.lambda_functions @@ -306,7 +419,6 @@ class ReverseSwitchingUnit(_BaseSwitchingUnit): """Runs one reverse NEQ switch (lambda 1->0).""" _direction = "reverse" - _snapshot_settings_key = "lambda1_snapshots" def _get_lambda_functions(self, settings): return _reversed_lambda_functions(settings.lambda_functions) @@ -457,17 +569,34 @@ def _create( name="setup", ) - # Forward and reverse units both depend only on setup — they can run - # in parallel with each other and across replicates. + # One equilibration unit per lambda endpoint; each produces num_switches + # snapshots and depends only on setup. + lambda0_eq = Lambda0EquilibrationUnit( + protocol=self, setup=setup, name="eq_lambda0" + ) + lambda1_eq = Lambda1EquilibrationUnit( + protocol=self, setup=setup, name="eq_lambda1" + ) + + # Switching units depend on setup (system/indices) and the appropriate + # equilibration unit (starting snapshot by index). forward_switches = [ ForwardSwitchingUnit( - protocol=self, setup=setup, index=i, name=f"forward_{i}" + protocol=self, + setup=setup, + equilibration=lambda0_eq, + index=i, + name=f"forward_{i}", ) for i in range(num_switches) ] reverse_switches = [ ReverseSwitchingUnit( - protocol=self, setup=setup, index=i, name=f"reverse_{i}" + protocol=self, + setup=setup, + equilibration=lambda1_eq, + index=i, + name=f"reverse_{i}", ) for i in range(num_switches) ] @@ -478,7 +607,7 @@ def _create( reverse_switches=reverse_switches, ) - return [setup, *forward_switches, *reverse_switches, end] + return [setup, lambda0_eq, lambda1_eq, *forward_switches, *reverse_switches, end] @staticmethod def _check_mappings_consistency(mapping, chemical_system_a, chemical_system_b): From 7467aeabe3b7d68dd4c65f41c913b34028b6ddf7 Mon Sep 17 00:00:00 2001 From: ijpulidos <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 23 Jun 2026 20:28:07 -0400 Subject: [PATCH 16/21] using pydantic v2 model_validator --- feflow/settings/nonequilibrium_switching.py | 60 +++++++++------------ 1 file changed, 26 insertions(+), 34 deletions(-) diff --git a/feflow/settings/nonequilibrium_switching.py b/feflow/settings/nonequilibrium_switching.py index f75ed62e..022be1a1 100644 --- a/feflow/settings/nonequilibrium_switching.py +++ b/feflow/settings/nonequilibrium_switching.py @@ -10,14 +10,13 @@ ) from gufe.settings import Settings, OpenMMSystemGeneratorFFSettings, SettingsBaseModel -from pydantic.v1 import root_validator from openfe.protocols.openmm_utils.omm_settings import ( OpenMMSolvationSettings, OpenMMEngineSettings, ThermoSettings, ) from openfe.protocols.openmm_rfe.equil_rfe_settings import AlchemicalSettings -from pydantic import ConfigDict +from pydantic import ConfigDict, model_validator # Default settings for the lambda functions @@ -131,71 +130,64 @@ class NonEquilibriumSwitchingSettings(Settings): store_minimized_pdb: bool = True """Setting for storing pdb right after minimization (right before neq cycle)""" - @root_validator - def set_and_validate_save_frequencies(cls, values): + @model_validator(mode="after") + def set_and_validate_save_frequencies(self): """ Derive save-frequency defaults from nonequilibrium_steps when not set, then check consistency. """ - integrator_settings = values.get("integrator_settings") neq_steps = ( - integrator_settings.nonequilibrium_steps if integrator_settings else 2500 + self.integrator_settings.nonequilibrium_steps + if self.integrator_settings + else 2500 ) - work_freq = values.get("work_save_frequency") - traj_freq = values.get("traj_save_frequency") + if self.work_save_frequency is None: + self.work_save_frequency = max(1, neq_steps // 50) + if self.traj_save_frequency is None: + self.traj_save_frequency = self.work_save_frequency * 5 - if work_freq is None: - work_freq = max(1, neq_steps // 50) - values["work_save_frequency"] = work_freq - if traj_freq is None: - values["traj_save_frequency"] = work_freq * 5 - traj_freq = values["traj_save_frequency"] - - if traj_freq % work_freq != 0: + if self.traj_save_frequency % self.work_save_frequency != 0: raise ValueError( "traj_save_frequency must be a multiple of work_save_frequency. " "Please specify consistent values." ) - return values + return self - @root_validator - def snapshots_require_no_internal_equilibration(cls, values): + @model_validator(mode="after") + def snapshots_require_no_internal_equilibration(self): """ When snapshot settings are provided equilibrium_steps must be 0 — the snapshots are already equilibrated. """ - integrator_settings = values.get("integrator_settings") - if integrator_settings is None: - return values + if self.integrator_settings is None: + return self - has_snapshots = values.get("lambda0_snapshots") or values.get( - "lambda1_snapshots" - ) - if has_snapshots and integrator_settings.equilibrium_steps != 0: + has_snapshots = self.lambda0_snapshots or self.lambda1_snapshots + if has_snapshots and self.integrator_settings.equilibrium_steps != 0: raise ValueError( "When lambda0_snapshots or lambda1_snapshots are provided the " "snapshots are assumed to be pre-equilibrated. " "Set integrator_settings.equilibrium_steps = 0 to avoid " "running redundant equilibration on top of them." ) - return values + return self - @root_validator - def snapshot_trajectories_have_enough_frames(cls, values): + @model_validator(mode="after") + def snapshot_trajectories_have_enough_frames(self): """ When both lambda0_snapshots and lambda1_snapshots are provided, check that each trajectory contains enough frames to cover all replicates (i.e. at least num_switches * stride frames), and that both trajectories expose the same number of usable snapshots. """ - snap0: Optional[SnapshotSettings] = values.get("lambda0_snapshots") - snap1: Optional[SnapshotSettings] = values.get("lambda1_snapshots") + snap0 = self.lambda0_snapshots + snap1 = self.lambda1_snapshots if snap0 is None or snap1 is None: - return values + return self - num_switches: int = values.get("num_switches", 100) + num_switches = self.num_switches try: import MDAnalysis as mda @@ -230,4 +222,4 @@ def snapshot_trajectories_have_enough_frames(cls, values): f"usable snapshots ({usable0} vs {usable1}). " "Provide trajectories of equal length or adjust the stride values." ) - return values + return self From d2019a92c69ab09ddf14a0adedc519100a141ad5 Mon Sep 17 00:00:00 2001 From: ijpulidos <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 23 Jun 2026 20:28:38 -0400 Subject: [PATCH 17/21] integrator settings with barostat fields to meet utils expectations --- feflow/settings/integrators.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/feflow/settings/integrators.py b/feflow/settings/integrators.py index 54436fc4..65400097 100644 --- a/feflow/settings/integrators.py +++ b/feflow/settings/integrators.py @@ -69,6 +69,12 @@ class AlchemicalNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): """Number of steps for the non-equilibrium switching (lambda 0->1 or 1->0). Default 2500 (10 ps at 4 fs).""" equilibrium_steps: int = 1000 """Number of equilibration steps at the endpoint before each switch. Default 1000.""" + barostat_frequency: TimestepQuantity = 25 * unit.timestep + """ + Frequency at which volume scaling changes should be attempted. + Note: The barostat frequency is ignored for gas-phase simulations. + Default 25 * unit.timestep. + """ @field_validator("nonequilibrium_steps", "equilibrium_steps") @classmethod From a950e243a1cbca37d630e1a92f46eb36dfbdf390 Mon Sep 17 00:00:00 2001 From: ijpulidos <2949729+ijpulidos@users.noreply.github.com> Date: Tue, 23 Jun 2026 20:28:53 -0400 Subject: [PATCH 18/21] minimization on each switch unit. Avoids NaNs --- feflow/protocols/nonequilibrium_switching.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py index c6f7cfff..581d6a37 100644 --- a/feflow/protocols/nonequilibrium_switching.py +++ b/feflow/protocols/nonequilibrium_switching.py @@ -353,6 +353,17 @@ def _execute(self, ctx, *, protocol, setup, equilibration, index, **inputs): neq_ctx = openmm.Context(system, neq_integrator, platform) neq_ctx.setState(snap_state) + # Adding minimization in the base switching unit. Helped to avoid NaNs in cases. + t_min0 = time.perf_counter() + openmm.LocalEnergyMinimizer.minimize(neq_ctx) + timing_info["minimization_time_in_s"] = ( + time.perf_counter() - t_min0 + ) + file_logger.info( + f"{self.name}: minimized starting snapshot " + f"({timing_info['minimization_time_in_s']:.1f} s)" + ) + works = [neq_integrator.get_protocol_work(dimensionless=True)] initial_traj, final_traj = [], [] From 762d80b1c2f8fa5889c3c777e36f98ae3b4942ed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Wed, 8 Apr 2026 19:25:12 -0400 Subject: [PATCH 19/21] Remove unneeded imports --- feflow/protocols/nonequilibrium_switching.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py index 581d6a37..ad94051a 100644 --- a/feflow/protocols/nonequilibrium_switching.py +++ b/feflow/protocols/nonequilibrium_switching.py @@ -60,12 +60,11 @@ def _load_snapshot(snapshot_settings: SnapshotSettings, index: int): box_vectors_nm : np.ndarray, shape (3, 3) or None Triclinic box vectors in nanometres, or None for vacuum systems. """ - import numpy as np import MDAnalysis as mda u = mda.Universe(snapshot_settings.topology_file, snapshot_settings.trajectory_file) frame_idx = index * snapshot_settings.stride - u.trajectory[frame_idx] + u.trajectory[frame_idx] # We need to place ourselves in the specified frame positions_nm = u.atoms.positions * 0.1 # Angstrom -> nm @@ -269,7 +268,6 @@ class _BaseSwitchingUnit(ProtocolUnit): @staticmethod def extract_positions(context, initial_atom_indices, final_atom_indices): - import numpy as np positions = context.getState(getPositions=True).getPositions(asNumpy=True) return ( From 873b201c96643ed3e5fb751a5ac81dfb48b5e136 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:01:24 +0000 Subject: [PATCH 20/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- feflow/protocols/nonequilibrium_switching.py | 13 +++++++++---- feflow/settings/nonequilibrium_switching.py | 1 - 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py index ad94051a..8aac9a5c 100644 --- a/feflow/protocols/nonequilibrium_switching.py +++ b/feflow/protocols/nonequilibrium_switching.py @@ -354,9 +354,7 @@ def _execute(self, ctx, *, protocol, setup, equilibration, index, **inputs): # Adding minimization in the base switching unit. Helped to avoid NaNs in cases. t_min0 = time.perf_counter() openmm.LocalEnergyMinimizer.minimize(neq_ctx) - timing_info["minimization_time_in_s"] = ( - time.perf_counter() - t_min0 - ) + timing_info["minimization_time_in_s"] = time.perf_counter() - t_min0 file_logger.info( f"{self.name}: minimized starting snapshot " f"({timing_info['minimization_time_in_s']:.1f} s)" @@ -616,7 +614,14 @@ def _create( reverse_switches=reverse_switches, ) - return [setup, lambda0_eq, lambda1_eq, *forward_switches, *reverse_switches, end] + return [ + setup, + lambda0_eq, + lambda1_eq, + *forward_switches, + *reverse_switches, + end, + ] @staticmethod def _check_mappings_consistency(mapping, chemical_system_a, chemical_system_b): diff --git a/feflow/settings/nonequilibrium_switching.py b/feflow/settings/nonequilibrium_switching.py index 022be1a1..542404da 100644 --- a/feflow/settings/nonequilibrium_switching.py +++ b/feflow/settings/nonequilibrium_switching.py @@ -18,7 +18,6 @@ from openfe.protocols.openmm_rfe.equil_rfe_settings import AlchemicalSettings from pydantic import ConfigDict, model_validator - # Default settings for the lambda functions x = "lambda" DEFAULT_ALCHEMICAL_FUNCTIONS = { From 4a43312bffe3ca79766b6f76d7d295d5104ab83b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Pulido?= <2949729+ijpulidos@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:59:51 -0400 Subject: [PATCH 21/21] Adding barostat field to integrators --- feflow/settings/integrators.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/feflow/settings/integrators.py b/feflow/settings/integrators.py index 5698504f..eb0f17c9 100644 --- a/feflow/settings/integrators.py +++ b/feflow/settings/integrators.py @@ -75,6 +75,13 @@ class AlchemicalNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): Note: The barostat frequency is ignored for gas-phase simulations. Default 25 * unit.timestep. """ + barostat: Literal["MonteCarloBarostat"] = "MonteCarloBarostat" + """ + The barostat to be used in the simulations. Default MonteCarloBarostat. + Notes + ----- + If the system contains a membrane, use the `MonteCarloMembraneBarostat`, if supported. + """ @field_validator("nonequilibrium_steps", "equilibrium_steps") @classmethod @@ -102,6 +109,13 @@ class PeriodicNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): """Number of steps for the equilibrium parts of the cycle. Default 12500""" nonequilibrium_steps: int = 12500 """Number of steps for the non-equilibrium parts of the cycle. Default 12500""" + barostat: Literal["MonteCarloBarostat"] = "MonteCarloBarostat" + """ + The barostat to be used in the simulations. Default MonteCarloBarostat. + Notes + ----- + If the system contains a membrane, use the `MonteCarloMembraneBarostat`, if supported. + """ barostat_frequency: TimestepQuantity = 25 * unit.timestep """ Frequency at which volume scaling changes should be attempted.