diff --git a/feflow/protocols/__init__.py b/feflow/protocols/__init__.py index b7f7807..86e91be 100644 --- a/feflow/protocols/__init__.py +++ b/feflow/protocols/__init__.py @@ -7,3 +7,9 @@ NonEquilibriumCyclingProtocol, NonEquilibriumCyclingProtocolResult, ) +from .nonequilibrium_switching import ( + NonEquilibriumSwitchingProtocol, + NonEquilibriumSwitchingProtocolResult, + ForwardSwitchingUnit, + ReverseSwitchingUnit, +) diff --git a/feflow/protocols/nonequilibrium_switching.py b/feflow/protocols/nonequilibrium_switching.py new file mode 100644 index 0000000..8aac9a5 --- /dev/null +++ b/feflow/protocols/nonequilibrium_switching.py @@ -0,0 +1,680 @@ +# 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, serialize +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() + } + + +def _load_snapshot(snapshot_settings: SnapshotSettings, index: int): + """ + 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 MDAnalysis as mda + + u = mda.Universe(snapshot_settings.topology_file, snapshot_settings.trajectory_file) + frame_idx = index * snapshot_settings.stride + u.trajectory[frame_idx] # We need to place ourselves in the specified frame + + 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 _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`` and ``_get_lambda_functions``. + """ + + @staticmethod + def extract_positions(context, initial_atom_indices, final_atom_indices): + + positions = context.getState(getPositions=True).getPositions(asNumpy=True) + return ( + positions[list(initial_atom_indices), :], + positions[list(final_atom_indices), :], + ) + + # --- to be overridden by subclasses --- + _direction: str = "" # "forward" or "reverse" + + def _get_lambda_functions(self, settings: NonEquilibriumSwitchingSettings) -> dict: + raise NotImplementedError + + # -------------------------------------------------------------- + + def _execute(self, ctx, *, protocol, setup, equilibration, index, **inputs): + """ + 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 atom indices. + equilibration : EquilibrationUnit result with the starting snapshot. + index : int + Replicate index; used for output file naming. + """ + import numpy as np + import openmm + 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_logger.addHandler(file_handler) + + 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 + 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_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 = {} + + # ------------------------------------------------------------------ + # 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) + + # 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 = [], [] + + 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 + ) + 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 + ) + 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" + + def _get_lambda_functions(self, settings): + return settings.lambda_functions + + +class ReverseSwitchingUnit(_BaseSwitchingUnit): + """Runs one reverse NEQ switch (lambda 1->0).""" + + _direction = "reverse" + + 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", + ) + + # 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, + equilibration=lambda0_eq, + index=i, + name=f"forward_{i}", + ) + for i in range(num_switches) + ] + reverse_switches = [ + ReverseSwitchingUnit( + protocol=self, + setup=setup, + equilibration=lambda1_eq, + index=i, + name=f"reverse_{i}", + ) + for i in range(num_switches) + ] + + end = ResultUnit( + name="result", + forward_switches=forward_switches, + reverse_switches=reverse_switches, + ) + + return [ + setup, + lambda0_eq, + lambda1_eq, + *forward_switches, + *reverse_switches, + end, + ] + + @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 + + 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 diff --git a/feflow/settings/__init__.py b/feflow/settings/__init__.py index 06fd6a4..d810455 100644 --- a/feflow/settings/__init__.py +++ b/feflow/settings/__init__.py @@ -1,3 +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, SnapshotSettings diff --git a/feflow/settings/integrators.py b/feflow/settings/integrators.py index 26a0dc4..eb0f17c 100644 --- a/feflow/settings/integrators.py +++ b/feflow/settings/integrators.py @@ -8,11 +8,10 @@ from typing import Annotated, TypeAlias, Literal -from pydantic import ConfigDict, field_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") @@ -20,34 +19,19 @@ TimestepQuantity: TypeAlias = Annotated[ GufeQuantity, specify_quantity_units("timestep") ] +InversePicosecondQuantity: TypeAlias = Annotated[ + GufeQuantity, specify_quantity_units("1/picoseconds") +] -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: Literal["MonteCarloBarostat"] = "MonteCarloBarostat" - """ - The barostat to be used in the simulations. Default MonteCarloBarostat. - Notes - ----- - If the system contains a membrane, use the `MonteCarloMembraneBarostat`. - """ - 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. @@ -71,6 +55,74 @@ def is_time(cls, v): raise ValueError("timestep must be in time units " "(i.e. picoseconds)") return v + +class AlchemicalNonequilibriumIntegratorSettings(BaseNonequilibriumIntegrator): + """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.""" + 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. + """ + 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 + 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): + """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: 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. + Note: The barostat frequency is ignored for gas-phase simulations. + Default 25 * unit.timestep. + """ + # TODO: This validator is used in other settings, better create a new Type @field_validator("equilibrium_steps", "nonequilibrium_steps") @classmethod diff --git a/feflow/settings/nonequilibrium_cycling.py b/feflow/settings/nonequilibrium_cycling.py index b41f556..2542a6a 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 x = "lambda" diff --git a/feflow/settings/nonequilibrium_switching.py b/feflow/settings/nonequilibrium_switching.py new file mode 100644 index 0000000..542404d --- /dev/null +++ b/feflow/settings/nonequilibrium_switching.py @@ -0,0 +1,224 @@ +""" +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 openfe.protocols.openmm_utils.omm_settings import ( + OpenMMSolvationSettings, + OpenMMEngineSettings, + ThermoSettings, +) +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 = { + "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. + """ + + # Debugging settings + store_minimized_pdb: bool = True + """Setting for storing pdb right after minimization (right before neq cycle)""" + + @model_validator(mode="after") + def set_and_validate_save_frequencies(self): + """ + Derive save-frequency defaults from nonequilibrium_steps when not set, + then check consistency. + """ + neq_steps = ( + self.integrator_settings.nonequilibrium_steps + if self.integrator_settings + else 2500 + ) + + 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 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 self + + @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. + """ + if self.integrator_settings is None: + return self + + 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 self + + @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 = self.lambda0_snapshots + snap1 = self.lambda1_snapshots + + if snap0 is None or snap1 is None: + return self + + num_switches = self.num_switches + + 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 self diff --git a/feflow/tests/conftest.py b/feflow/tests/conftest.py index 324e159..611eea2 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 0000000..21f2b04 --- /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