diff --git a/CHANGELOG.md b/CHANGELOG.md index 65acf54..04bb039 100644 --- a/CHANGELOG.md +++ b/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 -------------------------------------------------------------------------------------- diff --git a/README.md b/README.md index b9d0b62..b9a60a0 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,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. + + +### 🎯 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. + ## 🌡️ REST2 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} 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.") diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index 6056fa9..4ad19a2 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -89,11 +89,22 @@ 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") _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.]+)" +) +_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.]+)" @@ -178,6 +189,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 +269,12 @@ 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 = [] + # 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 @@ -261,6 +285,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 @@ -306,17 +336,25 @@ 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: 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 @@ -334,6 +372,18 @@ 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 := _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)), []) @@ -509,23 +559,37 @@ 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) + # Adaptive replica exchange samples cold windows less, so progress is + # the simulation time of the hot windows rather than the samples. + is_adaptive = _is_adaptive(config) + 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 = clamp(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, + "saved_ns": completed, "fraction": completed / total if total else None, "eta_seconds": None, "throughput_ns_day": None, @@ -559,6 +623,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) @@ -566,7 +643,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,13 +957,16 @@ 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), "terminal_flip_stats": getattr( state, "_terminal_flip_stats", None ), + "hot_replicas": getattr(state, "_hot_replicas", None), } except Exception as e: error = error or e @@ -925,9 +1007,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"] @@ -943,6 +1028,45 @@ 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) + ], + } + 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) + + return result + @staticmethod def _replica_paths(states, num_states, max_points=500): """ @@ -1742,7 +1866,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..886f6e7 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}
`; @@ -931,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}`); @@ -941,7 +962,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); } @@ -1422,10 +1444,70 @@

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 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 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.
`; + } + 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 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 ph = cellH * numStates; + const height = m.t + ph + m.b; + 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.max(1, Math.floor(ph / 14))); + for (let k = 0; k < numStates; k += rowStep) { + svg += `${k}`; + } + 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 += `cycle`; + svg += `window`; + 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; @@ -1508,19 +1590,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); } @@ -1614,6 +1696,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; @@ -2378,9 +2462,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); diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index da9b69d..9cb286a 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -117,6 +117,10 @@ class Config: "timeout", "oversubscription_factor", "max_contexts", + "adaptive_repex", + "adaptive_mix_interval", + "adaptive_burn_in", + "adaptive_min_samples", "restraint_search_time", "restraint_search_frequency", "gcmc_batch_size", @@ -187,6 +191,10 @@ 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_min_samples=50, randomise_velocities=False, perturbed_system=None, terminal_flip_frequency=None, @@ -494,6 +502,32 @@ 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 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 + 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_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 +819,10 @@ 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_min_samples = adaptive_min_samples self.randomise_velocities = randomise_velocities self.perturbed_system = perturbed_system self.terminal_flip_frequency = terminal_flip_frequency @@ -2136,6 +2174,69 @@ 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_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..c6e1aea 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. @@ -2276,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): """ @@ -2489,6 +2515,7 @@ def _checkpoint( is_final_block=False, context=None, gcmc_sampler=None, + num_energy_rows=None, ): """ Save a checkpoint file. @@ -2517,6 +2544,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 +2569,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 +2668,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 +2708,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..7771462 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,8 @@ def __setstate__(self, state): self._update_constraints = True if not hasattr(self, "_virtual_sites"): self._virtual_sites = None + # 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 @@ -280,6 +390,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. @@ -360,14 +474,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.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. @@ -402,7 +511,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: @@ -646,6 +755,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. @@ -798,19 +943,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 +1097,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 +1412,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 +1515,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 +1602,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 +1670,13 @@ 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") + + 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: @@ -1568,7 +1738,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 +1757,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 +1768,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 @@ -1608,39 +1786,60 @@ def _replica_passes(self, cycle): 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. An interrupt exits + straight away, without waiting. 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. An exception it raises is treated the same way. - list of int - A batch of replicas that is safe to process concurrently. + replicas: list of int + The replicas to run the step for. + + args: iterable + The remaining arguments, an iterable of one per replica for each. """ - from math import ceil + from concurrent.futures import wait - 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] + try: + 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( + f"{step} failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}: {error}" + ) + raise error def _set_num_slots(self): """ @@ -1648,7 +1847,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 +1862,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 @@ -1752,11 +1952,53 @@ def run(self): """ Run the replica exchange simulation. """ + 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) + + executors = (dynamics_executor, checkpoint_executor) + try: + self._run(dynamics_executor, checkpoint_executor) + except Exception: + # Wait for running blocks, so that nothing still uses the contexts. + 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): + """ + 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. + """ from math import ceil from time import time - from concurrent.futures import ThreadPoolExecutor from itertools import repeat # Record the start time. @@ -1804,64 +2046,51 @@ 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) - # Create the replica list. replica_list = list(range(self._config.num_lambda)) + # 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: - for batch in self._safe_batches(num_workers): - 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) + for batch in self._dynamics_batches(replica_devices): + 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._safe_batches(num_workers): - 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) + for batch in self._dynamics_batches(replica_devices): + 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. @@ -1870,51 +2099,25 @@ def run(self): # 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 - ] - 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 - except KeyboardInterrupt: - checkpoint_executor.shutdown(wait=False, cancel_futures=True) - _logger.error( - "Post-equilibration checkpoint cancelled. Exiting." - ) - _sys.exit(1) + 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 @@ -1955,16 +2158,47 @@ 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 + 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() + # 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 + 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. @@ -1978,56 +2212,71 @@ 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 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 to checkpoint. + is_checkpoint = (i + 1) in checkpoint_set - # 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 - # 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 + # Only the hot replicas are propagated, except on mix and checkpoint + # cycles. + hot = self._dynamics_cache._hot_replicas + is_full = ( + hot is None + or (i + 1) % self._config.adaptive_mix_interval == 0 + or do_checkpoint ) + active = list(range(num_replicas)) if is_full else hot - # Whether the checkpoint files are written at the end of this cycle. - do_checkpoint = is_checkpoint or i == cycles - 1 + if not is_full: + _logger.info(f"Propagating {len(active)} hot replicas: {active}") - # 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. + # 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] + 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 = [ + 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 + ] + + # 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 + ] + try: for result, index, energies in dynamics_executor.map( self._run_block, replicas, repeat(self._lambda_values), - repeat(is_gcmc), - repeat(write_gcmc_ghosts), - repeat(is_terminal_flip), - repeat(time_ns), + is_gcmc, + is_frame, + is_terminal_flip, + repeat(self._config.save_energy_components), repeat(clock), + ( + is_frame + if is_adaptive and self._save_frames + else repeat(None) + ), ): if not result: _logger.error( @@ -2047,38 +2296,50 @@ def run(self): # 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 + 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 + ): + 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) + + # 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 +2356,23 @@ 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: + cold_rate = self._cold_rate( + i + 1, + cycles, + checkpoint_cycles, + self._config.adaptive_mix_interval, + ) + self._rank_windows( + replica_devices, + device_capacity, + self._window_samples(), + cycles - (i + 1), + cold_rate, + ) + # Create the lock. lock = _FileLock(self._lock_file) @@ -2102,58 +2380,31 @@ def run(self): # 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 - ] - 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 - except KeyboardInterrupt: - checkpoint_executor.shutdown( - wait=False, cancel_futures=True - ) - _logger.error("Backup cancelled. Exiting.") - _sys.exit(1) + 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 - ] - 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), - ): - if error: - _logger.error( - f"Checkpoint 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("Checkpoint cancelled. Exiting.") - _sys.exit(1) + 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 + + # 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) @@ -2162,12 +2413,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) - # Record the end time for the production block. prod_end = time() @@ -2202,8 +2447,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 +2482,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 +2512,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 +2605,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() @@ -2417,7 +2687,7 @@ def _minimise(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = slot % self._num_gpus + device = self._dynamics_cache.device_for(slot) # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -2527,7 +2797,7 @@ def _equilibrate(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = slot % self._num_gpus + device = self._dynamics_cache.device_for(slot) # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -2573,7 +2843,7 @@ def _equilibrate(self, index): self._dynamics_cache.delete(slot) # Work out the device index. - device = slot % self._num_gpus + device = self._dynamics_cache.device_for(slot) # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -2670,8 +2940,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 +2953,736 @@ def _assemble_results(self, results): return matrix + @staticmethod + 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. + """ + 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) + 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. + + Parameters + ---------- + + first_cycle: int + The number of cycles run so far. + + num_cycles: int + The number of cycles in the run. + + checkpoint_cycles: np.ndarray + The checkpoint cycles, numbered from one, in increasing order. + + mix_interval: int + The number of cycles between mix cycles. + + Returns + ------- + + float + The share of cycles. + """ + if first_cycle >= num_cycles: + return 1.0 + + # 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" + ) + 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 + + 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. + """ + cache = self._dynamics_cache + return [cache.device_for(cache.slot_for(r)) 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. + """ + return [ + max( + 1, + min( + self._config.oversubscription_factor, + self._dynamics_cache.contexts_on(d), + ), + ) + for d in range(self._num_gpus) + ] + + 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 + ---------- + + 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. + """ + 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 _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. + + Parameters + ---------- + + devices: list of int + The GPU that propagates each replica. + + Yields + ------ + + list of int + A batch of replicas. + """ + for batch in self._replica_passes(0): + yield from self._device_batches(batch, devices) + + 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 + effect, before any time is spent on minimisation and equilibration. + """ + from math import floor + + 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." + ) + # 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 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") + <= 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 + 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. + """ + # 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 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, 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. + + 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 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. 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. + + Parameters + ---------- + + 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. + + cycles_left: int + The number of cycles left in the run. + + cold_rate: float + The share of the cycles left 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: + _logger.info("Not ranking windows: no cycles left, keeping hot replicas") + 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] + + 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 + + # 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] + 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] + 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) + + weights = _np.sqrt(fwd_variance + rev_variance) + + forced = set(undersampled) | set(slow_distance) + for k in flagged: + 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 + } + + 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 each GPU + # 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 + 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]): + d = window_devices[r] + if counts[d] < batches * device_capacity[d]: + filled.add(r) + counts[d] += 1 + return filled + + total_variance = owned.sum() + samples = _np.maximum(_np.asarray(window_samples, dtype=float), 1.0) + + def rates(windows): + # 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 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] + rev_variance / final[1:]).sum() + return cost, variance + + # A set must beat all hot by more than rounding error. + 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 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 + 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 + efficiency = _np.prod(score(fill(windows))) + if efficiency < best: + chosen, best = windows, efficiency + + votes = self._dynamics_cache._cold_votes + 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 chosen: + votes[r] = 0 + elif r in previous: + votes[r] += 1 + if votes[r] < 2: + hot.add(r) + + # Windows kept hot can leave no gain over all hot. + hot = fill(hot) + if _np.prod(score(hot)) >= threshold: + hot = all_hot + + hot_variance = score(hot)[1] + hot = sorted(hot) + 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" + ) + # 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 = 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 every window being hot" + ) + def _check_restart(self): """ Check the output directory for a valid restart state. @@ -2789,7 +3791,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 +3821,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,16 +3860,22 @@ 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: 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( @@ -2942,70 +3962,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 +4011,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/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) diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 2d688e2..57a93cb 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) @@ -324,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", [ @@ -709,6 +773,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", @@ -893,9 +1009,11 @@ 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, 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) + 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" @@ -1319,3 +1437,1018 @@ 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 + runner._is_cached = False + 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 _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, 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, [capacity] * num_gpus, window_samples, cycles_left, cold_rate + ) + + +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, capacity=capacity) + _rank(runner, num_gpus=num_gpus, capacity=capacity) + + +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_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) + _rank(runner, num_gpus=2) + assert cache._hot_replicas is None + + _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_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_gpus=2) + assert cache._hot_replicas == [4, 5, 6, 7] + + +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. + """ + runner = _ranking_runner(8, adaptive_min_samples=10) + + _fill_ranking_data(runner, hard_pair=3) + _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_gpus=8) + 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), + # 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), + ], +) +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 + 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)) + _rank_twice(runner, num_gpus=2) + + 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_gpus=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_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_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] + + +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, 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 - tolerance + if is_checkpoint: + 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 + + +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), + # 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. + (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 + + +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): + """ + Validate the share of cycles that propagate the cold windows against the + run loop's schedule. + """ + 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) + ) + + +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] + 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) + _rank(runner, num_gpus=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 + 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 + _rank(runner, num_gpus=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_min_samples=50) + + _fill_ranking_data(runner, hard_pair=3, num_samples=20) + _rank(runner, num_gpus=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_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) + _rank_twice(runner, num_gpus=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_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) + _rank_twice(runner, num_gpus=2) + + assert 7 in runner._dynamics_cache._hot_replicas + assert 0 not in runner._dynamics_cache._hot_replicas + + +@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, of as many + replicas as the GPU can propagate at once. + """ + 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, capacity=capacity) + + 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_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 waits for the blocks still running + in the thread pools before it is raised. + """ + import time + + runner = object.__new__(RepexRunner) + runner._config = Config() + runner._num_gpus = 1 + + finished = [] + + def fail(dynamics_executor, checkpoint_executor): + dynamics_executor.submit(lambda: (time.sleep(0.2), finished.append(True))) + raise RuntimeError("dynamics failed") + + monkeypatch.setattr(runner, "_run", fail) + + with pytest.raises(RuntimeError, match="dynamics failed"): + runner.run() + + 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 + 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( + "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_gpus = num_gpus + runner._dynamics_cache = _layout(num_slots, num_slots, 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. + """ + runner = object.__new__(RepexRunner) + runner._num_gpus = 2 + runner._num_replicas = 10 + runner._dynamics_cache = _layout(10, 4, 2) + + 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 + oversubscription factor, so that hot replicas that aren't consecutive + still spread over the GPUs. + """ + runner = object.__new__(RepexRunner) + runner._config = Config(num_lambda=25) + runner._num_gpus = 3 + runner._num_replicas = 25 + runner._dynamics_cache = _layout(25, 25, 3) + devices = runner._replica_devices() + + # Every replica is batched as consecutive groups of three. + 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, devices) + 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. + """ + 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) + + # 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, + and that swaps are only counted from the end of the burn-in. + """ + 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": "8fs", + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + calls = [] + 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)] + assert capacity == [1] * runner._num_gpus + first = 12 - cycles_left + assert samples == [first] * num_lambda + if cycles_left > 0: + expected = _full_cycle_share(first, 12, 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, 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. + """ + 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, + "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() + + # 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): + """ + 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_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)) + _rank_twice(runner, num_gpus=2) + + hot = runner._dynamics_cache._hot_replicas + assert hot is not None and {0, 1, 2}.issubset(hot) + + +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, 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, 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 + + 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, + "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. + 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}" + + # The sample counts depend on the energy frequency. + config["runtime"] = "80fs" + 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() + + 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}" + + # 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() == 6 diff --git a/tests/viewer/test_adaptive.py b/tests/viewer/test_adaptive.py new file mode 100644 index 0000000..347a4fe --- /dev/null +++ b/tests/viewer/test_adaptive.py @@ -0,0 +1,141 @@ +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 per cycle 0.75 and final variance 1.10"), + ] + + _cycle(2, 4, hot=[1, 2]) + + _cycle(3, 4) + + [ + _line("All replicas hot"), + _line( + "Estimated cost per cycle 1.00 and final 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_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 + 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 per cycle 0.75 and final variance 1.10"), + _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 per cycle 0.75 and final variance 1.10"), + ] + + _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 per cycle 0.75 and final variance 1.10"), + ] + 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 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