diff --git a/CHANGELOG.md b/CHANGELOG.md index a8b37a7..c262c30 100644 --- a/CHANGELOG.md +++ b/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 -------------------------------------------------------------------------------------- diff --git a/README.md b/README.md index ed1868b..9c98f7d 100644 --- a/README.md +++ b/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/). + + +## ⏳ 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. + ## 💾 Restarting diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 00d9bb0..679e5fd 100644 --- a/src/somd2/config/_config.py +++ b/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): diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 5c20905..22c7d91 100644 --- a/src/somd2/runner/_base.py +++ b/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: diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 103510d..c9b855c 100644 --- a/src/somd2/runner/_repex.py +++ b/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) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 0bfd0cb..aa2d350 100644 --- a/src/somd2/runner/_runner.py +++ b/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, diff --git a/tests/runner/test_config.py b/tests/runner/test_config.py index ea7e983..2113e7e 100644 --- a/tests/runner/test_config.py +++ b/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() diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 4ee74a7..0fc7227 100644 --- a/tests/runner/test_repex.py +++ b/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, diff --git a/tests/runner/test_runner.py b/tests/runner/test_runner.py index 4ef47f8..bfac06b 100644 --- a/tests/runner/test_runner.py +++ b/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