From a0b14b46d1238268626cf2a50fec09187c065964 Mon Sep 17 00:00:00 2001 From: Josh Horton Date: Fri, 28 Aug 2026 13:58:36 +0100 Subject: [PATCH 1/2] prototype trailblazing equilibration unit --- .../protocols/openmm_utils/trailblazing.py | 367 ++++++++++++++++++ 1 file changed, 367 insertions(+) create mode 100644 src/openfe/protocols/openmm_utils/trailblazing.py diff --git a/src/openfe/protocols/openmm_utils/trailblazing.py b/src/openfe/protocols/openmm_utils/trailblazing.py new file mode 100644 index 000000000..18ca3763e --- /dev/null +++ b/src/openfe/protocols/openmm_utils/trailblazing.py @@ -0,0 +1,367 @@ +""" +An implementation of a general trailblazing algorithm for alchemical OpenMM simulations based on the work of the Chodera lab in the Yank software. +The algorithm is designed to find an optimal lambda schedule for alchemical transformations, which can improve the efficiency and accuracy of free energy calculations. + +Basic outline of the trailblazing algorithm: +1. Start at lambda=0 end state, run equilibration and extract N samples. +2. For the next candidate lambda window reweight all samples to the new lambda window and calculate the std of the reduced potential energy differences. +3. If the std is below a threshold and within the chosen tolerance, accept the new lambda window and move to the next candidate window. Go back to step 2. +4. If the std is above the threshold, reject the new lambda window and try a different candidate window. Go back to step 2. +5. Repeat until the entire lambda schedule is constructed, return the final lambda schedule and the equilibrated input samples for each window for the production phase of the simulation. +""" +import pathlib +from dataclasses import dataclass + +from openff.units.units import Quantity +from openff.units.openmm import to_openmm +from openmm import openmm, app +from openmm import unit as ommunit +from openfe.protocols.openmm_md.plain_md_methods import PlainMDSimulationUnit +from openfe.protocols.openmm_utils import omm_compute, settings_validation +import numpy as np +import logging +import pandas as pd + + +logger = logging.getLogger(__name__) + + +@dataclass +class _StateSnapshot: + step: int + positions_nm: np.ndarray + box_vectors_nm: np.ndarray + + +class _InMemorySnapshotReporter: + def __init__(self, report_interval: int): + self._report_interval = report_interval + self.snapshots: list[_StateSnapshot] = [] + + def describeNextReport(self, simulation: app.Simulation) -> tuple[int, bool, bool, bool, bool, bool]: + steps = self._report_interval - simulation.currentStep % self._report_interval + return (steps, True, False, False, False, False) + + def report(self, simulation: app.Simulation, state: openmm.State) -> None: + self.snapshots.append( + _StateSnapshot( + step=simulation.currentStep, + positions_nm=state.getPositions(asNumpy=True).value_in_unit(ommunit.nanometer), + box_vectors_nm=state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(ommunit.nanometer), + ) + ) + + +class TrailblazingMixin: + """ + Mixin class for trailblazing algorithm for alchemical OpenMM simulations. + """ + + def __init__(self, settings: dict, shared_basepath: pathlib.Path, verbose: bool = False): + """ + Initialize the TrailblazingMixin class. + + Parameters + ---------- + inputs : dict + A dictionary of inputs for the trailblazing algorithm. + shared_basepath : pathlib.Path + The base path for storing simulation outputs. + verbose : bool, optional + If True, print detailed logs during execution. Default is False. + """ + self._settings = settings + self.shared_basepath = shared_basepath + self.verbose = verbose + + def _get_settings(self): + """ + Extract the relevant settings for a solvent transformation. + + Returns + ------- + settings : dict[str, SettingsBaseModel] + A dictionary with the following entries: + * forcefield_settings : OpenMMSystemGeneratorFFSettings + * thermo_settings : ThermoSettings + * charge_settings : OpenFFPartialChargeSettings + * solvation_settings : OpenMMSolvationSettings + * alchemical_settings : AlchemicalSettings + * lambda_settings : LambdaSettings + * engine_settings : OpenMMEngineSettings + * integrator_settings : IntegratorSettings + * equil_simulation_settings : MDSimulationSettings + * equil_output_settings : MDOutputSettings + * simulation_settings : MultiStateSimulationSettings + * output_settings: MultiStateOutputSettings + """ + prot_settings = self._settings # type: ignore[attr-defined] + + settings = {} + settings["forcefield_settings"] = prot_settings.solvent_forcefield_settings.unfrozen_copy() + settings["thermo_settings"] = prot_settings.thermo_settings.unfrozen_copy() + settings["charge_settings"] = prot_settings.partial_charge_settings.unfrozen_copy() + settings["solvation_settings"] = prot_settings.solvation_settings.unfrozen_copy() + settings["alchemical_settings"] = prot_settings.alchemical_settings.unfrozen_copy() + settings["lambda_settings"] = prot_settings.lambda_settings.unfrozen_copy() + settings["engine_settings"] = prot_settings.solvent_engine_settings.unfrozen_copy() + settings["integrator_settings"] = prot_settings.integrator_settings.unfrozen_copy() + settings["equil_simulation_settings"] = prot_settings.solvent_equil_simulation_settings.unfrozen_copy() + settings["equil_output_settings"] = prot_settings.solvent_equil_output_settings.unfrozen_copy() + settings["simulation_settings"] = prot_settings.solvent_simulation_settings.unfrozen_copy() + settings["output_settings"] = prot_settings.solvent_output_settings.unfrozen_copy() + + return settings + + + def _lambda_converter(self, lambda_value: float) -> dict[str, float]: + """ + Convert a global lambda value to a dictionary of alchemical parameters for the OpenMM simulation. + + Parameters + ---------- + lambda_value : float + The lambda value to convert. + + Returns + ------- + dict[str, float] + A dictionary of alchemical parameters. + """ + # make sure its bounded between 0 and 1 + lambda_value = np.clip(lambda_value, 0.0, 1.0) + # Yank uses the lambda bounds of 1.0 -> 0.0 for the end state, so we need to invert the lambda value + lambda_value = 1.0 - lambda_value + return { + # sterics are only scaled for lambda >= 0.5, electrostatics are only scaled for lambda < 0.5 + # method taken from Yank: + "lambda_electrostatics": 2 * (lambda_value - 0.5) * np.heaviside(lambda_value - 0.5, 0), + "lambda_sterics": np.heaviside(lambda_value -0.5, 0.5) + 2 * lambda_value * np.heaviside(0.5 - lambda_value, 0.5), + } + + def _get_reduced_potential(self, simulation: app.Simulation, settings) -> Quantity: + state = simulation.context.getState(getEnergy=True) + potential_energy = state.getPotentialEnergy() + pressure = to_openmm(settings["thermo_settings"].pressure) + volume = state.getPeriodicBoxVolume() + return potential_energy / self._kbT + (pressure * volume * ommunit.AVOGADRO_CONSTANT_NA) / self._kbT + + def _get_snapshot_energies(self, simulation: app.Simulation, snapshots: list[_StateSnapshot], settings) -> list[Quantity]: + energies = [] + for snapshot in snapshots: + simulation.context.setPositions(snapshot.positions_nm * ommunit.nanometer) + a_vec, b_vec, c_vec = snapshot.box_vectors_nm + simulation.context.setPeriodicBoxVectors( + openmm.Vec3(*a_vec) * ommunit.nanometer, + openmm.Vec3(*b_vec) * ommunit.nanometer, + openmm.Vec3(*c_vec) * ommunit.nanometer, + ) + energies.append(self._get_reduced_potential(simulation=simulation, settings=settings)) + + return energies + + def _sample_state(self, simulation: app.Simulation, lambda_value: float, positions: ommunit.Quantity, settings, equil_steps_nvt: int | None, equil_steps_npt: int, prod_steps_npt: int) -> tuple[list[Quantity], list[_StateSnapshot]]: + """ + Equilibrate the simulation and extract N in-memory state snapshots. + + Notes + ----- + - The alchemical lambda values should already be set before being passed to this function. + """ + # make a new folder in the shared basepath for this lambda value + lambda_folder = self.shared_basepath / f"lambda_{lambda_value:.3f}" + lambda_folder.mkdir(parents=True, exist_ok=True) + write_interval = settings_validation.divmod_time_and_check( + settings["equil_output_settings"].trajectory_write_interval, + settings["integrator_settings"].timestep, + "trajectory_write_interval", + "timestep", + ) + snapshot_reporter = _InMemorySnapshotReporter(report_interval=write_interval) + simulation.reporters.append(snapshot_reporter) + + PlainMDSimulationUnit._run_MD( + simulation=simulation, + positions=positions, + simulation_settings=settings["equil_simulation_settings"], + output_settings=settings["equil_output_settings"], + temperature=settings["thermo_settings"].temperature, + barostat_frequency=settings["integrator_settings"].barostat_frequency, + timestep=settings["integrator_settings"].timestep, + equil_steps_nvt=equil_steps_nvt, + equil_steps_npt=equil_steps_npt, + prod_steps=prod_steps_npt, + verbose=self.verbose, + shared_basepath=lambda_folder + ) + # clean up the simulation object so it can be reused + simulation.reporters = [] + prod_start_step = (equil_steps_nvt or 0) + equil_steps_npt + prod_snapshots = [snapshot for snapshot in snapshot_reporter.snapshots if snapshot.step > prod_start_step] + # reset the current step of the simulation to 0 so it can be reused for the next lambda value + simulation.currentStep = 0 + if not prod_snapshots: + errmsg = ( + "No production snapshots captured for trailblazing reweighting. " + "Ensure trajectory_write_interval is shorter than production_length." + ) + raise ValueError(errmsg) + + # calculate the reduced potential energies for the production snapshots using the log data + state_data = pd.read_csv(lambda_folder / settings["equil_output_settings"].log_output) + potential_energies = state_data["Potential Energy (kJ/mole)"].values * ommunit.kilojoule_per_mole + volumes = state_data["Box Volume (nm^3)"].values * ommunit.nanometer ** 3 + pressure = to_openmm(settings["thermo_settings"].pressure) + simulated_energies = potential_energies / self._kbT + (pressure * volumes * ommunit.AVOGADRO_CONSTANT_NA) / self._kbT + + # save the box vectors and the positions to a numpy file + box_vectors = np.array([snapshot.box_vectors_nm for snapshot in prod_snapshots]) + positions = np.array([snapshot.positions_nm for snapshot in prod_snapshots]) + np.savez_compressed(lambda_folder / "box_vectors", box_vectors) + np.savez_compressed(lambda_folder / "positions", positions) + + return simulated_energies, prod_snapshots + + def _reweight_samples(self, simulation: app.Simulation, snapshots: list[_StateSnapshot], reweight_lambda_value: float, settings) -> list[Quantity]: + """ + Calculate the reweighted reduced potential energies for a new lambda value based on the simulated samples. + """ + reweight_alchemical_parameters = self._lambda_converter(reweight_lambda_value) + for param_name, param_value in reweight_alchemical_parameters.items(): + simulation.context.setParameter(param_name, param_value) + return self._get_snapshot_energies( + simulation=simulation, + snapshots=snapshots, + settings=settings, + ) + + def _run_trailblazing_method(self, system: openmm.System, topology: app.Topology, positions: ommunit.Quantity, thermodynamic_distance: float = 1.0, distance_tolerance: float = 0.05) -> list[float]: + """ + Run the trailblazing algorithm to find an optimal lambda schedule for alchemical transformations. + """ + settings = self._get_settings() # type: ignore[attr-defined] + # save all particles if any structure output is produced and disable trajectory output + settings["equil_output_settings"].output_indices = "all" + settings["equil_output_settings"].production_trajectory_filename = None + + # extract the settings we need for the equilibration + if settings["equil_simulation_settings"].equilibration_length_nvt is not None: + equil_steps_nvt = settings_validation.get_simsteps( + sim_length=settings["equil_simulation_settings"].equilibration_length_nvt, + timestep=settings["integrator_settings"].timestep, + mc_steps=1, + ) + else: + equil_steps_nvt = None + + equil_steps_npt = settings_validation.get_simsteps( + sim_length=settings["equil_simulation_settings"].equilibration_length, + timestep=settings["integrator_settings"].timestep, + mc_steps=1, + ) + + prod_steps_npt = settings_validation.get_simsteps( + sim_length=settings["equil_simulation_settings"].production_length, + timestep=settings["integrator_settings"].timestep, + mc_steps=1, + ) + + # build the simulation from the alchemical system and the initial lambda value + restrict_cpu = settings["forcefield_settings"].nonbonded_method.lower() == "nocutoff" + platform = omm_compute.get_openmm_platform( + platform_name=settings["engine_settings"].compute_platform, + gpu_device_index=settings["engine_settings"].gpu_device_index, + restrict_cpu_count=restrict_cpu, + ) + integrator = openmm.LangevinMiddleIntegrator( + to_openmm(settings["thermo_settings"].temperature), + to_openmm(settings["integrator_settings"].langevin_collision_rate), + to_openmm(settings["integrator_settings"].timestep), + ) + simulation = app.Simulation( + topology=topology, + system=system, + integrator=integrator, + platform=platform, + ) + # calculate and store kbT for the simulation + self._kbT = (ommunit.MOLAR_GAS_CONSTANT_R * to_openmm(settings["thermo_settings"].temperature)) + + optimal_lambda = [0.0] + state_stds = [] + + # start the trailblazing algorithm by iteratively finding the next optimal lambda window + while optimal_lambda[-1] < 1.0: + # set the lambda value in the simulation context + current_lambda = optimal_lambda[-1] + alchemical_parameters = self._lambda_converter(current_lambda) + for param_name, param_value in alchemical_parameters.items(): + simulation.context.setParameter(param_name, param_value) + + # equilibrate and sample the simulation, samples are saved in the trajectory reporter + simulated_energies, simulated_snapshots = self._sample_state( + simulation=simulation, + lambda_value=current_lambda, + positions=positions, + settings=settings, + equil_steps_nvt=equil_steps_nvt, + equil_steps_npt=equil_steps_npt, + prod_steps_npt=prod_steps_npt + ) + + # generate the next lambda value candidate + std_energy = 0.0 + old_std_energy = 0.0 + old_lambda = current_lambda + print(simulated_energies) + + while abs(std_energy - thermodynamic_distance) > distance_tolerance and not (current_lambda == 1.0 and std_energy < thermodynamic_distance): + if np.isclose(std_energy, 0.0): + # This is the first iteration or the two states overlap significantly + current_lambda += 0.05 + else: + # Assume std_energy is linear to determine the next value to try + derivative_std_energy = ((std_energy - old_std_energy) / (current_lambda - old_lambda)) + old_lambda = current_lambda + current_lambda += (thermodynamic_distance - std_energy) / derivative_std_energy + + # clip the value if we go over 1.0 + current_lambda = min(current_lambda, 1.0) + + reweighted_energies = self._reweight_samples( + simulation=simulation, + snapshots=simulated_snapshots, + reweight_lambda_value=current_lambda, + settings=settings + ) + + # update the std_energies with the new value + old_std_energy = std_energy + energy_diffs = np.array(reweighted_energies) - np.array(simulated_energies) + std_energy = np.std(energy_diffs, ddof=1) + logger.info(f"Trailblazing: simulated_lambda={optimal_lambda[-1]}, current_lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}") + print(f"Trailblazing: simulated_lambda={optimal_lambda[-1]}, current_lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}") + + + # add the value to the optimal lambda schedule + optimal_lambda.append(current_lambda) + state_stds.append(std_energy) + logger.info(f"Trailblazing: accepted lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}") + print(f"Trailblazing: accepted lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}") + # todo + # if we request a bidirectional optimization at the end then we need to reweight to the previous lambda value + # if this is not the first value + + # todo + # if we request a bidirectional optimization this should be done here + + logger.info(f"Trailblazing: completed with optimal lambda schedule: {optimal_lambda}") + print(f"Trailblazing: completed with optimal lambda schedule: {optimal_lambda}") + return optimal_lambda + # # if there is a previous lambda value, calculate the reverse direction simulated reweighted to the previous lambda value + # if len(optimal_lambda) > 2: + # previous_lambda = optimal_lambda[-2] + # previous_alchemical_parameters = self._lambda_converter(previous_lambda) + # reverse_reweighted_energies = self._reweight_samples(samples, alchemical_parameters, previous_alchemical_parameters) + + From fe09da1c0c2925d8e8733ecad42331c879a929e3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 13:02:14 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../protocols/openmm_utils/trailblazing.py | 134 ++++++++++++------ 1 file changed, 94 insertions(+), 40 deletions(-) diff --git a/src/openfe/protocols/openmm_utils/trailblazing.py b/src/openfe/protocols/openmm_utils/trailblazing.py index 18ca3763e..3693f95d0 100644 --- a/src/openfe/protocols/openmm_utils/trailblazing.py +++ b/src/openfe/protocols/openmm_utils/trailblazing.py @@ -9,19 +9,20 @@ 4. If the std is above the threshold, reject the new lambda window and try a different candidate window. Go back to step 2. 5. Repeat until the entire lambda schedule is constructed, return the final lambda schedule and the equilibrated input samples for each window for the production phase of the simulation. """ + +import logging import pathlib from dataclasses import dataclass -from openff.units.units import Quantity +import numpy as np +import pandas as pd from openff.units.openmm import to_openmm -from openmm import openmm, app +from openff.units.units import Quantity +from openmm import app, openmm from openmm import unit as ommunit + from openfe.protocols.openmm_md.plain_md_methods import PlainMDSimulationUnit from openfe.protocols.openmm_utils import omm_compute, settings_validation -import numpy as np -import logging -import pandas as pd - logger = logging.getLogger(__name__) @@ -38,7 +39,9 @@ def __init__(self, report_interval: int): self._report_interval = report_interval self.snapshots: list[_StateSnapshot] = [] - def describeNextReport(self, simulation: app.Simulation) -> tuple[int, bool, bool, bool, bool, bool]: + def describeNextReport( + self, simulation: app.Simulation + ) -> tuple[int, bool, bool, bool, bool, bool]: steps = self._report_interval - simulation.currentStep % self._report_interval return (steps, True, False, False, False, False) @@ -47,7 +50,9 @@ def report(self, simulation: app.Simulation, state: openmm.State) -> None: _StateSnapshot( step=simulation.currentStep, positions_nm=state.getPositions(asNumpy=True).value_in_unit(ommunit.nanometer), - box_vectors_nm=state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(ommunit.nanometer), + box_vectors_nm=state.getPeriodicBoxVectors(asNumpy=True).value_in_unit( + ommunit.nanometer + ), ) ) @@ -106,14 +111,17 @@ def _get_settings(self): settings["lambda_settings"] = prot_settings.lambda_settings.unfrozen_copy() settings["engine_settings"] = prot_settings.solvent_engine_settings.unfrozen_copy() settings["integrator_settings"] = prot_settings.integrator_settings.unfrozen_copy() - settings["equil_simulation_settings"] = prot_settings.solvent_equil_simulation_settings.unfrozen_copy() - settings["equil_output_settings"] = prot_settings.solvent_equil_output_settings.unfrozen_copy() + settings["equil_simulation_settings"] = ( + prot_settings.solvent_equil_simulation_settings.unfrozen_copy() + ) + settings["equil_output_settings"] = ( + prot_settings.solvent_equil_output_settings.unfrozen_copy() + ) settings["simulation_settings"] = prot_settings.solvent_simulation_settings.unfrozen_copy() settings["output_settings"] = prot_settings.solvent_output_settings.unfrozen_copy() return settings - def _lambda_converter(self, lambda_value: float) -> dict[str, float]: """ Convert a global lambda value to a dictionary of alchemical parameters for the OpenMM simulation. @@ -136,7 +144,8 @@ def _lambda_converter(self, lambda_value: float) -> dict[str, float]: # sterics are only scaled for lambda >= 0.5, electrostatics are only scaled for lambda < 0.5 # method taken from Yank: "lambda_electrostatics": 2 * (lambda_value - 0.5) * np.heaviside(lambda_value - 0.5, 0), - "lambda_sterics": np.heaviside(lambda_value -0.5, 0.5) + 2 * lambda_value * np.heaviside(0.5 - lambda_value, 0.5), + "lambda_sterics": np.heaviside(lambda_value - 0.5, 0.5) + + 2 * lambda_value * np.heaviside(0.5 - lambda_value, 0.5), } def _get_reduced_potential(self, simulation: app.Simulation, settings) -> Quantity: @@ -144,9 +153,14 @@ def _get_reduced_potential(self, simulation: app.Simulation, settings) -> Quanti potential_energy = state.getPotentialEnergy() pressure = to_openmm(settings["thermo_settings"].pressure) volume = state.getPeriodicBoxVolume() - return potential_energy / self._kbT + (pressure * volume * ommunit.AVOGADRO_CONSTANT_NA) / self._kbT + return ( + potential_energy / self._kbT + + (pressure * volume * ommunit.AVOGADRO_CONSTANT_NA) / self._kbT + ) - def _get_snapshot_energies(self, simulation: app.Simulation, snapshots: list[_StateSnapshot], settings) -> list[Quantity]: + def _get_snapshot_energies( + self, simulation: app.Simulation, snapshots: list[_StateSnapshot], settings + ) -> list[Quantity]: energies = [] for snapshot in snapshots: simulation.context.setPositions(snapshot.positions_nm * ommunit.nanometer) @@ -160,7 +174,16 @@ def _get_snapshot_energies(self, simulation: app.Simulation, snapshots: list[_St return energies - def _sample_state(self, simulation: app.Simulation, lambda_value: float, positions: ommunit.Quantity, settings, equil_steps_nvt: int | None, equil_steps_npt: int, prod_steps_npt: int) -> tuple[list[Quantity], list[_StateSnapshot]]: + def _sample_state( + self, + simulation: app.Simulation, + lambda_value: float, + positions: ommunit.Quantity, + settings, + equil_steps_nvt: int | None, + equil_steps_npt: int, + prod_steps_npt: int, + ) -> tuple[list[Quantity], list[_StateSnapshot]]: """ Equilibrate the simulation and extract N in-memory state snapshots. @@ -192,12 +215,14 @@ def _sample_state(self, simulation: app.Simulation, lambda_value: float, positio equil_steps_npt=equil_steps_npt, prod_steps=prod_steps_npt, verbose=self.verbose, - shared_basepath=lambda_folder + shared_basepath=lambda_folder, ) # clean up the simulation object so it can be reused simulation.reporters = [] prod_start_step = (equil_steps_nvt or 0) + equil_steps_npt - prod_snapshots = [snapshot for snapshot in snapshot_reporter.snapshots if snapshot.step > prod_start_step] + prod_snapshots = [ + snapshot for snapshot in snapshot_reporter.snapshots if snapshot.step > prod_start_step + ] # reset the current step of the simulation to 0 so it can be reused for the next lambda value simulation.currentStep = 0 if not prod_snapshots: @@ -209,10 +234,15 @@ def _sample_state(self, simulation: app.Simulation, lambda_value: float, positio # calculate the reduced potential energies for the production snapshots using the log data state_data = pd.read_csv(lambda_folder / settings["equil_output_settings"].log_output) - potential_energies = state_data["Potential Energy (kJ/mole)"].values * ommunit.kilojoule_per_mole - volumes = state_data["Box Volume (nm^3)"].values * ommunit.nanometer ** 3 + potential_energies = ( + state_data["Potential Energy (kJ/mole)"].values * ommunit.kilojoule_per_mole + ) + volumes = state_data["Box Volume (nm^3)"].values * ommunit.nanometer**3 pressure = to_openmm(settings["thermo_settings"].pressure) - simulated_energies = potential_energies / self._kbT + (pressure * volumes * ommunit.AVOGADRO_CONSTANT_NA) / self._kbT + simulated_energies = ( + potential_energies / self._kbT + + (pressure * volumes * ommunit.AVOGADRO_CONSTANT_NA) / self._kbT + ) # save the box vectors and the positions to a numpy file box_vectors = np.array([snapshot.box_vectors_nm for snapshot in prod_snapshots]) @@ -222,7 +252,13 @@ def _sample_state(self, simulation: app.Simulation, lambda_value: float, positio return simulated_energies, prod_snapshots - def _reweight_samples(self, simulation: app.Simulation, snapshots: list[_StateSnapshot], reweight_lambda_value: float, settings) -> list[Quantity]: + def _reweight_samples( + self, + simulation: app.Simulation, + snapshots: list[_StateSnapshot], + reweight_lambda_value: float, + settings, + ) -> list[Quantity]: """ Calculate the reweighted reduced potential energies for a new lambda value based on the simulated samples. """ @@ -235,11 +271,18 @@ def _reweight_samples(self, simulation: app.Simulation, snapshots: list[_StateSn settings=settings, ) - def _run_trailblazing_method(self, system: openmm.System, topology: app.Topology, positions: ommunit.Quantity, thermodynamic_distance: float = 1.0, distance_tolerance: float = 0.05) -> list[float]: + def _run_trailblazing_method( + self, + system: openmm.System, + topology: app.Topology, + positions: ommunit.Quantity, + thermodynamic_distance: float = 1.0, + distance_tolerance: float = 0.05, + ) -> list[float]: """ Run the trailblazing algorithm to find an optimal lambda schedule for alchemical transformations. """ - settings = self._get_settings() # type: ignore[attr-defined] + settings = self._get_settings() # type: ignore[attr-defined] # save all particles if any structure output is produced and disable trajectory output settings["equil_output_settings"].output_indices = "all" settings["equil_output_settings"].production_trajectory_filename = None @@ -285,7 +328,9 @@ def _run_trailblazing_method(self, system: openmm.System, topology: app.Topology platform=platform, ) # calculate and store kbT for the simulation - self._kbT = (ommunit.MOLAR_GAS_CONSTANT_R * to_openmm(settings["thermo_settings"].temperature)) + self._kbT = ommunit.MOLAR_GAS_CONSTANT_R * to_openmm( + settings["thermo_settings"].temperature + ) optimal_lambda = [0.0] state_stds = [] @@ -306,7 +351,7 @@ def _run_trailblazing_method(self, system: openmm.System, topology: app.Topology settings=settings, equil_steps_nvt=equil_steps_nvt, equil_steps_npt=equil_steps_npt, - prod_steps_npt=prod_steps_npt + prod_steps_npt=prod_steps_npt, ) # generate the next lambda value candidate @@ -315,13 +360,17 @@ def _run_trailblazing_method(self, system: openmm.System, topology: app.Topology old_lambda = current_lambda print(simulated_energies) - while abs(std_energy - thermodynamic_distance) > distance_tolerance and not (current_lambda == 1.0 and std_energy < thermodynamic_distance): + while abs(std_energy - thermodynamic_distance) > distance_tolerance and not ( + current_lambda == 1.0 and std_energy < thermodynamic_distance + ): if np.isclose(std_energy, 0.0): # This is the first iteration or the two states overlap significantly current_lambda += 0.05 else: # Assume std_energy is linear to determine the next value to try - derivative_std_energy = ((std_energy - old_std_energy) / (current_lambda - old_lambda)) + derivative_std_energy = (std_energy - old_std_energy) / ( + current_lambda - old_lambda + ) old_lambda = current_lambda current_lambda += (thermodynamic_distance - std_energy) / derivative_std_energy @@ -332,22 +381,29 @@ def _run_trailblazing_method(self, system: openmm.System, topology: app.Topology simulation=simulation, snapshots=simulated_snapshots, reweight_lambda_value=current_lambda, - settings=settings + settings=settings, ) # update the std_energies with the new value old_std_energy = std_energy energy_diffs = np.array(reweighted_energies) - np.array(simulated_energies) std_energy = np.std(energy_diffs, ddof=1) - logger.info(f"Trailblazing: simulated_lambda={optimal_lambda[-1]}, current_lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}") - print(f"Trailblazing: simulated_lambda={optimal_lambda[-1]}, current_lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}") - + logger.info( + f"Trailblazing: simulated_lambda={optimal_lambda[-1]}, current_lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}" + ) + print( + f"Trailblazing: simulated_lambda={optimal_lambda[-1]}, current_lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}" + ) # add the value to the optimal lambda schedule optimal_lambda.append(current_lambda) state_stds.append(std_energy) - logger.info(f"Trailblazing: accepted lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}") - print(f"Trailblazing: accepted lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}") + logger.info( + f"Trailblazing: accepted lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}" + ) + print( + f"Trailblazing: accepted lambda={current_lambda:.3f}, std_energy={std_energy:.3f}, target={thermodynamic_distance:.3f}, tolerance={distance_tolerance:.3f}" + ) # todo # if we request a bidirectional optimization at the end then we need to reweight to the previous lambda value # if this is not the first value @@ -358,10 +414,8 @@ def _run_trailblazing_method(self, system: openmm.System, topology: app.Topology logger.info(f"Trailblazing: completed with optimal lambda schedule: {optimal_lambda}") print(f"Trailblazing: completed with optimal lambda schedule: {optimal_lambda}") return optimal_lambda - # # if there is a previous lambda value, calculate the reverse direction simulated reweighted to the previous lambda value - # if len(optimal_lambda) > 2: - # previous_lambda = optimal_lambda[-2] - # previous_alchemical_parameters = self._lambda_converter(previous_lambda) - # reverse_reweighted_energies = self._reweight_samples(samples, alchemical_parameters, previous_alchemical_parameters) - - + # # if there is a previous lambda value, calculate the reverse direction simulated reweighted to the previous lambda value + # if len(optimal_lambda) > 2: + # previous_lambda = optimal_lambda[-2] + # previous_alchemical_parameters = self._lambda_converter(previous_lambda) + # reverse_reweighted_energies = self._reweight_samples(samples, alchemical_parameters, previous_alchemical_parameters)