Allow a list of equilibration times, one for each lambda window

CHANGELOG.md

@@ -33,6 +33,7 @@ Changelog
 * Fixed the Boresch restraint for ABFE and the Morse restraints for ring-breaking perturbations not being generated when any other restraint was passed. They are now generated unless restraints of the same type are passed, and added to the others [#259](https://github.com/OpenBioSim/somd2/pull/259).
 * Fixed alchemical ions changing their charge linearly over the whole lambda range for the `annihilate`, `decouple`, `charge_scaled_morph`, `ring_break_morph` and `reverse_ring_break_morph` lambda schedules, rather than alongside the perturbable molecule. A warning is raised for custom lambda schedules, where alchemical ions still use the `standard_morph` schedule [#263](https://github.com/OpenBioSim/somd2/pull/263).
 * Add a `ghostly_kwargs` option to pass options to `ghostly` when applying ghost atom modifications, e.g. `--ghostly-kwargs k_soft=3 stiffen_rotamers=true` on the command line. Options are validated against those supported by `ghostly`. Systems that have already been modified with `ghostly` raise an error unless `ghost_modifications` is disabled, or `force=True` is passed in `ghostly_kwargs` [#264](https://github.com/OpenBioSim/somd2/pull/264).
+* Allow `equilibration_time` to be a list of times, one for each lambda window, e.g. `--equilibration-time 0ps 100ps 100ps 0ps` on the command line. Windows with a time of zero are not equilibrated. A single time is still used for every window [#265](https://github.com/OpenBioSim/somd2/pull/265).
 
 [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026
 --------------------------------------------------------------------------------------

README.md

@@ -287,6 +287,23 @@ The number of contexts that can be run in parallel is then controlled by the
 More details on MPS, including tuning options, can be found in the following
 [technical blog](https://developer.nvidia.com/blog/maximizing-openmm-molecular-dynamics-throughput-with-nvidia-multi-process-service/).
 
+<a id="equilibration"></a>
+
+## ⏳ Equilibration
+
+Each λ window can be equilibrated after minimisation, which is off by default.
+Use `--equilibration-time` to set how long for, e.g. `--equilibration-time
+100ps`. Equilibration uses a timestep of 2 fs by default, set with
+`--equilibration-timestep`. No energies or frames are saved during
+equilibration, and it doesn't count towards the `--runtime`. Only simulations
+starting from scratch are equilibrated, not restarts.
+
+Passing multiple values to `--equilibration-time` gives each λ window its own
+time, e.g. `--equilibration-time 0ps 100ps 100ps 0ps` for four windows, which
+can be used to equilibrate only the windows that need it. The number of values
+must match the number of λ windows, and windows with a time of zero are not
+equilibrated.
+
 <a id="restarting"></a>
 
 ## 💾 Restarting

src/somd2/config/_config.py

@@ -90,6 +90,7 @@ class Config:
         "pme_grid": "+",
         "rest2_scale": "+",
         "restraints": "+",
+        "equilibration_time": "+",
         "ghostly_kwargs": "+",
     }
 
@@ -406,9 +407,11 @@ def __init__(
         minimisation_errors: bool
             Whether to raise an exception if a minimisation fails to converge.
 
-        equilibration_time: str
+        equilibration_time: str, list(str)
             Time interval for equilibration. Only simulations starting from
-            scratch will be equilibrated.
+            scratch will be equilibrated. This can either be a single time,
+            used for all lambda windows, or a list of times, one for each
+            lambda window. A window with a time of zero is not equilibrated.
 
         equilibration_timestep: str
             Equilibration timestep. (Can be different to simulation timestep.)
@@ -992,6 +995,9 @@ def as_dict(self, sire_compatible=False):
             else:
                 d["lambda_schedule"] = self._to_hex(self.lambda_schedule)
 
+        if isinstance(self.equilibration_time, list):
+            d["equilibration_time"] = [t.to_string() for t in self.equilibration_time]
+
         # Serialise restraints.
         if self.restraints is not None:
             d["restraints"] = [self._to_hex(restraint) for restraint in self.restraints]
@@ -1842,22 +1848,37 @@ def equilibration_time(self):
 
     @equilibration_time.setter
     def equilibration_time(self, equilibration_time):
-        if not isinstance(equilibration_time, str):
-            raise TypeError("'equilibration_time' must be of type 'str'")
+        if isinstance(equilibration_time, str):
+            equilibration_time = [equilibration_time]
+        elif not isinstance(equilibration_time, (list, tuple)) or not all(
+            isinstance(x, str) for x in equilibration_time
+        ):
+            raise TypeError(
+                "'equilibration_time' must be of type 'str', or a list of 'str'"
+            )
+
+        if len(equilibration_time) == 0:
+            raise ValueError("'equilibration_time' must not be empty.")
 
         from sire.units import picosecond
 
-        try:
-            t = _sr.u(equilibration_time)
-        except:
-            raise ValueError(
-                f"Unable to parse 'equilibration_time' as a Sire GeneralUnit: {equilibration_time}"
-            )
+        times = []
+        for time in equilibration_time:
+            try:
+                t = _sr.u(time)
+            except:
+                raise ValueError(
+                    f"Unable to parse 'equilibration_time' as a Sire GeneralUnit: {time}"
+                )
 
-        if t.value() != 0 and not t.has_same_units(picosecond):
-            raise ValueError("'equilibration_time' units are invalid.")
+            if t.value() != 0 and not t.has_same_units(picosecond):
+                raise ValueError("'equilibration_time' units are invalid.")
+
+            times.append(t)
 
-        self._equilibration_time = t
+        if len(times) == 1:
+            times = times[0]
+        self._equilibration_time = times
 
     @property
     def equilibration_timestep(self):

src/somd2/runner/_base.py

@@ -685,6 +685,21 @@ def _set_schedule(self, pert_idxs, ion_indices):
             _logger.error(msg)
             raise ValueError(msg)
 
+        # Set the equilibration time for each window.
+        if isinstance(self._config.equilibration_time, list):
+            if len(self._config.equilibration_time) != len(self._lambda_values):
+                msg = (
+                    "Length of 'equilibration_time' must match the number of "
+                    f"{_lam_sym} values."
+                )
+                _logger.error(msg)
+                raise ValueError(msg)
+            self._equilibration_times = self._config.equilibration_time
+        else:
+            self._equilibration_times = [self._config.equilibration_time] * len(
+                self._lambda_values
+            )
+
         # Set the REST2 scale factors.
         is_rest2 = False
         if self._config.rest2_scale is not None:

src/somd2/runner/_repex.py

@@ -240,7 +240,9 @@ def __init__(self, system, config):
             constraint = "none"
             perturbable_constraint = "none"
 
-        if not self._is_restart and self._config.equilibration_time.value() > 0.0:
+        if not self._is_restart and any(
+            t.value() > 0.0 for t in self._equilibration_times
+        ):
             self._is_equilibration = True
 
             # Don't use constraints during equilibration.
@@ -1819,7 +1821,15 @@ def _equilibrate(self, index):
         exception: Exception
             The exception if the equilibration failed.
         """
-        _logger.info(f"Equilibrating at {_lam_sym} = {self._lambda_values[index]:.5f}")
+        time = self._equilibration_times[index]
+
+        # The replica is still loaded and stored when its time is zero, so
+        # that the slot's resident replica is the same as for other windows.
+        is_equilibrated = time.value() > 0.0
+        if is_equilibrated:
+            _logger.info(
+                f"Equilibrating at {_lam_sym} = {self._lambda_values[index]:.5f}"
+            )
 
         try:
             # Make the replica resident in its slot, with its own clock.
@@ -1835,7 +1845,7 @@ def _equilibrate(self, index):
             if clock is None:
                 clock = dynamics._get_clock()
 
-            if gcmc_sampler is not None:
+            if is_equilibrated and gcmc_sampler is not None:
                 gcmc_sampler.push()
                 try:
                     _logger.info(
@@ -1847,14 +1857,15 @@ def _equilibrate(self, index):
                     gcmc_sampler.pop()
 
             # Equilibrate.
-            dynamics.run(
-                self._config.equilibration_time,
-                energy_frequency=0,
-                frame_frequency=0,
-                save_velocities=False,
-                auto_fix_minimise=self._config.auto_fix_minimise,
-                save_crash_report=self._config.save_crash_report,
-            )
+            if is_equilibrated:
+                dynamics.run(
+                    time,
+                    energy_frequency=0,
+                    frame_frequency=0,
+                    save_velocities=False,
+                    auto_fix_minimise=self._config.auto_fix_minimise,
+                    save_crash_report=self._config.save_crash_report,
+                )
 
             dynamics._set_clock(clock)
 

src/somd2/runner/_runner.py

@@ -539,7 +539,7 @@ def _run(
             system = self._minimise_window(system, window)
 
         is_equilibrated = (
-            not is_restart and self._config.equilibration_time.value() > 0.0
+            not is_restart and self._equilibration_times[window.index].value() > 0.0
         )
         if is_equilibrated:
             system = self._equilibrate_window(system, window)
@@ -730,7 +730,7 @@ def _minimise_window(self, system, window):
             constraint = "none"
             perturbable_constraint = "none"
         # We will be performing an equilibration stage.
-        elif self._config.equilibration_time.value() > 0.0:
+        elif self._equilibration_times[window.index].value() > 0.0:
             # Don't use constraints during equilibration.
             if not self._config.equilibration_constraints:
                 constraint = "none"
@@ -838,7 +838,7 @@ def _equilibrate_window(self, system, window):
 
             # Run without saving energies or frames.
             dynamics.run(
-                self._config.equilibration_time,
+                self._equilibration_times[window.index],
                 energy_frequency=0,
                 frame_frequency=0,
                 save_velocities=False,

tests/runner/test_config.py

@@ -355,6 +355,46 @@ def test_ghostly_premodified():
         )
 
 
+def test_equilibration_time(ethane_methanol):
+    """Validate a single equilibration time, and one for each lambda window."""
+    import pytest
+
+    times = ["0ps", "1ps", "2ps"]
+
+    # A single time, as passed on the command line.
+    parser = Config._create_parser()
+    args = parser.parse_args(["--equilibration-time", "1ps"])
+    config = Config(equilibration_time=args.equilibration_time)
+    assert config.equilibration_time == sr.u("1ps")
+    assert Config(**config.as_dict()) == config
+
+    # A time for each window.
+    args = parser.parse_args(["--equilibration-time", *times])
+    config = Config(equilibration_time=args.equilibration_time)
+    assert config.equilibration_time == [sr.u(t) for t in times]
+    assert Config(**config.as_dict()) == config
+
+    with pytest.raises(ValueError, match="units are invalid"):
+        Config(equilibration_time=["1ps", "1A"])
+    with pytest.raises(TypeError):
+        Config(equilibration_time=1)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "platform": "cpu",
+            "output_directory": tmpdir,
+            "num_lambda": 3,
+            "equilibration_time": times,
+        }
+        runner = Runner(ethane_methanol, Config(**config))
+        assert runner._equilibration_times == [sr.u(t) for t in times]
+
+        config["num_lambda"] = 4
+        config["overwrite"] = True
+        with pytest.raises(ValueError, match="Length of 'equilibration_time'"):
+            Runner(ethane_methanol, Config(**config))
+
+
 def test_help_text_scraping():
     """Validate that help text isn't truncated by the parameter name."""
     parser = Config._create_parser()

tests/runner/test_repex.py

@@ -1035,7 +1035,8 @@ def test_repex_concurrent_slots(ethane_methanol, max_contexts):
 
     This is also the only test that equilibrates, so it covers moving replicas
     in and out of their slots during equilibration, the context rebuild when
-    the constraints change, and the post-equilibration checkpoint.
+    the constraints change, windows that aren't equilibrated, and the
+    post-equilibration checkpoint.
     """
     num_lambda = 4
 
@@ -1047,7 +1048,7 @@ def test_repex_concurrent_slots(ethane_methanol, max_contexts):
             "energy_frequency": "4fs",
             "checkpoint_frequency": "4fs",
             "frame_frequency": "4fs",
-            "equilibration_time": "4fs",
+            "equilibration_time": ["4fs", "0fs", "4fs", "0fs"],
             "platform": "cuda",
             "max_threads": 1,
             "num_lambda": num_lambda,

tests/runner/test_runner.py

@@ -41,6 +41,30 @@ def test_runtime_remainder(ethane_methanol):
             assert checkpoint_state["time_ps"].item() == pytest.approx(0.020)
 
 
+def test_equilibration_time_per_window(ethane_methanol):
+    """Validate that only windows with a non-zero time are equilibrated."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "4fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "equilibration_time": ["0fs", "4fs"],
+            "platform": "CPU",
+            "max_threads": 1,
+            "num_lambda": 2,
+        }
+
+        runner = Runner(ethane_methanol, Config(**config))
+        runner.run()
+
+        log = (Path(tmpdir) / runner._config.log_file).read_text()
+        assert "Equilibrating at λ = 0.00000" not in log
+        assert "Equilibrating at λ = 1.00000" in log
+
+
 def test_runtime_remainder_with_terminal_flips(phenethyl_mols):
     """
     Validate that blocks which aren't a whole number of terminal flip intervals