From cba79b9d757d3f2da7a78a85e44ecf3215edc28e Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Oct 2026 14:38:35 +0100 Subject: [PATCH 01/32] Add adaptive replica exchange. --- CHANGELOG.md | 1 + README.md | 55 +++ src/somd2/_viewer/_data.py | 21 +- src/somd2/config/_config.py | 123 +++++ src/somd2/runner/_base.py | 15 +- src/somd2/runner/_repex.py | 856 +++++++++++++++++++++++++++++----- tests/runner/test_repex.py | 472 ++++++++++++++++++- tests/viewer/test_progress.py | 105 +++++ 8 files changed, 1521 insertions(+), 127 deletions(-) create mode 100644 tests/viewer/test_progress.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 65acf54..1483e00 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ Changelog * Only fall back to the Aldeghi Boresch restraint search protocol when the system can't be used with the RXRX protocol. Systems that can't be used with either protocol, or restraint search settings that would save too few frames, now raise an error before the restraint search trajectory is run, rather than after a misleading fallback warning [#241](https://github.com/OpenBioSim/somd2/pull/241). * Save the GCMC end-state topologies, which include the additional waters, at the first checkpoint of each run as well as the last, and keep them on restart, so that they match the trajectory and checkpoints while a simulation is running, if it crashes, or if a restart never reaches a checkpoint [#244](https://github.com/OpenBioSim/somd2/pull/244). * Add `gcmc_batch_size` and `gcmc_num_attempts` options to set the number of GCMC trials evaluated in parallel per batch and in total per move, which were previously fixed at the defaults of the GCMC sampler. Both can be changed on restart, e.g. to reduce wasted trials when more than one is accepted per move [#246](https://github.com/OpenBioSim/somd2/pull/246). +* Add adaptive replica exchange, enabled with `adaptive_repex`, which concentrates sampling on the lambda windows that are slowest to converge. Windows are ranked at checkpoints by the uncertainty they contribute to the free energy, and only the highest ranked are propagated and exchanged every cycle, with the rest joining every `adaptive_mix_interval` cycles [#247](https://github.com/OpenBioSim/somd2/pull/247). [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 b9d0b62..b51c5e3 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,8 @@ simulations. Built on top of [Sire](https://github.com/OpenBioSim/sire) and [charge-change](#charge-change-perturbations), and protein mutations. - 🔁 **[Replica exchange](#replica-exchange)**: Hamiltonian replica exchange between λ windows. +- 🎯 **[Adaptive replica exchange](#adaptive-replica-exchange)**: sampling + concentrated on the λ windows that are slowest to converge. - 🌡️ **[REST2](#rest2)**: replica exchange with solute scaling. - 💧 **[GCMC](#gcmc)**: grand canonical Monte Carlo water sampling. - 🔄 **[Terminal ring flips](#terminal-ring-flip-monte-carlo)**: Monte Carlo @@ -414,6 +416,59 @@ on either side of it, and the null energy with `--null-energy`. The number of neighbours is a trade-off between accuracy and computational cost, and around 20% of the number of replicas has been found to be a good starting point. + + +### 🎯 Adaptive replica exchange + +For difficult perturbations, most of the uncertainty in the free energy often +comes from a few λ windows, while the rest converge quickly and then contribute +little as the simulation is extended. Adaptive replica exchange, enabled with +`--adaptive-repex`, concentrates the sampling on the slow windows. After a +burn-in period of standard replica exchange, set with `--adaptive-burn-in`, the +windows are ranked at each checkpoint. Only the highest ranked ("hot") windows +are propagated and exchanged every cycle, with every window joining in every +`--adaptive-mix-interval` cycles, as well as on checkpoint and frame cycles. + +Each pair of neighbouring windows is ranked by σ√g of the reduced energy +differences between them, where σ is the standard deviation and g the +statistical inefficiency, which sets the uncertainty of its free energy +difference. The highest ranked pairs that together account for +`--adaptive-coverage` of the total are made hot, along with: + +- pairs whose forward and reverse free energy estimates disagree, or whose + estimate drifts during the simulation, +- pairs that rarely exchange, +- windows with fewer than `--adaptive-min-samples` statistically independent + samples since the burn-in, +- for ring-breaking perturbations that break or make a single bond, windows in + which the bond distance has fewer than `--adaptive-min-samples` statistically + independent samples, since a slow distance can bias the free energy without + showing in the energy differences. + +The windows either side of each hot window are also made hot, so that exchange +carries configurations into and out of them, and a window only stops being hot +after it has been found cold at two consecutive checkpoints. The number of hot +windows is rounded up to a multiple of the number of GPUs multiplied by the +oversubscription factor, so that the GPUs are evenly loaded. When contexts are +re-used, replicas are assigned to them in a strided rather than contiguous +pattern, so that the hot windows are spread evenly over the GPUs. + +The `--runtime` is the simulation time of the hot windows, so the remaining +windows are sampled for less time and the windows have different numbers of +energy samples. For the same `--runtime`, adaptive replica exchange gives +roughly the same uncertainty as standard replica exchange for less compute. +To reduce the uncertainty instead, increase the `--runtime` until the run takes +as long as standard replica exchange would, so that the slow windows are +sampled for longer. It doesn't help a window that is stuck to escape. + +The windows are ranked at checkpoints, so the `--checkpoint-frequency` should +be several times the `--adaptive-mix-interval` multiplied by the +`--energy-frequency`. The default burn-in of 500 ps also allows for +equilibration, which is off by default. If the system has been equilibrated +with `--equilibration-time`, a burn-in of around 200 ps is enough. The saving +is largest when few windows are slow. If most of the windows are hot, adaptive +replica exchange costs almost as much as standard replica exchange. + ## 🌡️ REST2 diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index 6056fa9..073531a 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -517,12 +517,25 @@ def _progress(self, config, lambda_values, windows, log): total = runtime * len(lambda_values) if runtime else None completed = sum(done.get(lam, 0.0) for lam in lambda_values) + # Adaptive replica exchange samples cold windows less, so progress is + # the simulation time of the hot windows rather than the samples. + is_adaptive = bool( + config.get("replica_exchange") and config.get("adaptive_repex") + ) + if is_adaptive: + state = self._repex_state() + time = state.get("time_ns") if state else None + if time is None: + time = max(done.values(), default=0.0) + time = min(time, runtime) if runtime else time + completed = time * len(lambda_values) + now = _time.time() progress = { "runtime_ns": runtime, "expected_samples": ( round(runtime / energy_frequency) - if runtime and energy_frequency + if runtime and energy_frequency and not is_adaptive else None ), "completed_ns": completed, @@ -566,7 +579,9 @@ def _progress(self, config, lambda_values, windows, log): progress["eta_seconds"] = max(remaining, 0.0) if energy_frequency is not None and interval > 0: progress["throughput_ns_day"] = energy_frequency / interval * 86400 - progress["throughput_label"] = "per replica" + progress["throughput_label"] = ( + "per hot replica" if is_adaptive else "per replica" + ) elif log.blocks and log.production_start is not None: elapsed = log.blocks[-1][0] - log.production_start if elapsed > 0 and checkpoint_frequency is not None: @@ -878,7 +893,9 @@ def load(): try: with open(p, "rb") as f: state = _RepexUnpickler(f).load() + time = getattr(state, "_time", None) return { + "time_ns": time.to("ns") if time is not None else None, "proposed": _np.asarray(state._num_proposed), "accepted": _np.asarray(state._num_accepted), "gcmc_stats": getattr(state, "_gcmc_stats", None), diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index da9b69d..412ccc3 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -117,6 +117,11 @@ class Config: "timeout", "oversubscription_factor", "max_contexts", + "adaptive_repex", + "adaptive_mix_interval", + "adaptive_burn_in", + "adaptive_coverage", + "adaptive_min_samples", "restraint_search_time", "restraint_search_frequency", "gcmc_batch_size", @@ -187,6 +192,11 @@ def __init__( max_contexts=None, update_constraints=True, constraint_lambda_index=0, + adaptive_repex=False, + adaptive_mix_interval=4, + adaptive_burn_in="500 ps", + adaptive_coverage=0.8, + adaptive_min_samples=50, randomise_velocities=False, perturbed_system=None, terminal_flip_frequency=None, @@ -494,6 +504,36 @@ def __init__( number of replicas, 'update_constraints' is False, and a constrained bond length actually perturbs with lambda. + adaptive_repex: bool + Whether to concentrate sampling on the lambda windows that are slowest to + converge in a replica exchange simulation. After a burn-in period of + standard replica exchange, the windows are periodically ranked by the + statistical uncertainty that they contribute to the free energy. Only the + highest ranked ("hot") windows are propagated and exchanged every cycle, + with the remaining windows joining in every 'adaptive_mix_interval' + cycles. The 'runtime' is the simulation time of the hot windows, so the + other windows are sampled for less time. Checkpoint and frame cycles always + include every window. + + adaptive_mix_interval: int + The number of cycles between those that propagate and exchange every + window when using adaptive replica exchange. + + adaptive_burn_in: str + The simulation time of standard replica exchange to run before windows are + first ranked when using adaptive replica exchange. + + adaptive_coverage: float + The fraction of the total uncertainty, summed over neighbouring pairs of + lambda windows, that the hot windows must account for when using adaptive + replica exchange. + + adaptive_min_samples: int + The minimum number of statistically independent energy samples that a + window must have collected since the burn-in before it can stop being hot + when using adaptive replica exchange. Windows are not ranked until each has + collected this many samples. + randomise_velocities: bool Whether to randomise velocities at the start of each replica exchange cycle or following a terminal flip Monte Carlo move. @@ -785,6 +825,11 @@ def __init__( self.max_contexts = max_contexts self.update_constraints = update_constraints self.constraint_lambda_index = constraint_lambda_index + self.adaptive_repex = adaptive_repex + self.adaptive_mix_interval = adaptive_mix_interval + self.adaptive_burn_in = adaptive_burn_in + self.adaptive_coverage = adaptive_coverage + self.adaptive_min_samples = adaptive_min_samples self.randomise_velocities = randomise_velocities self.perturbed_system = perturbed_system self.terminal_flip_frequency = terminal_flip_frequency @@ -2136,6 +2181,84 @@ def constraint_lambda_index(self, constraint_lambda_index): ) self._constraint_lambda_index = constraint_lambda_index + @property + def adaptive_repex(self): + return self._adaptive_repex + + @adaptive_repex.setter + def adaptive_repex(self, adaptive_repex): + if not isinstance(adaptive_repex, bool): + raise ValueError("'adaptive_repex' must be of type 'bool'") + self._adaptive_repex = adaptive_repex + + @property + def adaptive_mix_interval(self): + return self._adaptive_mix_interval + + @adaptive_mix_interval.setter + def adaptive_mix_interval(self, adaptive_mix_interval): + if not isinstance(adaptive_mix_interval, int): + try: + adaptive_mix_interval = int(adaptive_mix_interval) + except Exception: + raise ValueError("'adaptive_mix_interval' must be of type 'int'") + if adaptive_mix_interval < 1: + raise ValueError("'adaptive_mix_interval' must be greater than 0") + self._adaptive_mix_interval = adaptive_mix_interval + + @property + def adaptive_burn_in(self): + return self._adaptive_burn_in + + @adaptive_burn_in.setter + def adaptive_burn_in(self, adaptive_burn_in): + if not isinstance(adaptive_burn_in, str): + raise TypeError("'adaptive_burn_in' must be of type 'str'") + + from sire.units import picosecond + + try: + t = _sr.u(adaptive_burn_in) + except: + raise ValueError( + f"Unable to parse 'adaptive_burn_in' as a Sire GeneralUnit: {adaptive_burn_in}" + ) + + if t.value() != 0 and not t.has_same_units(picosecond): + raise ValueError("'adaptive_burn_in' units are invalid.") + + self._adaptive_burn_in = t + + @property + def adaptive_coverage(self): + return self._adaptive_coverage + + @adaptive_coverage.setter + def adaptive_coverage(self, adaptive_coverage): + if not isinstance(adaptive_coverage, float): + try: + adaptive_coverage = float(adaptive_coverage) + except Exception: + raise ValueError("'adaptive_coverage' must be of type 'float'") + if not 0.0 < adaptive_coverage <= 1.0: + raise ValueError("'adaptive_coverage' must be in the range (0, 1]") + self._adaptive_coverage = adaptive_coverage + + @property + def adaptive_min_samples(self): + return self._adaptive_min_samples + + @adaptive_min_samples.setter + def adaptive_min_samples(self, adaptive_min_samples): + if not isinstance(adaptive_min_samples, int): + try: + adaptive_min_samples = int(adaptive_min_samples) + except Exception: + raise ValueError("'adaptive_min_samples' must be of type 'int'") + if adaptive_min_samples < 1: + raise ValueError("'adaptive_min_samples' must be greater than 0") + self._adaptive_min_samples = adaptive_min_samples + @property def randomise_velocities(self): return self._randomise_velocities diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 74d2460..f24c130 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -2489,6 +2489,7 @@ def _checkpoint( is_final_block=False, context=None, gcmc_sampler=None, + num_energy_rows=None, ): """ Save a checkpoint file. @@ -2517,6 +2518,10 @@ def _checkpoint( is_final_block: bool Whether this is the final block of the simulation. + num_energy_rows: int + The number of energy samples collected since the last checkpoint. If + None, every window is assumed to have collected a full block. + Returns ------- @@ -2538,6 +2543,9 @@ def _checkpoint( # Get the lambda value. lam = self._lambda_values[index] + if num_energy_rows is None: + num_energy_rows = self._energy_per_block + # -1 is the sentinel for a post-equilibration checkpoint. No # energies are collected during equilibration, so skip all # parquet-related work in this case. @@ -2634,7 +2642,8 @@ def _checkpoint( # yet (e.g. checkpoint_frequency=0), create one from scratch. _energy_traj = self._filenames[index]["energy_traj"] if _Path(_energy_traj).exists(): - _parquet_append(_energy_traj, df.iloc[-self._energy_per_block :]) + if num_energy_rows > 0: + _parquet_append(_energy_traj, df.iloc[-num_energy_rows:]) else: _dataframe_to_parquet(df, metadata=metadata, filename=_energy_traj) @@ -2673,10 +2682,10 @@ def _checkpoint( self._is_restart and _Path(filename).exists() ): _dataframe_to_parquet(df, metadata=metadata, filename=filename) - else: + elif num_energy_rows > 0: _parquet_append( filename, - df.iloc[-self._energy_per_block :], + df.iloc[-num_energy_rows:], ) if not is_post_equilibration: diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 4604947..83b78e3 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -38,6 +38,82 @@ from ._base import RunnerBase as _RunnerBase +@_njit +def _mix_subset(replicas, energy_matrix, proposed, accepted): + """ + Mix a subset of the replicas, leaving the others in place. + + Swaps are only proposed between the replicas in the subset, which is a + valid Gibbs move for any subset, since the others are held fixed. + + Parameters + ---------- + + replicas: np.ndarray + The indices of the replicas to mix. + + energy_matrix: np.ndarray + The reduced energy of each replica at each state. Only the rows of the + replicas being mixed are used. + + proposed: np.ndarray + The number of proposed swaps between states, updated in place. + + accepted: np.ndarray + The number of accepted swaps between states, updated in place. + + Returns + ------- + + states: np.ndarray + The new states for every replica, where states[i] is the replica whose + configuration seeds replica i. + """ + + # Adapted from OpenMMTools: https://github.com/choderalab/openmmtools + + num_mixed = len(replicas) + + # Set the states to the initial order. + states = _np.arange(energy_matrix.shape[0]) + + # Attempt swaps. + for swap in range(num_mixed**3): + # Choose two replicas to swap. + replica_i = replicas[_np.random.randint(num_mixed)] + replica_j = replicas[_np.random.randint(num_mixed)] + + # Get the current state. + state_i = states[replica_i] + state_j = states[replica_j] + + # Record that we have proposed a swap. + proposed[state_i, state_j] += 1 + proposed[state_j, state_i] += 1 + + # Get the energies. + energy_ii = energy_matrix[replica_i, state_i] + energy_jj = energy_matrix[replica_j, state_j] + energy_ij = energy_matrix[replica_i, state_j] + energy_ji = energy_matrix[replica_j, state_i] + + # Compute the log probability of the swap. + log_p_swap = -(energy_ij + energy_ji) + energy_ii + energy_jj + + # Accept the swap and update the states. + if log_p_swap >= 0 or _np.random.rand() < _np.exp(log_p_swap): + # Swap the states. + states[replica_i] = state_j + states[replica_j] = state_i + # Record the swap. + accepted[state_i, state_j] += 1 + accepted[state_j, state_i] += 1 + + # Here states[i] is the state that replica i's configuration moves to, + # whereas the configurations are moved into fixed states, so invert. + return _np.argsort(states) + + class DynamicsCache: """ A class for caching dynamics objects. @@ -58,6 +134,7 @@ def __init__( update_constraints=True, constraint_lambda_index=None, gpu_devices=None, + strided=False, ): """ Constructor. @@ -117,6 +194,11 @@ def __init__( The physical devices backing each OpenMM device index, i.e. the entries of CUDA_VISIBLE_DEVICES. Used to query the memory of the right device. If None, the OpenMM index is used directly. + + strided: bool + Whether to assign replicas to slots in a strided rather than a + contiguous pattern, so that any contiguous range of replicas is + spread evenly over the slots. """ num_replicas = len(lambdas) @@ -132,8 +214,15 @@ def __init__( self._update_constraints = update_constraints self._constraint_lambda_index = constraint_lambda_index self._gpu_devices = gpu_devices + self._strided = strided self._states = _np.array(range(num_replicas)) self._time = None + # The simulation clock of each replica. + self._clocks = [None] * num_replicas + # The hot replicas of adaptive replica exchange, or None for all. + self._hot_replicas = None + self._cold_votes = [0] * num_replicas + self._adaptive_data = None self._openmm_states = [None] * num_replicas # The molecules with virtual sites, so that each molecule's particles # can be found in the saved states. Set by the runner. @@ -179,18 +268,28 @@ def _build_slot_layout(self): force field parameters, and hence the chance of a constraint update forcing the context to be reinitialised, as small as possible. + When strided, replica r is hosted by slot r % num_slots instead, so + that the hot replicas of adaptive replica exchange, which form + contiguous ranges, are spread evenly over the slots and GPUs. + The layout is derived from the number of replicas and slots, so it is rebuilt rather than stored in a checkpoint. """ self._gcmc_samplers = [None] * self._num_slots self._slot_replica = [None] * self._num_slots - self._groups = [ - [int(r) for r in group] - for group in _np.array_split( - _np.arange(self._num_replicas), self._num_slots - ) - ] + if self._strided: + self._groups = [ + list(range(slot, self._num_replicas, self._num_slots)) + for slot in range(self._num_slots) + ] + else: + self._groups = [ + [int(r) for r in group] + for group in _np.array_split( + _np.arange(self._num_replicas), self._num_slots + ) + ] # The slot that hosts each replica. self._replica_slot = [None] * self._num_replicas @@ -228,6 +327,15 @@ def __setstate__(self, state): self._energy_trajectories = [None] * n if not hasattr(self, "_ghost_files"): self._ghost_files = [None] * n + # Without these, replicas start from the shared clock with all hot. + if not hasattr(self, "_clocks"): + self._clocks = [None] * n + if not hasattr(self, "_hot_replicas"): + self._hot_replicas = None + if not hasattr(self, "_cold_votes"): + self._cold_votes = [0] * n + if not hasattr(self, "_adaptive_data"): + self._adaptive_data = None self._gcmc_num_waters = [None] * n # The slot layout is not pickled, since it is rebuilt by @@ -239,6 +347,9 @@ def __setstate__(self, state): self._update_constraints = True if not hasattr(self, "_virtual_sites"): self._virtual_sites = None + # Set by the runner on a restart, since it follows the configuration. + if not hasattr(self, "_strided"): + self._strided = False # Convert a legacy checkpoint to the current convention, in which the # stored state of a replica is its own, with the last mix already @@ -280,6 +391,10 @@ def __getstate__(self): "_update_constraints": self._update_constraints, "_states": self._states, "_time": self._time, + "_clocks": self._clocks, + "_hot_replicas": self._hot_replicas, + "_cold_votes": self._cold_votes, + "_adaptive_data": self._adaptive_data, "_openmm_states": self._openmm_states, "_virtual_sites": self._virtual_sites, # Don't pickle the GCMC samplers since they need to be recreated. @@ -798,19 +913,28 @@ def save_gcmc_state(self, slot, replica): # Store the state. self._gcmc_states[replica] = gcmc_sampler.water_state() - def get_clock(self): + def get_clock(self, replica=None): """ Get the simulation clock. - Every replica advances by the same amount each cycle, so the clock is - common to all of them and is read from the first slot. + Parameters + ---------- + + replica: int + The index of the replica. If None, the clock of the first slot is + returned, which is common to every replica until they have each + been propagated and stored a clock of their own. Returns ------- dict - The clock, as returned by Dynamics._get_clock(). + The clock, as returned by Dynamics._get_clock(), or None if the + replica has no clock of its own. """ + if replica is not None: + return self._clocks[replica] + return self._dynamics[0]._get_clock() def get_states(self): @@ -943,6 +1067,7 @@ def store_replica(self, replica): slot = self._replica_slot[replica] self.save_openmm_state(slot, replica) + self._clocks[replica] = self._dynamics[slot]._get_clock() if self._gcmc_samplers[slot] is not None: self.save_gcmc_state(slot, replica) @@ -1257,6 +1382,14 @@ def __init__(self, system, config): self._set_num_slots() + # The particles of the bond that is broken or made, whose distance is + # used to rank the windows for adaptive replica exchange. + self._ring_pair = ( + self._find_ring_pair() + if self._config.adaptive_repex and self._is_ring_break + else None + ) + # 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 @@ -1352,6 +1485,7 @@ def __init__(self, system, config): update_constraints=self._config.update_constraints, constraint_lambda_index=self._constraint_lambda_index, gpu_devices=self._gpu_devices, + strided=self._is_strided, ) self._dynamics_cache._virtual_sites = self._virtual_sites @@ -1438,6 +1572,7 @@ def __init__(self, system, config): self._dynamics_cache._constraint_lambda_index = ( self._constraint_lambda_index ) + self._dynamics_cache._strided = self._is_strided self._dynamics_cache._build_slot_layout() # Create the dynamics objects. @@ -1505,8 +1640,10 @@ def __init__(self, system, config): ) else: self._start_block = 0 + self._start_time = time else: self._start_block = 0 + self._start_time = _sr.u("0 ps") # Create a terminal flip sampler per replica (if terminal groups were detected). if self._terminal_groups: @@ -1568,7 +1705,7 @@ def _build_lambda(self, replica): return self._lambda_values[self._constraint_lambda_index] - def _replica_passes(self, cycle): + def _replica_passes(self, cycle, replicas=None): """ Work out which replicas to propagate in each pass of a cycle. @@ -1587,6 +1724,9 @@ def _replica_passes(self, cycle): cycle: int The index of the current cycle. + replicas: list of int + The replicas to propagate. If None, every replica is propagated. + Returns ------- @@ -1595,6 +1735,11 @@ def _replica_passes(self, cycle): """ groups = self._dynamics_cache._groups + if replicas is not None: + replicas = set(replicas) + groups = [[r for r in group if r in replicas] for group in groups] + groups = [group for group in groups if group] + # Traverse the groups backwards on odd cycles. is_reversed = cycle % 2 == 1 @@ -1648,7 +1793,7 @@ def _set_num_slots(self): the configuration options that only apply when contexts are re-used across lambda values. - Sets self._num_slots and self._is_cached. + Sets self._num_slots, self._is_cached and self._is_strided. """ num_replicas = self._num_replicas @@ -1663,6 +1808,7 @@ def _set_num_slots(self): # value it was created at, so there is no need to fix the constraints # at a common one. self._is_cached = self._num_slots < num_replicas + self._is_strided = self._is_cached and self._config.adaptive_repex self._constraint_lambda_index = None # Contexts run concurrently, so it's the number of them that determines @@ -1962,6 +2108,53 @@ def run(self): # Whether the most recent cycle saved the replica exchange state. is_checkpoint = False + num_replicas = self._config.num_lambda + is_adaptive = self._config.adaptive_repex + + # The cycle on which each replica was last propagated, used to catch up + # on the MC moves that fell due while it wasn't. + last_cycle = [-1] * num_replicas + + # The number of energy samples each replica has collected since the + # last checkpoint. + num_samples = [0] * num_replicas + + if is_adaptive: + self._reset_adaptive_data() + self._restore_adaptive_data() + + end_time = self._start_time + self._config.runtime + if end_time <= self._config.adaptive_burn_in: + _logger.warning( + "The 'runtime' is no longer than 'adaptive_burn_in', so the " + "windows will never be ranked. Every window will be " + "propagated every cycle." + ) + elif checkpoint_frequency.value() == 0.0: + _logger.warning( + "Adaptive replica exchange ranks the windows at checkpoints, " + "but 'checkpoint_frequency' is zero. Every window will be " + "propagated every cycle." + ) + elif cycles_per_checkpoint <= self._config.adaptive_mix_interval: + _logger.warning( + "Every window is propagated on checkpoint cycles, which are " + "no further apart than 'adaptive_mix_interval'. Adaptive " + "replica exchange will have little effect." + ) + + if ( + self._save_frames or self._config.gcmc + ) and cycles_per_frame <= self._config.adaptive_mix_interval: + _logger.warning( + "Every window is propagated on frame cycles, which are no " + "further apart than 'adaptive_mix_interval'. Adaptive " + "replica exchange will have little effect." + ) + else: + # Adaptive replica exchange may have been switched off on restart. + self._dynamics_cache._hot_replicas = None + # Perform the replica exchange simulation. for i in range(cycles): _logger.info(f"Running dynamics for cycle {i + 1} of {cycles}") @@ -1982,52 +2175,68 @@ def run(self): # non-integer ratios between the checkpoint and energy frequencies. is_checkpoint = (i + 1) >= next_checkpoint - 1e-10 - # Whether to perform a GCMC move before the dynamics block. - is_gcmc = (i + 1) % cycles_per_gcmc == 0 - - # Whether to perform a terminal flip move before the dynamics block. - is_terminal_flip = (i + 1) % cycles_per_flip == 0 - # Whether a frame is saved at the end of the cycle. write_gcmc_ghosts = (i + 1) % cycles_per_frame == 0 - # Current simulation time in ns for energy components saving. - time_ns = ( - ( - self._start_block * checkpoint_frequency - + (i + 1) * self._config.energy_frequency - ).to("ns") - if self._config.save_energy_components - else None - ) - # Whether the checkpoint files are written at the end of this cycle. do_checkpoint = is_checkpoint or i == cycles - 1 - # Capture the simulation clock at the start of the cycle. Every - # replica advances by the same amount each cycle, so a slot has to - # be rewound to this point before it propagates the next replica. + # Only the hot replicas are propagated, except on mix cycles and + # cycles that write output for every replica. + hot = self._dynamics_cache._hot_replicas + is_full = ( + hot is None + or (i + 1) % self._config.adaptive_mix_interval == 0 + or do_checkpoint + or (write_gcmc_ghosts and (self._save_frames or self._config.gcmc)) + ) + active = list(range(num_replicas)) if is_full else hot + + if not is_full: + _logger.info(f"Propagating {len(active)} hot replicas: {active}") + + # Frames are scheduled by the cycle rather than by each replica's + # own clock, which falls behind while the replica is cold. + save_frame = ( + write_gcmc_ghosts if is_adaptive and self._save_frames else None + ) + + # Capture the simulation clock at the start of the cycle. This seeds + # any replica that doesn't yet have a clock of its own. clock = self._dynamics_cache.get_clock() # Propagate the replicas, one pass at a time. Each pass takes at # most one replica from each slot, and there is a single pass when # there is a slot per replica. - for batch in self._replica_passes(i): + for batch in self._replica_passes(i, replicas=None if is_full else active): # Run a dynamics block for each replica in the pass, making sure # each GPU is only oversubscribed by a factor of # self._config.oversubscription_factor. for j in range(ceil(len(batch) / num_workers)): replicas = batch[j * num_workers : (j + 1) * num_workers] + + # Whether each replica is due a GCMC or terminal flip move, + # including any that fell due while it wasn't propagated. + is_gcmc = [ + self._is_due(last_cycle[r], i, cycles_per_gcmc) + for r in replicas + ] + is_terminal_flip = [ + self._is_due(last_cycle[r], i, cycles_per_flip) + for r in replicas + ] + try: for result, index, energies in dynamics_executor.map( self._run_block, replicas, repeat(self._lambda_values), - repeat(is_gcmc), + is_gcmc, repeat(write_gcmc_ghosts), - repeat(is_terminal_flip), - repeat(time_ns), + is_terminal_flip, + repeat(self._config.save_energy_components), repeat(clock), + repeat(save_frame), ): if not result: _logger.error( @@ -2070,15 +2279,31 @@ def run(self): _logger.error("Commit cancelled. Exiting.") _sys.exit(1) + for index in active: + last_cycle[index] = i + num_samples[index] += 1 + # Assemble an energy matrix from the results. _logger.info("Assembling energy matrix") energy_matrix = self._assemble_results(results) - # Mix the replicas. + # Record the energy differences used to rank the windows, once the + # burn-in is complete. + cycle_time = self._start_time + (i + 1) * self._config.energy_frequency + is_burnt_in = cycle_time > self._config.adaptive_burn_in + if is_adaptive and is_burnt_in: + if ( + cycle_time - self._config.energy_frequency + <= self._config.adaptive_burn_in + ): + _logger.info("Adaptive replica exchange burn-in complete") + self._record_adaptive_data(energy_matrix, active) + + # Mix the replicas that were propagated. _logger.info("Mixing replicas") self._dynamics_cache.set_states( - self._mix_replicas( - self._config.num_lambda, + _mix_subset( + _np.array(active, dtype=_np.int64), energy_matrix, self._dynamics_cache.get_proposed(), self._dynamics_cache.get_accepted(), @@ -2095,6 +2320,11 @@ def run(self): # external process reading the output directory always sees a # coherent set rather than a mixture of new and old files. if do_checkpoint: + # Rank the windows before the state is saved, so that a restart + # carries on with the same hot replicas. + if is_adaptive and is_checkpoint and is_burnt_in: + self._rank_windows(num_workers) + # Create the lock. lock = _FileLock(self._lock_file) @@ -2141,6 +2371,7 @@ def run(self): repeat(block), repeat(num_blocks + int(rem > 0)), repeat(i == cycles - 1), + [num_samples[r] for r in replicas], ): if error: _logger.error( @@ -2155,6 +2386,12 @@ def run(self): _logger.error("Checkpoint cancelled. Exiting.") _sys.exit(1) + num_samples = [0] * num_replicas + + # Replicas' own clocks fall behind while they are cold, so + # the simulation time is that of the cycle. + self._dynamics_cache._time = cycle_time + self._save_repex_state(final=not is_checkpoint) # This is a checkpoint cycle. @@ -2202,8 +2439,9 @@ def _run_block( is_gcmc=False, write_gcmc_ghosts=False, is_terminal_flip=False, - time_ns=None, + save_energy_components=False, clock=None, + save_frame=None, ): """ Run a dynamics block for a given replica. @@ -2236,15 +2474,17 @@ def _run_block( Whether a terminal flip MC move should be performed before the dynamics block. - time_ns: float or None - The current simulation time in nanoseconds, used when saving energy - components. If None, energy components are not saved. + save_energy_components: bool + Whether to save the energy contribution of each force. clock: dict The simulation clock at the start of the cycle, as returned by - Dynamics._get_clock(). Every replica advances by the same amount - each cycle, so a slot must be rewound to the start of the cycle - before propagating the next replica. + Dynamics._get_clock(). Only used if the replica doesn't yet have a + clock of its own. + + save_frame: bool + Whether to save a frame at the end of the block. If None, frames + are saved every 'frame_frequency' of the replica's own clock. Returns ------- @@ -2264,14 +2504,34 @@ def _run_block( lam = lambdas[index] try: - # Make the replica resident in its slot. - self._dynamics_cache.load_replica(index, clock=clock) + # Make the replica resident in its slot, restoring its own clock so + # that its energies and frames are never written over. + own_clock = self._dynamics_cache.get_clock(index) + self._dynamics_cache.load_replica( + index, clock=clock if own_clock is None else own_clock + ) # Get the dynamics object (and GCMC sampler). dynamics, gcmc_sampler = self._dynamics_cache.get( self._dynamics_cache.slot_for(index) ) + # Schedule the next frame for the end of this block, or never. + if save_frame is not None: + steps = max( + 1, + int( + self._config.energy_frequency.to("ps") + / dynamics.timestep().to("ps") + ), + ) + dynamics._d._next_save_frame = ( + dynamics.current_step() + steps if save_frame else _sys.maxsize + ) + # Let Sire reschedule frames if adaptive scheduling was switched off. + elif dynamics._d._next_save_frame == _sys.maxsize: + dynamics._d._next_save_frame = None + auto_fix_minimise = self._config.auto_fix_minimise # Perform the GCMC move before dynamics so that the energies @@ -2337,8 +2597,10 @@ def _run_block( self._dynamics_cache.store_replica(index) # Save the energy contribution for each force. - if self._config.save_energy_components and time_ns is not None: - self._save_energy_components(index, dynamics.context(), time_ns) + if save_energy_components: + self._save_energy_components( + index, dynamics.context(), dynamics.current_time().to("ns") + ) # Get the energy at each lambda value. energies = dynamics._current_energy_array() @@ -2670,8 +2932,10 @@ def _assemble_results(self, results): results: list The results from the repex dynamics block. """ - # Create the matrix. - matrix = _np.zeros((len(results), len(results))) + # Create the matrix. Rows of replicas that weren't propagated are left + # empty, since they aren't mixed. + num_replicas = self._config.num_lambda + matrix = _np.zeros((num_replicas, num_replicas)) # Fill the matrix. The energy returned by the dynamics block already # includes the pressure and grand canonical contributions. @@ -2681,6 +2945,407 @@ def _assemble_results(self, results): return matrix + def _find_ring_pair(self): + """ + Find the OpenMM particle indices of the bond that is broken or made by + a ring-breaking perturbation. + + Returns + ------- + + tuple of int + The particle indices, or None unless exactly one bond is broken or + made, since the distance of a single bond is used. + """ + atoms = self._system.atoms() + pairs = [] + + for mol in self._system["property is_perturbable"].molecules(): + for prop in ("ring_breaking_bonds", "ring_making_bonds"): + if mol.has_property(prop): + flat = mol.property(prop).to_list() + for i in range(0, len(flat) - 1, 2): + pairs.append( + tuple( + atoms.find(mol.atoms()[int(idx)]) + for idx in flat[i : i + 2] + ) + ) + + if len(pairs) != 1: + if len(pairs) > 1: + _logger.info( + f"{len(pairs)} bonds are broken or made, so the bond distance " + "is not used to rank windows for adaptive replica exchange." + ) + return None + + # Virtual sites are added after the atoms of their molecule, so shift + # the particles of any later molecule along. + first = atoms.find(atoms[pairs[0][0]].molecule().atoms()[0]) + offset = 0 + if self._virtual_sites is not None: + offset = int(sum(n for start, n in self._virtual_sites if start < first)) + + pair = tuple(int(idx) + offset for idx in pairs[0]) + _logger.info( + f"Ranking windows for adaptive replica exchange using the distance " + f"between particles {pair[0]} and {pair[1]}" + ) + + return pair + + @staticmethod + def _is_due(last_cycle, cycle, cycles_per_move): + """ + Whether a periodic move falls due on a cycle, or fell due on one of + the cycles since a replica was last propagated. + + Parameters + ---------- + + last_cycle: int + The cycle on which the replica was last propagated, or -1. + + cycle: int + The current cycle. + + cycles_per_move: int + The number of cycles between moves. + + Returns + ------- + + bool + Whether the move is due. + """ + return (cycle + 1) // cycles_per_move > (last_cycle + 1) // cycles_per_move + + def _reset_adaptive_data(self): + """ + Clear the data used to rank the windows for adaptive replica exchange. + """ + num_replicas = self._config.num_lambda + + # The reduced energy difference to the lower and upper neighbouring + # window, sampled in each window. + self._adaptive_lower = [[] for _ in range(num_replicas)] + self._adaptive_upper = [[] for _ in range(num_replicas)] + + # The distance of the bond that is broken or made, sampled in each + # window. + self._adaptive_distance = [[] for _ in range(num_replicas)] + + # The swap counts when the data was cleared. + self._adaptive_proposed = self._dynamics_cache.get_proposed().copy() + self._adaptive_accepted = self._dynamics_cache.get_accepted().copy() + + def _store_adaptive_data(self): + """ + Store the data used to rank the windows in the dynamics cache, so that + it is kept on restart. + """ + self._dynamics_cache._adaptive_data = { + "lower": [_np.array(x) for x in self._adaptive_lower], + "upper": [_np.array(x) for x in self._adaptive_upper], + "distance": [_np.array(x) for x in self._adaptive_distance], + "proposed": self._adaptive_proposed, + "accepted": self._adaptive_accepted, + } + + def _restore_adaptive_data(self): + """ + Restore the data used to rank the windows from the dynamics cache. + """ + data = self._dynamics_cache._adaptive_data + if data is None: + return + + self._adaptive_lower = [list(x) for x in data["lower"]] + self._adaptive_upper = [list(x) for x in data["upper"]] + self._adaptive_distance = [list(x) for x in data["distance"]] + self._adaptive_proposed = data["proposed"] + self._adaptive_accepted = data["accepted"] + + def _record_adaptive_data(self, energy_matrix, replicas): + """ + Record the reduced energy differences between neighbouring windows, + and the distance of the bond that is broken or made, if any. Must be + called before the replicas are mixed. + + Parameters + ---------- + + energy_matrix: np.ndarray + The reduced energy of each replica at each window. + + replicas: list of int + The replicas that were propagated. + """ + num_replicas = self._config.num_lambda + + for k in replicas: + row = energy_matrix[k] + if k > 0: + self._adaptive_lower[k].append(row[k - 1] - row[k]) + if k < num_replicas - 1: + self._adaptive_upper[k].append(row[k + 1] - row[k]) + + if self._ring_pair is not None: + import openmm.unit as _omm_unit + + positions = DynamicsCache._get_positions( + self._dynamics_cache._openmm_states[k] + ) + i, j = self._ring_pair + pair = _np.asarray(positions[[i, j]].value_in_unit(_omm_unit.angstrom)) + self._adaptive_distance[k].append( + float(_np.linalg.norm(pair[1] - pair[0])) + ) + + @staticmethod + def _statistical_inefficiency(x): + """ + Estimate the statistical inefficiency of a time series, i.e. the number + of samples per independent sample. + + Parameters + ---------- + + x: np.ndarray + The time series. + + Returns + ------- + + float + The statistical inefficiency. + """ + n = len(x) + if n < 2: + return 1.0 + + dx = x - x.mean() + var = dx.dot(dx) / n + if var == 0.0: + return 1.0 + + # The unnormalised autocovariance at every lag. + f = _np.fft.rfft(dx, 2 * n) + acov = _np.fft.irfft(f * _np.conj(f))[:n] + + # Sum the autocorrelation function up to its first non-positive value. + g = 1.0 + for t in range(1, n): + c = acov[t] / ((n - t) * var) + if c <= 0.0: + break + g += 2.0 * c * (1.0 - t / n) + + return max(1.0, g) + + @staticmethod + def _exp_estimate(du): + """ + The free energy difference from exponential averaging of reduced + energy differences. + + Parameters + ---------- + + du: np.ndarray + The reduced energy differences. + + Returns + ------- + + float + The reduced free energy difference. + """ + m = (-du).max() + return -(m + _np.log(_np.mean(_np.exp(-du - m)))) + + @classmethod + def _is_unconverged(cls, fwd, rev, g_fwd, g_rev): + """ + Whether the free energy difference between a pair of neighbouring + windows shows signs of poor sampling, regardless of its uncertainty. + + Parameters + ---------- + + fwd: np.ndarray + The reduced energy difference to the upper window, sampled in the + lower window. + + rev: np.ndarray + The reduced energy difference to the lower window, sampled in the + upper window. + + g_fwd, g_rev: float + The statistical inefficiencies of fwd and rev. + + Returns + ------- + + bool + Whether the pair is unconverged. + """ + # The forward and reverse estimates disagree by more than 1 kT. + if abs(cls._exp_estimate(fwd) + cls._exp_estimate(rev)) > 1.0: + return True + + # The estimates from the two halves of the data differ by more than + # 0.25 kT, and by more than twice their combined standard error. + estimates = [] + errors = [] + for half in range(2): + f = _np.array_split(fwd, 2)[half] + r = _np.array_split(rev, 2)[half] + estimates.append(0.5 * (cls._exp_estimate(f) - cls._exp_estimate(r))) + errors.append( + 0.5 * _np.sqrt(f.var() * g_fwd / len(f) + r.var() * g_rev / len(r)) + ) + + drift = abs(estimates[0] - estimates[1]) + return drift > 0.25 and drift > 2.0 * _np.hypot(*errors) + + def _rank_windows(self, num_workers): + """ + Rank the windows by the uncertainty they contribute to the free energy + and choose the hot replicas for adaptive replica exchange. + + Each pair of neighbouring windows is weighted by sigma * sqrt(g) of + its reduced energy differences, which sets the uncertainty of its free + energy difference for a fixed amount of sampling. Pairs are taken in + order of weight until they cover 'adaptive_coverage' of the total. + Pairs with signs of poor sampling are always taken, as are windows with + fewer than 'adaptive_min_samples' independent samples of either the + energy differences or, for a perturbation that breaks or makes a single + bond, the bond distance, and the windows either side of each one taken, + so that exchange carries configurations into and out of them. A window + only stops being hot once it has been found cold at two consecutive + rankings. The hot replicas are padded to a multiple of the number of + workers, so that the GPUs are evenly loaded. + + The ranking uses all of the data since the burn-in, since the slowest + windows decorrelate over far longer than a checkpoint interval. + + Parameters + ---------- + + num_workers: int + The number of replicas that are propagated concurrently. + """ + from math import ceil + + num_replicas = self._config.num_lambda + lower = [_np.array(x) for x in self._adaptive_lower] + upper = [_np.array(x) for x in self._adaptive_upper] + + min_samples = self._config.adaptive_min_samples + + # Every window needs enough samples to estimate its fluctuations. + num_samples = min(len(x) for x in lower[1:] + upper[:-1]) + if num_samples < min_samples: + _logger.info(f"Not ranking windows: {num_samples} samples of {min_samples}") + return + + g_lower = [self._statistical_inefficiency(x) for x in lower] + g_upper = [self._statistical_inefficiency(x) for x in upper] + + # Windows with too few independent samples can't be made cold. + undersampled = [] + for k in range(num_replicas): + series = [] + if k > 0: + series.append((lower[k], g_lower[k])) + if k < num_replicas - 1: + series.append((upper[k], g_upper[k])) + if min(len(x) / g for x, g in series) < min_samples: + undersampled.append(k) + + # Nor can windows with a slow bond distance, which can bias the free + # energy without showing in the energy differences. + slow_distance = [] + if self._ring_pair is not None: + for k, x in enumerate(self._adaptive_distance): + x = _np.array(x) + if len(x) / self._statistical_inefficiency(x) < min_samples: + slow_distance.append(k) + + proposed = self._dynamics_cache.get_proposed() - self._adaptive_proposed + accepted = self._dynamics_cache.get_accepted() - self._adaptive_accepted + + weights = _np.zeros(num_replicas - 1) + flagged = [] + for k in range(num_replicas - 1): + fwd = upper[k] + rev = lower[k + 1] + weights[k] = _np.sqrt(fwd.var() * g_upper[k] + rev.var() * g_lower[k + 1]) + + # Flag pairs that rarely exchange, or whose estimates disagree. + num_proposed = proposed[k, k + 1] + if ( + num_proposed > 0 and accepted[k, k + 1] / num_proposed < 0.05 + ) or self._is_unconverged(fwd, rev, g_upper[k], g_lower[k + 1]): + flagged.append(k) + + # Take pairs in order of weight until they cover enough of the total. + hot = set() + covered = 0.0 + for k in _np.argsort(-weights): + if covered >= self._config.adaptive_coverage * weights.sum(): + break + hot.update((int(k), int(k) + 1)) + covered += weights[k] + + for k in flagged: + hot.update((k, k + 1)) + + hot.update(undersampled) + hot.update(slow_distance) + + hot.update([r + d for r in hot for d in (-1, 1) if 0 <= r + d < num_replicas]) + + # Keep windows hot until they have been found cold twice in a row. + previous = self._dynamics_cache._hot_replicas + votes = self._dynamics_cache._cold_votes + for r in range(num_replicas): + if r in hot: + votes[r] = 0 + elif previous is not None and r in previous: + votes[r] += 1 + if votes[r] < 2: + hot.add(r) + + # Pad with the highest weighted of the remaining windows. + window_weights = [ + max(weights[k] for k in (r - 1, r) if 0 <= k < num_replicas - 1) + for r in range(num_replicas) + ] + target = min(num_replicas, ceil(len(hot) / num_workers) * num_workers) + cold = sorted(set(range(num_replicas)) - hot, key=lambda r: -window_weights[r]) + hot.update(cold[: target - len(hot)]) + + # An empty ranking carries no information, so every replica stays hot. + hot = sorted(hot) if hot else list(range(num_replicas)) + self._dynamics_cache._hot_replicas = None if len(hot) == num_replicas else hot + + _logger.info( + "Window pair weights (sigma * sqrt(g)): " + + ", ".join(f"{w:.3f}" for w in weights) + ) + if flagged: + _logger.info(f"Unconverged window pairs: {flagged}") + if undersampled: + _logger.info(f"Undersampled windows: {undersampled}") + if slow_distance: + _logger.info(f"Windows with a slow bond distance: {slow_distance}") + _logger.info( + f"Hot replicas: {hot}" if len(hot) < num_replicas else "All replicas hot" + ) + def _check_restart(self): """ Check the output directory for a valid restart state. @@ -2789,7 +3454,15 @@ def _commit_replica(self, index): return index, None - def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): + def _checkpoint( + self, + index, + lambdas, + block, + num_blocks, + is_final_block=False, + num_energy_rows=None, + ): """ Checkpoint the simulation. @@ -2811,6 +3484,10 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): is_final_block: bool Whether this is the final block. + num_energy_rows: int + The number of energy samples the replica has collected since the + last checkpoint. If None, a full block is assumed. + Returns ------- @@ -2846,7 +3523,12 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): # Call the base class checkpoint method to save the system state. with self._lock: index, error = super()._checkpoint( - system, index, block, speed, is_final_block=is_final_block + system, + index, + block, + speed, + is_final_block=is_final_block, + num_energy_rows=num_energy_rows, ) if error is not None: @@ -2942,70 +3624,6 @@ def _consolidate_trajectories(self): for chunk in traj_chunks: _Path_local(chunk).unlink() - @staticmethod - @_njit - def _mix_replicas(num_replicas, energy_matrix, proposed, accepted): - """ - Mix the replicas. - - Parameters - ---------- - - num_replicas: int - The number of replicas. - - energy_matrix: np.ndarray - The energy matrix for the replicas. - - Returns - ------- - - states: np.ndarray - The new states, where states[i] is the replica whose configuration - seeds replica i. - """ - - # Adapted from OpenMMTools: https://github.com/choderalab/openmmtools - - # Set the states to the initial order. - states = _np.arange(num_replicas) - - # Attempt swaps. - for swap in range(num_replicas**3): - # Choose two replicas to swap. - replica_i = _np.random.randint(num_replicas) - replica_j = _np.random.randint(num_replicas) - - # Get the current state. - state_i = states[replica_i] - state_j = states[replica_j] - - # Record that we have proposed a swap. - proposed[state_i, state_j] += 1 - proposed[state_j, state_i] += 1 - - # Get the energies. - energy_ii = energy_matrix[replica_i, state_i] - energy_jj = energy_matrix[replica_j, state_j] - energy_ij = energy_matrix[replica_i, state_j] - energy_ji = energy_matrix[replica_j, state_i] - - # Compute the log probability of the swap. - log_p_swap = -(energy_ij + energy_ji) + energy_ii + energy_jj - - # Accept the swap and update the states. - if log_p_swap >= 0 or _np.random.rand() < _np.exp(log_p_swap): - # Swap the states. - states[replica_i] = state_j - states[replica_j] = state_i - # Record the swap. - accepted[state_i, state_j] += 1 - accepted[state_j, state_i] += 1 - - # Here states[i] is the state that replica i's configuration moves to, - # whereas the configurations are moved into fixed states, so invert. - return _np.argsort(states) - def _merge_gcmc_stats(self): """ Merge the GCMC sampling statistics from every sampler. @@ -3055,6 +3673,10 @@ def _save_repex_state(self, final=False): _logger.info(f"Saving {label} state") self._save_sampler_stats() + if self._config.adaptive_repex: + self._store_adaptive_data() + else: + self._dynamics_cache._adaptive_data = None with open(self._repex_state, "wb") as f: _pickle.dump(self._dynamics_cache, f) diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 2d688e2..813b77c 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -6,6 +6,7 @@ from somd2.runner import RepexRunner from somd2.runner._base import RunnerBase +from somd2.runner._repex import _mix_subset from somd2.config import Config from tests.conftest import has_cuda @@ -53,7 +54,7 @@ def test_repex_mixing(): accepted = np.zeros((10, 10), dtype=np.int32) # Perform the mixing. - states = RepexRunner._mix_replicas(10, energy_matrix, proposed, accepted) + states = _mix_subset(np.arange(10), energy_matrix, proposed, accepted) # Make sure that exchanges are always accepted. assert (proposed == accepted).all() @@ -68,7 +69,7 @@ def test_repex_mixing(): accepted = np.zeros((10, 10), dtype=np.int32) # Perform the mixing. - states = RepexRunner._mix_replicas(10, energy_matrix, proposed, accepted) + states = _mix_subset(np.arange(10), energy_matrix, proposed, accepted) # Get the off-diagonal elements of the accepted matrix. off_diagonal = accepted - np.diag(np.diag(accepted)) @@ -95,7 +96,7 @@ def test_repex_mixing_moves_configurations(): accepted = np.zeros((num_replicas, num_replicas), dtype=np.int32) np.random.seed(42) - states = RepexRunner._mix_replicas(num_replicas, energy_matrix, proposed, accepted) + states = _mix_subset(np.arange(num_replicas), energy_matrix, proposed, accepted) cache = object.__new__(DynamicsCache) cache._openmm_states = list(range(num_replicas)) @@ -164,6 +165,9 @@ def clear_energy_cache(self): def set_lambda(self, lam, rest2_scale=None, update_constraints=None): self._context.lam = lam + def _get_clock(self): + return None + def _set_clock(self, clock): pass @@ -241,8 +245,13 @@ def test_repex_mixing_samples_boltzmann(slot_schedule, inverted): cache._rest2_scale_factors = [1.0] * n cache._num_replicas = n cache._update_constraints = True + cache._strided = False cache._states = np.arange(n) cache._time = None + cache._clocks = [None] * n + cache._hot_replicas = None + cache._cold_votes = [0] * n + cache._adaptive_data = None cache._virtual_sites = None cache._gcmc_stats = None cache._terminal_flip_stats = [[0, 0] for _ in range(n)] @@ -294,8 +303,8 @@ def test_repex_mixing_samples_boltzmann(slot_schedule, inverted): energies[replica] = U[:, config] cache.store_replica(replica) - states = RepexRunner._mix_replicas( - n, energies, cache._num_proposed, cache._num_accepted + states = _mix_subset( + np.arange(n), energies, cache._num_proposed, cache._num_accepted ) if inverted: states = np.argsort(states) @@ -1319,3 +1328,456 @@ def test_repex_perturbed_system_seeding(ethane_methanol, swap_end_states): f"group {group} has {len(wrong)} replicas seeded from the " f"wrong end state: {wrong}" ) + + +@pytest.mark.parametrize( + "last_cycle, cycle, cycles_per_move, expected", + [ + (2, 3, 4, True), + (1, 2, 4, False), + # A move that fell due while the replica wasn't propagated. + (0, 5, 4, True), + (4, 5, 4, False), + (-1, 0, 1, True), + ], +) +def test_is_due(last_cycle, cycle, cycles_per_move, expected): + """ + Validate that periodic moves fall due on schedule, and that a replica + catches up on one that fell due while it wasn't propagated. + """ + assert RepexRunner._is_due(last_cycle, cycle, cycles_per_move) == expected + + +def test_mix_subset_leaves_others_in_place(): + """ + Validate that mixing a subset of the replicas only swaps between them. + """ + n = 6 + subset = np.array([1, 2, 4]) + energy_matrix = np.zeros((n, n)) + proposed = np.zeros((n, n)) + accepted = np.zeros((n, n)) + + # Every swap is accepted, so the subset is thoroughly shuffled. + states = _mix_subset(subset, energy_matrix, proposed, accepted) + + for r in set(range(n)) - set(subset): + assert states[r] == r + assert sorted(states[subset]) == sorted(subset) + + # Only swaps between members of the subset are recorded. + outside = [r for r in range(n) if r not in subset] + assert proposed[outside, :].sum() == 0 + assert proposed[:, outside].sum() == 0 + + +def test_strided_slot_layout(): + """ + Validate that a strided layout spreads any contiguous range of replicas + evenly over the slots. + """ + from somd2.runner._repex import DynamicsCache + + cache = object.__new__(DynamicsCache) + cache._num_replicas = 10 + cache._num_slots = 4 + cache._strided = True + cache._build_slot_layout() + + assert cache._groups == [[0, 4, 8], [1, 5, 9], [2, 6], [3, 7]] + + for start in range(7): + slots = [cache.slot_for(r) for r in range(start, start + 4)] + assert sorted(slots) == [0, 1, 2, 3] + + +def test_statistical_inefficiency(): + """ + Validate the statistical inefficiency against an AR(1) process, for which + it is (1 + phi) / (1 - phi). + """ + rng = np.random.default_rng(42) + + x = rng.normal(size=20000) + assert RepexRunner._statistical_inefficiency(x) == pytest.approx(1.0, abs=0.1) + + phi = 0.8 + y = np.zeros(20000) + for i in range(1, len(y)): + y[i] = phi * y[i - 1] + rng.normal() + expected = (1 + phi) / (1 - phi) + assert RepexRunner._statistical_inefficiency(y) == pytest.approx(expected, rel=0.15) + + +def _ranking_runner(num_lambda, **config): + """ + Create a runner holding only what is needed to rank windows. + """ + from types import SimpleNamespace + + runner = object.__new__(RepexRunner) + runner._config = Config(num_lambda=num_lambda, **config) + runner._lambda_values = list(np.linspace(0, 1, num_lambda)) + runner._ring_pair = None + zeros = np.zeros((num_lambda, num_lambda)) + runner._dynamics_cache = SimpleNamespace( + get_proposed=lambda: zeros, + get_accepted=lambda: zeros, + _hot_replicas=None, + _cold_votes=[0] * num_lambda, + ) + runner._reset_adaptive_data() + return runner + + +def _fill_ranking_data(runner, hard_pair, num_samples=200, seed=7): + """ + Record converged energy differences between every pair of neighbouring + windows, apart from one pair with a large spread. + """ + rng = np.random.default_rng(seed) + num_lambda = runner._config.num_lambda + + for k in range(num_lambda - 1): + sigma = 1.5 if k == hard_pair else 0.1 + runner._adaptive_upper[k] = list(rng.normal(0.5, sigma, num_samples)) + runner._adaptive_lower[k + 1] = list(rng.normal(-0.5, sigma, num_samples)) + + +def test_rank_windows(): + """ + Validate that the hot replicas are the windows either side of the pair with + the largest uncertainty, padded by a window each side, and that windows + only go cold after being found cold at two consecutive rankings. + """ + runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=10) + cache = runner._dynamics_cache + + _fill_ranking_data(runner, hard_pair=3) + runner._rank_windows(num_workers=2) + assert cache._hot_replicas == [2, 3, 4, 5] + + # Windows 2 and 3 are found cold once, so stay hot. + _fill_ranking_data(runner, hard_pair=5) + runner._rank_windows(num_workers=2) + assert cache._hot_replicas == [2, 3, 4, 5, 6, 7] + + # Found cold a second time. + _fill_ranking_data(runner, hard_pair=5) + runner._rank_windows(num_workers=2) + assert cache._hot_replicas == [4, 5, 6, 7] + + +def test_rank_windows_pads_to_workers(): + """ + Validate that the hot replicas are padded to a multiple of the number of + workers, and that every replica being hot is recorded as None. + """ + runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=10) + + _fill_ranking_data(runner, hard_pair=3) + runner._rank_windows(num_workers=3) + assert len(runner._dynamics_cache._hot_replicas) == 6 + + runner._dynamics_cache._hot_replicas = None + _fill_ranking_data(runner, hard_pair=3) + runner._rank_windows(num_workers=8) + assert runner._dynamics_cache._hot_replicas is None + + +def test_rank_windows_without_uncertainty(): + """ + Validate that every replica stays hot when no window contributes any + uncertainty, rather than none of them. + """ + runner = _ranking_runner(4, adaptive_min_samples=10) + + for k in range(3): + runner._adaptive_upper[k] = [0.5] * 100 + runner._adaptive_lower[k + 1] = [-0.5] * 100 + runner._rank_windows(num_workers=1) + + assert runner._dynamics_cache._hot_replicas is None + + +def test_adaptive_data_is_restored(): + """ + Validate that the data used to rank the windows survives being stored in + the dynamics cache and restored, as it is across a restart. + """ + import pickle + + runner = _ranking_runner(8, adaptive_min_samples=10) + runner._dynamics_cache._adaptive_data = None + _fill_ranking_data(runner, hard_pair=3) + runner._adaptive_distance[2] = [1.0, 2.0] + + expected = (runner._adaptive_lower, runner._adaptive_upper) + runner._store_adaptive_data() + data = pickle.loads(pickle.dumps(runner._dynamics_cache._adaptive_data)) + + runner._reset_adaptive_data() + runner._dynamics_cache._adaptive_data = data + runner._restore_adaptive_data() + + assert runner._adaptive_lower == expected[0] + assert runner._adaptive_upper == expected[1] + assert runner._adaptive_distance[2] == [1.0, 2.0] + + +def test_rank_windows_needs_enough_samples(): + """ + Validate that windows aren't ranked until every one has enough samples, + and that the data keeps accumulating until they do. + """ + runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=50) + + _fill_ranking_data(runner, hard_pair=3, num_samples=20) + runner._rank_windows(num_workers=2) + + assert runner._dynamics_cache._hot_replicas is None + assert len(runner._adaptive_upper[0]) == 20 + + +def test_rank_windows_keeps_undersampled_windows_hot(): + """ + Validate that a window with too few independent samples stays hot, even + though its uncertainty is small. + """ + runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=50) + + _fill_ranking_data(runner, hard_pair=3) + + # A slowly decorrelating series with a small spread. + rng = np.random.default_rng(5) + y = np.zeros(200) + for i in range(1, len(y)): + y[i] = 0.95 * y[i - 1] + rng.normal(0.0, 0.01) + runner._adaptive_upper[6] = list(0.5 + y) + runner._rank_windows(num_workers=2) + + assert 6 in runner._dynamics_cache._hot_replicas + + +def test_rank_windows_keeps_slow_bond_distance_hot(): + """ + Validate that a window in which the distance of the bond that is broken + decorrelates slowly stays hot, even though its energies are converged. + """ + runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=50) + runner._ring_pair = (0, 1) + + _fill_ranking_data(runner, hard_pair=3) + + rng = np.random.default_rng(9) + for k in range(8): + runner._adaptive_distance[k] = list(rng.normal(1.5, 0.05, 200)) + + y = np.zeros(200) + for i in range(1, len(y)): + y[i] = 0.95 * y[i - 1] + rng.normal() + runner._adaptive_distance[7] = list(10.0 + y) + runner._rank_windows(num_workers=2) + + assert 7 in runner._dynamics_cache._hot_replicas + assert 0 not in runner._dynamics_cache._hot_replicas + + +def _ring_pair_runner(system, virtual_sites=None): + """ + Create a runner holding only what is needed to find the ring pair. + """ + runner = object.__new__(RepexRunner) + runner._system = system + runner._virtual_sites = virtual_sites + return runner + + +def test_find_ring_pair(syk_ring_break_mols): + """ + Validate that the bond that is broken is found from the perturbable + molecule, and isn't shifted by its own virtual sites, which come after its + atoms, or by those of later molecules. + """ + runner = _ring_pair_runner(syk_ring_break_mols) + assert runner._find_ring_pair() == (1, 7) + + runner = _ring_pair_runner(syk_ring_break_mols, np.array([[0, 3], [100, 4]])) + assert runner._find_ring_pair() == (1, 7) + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_adaptive_repex_records_bond_distance(syk_ring_break_mols): + """ + Validate that the distance of the bond that is broken is recorded in each + window during an adaptive replica exchange simulation. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "40fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "40fs", + "frame_frequency": "40fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": 3, + "replica_exchange": True, + "lambda_schedule": "ring_break_morph", + "adaptive_repex": True, + "adaptive_burn_in": "0 ps", + } + + runner = RepexRunner(syk_ring_break_mols, Config(**config)) + assert runner._ring_pair == (1, 7) + runner.run() + + for distances in runner._adaptive_distance: + assert len(distances) == 10 + + # The S-C bond is intact at lambda = 0. + assert np.mean(runner._adaptive_distance[0]) == pytest.approx(1.8, abs=0.2) + + +def test_find_ring_pair_needs_a_single_bond(syk_ring_break_mols): + """ + Validate that the distance isn't used when more than one bond is broken or + made. + """ + mols = syk_ring_break_mols.clone() + mol = mols["property is_perturbable"].molecules()[0] + mol = mol.edit().set_property("ring_making_bonds", [2, 8]).molecule().commit() + mols.update(mol) + + assert _ring_pair_runner(mols)._find_ring_pair() is None + + +def test_rank_windows_flags_unconverged_pairs(): + """ + Validate that a pair whose forward and reverse estimates disagree is made + hot, even though its uncertainty is small. + """ + runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=10) + + _fill_ranking_data(runner, hard_pair=3) + rng = np.random.default_rng(3) + runner._adaptive_lower[1] = list(rng.normal(-3.0, 0.1, 200)) + runner._rank_windows(num_workers=2) + + assert {0, 1, 2}.issubset(runner._dynamics_cache._hot_replicas) + + +def test_checkpoint_without_adaptive_state(): + """ + Validate that a checkpoint written before adaptive replica exchange was + added restores with every replica hot and no clocks of its own. + """ + from somd2.runner._repex import DynamicsCache + from tests.runner.test_gcmc_stats import TestLegacyRepexCheckpoint + + cache = object.__new__(DynamicsCache) + cache.__setstate__(TestLegacyRepexCheckpoint.make_state(None)) + + assert cache._hot_replicas is None + assert cache._clocks == [None] * 3 + assert cache._cold_votes == [0] * 3 + assert cache._adaptive_data is None + assert cache._strided is False + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +@pytest.mark.parametrize("max_contexts", [None, 2]) +def test_adaptive_repex_hot_replicas(ethane_methanol, max_contexts): + """ + Validate that only the hot replicas are propagated between full cycles, + that each replica records energies against its own clock without gaps or + repeats, that frames are saved for every replica on the same cycles, and + that a restart carries on with the same hot replicas. + """ + import pandas as pd + import sire as sr + + num_lambda = 4 + hot = [0, 1] + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "48fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "16fs", + "frame_frequency": "16fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": max_contexts, + "adaptive_repex": True, + "adaptive_mix_interval": 2, + # Never rank, so that the hot replicas stay as set. + "adaptive_burn_in": "1 ns", + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + # Not a ring-breaking perturbation, so there is no bond distance. + assert runner._ring_pair is None + + runner._dynamics_cache._hot_replicas = hot + runner.run() + + # Hot replicas run every cycle. The others only run on even cycles, + # which are every 'adaptive_mix_interval', or a checkpoint. + def expected_times(num_hot, num_cold, r): + num = num_hot if r in hot else num_cold + return [round(0.004 * (k + 1), 6) for k in range(num)] + + for i in range(num_lambda): + records = pd.read_parquet(runner._filenames[i]["energy_traj"]) + times = [round(t, 6) for t in records.index.get_level_values(0)] + assert times == expected_times(12, 6, i), f"replica {i}: {times}" + + # A frame on every checkpoint cycle, for every replica. + mols = sr.load( + [runner._filenames["topology0"], runner._filenames[i]["trajectory"]] + ) + assert mols.num_frames() == 3 + + assert runner._dynamics_cache._time.to("fs") == pytest.approx(48) + + # Restart, extending the runtime. + config["runtime"] = "64fs" + config["restart"] = True + + runner = RepexRunner(ethane_methanol, Config(**config)) + assert runner._dynamics_cache._hot_replicas == hot + runner.run() + + for i in range(num_lambda): + records = pd.read_parquet(runner._filenames[i]["energy_traj"]) + times = [round(t, 6) for t in records.index.get_level_values(0)] + assert times == expected_times(16, 8, i), f"replica {i}: {times}" + + # Restart with adaptive replica exchange switched off, so that every + # replica is propagated every cycle. + config["runtime"] = "80fs" + config["adaptive_repex"] = False + + runner = RepexRunner(ethane_methanol, Config(**config)) + runner.run() + + assert runner._dynamics_cache._hot_replicas is None + + for i in range(num_lambda): + records = pd.read_parquet(runner._filenames[i]["energy_traj"]) + times = [round(t, 6) for t in records.index.get_level_values(0)] + assert times == expected_times(20, 12, i), f"replica {i}: {times}" + + # Frames are scheduled by each replica's own clock again. + mols = sr.load( + [runner._filenames["topology0"], runner._filenames[i]["trajectory"]] + ) + assert mols.num_frames() == 5 diff --git a/tests/viewer/test_progress.py b/tests/viewer/test_progress.py new file mode 100644 index 0000000..44e106b --- /dev/null +++ b/tests/viewer/test_progress.py @@ -0,0 +1,105 @@ +import pickle + +import numpy as np +import pytest +import sire as sr + +from somd2._viewer._data import Simulation +from somd2.runner._repex import DynamicsCache + +LAMBDAS = [0.0, 0.5, 1.0] + +# A hot window at either end and a cold one in the middle. +WINDOWS = [ + {"lambda": 0.0, "samples": 2000}, + {"lambda": 0.5, "samples": 800}, + {"lambda": 1.0, "samples": 2000}, +] + + +def _simulation(tmp_path, time): + """ + Create a simulation whose replica exchange state is at the given time. + """ + n = len(LAMBDAS) + cache = object.__new__(DynamicsCache) + cache.__setstate__( + { + "_lambdas": LAMBDAS, + "_rest2_scale_factors": [1.0] * n, + "_num_replicas": n, + "_num_slots": n, + "_update_constraints": True, + "_states": np.arange(n), + "_time": sr.u(time), + "_clocks": [None] * n, + "_hot_replicas": [0, 2], + "_cold_votes": [0] * n, + "_adaptive_data": None, + "_openmm_states": [None] * n, + "_virtual_sites": None, + "_gcmc_samplers": [None] * n, + "_gcmc_states": [None] * n, + "_gcmc_stats": None, + "_terminal_flip_stats": [[0, 0]] * n, + "_num_proposed": np.zeros((n, n)), + "_num_accepted": np.zeros((n, n)), + "_num_swaps": np.zeros((n, n)), + } + ) + + with open(tmp_path / "repex_state.pkl", "wb") as f: + pickle.dump(cache, f) + + return Simulation(tmp_path, "test") + + +def _config(adaptive): + return { + "replica_exchange": True, + "adaptive_repex": adaptive, + "runtime": "5 ns", + "energy_frequency": "1 ps", + "checkpoint_frequency": "100 ps", + } + + +def test_adaptive_progress_uses_simulation_time(tmp_path): + """ + Validate that the progress of an adaptive replica exchange simulation is + the simulation time of the hot windows, not the samples of every window. + """ + simulation = _simulation(tmp_path, "2 ns") + progress = simulation._progress(_config(True), LAMBDAS, WINDOWS, None) + + assert progress["fraction"] == pytest.approx(0.4) + assert progress["completed_ns"] == pytest.approx(6.0) + + # Cold windows have fewer samples by design, so there is no expectation. + assert progress["expected_samples"] is None + + +def test_adaptive_progress_finishes_without_a_log(tmp_path): + """ + Validate that a finished adaptive replica exchange simulation is reported + as finished when there is no log, even though its cold windows have fewer + samples than the runtime implies. + """ + simulation = _simulation(tmp_path, "5 ns") + windows = [dict(w, samples=5000 if w["samples"] == 2000 else 1500) for w in WINDOWS] + progress = simulation._progress(_config(True), LAMBDAS, windows, None) + + assert progress["fraction"] == pytest.approx(1.0) + assert progress["status"] == "finished" + + +def test_standard_progress_uses_samples(tmp_path): + """ + Validate that the progress of a standard replica exchange simulation is + still taken from the samples of every window. + """ + simulation = _simulation(tmp_path, "2 ns") + progress = simulation._progress(_config(False), LAMBDAS, WINDOWS, None) + + assert progress["fraction"] == pytest.approx(4.8 / 15) + assert progress["expected_samples"] == 5000 From 0573c4c70a5520860fd2102b0655fd8f6e025684 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Oct 2026 16:10:18 +0100 Subject: [PATCH 02/32] Silence the expected RDKit warning when a conformer template does not fit. --- src/somd2/_viewer/_conformers.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/somd2/_viewer/_conformers.py b/src/somd2/_viewer/_conformers.py index fef6984..204ab8b 100644 --- a/src/somd2/_viewer/_conformers.py +++ b/src/somd2/_viewer/_conformers.py @@ -192,6 +192,7 @@ def _embed(mol, template=None, mapping=None): mapped onto it are placed and kept at its coordinates, so that the two overlay. """ + from rdkit import rdBase from rdkit.Chem import AllChem, rdMolAlign, rdMolTransforms coord_map = None @@ -202,7 +203,10 @@ def _embed(mol, template=None, mapping=None): kwargs = {"randomSeed": _SEED} if coord_map: kwargs["coordMap"] = coord_map - if AllChem.EmbedMolecule(mol, **kwargs) < 0: + # The template's coordinates may not fit, which the fallback below handles. + with rdBase.BlockLogs(): + embedded = AllChem.EmbedMolecule(mol, **kwargs) >= 0 + if not embedded: kwargs.pop("coordMap", None) if AllChem.EmbedMolecule(mol, useRandomCoords=True, **kwargs) < 0: raise RuntimeError("Couldn't generate 3D coordinates.") From 957eba6107275b6697688779a42f312736cbd6dd Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Oct 2026 18:20:30 +0100 Subject: [PATCH 03/32] Choose hot windows by variance per unit compute. --- README.md | 6 ++-- src/somd2/config/_config.py | 23 ------------ src/somd2/runner/_repex.py | 70 +++++++++++++++++++++++++------------ tests/runner/test_repex.py | 31 ++++++++++++---- 4 files changed, 76 insertions(+), 54 deletions(-) diff --git a/README.md b/README.md index b51c5e3..0c919f7 100644 --- a/README.md +++ b/README.md @@ -432,8 +432,10 @@ are propagated and exchanged every cycle, with every window joining in every Each pair of neighbouring windows is ranked by σ√g of the reduced energy differences between them, where σ is the standard deviation and g the statistical inefficiency, which sets the uncertainty of its free energy -difference. The highest ranked pairs that together account for -`--adaptive-coverage` of the total are made hot, along with: +difference. Pairs are made hot in order of rank for as long as doing so reduces +the variance of the free energy for a fixed amount of compute, accounting for +cold windows costing and sampling `--adaptive-mix-interval` times less than hot +ones. The following are always hot: - pairs whose forward and reverse free energy estimates disagree, or whose estimate drifts during the simulation, diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 412ccc3..95b50e9 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -120,7 +120,6 @@ class Config: "adaptive_repex", "adaptive_mix_interval", "adaptive_burn_in", - "adaptive_coverage", "adaptive_min_samples", "restraint_search_time", "restraint_search_frequency", @@ -195,7 +194,6 @@ def __init__( adaptive_repex=False, adaptive_mix_interval=4, adaptive_burn_in="500 ps", - adaptive_coverage=0.8, adaptive_min_samples=50, randomise_velocities=False, perturbed_system=None, @@ -523,11 +521,6 @@ def __init__( The simulation time of standard replica exchange to run before windows are first ranked when using adaptive replica exchange. - adaptive_coverage: float - The fraction of the total uncertainty, summed over neighbouring pairs of - lambda windows, that the hot windows must account for when using adaptive - replica exchange. - adaptive_min_samples: int The minimum number of statistically independent energy samples that a window must have collected since the burn-in before it can stop being hot @@ -828,7 +821,6 @@ def __init__( self.adaptive_repex = adaptive_repex self.adaptive_mix_interval = adaptive_mix_interval self.adaptive_burn_in = adaptive_burn_in - self.adaptive_coverage = adaptive_coverage self.adaptive_min_samples = adaptive_min_samples self.randomise_velocities = randomise_velocities self.perturbed_system = perturbed_system @@ -2229,21 +2221,6 @@ def adaptive_burn_in(self, adaptive_burn_in): self._adaptive_burn_in = t - @property - def adaptive_coverage(self): - return self._adaptive_coverage - - @adaptive_coverage.setter - def adaptive_coverage(self, adaptive_coverage): - if not isinstance(adaptive_coverage, float): - try: - adaptive_coverage = float(adaptive_coverage) - except Exception: - raise ValueError("'adaptive_coverage' must be of type 'float'") - if not 0.0 < adaptive_coverage <= 1.0: - raise ValueError("'adaptive_coverage' must be in the range (0, 1]") - self._adaptive_coverage = adaptive_coverage - @property def adaptive_min_samples(self): return self._adaptive_min_samples diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 83b78e3..96ba81e 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -3217,16 +3217,18 @@ def _rank_windows(self, num_workers): Each pair of neighbouring windows is weighted by sigma * sqrt(g) of its reduced energy differences, which sets the uncertainty of its free - energy difference for a fixed amount of sampling. Pairs are taken in - order of weight until they cover 'adaptive_coverage' of the total. - Pairs with signs of poor sampling are always taken, as are windows with - fewer than 'adaptive_min_samples' independent samples of either the - energy differences or, for a perturbation that breaks or makes a single - bond, the bond distance, and the windows either side of each one taken, - so that exchange carries configurations into and out of them. A window - only stops being hot once it has been found cold at two consecutive - rankings. The hot replicas are padded to a multiple of the number of - workers, so that the GPUs are evenly loaded. + energy difference for a fixed amount of sampling. Pairs are added in + order of weight, and the hot windows chosen are those that minimise the + variance of the free energy multiplied by the cost, i.e. the variance + for a fixed amount of compute. Pairs with signs of poor sampling are + always hot, as are windows with fewer than 'adaptive_min_samples' + independent samples of either the energy differences or, for a + perturbation that breaks or makes a single bond, the bond distance, and + the windows either side of each hot one, so that exchange carries + configurations into and out of them. A window only stops being hot + once it has been found cold at two consecutive rankings. The hot + replicas are padded to a multiple of the number of workers, so that the + GPUs are evenly loaded. The ranking uses all of the data since the burn-in, since the slowest windows decorrelate over far longer than a checkpoint interval. @@ -3291,22 +3293,39 @@ def _rank_windows(self, num_workers): ) or self._is_unconverged(fwd, rev, g_upper[k], g_lower[k + 1]): flagged.append(k) - # Take pairs in order of weight until they cover enough of the total. - hot = set() - covered = 0.0 - for k in _np.argsort(-weights): - if covered >= self._config.adaptive_coverage * weights.sum(): - break - hot.update((int(k), int(k) + 1)) - covered += weights[k] - + forced = set(undersampled) | set(slow_distance) for k in flagged: - hot.update((k, k + 1)) + forced.update((k, k + 1)) + + def pad(windows): + return windows | { + r + d for r in windows for d in (-1, 1) if 0 <= r + d < num_replicas + } - hot.update(undersampled) - hot.update(slow_distance) + # A cold window costs 1 / 'adaptive_mix_interval' of a hot one, and + # multiplies the variance of its pairs by 'adaptive_mix_interval'. + interval = self._config.adaptive_mix_interval + total_variance = (weights**2).sum() or 1.0 - hot.update([r + d for r in hot for d in (-1, 1) if 0 <= r + d < num_replicas]) + def cost_and_variance(windows): + cost = len(windows) + (num_replicas - len(windows)) / interval + variance = sum( + w**2 * (1 if k in windows and k + 1 in windows else interval) + for k, w in enumerate(weights) + ) + return cost / num_replicas, variance / total_variance + + # Add pairs in order of weight, keeping the hot windows that minimise + # the variance for a fixed amount of compute. + candidate = set(forced) + hot = pad(candidate) + best = _np.prod(cost_and_variance(hot)) + for k in _np.argsort(-weights): + candidate.update((int(k), int(k) + 1)) + windows = pad(candidate) + efficiency = _np.prod(cost_and_variance(windows)) + if efficiency < best: + hot, best = windows, efficiency # Keep windows hot until they have been found cold twice in a row. previous = self._dynamics_cache._hot_replicas @@ -3345,6 +3364,11 @@ def _rank_windows(self, num_workers): _logger.info( f"Hot replicas: {hot}" if len(hot) < num_replicas else "All replicas hot" ) + cost, variance = cost_and_variance(set(hot)) + _logger.info( + f"Estimated cost {cost:.2f} and variance {variance:.2f} relative to " + "standard replica exchange" + ) def _check_restart(self): """ diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 813b77c..d23808c 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1451,7 +1451,7 @@ def test_rank_windows(): the largest uncertainty, padded by a window each side, and that windows only go cold after being found cold at two consecutive rankings. """ - runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=10) + runner = _ranking_runner(8, adaptive_min_samples=10) cache = runner._dynamics_cache _fill_ranking_data(runner, hard_pair=3) @@ -1474,7 +1474,7 @@ def test_rank_windows_pads_to_workers(): Validate that the hot replicas are padded to a multiple of the number of workers, and that every replica being hot is recorded as None. """ - runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=10) + runner = _ranking_runner(8, adaptive_min_samples=10) _fill_ranking_data(runner, hard_pair=3) runner._rank_windows(num_workers=3) @@ -1486,6 +1486,25 @@ def test_rank_windows_pads_to_workers(): assert runner._dynamics_cache._hot_replicas is None +def test_rank_windows_minimises_variance_for_compute(): + """ + Validate that a pair with a large uncertainty is made hot without being + flagged, and that pairs which add more cost than they remove variance are + left cold. + """ + runner = _ranking_runner(8, adaptive_min_samples=10) + + _fill_ranking_data(runner, hard_pair=-1) + + # Forward and reverse estimates that agree, so the pair isn't flagged. + rng = np.random.default_rng(11) + runner._adaptive_upper[3] = list(rng.normal(1.0, 1.0, 200)) + runner._adaptive_lower[4] = list(rng.normal(0.0, 1.0, 200)) + runner._rank_windows(num_workers=2) + + assert runner._dynamics_cache._hot_replicas == [2, 3, 4, 5] + + def test_rank_windows_without_uncertainty(): """ Validate that every replica stays hot when no window contributes any @@ -1531,7 +1550,7 @@ def test_rank_windows_needs_enough_samples(): Validate that windows aren't ranked until every one has enough samples, and that the data keeps accumulating until they do. """ - runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=50) + runner = _ranking_runner(8, adaptive_min_samples=50) _fill_ranking_data(runner, hard_pair=3, num_samples=20) runner._rank_windows(num_workers=2) @@ -1545,7 +1564,7 @@ def test_rank_windows_keeps_undersampled_windows_hot(): Validate that a window with too few independent samples stays hot, even though its uncertainty is small. """ - runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=50) + runner = _ranking_runner(8, adaptive_min_samples=50) _fill_ranking_data(runner, hard_pair=3) @@ -1565,7 +1584,7 @@ def test_rank_windows_keeps_slow_bond_distance_hot(): Validate that a window in which the distance of the bond that is broken decorrelates slowly stays hot, even though its energies are converged. """ - runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=50) + runner = _ranking_runner(8, adaptive_min_samples=50) runner._ring_pair = (0, 1) _fill_ranking_data(runner, hard_pair=3) @@ -1659,7 +1678,7 @@ def test_rank_windows_flags_unconverged_pairs(): Validate that a pair whose forward and reverse estimates disagree is made hot, even though its uncertainty is small. """ - runner = _ranking_runner(8, adaptive_coverage=0.5, adaptive_min_samples=10) + runner = _ranking_runner(8, adaptive_min_samples=10) _fill_ranking_data(runner, hard_pair=3) rng = np.random.default_rng(3) From 866a12361ff76e38ab15aa2511afc60ddffee28d Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Oct 2026 20:19:15 +0100 Subject: [PATCH 04/32] Let cold replicas catch up on missed frames rather than forcing full cycles. --- README.md | 3 +- src/somd2/config/_config.py | 5 +- src/somd2/runner/_repex.py | 93 ++++++++++++++++++++----------------- tests/runner/test_repex.py | 47 +++++++++++++++++++ 4 files changed, 102 insertions(+), 46 deletions(-) diff --git a/README.md b/README.md index 0c919f7..788b37e 100644 --- a/README.md +++ b/README.md @@ -427,7 +427,8 @@ little as the simulation is extended. Adaptive replica exchange, enabled with burn-in period of standard replica exchange, set with `--adaptive-burn-in`, the windows are ranked at each checkpoint. Only the highest ranked ("hot") windows are propagated and exchanged every cycle, with every window joining in every -`--adaptive-mix-interval` cycles, as well as on checkpoint and frame cycles. +`--adaptive-mix-interval` cycles, as well as on checkpoint cycles. A window that +misses a frame cycle saves its frame the next time it is propagated. Each pair of neighbouring windows is ranked by σ√g of the reduced energy differences between them, where σ is the standard deviation and g the diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 95b50e9..9cb286a 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -510,8 +510,9 @@ def __init__( highest ranked ("hot") windows are propagated and exchanged every cycle, with the remaining windows joining in every 'adaptive_mix_interval' cycles. The 'runtime' is the simulation time of the hot windows, so the - other windows are sampled for less time. Checkpoint and frame cycles always - include every window. + other windows are sampled for less time. Checkpoint cycles always include + every window, and a window that misses a frame cycle saves its frame the + next time it is propagated. adaptive_mix_interval: int The number of cycles between those that propagate and exchange every diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 96ba81e..545837e 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1645,6 +1645,9 @@ def __init__(self, system, config): self._start_block = 0 self._start_time = _sr.u("0 ps") + if self._config.adaptive_repex: + self._check_adaptive_config() + # Create a terminal flip sampler per replica (if terminal groups were detected). if self._terminal_groups: from ._samplers import TerminalFlipSampler @@ -2122,35 +2125,6 @@ def run(self): if is_adaptive: self._reset_adaptive_data() self._restore_adaptive_data() - - end_time = self._start_time + self._config.runtime - if end_time <= self._config.adaptive_burn_in: - _logger.warning( - "The 'runtime' is no longer than 'adaptive_burn_in', so the " - "windows will never be ranked. Every window will be " - "propagated every cycle." - ) - elif checkpoint_frequency.value() == 0.0: - _logger.warning( - "Adaptive replica exchange ranks the windows at checkpoints, " - "but 'checkpoint_frequency' is zero. Every window will be " - "propagated every cycle." - ) - elif cycles_per_checkpoint <= self._config.adaptive_mix_interval: - _logger.warning( - "Every window is propagated on checkpoint cycles, which are " - "no further apart than 'adaptive_mix_interval'. Adaptive " - "replica exchange will have little effect." - ) - - if ( - self._save_frames or self._config.gcmc - ) and cycles_per_frame <= self._config.adaptive_mix_interval: - _logger.warning( - "Every window is propagated on frame cycles, which are no " - "further apart than 'adaptive_mix_interval'. Adaptive " - "replica exchange will have little effect." - ) else: # Adaptive replica exchange may have been switched off on restart. self._dynamics_cache._hot_replicas = None @@ -2175,32 +2149,23 @@ def run(self): # non-integer ratios between the checkpoint and energy frequencies. is_checkpoint = (i + 1) >= next_checkpoint - 1e-10 - # Whether a frame is saved at the end of the cycle. - write_gcmc_ghosts = (i + 1) % cycles_per_frame == 0 - # Whether the checkpoint files are written at the end of this cycle. do_checkpoint = is_checkpoint or i == cycles - 1 - # Only the hot replicas are propagated, except on mix cycles and - # cycles that write output for every replica. + # Only the hot replicas are propagated, except on mix and checkpoint + # cycles, and frame cycles when frames are held in shared contexts. hot = self._dynamics_cache._hot_replicas is_full = ( hot is None or (i + 1) % self._config.adaptive_mix_interval == 0 or do_checkpoint - or (write_gcmc_ghosts and (self._save_frames or self._config.gcmc)) + or (self._is_cached and (i + 1) % cycles_per_frame == 0) ) active = list(range(num_replicas)) if is_full else hot if not is_full: _logger.info(f"Propagating {len(active)} hot replicas: {active}") - # Frames are scheduled by the cycle rather than by each replica's - # own clock, which falls behind while the replica is cold. - save_frame = ( - write_gcmc_ghosts if is_adaptive and self._save_frames else None - ) - # Capture the simulation clock at the start of the cycle. This seeds # any replica that doesn't yet have a clock of its own. clock = self._dynamics_cache.get_clock() @@ -2226,17 +2191,30 @@ def run(self): for r in replicas ] + # Whether each replica saves a frame and writes its ghost + # waters, scheduled by the cycle since a cold replica's own + # clock falls behind. + is_frame = [ + self._is_due(last_cycle[r], i, cycles_per_frame) + for r in replicas + ] + save_frame = ( + is_frame + if is_adaptive and self._save_frames + else [None] * len(replicas) + ) + try: for result, index, energies in dynamics_executor.map( self._run_block, replicas, repeat(self._lambda_values), is_gcmc, - repeat(write_gcmc_ghosts), + is_frame, is_terminal_flip, repeat(self._config.save_energy_components), repeat(clock), - repeat(save_frame), + save_frame, ): if not result: _logger.error( @@ -2945,6 +2923,35 @@ def _assemble_results(self, results): return matrix + def _check_adaptive_config(self): + """ + Warn about settings that stop adaptive replica exchange having an + effect, before any time is spent on minimisation and equilibration. + """ + checkpoint_frequency = self._config.checkpoint_frequency + end_time = self._start_time + self._config.runtime + + if end_time <= self._config.adaptive_burn_in: + _logger.warning( + "The 'runtime' is no longer than 'adaptive_burn_in', so the " + "windows will never be ranked. Every window will be propagated " + "every cycle." + ) + elif checkpoint_frequency.value() == 0.0: + _logger.warning( + "Adaptive replica exchange ranks the windows at checkpoints, but " + "'checkpoint_frequency' is zero. Every window will be propagated " + "every cycle." + ) + elif ( + checkpoint_frequency / self._config.energy_frequency + ).value() <= self._config.adaptive_mix_interval: + _logger.warning( + "Every window is propagated on checkpoint cycles, which are no " + "further apart than 'adaptive_mix_interval'. Adaptive replica " + "exchange will have little effect." + ) + def _find_ring_pair(self): """ Find the OpenMM particle indices of the bond that is broken or made by diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index d23808c..c394d0b 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1626,6 +1626,53 @@ def test_find_ring_pair(syk_ring_break_mols): assert runner._find_ring_pair() == (1, 7) +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_adaptive_repex_frames_every_cycle(ethane_methanol): + """ + Validate that frequent frames don't stop cold replicas being skipped, and + that a cold replica saves the frame it missed when it is next propagated. + """ + import pandas as pd + import sire as sr + + num_lambda = 4 + hot = [0, 1] + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "48fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "16fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "adaptive_repex": True, + "adaptive_mix_interval": 2, + # Never rank, so that the hot replicas stay as set. + "adaptive_burn_in": "1 ns", + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + runner._dynamics_cache._hot_replicas = hot + runner.run() + + # Cold replicas only run on even cycles, saving a frame each time. + for i in range(num_lambda): + expected = 12 if i in hot else 6 + + records = pd.read_parquet(runner._filenames[i]["energy_traj"]) + assert len(records) == expected, f"replica {i}" + + mols = sr.load( + [runner._filenames["topology0"], runner._filenames[i]["trajectory"]] + ) + assert mols.num_frames() == expected, f"replica {i}" + + @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") def test_adaptive_repex_records_bond_distance(syk_ring_break_mols): """ From 4a1cdc5b822cbb983d607fb756b400a46a433bfa Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Oct 2026 20:29:20 +0100 Subject: [PATCH 05/32] Fix adaptive frame scheduling and startup checks. --- src/somd2/runner/_base.py | 9 +++++++++ src/somd2/runner/_repex.py | 41 ++++++++++++++++++++------------------ tests/runner/test_repex.py | 41 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 19 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index f24c130..cfefe76 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -802,6 +802,15 @@ def __init__(self, system, config): _logger.error(msg) raise ValueError(msg) + # Ghost waters are written by GCMC moves, so frames must coincide + # with them. + ratio = (self._config.frame_frequency / self._config.gcmc_frequency).value() + + if not isclose(ratio, round(ratio), abs_tol=1e-4): + msg = "'frame_frequency' must be a multiple of 'gcmc_frequency'." + _logger.error(msg) + raise ValueError(msg) + # Make sure the checkpoint frequency is a multiple of the frame frequency. # Get the ratio. diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 545837e..89fbd9a 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2153,13 +2153,12 @@ def run(self): do_checkpoint = is_checkpoint or i == cycles - 1 # Only the hot replicas are propagated, except on mix and checkpoint - # cycles, and frame cycles when frames are held in shared contexts. + # cycles. hot = self._dynamics_cache._hot_replicas is_full = ( hot is None or (i + 1) % self._config.adaptive_mix_interval == 0 or do_checkpoint - or (self._is_cached and (i + 1) % cycles_per_frame == 0) ) active = list(range(num_replicas)) if is_full else hot @@ -2198,11 +2197,6 @@ def run(self): self._is_due(last_cycle[r], i, cycles_per_frame) for r in replicas ] - save_frame = ( - is_frame - if is_adaptive and self._save_frames - else [None] * len(replicas) - ) try: for result, index, energies in dynamics_executor.map( @@ -2214,7 +2208,11 @@ def run(self): is_terminal_flip, repeat(self._config.save_energy_components), repeat(clock), - save_frame, + ( + is_frame + if is_adaptive and self._save_frames + else repeat(None) + ), ): if not result: _logger.error( @@ -2928,24 +2926,29 @@ def _check_adaptive_config(self): Warn about settings that stop adaptive replica exchange having an effect, before any time is spent on minimisation and equilibration. """ - checkpoint_frequency = self._config.checkpoint_frequency - end_time = self._start_time + self._config.runtime + from math import floor - if end_time <= self._config.adaptive_burn_in: - _logger.warning( - "The 'runtime' is no longer than 'adaptive_burn_in', so the " - "windows will never be ranked. Every window will be propagated " - "every cycle." - ) - elif checkpoint_frequency.value() == 0.0: + checkpoint = self._config.checkpoint_frequency.to("ps") + burn_in = self._config.adaptive_burn_in.to("ps") + end_time = (self._start_time + self._config.runtime).to("ps") + + if checkpoint == 0.0: _logger.warning( "Adaptive replica exchange ranks the windows at checkpoints, but " "'checkpoint_frequency' is zero. Every window will be propagated " "every cycle." ) + # The first ranking is at the first checkpoint after the burn-in. + elif (floor(burn_in / checkpoint + 1e-9) + 1) * checkpoint >= end_time - 1e-6: + _logger.warning( + "There is no checkpoint after 'adaptive_burn_in' and before the " + "end of the run, so the windows will never be ranked in time to " + "be used. Every window will be propagated every cycle." + ) elif ( - checkpoint_frequency / self._config.energy_frequency - ).value() <= self._config.adaptive_mix_interval: + checkpoint / self._config.energy_frequency.to("ps") + <= self._config.adaptive_mix_interval + ): _logger.warning( "Every window is propagated on checkpoint cycles, which are no " "further apart than 'adaptive_mix_interval'. Adaptive replica " diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index c394d0b..454c2c1 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1486,6 +1486,47 @@ def test_rank_windows_pads_to_workers(): assert runner._dynamics_cache._hot_replicas is None +@pytest.mark.parametrize( + "runtime, checkpoint_frequency, burn_in, warns", + [ + ("5 ns", "100 ps", "500 ps", False), + # The only checkpoint is the end of the run. + ("1 ns", "2 ns", "0 ps", True), + # The first checkpoint after the burn-in is the end of the run. + ("12 ns", "4 ns", "9 ns", True), + ("5 ns", "0 ps", "500 ps", True), + # Checkpoints are no further apart than 'adaptive_mix_interval'. + ("5 ns", "4 ps", "500 ps", True), + ], +) +def test_check_adaptive_config( + monkeypatch, runtime, checkpoint_frequency, burn_in, warns +): + """ + Validate that settings which stop adaptive replica exchange having an + effect are warned about. + """ + import sire as sr + + from somd2.runner import _repex + + warnings = [] + monkeypatch.setattr(_repex._logger, "warning", warnings.append) + + runner = object.__new__(RepexRunner) + runner._config = Config( + runtime=runtime, + energy_frequency="1 ps", + checkpoint_frequency=checkpoint_frequency, + adaptive_burn_in=burn_in, + adaptive_mix_interval=4, + ) + runner._start_time = sr.u("0 ps") + runner._check_adaptive_config() + + assert bool(warnings) == warns + + def test_rank_windows_minimises_variance_for_compute(): """ Validate that a pair with a large uncertainty is made hot without being From 777b80abb60f842e95088051f56fbd95c5dffdac Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Oct 2026 20:38:22 +0100 Subject: [PATCH 06/32] Fix the cost and variance model for choosing hot windows. --- src/somd2/runner/_repex.py | 110 +++++++++++++++++++++---------------- tests/runner/test_repex.py | 17 ++++++ 2 files changed, 80 insertions(+), 47 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 89fbd9a..3cb9694 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -3227,18 +3227,18 @@ def _rank_windows(self, num_workers): Each pair of neighbouring windows is weighted by sigma * sqrt(g) of its reduced energy differences, which sets the uncertainty of its free - energy difference for a fixed amount of sampling. Pairs are added in - order of weight, and the hot windows chosen are those that minimise the - variance of the free energy multiplied by the cost, i.e. the variance - for a fixed amount of compute. Pairs with signs of poor sampling are - always hot, as are windows with fewer than 'adaptive_min_samples' - independent samples of either the energy differences or, for a - perturbation that breaks or makes a single bond, the bond distance, and - the windows either side of each hot one, so that exchange carries - configurations into and out of them. A window only stops being hot - once it has been found cold at two consecutive rankings. The hot - replicas are padded to a multiple of the number of workers, so that the - GPUs are evenly loaded. + energy difference for a fixed amount of sampling. Pairs with signs of + poor sampling are always hot, as are windows with fewer than + 'adaptive_min_samples' independent samples of either the energy + differences or, for a perturbation that breaks or makes a single bond, + the bond distance. The remaining pairs are added in order of weight, + with the windows either side of each so that exchange carries + configurations into and out of them, and of the sets this gives, the + one with the lowest variance of the free energy multiplied by the cost + is chosen. If none is better than every window being hot, every window + is hot. A window only stops being hot once it has been found cold at + two consecutive rankings, and the cold windows with the highest weights + are made hot to fill the workers, so that the GPUs are evenly loaded. The ranking uses all of the data since the burn-in, since the slowest windows decorrelate over far longer than a checkpoint interval. @@ -3289,12 +3289,16 @@ def _rank_windows(self, num_workers): proposed = self._dynamics_cache.get_proposed() - self._adaptive_proposed accepted = self._dynamics_cache.get_accepted() - self._adaptive_accepted - weights = _np.zeros(num_replicas - 1) + # The variance per sample of each pair's free energy, from the samples + # in its lower (forward) and upper (reverse) windows. + fwd_variance = _np.zeros(num_replicas - 1) + rev_variance = _np.zeros(num_replicas - 1) flagged = [] for k in range(num_replicas - 1): fwd = upper[k] rev = lower[k + 1] - weights[k] = _np.sqrt(fwd.var() * g_upper[k] + rev.var() * g_lower[k + 1]) + fwd_variance[k] = fwd.var() * g_upper[k] + rev_variance[k] = rev.var() * g_lower[k + 1] # Flag pairs that rarely exchange, or whose estimates disagree. num_proposed = proposed[k, k + 1] @@ -3303,6 +3307,8 @@ def _rank_windows(self, num_workers): ) or self._is_unconverged(fwd, rev, g_upper[k], g_lower[k + 1]): flagged.append(k) + weights = _np.sqrt(fwd_variance + rev_variance) + forced = set(undersampled) | set(slow_distance) for k in flagged: forced.update((k, k + 1)) @@ -3313,52 +3319,62 @@ def pad(windows): } # A cold window costs 1 / 'adaptive_mix_interval' of a hot one, and - # multiplies the variance of its pairs by 'adaptive_mix_interval'. + # multiplies the variance of its side of each pair by the same factor. interval = self._config.adaptive_mix_interval - total_variance = (weights**2).sum() or 1.0 + total_variance = (fwd_variance + rev_variance).sum() def cost_and_variance(windows): - cost = len(windows) + (num_replicas - len(windows)) / interval + # The hot windows are padded to a multiple of the number of workers. + num_hot = min(num_replicas, ceil(len(windows) / num_workers) * num_workers) + cost = num_hot + (num_replicas - num_hot) / interval variance = sum( - w**2 * (1 if k in windows and k + 1 in windows else interval) - for k, w in enumerate(weights) + fwd_variance[k] * (1 if k in windows else interval) + + rev_variance[k] * (1 if k + 1 in windows else interval) + for k in range(num_replicas - 1) ) - return cost / num_replicas, variance / total_variance + return cost / num_replicas, variance / (total_variance or 1.0) - # Add pairs in order of weight, keeping the hot windows that minimise - # the variance for a fixed amount of compute. + # Add pairs in order of weight, keeping the hot windows that give the + # lowest variance for a fixed amount of compute, if better than all hot. candidate = set(forced) - hot = pad(candidate) - best = _np.prod(cost_and_variance(hot)) - for k in _np.argsort(-weights): - candidate.update((int(k), int(k) + 1)) + hot = None + best = 1.0 - 1e-6 + order = list(_np.argsort(-weights)) if total_variance > 0 else [] + for k in [None] + order: + if k is not None: + candidate.update((int(k), int(k) + 1)) windows = pad(candidate) efficiency = _np.prod(cost_and_variance(windows)) - if efficiency < best: + if windows and efficiency < best: hot, best = windows, efficiency - # Keep windows hot until they have been found cold twice in a row. - previous = self._dynamics_cache._hot_replicas votes = self._dynamics_cache._cold_votes - for r in range(num_replicas): - if r in hot: - votes[r] = 0 - elif previous is not None and r in previous: - votes[r] += 1 - if votes[r] < 2: - hot.add(r) - - # Pad with the highest weighted of the remaining windows. - window_weights = [ - max(weights[k] for k in (r - 1, r) if 0 <= k < num_replicas - 1) - for r in range(num_replicas) - ] - target = min(num_replicas, ceil(len(hot) / num_workers) * num_workers) - cold = sorted(set(range(num_replicas)) - hot, key=lambda r: -window_weights[r]) - hot.update(cold[: target - len(hot)]) + if hot is None: + hot = set(range(num_replicas)) + votes[:] = [0] * num_replicas + else: + # Keep windows hot until they have been found cold twice in a row. + previous = self._dynamics_cache._hot_replicas + for r in range(num_replicas): + if r in hot: + votes[r] = 0 + elif previous is not None and r in previous: + votes[r] += 1 + if votes[r] < 2: + hot.add(r) + + # Fill the workers with the highest weighted of the cold windows. + window_weights = [ + max(weights[k] for k in (r - 1, r) if 0 <= k < num_replicas - 1) + for r in range(num_replicas) + ] + target = min(num_replicas, ceil(len(hot) / num_workers) * num_workers) + cold = sorted( + set(range(num_replicas)) - hot, key=lambda r: -window_weights[r] + ) + hot.update(cold[: target - len(hot)]) - # An empty ranking carries no information, so every replica stays hot. - hot = sorted(hot) if hot else list(range(num_replicas)) + hot = sorted(hot) self._dynamics_cache._hot_replicas = None if len(hot) == num_replicas else hot _logger.info( diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 454c2c1..6f568f6 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1546,6 +1546,23 @@ def test_rank_windows_minimises_variance_for_compute(): assert runner._dynamics_cache._hot_replicas == [2, 3, 4, 5] +def test_rank_windows_without_gain(): + """ + Validate that every replica is made hot straight away when no subset is + better than every replica being hot, rather than keeping the previous hot + replicas for another ranking. + """ + runner = _ranking_runner(8, adaptive_min_samples=10) + runner._dynamics_cache._hot_replicas = [2, 3, 4, 5] + + # Every pair has the same uncertainty. + _fill_ranking_data(runner, hard_pair=-1) + runner._rank_windows(num_workers=2) + + assert runner._dynamics_cache._hot_replicas is None + assert runner._dynamics_cache._cold_votes == [0] * 8 + + def test_rank_windows_without_uncertainty(): """ Validate that every replica stays hot when no window contributes any From 98d999100c811bd8a7488c66c1c9e6baeb763627 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Oct 2026 20:45:20 +0100 Subject: [PATCH 07/32] Score hot window candidates as they will run. --- src/somd2/runner/_repex.py | 79 +++++++++++++++++++++++--------------- tests/runner/test_repex.py | 27 ++++++++++--- 2 files changed, 68 insertions(+), 38 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 3cb9694..519de8b 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -3233,12 +3233,13 @@ def _rank_windows(self, num_workers): differences or, for a perturbation that breaks or makes a single bond, the bond distance. The remaining pairs are added in order of weight, with the windows either side of each so that exchange carries - configurations into and out of them, and of the sets this gives, the - one with the lowest variance of the free energy multiplied by the cost - is chosen. If none is better than every window being hot, every window - is hot. A window only stops being hot once it has been found cold at - two consecutive rankings, and the cold windows with the highest weights - are made hot to fill the workers, so that the GPUs are evenly loaded. + configurations into and out of them. Each set this gives is filled to a + multiple of the number of workers with the cold windows owning the most + variance, so that the GPUs are evenly loaded, and the set with the + lowest variance of the free energy multiplied by the cost is chosen. A + window only stops being hot once it has been found cold at two + consecutive rankings. Every window is hot if no set, including the + windows kept hot, is better than that. The ranking uses all of the data since the burn-in, since the slowest windows decorrelate over far longer than a checkpoint interval. @@ -3318,61 +3319,75 @@ def pad(windows): r + d for r in windows for d in (-1, 1) if 0 <= r + d < num_replicas } + all_hot = set(range(num_replicas)) + + # The variance owned by each window: the forward side of the pair above + # it and the reverse side of the pair below it. + owned = _np.zeros(num_replicas) + owned[:-1] += fwd_variance + owned[1:] += rev_variance + + def fill(windows): + # Make the cold windows owning the most variance hot, until the hot + # windows are a multiple of the number of workers. + target = min(num_replicas, ceil(len(windows) / num_workers) * num_workers) + cold = sorted(all_hot - windows, key=lambda r: -owned[r]) + return windows | set(cold[: target - len(windows)]) + # A cold window costs 1 / 'adaptive_mix_interval' of a hot one, and - # multiplies the variance of its side of each pair by the same factor. + # multiplies the variance it owns by the same factor. interval = self._config.adaptive_mix_interval - total_variance = (fwd_variance + rev_variance).sum() + total_variance = owned.sum() def cost_and_variance(windows): - # The hot windows are padded to a multiple of the number of workers. - num_hot = min(num_replicas, ceil(len(windows) / num_workers) * num_workers) - cost = num_hot + (num_replicas - num_hot) / interval + cost = len(windows) + (num_replicas - len(windows)) / interval variance = sum( - fwd_variance[k] * (1 if k in windows else interval) - + rev_variance[k] * (1 if k + 1 in windows else interval) - for k in range(num_replicas - 1) + v * (1 if r in windows else interval) for r, v in enumerate(owned) ) return cost / num_replicas, variance / (total_variance or 1.0) + baseline = _np.prod(cost_and_variance(all_hot)) + # Add pairs in order of weight, keeping the hot windows that give the # lowest variance for a fixed amount of compute, if better than all hot. candidate = set(forced) - hot = None - best = 1.0 - 1e-6 + chosen = None + best = baseline order = list(_np.argsort(-weights)) if total_variance > 0 else [] for k in [None] + order: if k is not None: + if k in candidate and k + 1 in candidate: + continue candidate.update((int(k), int(k) + 1)) windows = pad(candidate) + if not windows: + continue + windows = fill(windows) efficiency = _np.prod(cost_and_variance(windows)) - if windows and efficiency < best: - hot, best = windows, efficiency + if efficiency < best: + chosen, best = windows, efficiency votes = self._dynamics_cache._cold_votes - if hot is None: - hot = set(range(num_replicas)) + if chosen is None: + hot = all_hot votes[:] = [0] * num_replicas else: # Keep windows hot until they have been found cold twice in a row. previous = self._dynamics_cache._hot_replicas + previous = all_hot if previous is None else set(previous) + hot = set(chosen) for r in range(num_replicas): - if r in hot: + if r in chosen: votes[r] = 0 - elif previous is not None and r in previous: + elif r in previous: votes[r] += 1 if votes[r] < 2: hot.add(r) - # Fill the workers with the highest weighted of the cold windows. - window_weights = [ - max(weights[k] for k in (r - 1, r) if 0 <= k < num_replicas - 1) - for r in range(num_replicas) - ] - target = min(num_replicas, ceil(len(hot) / num_workers) * num_workers) - cold = sorted( - set(range(num_replicas)) - hot, key=lambda r: -window_weights[r] - ) - hot.update(cold[: target - len(hot)]) + # Windows kept hot can leave no gain over all hot. + hot = fill(hot) + if _np.prod(cost_and_variance(hot)) >= baseline: + hot = all_hot hot = sorted(hot) self._dynamics_cache._hot_replicas = None if len(hot) == num_replicas else hot diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 6f568f6..5050c60 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1431,6 +1431,15 @@ def _ranking_runner(num_lambda, **config): return runner +def _rank_twice(runner, num_workers): + """ + Rank the windows at two consecutive checkpoints, so that the windows found + cold at both stop being hot. + """ + runner._rank_windows(num_workers=num_workers) + runner._rank_windows(num_workers=num_workers) + + def _fill_ranking_data(runner, hard_pair, num_samples=200, seed=7): """ Record converged energy differences between every pair of neighbouring @@ -1454,7 +1463,11 @@ def test_rank_windows(): runner = _ranking_runner(8, adaptive_min_samples=10) cache = runner._dynamics_cache + # Starting from every window hot, those found cold once stay hot. _fill_ranking_data(runner, hard_pair=3) + runner._rank_windows(num_workers=2) + assert cache._hot_replicas is None + runner._rank_windows(num_workers=2) assert cache._hot_replicas == [2, 3, 4, 5] @@ -1477,7 +1490,7 @@ def test_rank_windows_pads_to_workers(): runner = _ranking_runner(8, adaptive_min_samples=10) _fill_ranking_data(runner, hard_pair=3) - runner._rank_windows(num_workers=3) + _rank_twice(runner, num_workers=3) assert len(runner._dynamics_cache._hot_replicas) == 6 runner._dynamics_cache._hot_replicas = None @@ -1541,7 +1554,7 @@ def test_rank_windows_minimises_variance_for_compute(): rng = np.random.default_rng(11) runner._adaptive_upper[3] = list(rng.normal(1.0, 1.0, 200)) runner._adaptive_lower[4] = list(rng.normal(0.0, 1.0, 200)) - runner._rank_windows(num_workers=2) + _rank_twice(runner, num_workers=2) assert runner._dynamics_cache._hot_replicas == [2, 3, 4, 5] @@ -1554,6 +1567,7 @@ def test_rank_windows_without_gain(): """ runner = _ranking_runner(8, adaptive_min_samples=10) runner._dynamics_cache._hot_replicas = [2, 3, 4, 5] + runner._dynamics_cache._cold_votes = [1, 1, 0, 0, 0, 0, 1, 1] # Every pair has the same uncertainty. _fill_ranking_data(runner, hard_pair=-1) @@ -1632,7 +1646,7 @@ def test_rank_windows_keeps_undersampled_windows_hot(): for i in range(1, len(y)): y[i] = 0.95 * y[i - 1] + rng.normal(0.0, 0.01) runner._adaptive_upper[6] = list(0.5 + y) - runner._rank_windows(num_workers=2) + _rank_twice(runner, num_workers=2) assert 6 in runner._dynamics_cache._hot_replicas @@ -1655,7 +1669,7 @@ def test_rank_windows_keeps_slow_bond_distance_hot(): for i in range(1, len(y)): y[i] = 0.95 * y[i - 1] + rng.normal() runner._adaptive_distance[7] = list(10.0 + y) - runner._rank_windows(num_workers=2) + _rank_twice(runner, num_workers=2) assert 7 in runner._dynamics_cache._hot_replicas assert 0 not in runner._dynamics_cache._hot_replicas @@ -1788,9 +1802,10 @@ def test_rank_windows_flags_unconverged_pairs(): _fill_ranking_data(runner, hard_pair=3) rng = np.random.default_rng(3) runner._adaptive_lower[1] = list(rng.normal(-3.0, 0.1, 200)) - runner._rank_windows(num_workers=2) + _rank_twice(runner, num_workers=2) - assert {0, 1, 2}.issubset(runner._dynamics_cache._hot_replicas) + hot = runner._dynamics_cache._hot_replicas + assert hot is not None and {0, 1, 2}.issubset(hot) def test_checkpoint_without_adaptive_state(): From 67fad7c553eba6477a497e328bbc22c4d1ba2c1c Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Oct 2026 20:51:23 +0100 Subject: [PATCH 08/32] Fill hot windows once and fix startup and restart checks. --- README.md | 12 +++++++----- src/somd2/runner/_repex.py | 22 ++++++++++++---------- tests/runner/test_repex.py | 21 +++++++++++++++++++++ 3 files changed, 40 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 788b37e..9a00486 100644 --- a/README.md +++ b/README.md @@ -450,11 +450,13 @@ ones. The following are always hot: The windows either side of each hot window are also made hot, so that exchange carries configurations into and out of them, and a window only stops being hot -after it has been found cold at two consecutive checkpoints. The number of hot -windows is rounded up to a multiple of the number of GPUs multiplied by the -oversubscription factor, so that the GPUs are evenly loaded. When contexts are -re-used, replicas are assigned to them in a strided rather than contiguous -pattern, so that the hot windows are spread evenly over the GPUs. +after it has been found cold at two consecutive checkpoints, so no window is +cold until the second checkpoint after the burn-in. The number of hot windows +is rounded up to a multiple of the number of GPUs multiplied by the +oversubscription factor, using the cold windows that contribute the most +uncertainty, so that the GPUs are evenly loaded. When contexts are re-used, +replicas are assigned to them in a strided rather than contiguous pattern, so +that the hot windows are spread evenly over the GPUs. The `--runtime` is the simulation time of the hot windows, so the remaining windows are sampled for less time and the windows have different numbers of diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 519de8b..50363e8 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2128,6 +2128,7 @@ def run(self): else: # Adaptive replica exchange may have been switched off on restart. self._dynamics_cache._hot_replicas = None + self._dynamics_cache._cold_votes = [0] * num_replicas # Perform the replica exchange simulation. for i in range(cycles): @@ -2938,12 +2939,12 @@ def _check_adaptive_config(self): "'checkpoint_frequency' is zero. Every window will be propagated " "every cycle." ) - # The first ranking is at the first checkpoint after the burn-in. - elif (floor(burn_in / checkpoint + 1e-9) + 1) * checkpoint >= end_time - 1e-6: + # Windows first go cold at the second checkpoint after the burn-in. + elif (floor(burn_in / checkpoint + 1e-9) + 2) * checkpoint >= end_time - 1e-6: _logger.warning( - "There is no checkpoint after 'adaptive_burn_in' and before the " - "end of the run, so the windows will never be ranked in time to " - "be used. Every window will be propagated every cycle." + "There are fewer than two checkpoints after 'adaptive_burn_in' " + "and before the end of the run, so no window will stop being " + "hot. Every window will be propagated every cycle." ) elif ( checkpoint / self._config.energy_frequency.to("ps") @@ -3346,13 +3347,15 @@ def cost_and_variance(windows): ) return cost / num_replicas, variance / (total_variance or 1.0) - baseline = _np.prod(cost_and_variance(all_hot)) + # A set must beat all hot by more than rounding error. + threshold = _np.prod(cost_and_variance(all_hot)) * (1 - 1e-9) # Add pairs in order of weight, keeping the hot windows that give the # lowest variance for a fixed amount of compute, if better than all hot. + # Sets are scored with the windows that will fill the workers. candidate = set(forced) chosen = None - best = baseline + best = threshold order = list(_np.argsort(-weights)) if total_variance > 0 else [] for k in [None] + order: if k is not None: @@ -3362,8 +3365,7 @@ def cost_and_variance(windows): windows = pad(candidate) if not windows: continue - windows = fill(windows) - efficiency = _np.prod(cost_and_variance(windows)) + efficiency = _np.prod(cost_and_variance(fill(windows))) if efficiency < best: chosen, best = windows, efficiency @@ -3386,7 +3388,7 @@ def cost_and_variance(windows): # Windows kept hot can leave no gain over all hot. hot = fill(hot) - if _np.prod(cost_and_variance(hot)) >= baseline: + if _np.prod(cost_and_variance(hot)) >= threshold: hot = all_hot hot = sorted(hot) diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 5050c60..2e29187 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1507,6 +1507,10 @@ def test_rank_windows_pads_to_workers(): ("1 ns", "2 ns", "0 ps", True), # The first checkpoint after the burn-in is the end of the run. ("12 ns", "4 ns", "9 ns", True), + # The second checkpoint after the burn-in is the end of the run, so + # no window goes cold. + ("1 ns", "500 ps", "400 ps", True), + ("1.2 ns", "500 ps", "400 ps", False), ("5 ns", "0 ps", "500 ps", True), # Checkpoints are no further apart than 'adaptive_mix_interval'. ("5 ns", "4 ps", "500 ps", True), @@ -1559,6 +1563,23 @@ def test_rank_windows_minimises_variance_for_compute(): assert runner._dynamics_cache._hot_replicas == [2, 3, 4, 5] +def test_rank_windows_fill_windows_stay_cold_candidates(): + """ + Validate that a window made hot only to fill the workers keeps counting + as found cold, rather than as chosen. + """ + runner = _ranking_runner(8, adaptive_min_samples=10) + + # The flagged pair 0 gives windows 0, 1 and 2, so one window is filled. + _fill_ranking_data(runner, hard_pair=0) + _rank_twice(runner, num_workers=2) + + hot = runner._dynamics_cache._hot_replicas + assert hot is not None and len(hot) == 4 + (filler,) = set(hot) - {0, 1, 2} + assert runner._dynamics_cache._cold_votes[filler] == 2 + + def test_rank_windows_without_gain(): """ Validate that every replica is made hot straight away when no subset is From 874fca1e4eff4da685bd44c683eb8f2fb1d80eb2 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 09:22:21 +0100 Subject: [PATCH 09/32] Show adaptive replica exchange status in the viewer. --- src/somd2/_viewer/_data.py | 84 +++++++++++++++++-- src/somd2/_viewer/_static/index.html | 56 +++++++++++++ tests/viewer/test_adaptive.py | 119 +++++++++++++++++++++++++++ 3 files changed, 253 insertions(+), 6 deletions(-) create mode 100644 tests/viewer/test_adaptive.py diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index 073531a..c9d826e 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -94,6 +94,8 @@ def _run(self): _block_re = _re.compile(r"^Finished block (\d+) of (\d+) for \S+ = ([\d.]+)") _perf_re = _re.compile(r"^Overall performance: ([\d.]+) ns day-1") _correction_re = _re.compile(r"Standard state correction: (-?[\d.]+) kcal mol-1") +_hot_re = _re.compile(r"^Hot replicas: \[([\d, ]*)\]") +_estimate_re = _re.compile(r"^Estimated cost ([\d.]+) and variance ([\d.]+)") # Older versions don't log the time. _waters_re = _re.compile( r"^Current number of waters in GCMC volume at \S+ = ([\d.]+)" @@ -178,6 +180,13 @@ def _checkpoint_lambda(path): return None +def _is_adaptive(config): + """ + Whether a run uses adaptive replica exchange. + """ + return bool(config.get("replica_exchange") and config.get("adaptive_repex")) + + def _to_ns(value): """ Convert a time string from the config, e.g. '2 ps', to nanoseconds. @@ -251,6 +260,10 @@ def _reset(self): self._pending_states = None self.states = [] self._saved_states = 0 + # Adaptive replica exchange rankings, as [cycle, hot or None, cost, + # variance], and the cycles that propagated only the hot replicas. + self.rankings = [] + self.hot_only_cycles = [] self.standard_state_correction = None self.messages = _deque(maxlen=20) # The number of waters in the GCMC region at each checkpoint, by λ, as @@ -306,8 +319,13 @@ def _parse(self, line): # A restart resumes from the last checkpoint, so the permutations # logged after it were discarded. del self.states[self._saved_states :] + saved = self._saved_states + self.rankings = [r for r in self.rankings if r[0] <= saved] + self.hot_only_cycles = [c for c in self.hot_only_cycles if c < saved] self._new_session(_parse_timestamp(stamp)) - elif message.startswith("Saving replica exchange state"): + elif message.startswith( + ("Saving replica exchange state", "Saving final replica exchange state") + ): self._saved_states = len(self.states) elif message.startswith("States: ["): if "]" in message: @@ -334,6 +352,16 @@ def _parse(self, line): self.finished = True elif m := _correction_re.search(message): self.standard_state_correction = float(m.group(1)) + elif message.startswith("Propagating") and "hot replicas" in message: + # Logged after the cycle's states, so this is the latest cycle. + self.hot_only_cycles.append(len(self.states) - 1) + elif m := _hot_re.match(message): + hot = [int(x) for x in m.group(1).split(",")] + self.rankings.append([len(self.states), hot, None, None]) + elif message.startswith("All replicas hot"): + self.rankings.append([len(self.states), None, None, None]) + elif (m := _estimate_re.match(message)) and self.rankings: + self.rankings[-1][2:] = [float(m.group(1)), float(m.group(2))] elif m := _waters_re.match(message): time = float(m.group(2)) if m.group(2) is not None else None counts = self.gcmc_waters.setdefault(float(m.group(1)), []) @@ -519,9 +547,7 @@ def _progress(self, config, lambda_values, windows, log): # Adaptive replica exchange samples cold windows less, so progress is # the simulation time of the hot windows rather than the samples. - is_adaptive = bool( - config.get("replica_exchange") and config.get("adaptive_repex") - ) + is_adaptive = _is_adaptive(config) if is_adaptive: state = self._repex_state() time = state.get("time_ns") if state else None @@ -902,6 +928,7 @@ def load(): "terminal_flip_stats": getattr( state, "_terminal_flip_stats", None ), + "hot_replicas": getattr(state, "_hot_replicas", None), } except Exception as e: error = error or e @@ -942,9 +969,12 @@ def _correction(windows, log): return float(corrections.pop()) return log.standard_state_correction if log is not None else None - def _repex(self, log, lambda_values, repex_state): + def _repex(self, log, lambda_values, repex_state, is_adaptive=False): result = {"transition_matrix": self._transition_matrix()} + if is_adaptive: + result["adaptive"] = self._adaptive(log, len(lambda_values), repex_state) + if repex_state is not None: proposed = repex_state["proposed"] accepted = repex_state["accepted"] @@ -960,6 +990,43 @@ def _repex(self, log, lambda_values, repex_state): return result + @staticmethod + def _adaptive(log, num_states, repex_state, max_rankings=200): + """ + Summarise adaptive replica exchange: the current hot replicas, the + estimates from the latest ranking, the share of cycles that only + propagated the hot replicas, and the hot replicas over time. + """ + result = {} + + # The state holds the hot replicas in use, which a restart keeps. + if repex_state is not None: + hot = repex_state["hot_replicas"] + result["hot"] = list(range(num_states)) if hot is None else list(hot) + + if log is None: + return result + + if log.rankings: + cycle, _, cost, variance = log.rankings[-1] + result.update({"ranked_cycle": cycle, "cost": cost, "variance": variance}) + + # Keep the latest ranking when thinning a long history. + stride = max(1, -(-len(log.rankings) // max_rankings)) + rankings = log.rankings[::-1][::stride][::-1] + result["history"] = { + "cycle": [r[0] for r in rankings], + "hot": [ + [int(r[1] is None or k in r[1]) for r in rankings] + for k in range(num_states) + ], + } + + if log.states: + result["hot_only_fraction"] = len(log.hot_only_cycles) / len(log.states) + + return result + @staticmethod def _replica_paths(states, num_states, max_points=500): """ @@ -1759,7 +1826,12 @@ def _summary(self): "repex": ( self._section( "summarise the replica exchange", - lambda: self._repex(log, lambda_values, repex_state), + lambda: self._repex( + log, + lambda_values, + repex_state, + is_adaptive=_is_adaptive(config), + ), ) if is_repex else None diff --git a/src/somd2/_viewer/_static/index.html b/src/somd2/_viewer/_static/index.html index 555dcb6..28a6afa 100644 --- a/src/somd2/_viewer/_static/index.html +++ b/src/somd2/_viewer/_static/index.html @@ -1422,10 +1422,64 @@

SOMD2 Viewer

html += `

Round trips

${total} complete round trips (${LAMBDA} = 0 → 1 → 0) across all replicas, from ${r.num_cycles} logged cycles.
`; } html += ``; + if (r.adaptive) html += renderAdaptive(s, r.adaptive); if (r.replica_states) html += `

Replica state trajectories

`; return card("Replica exchange", html); } +function renderAdaptive(s, a) { + const n = s.lambda_values.length; + let html = `

Adaptive replica exchange

`; + if (a.hot) { + const hot = a.hot.length === n ? "all" : a.hot.map((k) => fmtLambda(s.lambda_values[k])).join(", "); + html += stat("Hot windows", `${a.hot.length} of ${n}`); + html += `
Hot ${LAMBDA}: ${esc(hot)}
`; + } + if (a.cost !== undefined && a.cost !== null) { + html += stat("Estimated cost", `${fmt(a.cost, 2)}`); + html += stat("Estimated variance", `${fmt(a.variance, 2)}`); + } + if (a.hot_only_fraction !== undefined) html += stat("Hot-only cycles", `${fmt(100 * a.hot_only_fraction, 1)}%`); + html += `
`; + if (a.cost !== undefined && a.cost !== null) { + html += `
Estimates from the ranking at cycle ${a.ranked_cycle}, relative to standard replica exchange: cost is the compute per cycle and variance that of the free energy for the same simulation time.
`; + } else if (a.ranked_cycle === undefined && (!a.hot || a.hot.length === n)) { + html += `
The windows haven't been ranked yet, so every window is hot until the burn-in is complete.
`; + } + if (a.history) html += `

Hot windows at each ranking

`; + return html; +} + +// Which windows were hot (filled) at each ranking, by cycle. +function hotChart(el, history, numStates) { + const cols = history.cycle.length; + const pad = 34; + const width = Math.max(240, el.clientWidth || 520); + const cellW = (width - pad) / cols; + const cellH = Math.max(6, Math.min(16, 240 / numStates)); + const height = pad + cellH * numStates; + const rgb = css("--accent-rgb"), border = css("--border"); + let svg = ``; + for (let k = 0; k < numStates; k++) { + for (let j = 0; j < cols; j++) { + const hot = history.hot[k][j]; + svg += `window ${k}, cycle ${history.cycle[j]}: ${hot ? "hot" : "cold"}`; + } + } + const rowStep = Math.ceil(numStates / Math.floor((height - pad) / 14)); + for (let k = 0; k < numStates; k += rowStep) { + svg += `${k}`; + } + const colStep = Math.ceil(cols / Math.floor((width - pad) / 60)); + for (let j = 0; j < cols; j += colStep) { + svg += `${history.cycle[j]}`; + } + svg += ""; + el.innerHTML = svg + SAVE_BUTTONS; + el.chartDraw = (target) => hotChart(target, history, numStates); + el.chartLegend = []; +} + function componentsLambda(s) { const chosen = state.componentsLambda[s.id]; if (chosen !== undefined && s.lambda_values.includes(chosen)) return chosen; @@ -1614,6 +1668,8 @@

