Add support for adaptive replica exchange

CHANGELOG.md

@@ -23,6 +23,8 @@ 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).
+* 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
 --------------------------------------------------------------------------------------

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,82 @@ 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.
 
+<a id="adaptive-replica-exchange"></a>
+
+### 🎯 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 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
+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 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,
+- 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, 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
+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
+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`, 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
+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.
+
 <a id="rest2"></a>
 
 ## 🌡️ REST2

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}

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.")