Add support for adaptive replica exchange

[lohedges]

This PR adds an option to concentrate replica exchange sampling on the λ windows that are slowest to converge, enabled with --adaptive-repex. This was written with the aid of Claude Code (Opus 5.5).

Why

For difficult perturbations, most of the uncertainty in a free energy usually comes from a handful of λ windows. In a ring-breaking calculation, for example, the windows that swap the hard Morse potential for the soft one converge quickly and then contribute almost nothing as the simulation is extended, while the windows around the point where the bond breaks are still far from converged. Standard replica exchange propagates every window every cycle regardless, so much of the compute goes on windows that no longer need it.

Adaptive replica exchange spends less time on the windows that have converged. --runtime sets the number of cycles, which is the simulation time of the hot windows, so how the saving is used depends on the --runtime chosen:

• The same --runtime as standard replica exchange. The hot windows get the same sampling as before, the cold windows get less, and the run finishes sooner. In the PCSK9 example below, 5 ns took 60 minutes rather than 75.
• A longer --runtime, so that the run takes as long as standard replica exchange would. The hot windows get more sampling in the same time, which lowers the uncertainty. In the same example, about 6.3 ns fits in the 75 minutes that standard replica exchange needs for 5 ns.

It doesn't help a window that is genuinely stuck to escape, which needs a better λ schedule or enhanced sampling.

How it works

The run starts with a burn-in period of standard replica exchange. After that, the windows are ranked at each checkpoint, and split into two groups:

• Hot windows are propagated and exchanged every cycle, as now.
• Cold windows only join in every --adaptive-mix-interval cycles, and on checkpoint cycles.

Exchanges on the cycles in between only involve the hot windows. Mixing a subset of the replicas, with the rest held fixed, is still a valid replica exchange move, so the sampled distributions are unchanged. Each replica keeps its own simulation clock, so energies and frames from cold windows are recorded against the time they were actually simulated for, without gaps or overwrites.

Choosing the hot windows

Each pair of neighbouring windows is weighted by how much uncertainty it contributes to the free energy for a given amount of sampling: the spread of its energy differences, σ, multiplied by the square root of how correlated they are, the statistical inefficiency g. Pairs are added to the hot set in order of weight, together with the windows either side of them, so that exchange can still carry configurations in and out. Of the sets this gives, the one chosen gives the lowest uncertainty in the final free energy for the compute of the whole run, accounting for the samples each window has already collected and for cold windows being sampled less often. If no set does better than keeping every window hot, every window stays hot.

Some windows are always kept hot, because their uncertainty can't be trusted to be small:

• pairs whose forward and reverse free energy estimates disagree by more than 1 kT, or whose estimate drifts during the run,
• pairs that rarely exchange,
• windows with fewer than --adaptive-min-samples statistically independent samples,
• for ring-breaking perturbations that break or make a single bond, windows in which the bond distance itself decorrelates slowly. A slow change in the bond length can bias the free energy without showing up in the energy differences, so this is checked separately.

To avoid flipping back and forth, a window only goes cold after being found cold at two consecutive rankings.

Keeping the GPUs busy

Each replica's context lives on a fixed GPU, so 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. The time a cycle takes is set by the number of batches the busiest GPU needs, so with a context per replica the hot set is topped up with the cold windows that contribute the most uncertainty until every GPU fills that many batches. Those windows cost no extra time, and the ranking accounts for them when choosing the hot set.

When contexts are re-used with --max-contexts, replicas are assigned to contexts in a strided pattern, so that the hot windows are spread across the GPUs. The hot set isn't topped up in that case, since a window added to a context that is already busy would need another pass.

Usage

somd2 system.bss --replica-exchange --adaptive-repex

Option	Default	Meaning
--adaptive-repex	off	Enable adaptive replica exchange.
--adaptive-mix-interval	4	Cycles between those that propagate every window.
--adaptive-burn-in	500 ps	Standard replica exchange before the windows are first ranked.
--adaptive-min-samples	50	Independent samples a window needs before it can go cold.

Things to be aware of:

• --runtime is the simulation time of the hot windows, so the cold windows are sampled for less time and the windows have different numbers of energy samples. MBAR handles this as normal.
• Windows are ranked at checkpoints, so --checkpoint-frequency should be several times --adaptive-mix-interval × --energy-frequency. No window goes cold until the second checkpoint after the burn-in. Settings that would stop adaptive replica exchange having any effect are warned about at startup.
• The default burn-in of 500 ps also allows for equilibration, which is off by default. With --equilibration-time, around 200 ps is enough.
• The hot set and the data used to rank the windows are saved in the checkpoint, so a restart carries on where it left off. Adaptive replica exchange can be switched on or off on restart, but --energy-frequency can't be changed while it is on. Raising --adaptive-burn-in beyond the current time on restart makes every window hot again until it is reached.
• Each ranking is logged with the hot windows and the estimated compute per cycle and final variance relative to keeping every window hot, e.g. Estimated cost per cycle 0.70 and final variance 1.09 relative to every window being hot.

Viewer