SOMD2 Viewer

xLabel: "replica", yLabel: "round trips", xFormat: (v) => String(Math.round(v)), }); } + const hc = $("#hot-chart"); + if (hc && r.adaptive && r.adaptive.history) hotChart(hc, r.adaptive.history, s.lambda_values.length); const rc = $("#replica-chart"); if (rc && r.replica_states) { const n = r.replica_states.length; diff --git a/tests/viewer/test_adaptive.py b/tests/viewer/test_adaptive.py new file mode 100644 index 0000000..a29500a --- /dev/null +++ b/tests/viewer/test_adaptive.py @@ -0,0 +1,119 @@ +from somd2._viewer._data import Simulation, _LogMonitor + + +def _line(message, level="INFO"): + return ( + f"2026-10-13 21:00:00.000 | {level:<8} | somd2.runner._repex:run:1 - {message}" + ) + + +def _cycle(n, total, hot=None): + lines = [ + _line(f"Running dynamics for cycle {n} of {total}"), + _line("States: [0 1 2]"), + ] + if hot is not None: + lines.append(_line(f"Propagating {len(hot)} hot replicas: {hot}")) + return lines + + +def _monitor(tmp_path, lines): + path = tmp_path / "log.txt" + path.write_text("\n".join(lines) + "\n") + log = _LogMonitor(path) + log.update() + return log + + +def test_adaptive_summary(tmp_path): + """ + Validate that the rankings and hot-only cycles of an adaptive replica + exchange run are read from the log and summarised. + """ + lines = ( + _cycle(1, 4) + + [ + _line("Hot replicas: [1, 2]"), + _line("Estimated cost 0.75 and variance 1.10 relative to standard"), + ] + + _cycle(2, 4, hot=[1, 2]) + + _cycle(3, 4) + + [ + _line("All replicas hot"), + _line("Estimated cost 1.00 and variance 1.00 relative to standard"), + ] + + _cycle(4, 4) + ) + log = _monitor(tmp_path, lines) + + summary = Simulation._adaptive(log, 3, {"hot_replicas": None}) + + assert summary["hot"] == [0, 1, 2] + assert summary["ranked_cycle"] == 3 + assert summary["cost"] == 1.0 and summary["variance"] == 1.0 + assert summary["hot_only_fraction"] == 0.25 + assert summary["history"] == {"cycle": [1, 3], "hot": [[0, 1], [1, 1], [1, 1]]} + + +def test_adaptive_rankings_after_restart(tmp_path): + """ + Validate that rankings and hot-only cycles logged after the last saved + state are discarded on restart, since the run resumes from that state. + """ + lines = ( + _cycle(1, 3) + + [_line("Saving replica exchange state")] + + _cycle(2, 3, hot=[1, 2]) + + [ + _line("Hot replicas: [1, 2]"), + _line("Estimated cost 0.75 and variance 1.10 relative to standard"), + _line("somd2 version: test"), + ] + + _cycle(2, 3) + ) + log = _monitor(tmp_path, lines) + + assert log.rankings == [] + assert log.hot_only_cycles == [] + + +def test_adaptive_rankings_kept_after_final_save(tmp_path): + """ + Validate that a final save, made when the last cycle isn't a checkpoint, + keeps the rankings and hot-only cycles before it when the run is extended. + """ + lines = ( + _cycle(1, 2) + + [ + _line("Hot replicas: [1, 2]"), + _line("Estimated cost 0.75 and variance 1.10 relative to standard"), + ] + + _cycle(2, 2, hot=[1, 2]) + + [ + _line("Saving final replica exchange state"), + _line("somd2 version: test"), + ] + + _cycle(1, 1) + ) + log = _monitor(tmp_path, lines) + + assert len(log.rankings) == 1 + assert log.hot_only_cycles == [1] + + +def test_adaptive_history_is_thinned(tmp_path): + """ + Validate that a long ranking history is thinned, keeping the latest. + """ + lines = [] + for n in range(1, 11): + lines += _cycle(n, 10) + [ + _line("Hot replicas: [1, 2]"), + _line("Estimated cost 0.75 and variance 1.10 relative to standard"), + ] + log = _monitor(tmp_path, lines) + + summary = Simulation._adaptive(log, 3, None, max_rankings=4) + + assert len(summary["history"]["cycle"]) <= 4 + assert summary["history"]["cycle"][-1] == 10 From a1964001203b374d995a45dd997a88aa0ffa4b44 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 09:30:31 +0100 Subject: [PATCH 10/32] Label the adaptive hot window chart. --- src/somd2/_viewer/_static/index.html | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/src/somd2/_viewer/_static/index.html b/src/somd2/_viewer/_static/index.html index 28a6afa..584d39d 100644 --- a/src/somd2/_viewer/_static/index.html +++ b/src/somd2/_viewer/_static/index.html @@ -1446,34 +1446,38 @@

SOMD2 Viewer

} else if (a.ranked_cycle === undefined && (!a.hot || a.hot.length === n)) { html += `
The windows haven't been ranked yet, so every window is hot until the burn-in is complete.
`; } - if (a.history) html += `

Hot windows at each ranking

`; + if (a.history) html += `

Hot windows at each ranking

Filled cells are hot windows, propagated every cycle. The others are cold, propagated every adaptive mix interval.
`; return html; } // Which windows were hot (filled) at each ranking, by cycle. function hotChart(el, history, numStates) { const cols = history.cycle.length; - const pad = 34; - const width = Math.max(240, el.clientWidth || 520); - const cellW = (width - pad) / cols; + const m = { l: 54, r: 14, t: 10, b: 40 }; + const width = Math.max(280, el.clientWidth || 560); + const pw = width - m.l - m.r; + const cellW = pw / cols; const cellH = Math.max(6, Math.min(16, 240 / numStates)); - const height = pad + cellH * numStates; + const ph = cellH * numStates; + const height = m.t + ph + m.b; const rgb = css("--accent-rgb"), border = css("--border"); - let svg = ``; + let svg = ``; for (let k = 0; k < numStates; k++) { for (let j = 0; j < cols; j++) { const hot = history.hot[k][j]; - svg += `window ${k}, cycle ${history.cycle[j]}: ${hot ? "hot" : "cold"}`; + svg += `window ${k}, cycle ${history.cycle[j]}: ${hot ? "hot" : "cold"}`; } } - const rowStep = Math.ceil(numStates / Math.floor((height - pad) / 14)); + const rowStep = Math.ceil(numStates / Math.max(1, Math.floor(ph / 14))); for (let k = 0; k < numStates; k += rowStep) { - svg += `${k}`; + svg += `${k}`; } - const colStep = Math.ceil(cols / Math.floor((width - pad) / 60)); + const colStep = Math.ceil(cols / Math.max(1, Math.floor(pw / 60))); for (let j = 0; j < cols; j += colStep) { - svg += `${history.cycle[j]}`; + svg += `${history.cycle[j]}`; } + svg += `cycle`; + svg += `window`; svg += ""; el.innerHTML = svg + SAVE_BUTTONS; el.chartDraw = (target) => hotChart(target, history, numStates); From 877b443afb3b42b0fe7947c959642d999ecc8bc1 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 09:39:47 +0100 Subject: [PATCH 11/32] Rank adaptive windows by final variance given samples collected. --- README.md | 7 ++-- src/somd2/runner/_repex.py | 72 ++++++++++++++++++++++++++++++-------- tests/runner/test_repex.py | 58 ++++++++++++++++++++++++------ 3 files changed, 109 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 9a00486..b8b32b1 100644 --- a/README.md +++ b/README.md @@ -434,9 +434,10 @@ Each pair of neighbouring windows is ranked by σ√g of the reduced energy differences between them, where σ is the standard deviation and g the statistical inefficiency, which sets the uncertainty of its free energy difference. Pairs are made hot in order of rank for as long as doing so reduces -the variance of the free energy for a fixed amount of compute, accounting for -cold windows costing and sampling `--adaptive-mix-interval` times less than hot -ones. The following are always hot: +the variance of the final free energy for the compute of the whole run, +accounting for the samples each window has already collected and for cold +windows costing and sampling `--adaptive-mix-interval` times less than hot ones. +The following are always hot: - pairs whose forward and reverse free energy estimates disagree, or whose estimate drifts during the simulation, diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 50363e8..7136b21 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2300,7 +2300,9 @@ def run(self): # Rank the windows before the state is saved, so that a restart # carries on with the same hot replicas. if is_adaptive and is_checkpoint and is_burnt_in: - self._rank_windows(num_workers) + self._rank_windows( + num_workers, self._window_samples(), cycles - (i + 1) + ) # Create the lock. lock = _FileLock(self._lock_file) @@ -2922,6 +2924,24 @@ def _assemble_results(self, results): return matrix + def _window_samples(self): + """ + The number of energy samples each window has collected, from its clock. + + Returns + ------- + + list of int + The number of samples for each window. + """ + energy_frequency = self._config.energy_frequency.to("ps") + return [ + 0 + if clock is None + else round(clock["_current_time"].to("ps") / energy_frequency) + for clock in self._dynamics_cache._clocks + ] + def _check_adaptive_config(self): """ Warn about settings that stop adaptive replica exchange having an @@ -3221,7 +3241,7 @@ def _is_unconverged(cls, fwd, rev, g_fwd, g_rev): drift = abs(estimates[0] - estimates[1]) return drift > 0.25 and drift > 2.0 * _np.hypot(*errors) - def _rank_windows(self, num_workers): + def _rank_windows(self, num_workers, window_samples, cycles_left): """ Rank the windows by the uncertainty they contribute to the free energy and choose the hot replicas for adaptive replica exchange. @@ -3236,11 +3256,13 @@ def _rank_windows(self, num_workers): with the windows either side of each so that exchange carries configurations into and out of them. Each set this gives is filled to a multiple of the number of workers with the cold windows owning the most - variance, so that the GPUs are evenly loaded, and the set with the - lowest variance of the free energy multiplied by the cost is chosen. A - window only stops being hot once it has been found cold at two - consecutive rankings. Every window is hot if no set, including the - windows kept hot, is better than that. + variance, so that the GPUs are evenly loaded, and scored by the + variance of the final free energy, given the samples each window has + and will collect in the cycles left, multiplied by the compute of the + whole run. The set with the lowest score is chosen. A window only stops + being hot once it has been found cold at two consecutive rankings. + Every window is hot if no set, including the windows kept hot, is + better than that. The ranking uses all of the data since the burn-in, since the slowest windows decorrelate over far longer than a checkpoint interval. @@ -3250,6 +3272,12 @@ def _rank_windows(self, num_workers): num_workers: int The number of replicas that are propagated concurrently. + + window_samples: list of int + The number of energy samples each window has collected. + + cycles_left: int + The number of cycles left in the run. """ from math import ceil @@ -3335,17 +3363,26 @@ def fill(windows): cold = sorted(all_hot - windows, key=lambda r: -owned[r]) return windows | set(cold[: target - len(windows)]) - # A cold window costs 1 / 'adaptive_mix_interval' of a hot one, and - # multiplies the variance it owns by the same factor. + # A cold window is propagated every 'adaptive_mix_interval' cycles. interval = self._config.adaptive_mix_interval total_variance = owned.sum() + samples = _np.maximum(_np.asarray(window_samples, dtype=float), 1.0) - def cost_and_variance(windows): - cost = len(windows) + (num_replicas - len(windows)) / interval - variance = sum( - v * (1 if r in windows else interval) for r, v in enumerate(owned) + def rates(windows): + return _np.array( + [1.0 if r in windows else 1.0 / interval for r in range(num_replicas)] ) - return cost / num_replicas, variance / (total_variance or 1.0) + + def cost_and_variance(windows): + # The compute of the whole run, in blocks of dynamics, and the + # variance of the final free energy, from the samples at its end. + rate = rates(windows) + final = samples + rate * cycles_left + cost = samples.sum() + rate.sum() * cycles_left + variance = (fwd_variance / final[:-1]).sum() + ( + rev_variance / final[1:] + ).sum() + return cost, variance # A set must beat all hot by more than rounding error. threshold = _np.prod(cost_and_variance(all_hot)) * (1 - 1e-9) @@ -3407,7 +3444,12 @@ def cost_and_variance(windows): _logger.info( f"Hot replicas: {hot}" if len(hot) < num_replicas else "All replicas hot" ) - cost, variance = cost_and_variance(set(hot)) + # The cost per cycle from now on, and the final variance, relative to + # every window being hot for the rest of the run. + cost = rates(set(hot)).sum() / num_replicas + variance = cost_and_variance(set(hot))[1] / ( + cost_and_variance(all_hot)[1] or 1.0 + ) _logger.info( f"Estimated cost {cost:.2f} and variance {variance:.2f} relative to " "standard replica exchange" diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 2e29187..644359c 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1431,13 +1431,23 @@ def _ranking_runner(num_lambda, **config): return runner +def _rank(runner, num_workers, window_samples=None, cycles_left=10**6): + """ + Rank the windows. By default every window has the same samples and the run + is far from its end, so that only the sampling rates matter. + """ + if window_samples is None: + window_samples = [200] * runner._config.num_lambda + runner._rank_windows(num_workers, window_samples, cycles_left) + + def _rank_twice(runner, num_workers): """ Rank the windows at two consecutive checkpoints, so that the windows found cold at both stop being hot. """ - runner._rank_windows(num_workers=num_workers) - runner._rank_windows(num_workers=num_workers) + _rank(runner, num_workers=num_workers) + _rank(runner, num_workers=num_workers) def _fill_ranking_data(runner, hard_pair, num_samples=200, seed=7): @@ -1465,20 +1475,20 @@ def test_rank_windows(): # Starting from every window hot, those found cold once stay hot. _fill_ranking_data(runner, hard_pair=3) - runner._rank_windows(num_workers=2) + _rank(runner, num_workers=2) assert cache._hot_replicas is None - runner._rank_windows(num_workers=2) + _rank(runner, num_workers=2) assert cache._hot_replicas == [2, 3, 4, 5] # Windows 2 and 3 are found cold once, so stay hot. _fill_ranking_data(runner, hard_pair=5) - runner._rank_windows(num_workers=2) + _rank(runner, num_workers=2) assert cache._hot_replicas == [2, 3, 4, 5, 6, 7] # Found cold a second time. _fill_ranking_data(runner, hard_pair=5) - runner._rank_windows(num_workers=2) + _rank(runner, num_workers=2) assert cache._hot_replicas == [4, 5, 6, 7] @@ -1495,7 +1505,7 @@ def test_rank_windows_pads_to_workers(): runner._dynamics_cache._hot_replicas = None _fill_ranking_data(runner, hard_pair=3) - runner._rank_windows(num_workers=8) + _rank(runner, num_workers=8) assert runner._dynamics_cache._hot_replicas is None @@ -1580,6 +1590,34 @@ def test_rank_windows_fill_windows_stay_cold_candidates(): assert runner._dynamics_cache._cold_votes[filler] == 2 +def test_rank_windows_accounts_for_samples_collected(): + """ + Validate that, of two pairs with the same uncertainty per sample, only the + one with fewer samples so far is made hot, since more sampling of the + other would barely change the final free energy. + """ + + def rank(samples): + runner = _ranking_runner(8, adaptive_min_samples=10) + _fill_ranking_data(runner, hard_pair=-1) + + # Forward and reverse estimates that agree, so neither pair is flagged. + rng = np.random.default_rng(11) + for k in (1, 5): + runner._adaptive_upper[k] = list(rng.normal(1.0, 1.0, 200)) + runner._adaptive_lower[k + 1] = list(rng.normal(0.0, 1.0, 200)) + + for _ in range(2): + _rank(runner, 2, window_samples=samples, cycles_left=1000) + return runner._dynamics_cache._hot_replicas + + # With the same samples, both pairs are worth sampling, i.e. every window. + assert rank([1000] * 8) is None + + # Pair 1 has been sampled more, so only pair 5 is. + assert rank([1000, 4000, 4000, 1000, 1000, 1000, 1000, 1000]) == [4, 5, 6, 7] + + def test_rank_windows_without_gain(): """ Validate that every replica is made hot straight away when no subset is @@ -1592,7 +1630,7 @@ def test_rank_windows_without_gain(): # Every pair has the same uncertainty. _fill_ranking_data(runner, hard_pair=-1) - runner._rank_windows(num_workers=2) + _rank(runner, num_workers=2) assert runner._dynamics_cache._hot_replicas is None assert runner._dynamics_cache._cold_votes == [0] * 8 @@ -1608,7 +1646,7 @@ def test_rank_windows_without_uncertainty(): for k in range(3): runner._adaptive_upper[k] = [0.5] * 100 runner._adaptive_lower[k + 1] = [-0.5] * 100 - runner._rank_windows(num_workers=1) + _rank(runner, num_workers=1) assert runner._dynamics_cache._hot_replicas is None @@ -1646,7 +1684,7 @@ def test_rank_windows_needs_enough_samples(): runner = _ranking_runner(8, adaptive_min_samples=50) _fill_ranking_data(runner, hard_pair=3, num_samples=20) - runner._rank_windows(num_workers=2) + _rank(runner, num_workers=2) assert runner._dynamics_cache._hot_replicas is None assert len(runner._adaptive_upper[0]) == 20 From 58f70f15d146f42a3beed0037ad5f665cae137af Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 09:46:08 +0100 Subject: [PATCH 12/32] Fix the adaptive cold rate, end-of-run ranking and estimate logging. --- src/somd2/_viewer/_data.py | 4 +- src/somd2/_viewer/_static/index.html | 6 +- src/somd2/runner/_repex.py | 94 +++++++++++++++++++++------- tests/runner/test_repex.py | 40 +++++++++++- tests/viewer/test_adaptive.py | 5 +- 5 files changed, 119 insertions(+), 30 deletions(-) diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index c9d826e..519057e 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -95,7 +95,9 @@ def _run(self): _perf_re = _re.compile(r"^Overall performance: ([\d.]+) ns day-1") _correction_re = _re.compile(r"Standard state correction: (-?[\d.]+) kcal mol-1") _hot_re = _re.compile(r"^Hot replicas: \[([\d, ]*)\]") -_estimate_re = _re.compile(r"^Estimated cost ([\d.]+) and variance ([\d.]+)") +_estimate_re = _re.compile( + r"^Estimated cost (?:per cycle )?([\d.]+) and (?:final )?variance ([\d.]+)" +) # Older versions don't log the time. _waters_re = _re.compile( r"^Current number of waters in GCMC volume at \S+ = ([\d.]+)" diff --git a/src/somd2/_viewer/_static/index.html b/src/somd2/_viewer/_static/index.html index 584d39d..0bb9a6f 100644 --- a/src/somd2/_viewer/_static/index.html +++ b/src/somd2/_viewer/_static/index.html @@ -1436,13 +1436,13 @@

