Add support for autogenerating ABFE restraints

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
 --------------------------------------------------------------------------------------

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<revisionid>" 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<revisionid>" 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,

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,20 +162,35 @@ 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
     # equations as needed. This will be folded into Sire in future, but
     # 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
 

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.