The replica exchange card has a new section for adaptive runs, showing the current hot windows, the estimates from the latest ranking, the share of cycles that only propagated the hot windows, and a chart of which windows were hot at each ranking. Progress for adaptive runs is now measured from the simulation time rather than the number of energy samples, since cold windows have fewer.

Other changes

Without --adaptive-repex, replica exchange runs as before, with the same mixing and the same contiguous slot layout with --max-contexts, apart from the following, which apply to every run:

• GCMC runs now require frame_frequency to be a multiple of gcmc_frequency, so that every frame coincides with a GCMC move and has its ghost water record. Previously some combinations silently saved frames without one.
• The energy trajectory for each window is appended with exactly the samples collected since the last checkpoint, rather than a fixed number per block, which was wrong when the checkpoint and energy frequencies weren't whole multiples of each other.
• Checkpoint cycles are computed directly as the first cycle at or after each multiple of the checkpoint interval, rather than by repeatedly adding the interval, which could drift by a cycle for fractional intervals in long runs.
• Replicas propagated together are now batched by GPU, with at most --oversubscription-factor on any GPU, rather than in index order, which could leave a GPU idle. This applies to minimisation and equilibration as well as production, which also share one persistent thread pool rather than creating a new pool for every batch.
• If one replica's checkpoint fails, the others now finish writing before the error is raised, while the file lock is still held, so a reader never sees a mixture of old and new checkpoint files. All checkpoint steps go through one helper with this behaviour.
• If an error stops a run, it waits for any dynamics or checkpoint work still running before the error is raised, so nothing is still using the OpenMM contexts. Interrupting a run with Ctrl-C still exits immediately.
• The viewer now recognises the final save written when a run doesn't end on a checkpoint, so extending such a run no longer discards the replica history logged before it.
• An expected RDKit warning, printed when the viewer overlays end states whose geometries don't fit each other, is silenced.

Validation

The PCSK9 C35 → C25 macrocycle ring-break (free leg, 25 windows, 5 ns, 3 GPUs) was run three times with adaptive replica exchange, and compared with an earlier standard run of the same system. This is a hard case for adaptive replica exchange, since about two thirds of the schedule needs extensive sampling.

Which windows are hot. The ranking matched an independent analysis of the earlier run: the first third of the schedule, where the hard Morse potential is swapped for the soft one, went cold, while the windows around the bond break and towards λ = 1 stayed hot. After about 1.5 ns, 15–18 of the 25 windows were hot.

Speed. With the final code, the run took 60 minutes against about 75 minutes had every cycle propagated every window, a saving of 21%. A cycle propagating 15 hot windows took 0.47 s against 0.84 s for all 25. Earlier runs saved only 13–14%, because hot replicas were batched without regard to their GPU, so some batches put two replicas on one GPU while leaving another idle. Every batch now takes the same time, about 0.09 s. Transformations where fewer windows are slow, such as small ring breaks, should save more.

Free energy. Decorrelated MBAR free energies at 5 ns (kcal/mol):

Run	ΔG	λ = 0.5
Standard	-18.98 ± 0.34
Adaptive 1	-19.34 ± 0.30	bond broke
Adaptive 2	-18.59 ± 0.22	bond stayed formed
Adaptive 3	-18.15 ± 0.24	bond stayed formed

The results span about 1.2 kcal/mol, several times their stated errors. The bond distance shows why: at λ = 0.5 the bond is bistable, and whether it breaks within a 5 ns run largely decides the answer. When it stays formed, the window looks well converged by every statistic, which no sampling scheme can detect. The run where it broke gave the most negative result. Extra windows between λ = 0.45 and 0.525 are needed to converge this system, with either standard or adaptive replica exchange. Nothing points to adaptive replica exchange changing the result.

The earlier standard run predates the replica exchange mixing fix in #232, so a standard run with the current code is needed for a clean comparison, and repeats of both methods on an improved schedule are needed to show whether adaptive replica exchange reduces the spread of results at equal compute.

Tests

New tests cover subset mixing, including that it samples the Boltzmann distribution over assignments with a changing hot set, per-replica clocks across restarts, frames and energy records for hot and cold windows (with one context per replica and with shared contexts), the window ranking and its edge cases, the cold sampling rate against the run loop's schedule, the arguments the run loop passes to the ranking, batching across GPUs and contexts, the bond distance check, startup warnings, restoring older checkpoints, checkpoint failures under the file lock, waiting for running work on error, and the viewer summary.

[lohedges]

I did a 5 ns ABFE bound-leg run overnight (24 windows, 3 GPUs) and it worked well. The plot below shows which windows were hot at each ranking. All windows start hot, and windows drop out as they stop contributing much to the uncertainty, starting with the decharge stage. The hot set adapts as the run goes on: a couple of decharge windows briefly came back mid-run. By the end only 6 of the 24 windows were hot, and each cycle ran at a bit over twice the speed of standard replica exchange. Over the whole run it took about 0.56 of the time a standard run would have. The bound-leg ΔG was 221.75 kcal/mol, within the range of three previous standard runs (221.75–225.02).

[lohedges]

For reference, the windows that temporarily came back as hot did so for free. These were added to fill the GPUs on those cycles, i.e. to ensure that the load wasn't unbalanced.