SOMD2 Viewer

html += `
Hot ${LAMBDA}: ${esc(hot)}
`; } if (a.cost !== undefined && a.cost !== null) { - html += stat("Estimated cost", `${fmt(a.cost, 2)}`); - html += stat("Estimated variance", `${fmt(a.variance, 2)}`); + html += stat("Estimated cost per cycle", `${fmt(a.cost, 2)}`); + html += stat("Estimated final variance", `${fmt(a.variance, 2)}`); } if (a.hot_only_fraction !== undefined) html += stat("Hot-only cycles", `${fmt(100 * a.hot_only_fraction, 1)}%`); html += `
`; if (a.cost !== undefined && a.cost !== null) { - html += `
Estimates from the ranking at cycle ${a.ranked_cycle}, relative to standard replica exchange: cost is the compute per cycle and variance that of the free energy for the same simulation time.
`; + html += `
Estimates from the ranking at cycle ${a.ranked_cycle}, relative to every window being hot for the rest of the run: cost is the compute per cycle from then on, and variance that of the final free energy.
`; } else if (a.ranked_cycle === undefined && (!a.hot || a.hot.length === n)) { html += `
The windows haven't been ranked yet, so every window is hot until the burn-in is complete.
`; } diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 7136b21..0fce75d 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2125,6 +2125,9 @@ def run(self): if is_adaptive: self._reset_adaptive_data() self._restore_adaptive_data() + cold_rate = self._cold_rate( + self._config.adaptive_mix_interval, cycles_per_checkpoint + ) else: # Adaptive replica exchange may have been switched off on restart. self._dynamics_cache._hot_replicas = None @@ -2301,7 +2304,10 @@ def run(self): # carries on with the same hot replicas. if is_adaptive and is_checkpoint and is_burnt_in: self._rank_windows( - num_workers, self._window_samples(), cycles - (i + 1) + num_workers, + self._window_samples(), + cycles - (i + 1), + cold_rate, ) # Create the lock. @@ -2924,6 +2930,42 @@ def _assemble_results(self, results): return matrix + @staticmethod + def _cold_rate(mix_interval, cycles_per_checkpoint): + """ + The share of cycles that propagate every window, i.e. those on which + the cold windows are propagated: the mix and checkpoint cycles. + + Parameters + ---------- + + mix_interval: int + The number of cycles between mix cycles. + + cycles_per_checkpoint: float + The number of cycles between checkpoints. + + Returns + ------- + + float + The share of cycles. + """ + from math import gcd + + mix = 1.0 / mix_interval + checkpoint = 1.0 / cycles_per_checkpoint + + # Count cycles that are both once, which for whole numbers of cycles + # are those at multiples of the lowest common multiple. + if abs(cycles_per_checkpoint - round(cycles_per_checkpoint)) < 1e-6: + c = round(cycles_per_checkpoint) + both = gcd(mix_interval, c) / (mix_interval * c) + else: + both = mix * checkpoint + + return min(1.0, mix + checkpoint - both) + def _window_samples(self): """ The number of energy samples each window has collected, from its clock. @@ -3241,7 +3283,7 @@ def _is_unconverged(cls, fwd, rev, g_fwd, g_rev): drift = abs(estimates[0] - estimates[1]) return drift > 0.25 and drift > 2.0 * _np.hypot(*errors) - def _rank_windows(self, num_workers, window_samples, cycles_left): + def _rank_windows(self, num_workers, window_samples, cycles_left, cold_rate): """ Rank the windows by the uncertainty they contribute to the free energy and choose the hot replicas for adaptive replica exchange. @@ -3278,9 +3320,17 @@ def _rank_windows(self, num_workers, window_samples, cycles_left): cycles_left: int The number of cycles left in the run. + + cold_rate: float + The share of cycles that propagate every window, and so the cold + ones. """ from math import ceil + # There is nothing left to allocate, so keep the current hot replicas. + if cycles_left <= 0: + return + num_replicas = self._config.num_lambda lower = [_np.array(x) for x in self._adaptive_lower] upper = [_np.array(x) for x in self._adaptive_upper] @@ -3363,33 +3413,31 @@ def fill(windows): cold = sorted(all_hot - windows, key=lambda r: -owned[r]) return windows | set(cold[: target - len(windows)]) - # A cold window is propagated every 'adaptive_mix_interval' cycles. - interval = self._config.adaptive_mix_interval total_variance = owned.sum() samples = _np.maximum(_np.asarray(window_samples, dtype=float), 1.0) def rates(windows): - return _np.array( - [1.0 if r in windows else 1.0 / interval for r in range(num_replicas)] - ) + # The samples per cycle of each window for the rest of the run. + is_hot = _np.zeros(num_replicas, dtype=bool) + is_hot[list(windows)] = True + return _np.where(is_hot, 1.0, cold_rate) - def cost_and_variance(windows): + def score(windows): # The compute of the whole run, in blocks of dynamics, and the # variance of the final free energy, from the samples at its end. rate = rates(windows) final = samples + rate * cycles_left cost = samples.sum() + rate.sum() * cycles_left - variance = (fwd_variance / final[:-1]).sum() + ( - rev_variance / final[1:] - ).sum() + variance = (fwd_variance / final[:-1] + rev_variance / final[1:]).sum() return cost, variance # A set must beat all hot by more than rounding error. - threshold = _np.prod(cost_and_variance(all_hot)) * (1 - 1e-9) + all_hot_cost, all_hot_variance = score(all_hot) + threshold = all_hot_cost * all_hot_variance * (1 - 1e-9) - # Add pairs in order of weight, keeping the hot windows that give the - # lowest variance for a fixed amount of compute, if better than all hot. - # Sets are scored with the windows that will fill the workers. + # Add pairs in order of weight, keeping the set with the lowest final + # variance multiplied by the compute of the whole run, if better than + # all hot. Sets are scored with the windows that will fill the workers. candidate = set(forced) chosen = None best = threshold @@ -3402,7 +3450,7 @@ def cost_and_variance(windows): windows = pad(candidate) if not windows: continue - efficiency = _np.prod(cost_and_variance(fill(windows))) + efficiency = _np.prod(score(fill(windows))) if efficiency < best: chosen, best = windows, efficiency @@ -3425,7 +3473,7 @@ def cost_and_variance(windows): # Windows kept hot can leave no gain over all hot. hot = fill(hot) - if _np.prod(cost_and_variance(hot)) >= threshold: + if _np.prod(score(hot)) >= threshold: hot = all_hot hot = sorted(hot) @@ -3444,15 +3492,13 @@ def cost_and_variance(windows): _logger.info( f"Hot replicas: {hot}" if len(hot) < num_replicas else "All replicas hot" ) - # The cost per cycle from now on, and the final variance, relative to + # The compute per cycle from now on and the final variance, relative to # every window being hot for the rest of the run. - cost = rates(set(hot)).sum() / num_replicas - variance = cost_and_variance(set(hot))[1] / ( - cost_and_variance(all_hot)[1] or 1.0 - ) + cost = rates(hot).sum() / num_replicas + variance = score(hot)[1] / (all_hot_variance or 1.0) _logger.info( - f"Estimated cost {cost:.2f} and variance {variance:.2f} relative to " - "standard replica exchange" + f"Estimated cost per cycle {cost:.2f} and final variance " + f"{variance:.2f} relative to standard replica exchange" ) def _check_restart(self): diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 644359c..cfff5aa 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1438,7 +1438,8 @@ def _rank(runner, num_workers, window_samples=None, cycles_left=10**6): """ if window_samples is None: window_samples = [200] * runner._config.num_lambda - runner._rank_windows(num_workers, window_samples, cycles_left) + cold_rate = 1.0 / runner._config.adaptive_mix_interval + runner._rank_windows(num_workers, window_samples, cycles_left, cold_rate) def _rank_twice(runner, num_workers): @@ -1618,6 +1619,43 @@ def rank(samples): assert rank([1000, 4000, 4000, 1000, 1000, 1000, 1000, 1000]) == [4, 5, 6, 7] +def test_rank_windows_at_end_of_run(): + """ + Validate that a ranking with no cycles left keeps the hot replicas and the + cold votes, so that an extended run carries on with them. + """ + runner = _ranking_runner(8, adaptive_min_samples=10) + runner._dynamics_cache._hot_replicas = [2, 3, 4, 5] + runner._dynamics_cache._cold_votes = [1, 0, 0, 0, 0, 0, 0, 1] + + _fill_ranking_data(runner, hard_pair=-1) + _rank(runner, 2, cycles_left=0) + + assert runner._dynamics_cache._hot_replicas == [2, 3, 4, 5] + assert runner._dynamics_cache._cold_votes == [1, 0, 0, 0, 0, 0, 0, 1] + + +@pytest.mark.parametrize( + "mix_interval, cycles_per_checkpoint, expected", + [ + # Every checkpoint cycle is also a mix cycle. + (4, 100, 0.25), + # Mix cycles at 4, 8, ..., checkpoints at 10, 20, ..., both at 20, 40. + (4, 10, 0.25 + 0.1 - 0.05), + (1, 100, 1.0), + # Not a whole number of cycles between checkpoints. + (4, 2.5, 0.25 + 0.4 - 0.1), + ], +) +def test_cold_rate(mix_interval, cycles_per_checkpoint, expected): + """ + Validate the share of cycles that propagate the cold windows. + """ + assert RepexRunner._cold_rate(mix_interval, cycles_per_checkpoint) == ( + pytest.approx(expected) + ) + + def test_rank_windows_without_gain(): """ Validate that every replica is made hot straight away when no subset is diff --git a/tests/viewer/test_adaptive.py b/tests/viewer/test_adaptive.py index a29500a..628e13f 100644 --- a/tests/viewer/test_adaptive.py +++ b/tests/viewer/test_adaptive.py @@ -40,7 +40,10 @@ def test_adaptive_summary(tmp_path): + _cycle(3, 4) + [ _line("All replicas hot"), - _line("Estimated cost 1.00 and variance 1.00 relative to standard"), + _line( + "Estimated cost per cycle 1.00 and final variance 1.00 relative " + "to standard" + ), ] + _cycle(4, 4) ) From dea2cd8bd9ce18378bbd715c5faa7a6785766787 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 09:53:14 +0100 Subject: [PATCH 13/32] Compute the adaptive cold rate from the remaining schedule. --- src/somd2/_viewer/_data.py | 2 +- src/somd2/runner/_repex.py | 66 +++++++++++++++++--------- tests/runner/test_repex.py | 87 +++++++++++++++++++++++++++++++---- tests/viewer/test_adaptive.py | 8 ++-- 4 files changed, 127 insertions(+), 36 deletions(-) diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index 519057e..ac8bb79 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -96,7 +96,7 @@ def _run(self): _correction_re = _re.compile(r"Standard state correction: (-?[\d.]+) kcal mol-1") _hot_re = _re.compile(r"^Hot replicas: \[([\d, ]*)\]") _estimate_re = _re.compile( - r"^Estimated cost (?:per cycle )?([\d.]+) and (?:final )?variance ([\d.]+)" + r"^Estimated cost per cycle ([\d.]+) and final variance ([\d.]+)" ) # Older versions don't log the time. _waters_re = _re.compile( diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 0fce75d..5b9b29a 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2125,9 +2125,6 @@ def run(self): if is_adaptive: self._reset_adaptive_data() self._restore_adaptive_data() - cold_rate = self._cold_rate( - self._config.adaptive_mix_interval, cycles_per_checkpoint - ) else: # Adaptive replica exchange may have been switched off on restart. self._dynamics_cache._hot_replicas = None @@ -2303,6 +2300,14 @@ def run(self): # Rank the windows before the state is saved, so that a restart # carries on with the same hot replicas. if is_adaptive and is_checkpoint and is_burnt_in: + # The threshold isn't advanced past this checkpoint yet. + cold_rate = self._cold_rate( + i + 1, + cycles, + next_checkpoint + cycles_per_checkpoint, + cycles_per_checkpoint, + self._config.adaptive_mix_interval, + ) self._rank_windows( num_workers, self._window_samples(), @@ -2931,40 +2936,57 @@ def _assemble_results(self, results): return matrix @staticmethod - def _cold_rate(mix_interval, cycles_per_checkpoint): + def _cold_rate( + first_cycle, num_cycles, next_checkpoint, cycles_per_checkpoint, mix_interval + ): """ - The share of cycles that propagate every window, i.e. those on which - the cold windows are propagated: the mix and checkpoint cycles. + The share of the cycles left that propagate every window, and so the + cold ones: the mix and checkpoint cycles, and the last cycle. These + follow the same rules as the run loop. Parameters ---------- - mix_interval: int - The number of cycles between mix cycles. + first_cycle: int + The number of cycles run so far. + + num_cycles: int + The number of cycles in the run. + + next_checkpoint: float + The threshold for the next checkpoint, as used by the run loop. cycles_per_checkpoint: float The number of cycles between checkpoints. + mix_interval: int + The number of cycles between mix cycles. + Returns ------- float The share of cycles. """ - from math import gcd + if first_cycle >= num_cycles: + return 1.0 - mix = 1.0 / mix_interval - checkpoint = 1.0 / cycles_per_checkpoint + # The cycles left, numbered from one as in the run loop. + cycles = _np.arange(first_cycle + 1, num_cycles + 1) + full = cycles % mix_interval == 0 - # Count cycles that are both once, which for whole numbers of cycles - # are those at multiples of the lowest common multiple. - if abs(cycles_per_checkpoint - round(cycles_per_checkpoint)) < 1e-6: - c = round(cycles_per_checkpoint) - both = gcd(mix_interval, c) / (mix_interval * c) - else: - both = mix * checkpoint + # A checkpoint is on the first cycle that reaches each threshold. + num_checkpoints = int( + (num_cycles - next_checkpoint + 1e-10) // cycles_per_checkpoint + 1 + ) + if num_checkpoints > 0: + thresholds = next_checkpoint + _np.arange(num_checkpoints) * ( + cycles_per_checkpoint + ) + full |= _np.isin(cycles, _np.ceil(thresholds - 1e-10).astype(int)) - return min(1.0, mix + checkpoint - both) + full[-1] = True + return full.mean() def _window_samples(self): """ @@ -3329,6 +3351,7 @@ def _rank_windows(self, num_workers, window_samples, cycles_left, cold_rate): # There is nothing left to allocate, so keep the current hot replicas. if cycles_left <= 0: + _logger.info("Not ranking windows: no cycles left, keeping hot replicas") return num_replicas = self._config.num_lambda @@ -3476,6 +3499,7 @@ def score(windows): if _np.prod(score(hot)) >= threshold: hot = all_hot + hot_variance = all_hot_variance if hot == all_hot else score(hot)[1] hot = sorted(hot) self._dynamics_cache._hot_replicas = None if len(hot) == num_replicas else hot @@ -3495,10 +3519,10 @@ def score(windows): # The compute per cycle from now on and the final variance, relative to # every window being hot for the rest of the run. cost = rates(hot).sum() / num_replicas - variance = score(hot)[1] / (all_hot_variance or 1.0) + variance = hot_variance / all_hot_variance if all_hot_variance > 0 else 1.0 _logger.info( f"Estimated cost per cycle {cost:.2f} and final variance " - f"{variance:.2f} relative to standard replica exchange" + f"{variance:.2f} relative to every window being hot" ) def _check_restart(self): diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index cfff5aa..cec6659 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1635,25 +1635,47 @@ def test_rank_windows_at_end_of_run(): assert runner._dynamics_cache._cold_votes == [1, 0, 0, 0, 0, 0, 0, 1] +def _full_cycle_share(first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix): + """ + The share of the cycles left that propagate every window, stepping + through them with the same rules as the run loop. + """ + full = 0 + for i in range(first, num_cycles): + is_checkpoint = (i + 1) >= next_checkpoint - 1e-10 + if (i + 1) % mix == 0 or is_checkpoint or i == num_cycles - 1: + full += 1 + if is_checkpoint: + next_checkpoint += cycles_per_checkpoint + return full / (num_cycles - first) + + @pytest.mark.parametrize( - "mix_interval, cycles_per_checkpoint, expected", + "first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix", [ # Every checkpoint cycle is also a mix cycle. - (4, 100, 0.25), - # Mix cycles at 4, 8, ..., checkpoints at 10, 20, ..., both at 20, 40. - (4, 10, 0.25 + 0.1 - 0.05), - (1, 100, 1.0), + (100, 1000, 200, 100, 4), + (100, 1000, 110, 10, 4), + (0, 1000, 100, 100, 1), # Not a whole number of cycles between checkpoints. - (4, 2.5, 0.25 + 0.4 - 0.1), + (5, 1000, 7.5, 2.5, 4), + (100, 997, 133.3, 33.3, 7), + # Close to the end, where the final cycle matters. + (990, 1000, 1100, 100, 4), + (998, 1000, 1100, 100, 4), ], ) -def test_cold_rate(mix_interval, cycles_per_checkpoint, expected): +def test_cold_rate(first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix): """ - Validate the share of cycles that propagate the cold windows. + Validate the share of cycles that propagate the cold windows against the + run loop's schedule. """ - assert RepexRunner._cold_rate(mix_interval, cycles_per_checkpoint) == ( - pytest.approx(expected) + expected = _full_cycle_share( + first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix ) + assert RepexRunner._cold_rate( + first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix + ) == pytest.approx(expected) def test_rank_windows_without_gain(): @@ -1795,6 +1817,51 @@ def test_find_ring_pair(syk_ring_break_mols): assert runner._find_ring_pair() == (1, 7) +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_adaptive_repex_ranking_arguments(ethane_methanol, monkeypatch): + """ + Validate that each ranking is given the samples each window has, the + cycles left, and the share of those cycles that propagate every window. + """ + num_lambda = 4 + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "48fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "16fs", + "frame_frequency": "16fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "adaptive_repex": True, + "adaptive_mix_interval": 2, + "adaptive_burn_in": "0 ps", + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + calls = [] + monkeypatch.setattr( + runner, + "_rank_windows", + lambda *args: calls.append(args), + ) + runner.run() + + # 12 cycles with a checkpoint every 4, and every window always hot. + assert [c[2] for c in calls] == [8, 4, 0] + for num_workers, samples, cycles_left, cold_rate in calls: + first = 12 - cycles_left + assert samples == [first] * num_lambda + if cycles_left > 0: + expected = _full_cycle_share(first, 12, first + 4, 4, 2) + assert cold_rate == pytest.approx(expected) + + @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") def test_adaptive_repex_frames_every_cycle(ethane_methanol): """ diff --git a/tests/viewer/test_adaptive.py b/tests/viewer/test_adaptive.py index 628e13f..56e956c 100644 --- a/tests/viewer/test_adaptive.py +++ b/tests/viewer/test_adaptive.py @@ -34,7 +34,7 @@ def test_adaptive_summary(tmp_path): _cycle(1, 4) + [ _line("Hot replicas: [1, 2]"), - _line("Estimated cost 0.75 and variance 1.10 relative to standard"), + _line("Estimated cost per cycle 0.75 and final variance 1.10"), ] + _cycle(2, 4, hot=[1, 2]) + _cycle(3, 4) @@ -69,7 +69,7 @@ def test_adaptive_rankings_after_restart(tmp_path): + _cycle(2, 3, hot=[1, 2]) + [ _line("Hot replicas: [1, 2]"), - _line("Estimated cost 0.75 and variance 1.10 relative to standard"), + _line("Estimated cost per cycle 0.75 and final variance 1.10"), _line("somd2 version: test"), ] + _cycle(2, 3) @@ -89,7 +89,7 @@ def test_adaptive_rankings_kept_after_final_save(tmp_path): _cycle(1, 2) + [ _line("Hot replicas: [1, 2]"), - _line("Estimated cost 0.75 and variance 1.10 relative to standard"), + _line("Estimated cost per cycle 0.75 and final variance 1.10"), ] + _cycle(2, 2, hot=[1, 2]) + [ @@ -112,7 +112,7 @@ def test_adaptive_history_is_thinned(tmp_path): for n in range(1, 11): lines += _cycle(n, 10) + [ _line("Hot replicas: [1, 2]"), - _line("Estimated cost 0.75 and variance 1.10 relative to standard"), + _line("Estimated cost per cycle 0.75 and final variance 1.10"), ] log = _monitor(tmp_path, lines) From 480879b12ec6e3bec60107634b5cf8dedc3de3f6 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 10:00:26 +0100 Subject: [PATCH 14/32] Share one checkpoint schedule between the run loop and cold rate. --- src/somd2/runner/_repex.py | 92 ++++++++++++++++++++++---------------- tests/runner/test_repex.py | 89 +++++++++++++++++++++++------------- 2 files changed, 110 insertions(+), 71 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 5b9b29a..6ede381 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2104,9 +2104,9 @@ def run(self): else: cycles_per_flip = cycles + 1 - # Initialise the threshold for the next checkpoint cycle. This is a float - # to handle non-integer ratios between the checkpoint and energy frequencies. - next_checkpoint = cycles_per_checkpoint + # The cycles, numbered from one, that end with a checkpoint. + checkpoint_cycles = self._checkpoint_cycles(cycles, cycles_per_checkpoint) + checkpoint_set = set(checkpoint_cycles.tolist()) # Whether the most recent cycle saved the replica exchange state. is_checkpoint = False @@ -2146,9 +2146,8 @@ def run(self): # Clear the results list. results = [] - # Whether to checkpoint. Use a float threshold to correctly handle - # non-integer ratios between the checkpoint and energy frequencies. - is_checkpoint = (i + 1) >= next_checkpoint - 1e-10 + # Whether to checkpoint. + is_checkpoint = (i + 1) in checkpoint_set # Whether the checkpoint files are written at the end of this cycle. do_checkpoint = is_checkpoint or i == cycles - 1 @@ -2300,12 +2299,10 @@ def run(self): # Rank the windows before the state is saved, so that a restart # carries on with the same hot replicas. if is_adaptive and is_checkpoint and is_burnt_in: - # The threshold isn't advanced past this checkpoint yet. cold_rate = self._cold_rate( i + 1, cycles, - next_checkpoint + cycles_per_checkpoint, - cycles_per_checkpoint, + checkpoint_cycles, self._config.adaptive_mix_interval, ) self._rank_windows( @@ -2389,9 +2386,6 @@ def run(self): # Update the block number. block += 1 - # Advance the checkpoint threshold. - next_checkpoint += cycles_per_checkpoint - dynamics_executor.shutdown(wait=True) checkpoint_executor.shutdown(wait=True) @@ -2936,13 +2930,37 @@ def _assemble_results(self, results): return matrix @staticmethod - def _cold_rate( - first_cycle, num_cycles, next_checkpoint, cycles_per_checkpoint, mix_interval - ): + def _checkpoint_cycles(num_cycles, cycles_per_checkpoint): + """ + The cycles that end with a checkpoint: the first to reach each multiple + of the number of cycles between checkpoints, which may not be a whole + number. + + Parameters + ---------- + + num_cycles: int + The number of cycles in the run. + + cycles_per_checkpoint: float + The number of cycles between checkpoints. + + Returns + ------- + + np.ndarray + The checkpoint cycles, numbered from one, in increasing order. + """ + num_checkpoints = int(num_cycles / cycles_per_checkpoint) + 1 + multiples = _np.arange(1, num_checkpoints + 1) * cycles_per_checkpoint + cycles = _np.ceil(multiples - 1e-10).astype(int) + return cycles[cycles <= num_cycles] + + @staticmethod + def _cold_rate(first_cycle, num_cycles, checkpoint_cycles, mix_interval): """ The share of the cycles left that propagate every window, and so the - cold ones: the mix and checkpoint cycles, and the last cycle. These - follow the same rules as the run loop. + cold ones: the mix and checkpoint cycles, and the last cycle. Parameters ---------- @@ -2953,11 +2971,8 @@ def _cold_rate( num_cycles: int The number of cycles in the run. - next_checkpoint: float - The threshold for the next checkpoint, as used by the run loop. - - cycles_per_checkpoint: float - The number of cycles between checkpoints. + checkpoint_cycles: np.ndarray + The checkpoint cycles, numbered from one, in increasing order. mix_interval: int The number of cycles between mix cycles. @@ -2971,22 +2986,21 @@ def _cold_rate( if first_cycle >= num_cycles: return 1.0 - # The cycles left, numbered from one as in the run loop. - cycles = _np.arange(first_cycle + 1, num_cycles + 1) - full = cycles % mix_interval == 0 - - # A checkpoint is on the first cycle that reaches each threshold. - num_checkpoints = int( - (num_cycles - next_checkpoint + 1e-10) // cycles_per_checkpoint + 1 + # Mix cycles, then checkpoint cycles that aren't also mix cycles. + full = num_cycles // mix_interval - first_cycle // mix_interval + start, end = _np.searchsorted( + checkpoint_cycles, [first_cycle, num_cycles], side="right" ) - if num_checkpoints > 0: - thresholds = next_checkpoint + _np.arange(num_checkpoints) * ( - cycles_per_checkpoint - ) - full |= _np.isin(cycles, _np.ceil(thresholds - 1e-10).astype(int)) + checkpoints = checkpoint_cycles[start:end] + full += int((checkpoints % mix_interval != 0).sum()) + + # The last cycle always propagates every window. + if num_cycles % mix_interval != 0 and ( + len(checkpoints) == 0 or checkpoints[-1] != num_cycles + ): + full += 1 - full[-1] = True - return full.mean() + return full / (num_cycles - first_cycle) def _window_samples(self): """ @@ -3344,8 +3358,8 @@ def _rank_windows(self, num_workers, window_samples, cycles_left, cold_rate): The number of cycles left in the run. cold_rate: float - The share of cycles that propagate every window, and so the cold - ones. + The share of the cycles left that propagate every window, and so + the cold ones. """ from math import ceil @@ -3499,7 +3513,7 @@ def score(windows): if _np.prod(score(hot)) >= threshold: hot = all_hot - hot_variance = all_hot_variance if hot == all_hot else score(hot)[1] + hot_variance = score(hot)[1] hot = sorted(hot) self._dynamics_cache._hot_replicas = None if len(hot) == num_replicas else hot diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index cec6659..4a6c7d7 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1635,47 +1635,72 @@ def test_rank_windows_at_end_of_run(): assert runner._dynamics_cache._cold_votes == [1, 0, 0, 0, 0, 0, 0, 1] -def _full_cycle_share(first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix): +def _loop_schedule(num_cycles, cycles_per_checkpoint, mix): """ - The share of the cycles left that propagate every window, stepping - through them with the same rules as the run loop. + The checkpoint cycles and the cycles that propagate every window, from + stepping through the cycles with a running checkpoint threshold, kept + exact so that it doesn't drift. """ - full = 0 - for i in range(first, num_cycles): - is_checkpoint = (i + 1) >= next_checkpoint - 1e-10 - if (i + 1) % mix == 0 or is_checkpoint or i == num_cycles - 1: - full += 1 + from fractions import Fraction + + checkpoints, full = [], [] + cycles_per_checkpoint = Fraction(cycles_per_checkpoint) + threshold = cycles_per_checkpoint + for i in range(num_cycles): + is_checkpoint = (i + 1) >= threshold if is_checkpoint: - next_checkpoint += cycles_per_checkpoint - return full / (num_cycles - first) + checkpoints.append(i + 1) + threshold += cycles_per_checkpoint + if (i + 1) % mix == 0 or is_checkpoint or i == num_cycles - 1: + full.append(i + 1) + return checkpoints, full -@pytest.mark.parametrize( - "first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix", - [ - # Every checkpoint cycle is also a mix cycle. - (100, 1000, 200, 100, 4), - (100, 1000, 110, 10, 4), - (0, 1000, 100, 100, 1), - # Not a whole number of cycles between checkpoints. - (5, 1000, 7.5, 2.5, 4), - (100, 997, 133.3, 33.3, 7), - # Close to the end, where the final cycle matters. - (990, 1000, 1100, 100, 4), - (998, 1000, 1100, 100, 4), - ], -) -def test_cold_rate(first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix): +def _full_cycle_share(first, num_cycles, cycles_per_checkpoint, mix): + """ + The share of the cycles after 'first' that propagate every window. + """ + _, full = _loop_schedule(num_cycles, cycles_per_checkpoint, mix) + return sum(c > first for c in full) / (num_cycles - first) + + +SCHEDULES = [ + # Every checkpoint cycle is also a mix cycle. + (100, 1000, 100, 4), + (100, 1000, 10, 4), + (0, 1000, 100, 1), + # Not a whole number of cycles between checkpoints. + (5, 1000, 2.5, 4), + (100, 997, 33.3, 7), + (0, 100000, 33.3, 3), + # A checkpoint interval longer than the run. + (0, 250, 250.0, 4), + # Close to the end, where the final cycle matters. + (990, 1000, 100, 4), + (998, 1000, 100, 4), +] + + +@pytest.mark.parametrize("first, num_cycles, cycles_per_checkpoint, mix", SCHEDULES) +def test_checkpoint_cycles(first, num_cycles, cycles_per_checkpoint, mix): + """ + Validate the checkpoint cycles against a running checkpoint threshold. + """ + expected, _ = _loop_schedule(num_cycles, cycles_per_checkpoint, mix) + cycles = RepexRunner._checkpoint_cycles(num_cycles, cycles_per_checkpoint) + assert cycles.tolist() == expected + + +@pytest.mark.parametrize("first, num_cycles, cycles_per_checkpoint, mix", SCHEDULES) +def test_cold_rate(first, num_cycles, cycles_per_checkpoint, mix): """ Validate the share of cycles that propagate the cold windows against the run loop's schedule. """ - expected = _full_cycle_share( - first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix + checkpoints = RepexRunner._checkpoint_cycles(num_cycles, cycles_per_checkpoint) + assert RepexRunner._cold_rate(first, num_cycles, checkpoints, mix) == pytest.approx( + _full_cycle_share(first, num_cycles, cycles_per_checkpoint, mix) ) - assert RepexRunner._cold_rate( - first, num_cycles, next_checkpoint, cycles_per_checkpoint, mix - ) == pytest.approx(expected) def test_rank_windows_without_gain(): @@ -1858,7 +1883,7 @@ def test_adaptive_repex_ranking_arguments(ethane_methanol, monkeypatch): first = 12 - cycles_left assert samples == [first] * num_lambda if cycles_left > 0: - expected = _full_cycle_share(first, 12, first + 4, 4, 2) + expected = _full_cycle_share(first, 12, 4, 2) assert cold_rate == pytest.approx(expected) From 626ac3d3bf37a879d6384647655322ab4e10aeec Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 10:17:04 +0100 Subject: [PATCH 15/32] Handle runs without checkpoint cycles and fix the schedule test. --- src/somd2/runner/_repex.py | 3 +++ tests/runner/test_repex.py | 15 +++++++++++++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 6ede381..40562dc 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2951,6 +2951,9 @@ def _checkpoint_cycles(num_cycles, cycles_per_checkpoint): np.ndarray The checkpoint cycles, numbered from one, in increasing order. """ + if num_cycles <= 0 or cycles_per_checkpoint <= 0: + return _np.array([], dtype=int) + num_checkpoints = int(num_cycles / cycles_per_checkpoint) + 1 multiples = _np.arange(1, num_checkpoints + 1) * cycles_per_checkpoint cycles = _np.ceil(multiples - 1e-10).astype(int) diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 4a6c7d7..e5aedc5 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1639,15 +1639,16 @@ def _loop_schedule(num_cycles, cycles_per_checkpoint, mix): """ The checkpoint cycles and the cycles that propagate every window, from stepping through the cycles with a running checkpoint threshold, kept - exact so that it doesn't drift. + exact so that it doesn't drift, with the same tolerance as the runner. """ from fractions import Fraction checkpoints, full = [], [] cycles_per_checkpoint = Fraction(cycles_per_checkpoint) + tolerance = Fraction(1, 10**10) threshold = cycles_per_checkpoint for i in range(num_cycles): - is_checkpoint = (i + 1) >= threshold + is_checkpoint = (i + 1) >= threshold - tolerance if is_checkpoint: checkpoints.append(i + 1) threshold += cycles_per_checkpoint @@ -1673,6 +1674,8 @@ def _full_cycle_share(first, num_cycles, cycles_per_checkpoint, mix): (5, 1000, 2.5, 4), (100, 997, 33.3, 7), (0, 100000, 33.3, 3), + # Multiples that are whole numbers, but not exactly so in binary. + (0, 1000, 1.1, 4), # A checkpoint interval longer than the run. (0, 250, 250.0, 4), # Close to the end, where the final cycle matters. @@ -1691,6 +1694,14 @@ def test_checkpoint_cycles(first, num_cycles, cycles_per_checkpoint, mix): assert cycles.tolist() == expected +def test_checkpoint_cycles_without_cycles(): + """ + Validate that a run with no cycles, and so no checkpoint interval, has no + checkpoints. + """ + assert RepexRunner._checkpoint_cycles(0, 0.0).tolist() == [] + + @pytest.mark.parametrize("first, num_cycles, cycles_per_checkpoint, mix", SCHEDULES) def test_cold_rate(first, num_cycles, cycles_per_checkpoint, mix): """ From a7a475e5eb01a5037f6bdea39fc5f05f34b9f5d5 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 11:18:38 +0100 Subject: [PATCH 16/32] Balance hot replicas across GPUs when batching and filling. --- src/somd2/runner/_repex.py | 98 ++++++++++++++++++++++++++++++-------- tests/runner/test_repex.py | 86 +++++++++++++++++++++++---------- 2 files changed, 137 insertions(+), 47 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 40562dc..d071384 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2176,9 +2176,7 @@ def run(self): # Run a dynamics block for each replica in the pass, making sure # each GPU is only oversubscribed by a factor of # self._config.oversubscription_factor. - for j in range(ceil(len(batch) / num_workers)): - replicas = batch[j * num_workers : (j + 1) * num_workers] - + for replicas in self._device_batches(batch): # Whether each replica is due a GCMC or terminal flip move, # including any that fell due while it wasn't propagated. is_gcmc = [ @@ -2306,7 +2304,7 @@ def run(self): self._config.adaptive_mix_interval, ) self._rank_windows( - num_workers, + self._replica_devices(), self._window_samples(), cycles - (i + 1), cold_rate, @@ -3005,6 +3003,55 @@ def _cold_rate(first_cycle, num_cycles, checkpoint_cycles, mix_interval): return full / (num_cycles - first_cycle) + def _replica_devices(self): + """ + The GPU that propagates each replica, which is that of its slot. + + Returns + ------- + + list of int + The device index for each replica. + """ + return [ + self._dynamics_cache.slot_for(r) % self._num_gpus + for r in range(self._num_replicas) + ] + + def _device_batches(self, replicas): + """ + Split replicas into batches that are propagated concurrently, each with + at most 'oversubscription_factor' replicas on any GPU, so that every + batch spreads its work over the GPUs. + + Parameters + ---------- + + replicas: list of int + The replicas to propagate, at most one per slot. + + Returns + ------- + + list of list of int + The batches. + """ + devices = self._replica_devices() + per_device = self._config.oversubscription_factor + + queues = {} + for r in replicas: + queues.setdefault(devices[r], []).append(r) + + batches = [] + while any(queues.values()): + batch = [] + for queue in queues.values(): + batch += queue[:per_device] + del queue[:per_device] + batches.append(sorted(batch)) + return batches + def _window_samples(self): """ The number of energy samples each window has collected, from its clock. @@ -3322,7 +3369,7 @@ def _is_unconverged(cls, fwd, rev, g_fwd, g_rev): drift = abs(estimates[0] - estimates[1]) return drift > 0.25 and drift > 2.0 * _np.hypot(*errors) - def _rank_windows(self, num_workers, window_samples, cycles_left, cold_rate): + def _rank_windows(self, window_devices, window_samples, cycles_left, cold_rate): """ Rank the windows by the uncertainty they contribute to the free energy and choose the hot replicas for adaptive replica exchange. @@ -3335,15 +3382,15 @@ def _rank_windows(self, num_workers, window_samples, cycles_left, cold_rate): differences or, for a perturbation that breaks or makes a single bond, the bond distance. The remaining pairs are added in order of weight, with the windows either side of each so that exchange carries - configurations into and out of them. Each set this gives is filled to a - multiple of the number of workers with the cold windows owning the most - variance, so that the GPUs are evenly loaded, and scored by the - variance of the final free energy, given the samples each window has - and will collect in the cycles left, multiplied by the compute of the - whole run. The set with the lowest score is chosen. A window only stops - being hot once it has been found cold at two consecutive rankings. - Every window is hot if no set, including the windows kept hot, is - better than that. + configurations into and out of them. Each set this gives is filled with + the cold windows owning the most variance, until every GPU has as many + hot windows as the busiest, since those take no extra time, and scored + by the variance of the final free energy, given the samples each window + has and will collect in the cycles left, multiplied by the compute of + the whole run. The set with the lowest score is chosen. A window only + stops being hot once it has been found cold at two consecutive + rankings. Every window is hot if no set, including the windows kept + hot, is better than that. The ranking uses all of the data since the burn-in, since the slowest windows decorrelate over far longer than a checkpoint interval. @@ -3351,8 +3398,8 @@ def _rank_windows(self, num_workers, window_samples, cycles_left, cold_rate): Parameters ---------- - num_workers: int - The number of replicas that are propagated concurrently. + window_devices: list of int + The GPU that propagates each window. window_samples: list of int The number of energy samples each window has collected. @@ -3446,12 +3493,21 @@ def pad(windows): owned[:-1] += fwd_variance owned[1:] += rev_variance + per_device = self._config.oversubscription_factor + def fill(windows): - # Make the cold windows owning the most variance hot, until the hot - # windows are a multiple of the number of workers. - target = min(num_replicas, ceil(len(windows) / num_workers) * num_workers) - cold = sorted(all_hot - windows, key=lambda r: -owned[r]) - return windows | set(cold[: target - len(windows)]) + # Make the cold windows owning the most variance hot, until each GPU + # has as many hot windows as a whole number of batches on the busiest. + counts = {d: 0 for d in window_devices} + for r in windows: + counts[window_devices[r]] += 1 + target = ceil(max(counts.values()) / per_device) * per_device + filled = set(windows) + for r in sorted(all_hot - windows, key=lambda r: -owned[r]): + if counts[window_devices[r]] < target: + filled.add(r) + counts[window_devices[r]] += 1 + return filled total_variance = owned.sum() samples = _np.maximum(_np.asarray(window_samples, dtype=float), 1.0) diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index e5aedc5..d1efc8e 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1431,24 +1431,27 @@ def _ranking_runner(num_lambda, **config): return runner -def _rank(runner, num_workers, window_samples=None, cycles_left=10**6): +def _rank(runner, num_gpus, window_samples=None, cycles_left=10**6): """ - Rank the windows. By default every window has the same samples and the run - is far from its end, so that only the sampling rates matter. + Rank the windows, with window r on GPU r % num_gpus. By default every + window has the same samples and the run is far from its end, so that only + the sampling rates matter. """ + num_lambda = runner._config.num_lambda if window_samples is None: - window_samples = [200] * runner._config.num_lambda + window_samples = [200] * num_lambda + devices = [r % num_gpus for r in range(num_lambda)] cold_rate = 1.0 / runner._config.adaptive_mix_interval - runner._rank_windows(num_workers, window_samples, cycles_left, cold_rate) + runner._rank_windows(devices, window_samples, cycles_left, cold_rate) -def _rank_twice(runner, num_workers): +def _rank_twice(runner, num_gpus): """ Rank the windows at two consecutive checkpoints, so that the windows found cold at both stop being hot. """ - _rank(runner, num_workers=num_workers) - _rank(runner, num_workers=num_workers) + _rank(runner, num_gpus=num_gpus) + _rank(runner, num_gpus=num_gpus) def _fill_ranking_data(runner, hard_pair, num_samples=200, seed=7): @@ -1476,37 +1479,39 @@ def test_rank_windows(): # Starting from every window hot, those found cold once stay hot. _fill_ranking_data(runner, hard_pair=3) - _rank(runner, num_workers=2) + _rank(runner, num_gpus=2) assert cache._hot_replicas is None - _rank(runner, num_workers=2) + _rank(runner, num_gpus=2) assert cache._hot_replicas == [2, 3, 4, 5] # Windows 2 and 3 are found cold once, so stay hot. _fill_ranking_data(runner, hard_pair=5) - _rank(runner, num_workers=2) + _rank(runner, num_gpus=2) assert cache._hot_replicas == [2, 3, 4, 5, 6, 7] # Found cold a second time. _fill_ranking_data(runner, hard_pair=5) - _rank(runner, num_workers=2) + _rank(runner, num_gpus=2) assert cache._hot_replicas == [4, 5, 6, 7] def test_rank_windows_pads_to_workers(): """ - Validate that the hot replicas are padded to a multiple of the number of - workers, and that every replica being hot is recorded as None. + Validate that the hot replicas are padded so that every GPU has the same + number, and that every replica being hot is recorded as None. """ runner = _ranking_runner(8, adaptive_min_samples=10) _fill_ranking_data(runner, hard_pair=3) - _rank_twice(runner, num_workers=3) - assert len(runner._dynamics_cache._hot_replicas) == 6 + _rank_twice(runner, num_gpus=3) + hot = runner._dynamics_cache._hot_replicas + assert len(hot) == 6 + assert sorted(r % 3 for r in hot) == [0, 0, 1, 1, 2, 2] runner._dynamics_cache._hot_replicas = None _fill_ranking_data(runner, hard_pair=3) - _rank(runner, num_workers=8) + _rank(runner, num_gpus=8) assert runner._dynamics_cache._hot_replicas is None @@ -1569,7 +1574,7 @@ def test_rank_windows_minimises_variance_for_compute(): rng = np.random.default_rng(11) runner._adaptive_upper[3] = list(rng.normal(1.0, 1.0, 200)) runner._adaptive_lower[4] = list(rng.normal(0.0, 1.0, 200)) - _rank_twice(runner, num_workers=2) + _rank_twice(runner, num_gpus=2) assert runner._dynamics_cache._hot_replicas == [2, 3, 4, 5] @@ -1583,7 +1588,7 @@ def test_rank_windows_fill_windows_stay_cold_candidates(): # The flagged pair 0 gives windows 0, 1 and 2, so one window is filled. _fill_ranking_data(runner, hard_pair=0) - _rank_twice(runner, num_workers=2) + _rank_twice(runner, num_gpus=2) hot = runner._dynamics_cache._hot_replicas assert hot is not None and len(hot) == 4 @@ -1726,7 +1731,7 @@ def test_rank_windows_without_gain(): # Every pair has the same uncertainty. _fill_ranking_data(runner, hard_pair=-1) - _rank(runner, num_workers=2) + _rank(runner, num_gpus=2) assert runner._dynamics_cache._hot_replicas is None assert runner._dynamics_cache._cold_votes == [0] * 8 @@ -1742,7 +1747,7 @@ def test_rank_windows_without_uncertainty(): for k in range(3): runner._adaptive_upper[k] = [0.5] * 100 runner._adaptive_lower[k + 1] = [-0.5] * 100 - _rank(runner, num_workers=1) + _rank(runner, num_gpus=1) assert runner._dynamics_cache._hot_replicas is None @@ -1780,7 +1785,7 @@ def test_rank_windows_needs_enough_samples(): runner = _ranking_runner(8, adaptive_min_samples=50) _fill_ranking_data(runner, hard_pair=3, num_samples=20) - _rank(runner, num_workers=2) + _rank(runner, num_gpus=2) assert runner._dynamics_cache._hot_replicas is None assert len(runner._adaptive_upper[0]) == 20 @@ -1801,7 +1806,7 @@ def test_rank_windows_keeps_undersampled_windows_hot(): for i in range(1, len(y)): y[i] = 0.95 * y[i - 1] + rng.normal(0.0, 0.01) runner._adaptive_upper[6] = list(0.5 + y) - _rank_twice(runner, num_workers=2) + _rank_twice(runner, num_gpus=2) assert 6 in runner._dynamics_cache._hot_replicas @@ -1824,12 +1829,39 @@ def test_rank_windows_keeps_slow_bond_distance_hot(): for i in range(1, len(y)): y[i] = 0.95 * y[i - 1] + rng.normal() runner._adaptive_distance[7] = list(10.0 + y) - _rank_twice(runner, num_workers=2) + _rank_twice(runner, num_gpus=2) assert 7 in runner._dynamics_cache._hot_replicas assert 0 not in runner._dynamics_cache._hot_replicas +def test_device_batches(): + """ + Validate that a batch never has more replicas on a GPU than the + oversubscription factor, so that hot replicas that aren't consecutive + still spread over the GPUs. + """ + from types import SimpleNamespace + + runner = object.__new__(RepexRunner) + runner._config = Config(num_lambda=25) + runner._num_gpus = 3 + runner._num_replicas = 25 + runner._dynamics_cache = SimpleNamespace(slot_for=lambda r: r) + + # Every replica is batched as consecutive groups of three. + assert runner._device_batches(list(range(25))) == [ + list(range(i, min(i + 3, 25))) for i in range(0, 25, 3) + ] + + hot = [6, 7, 8, 9, 10, 13, 14, 15, 16, 17, 20, 21, 22, 23, 24] + batches = runner._device_batches(hot) + assert sorted(r for batch in batches for r in batch) == hot + assert len(batches) == 5 + for batch in batches: + assert len({r % 3 for r in batch}) == len(batch) + + def _ring_pair_runner(system, virtual_sites=None): """ Create a runner holding only what is needed to find the ring pair. @@ -1890,7 +1922,9 @@ def test_adaptive_repex_ranking_arguments(ethane_methanol, monkeypatch): # 12 cycles with a checkpoint every 4, and every window always hot. assert [c[2] for c in calls] == [8, 4, 0] - for num_workers, samples, cycles_left, cold_rate in calls: + for devices, samples, cycles_left, cold_rate in calls: + # A context per replica, assigned to the GPUs in turn. + assert devices == [r % runner._num_gpus for r in range(num_lambda)] first = 12 - cycles_left assert samples == [first] * num_lambda if cycles_left > 0: @@ -2002,7 +2036,7 @@ def test_rank_windows_flags_unconverged_pairs(): _fill_ranking_data(runner, hard_pair=3) rng = np.random.default_rng(3) runner._adaptive_lower[1] = list(rng.normal(-3.0, 0.1, 200)) - _rank_twice(runner, num_workers=2) + _rank_twice(runner, num_gpus=2) hot = runner._dynamics_cache._hot_replicas assert hot is not None and {0, 1, 2}.issubset(hot) From 1d7f6f3cc5b0c0c16aa7edb471b382447aa632a0 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 11:24:13 +0100 Subject: [PATCH 17/32] Use GPU-aware batching for all dynamics and test oversubscription. --- src/somd2/runner/_repex.py | 55 +++++++++++++++++++++++++++++--------- tests/runner/test_repex.py | 44 +++++++++++++++++++++++++++--- 2 files changed, 83 insertions(+), 16 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index d071384..8288cb0 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1974,9 +1974,12 @@ def run(self): # Create the replica list. replica_list = list(range(self._config.num_lambda)) + # The GPU that propagates each replica, which is fixed for the run. + replica_devices = self._replica_devices() + # Minimise at each lambda value. Don't minimise on restart. if self._config.minimise and not self._is_restart: - for batch in self._safe_batches(num_workers): + for batch in self._dynamics_batches(replica_devices): with ThreadPoolExecutor(max_workers=num_workers) as executor: try: for success, index, e in executor.map( @@ -1996,7 +1999,7 @@ def run(self): # Equilibrate the system. if self._is_equilibration and not self._is_restart: - for batch in self._safe_batches(num_workers): + for batch in self._dynamics_batches(replica_devices): with ThreadPoolExecutor(max_workers=num_workers) as executor: try: for success, index, e in executor.map( @@ -2176,7 +2179,7 @@ def run(self): # Run a dynamics block for each replica in the pass, making sure # each GPU is only oversubscribed by a factor of # self._config.oversubscription_factor. - for replicas in self._device_batches(batch): + for replicas in self._device_batches(batch, replica_devices): # Whether each replica is due a GCMC or terminal flip move, # including any that fell due while it wasn't propagated. is_gcmc = [ @@ -2304,7 +2307,7 @@ def run(self): self._config.adaptive_mix_interval, ) self._rank_windows( - self._replica_devices(), + replica_devices, self._window_samples(), cycles - (i + 1), cold_rate, @@ -3018,7 +3021,7 @@ def _replica_devices(self): for r in range(self._num_replicas) ] - def _device_batches(self, replicas): + def _device_batches(self, replicas, devices): """ Split replicas into batches that are propagated concurrently, each with at most 'oversubscription_factor' replicas on any GPU, so that every @@ -3030,13 +3033,15 @@ def _device_batches(self, replicas): replicas: list of int The replicas to propagate, at most one per slot. + devices: list of int + The GPU that propagates each replica. + Returns ------- list of list of int The batches. """ - devices = self._replica_devices() per_device = self._config.oversubscription_factor queues = {} @@ -3052,6 +3057,29 @@ def _device_batches(self, replicas): batches.append(sorted(batch)) return batches + def _dynamics_batches(self, devices, cycle=0): + """ + Yield batches of every replica that can be propagated concurrently: at + most one per slot, and at most 'oversubscription_factor' per GPU. + + Parameters + ---------- + + devices: list of int + The GPU that propagates each replica. + + cycle: int + The index of the current cycle, which sets the traversal order. + + Yields + ------ + + list of int + A batch of replicas. + """ + for batch in self._replica_passes(cycle): + yield from self._device_batches(batch, devices) + def _window_samples(self): """ The number of energy samples each window has collected, from its clock. @@ -3384,13 +3412,14 @@ def _rank_windows(self, window_devices, window_samples, cycles_left, cold_rate): with the windows either side of each so that exchange carries configurations into and out of them. Each set this gives is filled with the cold windows owning the most variance, until every GPU has as many - hot windows as the busiest, since those take no extra time, and scored - by the variance of the final free energy, given the samples each window - has and will collect in the cycles left, multiplied by the compute of - the whole run. The set with the lowest score is chosen. A window only - stops being hot once it has been found cold at two consecutive - rankings. Every window is hot if no set, including the windows kept - hot, is better than that. + hot windows as the busiest, rounded up to whole batches of + 'oversubscription_factor', since those take no extra time. It is then + scored by the variance of the final free energy, given the samples each + window has and will collect in the cycles left, multiplied by the + compute of the whole run. The set with the lowest score is chosen. A + window only stops being hot once it has been found cold at two + consecutive rankings. Every window is hot if no set, including the + windows kept hot, is better than that. The ranking uses all of the data since the burn-in, since the slowest windows decorrelate over far longer than a checkpoint interval. diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index d1efc8e..fc7ffcc 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1496,7 +1496,7 @@ def test_rank_windows(): assert cache._hot_replicas == [4, 5, 6, 7] -def test_rank_windows_pads_to_workers(): +def test_rank_windows_balances_gpus(): """ Validate that the hot replicas are padded so that every GPU has the same number, and that every replica being hot is recorded as None. @@ -1835,6 +1835,43 @@ def test_rank_windows_keeps_slow_bond_distance_hot(): assert 0 not in runner._dynamics_cache._hot_replicas +@pytest.mark.parametrize("oversubscription_factor, num_hot", [(1, 4), (2, 8)]) +def test_rank_windows_fills_whole_batches(oversubscription_factor, num_hot): + """ + Validate that the hot replicas fill whole batches on every GPU, so that + with two replicas per GPU at once, each GPU gets two hot windows. + """ + runner = _ranking_runner( + 8, adaptive_min_samples=10, oversubscription_factor=oversubscription_factor + ) + + # The flagged pair 0 gives windows 0, 1 and 2, one on each of three GPUs. + _fill_ranking_data(runner, hard_pair=0) + _rank_twice(runner, num_gpus=4) + + hot = runner._dynamics_cache._hot_replicas + hot = list(range(8)) if hot is None else hot + assert len(hot) == num_hot + assert {0, 1, 2}.issubset(hot) + assert sorted(r % 4 for r in hot) == sorted(list(range(4)) * (num_hot // 4)) + + +def test_replica_devices_with_shared_contexts(): + """ + Validate that a replica runs on the GPU of the context it shares. + """ + from types import SimpleNamespace + + runner = object.__new__(RepexRunner) + runner._num_gpus = 2 + runner._num_replicas = 10 + + # A strided layout over 4 contexts. + runner._dynamics_cache = SimpleNamespace(slot_for=lambda r: r % 4) + + assert runner._replica_devices() == [(r % 4) % 2 for r in range(10)] + + def test_device_batches(): """ Validate that a batch never has more replicas on a GPU than the @@ -1848,14 +1885,15 @@ def test_device_batches(): runner._num_gpus = 3 runner._num_replicas = 25 runner._dynamics_cache = SimpleNamespace(slot_for=lambda r: r) + devices = runner._replica_devices() # Every replica is batched as consecutive groups of three. - assert runner._device_batches(list(range(25))) == [ + assert runner._device_batches(list(range(25)), devices) == [ list(range(i, min(i + 3, 25))) for i in range(0, 25, 3) ] hot = [6, 7, 8, 9, 10, 13, 14, 15, 16, 17, 20, 21, 22, 23, 24] - batches = runner._device_batches(hot) + batches = runner._device_batches(hot, devices) assert sorted(r for batch in batches for r in batch) == hot assert len(batches) == 5 for batch in batches: From c86b797ffd7ac7fdb22cc0ad844a0b07236e0581 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 11:55:49 +0100 Subject: [PATCH 18/32] Respect contexts per GPU when filling and reuse the dynamics pool. --- src/somd2/runner/_repex.py | 131 ++++++++++++++++++++++--------------- tests/runner/test_repex.py | 69 ++++++++++++++----- 2 files changed, 129 insertions(+), 71 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 8288cb0..049410c 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1980,40 +1980,40 @@ def run(self): # Minimise at each lambda value. Don't minimise on restart. if self._config.minimise and not self._is_restart: for batch in self._dynamics_batches(replica_devices): - with ThreadPoolExecutor(max_workers=num_workers) as executor: - try: - for success, index, e in executor.map( - self._minimise, - batch, - ): - if not success: - msg = f"Minimisation failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" - if self._config.minimisation_errors: - _logger.error(msg) - raise e - else: - _logger.warning(msg) - except KeyboardInterrupt: - _logger.error("Minimisation cancelled. Exiting.") - _sys.exit(1) + try: + for success, index, e in dynamics_executor.map( + self._minimise, + batch, + ): + if not success: + msg = f"Minimisation failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" + if self._config.minimisation_errors: + _logger.error(msg) + raise e + else: + _logger.warning(msg) + except KeyboardInterrupt: + dynamics_executor.shutdown(wait=False, cancel_futures=True) + _logger.error("Minimisation cancelled. Exiting.") + _sys.exit(1) # Equilibrate the system. if self._is_equilibration and not self._is_restart: for batch in self._dynamics_batches(replica_devices): - with ThreadPoolExecutor(max_workers=num_workers) as executor: - try: - for success, index, e in executor.map( - self._equilibrate, - batch, - ): - if not success: - _logger.error( - f"Equilibration failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" - ) - raise e - except KeyboardInterrupt: - _logger.error("Equilibration cancelled. Exiting.") - _sys.exit(1) + try: + for success, index, e in dynamics_executor.map( + self._equilibrate, + batch, + ): + if not success: + _logger.error( + f"Equilibration failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" + ) + raise e + except KeyboardInterrupt: + dynamics_executor.shutdown(wait=False, cancel_futures=True) + _logger.error("Equilibration cancelled. Exiting.") + _sys.exit(1) # Write a checkpoint immediately after equilibration so that a restart # after an early production crash doesn't need to re-equilibrate. @@ -2308,6 +2308,7 @@ def run(self): ) self._rank_windows( replica_devices, + self._device_capacity(), self._window_samples(), cycles - (i + 1), cold_rate, @@ -3021,12 +3022,31 @@ def _replica_devices(self): for r in range(self._num_replicas) ] + def _device_capacity(self): + """ + The number of replicas each GPU can propagate at once: the + oversubscription factor, or fewer if the GPU has fewer contexts. + + Returns + ------- + + list of int + The capacity of each device. + """ + contexts = [0] * self._num_gpus + for slot in range(self._num_slots): + contexts[slot % self._num_gpus] += 1 + return [max(1, min(self._config.oversubscription_factor, n)) for n in contexts] + def _device_batches(self, replicas, devices): """ Split replicas into batches that are propagated concurrently, each with at most 'oversubscription_factor' replicas on any GPU, so that every batch spreads its work over the GPUs. + Two replicas that share a slot can't be propagated at once, so the + replicas must come from a single pass, i.e. at most one per slot. + Parameters ---------- @@ -3057,7 +3077,7 @@ def _device_batches(self, replicas, devices): batches.append(sorted(batch)) return batches - def _dynamics_batches(self, devices, cycle=0): + def _dynamics_batches(self, devices): """ Yield batches of every replica that can be propagated concurrently: at most one per slot, and at most 'oversubscription_factor' per GPU. @@ -3068,16 +3088,13 @@ def _dynamics_batches(self, devices, cycle=0): devices: list of int The GPU that propagates each replica. - cycle: int - The index of the current cycle, which sets the traversal order. - Yields ------ list of int A batch of replicas. """ - for batch in self._replica_passes(cycle): + for batch in self._replica_passes(0): yield from self._device_batches(batch, devices) def _window_samples(self): @@ -3092,9 +3109,11 @@ def _window_samples(self): """ energy_frequency = self._config.energy_frequency.to("ps") return [ - 0 - if clock is None - else round(clock["_current_time"].to("ps") / energy_frequency) + ( + 0 + if clock is None + else round(clock["_current_time"].to("ps") / energy_frequency) + ) for clock in self._dynamics_cache._clocks ] @@ -3397,7 +3416,9 @@ def _is_unconverged(cls, fwd, rev, g_fwd, g_rev): drift = abs(estimates[0] - estimates[1]) return drift > 0.25 and drift > 2.0 * _np.hypot(*errors) - def _rank_windows(self, window_devices, window_samples, cycles_left, cold_rate): + def _rank_windows( + self, window_devices, device_capacity, window_samples, cycles_left, cold_rate + ): """ Rank the windows by the uncertainty they contribute to the free energy and choose the hot replicas for adaptive replica exchange. @@ -3411,15 +3432,15 @@ def _rank_windows(self, window_devices, window_samples, cycles_left, cold_rate): the bond distance. The remaining pairs are added in order of weight, with the windows either side of each so that exchange carries configurations into and out of them. Each set this gives is filled with - the cold windows owning the most variance, until every GPU has as many - hot windows as the busiest, rounded up to whole batches of - 'oversubscription_factor', since those take no extra time. It is then - scored by the variance of the final free energy, given the samples each - window has and will collect in the cycles left, multiplied by the - compute of the whole run. The set with the lowest score is chosen. A - window only stops being hot once it has been found cold at two - consecutive rankings. Every window is hot if no set, including the - windows kept hot, is better than that. + the cold windows owning the most variance, until every GPU runs as many + full batches as the busiest needs, since those windows take no extra + time. A GPU's batches hold as many replicas as it can propagate at once. + The set is then scored by the variance of the final free energy, given + the samples each window has and will collect in the cycles left, + multiplied by the compute of the whole run. The set with the lowest + score is chosen. A window only stops being hot once it has been found + cold at two consecutive rankings. Every window is hot if no set, + including the windows kept hot, is better than that. The ranking uses all of the data since the burn-in, since the slowest windows decorrelate over far longer than a checkpoint interval. @@ -3430,6 +3451,9 @@ def _rank_windows(self, window_devices, window_samples, cycles_left, cold_rate): window_devices: list of int The GPU that propagates each window. + device_capacity: list of int + The number of replicas each GPU can propagate at once. + window_samples: list of int The number of energy samples each window has collected. @@ -3522,20 +3546,19 @@ def pad(windows): owned[:-1] += fwd_variance owned[1:] += rev_variance - per_device = self._config.oversubscription_factor - def fill(windows): # Make the cold windows owning the most variance hot, until each GPU - # has as many hot windows as a whole number of batches on the busiest. + # fills as many batches as the busiest needs. counts = {d: 0 for d in window_devices} for r in windows: counts[window_devices[r]] += 1 - target = ceil(max(counts.values()) / per_device) * per_device + batches = max(ceil(n / device_capacity[d]) for d, n in counts.items()) filled = set(windows) for r in sorted(all_hot - windows, key=lambda r: -owned[r]): - if counts[window_devices[r]] < target: + d = window_devices[r] + if counts[d] < batches * device_capacity[d]: filled.add(r) - counts[window_devices[r]] += 1 + counts[d] += 1 return filled total_variance = owned.sum() diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index fc7ffcc..7b965cc 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -902,9 +902,8 @@ def test_repex_concurrent_slots(ethane_methanol): assert not runner._config.minimisation_constraints assert runner._config.equilibration_constraints - # Every batch must contain at most one replica per slot. - num_workers = runner._num_gpus * config["oversubscription_factor"] - for batch in runner._safe_batches(num_workers): + # Every batch of dynamics must contain at most one replica per slot. + for batch in runner._dynamics_batches(runner._replica_devices()): slots = [runner._dynamics_cache.slot_for(r) for r in batch] assert len(slots) == len(set(slots)), f"batch {batch} shares a slot" @@ -1431,27 +1430,32 @@ def _ranking_runner(num_lambda, **config): return runner -def _rank(runner, num_gpus, window_samples=None, cycles_left=10**6): +def _rank(runner, num_gpus, window_samples=None, cycles_left=10**6, capacity=None): """ Rank the windows, with window r on GPU r % num_gpus. By default every - window has the same samples and the run is far from its end, so that only - the sampling rates matter. + window has the same samples, the run is far from its end, so that only + the sampling rates matter, and each GPU can propagate as many replicas at + once as the oversubscription factor. """ num_lambda = runner._config.num_lambda if window_samples is None: window_samples = [200] * num_lambda + if capacity is None: + capacity = runner._config.oversubscription_factor devices = [r % num_gpus for r in range(num_lambda)] cold_rate = 1.0 / runner._config.adaptive_mix_interval - runner._rank_windows(devices, window_samples, cycles_left, cold_rate) + runner._rank_windows( + devices, [capacity] * num_gpus, window_samples, cycles_left, cold_rate + ) -def _rank_twice(runner, num_gpus): +def _rank_twice(runner, num_gpus, capacity=None): """ Rank the windows at two consecutive checkpoints, so that the windows found cold at both stop being hot. """ - _rank(runner, num_gpus=num_gpus) - _rank(runner, num_gpus=num_gpus) + _rank(runner, num_gpus=num_gpus, capacity=capacity) + _rank(runner, num_gpus=num_gpus, capacity=capacity) def _fill_ranking_data(runner, hard_pair, num_samples=200, seed=7): @@ -1835,11 +1839,19 @@ def test_rank_windows_keeps_slow_bond_distance_hot(): assert 0 not in runner._dynamics_cache._hot_replicas -@pytest.mark.parametrize("oversubscription_factor, num_hot", [(1, 4), (2, 8)]) -def test_rank_windows_fills_whole_batches(oversubscription_factor, num_hot): +@pytest.mark.parametrize( + "oversubscription_factor, capacity, num_hot", + [ + (1, 1, 4), + (2, 2, 8), + # Two replicas per GPU at once, but only one context on each. + (2, 1, 4), + ], +) +def test_rank_windows_fills_whole_batches(oversubscription_factor, capacity, num_hot): """ - Validate that the hot replicas fill whole batches on every GPU, so that - with two replicas per GPU at once, each GPU gets two hot windows. + Validate that the hot replicas fill whole batches on every GPU, of as many + replicas as the GPU can propagate at once. """ runner = _ranking_runner( 8, adaptive_min_samples=10, oversubscription_factor=oversubscription_factor @@ -1847,7 +1859,7 @@ def test_rank_windows_fills_whole_batches(oversubscription_factor, num_hot): # The flagged pair 0 gives windows 0, 1 and 2, one on each of three GPUs. _fill_ranking_data(runner, hard_pair=0) - _rank_twice(runner, num_gpus=4) + _rank_twice(runner, num_gpus=4, capacity=capacity) hot = runner._dynamics_cache._hot_replicas hot = list(range(8)) if hot is None else hot @@ -1856,6 +1868,28 @@ def test_rank_windows_fills_whole_batches(oversubscription_factor, num_hot): assert sorted(r % 4 for r in hot) == sorted(list(range(4)) * (num_hot // 4)) +@pytest.mark.parametrize( + "num_slots, num_gpus, oversubscription_factor, expected", + [ + (25, 3, 2, [2, 2, 2]), + # One context per GPU, so only one replica at once. + (4, 4, 2, [1, 1, 1, 1]), + (5, 4, 2, [2, 1, 1, 1]), + ], +) +def test_device_capacity(num_slots, num_gpus, oversubscription_factor, expected): + """ + Validate that a GPU propagates no more replicas at once than it has + contexts. + """ + runner = object.__new__(RepexRunner) + runner._config = Config(oversubscription_factor=oversubscription_factor) + runner._num_slots = num_slots + runner._num_gpus = num_gpus + + assert runner._device_capacity() == expected + + def test_replica_devices_with_shared_contexts(): """ Validate that a replica runs on the GPU of the context it shares. @@ -1959,10 +1993,11 @@ def test_adaptive_repex_ranking_arguments(ethane_methanol, monkeypatch): runner.run() # 12 cycles with a checkpoint every 4, and every window always hot. - assert [c[2] for c in calls] == [8, 4, 0] - for devices, samples, cycles_left, cold_rate in calls: + assert [c[3] for c in calls] == [8, 4, 0] + for devices, capacity, samples, cycles_left, cold_rate in calls: # A context per replica, assigned to the GPUs in turn. assert devices == [r % runner._num_gpus for r in range(num_lambda)] + assert capacity == [1] * runner._num_gpus first = 12 - cycles_left assert samples == [first] * num_lambda if cycles_left > 0: From c7ec9e8e82dcc2cb82b83a4df7b2c297046680d9 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 12:05:07 +0100 Subject: [PATCH 19/32] Skip filling with shared contexts and wait for blocks on error. --- src/somd2/runner/_repex.py | 79 +++++++++++++++++++++++--------------- tests/runner/test_repex.py | 66 ++++++++++++++++++++++++++++--- 2 files changed, 110 insertions(+), 35 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 049410c..fc6dc8b 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -475,14 +475,9 @@ def _create_dynamics( # Per-device memory tracking for estimation. device_mem = {} - # Work out how many slots are assigned to each device. - # Slots are assigned round-robin, so the first (num_slots % num_gpus) - # devices get one extra slot. - base = floor(self._num_slots / num_gpus) - remainder = self._num_slots % num_gpus - contexts_per_device = [ - base + (1 if i < remainder else 0) for i in range(num_gpus) - ] + # The device of each slot. Slots are assigned round-robin. + self._slot_devices = [i % num_gpus for i in range(self._num_slots)] + contexts_per_device = [self._slot_devices.count(d) for d in range(num_gpus)] # Record the ghost file for each replica. A slot writes to the file of # whichever replica it currently hosts. @@ -517,7 +512,7 @@ def _create_dynamics( build_lam = lambdas[self._constraint_lambda_index] # Work out the device index. - device = i % num_gpus + device = self._slot_devices[i] # Record baseline memory before the first slot on this device. if device not in device_mem: @@ -1901,6 +1896,20 @@ def run(self): """ Run the replica exchange simulation. """ + self._executors = [] + try: + self._run() + except Exception: + # Wait for running blocks, so that nothing still uses the contexts. + for executor in self._executors: + executor.shutdown(wait=True, cancel_futures=True) + raise + + def _run(self): + """ + Run the replica exchange simulation, using thread pools that run() + shuts down if an error is raised. + """ from math import ceil from time import time @@ -1970,12 +1979,15 @@ def run(self): # Persistent thread pools, reused across every batch and cycle. dynamics_executor = ThreadPoolExecutor(max_workers=num_workers) checkpoint_executor = ThreadPoolExecutor(max_workers=num_checkpoint_workers) + self._executors = [dynamics_executor, checkpoint_executor] # Create the replica list. replica_list = list(range(self._config.num_lambda)) - # The GPU that propagates each replica, which is fixed for the run. + # The GPU that propagates each replica, and how many replicas each GPU + # can propagate at once, which are fixed for the run. replica_devices = self._replica_devices() + device_capacity = self._device_capacity() # Minimise at each lambda value. Don't minimise on restart. if self._config.minimise and not self._is_restart: @@ -2308,7 +2320,7 @@ def run(self): ) self._rank_windows( replica_devices, - self._device_capacity(), + device_capacity, self._window_samples(), cycles - (i + 1), cold_rate, @@ -2665,7 +2677,7 @@ def _minimise(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = slot % self._num_gpus + device = self._dynamics_cache._slot_devices[slot] # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -2775,7 +2787,7 @@ def _equilibrate(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = slot % self._num_gpus + device = self._dynamics_cache._slot_devices[slot] # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -2821,7 +2833,7 @@ def _equilibrate(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = slot % self._num_gpus + device = self._dynamics_cache._slot_devices[slot] # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -3017,8 +3029,9 @@ def _replica_devices(self): list of int The device index for each replica. """ + slot_devices = self._dynamics_cache._slot_devices return [ - self._dynamics_cache.slot_for(r) % self._num_gpus + slot_devices[self._dynamics_cache.slot_for(r)] for r in range(self._num_replicas) ] @@ -3033,10 +3046,11 @@ def _device_capacity(self): list of int The capacity of each device. """ - contexts = [0] * self._num_gpus - for slot in range(self._num_slots): - contexts[slot % self._num_gpus] += 1 - return [max(1, min(self._config.oversubscription_factor, n)) for n in contexts] + slot_devices = self._dynamics_cache._slot_devices + return [ + max(1, min(self._config.oversubscription_factor, slot_devices.count(d))) + for d in range(self._num_gpus) + ] def _device_batches(self, replicas, devices): """ @@ -3431,16 +3445,18 @@ def _rank_windows( differences or, for a perturbation that breaks or makes a single bond, the bond distance. The remaining pairs are added in order of weight, with the windows either side of each so that exchange carries - configurations into and out of them. Each set this gives is filled with - the cold windows owning the most variance, until every GPU runs as many - full batches as the busiest needs, since those windows take no extra - time. A GPU's batches hold as many replicas as it can propagate at once. - The set is then scored by the variance of the final free energy, given - the samples each window has and will collect in the cycles left, - multiplied by the compute of the whole run. The set with the lowest - score is chosen. A window only stops being hot once it has been found - cold at two consecutive rankings. Every window is hot if no set, - including the windows kept hot, is better than that. + configurations into and out of them. With a context per replica, each + set this gives is filled with the cold windows owning the most + variance, until every GPU runs as many full batches as the busiest + needs, since those windows take no extra time. A GPU's batches hold as + many replicas as it can propagate at once. With shared contexts the set + isn't filled, since an added window can need another pass. The set is + then scored by the variance of the final free energy, given the samples + each window has and will collect in the cycles left, multiplied by the + compute of the whole run. The set with the lowest score is chosen. A + window only stops being hot once it has been found cold at two + consecutive rankings. Every window is hot if no set, including the + windows kept hot, is better than that. The ranking uses all of the data since the burn-in, since the slowest windows decorrelate over far longer than a checkpoint interval. @@ -3548,7 +3564,10 @@ def pad(windows): def fill(windows): # Make the cold windows owning the most variance hot, until each GPU - # fills as many batches as the busiest needs. + # fills as many batches as the busiest needs. With shared contexts + # an added window can need another pass, so it isn't free. + if self._is_cached: + return set(windows) counts = {d: 0 for d in window_devices} for r in windows: counts[window_devices[r]] += 1 diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 7b965cc..98b5a81 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -902,8 +902,12 @@ def test_repex_concurrent_slots(ethane_methanol): assert not runner._config.minimisation_constraints assert runner._config.equilibration_constraints - # Every batch of dynamics must contain at most one replica per slot. - for batch in runner._dynamics_batches(runner._replica_devices()): + # Every batch of dynamics, and of commits, must contain at most one + # replica per slot. + num_workers = runner._num_gpus * config["oversubscription_factor"] + batches = list(runner._dynamics_batches(runner._replica_devices())) + batches += list(runner._safe_batches(num_workers)) + for batch in batches: slots = [runner._dynamics_cache.slot_for(r) for r in batch] assert len(slots) == len(set(slots)), f"batch {batch} shares a slot" @@ -1419,6 +1423,7 @@ def _ranking_runner(num_lambda, **config): runner._config = Config(num_lambda=num_lambda, **config) runner._lambda_values = list(np.linspace(0, 1, num_lambda)) runner._ring_pair = None + runner._is_cached = False zeros = np.zeros((num_lambda, num_lambda)) runner._dynamics_cache = SimpleNamespace( get_proposed=lambda: zeros, @@ -1868,6 +1873,49 @@ def test_rank_windows_fills_whole_batches(oversubscription_factor, capacity, num assert sorted(r % 4 for r in hot) == sorted(list(range(4)) * (num_hot // 4)) +def test_rank_windows_unfilled_with_shared_contexts(): + """ + Validate that the hot replicas aren't filled when contexts are shared, + since a window added on a busy context needs another pass. + """ + runner = _ranking_runner(8, adaptive_min_samples=10) + runner._is_cached = True + + # The flagged pair 0 gives windows 0, 1 and 2, one on each of three GPUs. + _fill_ranking_data(runner, hard_pair=0) + _rank_twice(runner, num_gpus=4) + + assert runner._dynamics_cache._hot_replicas == [0, 1, 2] + + +def test_run_waits_for_blocks_on_error(monkeypatch): + """ + Validate that an error during a run shuts down the thread pools and waits + for their running blocks before it is raised. + """ + + class Executor: + def __init__(self): + self.shutdown_args = None + + def shutdown(self, wait=True, cancel_futures=False): + self.shutdown_args = (wait, cancel_futures) + + executors = [Executor(), Executor()] + runner = object.__new__(RepexRunner) + + def fail(): + runner._executors = executors + raise RuntimeError("dynamics failed") + + monkeypatch.setattr(runner, "_run", fail) + + with pytest.raises(RuntimeError, match="dynamics failed"): + runner.run() + + assert [e.shutdown_args for e in executors] == [(True, True)] * 2 + + @pytest.mark.parametrize( "num_slots, num_gpus, oversubscription_factor, expected", [ @@ -1882,10 +1930,14 @@ def test_device_capacity(num_slots, num_gpus, oversubscription_factor, expected) Validate that a GPU propagates no more replicas at once than it has contexts. """ + from types import SimpleNamespace + runner = object.__new__(RepexRunner) runner._config = Config(oversubscription_factor=oversubscription_factor) - runner._num_slots = num_slots runner._num_gpus = num_gpus + runner._dynamics_cache = SimpleNamespace( + _slot_devices=[s % num_gpus for s in range(num_slots)] + ) assert runner._device_capacity() == expected @@ -1901,7 +1953,9 @@ def test_replica_devices_with_shared_contexts(): runner._num_replicas = 10 # A strided layout over 4 contexts. - runner._dynamics_cache = SimpleNamespace(slot_for=lambda r: r % 4) + runner._dynamics_cache = SimpleNamespace( + slot_for=lambda r: r % 4, _slot_devices=[s % 2 for s in range(4)] + ) assert runner._replica_devices() == [(r % 4) % 2 for r in range(10)] @@ -1918,7 +1972,9 @@ def test_device_batches(): runner._config = Config(num_lambda=25) runner._num_gpus = 3 runner._num_replicas = 25 - runner._dynamics_cache = SimpleNamespace(slot_for=lambda r: r) + runner._dynamics_cache = SimpleNamespace( + slot_for=lambda r: r, _slot_devices=[s % 3 for s in range(25)] + ) devices = runner._replica_devices() # Every replica is batched as consecutive groups of three. From c53c0f0ccdb8add0c4128b14cd86588a0fb8c8f8 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 12:35:06 +0100 Subject: [PATCH 20/32] Finish checkpoint writes under the lock and tidy the runner pools. --- src/somd2/runner/_repex.py | 207 ++++++++++++++++++++++++------------- tests/runner/test_repex.py | 108 +++++++++++++------ 2 files changed, 212 insertions(+), 103 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index fc6dc8b..5a8f9bb 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -477,7 +477,7 @@ def _create_dynamics( # The device of each slot. Slots are assigned round-robin. self._slot_devices = [i % num_gpus for i in range(self._num_slots)] - contexts_per_device = [self._slot_devices.count(d) for d in range(num_gpus)] + contexts_per_device = [self.contexts_on(d) for d in range(num_gpus)] # Record the ghost file for each replica. A slot writes to the file of # whichever replica it currently hosts. @@ -756,6 +756,42 @@ def slot_for(self, replica): """ return self._replica_slot[replica] + def device_for(self, slot): + """ + Return the index of the GPU that a slot's context is on. + + Parameters + ---------- + + slot: int + The index of the slot. + + Returns + ------- + + int + The device index. + """ + return self._slot_devices[slot] + + def contexts_on(self, device): + """ + Return the number of contexts (slots) on a GPU. + + Parameters + ---------- + + device: int + The device index. + + Returns + ------- + + int + The number of contexts. + """ + return self._slot_devices.count(device) + def get(self, slot): """ Get the dynamics object (and GCMC sampler) for a given slot. @@ -1896,25 +1932,52 @@ def run(self): """ Run the replica exchange simulation. """ - self._executors = [] + from concurrent.futures import ThreadPoolExecutor + + # Store the number of concurrent workers. + num_workers = self._num_gpus * self._config.oversubscription_factor + + # Store the number of workers to use for checkpointing. + if self._config.num_checkpoint_workers is None: + num_checkpoint_workers = num_workers + else: + num_checkpoint_workers = min( + self._config.num_checkpoint_workers, num_workers + ) + + # Persistent thread pools, reused across every batch and cycle. + dynamics_executor = ThreadPoolExecutor(max_workers=num_workers) + checkpoint_executor = ThreadPoolExecutor(max_workers=num_checkpoint_workers) + try: - self._run() + self._run(dynamics_executor, checkpoint_executor, num_checkpoint_workers) except Exception: # Wait for running blocks, so that nothing still uses the contexts. - for executor in self._executors: + for executor in (dynamics_executor, checkpoint_executor): executor.shutdown(wait=True, cancel_futures=True) raise - def _run(self): + def _run(self, dynamics_executor, checkpoint_executor, num_checkpoint_workers): """ - Run the replica exchange simulation, using thread pools that run() - shuts down if an error is raised. + Run the replica exchange simulation. + + Parameters + ---------- + + dynamics_executor: concurrent.futures.ThreadPoolExecutor + The thread pool for dynamics, with a worker for each replica that + can be propagated at once. + + checkpoint_executor: concurrent.futures.ThreadPoolExecutor + The thread pool for writing checkpoint files. + + num_checkpoint_workers: int + The number of workers in the checkpoint thread pool. """ from math import ceil from time import time - from concurrent.futures import ThreadPoolExecutor from itertools import repeat # Record the start time. @@ -1962,25 +2025,9 @@ def _run(self): num_blocks = 1 rem = 0 - # Store the number of concurrent workers. - num_workers = self._num_gpus * self._config.oversubscription_factor - - # Store the number of workers to use for checkpointing. - if self._config.num_checkpoint_workers is None: - num_checkpoint_workers = num_workers - else: - num_checkpoint_workers = min( - self._config.num_checkpoint_workers, num_workers - ) - # Work out the required number of batches. num_checkpoint_batches = ceil(self._config.num_lambda / num_checkpoint_workers) - # Persistent thread pools, reused across every batch and cycle. - dynamics_executor = ThreadPoolExecutor(max_workers=num_workers) - checkpoint_executor = ThreadPoolExecutor(max_workers=num_checkpoint_workers) - self._executors = [dynamics_executor, checkpoint_executor] - # Create the replica list. replica_list = list(range(self._config.num_lambda)) @@ -2058,27 +2105,31 @@ def _run(self): replicas = replica_list[ j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] + # Let every write finish under the lock before raising. try: - for index, error in checkpoint_executor.map( - self._checkpoint, - replicas, - repeat(self._lambda_values), - repeat(-1), - repeat(cycles), - ): - if error is not None: - msg = ( - f"Post-equilibration checkpoint failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}:\n{error}" - ) - _logger.error(msg) - raise error + results = list( + checkpoint_executor.map( + self._checkpoint, + replicas, + repeat(self._lambda_values), + repeat(-1), + repeat(cycles), + ) + ) except KeyboardInterrupt: checkpoint_executor.shutdown(wait=False, cancel_futures=True) _logger.error( "Post-equilibration checkpoint cancelled. Exiting." ) _sys.exit(1) + for index, error in results: + if error is not None: + msg = ( + f"Post-equilibration checkpoint failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}:\n{error}" + ) + _logger.error(msg) + raise error # Current block number. block = self._start_block @@ -2339,23 +2390,27 @@ def _run(self): j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] + # Let every write finish under the lock before raising. try: - for index, error in checkpoint_executor.map( - self._backup_checkpoint, - replicas, - ): - if error: - _logger.error( - f"Backup failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error + results = list( + checkpoint_executor.map( + self._backup_checkpoint, + replicas, + ) + ) except KeyboardInterrupt: checkpoint_executor.shutdown( wait=False, cancel_futures=True ) _logger.error("Backup cancelled. Exiting.") _sys.exit(1) + for index, error in results: + if error: + _logger.error( + f"Backup failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}: {error}" + ) + raise error # Now write the new checkpoint files. for j in range(num_checkpoint_batches): @@ -2364,28 +2419,32 @@ def _run(self): j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] + # Let every write finish under the lock before raising. try: - for index, error in checkpoint_executor.map( - self._checkpoint, - replicas, - repeat(self._lambda_values), - repeat(block), - repeat(num_blocks + int(rem > 0)), - repeat(i == cycles - 1), - [num_samples[r] for r in replicas], - ): - if error: - _logger.error( - f"Checkpoint failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error + results = list( + checkpoint_executor.map( + self._checkpoint, + replicas, + repeat(self._lambda_values), + repeat(block), + repeat(num_blocks + int(rem > 0)), + repeat(i == cycles - 1), + [num_samples[r] for r in replicas], + ) + ) except KeyboardInterrupt: checkpoint_executor.shutdown( wait=False, cancel_futures=True ) _logger.error("Checkpoint cancelled. Exiting.") _sys.exit(1) + for index, error in results: + if error: + _logger.error( + f"Checkpoint failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}: {error}" + ) + raise error num_samples = [0] * num_replicas @@ -2677,7 +2736,7 @@ def _minimise(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = self._dynamics_cache._slot_devices[slot] + device = self._dynamics_cache.device_for(slot) # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -2787,7 +2846,7 @@ def _equilibrate(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = self._dynamics_cache._slot_devices[slot] + device = self._dynamics_cache.device_for(slot) # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -2833,7 +2892,7 @@ def _equilibrate(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = self._dynamics_cache._slot_devices[slot] + device = self._dynamics_cache.device_for(slot) # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -3029,11 +3088,8 @@ def _replica_devices(self): list of int The device index for each replica. """ - slot_devices = self._dynamics_cache._slot_devices - return [ - slot_devices[self._dynamics_cache.slot_for(r)] - for r in range(self._num_replicas) - ] + cache = self._dynamics_cache + return [cache.device_for(cache.slot_for(r)) for r in range(self._num_replicas)] def _device_capacity(self): """ @@ -3046,9 +3102,14 @@ def _device_capacity(self): list of int The capacity of each device. """ - slot_devices = self._dynamics_cache._slot_devices return [ - max(1, min(self._config.oversubscription_factor, slot_devices.count(d))) + max( + 1, + min( + self._config.oversubscription_factor, + self._dynamics_cache.contexts_on(d), + ), + ) for d in range(self._num_gpus) ] diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 98b5a81..abf83c6 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -718,6 +718,58 @@ def acquire(self, *args, **kwargs): assert len(acquisitions) == 2 +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_checkpoint_failure_keeps_lock(ethane_methanol, monkeypatch): + """ + Validate that when one replica's checkpoint fails, the others finish + writing while the file lock is still held, before the error is raised. + """ + import time + + from filelock import FileLock, Timeout + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "4fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": 4, + "replica_exchange": True, + # Enough workers to checkpoint every replica at once. + "oversubscription_factor": 4, + "num_checkpoint_workers": 4, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + held = [] + + def checkpoint(index, *args, **kwargs): + if index == 0: + return index, RuntimeError("checkpoint failed") + time.sleep(0.5) + lock = FileLock(runner._lock_file) + try: + lock.acquire(timeout=0) + lock.release() + held.append(False) + except Timeout: + held.append(True) + return index, None + + monkeypatch.setattr(runner, "_checkpoint", checkpoint) + + with pytest.raises(RuntimeError, match="checkpoint failed"): + runner.run() + + assert held == [True, True, True] + + @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") @pytest.mark.parametrize( "runtime, checkpoint_frequency, expected", @@ -1890,22 +1942,19 @@ def test_rank_windows_unfilled_with_shared_contexts(): def test_run_waits_for_blocks_on_error(monkeypatch): """ - Validate that an error during a run shuts down the thread pools and waits - for their running blocks before it is raised. + Validate that an error during a run waits for the blocks still running + in the thread pools before it is raised. """ + import time - class Executor: - def __init__(self): - self.shutdown_args = None - - def shutdown(self, wait=True, cancel_futures=False): - self.shutdown_args = (wait, cancel_futures) - - executors = [Executor(), Executor()] runner = object.__new__(RepexRunner) + runner._config = Config() + runner._num_gpus = 1 - def fail(): - runner._executors = executors + finished = [] + + def fail(dynamics_executor, checkpoint_executor, num_checkpoint_workers): + dynamics_executor.submit(lambda: (time.sleep(0.2), finished.append(True))) raise RuntimeError("dynamics failed") monkeypatch.setattr(runner, "_run", fail) @@ -1913,7 +1962,20 @@ def fail(): with pytest.raises(RuntimeError, match="dynamics failed"): runner.run() - assert [e.shutdown_args for e in executors] == [(True, True)] * 2 + assert finished == [True] + + +def _layout(num_replicas, num_slots, num_gpus): + """ + A dynamics cache holding only a strided layout of replicas over slots, and + slots over GPUs. + """ + from somd2.runner._repex import DynamicsCache + + cache = object.__new__(DynamicsCache) + cache._replica_slot = [r % num_slots for r in range(num_replicas)] + cache._slot_devices = [s % num_gpus for s in range(num_slots)] + return cache @pytest.mark.parametrize( @@ -1930,14 +1992,10 @@ def test_device_capacity(num_slots, num_gpus, oversubscription_factor, expected) Validate that a GPU propagates no more replicas at once than it has contexts. """ - from types import SimpleNamespace - runner = object.__new__(RepexRunner) runner._config = Config(oversubscription_factor=oversubscription_factor) runner._num_gpus = num_gpus - runner._dynamics_cache = SimpleNamespace( - _slot_devices=[s % num_gpus for s in range(num_slots)] - ) + runner._dynamics_cache = _layout(num_slots, num_slots, num_gpus) assert runner._device_capacity() == expected @@ -1946,16 +2004,10 @@ def test_replica_devices_with_shared_contexts(): """ Validate that a replica runs on the GPU of the context it shares. """ - from types import SimpleNamespace - runner = object.__new__(RepexRunner) runner._num_gpus = 2 runner._num_replicas = 10 - - # A strided layout over 4 contexts. - runner._dynamics_cache = SimpleNamespace( - slot_for=lambda r: r % 4, _slot_devices=[s % 2 for s in range(4)] - ) + runner._dynamics_cache = _layout(10, 4, 2) assert runner._replica_devices() == [(r % 4) % 2 for r in range(10)] @@ -1966,15 +2018,11 @@ def test_device_batches(): oversubscription factor, so that hot replicas that aren't consecutive still spread over the GPUs. """ - from types import SimpleNamespace - runner = object.__new__(RepexRunner) runner._config = Config(num_lambda=25) runner._num_gpus = 3 runner._num_replicas = 25 - runner._dynamics_cache = SimpleNamespace( - slot_for=lambda r: r, _slot_devices=[s % 3 for s in range(25)] - ) + runner._dynamics_cache = _layout(25, 25, 3) devices = runner._replica_devices() # Every replica is batched as consecutive groups of three. From 86e954f02462a0235a40cb3c9f46ba7d13b63281 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 13:55:04 +0100 Subject: [PATCH 21/32] Run every checkpoint step through one helper that finishes all writes. --- README.md | 17 ++- src/somd2/runner/_repex.py | 236 ++++++++++++------------------------- tests/runner/test_repex.py | 9 +- 3 files changed, 90 insertions(+), 172 deletions(-) diff --git a/README.md b/README.md index b8b32b1..bd79b8d 100644 --- a/README.md +++ b/README.md @@ -452,12 +452,17 @@ The following are always hot: The windows either side of each hot window are also made hot, so that exchange carries configurations into and out of them, and a window only stops being hot after it has been found cold at two consecutive checkpoints, so no window is -cold until the second checkpoint after the burn-in. The number of hot windows -is rounded up to a multiple of the number of GPUs multiplied by the -oversubscription factor, using the cold windows that contribute the most -uncertainty, so that the GPUs are evenly loaded. When contexts are re-used, -replicas are assigned to them in a strided rather than contiguous pattern, so -that the hot windows are spread evenly over the GPUs. +cold until the second checkpoint after the burn-in. + +The replicas propagated in each cycle are split into batches with no more than +`--oversubscription-factor` replicas on any GPU, so that every GPU has work in +every batch. With a context per replica, the hot windows are topped up with the +cold windows that contribute the most uncertainty until every GPU has as many +hot windows as the busiest, since those take no extra time. When contexts are +re-used, replicas are assigned to them in a strided rather than contiguous +pattern, so that the hot windows are spread evenly over the GPUs, and the hot +windows aren't topped up, since a window added to a busy context would need +another pass. The `--runtime` is the simulation time of the hot windows, so the remaining windows are sampled for less time and the windows have different numbers of diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 5a8f9bb..71d966e 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1787,39 +1787,49 @@ def _replica_passes(self, cycle, replicas=None): return passes - def _safe_batches(self, num_workers, cycle=0): + def _map_replicas(self, executor, step, func, replicas, *args): """ - Yield batches of replicas that can be processed concurrently. + Run a step for each replica in a thread pool, letting every call + finish before raising the first error, so that files written under a + lock are all written before it is released. Two replicas that share a slot must never be processed at the same - time, since they would be using the same dynamics object and OpenMM - context. Batches are therefore taken from within a single pass, which - holds at most one replica per slot, and then split by the number of - workers. - - With one slot per replica there is a single pass containing every - replica, so this is just a split by the number of workers. + time, so the replicas must come from a single pass, unless the step + doesn't use the slot's context. Parameters ---------- - num_workers: int - The maximum number of replicas in a batch. + executor: concurrent.futures.ThreadPoolExecutor + The thread pool to run the step in. - cycle: int - The index of the current cycle, which sets the traversal order. + step: str + The name of the step, for logging. - Yields - ------ + func: callable + The step, which returns the index of the replica and the exception + raised, if any. - list of int - A batch of replicas that is safe to process concurrently. - """ - from math import ceil + replicas: list of int + The replicas to run the step for. - for batch in self._replica_passes(cycle): - for i in range(ceil(len(batch) / num_workers)): - yield batch[i * num_workers : (i + 1) * num_workers] + args: iterable + The remaining arguments, an iterable of one per replica for each. + """ + try: + outcomes = list(executor.map(func, replicas, *args)) + except KeyboardInterrupt: + executor.shutdown(wait=False, cancel_futures=True) + _logger.error(f"{step} cancelled. Exiting.") + _sys.exit(1) + + for index, error in outcomes: + if error is not None: + _logger.error( + f"{step} failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}: {error}" + ) + raise error def _set_num_slots(self): """ @@ -1949,15 +1959,19 @@ def run(self): dynamics_executor = ThreadPoolExecutor(max_workers=num_workers) checkpoint_executor = ThreadPoolExecutor(max_workers=num_checkpoint_workers) + executors = (dynamics_executor, checkpoint_executor) try: - self._run(dynamics_executor, checkpoint_executor, num_checkpoint_workers) + self._run(dynamics_executor, checkpoint_executor) except Exception: # Wait for running blocks, so that nothing still uses the contexts. - for executor in (dynamics_executor, checkpoint_executor): + for executor in executors: executor.shutdown(wait=True, cancel_futures=True) raise + else: + for executor in executors: + executor.shutdown(wait=True) - def _run(self, dynamics_executor, checkpoint_executor, num_checkpoint_workers): + def _run(self, dynamics_executor, checkpoint_executor): """ Run the replica exchange simulation. @@ -1970,9 +1984,6 @@ def _run(self, dynamics_executor, checkpoint_executor, num_checkpoint_workers): checkpoint_executor: concurrent.futures.ThreadPoolExecutor The thread pool for writing checkpoint files. - - num_checkpoint_workers: int - The number of workers in the checkpoint thread pool. """ from math import ceil @@ -2025,9 +2036,6 @@ def _run(self, dynamics_executor, checkpoint_executor, num_checkpoint_workers): num_blocks = 1 rem = 0 - # Work out the required number of batches. - num_checkpoint_batches = ceil(self._config.num_lambda / num_checkpoint_workers) - # Create the replica list. replica_list = list(range(self._config.num_lambda)) @@ -2081,55 +2089,25 @@ def _run(self, dynamics_executor, checkpoint_executor, num_checkpoint_workers): # resident so that every checkpoint file can still be written # under a single lock below. if self._is_cached: - for batch in self._safe_batches(num_checkpoint_workers): - try: - for index, error in checkpoint_executor.map( - self._load_and_commit, - batch, - ): - if error is not None: - msg = ( - f"Post-equilibration commit failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}:\n{error}" - ) - _logger.error(msg) - raise error - except KeyboardInterrupt: - checkpoint_executor.shutdown(wait=False, cancel_futures=True) - _logger.error("Post-equilibration commit cancelled. Exiting.") - _sys.exit(1) + for batch in self._replica_passes(0): + self._map_replicas( + checkpoint_executor, + "Post-equilibration commit", + self._load_and_commit, + batch, + ) lock = _FileLock(self._lock_file) with lock.acquire(timeout=self._config.timeout.to("seconds")): - for j in range(num_checkpoint_batches): - replicas = replica_list[ - j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers - ] - # Let every write finish under the lock before raising. - try: - results = list( - checkpoint_executor.map( - self._checkpoint, - replicas, - repeat(self._lambda_values), - repeat(-1), - repeat(cycles), - ) - ) - except KeyboardInterrupt: - checkpoint_executor.shutdown(wait=False, cancel_futures=True) - _logger.error( - "Post-equilibration checkpoint cancelled. Exiting." - ) - _sys.exit(1) - for index, error in results: - if error is not None: - msg = ( - f"Post-equilibration checkpoint failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}:\n{error}" - ) - _logger.error(msg) - raise error + self._map_replicas( + checkpoint_executor, + "Post-equilibration checkpoint", + self._checkpoint, + replica_list, + repeat(self._lambda_values), + repeat(-1), + repeat(cycles), + ) # Current block number. block = self._start_block @@ -2296,28 +2274,9 @@ def _run(self, dynamics_executor, checkpoint_executor, num_checkpoint_workers): # context per replica everything is still resident once the # cycle finishes, so committing is left to _checkpoint(). if do_checkpoint and self._is_cached: - for j in range(ceil(len(batch) / num_checkpoint_workers)): - replicas = batch[ - j * num_checkpoint_workers : (j + 1) - * num_checkpoint_workers - ] - try: - for index, error in checkpoint_executor.map( - self._commit_replica, - replicas, - ): - if error: - _logger.error( - f"Commit failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error - except KeyboardInterrupt: - checkpoint_executor.shutdown( - wait=False, cancel_futures=True - ) - _logger.error("Commit cancelled. Exiting.") - _sys.exit(1) + self._map_replicas( + checkpoint_executor, "Commit", self._commit_replica, batch + ) for index in active: last_cycle[index] = i @@ -2384,67 +2343,25 @@ def _run(self, dynamics_executor, checkpoint_executor, num_checkpoint_workers): # in a consistent state if read by another process. with lock.acquire(timeout=self._config.timeout.to("seconds")): # First backup existing checkpoint files. - for j in range(num_checkpoint_batches): - # Get the indices of the replicas in this batch. - replicas = replica_list[ - j * num_checkpoint_workers : (j + 1) - * num_checkpoint_workers - ] - # Let every write finish under the lock before raising. - try: - results = list( - checkpoint_executor.map( - self._backup_checkpoint, - replicas, - ) - ) - except KeyboardInterrupt: - checkpoint_executor.shutdown( - wait=False, cancel_futures=True - ) - _logger.error("Backup cancelled. Exiting.") - _sys.exit(1) - for index, error in results: - if error: - _logger.error( - f"Backup failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error + self._map_replicas( + checkpoint_executor, + "Backup", + self._backup_checkpoint, + replica_list, + ) # Now write the new checkpoint files. - for j in range(num_checkpoint_batches): - # Get the indices of the replicas in this batch. - replicas = replica_list[ - j * num_checkpoint_workers : (j + 1) - * num_checkpoint_workers - ] - # Let every write finish under the lock before raising. - try: - results = list( - checkpoint_executor.map( - self._checkpoint, - replicas, - repeat(self._lambda_values), - repeat(block), - repeat(num_blocks + int(rem > 0)), - repeat(i == cycles - 1), - [num_samples[r] for r in replicas], - ) - ) - except KeyboardInterrupt: - checkpoint_executor.shutdown( - wait=False, cancel_futures=True - ) - _logger.error("Checkpoint cancelled. Exiting.") - _sys.exit(1) - for index, error in results: - if error: - _logger.error( - f"Checkpoint failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error + self._map_replicas( + checkpoint_executor, + "Checkpoint", + self._checkpoint, + replica_list, + repeat(self._lambda_values), + repeat(block), + repeat(num_blocks + int(rem > 0)), + repeat(i == cycles - 1), + num_samples, + ) num_samples = [0] * num_replicas @@ -2459,9 +2376,6 @@ def _run(self, dynamics_executor, checkpoint_executor, num_checkpoint_workers): # Update the block number. block += 1 - dynamics_executor.shutdown(wait=True) - checkpoint_executor.shutdown(wait=True) - # Record the end time for the production block. prod_end = time() diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index abf83c6..73c0b33 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -954,11 +954,10 @@ def test_repex_concurrent_slots(ethane_methanol): assert not runner._config.minimisation_constraints assert runner._config.equilibration_constraints - # Every batch of dynamics, and of commits, must contain at most one - # replica per slot. - num_workers = runner._num_gpus * config["oversubscription_factor"] + # Every batch of dynamics, and every pass of commits, must contain at + # most one replica per slot. batches = list(runner._dynamics_batches(runner._replica_devices())) - batches += list(runner._safe_batches(num_workers)) + batches += runner._replica_passes(0) for batch in batches: slots = [runner._dynamics_cache.slot_for(r) for r in batch] assert len(slots) == len(set(slots)), f"batch {batch} shares a slot" @@ -1953,7 +1952,7 @@ def test_run_waits_for_blocks_on_error(monkeypatch): finished = [] - def fail(dynamics_executor, checkpoint_executor, num_checkpoint_workers): + def fail(dynamics_executor, checkpoint_executor): dynamics_executor.submit(lambda: (time.sleep(0.2), finished.append(True))) raise RuntimeError("dynamics failed") From e37087cf4b671ffcd9cacd0871e77ed362603c58 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 14:26:20 +0100 Subject: [PATCH 22/32] Gather every checkpoint result before raising and clarify GPU batching --- README.md | 18 +++++++++--------- src/somd2/runner/_repex.py | 26 +++++++++++++++++++------- tests/runner/test_repex.py | 33 ++++++++++++++++++++++++++++++++- 3 files changed, 60 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index bd79b8d..cab3888 100644 --- a/README.md +++ b/README.md @@ -454,15 +454,15 @@ carries configurations into and out of them, and a window only stops being hot after it has been found cold at two consecutive checkpoints, so no window is cold until the second checkpoint after the burn-in. -The replicas propagated in each cycle are split into batches with no more than -`--oversubscription-factor` replicas on any GPU, so that every GPU has work in -every batch. With a context per replica, the hot windows are topped up with the -cold windows that contribute the most uncertainty until every GPU has as many -hot windows as the busiest, since those take no extra time. When contexts are -re-used, replicas are assigned to them in a strided rather than contiguous -pattern, so that the hot windows are spread evenly over the GPUs, and the hot -windows aren't topped up, since a window added to a busy context would need -another pass. +The replicas propagated in each cycle are split into batches with no more +replicas on any GPU than it can propagate at once: `--oversubscription-factor`, +or fewer if the GPU has fewer contexts. With a context per replica, the hot +windows are topped up with the cold windows that contribute the most +uncertainty until every GPU fills as many batches as the busiest GPU needs, +since those windows take no extra time. When contexts are re-used, replicas +are assigned to them in a strided rather than contiguous pattern, so that the +hot windows are spread evenly over the GPUs, and the hot windows aren't topped +up, since a window added to a busy context would need another pass. The `--runtime` is the simulation time of the hot windows, so the remaining windows are sampled for less time and the windows have different numbers of diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 71d966e..6551372 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1791,7 +1791,8 @@ def _map_replicas(self, executor, step, func, replicas, *args): """ Run a step for each replica in a thread pool, letting every call finish before raising the first error, so that files written under a - lock are all written before it is released. + lock are all written before it is released. An interrupt exits + straight away, without waiting. Two replicas that share a slot must never be processed at the same time, so the replicas must come from a single pass, unless the step @@ -1808,7 +1809,7 @@ def _map_replicas(self, executor, step, func, replicas, *args): func: callable The step, which returns the index of the replica and the exception - raised, if any. + raised, if any. An exception it raises is treated the same way. replicas: list of int The replicas to run the step for. @@ -1816,13 +1817,23 @@ def _map_replicas(self, executor, step, func, replicas, *args): args: iterable The remaining arguments, an iterable of one per replica for each. """ + from concurrent.futures import wait + try: - outcomes = list(executor.map(func, replicas, *args)) + futures = [ + executor.submit(func, *arguments) for arguments in zip(replicas, *args) + ] + wait(futures) except KeyboardInterrupt: executor.shutdown(wait=False, cancel_futures=True) _logger.error(f"{step} cancelled. Exiting.") _sys.exit(1) + outcomes = [] + for replica, future in zip(replicas, futures): + error = future.exception() + outcomes.append((replica, error) if error else future.result()) + for index, error in outcomes: if error is not None: _logger.error( @@ -3832,10 +3843,11 @@ def _checkpoint( if error is not None: return index, error - # Delete all trajectory frames from the Sire system within the - # dynamics object. This is a no-op when the replica was committed - # earlier, since the frames were cleared then to free the slot. - dynamics._d._sire_mols.delete_all_frames() + # Delete the trajectory frames from the Sire system within the + # dynamics object. A replica committed earlier, with shared + # contexts, had them cleared then, so its slot isn't touched here. + if committed is None: + dynamics._d._sire_mols.delete_all_frames() if block == -1: _logger.info( diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 73c0b33..11dc266 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -957,7 +957,7 @@ def test_repex_concurrent_slots(ethane_methanol): # Every batch of dynamics, and every pass of commits, must contain at # most one replica per slot. batches = list(runner._dynamics_batches(runner._replica_devices())) - batches += runner._replica_passes(0) + batches += runner._replica_passes(0) + runner._replica_passes(1) for batch in batches: slots = [runner._dynamics_cache.slot_for(r) for r in batch] assert len(slots) == len(set(slots)), f"batch {batch} shares a slot" @@ -1964,6 +1964,37 @@ def fail(dynamics_executor, checkpoint_executor): assert finished == [True] +@pytest.mark.parametrize("is_raised", [False, True]) +def test_map_replicas_finishes_every_step(is_raised): + """ + Validate that a failing step is only raised once every other step has + finished, whether the step returns the error or raises it. + """ + import time + from concurrent.futures import ThreadPoolExecutor + + runner = object.__new__(RepexRunner) + runner._lambda_values = [0.0, 0.5, 1.0] + + finished = [] + + def step(replica): + if replica == 0: + error = RuntimeError("checkpoint failed") + if is_raised: + raise error + return replica, error + time.sleep(0.2) + finished.append(replica) + return replica, None + + with ThreadPoolExecutor(max_workers=3) as executor: + with pytest.raises(RuntimeError, match="checkpoint failed"): + runner._map_replicas(executor, "Checkpoint", step, [0, 1, 2]) + + assert sorted(finished) == [1, 2] + + def _layout(num_replicas, num_slots, num_gpus): """ A dynamics cache holding only a strided layout of replicas over slots, and From 5d1ba7d8a726eb95d03b020fd09552a8eeb78298 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 14:43:19 +0100 Subject: [PATCH 23/32] Count swaps from the end of the burn-in and fix legacy ring-break restarts --- CHANGELOG.md | 1 + src/somd2/runner/_repex.py | 13 ++++++++----- tests/runner/test_repex.py | 25 ++++++++++++++++++------- 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1483e00..04bb039 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ Changelog * Save the GCMC end-state topologies, which include the additional waters, at the first checkpoint of each run as well as the last, and keep them on restart, so that they match the trajectory and checkpoints while a simulation is running, if it crashes, or if a restart never reaches a checkpoint [#244](https://github.com/OpenBioSim/somd2/pull/244). * Add `gcmc_batch_size` and `gcmc_num_attempts` options to set the number of GCMC trials evaluated in parallel per batch and in total per move, which were previously fixed at the defaults of the GCMC sampler. Both can be changed on restart, e.g. to reduce wasted trials when more than one is accepted per move [#246](https://github.com/OpenBioSim/somd2/pull/246). * Add adaptive replica exchange, enabled with `adaptive_repex`, which concentrates sampling on the lambda windows that are slowest to converge. Windows are ranked at checkpoints by the uncertainty they contribute to the free energy, and only the highest ranked are propagated and exchanged every cycle, with the rest joining every `adaptive_mix_interval` cycles [#247](https://github.com/OpenBioSim/somd2/pull/247). +* Require `frame_frequency` to be a multiple of `gcmc_frequency` for GCMC simulations, so that every frame has a matching ghost water record. Previously some combinations silently saved frames without one [#247](https://github.com/OpenBioSim/somd2/pull/247). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 6551372..346dc21 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -347,9 +347,8 @@ def __setstate__(self, state): self._update_constraints = True if not hasattr(self, "_virtual_sites"): self._virtual_sites = None - # Set by the runner on a restart, since it follows the configuration. - if not hasattr(self, "_strided"): - self._strided = False + # Not pickled, since the runner sets it from the configuration. + self._strided = False # Convert a legacy checkpoint to the current convention, in which the # stored state of a replica is its own, with the last mix already @@ -2307,6 +2306,8 @@ def _run(self, dynamics_executor, checkpoint_executor): <= self._config.adaptive_burn_in ): _logger.info("Adaptive replica exchange burn-in complete") + self._adaptive_proposed = self._dynamics_cache.get_proposed().copy() + self._adaptive_accepted = self._dynamics_cache.get_accepted().copy() self._record_adaptive_data(energy_matrix, active) # Mix the replicas that were propagated. @@ -3163,10 +3164,12 @@ def _find_ring_pair(self): The particle indices, or None unless exactly one bond is broken or made, since the distance of a single bond is used. """ - atoms = self._system.atoms() + # Legacy restarts load a system per replica, which share a topology. + system = self._system[0] if isinstance(self._system, list) else self._system + atoms = system.atoms() pairs = [] - for mol in self._system["property is_perturbable"].molecules(): + for mol in system["property is_perturbable"].molecules(): for prop in ("ring_breaking_bonds", "ring_making_bonds"): if mol.has_property(prop): flat = mol.property(prop).to_list() diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 11dc266..01a407a 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -2090,12 +2090,17 @@ def test_find_ring_pair(syk_ring_break_mols): runner = _ring_pair_runner(syk_ring_break_mols, np.array([[0, 3], [100, 4]])) assert runner._find_ring_pair() == (1, 7) + # Legacy restarts load a system per replica. + runner = _ring_pair_runner([syk_ring_break_mols, syk_ring_break_mols]) + assert runner._find_ring_pair() == (1, 7) + @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") def test_adaptive_repex_ranking_arguments(ethane_methanol, monkeypatch): """ Validate that each ranking is given the samples each window has, the - cycles left, and the share of those cycles that propagate every window. + cycles left, and the share of those cycles that propagate every window, + and that swaps are only counted from the end of the burn-in. """ num_lambda = 4 @@ -2113,21 +2118,27 @@ def test_adaptive_repex_ranking_arguments(ethane_methanol, monkeypatch): "replica_exchange": True, "adaptive_repex": True, "adaptive_mix_interval": 2, - "adaptive_burn_in": "0 ps", + "adaptive_burn_in": "8fs", } runner = RepexRunner(ethane_methanol, Config(**config)) calls = [] - monkeypatch.setattr( - runner, - "_rank_windows", - lambda *args: calls.append(args), - ) + baselines = [] + + def rank_windows(*args): + calls.append(args) + baselines.append(runner._adaptive_proposed.sum()) + + monkeypatch.setattr(runner, "_rank_windows", rank_windows) runner.run() # 12 cycles with a checkpoint every 4, and every window always hot. assert [c[3] for c in calls] == [8, 4, 0] + + # The swaps proposed in the two burn-in cycles aren't counted. + assert baselines[0] > 0 + assert len(set(baselines)) == 1 for devices, capacity, samples, cycles_left, cold_rate in calls: # A context per replica, assigned to the GPUs in turn. assert devices == [r % runner._num_gpus for r in range(num_lambda)] From 4f029702d9bc48054da55bfd292043f53f832def Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 14:53:24 +0100 Subject: [PATCH 24/32] Test that subset mixing samples the Boltzmann distribution --- tests/runner/test_repex.py | 55 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 01a407a..a169808 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -333,6 +333,61 @@ def test_repex_mixing_samples_boltzmann(slot_schedule, inverted): assert tvd < 0.03 +def test_mix_subset_samples_boltzmann(): + """ + Validate that mixing only the hot replicas on most cycles, with a hot set + that changes during the run and every replica mixed periodically, samples + the same Boltzmann distribution over assignments as mixing them all. + """ + import itertools + + n = 4 + num_cycles = 40000 + mix_interval = 4 + hot_sets = [[0, 1, 2], [1, 2, 3], [0, 1], [2, 3], [1, 2]] + + rng = np.random.default_rng(1234) + U = rng.normal(0.0, 1.0, (n, n)) + proposed = np.zeros((n, n)) + accepted = np.zeros((n, n)) + + # The configuration in each window. + expected = list(range(n)) + + counts = {} + for cycle in range(num_cycles): + if (cycle + 1) % mix_interval == 0: + active = list(range(n)) + else: + active = hot_sets[(cycle // 1000) % len(hot_sets)] + + # Only the rows of the replicas being mixed may be used. + energies = rng.normal(0.0, 50.0, (n, n)) + for replica in active: + energies[replica] = U[:, expected[replica]] + + states = _mix_subset(np.array(active), energies, proposed, accepted) + for window in set(range(n)) - set(active): + assert states[window] == window + + expected = [expected[source] for source in states] + counts[tuple(expected)] = counts.get(tuple(expected), 0) + 1 + + perms = list(itertools.permutations(range(n))) + weights = np.array([np.exp(-sum(U[k, p[k]] for k in range(n))) for p in perms]) + weights /= weights.sum() + + def tvd(target): + return 0.5 * sum( + abs(counts.get(p, 0) / num_cycles - w) for p, w in zip(perms, target) + ) + + assert tvd(weights) < 0.03 + + # The energies must be far enough from uniform for the test to have power. + assert tvd(np.full(len(perms), 1 / len(perms))) > 0.1 + + @pytest.mark.parametrize( "rest2_scale, is_valid", [ From f6fae284c50f265911ce196b9ec92a549e9e9b91 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 15:01:33 +0100 Subject: [PATCH 25/32] Document the adaptive mix interval. --- README.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/README.md b/README.md index cab3888..187a815 100644 --- a/README.md +++ b/README.md @@ -480,6 +480,16 @@ with `--equilibration-time`, a burn-in of around 200 ps is enough. The saving is largest when few windows are slow. If most of the windows are hot, adaptive replica exchange costs almost as much as standard replica exchange. +A larger `--adaptive-mix-interval` saves more compute on the cold windows, but +samples them less and exchanges configurations between hot and cold windows less +often, which slows their movement across the λ range. The ranking accounts for +the reduced sampling, so fewer windows go cold as the interval grows, but not +for the slower exchange, so the interval isn't tuned automatically. Each ranking +logs the estimated compute per cycle and final variance relative to every window +being hot. If most windows stay hot, a smaller interval costs little, while if +most go cold, a larger interval saves more. Values between 2 and 8 are +reasonable, and the interval can be changed on restart. + ## 🌡️ REST2 From c5e4fa42d2e9aefbe43722974d61f603d674fced Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 15:14:39 +0100 Subject: [PATCH 26/32] Raise an error for adaptive replica exchange without replica exchange --- src/somd2/runner/_runner.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index bae8ec3..def9492 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -79,6 +79,11 @@ def __init__(self, system, config): _logger.error(msg) raise ValueError(msg) + if config.adaptive_repex: + msg = "'adaptive_repex' requires 'replica_exchange' to be enabled." + _logger.error(msg) + raise ValueError(msg) + # Call the base class constructor. super().__init__(system, config) From df9e733bc97ef668faca92338b467ea870581654 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 15:26:07 +0100 Subject: [PATCH 27/32] Keep nested details and list scrolling when the viewer refreshes. --- src/somd2/_viewer/_static/index.html | 31 ++++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/src/somd2/_viewer/_static/index.html b/src/somd2/_viewer/_static/index.html index 0bb9a6f..6b2f6b5 100644 --- a/src/somd2/_viewer/_static/index.html +++ b/src/somd2/_viewer/_static/index.html @@ -725,6 +725,25 @@

