From 65ccd41761cfa5720444d99edd243fb56de21193 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 1 Jul 2026 12:02:10 +0100 Subject: [PATCH 1/4] Add support for generating Boresch restraints for ABFE. --- CHANGELOG.md | 1 + src/somd2/config/_config.py | 110 ++++++++++++++++++++++++++++++- src/somd2/runner/_base.py | 127 +++++++++++++++++++++++++++++++++++- src/somd2/runner/_repex.py | 14 ++++ src/somd2/runner/_runner.py | 12 ++++ 5 files changed, 261 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ddf958..0190663 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ Changelog -------------------------------------------------------------------------------------- * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR. +* Add support for generating Boresch restraints for absolute binding free energy calculations [#166](https://github.com/OpenBioSim/somd2/pull/166). [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026 -------------------------------------------------------------------------------------- diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index b350295..9efd79c 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -172,6 +172,9 @@ def __init__( save_xml=False, page_size=None, timeout="300 s", + restraint_search_time="1 ns", + restraint_search_frequency="10 ps", + restraint_search_receptor_selection=None, ): """ Constructor. @@ -557,6 +560,21 @@ def __init__( null_energy: str The energy value to use for lambda windows that are not being computed as part of the energy trajectory. + + restraint_search_time: str + Length of the short pre-production trajectory used to auto-generate + a Boresch restraint when running an ABFE simulation without a + user-supplied restraint. Defaults to "1 ns". + + restraint_search_frequency: str + Frame-saving frequency during the restraint-search trajectory. + Defaults to "10 ps". Should be small enough to yield at least 50 + frames over ``restraint_search_time``. + + restraint_search_receptor_selection: str + Sire selection string for receptor anchor atom candidates used + during automatic Boresch restraint generation. If None, the default + backbone selection is used (CA, C, N atoms in non-water molecules). """ # Setup logger before doing anything else @@ -645,9 +663,10 @@ def __init__( self.num_energy_neighbours = num_energy_neighbours self.null_energy = null_energy self.page_size = page_size - + self.restraint_search_time = restraint_search_time + self.restraint_search_frequency = restraint_search_frequency + self.restraint_search_receptor_selection = restraint_search_receptor_selection self.write_config = write_config - self.overwrite = overwrite def __str__(self): @@ -2468,6 +2487,34 @@ def _from_hex(hex): return obj + def __getstate__(self): + """ + Hex-encode the same fields that to_yaml()/from_yaml() already + hex-encode (currently 'restraints' and 'lambda_schedule'), since + these legacy Sire objects are not guaranteed to have native pickle + support. This is needed so that a Config holding these can be sent + to a spawned worker process, e.g. via + concurrent.futures.ProcessPoolExecutor. + """ + state = self.__dict__.copy() + if state.get("_restraints") is not None: + state["_restraints"] = [ + self._to_hex(restraint) for restraint in state["_restraints"] + ] + if state.get("_lambda_schedule") is not None: + state["_lambda_schedule"] = self._to_hex(state["_lambda_schedule"]) + return state + + def __setstate__(self, state): + """Reverse the hex-encoding performed in __getstate__.""" + if state.get("_restraints") is not None: + state["_restraints"] = [ + self._from_hex(restraint) for restraint in state["_restraints"] + ] + if state.get("_lambda_schedule") is not None: + state["_lambda_schedule"] = self._from_hex(state["_lambda_schedule"]) + self.__dict__.update(state) + @classmethod def _create_parser(cls): """ @@ -2572,6 +2619,65 @@ def _create_parser(cls): return parser + @property + def restraint_search_time(self): + return self._restraint_search_time + + @restraint_search_time.setter + def restraint_search_time(self, restraint_search_time): + if not isinstance(restraint_search_time, str): + raise TypeError("'restraint_search_time' must be of type 'str'") + + from sire.units import picosecond + + try: + t = _sr.u(restraint_search_time) + except: + raise ValueError( + f"Unable to parse 'restraint_search_time' as a Sire GeneralUnit: {restraint_search_time}" + ) + + if not t.has_same_units(picosecond): + raise ValueError("'restraint_search_time' units are invalid.") + + self._restraint_search_time = t + + @property + def restraint_search_frequency(self): + return self._restraint_search_frequency + + @restraint_search_frequency.setter + def restraint_search_frequency(self, restraint_search_frequency): + if not isinstance(restraint_search_frequency, str): + raise TypeError("'restraint_search_frequency' must be of type 'str'") + + from sire.units import picosecond + + try: + t = _sr.u(restraint_search_frequency) + except: + raise ValueError( + f"Unable to parse 'restraint_search_frequency' as a Sire GeneralUnit: {restraint_search_frequency}" + ) + + if not t.has_same_units(picosecond): + raise ValueError("'restraint_search_frequency' units are invalid.") + + self._restraint_search_frequency = t + + @property + def restraint_search_receptor_selection(self): + return self._restraint_search_receptor_selection + + @restraint_search_receptor_selection.setter + def restraint_search_receptor_selection(self, restraint_search_receptor_selection): + if restraint_search_receptor_selection is not None: + if not isinstance(restraint_search_receptor_selection, str): + raise TypeError( + "'restraint_search_receptor_selection' must be of type 'str'" + ) + self._restraint_search_receptor_selection = restraint_search_receptor_selection + def _reset_logger(self, logger): """ Internal method to reset the logger. diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index c973400..d5101f2 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -223,6 +223,15 @@ def __init__(self, system, config): except: self._has_water = False + # Check for protein (distinguishes ABFE from AHFE). A protein+ligand + # system has more than one non-water molecule with at least 3 atoms; + # a solvated ligand alone has only one. + try: + non_water_mols = self._system["(not water) and (atomidx > 1)"].molecules() + self._has_protein = non_water_mols.num_molecules() > 1 + except: + self._has_protein = False + # Warn if dispersion correction is requested but can't be applied. if self._config.use_dispersion_correction and not self._has_water: msg = "Cannot use dispersion correction for vacuum simulations. Disabling!" @@ -274,7 +283,7 @@ def __init__(self, system, config): elif self._config.ghost_modifications: from ghostly import modify - _logger.info("Applying modifications to ghost atom bonded terms") + _logger.info("Applying modifications to ghost atom bonded terms.") try: self._system, self._modifications = modify(self._system) # Angle optimisation can sometimes fail. @@ -954,6 +963,120 @@ def __init__(self, system, config): # Update the maximum number of threads. _sr.legacy.Base.set_max_num_threads(sire_threads) + @property + def _is_abfe_bound(self): + """ + Whether this is the bound leg of an ABFE simulation: annihilate/decouple + schedule with a solvated protein present. False for the free leg (ligand + in solvent, no protein), even though it may use the same lambda schedule. + """ + return ( + self._config._lambda_schedule_name in ("annihilate", "decouple") + and self._has_protein + and self._has_water + ) + + def _generate_boresch_restraint(self, device=None): + """ + Return a Boresch restraint for the ABFE simulation, either by loading + one saved from a previous run or by running a short lambda=0 trajectory. + Called automatically before minimisation/equilibration of the production + windows/replicas when the simulation is ABFE and no restraint has been + supplied. + + The input system is assumed to already be equilibrated (SOMD2 does not + run a separate equilibration stage for the restraint search); it is + minimised using the production minimisation settings, then a short + trajectory is run at lambda=0, matching the dynamics settings used for + production, to derive the restraint geometry and force constants. + + Parameters + ---------- + + device : int, optional + GPU device number to use for the restraint-search run. + + Returns + ------- + + restraints : sire.mm.BoreschRestraints + """ + from sire.restraints import boresch_search + + restraint_file = str(self._config.output_directory / "abfe_restraint.s3") + + # On restart, load the restraint saved from the previous run. + if _Path(restraint_file).exists(): + _logger.info(f"Loading existing Boresch restraint from {restraint_file}") + return _sr.stream.load(restraint_file) + + _logger.info( + "No restraint supplied for ABFE. Running Boresch restraint search." + ) + + search_system = self._system + + if self._config.minimise: + constraint = self._config.constraint + perturbable_constraint = self._config.perturbable_constraint + + # Don't use constraints during minimisation. + if not self._config.minimisation_constraints: + constraint = "none" + perturbable_constraint = "none" + + min_dynamics_kwargs = self._dynamics_kwargs.copy() + min_dynamics_kwargs.update( + { + "device": device, + "lambda_value": 0.0, + "constraint": constraint, + "perturbable_constraint": perturbable_constraint, + } + ) + + min_dynamics = search_system.dynamics(**min_dynamics_kwargs) + min_dynamics.minimise(timeout=self._config.timeout) + search_system = min_dynamics.commit() + + dynamics_kwargs = self._dynamics_kwargs.copy() + dynamics_kwargs.update( + { + "device": device, + "lambda_value": 0.0, + } + ) + + dynamics = search_system.dynamics(**dynamics_kwargs) + dynamics.run( + self._config.restraint_search_time, + energy_frequency=0, + frame_frequency=self._config.restraint_search_frequency, + save_velocities=False, + auto_fix_minimise=self._config.auto_fix_minimise, + save_crash_report=self._config.save_crash_report, + ) + search_system = dynamics.commit() + + search_kwargs = {"temperature": self._config.temperature} + if self._config.restraint_search_receptor_selection is not None: + search_kwargs["receptor_selection"] = ( + self._config.restraint_search_receptor_selection + ) + + restraints, correction = boresch_search(search_system, **search_kwargs) + + correction_kcal_mol = float(correction.to(_sr.units.kcal_per_mol)) + _logger.info( + f"Boresch restraint generated. Standard state correction: " + f"{correction_kcal_mol:.4f} kcal mol-1" + ) + + # Save so that restarts can reload without re-generating. + _sr.stream.save(restraints, restraint_file) + + return restraints + def _check_space(self): """ Check if the system has a periodic space. @@ -1496,6 +1619,8 @@ def _compare_configs(config1, config2): "log_file", "overwrite", "timeout", + "restraint_search_time", + "restraint_search_frequency", ] for key in config1.keys(): if key not in allowed_diffs: diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 7da3ba3..56a1b28 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -740,6 +740,20 @@ def __init__(self, system, config): else: self._num_gpus = min(self._config.max_gpus, len(gpu_devices)) + # Auto-generate a Boresch restraint for ABFE runs with no user-supplied + # restraint. This must happen before the dynamics cache is built below, + # since the per-replica OpenMM contexts it creates are fixed at + # construction time and won't pick up a restraint added afterwards. + if self._is_abfe_bound and self._config.restraints is None: + try: + restraints = self._generate_boresch_restraint(device=0) + except Exception as e: + msg = f"Unable to generate Boresch restraint for ABFE simulation: {e}" + _logger.error(msg) + raise RuntimeError(msg) + self._config.restraints = restraints + self._dynamics_kwargs["restraints"] = restraints + # Store the name of the dynamics cache pickle file. self._repex_state = self._config.output_directory / "repex_state.pkl" diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index fb3980f..56dc31e 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -217,6 +217,18 @@ def run(self): else: self._max_workers = 1 + # Auto-generate a Boresch restraint for ABFE runs with no user-supplied restraint. + if self._is_abfe_bound and self._config.restraints is None: + device = self._gpu_pool[0] if self._is_gpu else None + try: + restraints = self._generate_boresch_restraint(device=device) + except Exception as e: + msg = f"Unable to generate Boresch restraint for ABFE simulation: {e}" + _logger.error(msg) + raise RuntimeError(msg) + self._config.restraints = restraints + self._dynamics_kwargs["restraints"] = restraints + import concurrent.futures as _futures import multiprocessing as _mp From 2c572de07595fe4b7a62ff91e84cfdaf69955d6d Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 2 Jul 2026 13:17:19 +0100 Subject: [PATCH 2/4] Add support for split restraint levers in ABFE schedules. --- src/somd2/_utils/_schedules.py | 97 +++++++++++++++++++-- tests/schedules/test_abfe.py | 153 +++++++++++++++++++++++++++++++++ 2 files changed, 244 insertions(+), 6 deletions(-) create mode 100644 tests/schedules/test_abfe.py diff --git a/src/somd2/_utils/_schedules.py b/src/somd2/_utils/_schedules.py index 4c49a0e..f40dd52 100644 --- a/src/somd2/_utils/_schedules.py +++ b/src/somd2/_utils/_schedules.py @@ -27,7 +27,44 @@ ] -def annihilate(fix_epsilon=True): +def _set_boresch_lever_equations(s, stage_dihedral, stage_distance_angle): + """ + Set the equations for a "split" restraint_lever Boresch restraint (see + sire.restraints.boresch's restraint_lever parameter), reproducing the + RXRX protocol's staged restraint turn-on (Table S1 of the RXRX paper's + SI): within each of the two named stages, the corresponding restraint + group ramps from ~0 to 1 following a geometric progression, while the + other group is held fixed. 'stage_dihedral' is the stage over which + the dihedral restraint group ramps on (with the distance/angle group + held at 0); 'stage_distance_angle' is the stage over which the + distance/angle group ramps on (with the dihedral group held at 1, + already fully on). + + Note: this aligns the restraint turn-on with SOMD2's own decharge/ + annihilate(or decouple) stage boundaries, rather than reproducing the + RXRX paper's exact global 50/50 window split (which falls partway + through the annihilate/decouple stage, since the paper's own decharge + stage is only 21 of 64 total bound-leg windows) - the relative sizes of + SOMD2's stages are controlled by lambda_values weighting, not fixed. + """ + from sire.legacy.CAS import Exp as _Exp + import math as _math + + # Geometric progression from ~0.01 (fully off) to 1.0 (fully on), matching + # the ratio observed in the RXRX paper's published lambda schedule. + ramp_on = _Exp((1 - s.lam()) * _math.log(0.01)) + + s.set_equation(stage=stage_dihedral, lever="restraint_dihedral", equation=ramp_on) + s.set_equation(stage=stage_dihedral, lever="restraint_distance_angle", equation=0) + s.set_equation(stage=stage_distance_angle, lever="restraint_dihedral", equation=1) + s.set_equation( + stage=stage_distance_angle, + lever="restraint_distance_angle", + equation=ramp_on, + ) + + +def annihilate(fix_epsilon=True, restraint_lever="split"): """ Build the ABFE lambda schedule using decharge → annihilate. @@ -44,12 +81,28 @@ def annihilate(fix_epsilon=True): If False, epsilon is scaled normally from initial to final and the LRC follows naturally. + restraint_lever : str, optional + How the Boresch restraint is controlled by this schedule, matching + sire.restraints.boresch's restraint_lever parameter. Either "split" + (default), where the dihedral restraint terms are turned on during + decharge and the distance/angle terms are turned on during + annihilate, reproducing the RXRX protocol's staged restraint + turn-on, or "combined", where the whole restraint is turned on + together during the decharge stage. The Boresch restraint object + passed to the simulation must have a matching restraint_lever value. + Returns ------- schedule : sire.legacy.CAS.LambdaSchedule The lambda schedule. """ + if restraint_lever not in ("combined", "split"): + raise ValueError( + "'restraint_lever' must be either 'combined' or 'split', " + f"got {restraint_lever!r}" + ) + from sire.cas import LambdaSchedule as _LambdaSchedule # Start with the standard decouple schedule and modify the stages and @@ -65,14 +118,22 @@ def annihilate(fix_epsilon=True): lever="charge", equation=s.lam() * s.final() + s.initial() * (1 - s.lam()), ) - s.set_equation(stage="decharge", force="restraint", equation=s.lam() * s.final()) s.add_stage( "annihilate", equation=(-s.lam() + 1) * s.initial() + s.lam() * s.final(), ) s.set_equation(stage="annihilate", lever="charge", equation=s.final()) - s.set_equation(stage="annihilate", force="restraint", equation=s.final()) + + if restraint_lever == "split": + _set_boresch_lever_equations( + s, stage_dihedral="decharge", stage_distance_angle="annihilate" + ) + else: + s.set_equation( + stage="decharge", lever="restraint", equation=s.lam() * s.final() + ) + s.set_equation(stage="annihilate", lever="restraint", equation=s.final()) if fix_epsilon: s.set_equation(stage="annihilate", lever="epsilon", equation=s.initial()) @@ -86,7 +147,7 @@ def annihilate(fix_epsilon=True): return s -def decouple(fix_epsilon=True): +def decouple(fix_epsilon=True, restraint_lever="split"): """ Build the ABFE lambda schedule using decharge → decouple. @@ -101,12 +162,28 @@ def decouple(fix_epsilon=True): ghost-LRC force is then explicitly scaled to zero over the stage. If False, epsilon is scaled normally and the LRC follows naturally. + restraint_lever : str, optional + How the Boresch restraint is controlled by this schedule, matching + sire.restraints.boresch's restraint_lever parameter. Either "split" + (default), where the dihedral restraint terms are turned on during + decharge and the distance/angle terms are turned on during decouple, + reproducing the RXRX protocol's staged restraint turn-on, or + "combined", where the whole restraint is turned on together during + the decharge stage. The Boresch restraint object passed to the + simulation must have a matching restraint_lever value. + Returns ------- schedule : sire.legacy.CAS.LambdaSchedule The lambda schedule. """ + if restraint_lever not in ("combined", "split"): + raise ValueError( + "'restraint_lever' must be either 'combined' or 'split', " + f"got {restraint_lever!r}" + ) + from sire.cas import LambdaSchedule as _LambdaSchedule # Start with the standard decouple schedule and modify the stages and @@ -114,7 +191,6 @@ def decouple(fix_epsilon=True): # we will use this approach for prototyping. s = _LambdaSchedule.standard_decouple() - s.set_equation(stage="decouple", lever="restraint", equation=s.final()) s.set_equation(stage="decouple", lever="kappa", force="ghost/ghost", equation=0) s.set_equation(stage="decouple", lever="kappa", force="ghost-14", equation=0) s.set_equation(stage="decouple", lever="charge", equation=s.final()) @@ -142,7 +218,16 @@ def decouple(fix_epsilon=True): s.set_equation( stage="decharge", lever="kappa", force="ghost-14", equation=-s.lam() + 1 ) - s.set_equation(stage="decharge", lever="restraint", equation=s.initial() * s.lam()) + + if restraint_lever == "split": + _set_boresch_lever_equations( + s, stage_dihedral="decharge", stage_distance_angle="decouple" + ) + else: + s.set_equation(stage="decouple", lever="restraint", equation=s.final()) + s.set_equation( + stage="decharge", lever="restraint", equation=s.initial() * s.lam() + ) return s diff --git a/tests/schedules/test_abfe.py b/tests/schedules/test_abfe.py new file mode 100644 index 0000000..e44a25a --- /dev/null +++ b/tests/schedules/test_abfe.py @@ -0,0 +1,153 @@ +import pytest + +from somd2._utils._schedules import annihilate, decouple + +# Lambda schedules are always symmetric between the two stages (decharge is +# the first stage, annihilate/decouple is the second), so both builders share +# identical expected lever values. +BUILDERS = [annihilate, decouple] + +_LAMBDA_VALUES = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] + +# Expected "restraint" lever values (restraint_lever="combined"): ramps +# linearly 0 -> 1 across the whole first stage (decharge), then held at 1. +_COMBINED_RESTRAINT = [0.0, 0.4, 0.8, 1.0, 1.0, 1.0] + +# Expected "restraint_dihedral"/"restraint_distance_angle" lever values +# (restraint_lever="split"): a geometric progression from ~0.01 to 1.0 across +# each stage in turn, reproducing the ratio in Table S1 of the RXRX paper's +# SI. dihedral ramps during the first stage then holds at 1; distance_angle +# holds at 0 during the first stage then ramps during the second. +_SPLIT_DIHEDRAL = [0.01, 0.06309573444801934, 0.39810717055349737, 1.0, 1.0, 1.0] +_SPLIT_DISTANCE_ANGLE = [0.0, 0.0, 0.0, 0.025118864315095805, 0.15848931924611143, 1.0] + + +def _morph(schedule, lever, lambda_value): + """ + Query a lever's value exactly as SireOpenMM's LambdaLever does at + runtime: lambda_schedule.morph("*", restraint_name, 1.0, 1.0, lambda_value). + """ + return schedule.morph("*", lever, 1.0, 1.0, lambda_value) + + +@pytest.mark.parametrize("builder", BUILDERS) +def test_restraint_lever_defaults_to_split(builder): + """ + annihilate()/decouple() default to restraint_lever="split", matching + boresch_search()'s own default of restraint_lever="split" for its + default protocol="rxrx" + """ + schedule = builder() + assert "restraint_dihedral" in schedule.get_levers() + assert "restraint_distance_angle" in schedule.get_levers() + assert "restraint" not in schedule.get_levers() + + +@pytest.mark.parametrize("builder", BUILDERS) +def test_restraint_lever_invalid_raises(builder): + with pytest.raises(ValueError, match="restraint_lever"): + builder(restraint_lever="not_a_real_lever") + + +@pytest.mark.parametrize("builder", BUILDERS) +def test_restraint_lever_combined(builder): + """ + restraint_lever="combined" sets a single "restraint" lever that ramps + 0 -> 1 across the first stage (decharge), then holds at 1. + """ + schedule = builder(restraint_lever="combined") + assert "restraint" in schedule.get_levers() + assert "restraint_dihedral" not in schedule.get_levers() + assert "restraint_distance_angle" not in schedule.get_levers() + + values = [_morph(schedule, "restraint", lv) for lv in _LAMBDA_VALUES] + assert values == pytest.approx(_COMBINED_RESTRAINT, abs=1e-6) + + +@pytest.mark.parametrize("builder", BUILDERS) +def test_restraint_lever_split(builder): + """ + restraint_lever="split" sets two independent levers ("restraint_dihedral" + and "restraint_distance_angle"), each following a geometric progression + across its own stage while the other is held fixed, reproducing the RXRX + protocol's staged restraint turn-on. + """ + schedule = builder(restraint_lever="split") + + dihedral_values = [ + _morph(schedule, "restraint_dihedral", lv) for lv in _LAMBDA_VALUES + ] + distance_angle_values = [ + _morph(schedule, "restraint_distance_angle", lv) for lv in _LAMBDA_VALUES + ] + + assert dihedral_values == pytest.approx(_SPLIT_DIHEDRAL, abs=1e-6) + assert distance_angle_values == pytest.approx(_SPLIT_DISTANCE_ANGLE, abs=1e-6) + + +@pytest.mark.skipif( + "openmm" not in __import__("sire").convert.supported_formats(), + reason="openmm support is not available", +) +def test_restraint_lever_split_openmm_system(): + """ + End-to-end regression test: build a real dynamics object with a "split" + Boresch restraint and confirm the two independently-lambda-addressable + OpenMM Forces (distance/angle and dihedral) exist, with 'rho' values + that follow the expected geometric progression as lambda changes, + matching what test_restraint_lever_split checks at the schedule level. + """ + import xml.etree.ElementTree as ET + + import sire as sr + + mols = sr.load_test_files("boresch_restraints.prm7", "boresch_restraints.dcd") + mols.update(sr.morph.decouple(mols.molecule(1), as_new_molecule=False)) + + restraints = sr.restraints.boresch( + mols, + receptor=[692, 702, 704], + ligand=[1496, 1498, 1499], + kr="1 kcal mol-1 A-2", + ktheta=["80 kcal mol-1 rad-2"] * 2, + kphi=["80 kcal mol-1 rad-2"] * 3, + r0="4.56908 A", + theta0=["82.5581 degrees", "94.9595 degrees"], + phi0=["27.9429 degrees", "125.68 degrees", "-107.008 degrees"], + angle_potential="restricted_bending", + restraint_lever="split", + ) + + schedule = decouple(restraint_lever="split") + + d = mols.dynamics( + timestep="2fs", + temperature="298 K", + schedule=schedule, + lambda_value=0.0, + map={"restraints": restraints}, + ) + + expected_rho = { + 0.0: {"distance_angle": 0.0, "dihedral": 0.01}, + 0.5: {"distance_angle": 0.01, "dihedral": 1.0}, + 1.0: {"distance_angle": 1.0, "dihedral": 1.0}, + } + + for lam, expected in expected_rho.items(): + d.set_lambda(lam) + root = ET.fromstring(d.to_xml()) + + rhos = {} + for force in root.iter("Force"): + if force.get("name") == "BoreschRestraintForce": + kind = ( + "distance_angle" if "e_bond" in force.get("energy") else "dihedral" + ) + rhos[kind] = float(force.find("Bonds")[0].get("param1")) + + assert set(rhos.keys()) == {"distance_angle", "dihedral"} + assert rhos["distance_angle"] == pytest.approx( + expected["distance_angle"], abs=1e-6 + ) + assert rhos["dihedral"] == pytest.approx(expected["dihedral"], abs=1e-6) From a496234564444a25c4d2b3de5b6cd5380f1eda01 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 3 Jul 2026 09:04:22 +0100 Subject: [PATCH 3/4] Store auto-generated standard state correction in parquet metadata. --- src/somd2/runner/_base.py | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index d5101f2..fa3932f 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -232,6 +232,11 @@ def __init__(self, system, config): except: self._has_protein = False + # Set by _generate_boresch_restraint() when a Boresch restraint is + # auto-generated, then written into the energy trajectory parquet + # metadata (see _checkpoint). + self._standard_state_correction = None + # Warn if dispersion correction is requested but can't be applied. if self._config.use_dispersion_correction and not self._has_water: msg = "Cannot use dispersion correction for vacuum simulations. Disabling!" @@ -1008,7 +1013,18 @@ def _generate_boresch_restraint(self, device=None): # On restart, load the restraint saved from the previous run. if _Path(restraint_file).exists(): _logger.info(f"Loading existing Boresch restraint from {restraint_file}") - return _sr.stream.load(restraint_file) + restraints = _sr.stream.load(restraint_file) + + from sire.restraints import get_standard_state_correction + + correction = get_standard_state_correction( + restraints[0], temperature=self._config.temperature + ) + self._standard_state_correction = float( + correction.to(_sr.units.kcal_per_mol) + ) + + return restraints _logger.info( "No restraint supplied for ABFE. Running Boresch restraint search." @@ -1066,10 +1082,13 @@ def _generate_boresch_restraint(self, device=None): restraints, correction = boresch_search(search_system, **search_kwargs) - correction_kcal_mol = float(correction.to(_sr.units.kcal_per_mol)) + # Cache so it can be written into the energy trajectory parquet + # metadata (see _checkpoint), letting analysis code automatically + # apply the correction without needing to scan the logs. + self._standard_state_correction = float(correction.to(_sr.units.kcal_per_mol)) _logger.info( f"Boresch restraint generated. Standard state correction: " - f"{correction_kcal_mol:.4f} kcal mol-1" + f"{self._standard_state_correction:.4f} kcal mol-1" ) # Save so that restarts can reload without re-generating. @@ -2070,6 +2089,13 @@ def _checkpoint( if lambda_grad is not None: metadata["lambda_grad"] = [f"{v:.5f}" for v in lambda_grad] + # Add the standard state correction, if a Boresch restraint + # was auto-generated for this ABFE run. + if self._standard_state_correction is not None: + metadata["standard_state_correction"] = ( + f"{self._standard_state_correction:.6f}" + ) + if is_final_block: # Save the end-state GCMC topologies for trajectory analysis and visualisation. # This topology contains additional water molecules that are used for GCMC From d276232fbe9a26f8e805d1d3f2fa20d032fdd42d Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 3 Jul 2026 09:16:20 +0100 Subject: [PATCH 4/4] Harmonise dependency version formatting. --- src/somd2/__init__.py | 14 ++++++++++++-- src/somd2/runner/_base.py | 3 +-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/somd2/__init__.py b/src/somd2/__init__.py index b3070df..7f0a24a 100644 --- a/src/somd2/__init__.py +++ b/src/somd2/__init__.py @@ -34,10 +34,20 @@ # Store the somd2 version. from ._version import __version__ -# Store the sire version. +# Store the sire version. Unlike somd2/BioSimSpace/ghostly/loch (which use +# versioningit and only append a "+g" local version segment +# for non-release (".dev") builds, omitting it entirely for a clean tagged +# release), sire exposes its version and revision id as separate attributes, +# with __revisionid__ always set regardless of release status. Build a +# composite string using the same "+g" convention as the other +# packages, only appending it for non-release (".dev") builds, so that all +# five version strings are formatted consistently. from sire import __version__ as _sire_version from sire import __revisionid__ as _sire_revisionid +if ".dev" in _sire_version: + _sire_version = f"{_sire_version}+g{_sire_revisionid}" + # Store the BioSimSpace version. from BioSimSpace import __version__ as _biosimspace_version @@ -61,7 +71,7 @@ def get_versions(): """ return { "somd2": __version__, - "sire": f"{_sire_version}+{_sire_revisionid}", + "sire": _sire_version, "biosimspace": _biosimspace_version, "ghostly": _ghostly_version, "loch": _loch_version, diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index fa3932f..85f6e82 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -2078,8 +2078,7 @@ def _checkpoint( if not is_post_equilibration: metadata = { "attrs": df.attrs, - "somd2 version": versions["somd2"], - "sire version": versions["sire"], + "versions": versions, "lambda": f"{lam:.5f}", "speed": speed, "temperature": str(self._config.temperature.value()),