SOMD2 Viewer

// ---------------------------------------------------------------- rendering +// Which details are open and how far lists are scrolled, for elements marked +// with data-keep, so that redrawing the page keeps them. +function keptState(root) { + const kept = {}; + root.querySelectorAll("[data-keep]").forEach((el) => { + kept[el.dataset.keep] = { open: el.open, scroll: el.scrollTop }; + }); + return kept; +} + +function restoreState(root, kept) { + root.querySelectorAll("[data-keep]").forEach((el) => { + const k = kept[el.dataset.keep]; + if (!k) return; + if (el.tagName === "DETAILS") el.open = k.open; + el.scrollTop = k.scroll; + }); +} + function card(title, body) { const open = state.collapsed.has(title) ? "" : " open"; return `

${esc(title)}

${body}
`; @@ -941,7 +960,8 @@

SOMD2 Viewer

html += ``; if (p.messages && p.messages.length) { const msgs = p.messages.slice().reverse().map((m) => `
${esc(new Date(m.time * 1000).toLocaleString())} · ${esc(m.level)} · ${esc(m.message)}
`).join(""); - html += `
Recent warnings and errors (${p.messages.length})
${msgs}
`; + // Keyed by status, so that an error still opens it. + html += `
Recent warnings and errors (${p.messages.length})
${msgs}
`; } return card("Progress", html); } @@ -1566,19 +1586,19 @@

SOMD2 Viewer

html += ``; } html += ``; - if (sc.description) html += `
Schedule definition
${esc(sc.description)}
`; + if (sc.description) html += `
Schedule definition
${esc(sc.description)}
`; return card(`${LAMBDA} schedule`, html); } function renderConfig(s) { let html = ``; - html += `
`; + html += `
`; for (const [k, v] of Object.entries(s.config)) { const value = v === null ? "null" : typeof v === "object" ? JSON.stringify(v) : String(v); html += ``; } html += `
${esc(k)}${esc(value)}
`; - if (s.pme) html += `
Tuned PME parameters${renderPME(s.pme)}
`; + if (s.pme) html += `
Tuned PME parameters${renderPME(s.pme)}
`; return card("Configuration", html); } @@ -2438,9 +2458,12 @@

SOMD2 Viewer

html += renderConfig(s); html += ``; const focused = document.activeElement && document.activeElement.id === "config-filter" ? document.activeElement.selectionStart : null; + const kept = main.dataset.run === s.id ? keptState(main) : null; main.innerHTML = html; + main.dataset.run = s.id; renderCharts(s); applyConfigFilter(); + if (kept) restoreState(main, kept); const repeatSelect = $("#repeat-select"); if (repeatSelect) repeatSelect.addEventListener("change", (e) => { location.hash = e.target.value; }); attachStructure(s); From dc4f16cf48d72fe63e6e95b675a8a93f7006ffdb Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 15:27:15 +0100 Subject: [PATCH 28/32] Fix energy frequency and burn-in changes on adaptive repex restart --- README.md | 12 ++++++---- src/somd2/runner/_base.py | 17 +++++++++++++ src/somd2/runner/_repex.py | 7 +++++- tests/runner/test_repex.py | 49 ++++++++++++++++++++++++++++---------- 4 files changed, 67 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 187a815..b9a60a0 100644 --- a/README.md +++ b/README.md @@ -474,11 +474,13 @@ sampled for longer. It doesn't help a window that is stuck to escape. The windows are ranked at checkpoints, so the `--checkpoint-frequency` should be several times the `--adaptive-mix-interval` multiplied by the -`--energy-frequency`. The default burn-in of 500 ps also allows for -equilibration, which is off by default. If the system has been equilibrated -with `--equilibration-time`, a burn-in of around 200 ps is enough. The saving -is largest when few windows are slow. If most of the windows are hot, adaptive -replica exchange costs almost as much as standard replica exchange. +`--energy-frequency`, which can't be changed on restart. The default burn-in of +500 ps also allows for equilibration, which is off by default. If the system has +been equilibrated with `--equilibration-time`, a burn-in of around 200 ps is +enough. Raising the burn-in on restart beyond the current time makes every +window hot again until it is reached. The saving is largest when few windows +are slow. If most of the windows are hot, adaptive replica exchange costs almost +as much as standard replica exchange. A larger `--adaptive-mix-interval` saves more compute on the cold windows, but samples them less and exchanges configurations between hot and cold windows less diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index cfefe76..c6e1aea 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -2285,6 +2285,23 @@ def get_last_config(output_directory): self._compare_configs(self._last_config, config) + # Adaptive replica exchange derives each window's sample count from its + # clock and the energy frequency, so the frequency can't change. + if self._config.adaptive_repex: + from math import isclose + + last = self._last_config.get("energy_frequency") + if last and not isclose( + _sr.u(str(last)).to("ps"), self._config.energy_frequency.to("ps") + ): + msg = ( + "'energy_frequency' can't be changed on restart when using " + f"'adaptive_repex'. Set energy_frequency='{last}' to continue " + "the run." + ) + _logger.error(msg) + raise ValueError(msg) + @staticmethod def _systems_are_same(system0, system1, num_gcmc_waters=0): """ diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 346dc21..1bc13c9 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2178,7 +2178,12 @@ def _run(self, dynamics_executor, checkpoint_executor): if is_adaptive: self._reset_adaptive_data() - self._restore_adaptive_data() + # A burn-in raised on restart starts the adaptation again. + if self._start_time >= self._config.adaptive_burn_in: + self._restore_adaptive_data() + else: + self._dynamics_cache._hot_replicas = None + self._dynamics_cache._cold_votes = [0] * num_replicas else: # Adaptive replica exchange may have been switched off on restart. self._dynamics_cache._hot_replicas = None diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index a169808..57a93cb 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -2206,7 +2206,7 @@ def rank_windows(*args): @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") -def test_adaptive_repex_frames_every_cycle(ethane_methanol): +def test_adaptive_repex_frames_every_cycle(ethane_methanol, monkeypatch): """ Validate that frequent frames don't stop cold replicas being skipped, and that a cold replica saves the frame it missed when it is next propagated. @@ -2231,10 +2231,12 @@ def test_adaptive_repex_frames_every_cycle(ethane_methanol): "replica_exchange": True, "adaptive_repex": True, "adaptive_mix_interval": 2, - # Never rank, so that the hot replicas stay as set. - "adaptive_burn_in": "1 ns", + "adaptive_burn_in": "0 ps", } + # Never rank, so that the hot replicas stay as set. + monkeypatch.setattr(RepexRunner, "_rank_windows", lambda *args: None) + runner = RepexRunner(ethane_methanol, Config(**config)) runner._dynamics_cache._hot_replicas = hot runner.run() @@ -2335,12 +2337,14 @@ def test_checkpoint_without_adaptive_state(): @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") @pytest.mark.parametrize("max_contexts", [None, 2]) -def test_adaptive_repex_hot_replicas(ethane_methanol, max_contexts): +def test_adaptive_repex_hot_replicas(ethane_methanol, max_contexts, monkeypatch): """ Validate that only the hot replicas are propagated between full cycles, that each replica records energies against its own clock without gaps or - repeats, that frames are saved for every replica on the same cycles, and - that a restart carries on with the same hot replicas. + repeats, that frames are saved for every replica on the same cycles, that + a restart carries on with the same hot replicas, that a restart can't + change the energy frequency, and that raising the burn-in on restart makes + every replica hot again. """ import pandas as pd import sire as sr @@ -2363,10 +2367,12 @@ def test_adaptive_repex_hot_replicas(ethane_methanol, max_contexts): "max_contexts": max_contexts, "adaptive_repex": True, "adaptive_mix_interval": 2, - # Never rank, so that the hot replicas stay as set. - "adaptive_burn_in": "1 ns", + "adaptive_burn_in": "0 ps", } + # Never rank, so that the hot replicas stay as set. + monkeypatch.setattr(RepexRunner, "_rank_windows", lambda *args: None) + runner = RepexRunner(ethane_methanol, Config(**config)) # Not a ring-breaking perturbation, so there is no bond distance. @@ -2407,10 +2413,15 @@ def expected_times(num_hot, num_cold, r): times = [round(t, 6) for t in records.index.get_level_values(0)] assert times == expected_times(16, 8, i), f"replica {i}: {times}" - # Restart with adaptive replica exchange switched off, so that every - # replica is propagated every cycle. + # The sample counts depend on the energy frequency. config["runtime"] = "80fs" - config["adaptive_repex"] = False + changed = {**config, "energy_frequency": "8fs"} + with pytest.raises(ValueError, match="'energy_frequency' can't be changed"): + RepexRunner(ethane_methanol, Config(**changed)) + + # Restart with a burn-in that hasn't been reached, so that every + # replica is propagated every cycle. + config["adaptive_burn_in"] = "1 ns" runner = RepexRunner(ethane_methanol, Config(**config)) runner.run() @@ -2422,8 +2433,22 @@ def expected_times(num_hot, num_cold, r): times = [round(t, 6) for t in records.index.get_level_values(0)] assert times == expected_times(20, 12, i), f"replica {i}: {times}" + # Restart with adaptive replica exchange switched off. + config["runtime"] = "96fs" + config["adaptive_repex"] = False + + runner = RepexRunner(ethane_methanol, Config(**config)) + runner.run() + + assert runner._dynamics_cache._hot_replicas is None + + for i in range(num_lambda): + records = pd.read_parquet(runner._filenames[i]["energy_traj"]) + times = [round(t, 6) for t in records.index.get_level_values(0)] + assert times == expected_times(24, 16, i), f"replica {i}: {times}" + # Frames are scheduled by each replica's own clock again. mols = sr.load( [runner._filenames["topology0"], runner._filenames[i]["trajectory"]] ) - assert mols.num_frames() == 5 + assert mols.num_frames() == 6 From 59961363fb4287355cf26787bf95184263415972 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 16:04:28 +0100 Subject: [PATCH 29/32] Brighten Blender renders with the Standard view transform. --- src/somd2/_viewer/_blender.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/somd2/_viewer/_blender.py b/src/somd2/_viewer/_blender.py index 9647d17..9f56a60 100644 --- a/src/somd2/_viewer/_blender.py +++ b/src/somd2/_viewer/_blender.py @@ -282,6 +282,10 @@ def _default_view(focus, protein): # Settings that can be changed. RESOLUTION = ({width}, {height}) TRANSPARENT = False +# Colour management, brightened since the render is lit more dimly than the +# viewer. +VIEW_TRANSFORM = "Standard" +EXPOSURE = 1.3 # How the GCMC region is drawn: "wireframe", "shell" or None to leave it out. GCMC_STYLE = "wireframe" @@ -305,6 +309,8 @@ def _default_view(focus, protein): canvas = mn.Canvas(resolution=RESOLUTION, transparent=TRANSPARENT) canvas.clear() +bpy.context.scene.view_settings.view_transform = VIEW_TRANSFORM +bpy.context.scene.view_settings.exposure = EXPOSURE mol = mn.Molecule.load(str(path)) {styles} From 4cd6b6b5b1adb6286016bac1ae8fe2ddb0cc6dfb Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 16:18:45 +0100 Subject: [PATCH 30/32] Show live replica exchange progress from the log. --- src/somd2/_viewer/_data.py | 37 +++++++++++++++++++++++++--- src/somd2/_viewer/_static/index.html | 4 ++- src/somd2/runner/_repex.py | 9 ++++++- 3 files changed, 44 insertions(+), 6 deletions(-) diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index ac8bb79..99ce232 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -89,7 +89,11 @@ def _run(self): _line_re = _re.compile( r"^(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3}) \| (\w+)\s*\| \S+ - (.*)$" ) -_cycle_re = _re.compile(r"^Running dynamics for cycle (\d+) of (\d+)") +# Older versions don't log the time. +_cycle_re = _re.compile( + r"^Running dynamics for cycle (\d+) of (\d+)(?: \(time = ([\d.]+) ns\))?" +) +_restart_re = _re.compile(r"^Restarting at time (.+?), time remaining") _states_re = _re.compile(r"^States: \[([\d\s]*)\]") _block_re = _re.compile(r"^Finished block (\d+) of (\d+) for \S+ = ([\d.]+)") _perf_re = _re.compile(r"^Overall performance: ([\d.]+) ns day-1") @@ -276,6 +280,12 @@ def _reset(self): def _new_session(self, timestamp): self.session_start = timestamp self.production_start = None + # The simulation time, in ns, that a replica exchange session restarted + # from, which its cycles are counted from. + self.start_ns = 0.0 + # The simulation time, in ns, at the start of the latest cycle, if + # logged. + self.cycle_ns = None self.cycles = [] self.blocks = [] self.performance = None @@ -334,9 +344,12 @@ def _parse(self, line): self._add_states(message) else: self._pending_states = message + elif m := _restart_re.match(message): + self.start_ns = _to_ns(m.group(1)) or 0.0 elif m := _cycle_re.match(message): timestamp = _parse_timestamp(stamp) self.cycles.append((timestamp, int(m.group(1)), int(m.group(2)))) + self.cycle_ns = float(m.group(3)) if m.group(3) else None self.last_progress = timestamp if self.production_start is None: self.production_start = timestamp @@ -539,11 +552,13 @@ def _progress(self, config, lambda_values, windows, log): energy_frequency = _to_ns(config.get("energy_frequency")) checkpoint_frequency = _to_ns(config.get("checkpoint_frequency")) + def clamp(t): + return min(t, runtime) if runtime else t + done = {} for w in windows: if energy_frequency is not None: - t = w["samples"] * energy_frequency - done[w["lambda"]] = min(t, runtime) if runtime else t + done[w["lambda"]] = clamp(w["samples"] * energy_frequency) total = runtime * len(lambda_values) if runtime else None completed = sum(done.get(lam, 0.0) for lam in lambda_values) @@ -555,7 +570,7 @@ def _progress(self, config, lambda_values, windows, log): time = state.get("time_ns") if state else None if time is None: time = max(done.values(), default=0.0) - time = min(time, runtime) if runtime else time + time = clamp(time) completed = time * len(lambda_values) now = _time.time() @@ -567,6 +582,7 @@ def _progress(self, config, lambda_values, windows, log): else None ), "completed_ns": completed, + "saved_ns": completed, "fraction": completed / total if total else None, "eta_seconds": None, "throughput_ns_day": None, @@ -600,6 +616,19 @@ def _progress(self, config, lambda_values, windows, log): last_time, cycle, num_cycles = log.cycles[-1] progress["cycle"] = cycle progress["num_cycles"] = num_cycles + # The time run, from the cycles logged as they start, since samples + # are only saved at checkpoints. Within half a sample of the saved + # time, it is taken to be the same. + if energy_frequency is not None: + started = log.cycle_ns + if started is None: + started = log.start_ns + (cycle - 1) * energy_frequency + time = clamp(started + log.finished * energy_frequency) + live = time * len(lambda_values) + if live - completed > 0.5 * energy_frequency * len(lambda_values): + progress["completed_ns"] = live + if total: + progress["fraction"] = live / total times = [c[0] for c in log.cycles[-51:]] if len(times) > 1: interval = (times[-1] - times[0]) / (len(times) - 1) diff --git a/src/somd2/_viewer/_static/index.html b/src/somd2/_viewer/_static/index.html index 6b2f6b5..53e58e0 100644 --- a/src/somd2/_viewer/_static/index.html +++ b/src/somd2/_viewer/_static/index.html @@ -950,7 +950,9 @@

SOMD2 Viewer

let html = `
`; html += `
`; html += stat("Progress", `${fmt(pct, 1)}%`); - html += stat("Simulated / window", `${simTime(p.completed_ns / Math.max(1, s.lambda_values.length))} of ${simTime(p.runtime_ns)}`); + const perWindow = (ns) => simTime(ns / Math.max(1, s.lambda_values.length)); + const saved = p.saved_ns < p.completed_ns ? ` (${perWindow(p.saved_ns)} saved)` : ""; + html += stat("Simulated / window", `${perWindow(p.completed_ns)} of ${simTime(p.runtime_ns)}${saved}`); html += stat("Time remaining", p.status === "finished" ? "done" : duration(p.eta_seconds)); html += stat(`Current speed${p.throughput_label ? ` (${p.throughput_label})` : ""}`, p.throughput_ns_day ? `${fmt(p.throughput_ns_day, 1)} ns/day` : "–"); if (p.cycle) html += stat("Cycle", `${p.cycle} of ${p.num_cycles}`); diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 1bc13c9..d3fd4b1 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2189,9 +2189,16 @@ def _run(self, dynamics_executor, checkpoint_executor): self._dynamics_cache._hot_replicas = None self._dynamics_cache._cold_votes = [0] * num_replicas + # Converted once, since unit conversions are slow. + start_ns = self._start_time.to("ns") + energy_frequency_ns = self._config.energy_frequency.to("ns") + # Perform the replica exchange simulation. for i in range(cycles): - _logger.info(f"Running dynamics for cycle {i + 1} of {cycles}") + _logger.info( + f"Running dynamics for cycle {i + 1} of {cycles} " + f"(time = {start_ns + i * energy_frequency_ns:.4f} ns)" + ) # Log the states. This is the replica index for the state (positions # and velocities) used to seed each replica for the current cycle. From 0c61ec7690abb92c1f90692a6fabd1a2338ee8fc Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 14 Oct 2026 16:40:05 +0100 Subject: [PATCH 31/32] Log and show when adaptive repex first ranks the windows --- src/somd2/_viewer/_data.py | 9 +++++++++ src/somd2/_viewer/_static/index.html | 2 ++ src/somd2/runner/_repex.py | 15 ++++++++++++++- tests/viewer/test_adaptive.py | 19 +++++++++++++++++++ 4 files changed, 44 insertions(+), 1 deletion(-) diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index 99ce232..4ad19a2 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -102,6 +102,9 @@ def _run(self): _estimate_re = _re.compile( r"^Estimated cost per cycle ([\d.]+) and final variance ([\d.]+)" ) +_burn_in_re = _re.compile( + r"^Adaptive replica exchange burn-in complete, first ranking at ([\d.]+) ns" +) # Older versions don't log the time. _waters_re = _re.compile( r"^Current number of waters in GCMC volume at \S+ = ([\d.]+)" @@ -270,6 +273,8 @@ def _reset(self): # variance], and the cycles that propagated only the hot replicas. self.rankings = [] self.hot_only_cycles = [] + # The time of the first ranking in ns, logged once the burn-in ends. + self.first_ranking = None self.standard_state_correction = None self.messages = _deque(maxlen=20) # The number of waters in the GCMC region at each checkpoint, by λ, as @@ -377,6 +382,8 @@ def _parse(self, line): self.rankings.append([len(self.states), None, None, None]) elif (m := _estimate_re.match(message)) and self.rankings: self.rankings[-1][2:] = [float(m.group(1)), float(m.group(2))] + elif m := _burn_in_re.match(message): + self.first_ranking = float(m.group(1)) elif m := _waters_re.match(message): time = float(m.group(2)) if m.group(2) is not None else None counts = self.gcmc_waters.setdefault(float(m.group(1)), []) @@ -1052,6 +1059,8 @@ def _adaptive(log, num_states, repex_state, max_rankings=200): for k in range(num_states) ], } + elif log.first_ranking is not None: + result["first_ranking"] = log.first_ranking if log.states: result["hot_only_fraction"] = len(log.hot_only_cycles) / len(log.states) diff --git a/src/somd2/_viewer/_static/index.html b/src/somd2/_viewer/_static/index.html index 53e58e0..1ecaec0 100644 --- a/src/somd2/_viewer/_static/index.html +++ b/src/somd2/_viewer/_static/index.html @@ -1465,6 +1465,8 @@

SOMD2 Viewer

html += `
`; if (a.cost !== undefined && a.cost !== null) { html += `
Estimates from the ranking at cycle ${a.ranked_cycle}, relative to every window being hot for the rest of the run: cost is the compute per cycle from then on, and variance that of the final free energy.
`; + } else if (a.first_ranking !== undefined) { + html += `
The burn-in is complete. The windows are first ranked at ${fmt(a.first_ranking, 3)} ns, and can go cold from the checkpoint after that.
`; } else if (a.ranked_cycle === undefined && (!a.hot || a.hot.length === n)) { html += `
The windows haven't been ranked yet, so every window is hot until the burn-in is complete.
`; } diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index d3fd4b1..7771462 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2317,7 +2317,20 @@ def _run(self, dynamics_executor, checkpoint_executor): cycle_time - self._config.energy_frequency <= self._config.adaptive_burn_in ): - _logger.info("Adaptive replica exchange burn-in complete") + ranked = checkpoint_cycles[checkpoint_cycles >= i + 1] + if len(ranked) > 0: + ranked_time = self._start_time + int(ranked[0]) * ( + self._config.energy_frequency + ) + _logger.info( + "Adaptive replica exchange burn-in complete, first " + f"ranking at {ranked_time.to('ns'):.4f} ns" + ) + else: + _logger.info( + "Adaptive replica exchange burn-in complete, but no " + "checkpoint remains to rank the windows" + ) self._adaptive_proposed = self._dynamics_cache.get_proposed().copy() self._adaptive_accepted = self._dynamics_cache.get_accepted().copy() self._record_adaptive_data(energy_matrix, active) diff --git a/tests/viewer/test_adaptive.py b/tests/viewer/test_adaptive.py index 56e956c..347a4fe 100644 --- a/tests/viewer/test_adaptive.py +++ b/tests/viewer/test_adaptive.py @@ -58,6 +58,25 @@ def test_adaptive_summary(tmp_path): assert summary["history"] == {"cycle": [1, 3], "hot": [[0, 1], [1, 1], [1, 1]]} +def test_adaptive_first_ranking(tmp_path): + """ + Validate that the time of the first ranking is shown once the burn-in is + complete, until the windows are ranked. + """ + lines = _cycle(1, 3) + [ + _line("Adaptive replica exchange burn-in complete, first ranking at 0.6000 ns") + ] + log = _monitor(tmp_path, lines) + + summary = Simulation._adaptive(log, 3, None) + assert summary["first_ranking"] == 0.6 + + log = _monitor(tmp_path, lines + [_line("All replicas hot")]) + + summary = Simulation._adaptive(log, 3, None) + assert "first_ranking" not in summary + + def test_adaptive_rankings_after_restart(tmp_path): """ Validate that rankings and hot-only cycles logged after the last saved From e6639423c8b825f9be15bad311c19479f9b541b8 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 15 Oct 2026 11:32:14 +0100 Subject: [PATCH 32/32] Put window 0 at the bottom of the hot window chart --- src/somd2/_viewer/_static/index.html | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/somd2/_viewer/_static/index.html b/src/somd2/_viewer/_static/index.html index 1ecaec0..886f6e7 100644 --- a/src/somd2/_viewer/_static/index.html +++ b/src/somd2/_viewer/_static/index.html @@ -1489,12 +1489,12 @@

SOMD2 Viewer

for (let k = 0; k < numStates; k++) { for (let j = 0; j < cols; j++) { const hot = history.hot[k][j]; - svg += `window ${k}, cycle ${history.cycle[j]}: ${hot ? "hot" : "cold"}`; + svg += `window ${k}, cycle ${history.cycle[j]}: ${hot ? "hot" : "cold"}`; } } const rowStep = Math.ceil(numStates / Math.max(1, Math.floor(ph / 14))); for (let k = 0; k < numStates; k += rowStep) { - svg += `${k}`; + svg += `${k}`; } const colStep = Math.ceil(cols / Math.max(1, Math.floor(pw / 60))); for (let j = 0; j < cols; j += colStep) {