diff --git a/CHANGELOG.md b/CHANGELOG.md index 04bb039..f753ae5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,8 @@ Changelog * 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). +* Fixed the regular runner's final block when the `runtime` isn't a multiple of the `checkpoint_frequency`, which ran a single integration step rather than the remaining time, without GCMC moves. Each checkpoint now appends exactly the energies sampled since the last one, so energies are no longer duplicated, or lost on restart when `checkpoint_frequency` is zero. Blocks that aren't a whole number of GCMC or terminal flip intervals no longer overrun, and energy components are time stamped correctly after a restart [#248](https://github.com/OpenBioSim/somd2/pull/248). +* Require the `runtime`, `checkpoint_frequency`, `energy_frequency`, `frame_frequency`, `gcmc_frequency` and `terminal_flip_frequency` to be multiples of the `timestep` for the regular runner, which runs in whole integration steps. GCMC and terminal flip moves now keep their frequency across checkpoints [#248](https://github.com/OpenBioSim/somd2/pull/248). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/src/somd2/_viewer/_data.py b/src/somd2/_viewer/_data.py index 4ad19a2..534102d 100644 --- a/src/somd2/_viewer/_data.py +++ b/src/somd2/_viewer/_data.py @@ -249,7 +249,11 @@ def __setstate__(self, state): class _RepexUnpickler(_pickle.Unpickler): def find_class(self, module, name): - if module == "somd2.runner._repex" and name == "DynamicsCache": + # The cache moved to its own module, so older checkpoints use _repex. + if name == "DynamicsCache" and module in ( + "somd2.runner._dynamics_cache", + "somd2.runner._repex", + ): return _RepexState return super().find_class(module, name) diff --git a/src/somd2/runner/_adaptive.py b/src/somd2/runner/_adaptive.py new file mode 100644 index 0000000..d86ce95 --- /dev/null +++ b/src/somd2/runner/_adaptive.py @@ -0,0 +1,618 @@ +###################################################################### +# SOMD2: GPU accelerated alchemical free-energy engine. +# +# Copyright: 2023-2026 +# +# Authors: The OpenBioSim Team +# +# SOMD2 is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# SOMD2 is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with SOMD2. If not, see . +##################################################################### + +__all__ = [ + "WindowRanker", + "cold_rate", + "exp_estimate", + "is_unconverged", + "statistical_inefficiency", +] + +from math import ceil as _ceil + +import numpy as _np + +from somd2 import _logger + + +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) + + +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)))) + + +def is_unconverged(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(exp_estimate(fwd) + 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 * (exp_estimate(f) - 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 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) + + +class WindowRanker: + """ + The data used to rank the lambda windows for adaptive replica exchange, and + the ranking that chooses the hot windows. + """ + + def __init__( + self, + num_windows, + min_samples, + proposed, + accepted, + is_ring_break=False, + is_cached=False, + ): + """ + Constructor. + + Parameters + ---------- + + num_windows: int + The number of lambda windows. + + min_samples: int + The number of independent samples a window needs before it can be + made cold. + + proposed: np.ndarray + The number of swaps proposed between each pair of windows so far, + from which swaps are counted for the ranking. + + accepted: np.ndarray + The number of swaps accepted between each pair of windows so far. + + is_ring_break: bool + Whether the distance of the bond that is broken or made is + recorded, in which case windows where it is slow are kept hot. + + is_cached: bool + Whether contexts are shared between replicas, in which case the hot + windows aren't topped up to fill the GPUs. + """ + self._num_windows = num_windows + self._min_samples = min_samples + self._is_ring_break = is_ring_break + self._is_cached = is_cached + self.reset(proposed, accepted) + + def reset(self, proposed, accepted): + """ + Clear the recorded data. + + Parameters + ---------- + + proposed: np.ndarray + The number of swaps proposed between each pair of windows so far. + + accepted: np.ndarray + The number of swaps accepted between each pair of windows so far. + """ + # The reduced energy difference to the lower and upper neighbouring + # window, sampled in each window. + self.lower = [[] for _ in range(self._num_windows)] + self.upper = [[] for _ in range(self._num_windows)] + + # The distance of the bond that is broken or made, sampled in each + # window. + self.distance = [[] for _ in range(self._num_windows)] + + self.set_swap_baseline(proposed, accepted) + + def set_swap_baseline(self, proposed, accepted): + """ + Set the swap counts from which swaps are counted for the ranking. + + Parameters + ---------- + + proposed: np.ndarray + The number of swaps proposed between each pair of windows so far. + + accepted: np.ndarray + The number of swaps accepted between each pair of windows so far. + """ + self.proposed = proposed.copy() + self.accepted = accepted.copy() + + def record(self, energy_matrix, windows, distances=None): + """ + 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. + + windows: list of int + The windows that were propagated. + + distances: dict + The bond distance in each window that was propagated, which is + required for a ring-breaking perturbation. + """ + if self._is_ring_break and distances is None: + raise ValueError( + "Bond distances are required for a ring-breaking perturbation" + ) + + for k in windows: + row = energy_matrix[k] + if k > 0: + self.lower[k].append(row[k - 1] - row[k]) + if k < self._num_windows - 1: + self.upper[k].append(row[k + 1] - row[k]) + if distances is not None: + self.distance[k].append(distances[k]) + + def to_dict(self): + """ + The recorded data, to be stored in a checkpoint. + + Returns + ------- + + dict + The recorded data. + """ + return { + "lower": [_np.array(x) for x in self.lower], + "upper": [_np.array(x) for x in self.upper], + "distance": [_np.array(x) for x in self.distance], + "proposed": self.proposed, + "accepted": self.accepted, + } + + def restore(self, data): + """ + Restore the recorded data from a checkpoint. + + Parameters + ---------- + + data: dict + The recorded data, as returned by to_dict(), or None. + """ + if data is None: + return + + self.lower = [list(x) for x in data["lower"]] + self.upper = [list(x) for x in data["upper"]] + self.distance = [list(x) for x in data["distance"]] + self.set_swap_baseline(data["proposed"], data["accepted"]) + + def rank( + self, + window_devices, + device_capacity, + window_samples, + cycles_left, + cold_rate, + proposed, + accepted, + hot, + votes, + ): + """ + Rank the windows by the uncertainty they contribute to the free energy + and choose the hot windows. + + 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 the + minimum number of 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. + + proposed: np.ndarray + The number of swaps proposed between each pair of windows so far. + + accepted: np.ndarray + The number of swaps accepted between each pair of windows so far. + + hot: list of int + The current hot windows, or None if every window is hot. + + votes: list of int + The number of consecutive rankings at which each window has been + found cold. + + Returns + ------- + + hot: list of int + The hot windows, or None if every window is hot. + + votes: list of int + The updated number of consecutive rankings at which each window + has been found cold. + """ + votes = list(votes) + + # There is nothing left to allocate, so keep the current hot windows. + if cycles_left <= 0: + _logger.info("Not ranking windows: no cycles left, keeping hot replicas") + return hot, votes + + num_windows = self._num_windows + lower = [_np.array(x) for x in self.lower] + upper = [_np.array(x) for x in self.upper] + + min_samples = self._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 hot, votes + + g_lower = [statistical_inefficiency(x) for x in lower] + g_upper = [statistical_inefficiency(x) for x in upper] + + # Windows with too few independent samples can't be made cold. + undersampled = [] + for k in range(num_windows): + series = [] + if k > 0: + series.append((lower[k], g_lower[k])) + if k < num_windows - 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._is_ring_break: + for k, x in enumerate(self.distance): + x = _np.array(x) + if len(x) / statistical_inefficiency(x) < min_samples: + slow_distance.append(k) + + proposed = proposed - self.proposed + accepted = accepted - self.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_windows - 1) + rev_variance = _np.zeros(num_windows - 1) + flagged = [] + for k in range(num_windows - 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 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_windows + } + + all_hot = set(range(num_windows)) + + # 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_windows) + 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_windows, 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 + + if chosen is None: + new_hot = all_hot + votes = [0] * num_windows + else: + # Keep windows hot until they have been found cold twice in a row. + previous = all_hot if hot is None else set(hot) + new_hot = set(chosen) + for r in range(num_windows): + if r in chosen: + votes[r] = 0 + elif r in previous: + votes[r] += 1 + if votes[r] < 2: + new_hot.add(r) + + # Windows kept hot can leave no gain over all hot. + new_hot = fill(new_hot) + if _np.prod(score(new_hot)) >= threshold: + new_hot = all_hot + + hot_variance = score(new_hot)[1] + new_hot = sorted(new_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: {new_hot}" + if len(new_hot) < num_windows + 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(new_hot).sum() / num_windows + 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" + ) + + return (None if len(new_hot) == num_windows else new_hot), votes diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index c6e1aea..df46a30 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -22,6 +22,7 @@ __all__ = ["RunnerBase"] from glob import glob as _glob +from math import isclose as _isclose from pathlib import Path as _Path from shutil import copyfile as _copyfile @@ -56,7 +57,38 @@ def __init__(self, system, config): config: :class: `Config ` The configuration options for the simulation. """ + self._load_system(system) + self._set_config(config) + self._check_inputs() + + # The system is fully prepared before the reference copy is stored and + # the restart is checked, since the checkpoints must match it. + pert_idxs, ion_indices = self._prepare_system() + self._set_schedule(pert_idxs, ion_indices) + self._apply_hmr() + + # Validate the samplers before anything is written to the output + # directory. + self._check_gcmc() + self._set_terminal_groups() + + self._prepare_output_and_restart() + self._finalise_systems() + self._set_dynamics_kwargs() + self._limit_sire_threads() + self._set_tuned_pme_parameters() + + def _load_system(self, system): + """ + Load the system to be simulated. + + Parameters + ---------- + system: str, :class: `System ` + The perturbable system to be simulated. This can be either a path + to a stream file, or a Sire system object. + """ if not isinstance(system, (str, _System)): msg = "'system' must be of type 'str' or 'sire.system.System'" _logger.error(msg) @@ -73,7 +105,16 @@ def __init__(self, system, config): else: self._system = system.clone() - # Validate the configuration. + def _set_config(self, config): + """ + Validate the configuration and set up the output directory. + + Parameters + ---------- + + config: :class: `Config ` + The configuration options for the simulation. + """ if not isinstance(config, _Config): msg = "'config' must be of type 'somd2.config.Config'" _logger.error(msg) @@ -90,6 +131,11 @@ def __init__(self, system, config): # blocks only exist to update the bar. _sr.base.ProgressBar.set_silent() + def _check_inputs(self): + """ + Check the perturbed system and the frame frequency, log the versions of + the packages used, and set whether frames are saved. + """ if self._config.replica_exchange and self._config.perturbed_system is not None: # Make sure the number of positions is correct. num_atoms = self._system.num_atoms() @@ -143,6 +189,21 @@ def __init__(self, system, config): _logger.error(msg) raise ValueError(msg) + def _prepare_system(self): + """ + Prepare the system for simulation: apply any perturbation file, check + the perturbable molecules, set the simulation options that depend on + the system, modify the bonded terms, and add alchemical ions. + + Returns + ------- + + pert_idxs: list of int + The indices of the perturbable atoms. + + ion_indices: list of int + The indices of the alchemical ions. + """ # Check whether we need to apply a perturbation to the reference system. if self._config.pert_file is not None: _logger.info( @@ -188,17 +249,8 @@ def __init__(self, system, config): # Link properties to the lambda = 0 end state. self._system = _sr.morph.link_to_reference(self._system) - # Whether this is a ring-breaking schedule. - if ( - self._config._lambda_schedule_name is not None - and "ring_break" in self._config._lambda_schedule_name - ): - self._is_ring_breaking = True - else: - self._is_ring_breaking = False - # Check to see if the end-state connectivities are the same. - if not self._is_ring_breaking: + if not self._is_ring_break: for mol in self._system["property is_perturbable"].molecules(): has_end_state_connectivity = False try: @@ -247,7 +299,7 @@ def __init__(self, system, config): # Check for water. try: # The search will fail if there are no water molecules. - water = self._system["water"].molecules() + self._system["water"].molecules() self._has_water = True except: self._has_water = False @@ -355,9 +407,21 @@ def __init__(self, system, config): _logger.error(msg) raise RuntimeError(msg) - # Convert water topology to AMBER format if not already done. AMBER - # format adds an explicit H-H bond, giving fully rigid water (O-H and - # H-H constraints) rather than just O-H constraints under h_bonds. + self._convert_water_to_amber() + + # Check the end state constraints. + self._check_end_state_constraints() + + ion_indices = self._add_alchemical_ions() + + return pert_idxs, ion_indices + + def _convert_water_to_amber(self): + """ + Convert the water topology to AMBER format if not already done. AMBER + format adds an explicit H-H bond, giving fully rigid water (O-H and H-H + constraints) rather than just O-H constraints under h_bonds. + """ try: waters = self._system["water"] except: @@ -394,9 +458,17 @@ def __init__(self, system, config): _logger.error("Unable to convert water topology to AMBER format.") raise e - # Check the end state constraints. - self._check_end_state_constraints() + def _add_alchemical_ions(self): + """ + Add alchemical ions to keep the charge constant when the end states + have different charges. + Returns + ------- + + ion_indices: list of int + The indices of the alchemical ions. + """ # Get the charge difference between the two end states. charge_diff = self._get_charge_difference(self._system) @@ -458,6 +530,22 @@ def __init__(self, system, config): coalchemical_restraints ) + return ion_indices + + def _set_schedule(self, pert_idxs, ion_indices): + """ + Set the soft-core form, the lambda schedule and values, and the REST2 + scaling factors. + + Parameters + ---------- + + pert_idxs: list of int + The indices of the perturbable atoms. + + ion_indices: list of int + The indices of the alchemical ions. + """ # Set the soft-core form. if self._config._softcore_form == "auto": schedule_name = self._config._lambda_schedule_name or "custom" @@ -538,8 +626,6 @@ def __init__(self, system, config): _logger.error(msg) raise ValueError(msg) - from math import isclose - # Set the REST2 scale factors. is_rest2 = False if self._config.rest2_scale is not None: @@ -558,18 +644,21 @@ def __init__(self, system, config): if len(self._config.rest2_scale) != len(self._lambda_energy): msg = f"Length of 'rest2_scale' must match the number of {_lam_sym} values." if is_missing: - msg += "If you have omitted some 'lambda_values` from `lambda_energy`, please " - "add them to `lambda_energy`, along with the corresponding `rest2_scale` values." + msg += ( + " If you have omitted some 'lambda_values' from 'lambda_energy', " + "please add them to 'lambda_energy', along with the " + "corresponding 'rest2_scale' values." + ) _logger.error(msg) raise ValueError(msg) # Make sure the end states are close to 1.0. - if isclose(self._lambda_energy[0], 0.0, abs_tol=1e-4): - if not isclose(self._config.rest2_scale[0], 1.0, abs_tol=1e-4): + if _isclose(self._lambda_energy[0], 0.0, abs_tol=1e-4): + if not _isclose(self._config.rest2_scale[0], 1.0, abs_tol=1e-4): msg = f"'rest2_scale' must be 1.0 at {_lam_sym}=0." _logger.error(msg) raise ValueError(msg) - if isclose(self._lambda_energy[-1], 1.0, abs_tol=1e-4): - if not isclose(self._config.rest2_scale[-1], 1.0, abs_tol=1e-4): + if _isclose(self._lambda_energy[-1], 1.0, abs_tol=1e-4): + if not _isclose(self._config.rest2_scale[-1], 1.0, abs_tol=1e-4): msg = f"'rest2_scale' must be 1.0 at {_lam_sym}=1." _logger.error(msg) raise ValueError(msg) @@ -577,7 +666,7 @@ def __init__(self, system, config): # If there are any non-zero REST2 scale factors, then log it. if any( - not isclose(factor, 1.0, abs_tol=1e-4) + not _isclose(factor, 1.0, abs_tol=1e-4) for factor in self._rest2_scale_factors ): is_rest2 = True @@ -613,7 +702,12 @@ def __init__(self, system, config): if is_rest2: _logger.info(f"REST2 selection contains {len(idxs)} atoms: {idxs}") - # Apply hydrogen mass repartitioning. + def _apply_hmr(self): + """ + Apply hydrogen mass repartitioning, if requested, accounting for any + repartitioning already applied to the system. + """ + if self._config.hmr: # Work out the current hydrogen mass factor. factor_non_water, factor_water = self._get_h_mass_factor(self._system) @@ -627,7 +721,7 @@ def __init__(self, system, config): # We don't support repartiioning water molecules, so check those first. if factor_water is not None: - if not isclose(factor_water, 1.0, abs_tol=1e-4): + if not _isclose(factor_water, 1.0, abs_tol=1e-4): msg = ( "Water molecules have already been repartitioned with " f"a factor of {factor_water:.3f}. We only support " @@ -638,12 +732,12 @@ def __init__(self, system, config): # HMR has already been applied. if factor_non_water is not None: - if not isclose(factor_non_water, 1.0, abs_tol=1e-4): + if not _isclose(factor_non_water, 1.0, abs_tol=1e-4): _logger.info( f"Detected existing hydrogen mass repartioning factor of {factor_non_water:.3f}" ) - if not isclose( + if not _isclose( factor_non_water, self._config.h_mass_factor, abs_tol=1e-4 ): new_factor = self._config.h_mass_factor / factor_non_water @@ -661,6 +755,11 @@ def __init__(self, system, config): self._system, self._config.h_mass_factor ) + def _prepare_output_and_restart(self): + """ + Prepare the output files, store the reference system, and check for a + valid restart. + """ # Flag whether this is a GPU simulation. self._is_gpu = self._config.platform in ["cuda", "opencl", "hip"] @@ -744,7 +843,11 @@ def __init__(self, system, config): self._config.checkpoint_frequency / self._config.energy_frequency ) - # GCMC specific validation. + def _check_gcmc(self): + """ + Validate the GCMC options, if GCMC is used. + """ + if self._config.gcmc: if self._config.platform not in ["cuda", "opencl"]: msg = "GCMC simulations require the CUDA or OpenCL platform." @@ -756,22 +859,17 @@ def __init__(self, system, config): _logger.error(msg) raise ValueError(msg) - if isinstance(self._system, list): - mols = self._system[0] - else: - mols = self._system - # Check that the system is solvated with water molecules. This # is required for GCMC simulations since the existing waters # provide a template for the ghost waters. try: - water = mols["water"].molecules()[0] + self._system["water"].molecules()[0] except: msg = "No water molecules in the system. Cannot perform GCMC." _logger.error(msg) raise ValueError(msg) - # Make sure the frame frequency is a multiple of the energy frequency. + # Make sure the GCMC frequency is a multiple of the energy frequency. if self._config.gcmc_frequency is None: self._config.gcmc_frequency = str(self._config.energy_frequency) msg = ( @@ -786,7 +884,7 @@ def __init__(self, system, config): ).value() # Make sure it's an integer. - if not isclose(ratio, round(ratio), abs_tol=1e-4): + if not _isclose(ratio, round(ratio), abs_tol=1e-4): msg = "'gcmc_frequency' must be a multiple of 'energy_frequency'." _logger.error(msg) raise ValueError(msg) @@ -797,7 +895,7 @@ def __init__(self, system, config): ).value() # Make sure it's an integer. - if not isclose(ratio, round(ratio), abs_tol=1e-4): + if not _isclose(ratio, round(ratio), abs_tol=1e-4): msg = "'frame_frequency' must be a multiple of 'energy_frequency'." _logger.error(msg) raise ValueError(msg) @@ -806,7 +904,7 @@ def __init__(self, system, config): # with them. ratio = (self._config.frame_frequency / self._config.gcmc_frequency).value() - if not isclose(ratio, round(ratio), abs_tol=1e-4): + 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) @@ -819,7 +917,7 @@ def __init__(self, system, config): ).value() # Make sure it's an integer. - if not isclose(ratio, round(ratio), abs_tol=1e-4): + if not _isclose(ratio, round(ratio), abs_tol=1e-4): msg = "'checkpoint_frequency' must be a multiple of 'frame_frequency'." _logger.error(msg) raise ValueError(msg) @@ -830,20 +928,16 @@ def __init__(self, system, config): ratio = (self._config.runtime / self._config.frame_frequency).value() # Make sure it's an integer. - if not isclose(ratio, round(ratio), abs_tol=1e-4): + if not _isclose(ratio, round(ratio), abs_tol=1e-4): msg = "'runtime' must be a multiple of 'frame_frequency'." _logger.error(msg) raise ValueError(msg) # Make sure the selection is valid. if self._config.gcmc_selection is not None: - if isinstance(self._system, list): - mols = self._system[0] - else: - mols = self._system try: - atoms = _sr.mol.selection_to_atoms( - mols, self._config.gcmc_selection + _sr.mol.selection_to_atoms( + self._system, self._config.gcmc_selection ) except: msg = "Invalid 'gcmc_selection' value." @@ -853,17 +947,19 @@ def __init__(self, system, config): # Store the excess chemcical potential value. self._mu_ex = self._config.gcmc_excess_chemical_potential.value() - # Terminal flip specific validation and setup. + def _set_terminal_groups(self): + """ + Validate the terminal flip options and detect the terminal ring groups, + if terminal flips are used. + """ if self._config.terminal_flip_frequency is not None: - from math import isclose - # Make sure the terminal flip frequency is a multiple of the # energy frequency. ratio = ( self._config.terminal_flip_frequency / self._config.energy_frequency ).value() - if not isclose(ratio, round(ratio), abs_tol=1e-4): + if not _isclose(ratio, round(ratio), abs_tol=1e-4): msg = "'terminal_flip_frequency' must be a multiple of 'energy_frequency'." _logger.error(msg) raise ValueError(msg) @@ -871,18 +967,13 @@ def __init__(self, system, config): # Auto-detect terminal ring groups using Sire connectivity. from ._samplers import detect_terminal_groups - if isinstance(self._system, list): - mols = self._system[0] - else: - mols = self._system - flip_angle = ( self._config.terminal_flip_angle.to("degrees").value() if self._config.terminal_flip_angle is not None else None ) self._terminal_groups = detect_terminal_groups( - mols, + self._system, flip_angle=flip_angle, max_mobile_atoms=self._config.terminal_flip_max_mobile_atoms, ) @@ -906,33 +997,21 @@ def __init__(self, system, config): else: self._terminal_groups = [] - # Store the initial system time. + def _finalise_systems(self): + """ + Store the initial time of each system, fill in any missing systems, + and write the end-state topologies. + """ + # Store the initial system time. Windows whose checkpoint couldn't be + # loaded start again from the reference system. if isinstance(self._system, list): - self._initial_time = [] - for system in self._system: + for i, system in enumerate(self._system): if system is None: - self._initial_time.append(_sr.u("0 ps")) - else: - self._initial_time.append(system.time()) + self._system[i] = self._reference_system.clone() + self._initial_time = [system.time() for system in self._system] else: self._initial_time = [self._system.time()] * len(self._lambda_values) - # Check for missing systems in a multi-system simulation. - if isinstance(self._system, list): - ref_system = None - missing_systems = [] - for i, system in enumerate(self._system): - if system is not None: - ref_system = None - else: - missing_systems.append(i) - if ref_system is None: - ref_system = self._reference_system - - # Fill in any missing systems. - for i in missing_systems: - self._system[i] = ref_system.clone() - # Create the lock file name. self._lock_file = str(self._config.output_directory / "somd2.lock") @@ -959,7 +1038,11 @@ def __init__(self, system, config): _logger.warning(f"Couldn't find the molecules with virtual sites: {e}") self._virtual_sites = None - # Update the tajectory page size. + def _set_dynamics_kwargs(self): + """ + Set the options used to create the dynamics objects and GCMC samplers. + """ + # Update the trajectory page size. if self._config.page_size is not None: # Convert from MB to bytes. page_size = int(self._config.page_size * 1024 * 1024) @@ -1044,6 +1127,11 @@ def __init__(self, system, config): "schedule" ].reverse() + def _limit_sire_threads(self): + """ + Find the GPU devices, and limit the number of CPU threads available to + Sire when running on GPUs in parallel. + """ # The physical GPU devices available to this run, as listed by the # visible-devices environment variable for the chosen platform. self._gpu_devices = [] @@ -1088,8 +1176,6 @@ def __init__(self, system, config): # Update the maximum number of threads. _sr.legacy.Base.set_max_num_threads(sire_threads) - self._set_tuned_pme_parameters() - def _set_tuned_pme_parameters(self): """ Tune the PME parameters for a new run, saving them so that restarts @@ -2288,10 +2374,8 @@ def get_last_config(output_directory): # 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( + if last and not _isclose( _sr.u(str(last)).to("ps"), self._config.energy_frequency.to("ps") ): msg = ( @@ -2486,10 +2570,8 @@ def _repartition_h_mass(system, factor=1.0): if not isinstance(factor, float): raise TypeError("'factor' must be of type 'float'") - from math import isclose - # Early exit if no repartitioning is required. - if isclose(factor, 1.0, abs_tol=1e-4): + if _isclose(factor, 1.0, abs_tol=1e-4): return system from sire.morph import ( diff --git a/src/somd2/runner/_dynamics_cache.py b/src/somd2/runner/_dynamics_cache.py new file mode 100644 index 0000000..5299c30 --- /dev/null +++ b/src/somd2/runner/_dynamics_cache.py @@ -0,0 +1,1485 @@ +###################################################################### +# SOMD2: GPU accelerated alchemical free-energy engine. +# +# Copyright: 2023-2026 +# +# Authors: The OpenBioSim Team +# +# SOMD2 is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# SOMD2 is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with SOMD2. If not, see . +##################################################################### + +__all__ = ["DynamicsCache"] + +import numpy as _np + +from somd2 import _logger + +from ._base import RunnerBase as _RunnerBase + + +class DynamicsCache: + """ + A class for caching dynamics objects. + """ + + def __init__( + self, + system, + lambdas, + rest2_scale_factors, + num_gpus, + dynamics_kwargs, + gcmc_kwargs=None, + output_directory=None, + perturbed_system=None, + xml_filenames=None, + num_slots=None, + update_constraints=True, + constraint_lambda_index=None, + gpu_devices=None, + strided=False, + virtual_sites=None, + ): + """ + Constructor. + + Parameters + ---------- + + system: :class: `System `, List[:class: `System `] + The perturbable system, or systems, to be simulated. + + lambdas: np.ndarray + The lambda value for each replica. + + rest2_scale_factors: np.ndarray + The REST2 scaling factor for each replica. + + num_gpus: int + The number of GPUs to use. + + dynamics_kwargs: dict + A dictionary of default dynamics keyword arguments. + + gcmc_kwargs: dict + GCMC specific keyword arguments. If None, then GCMC is not used. + + output_directory: pathlib.Path + The directory for simulation output. + + perturbed_system: :class: `System ` + The perturbed end-state system used to seed starting coordinates for + lambda > 0.5 replicas. If None, the perturbed state is not used. + + xml_filenames: list of str + A list of file paths for the OpenMM XML output, one per replica. + If None, XML files are not written. + + num_slots: int + The number of dynamics objects (slots) to create. If None, then one + is created per replica. If fewer, then each slot is re-used to + propagate several replicas per cycle, changing its lambda value as + it goes. + + update_constraints: bool + Whether to update the constraints when changing the lambda value of + a slot. + + constraint_lambda_index: int + The index of the lambda value to create every context at, so that + constrained bond lengths are the same for every replica. If None, + each context is created at the lambda value of the first replica it + hosts, which is only consistent between replicas when there is a + context each. Only meaningful when 'update_constraints' is False, + since the constraints are otherwise updated whenever lambda + changes. + + gpu_devices: list + 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. + + virtual_sites: np.ndarray + The molecules with virtual sites, so that each molecule's particles + can be found in the saved states. + """ + + num_replicas = len(lambdas) + + if num_slots is None: + num_slots = num_replicas + + # Initialise attributes. + self._lambdas = lambdas + self._rest2_scale_factors = rest2_scale_factors + self._num_replicas = num_replicas + self._num_slots = num_slots + 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 + self._virtual_sites = virtual_sites + self._gcmc_states = [None] * num_replicas + # GCMC statistics for the whole simulation, keyed by lambda value. + self._gcmc_stats = None + self._energy_trajectories = [None] * num_replicas + self._ghost_files = [None] * num_replicas + # Waters in the GCMC volume, recorded per replica while it is resident + # in its slot. Derived, so not stored in a checkpoint. + self._gcmc_num_waters = [None] * num_replicas + # Whether the last mix moved a replica's state, and so whether it must + # be pushed into the context before the next block. + self._state_moved = [False] * num_replicas + self._terminal_flip_stats = [[0, 0] for _ in range(num_replicas)] + self._num_proposed = _np.matrix(_np.zeros((num_replicas, num_replicas))) + self._num_accepted = _np.matrix(_np.zeros((num_replicas, num_replicas))) + self._num_swaps = _np.matrix(_np.zeros((num_replicas, num_replicas))) + + # Build the slot layout and the per-slot attributes. + self._build_slot_layout() + + # Create the dynamics objects. + self._create_dynamics( + system, + lambdas, + rest2_scale_factors, + num_gpus, + dynamics_kwargs, + gcmc_kwargs=gcmc_kwargs, + output_directory=output_directory, + perturbed_system=perturbed_system, + xml_filenames=xml_filenames, + ) + + def _build_slot_layout(self): + """ + Assign replicas to slots. + + Each slot is given a contiguous group of replicas, so that it only ever + moves between neighbouring lambda values. That keeps the change in the + 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 + + 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 + for slot, group in enumerate(self._groups): + for replica in group: + self._replica_slot[replica] = slot + + def __setstate__(self, state): + """ + Set the state of the object. + """ + + # Checkpoints written before slots were introduced stored the states + # unpermuted, with self._states holding the mapping to apply on + # restart. They are detected by the absence of "_num_slots". + is_legacy = "_num_slots" not in state + + for key, value in state.items(): + setattr(self, key, value) + + # Provide defaults for attributes added after the initial release, + # so that old checkpoint files can still be loaded. + n = len(self._lambdas) + if not hasattr(self, "_gcmc_stats"): + self._gcmc_stats = None + if not hasattr(self, "_gcmc_states"): + self._gcmc_states = [None] * n + if not hasattr(self, "_terminal_flip_stats"): + self._terminal_flip_stats = [[0, 0] for _ in range(n)] + if not hasattr(self, "_time"): + self._time = None + if not hasattr(self, "_num_replicas"): + self._num_replicas = n + if not hasattr(self, "_energy_trajectories"): + 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 + # _create_dynamics() when the run is restarted. Older checkpoints + # predate slots entirely, in which case there was one per replica. + if not hasattr(self, "_num_slots"): + self._num_slots = n + if not hasattr(self, "_update_constraints"): + 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 + # applied. Legacy states hold the destination of each replica's + # configuration, so are inverted. + if is_legacy: + sources = _np.argsort(self._states) + self._openmm_states = [self._openmm_states[s] for s in sources] + self._gcmc_states = [self._gcmc_states[s] for s in sources] + + # Every replica is seeded from its stored state on a restart, since the + # contexts are created from the input system rather than the checkpoint. + self._state_moved = [True] * n + + # Rebuild the slot layout, which is derived rather than stored. + self._build_slot_layout() + + # Checkpoints written before a sampler could be re-used across lambda + # values stored the GCMC statistics as a list of counters per replica. + # Convert these to a single dictionary keyed by lambda value. + if isinstance(self._gcmc_stats, list): + converted = {} + for lam, stats in zip(self._lambdas, self._gcmc_stats): + if _RunnerBase._is_legacy_gcmc_stats(stats): + converted.update(_RunnerBase._convert_legacy_gcmc_stats(stats, lam)) + self._gcmc_stats = converted if converted else None + + def __getstate__(self): + """ + Get the state of the object. + """ + + # Create the state dict. + d = { + "_lambdas": self._lambdas, + "_rest2_scale_factors": self._rest2_scale_factors, + "_num_replicas": self._num_replicas, + "_num_slots": self._num_slots, + "_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. + "_gcmc_samplers": len(self._gcmc_samplers) * [None], + "_gcmc_states": self._gcmc_states, + "_gcmc_stats": self._gcmc_stats, + "_terminal_flip_stats": self._terminal_flip_stats, + "_num_proposed": self._num_proposed, + "_num_accepted": self._num_accepted, + "_num_swaps": self._num_swaps, + } + + return d + + def _create_dynamics( + self, + system, + lambdas, + rest2_scale_factors, + num_gpus, + dynamics_kwargs, + gcmc_kwargs=None, + output_directory=None, + perturbed_system=None, + xml_filenames=None, + ): + """ + Create the dynamics objects. + + Parameters + ---------- + + system: :class: `System `, List[:class: `System `] + The perturbable system, or systems, to be simulated. + + lambdas: np.ndarray + The lambda value for each replica. + + rest2_scale_factors: np.ndarray + The REST2 scaling factor for each replica. + + num_gpus: int + The number of GPUs to use. + + dynamics_kwargs: dict + A dictionary of default dynamics keyword arguments. + + gcmc_kwargs: dict + GCMC specific keyword arguments. If None, then GCMC is not used. + + output_directory: pathlib.Path + The directory for simulation output. + + perturbed_system: :class: `System ` + The perturbed end-state system used to seed starting coordinates for + lambda > 0.5 replicas. If None, the perturbed state is not used. + + xml_filenames: list of str + A list of file paths for the OpenMM XML output, one per replica. + If None, XML files are not written. + """ + + from math import floor + + # Copy the dynamics keyword arguments. + dynamics_kwargs = dynamics_kwargs.copy() + + # Store the number of replicas. + num_replicas = len(lambdas) + + # Copy the GCMC keyword arguments. + if gcmc_kwargs is not None: + gcmc_kwargs = gcmc_kwargs.copy() + + # Initialise the dynamics object list. + self._dynamics = [] + + # Per-device memory tracking for estimation. + device_mem = {} + + # 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. + if gcmc_kwargs is not None: + self._ghost_files = [ + str(output_directory / f"gcmc_ghosts_{lam:.5f}.txt") for lam in lambdas + ] + + # Create the dynamics objects in serial. Each slot is created at the + # lambda value of the first replica that it hosts. + for i in range(self._num_slots): + # The replica that seeds this slot. + seed = self._groups[i][0] + + # The replica in the middle of the group, used to choose which end + # state the starting coordinates come from. A slot's context is + # created from a single system and every replica it hosts starts + # from that context, so taking the middle rather than the first + # keeps any mismatch to at most half a group, next to the lambda + # value at which the end state switches. + middle = self._groups[i][len(self._groups[i]) // 2] + + lam = lambdas[seed] + scale = rest2_scale_factors[seed] + + # Create the context at a common lambda value, so that constrained + # bond lengths are the same for every replica. Only set when the + # constraints aren't updated as the slot changes lambda. + if self._constraint_lambda_index is None: + build_lam = lam + else: + build_lam = lambdas[self._constraint_lambda_index] + + # Work out the device index. + device = self._slot_devices[i] + + # Record baseline memory before the first slot on this device. + if device not in device_mem: + used_before, _, total_mem = self._check_device_memory( + self._physical_device(device) + ) + device_mem[device] = { + "before": used_before, + "total": total_mem, + "count": 0, + } + + # Whether to seed from the perturbed end state. Swapping the end + # states reverses the lambda schedule, so the perturbed end state + # is then at lambda = 0 rather than lambda = 1. + seed_perturbed = (lambdas[middle] > 0.5) != dynamics_kwargs[ + "swap_end_states" + ] + + # This is a restart, get the system for the seeding replica. + if isinstance(system, list): + mols = system[seed] + # This is a new simulation. Use the perturbed system to seed the + # starting coordinates and periodic space. + elif perturbed_system is not None and seed_perturbed: + mols = perturbed_system + else: + mols = system + + # Delete an existing trajectory frames. + mols.delete_all_frames() + + # Overload the device and lambda value. + dynamics_kwargs["device"] = device + dynamics_kwargs["lambda_value"] = build_lam + dynamics_kwargs["rest2_scale"] = scale + + if gcmc_kwargs is not None: + try: + from loch import GCMCSampler + except: + msg = "loch is not installed. GCMC sampling cannot be performed." + _logger.error(msg) + + # Create the GCMC sampler, telling it every lambda value that + # this slot will host so that switching between them doesn't + # need to build an OpenMM context. + gcmc_sampler = GCMCSampler( + mols, + device=device, + lambda_value=lam, + rest2_scale=scale, + lambda_values=[lambdas[r] for r in self._groups[i]], + rest2_scales=[rest2_scale_factors[r] for r in self._groups[i]], + ghost_file=self._ghost_files[seed], + **gcmc_kwargs, + ) + + # Get the modified GCMC system. + mols = gcmc_sampler.system() + + # Store the GCMC sampler. + self._gcmc_samplers[i] = gcmc_sampler + + _logger.info( + f"Created GCMC sampler for lambda {lam:.5f} on device {device}" + ) + + # Create the dynamics object. + try: + dynamics = mols.dynamics(**dynamics_kwargs) + except Exception as e: + # Contexts are appended as they are created, so this is the + # number that fit before the failure. + num_created = len(self._dynamics) + msg = ( + f"Could not create dynamics object for lambda {lam:.5f} on " + f"device {device} after creating {num_created} of " + f"{self._num_slots} contexts: {e}" + ) + # Failing partway through points at exhausted resources. + # Failing on the first context does not, so don't advise on + # the number of contexts in that case. + if num_created > 0: + msg += f". {self._max_contexts_advice(num_created)}" + _logger.error(msg) + raise RuntimeError(msg) from e + + # Bind the GCMC sampler to the dynamics object. This allows the + # dynamics object to reset the water state in its internal OpenMM + # context following a crash recovery. + if gcmc_kwargs is not None: + gcmc_sampler.bind_dynamics(dynamics) + + # Append the dynamics object. + self._dynamics.append(dynamics) + + # Write the OpenMM XML file to the output directory. This is + # indexed by replica, so use the replica that seeded the slot. + if xml_filenames is not None: + _logger.info( + f"Writing OpenMM XML for lambda {lam:.5f} on device {device}" + ) + dynamics.to_xml(xml_filenames[seed]) + + # Track memory footprint for this device. + info = device_mem[device] + info["count"] += 1 + num_contexts = contexts_per_device[device] + + # Estimate memory after the first or second replica. + if info["count"] == 1: + used_mem, _, _ = self._check_device_memory( + self._physical_device(device) + ) + info["after_first"] = used_mem + + if num_contexts == 1: + # Only one replica on this device, use actual measurement. + est_total = used_mem + else: + # Wait for the second replica to get the marginal cost. + est_total = None + + elif info["count"] == 2: + used_mem, _, _ = self._check_device_memory( + self._physical_device(device) + ) + # The first replica includes one-time context overhead. + # The marginal cost of subsequent replicas is the difference + # between the second and first. + first_cost = info["after_first"] - info["before"] + marginal_cost = used_mem - info["after_first"] + est_total = ( + info["before"] + first_cost + marginal_cost * (num_contexts - 1) + ) + _logger.info( + f"Memory per replica on device {device}: " + f"first = {first_cost / (1024**2):.0f} MiB, " + f"marginal = {marginal_cost / (1024**2):.0f} MiB" + ) + else: + est_total = None + + if est_total is not None: + total_mem = info["total"] + + # If this exceeds the total memory, raise an error. + if est_total > total_mem: + baseline = info["before"] + replica_cost = first_cost + marginal_cost * (num_contexts - 1) + + # How many contexts fit on this device, given the measured + # cost of the first and the marginal cost of the rest. The + # slots are spread evenly over the devices, so the total is + # capped by what the busiest device can hold. + if marginal_cost > 0: + per_device = 1 + floor( + (total_mem - baseline - first_cost) / marginal_cost + ) + else: + per_device = info["count"] + + msg = ( + f"Not enough memory on device {device} for all assigned replicas. " + f"Baseline usage before simulation: {baseline / (1024**3):.2f} GB " + f"Estimated replica memory: {replica_cost / (1024**3):.2f} GB, " + f"Total estimated: {est_total / (1024**3):.2f} GB, " + f"Available memory: {total_mem / (1024**3):.2f} GB. " + f"{self._max_contexts_advice(per_device * num_gpus)}" + ) + _logger.error(msg) + raise MemoryError(msg) + + # If there's less than 20% free memory, raise a warning. + elif ((total_mem - est_total) / total_mem) < 0.2: + _logger.warning( + f"Device {device} will have less than 20% free memory " + f"after creating all assigned replicas. " + f"{est_total / (1024**3):.2f} GB, " + f"Available memory: {total_mem / (1024**3):.2f} GB." + ) + + else: + _logger.info( + f"Estimated memory usage on device {device} after creating all replicas: " + f"{est_total / (1024**3):.2f} GB, " + f"Available memory: {total_mem / (1024**3):.2f} GB." + ) + + _logger.info( + f"Created dynamics object for lambda {lam:.5f} on device {device}" + ) + + # Leave the slot marked as holding no replica, so that the first + # call to load_replica() does the full setup (lambda value, GCMC + # parameters, ghost file and sampling statistics) rather than + # assuming the seeding replica is already fully installed. + self._slot_replica[i] = None + + # Give each replica its own energy trajectory. These are seeded from a + # slot's own so that the "ensemble" property is carried over. A slot + # accumulates into the trajectory of whichever replica it hosts. + for replica in range(self._num_replicas): + slot = self._replica_slot[replica] + self._energy_trajectories[replica] = self._dynamics[ + slot + ]._d.energy_trajectory() + + # Seed the starting state for every replica from the context of the + # slot that hosts it. The GCMC water state must be seeded too, since + # load_replica() diffs against it: a replica with no stored state would + # be skipped, leaving the sampler holding the water configuration of + # whichever replica used the slot last. + for replica in range(self._num_replicas): + slot = self._replica_slot[replica] + + if self._openmm_states[replica] is None: + self.save_openmm_state(slot, replica) + + if self._gcmc_samplers[slot] is not None and ( + self._gcmc_states[replica] is None + ): + self.save_gcmc_state(slot, replica) + + def slot_for(self, replica): + """ + Return the index of the slot that hosts a given replica. + + Parameters + ---------- + + replica: int + The index of the replica. + + Returns + ------- + + int + The index of the slot. + """ + 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. + + When there is one slot per replica the slot and replica indices are + the same. Otherwise use slot_for() to map a replica to its slot. + + Parameters + ---------- + + slot: int + The index of the slot. + + Returns + ------- + + tuple + The dynamics object for the slot and its GCMC sampler. + """ + return self._dynamics[slot], self._gcmc_samplers[slot] + + def set(self, slot, dynamics): + """ + Set the dynamics object for a given slot. + + Parameters + ---------- + + slot: int + The index of the slot. + + dynamics: sire.legacy.Convert.SOMMContext + The dynamics object. + """ + self._dynamics[slot] = dynamics + + def delete(self, slot): + """ + Delete the dynamics object for a given slot. + + Parameters + ---------- + + slot: int + The index of the slot. + """ + self._dynamics[slot] = None + + def save_openmm_state(self, slot, replica): + """ + Save the state of a slot's dynamics object as the state of a replica. + + Parameters + ---------- + + slot: int + The index of the slot. + + replica: int + The index of the replica whose state this is. + """ + + # Get the current OpenMM state. + state = ( + self._dynamics[slot] + .context() + .getState(getPositions=True, getVelocities=True) + ) + + # Store positions, velocities, and box vectors as compact numpy arrays + # rather than the OpenMM State object, which serialises to XML when + # pickled and is orders of magnitude larger. + self._openmm_states[replica] = { + "positions": state.getPositions(asNumpy=True), + "velocities": state.getVelocities(asNumpy=True), + "box": state.getPeriodicBoxVectors(asNumpy=True), + } + + @staticmethod + def _get_positions(state): + """ + Return the positions from a saved OpenMM state. + + Parameters + ---------- + + state: dict or openmm.State + The state to read. Dicts (new format) hold the positions directly. + A bare openmm.State is accepted for backwards compatibility with + old checkpoint files, as it is by _apply_openmm_state(). + + Returns + ------- + + openmm.unit.Quantity + The positions. + """ + if isinstance(state, dict): + return state["positions"] + + return state.getPositions(asNumpy=True) + + @staticmethod + def _apply_openmm_state(context, state): + """ + Apply a saved OpenMM state to a context. + + Parameters + ---------- + + context: openmm.Context + The OpenMM context to update. + + state: dict or openmm.State + The state to apply. Dicts (new format) contain "positions", + "velocities", and "box" numpy arrays. A bare openmm.State is + accepted for backwards compatibility with old checkpoint files. + + Note that the step count and simulation time carried by an openmm.State + are deliberately not restored for the dict format. They are held + separately and applied by Dynamics._set_clock(). + """ + if isinstance(state, dict): + # Set the box before the positions, since a barostat may have + # changed it between the state being saved and restored. + context.setPeriodicBoxVectors(*state["box"]) + context.setPositions(state["positions"]) + context.setVelocities(state["velocities"]) + else: + # Legacy openmm.State from checkpoint files written before this + # format change. + context.setState(state) + + def save_gcmc_state(self, slot, replica): + """ + Save the current GCMC water state of a slot as that of a replica. + + Parameters + ---------- + + slot: int + The index of the slot. + + replica: int + The index of the replica whose state this is. + """ + # Get the GCMC sampler. + gcmc_sampler = self._gcmc_samplers[slot] + + # Store the state. + self._gcmc_states[replica] = gcmc_sampler.water_state() + + def get_clock(self, replica=None): + """ + Get the simulation clock. + + 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(), 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): + """ + Get the states of the dynamics objects. + + Returns + ------- + + np.ndarray + The states. + """ + return self._states.copy() + + def set_states(self, states): + """ + Set the states of the dynamics objects. + + Parameters + ---------- + + states: np.ndarray + The new states. + """ + self._states = states + + def load_replica(self, replica, clock=None): + """ + Make a replica resident in its slot, ready for a dynamics block. + + This sets the slot's lambda value, pushes the replica's state into the + OpenMM context, and points the slot at the replica's clock, energy + trajectory, GCMC water state and ghost file. + + Work that isn't needed is skipped. When there is one slot per replica + and no swap has taken place, the slot already holds everything the + replica needs and this reduces to a handful of comparisons. + + Only touches the replica's own slot, so this is safe to call + concurrently for replicas in different slots. + + Parameters + ---------- + + replica: int + The index of the replica. + + clock: dict + The simulation clock to restore, as returned by + Dynamics._get_clock(). If None, the slot's clock is left alone. + """ + slot = self._replica_slot[replica] + dynamics = self._dynamics[slot] + gcmc_sampler = self._gcmc_samplers[slot] + + # The replica that the slot currently holds. + resident = self._slot_replica[slot] + + # The slot needs new positions and velocities if it is being handed a + # different replica, or if the last mix moved this replica's state. + if resident != replica or self._state_moved[replica]: + self._apply_openmm_state(dynamics.context(), self._openmm_states[replica]) + + # Positions have changed underneath the context, so any cached + # energies are stale. + dynamics.clear_energy_cache() + + self._state_moved[replica] = False + + # Set the lambda value and REST2 scaling factor. This is a no-op if + # the slot is already at this lambda value. + if resident != replica: + dynamics.set_lambda( + self._lambdas[replica], + rest2_scale=self._rest2_scale_factors[replica], + update_constraints=self._update_constraints, + ) + + # Restore the clock and point the slot at this replica's energy + # trajectory, so that energies are accumulated against the replica + # rather than the slot. + if clock is not None: + dynamics._set_clock(clock) + dynamics.set_energy_trajectory(self._energy_trajectories[replica]) + + if gcmc_sampler is not None: + gcmc_sampler.push() + try: + # Swap the water state into the sampler. Diff against what the + # sampler currently holds, which is the state of whichever + # replica was last resident. + target = self._gcmc_states[replica] + if target is not None: + current = gcmc_sampler.water_state() + water_idxs = _np.where(current != target)[0] + + if len(water_idxs) > 0: + gcmc_sampler._set_water_state( + dynamics.context(), + indices=water_idxs, + states=target[water_idxs], + ) + + if resident != replica: + # Update the lambda dependent non-bonded parameters used to + # evaluate insertion and deletion energies, and append ghost + # residues to this replica's file. The sampler keeps its + # statistics per lambda value, so switching also switches + # to this replica's. + gcmc_sampler.set_lambda( + self._lambdas[replica], self._rest2_scale_factors[replica] + ) + gcmc_sampler.set_ghost_file(self._ghost_files[replica]) + finally: + gcmc_sampler.pop() + + self._slot_replica[slot] = replica + + def store_replica(self, replica): + """ + Save the state of a replica back out of its slot, so that the slot can + be handed to another replica. + + Parameters + ---------- + + replica: int + The index of the 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) + + # Count the waters against the slot's context rather than whichever + # one the sampler happens to be bound to, which is unset after the + # sampler has been reset. + gcmc_sampler = self._gcmc_samplers[slot] + gcmc_sampler.push() + try: + self._gcmc_num_waters[replica] = gcmc_sampler.num_waters( + context=self._dynamics[slot].context() + ) + finally: + gcmc_sampler.pop() + + def mix_states(self): + """ + Apply the result of a replica mix. + + The states are permuted here, but not pushed into the OpenMM contexts. + They are applied lazily by load_replica(), which is the only point at + which a slot is known to be free. This does the same amount of work as + applying them eagerly, since load_replica() pushes a state exactly when + the mix moved it. + + The permutation must happen here rather than being resolved lazily + through self._states. A slot is re-used within a cycle, so a replica + may be loaded after another replica has already stored its post-run + state; reading through the indirection at that point would pick up the + new state rather than the pre-mix one. + """ + # Permute the travelling state. This is a reference shuffle, so it is + # cheap even for large systems. Statistics and output files stay with + # the lambda window, so are not permuted. + self._openmm_states = [self._openmm_states[state] for state in self._states] + self._gcmc_states = [self._gcmc_states[state] for state in self._states] + + # Flag the replicas whose state moved, so that load_replica() knows it + # has to push new positions and velocities into the context. The flags + # accumulate, and are only cleared once the state has been pushed, so + # that a replica mixed twice without being loaded in between isn't left + # starting from whatever its context happens to hold. + self._state_moved = [ + moved or bool(state != i) + for i, (state, moved) in enumerate(zip(self._states, self._state_moved)) + ] + + # Update the swap matrix with each configuration's move. + for i, state in enumerate(self._states): + self._num_swaps[state, i] += 1 + + def get_proposed(self): + """ + Return the number of proposed swaps between replicas. + """ + return self._num_proposed + + def get_accepted(self): + """ + Return the number of accepted swaps between replicas. + """ + return self._num_accepted + + def get_swaps(self): + """ + Return the swap matrix. + """ + return self._num_swaps + + def prepare_restart( + self, + system, + lambdas, + rest2_scale_factors, + num_gpus, + dynamics_kwargs, + gcmc_kwargs=None, + output_directory=None, + num_slots=None, + update_constraints=True, + constraint_lambda_index=None, + gpu_devices=None, + strided=False, + virtual_sites=None, + ): + """ + Prepare a cache loaded from a checkpoint to carry on the simulation. + The settings that aren't stored in a checkpoint are taken from the + current configuration, so that the number of contexts and the GPUs can + change on restart. The slot layout is rebuilt, the dynamics objects are + created, and the GCMC statistics are handed to the new samplers. + + The OpenMM contexts are not reset here. Each replica's state is pushed + into its slot by load_replica() at the start of its first block, which + is the only point at which the slot is known to be free. + + Parameters + ---------- + + system: :class: `System `, List[:class: `System `] + The perturbable system, or systems, to be simulated. + + lambdas: np.ndarray + The lambda value for each replica. + + rest2_scale_factors: np.ndarray + The REST2 scaling factor for each replica. + + num_gpus: int + The number of GPUs to use. + + dynamics_kwargs: dict + A dictionary of default dynamics keyword arguments. + + gcmc_kwargs: dict + GCMC specific keyword arguments. If None, then GCMC is not used. + + output_directory: pathlib.Path + The directory for simulation output. + + num_slots: int + The number of dynamics objects (slots) to create. See the + constructor. + + update_constraints: bool + Whether to update the constraints when changing the lambda value of + a slot. + + constraint_lambda_index: int + The index of the lambda value to create every context at. See the + constructor. + + gpu_devices: list + The physical devices backing each OpenMM device index. + + strided: bool + Whether to assign replicas to slots in a strided pattern. + + virtual_sites: np.ndarray + The molecules with virtual sites. Set again, since checkpoints + written by older versions don't store them. + """ + self._gpu_devices = gpu_devices + self._num_slots = self._num_replicas if num_slots is None else num_slots + self._update_constraints = update_constraints + self._constraint_lambda_index = constraint_lambda_index + self._strided = strided + self._virtual_sites = virtual_sites + self._build_slot_layout() + + self._create_dynamics( + system, + lambdas, + rest2_scale_factors, + num_gpus, + dynamics_kwargs, + gcmc_kwargs=gcmc_kwargs, + output_directory=output_directory, + ) + + # A sampler keeps only the lambda values it visits, so each can be + # handed the whole simulation's statistics. + if self._gcmc_stats is not None: + for slot in range(self._num_slots): + _, gcmc_sampler = self.get(slot) + gcmc_sampler.restore_stats(self._gcmc_stats) + + def stored_positions(self, replica): + """ + Return the positions in the stored state of a replica. + + Parameters + ---------- + + replica: int + The index of the replica. + + Returns + ------- + + positions: openmm.unit.Quantity + The positions. + """ + return self._get_positions(self._openmm_states[replica]) + + def gcmc_num_waters(self, replica): + """ + Return the number of waters in the GCMC volume, recorded when a replica + was last propagated. + + Parameters + ---------- + + replica: int + The index of the replica. + + Returns + ------- + + int + The number of waters, or None if not recorded. + """ + return self._gcmc_num_waters[replica] + + def reset_hot_replicas(self): + """ + Make every replica hot, clearing the votes for adaptive replica + exchange. + """ + self._hot_replicas = None + self._cold_votes = [0] * self._num_replicas + + @property + def num_replicas(self): + """ + The number of replicas. + """ + return self._num_replicas + + @property + def num_slots(self): + """ + The number of dynamics objects (slots). + """ + return self._num_slots + + @property + def groups(self): + """ + The replicas hosted by each slot. + """ + return self._groups + + @property + def time(self): + """ + The simulation time at the last checkpoint, or None. + """ + return self._time + + @time.setter + def time(self, time): + self._time = time + + @property + def hot_replicas(self): + """ + The hot replicas of adaptive replica exchange, or None if every + replica is hot. + """ + return self._hot_replicas + + @hot_replicas.setter + def hot_replicas(self, hot_replicas): + self._hot_replicas = hot_replicas + + @property + def cold_votes(self): + """ + The number of consecutive rankings at which each replica has been + found cold. + """ + return self._cold_votes + + @cold_votes.setter + def cold_votes(self, cold_votes): + self._cold_votes = cold_votes + + @property + def adaptive_data(self): + """ + The data used to rank the windows for adaptive replica exchange, or + None. + """ + return self._adaptive_data + + @adaptive_data.setter + def adaptive_data(self, adaptive_data): + self._adaptive_data = adaptive_data + + @property + def gcmc_stats(self): + """ + The GCMC statistics for the whole simulation, keyed by lambda value, or + None. + """ + return self._gcmc_stats + + @gcmc_stats.setter + def gcmc_stats(self, gcmc_stats): + self._gcmc_stats = gcmc_stats + + @property + def terminal_flip_stats(self): + """ + The number of terminal flip moves attempted and accepted by each + replica. + """ + return self._terminal_flip_stats + + @terminal_flip_stats.setter + def terminal_flip_stats(self, terminal_flip_stats): + self._terminal_flip_stats = terminal_flip_stats + + def _physical_device(self, device): + """ + Map an OpenMM device index to the physical device backing it. + + Parameters + ---------- + + device: int + The OpenMM device index, which is relative to the visible set. + + Returns + ------- + + int, str + The physical device, i.e. the matching entry from + CUDA_VISIBLE_DEVICES. Falls back to the OpenMM index if the visible + set is unknown, which is correct whenever it starts at zero and is + contiguous. + """ + gpu_devices = getattr(self, "_gpu_devices", None) + + if gpu_devices is None or device >= len(gpu_devices): + return device + + return gpu_devices[device] + + @staticmethod + def _max_contexts_advice(num_contexts): + """ + Advice for a user whose replicas don't fit in GPU memory. + + Parameters + ---------- + + num_contexts: int + The number of OpenMM contexts that are known to fit. + """ + + if num_contexts < 1: + return ( + "A single OpenMM context does not fit on this device. Reduce " + "the size of the system, or use a device with more memory." + ) + + return ( + f"Set 'max_contexts' to {num_contexts} or fewer to re-use contexts " + "across lambda values, which removes the memory limit on the " + "number of replicas. This also requires 'frame_frequency' to equal " + "'checkpoint_frequency'." + ) + + @staticmethod + def _check_device_memory(device=0): + """ + Check the memory usage of the specified GPU device. + + Parameters + ---------- + + device: int, str + The device to query. This is the physical device, i.e. an entry + from CUDA_VISIBLE_DEVICES (or the equivalent for other platforms), + not the index used by OpenMM. OpenMM numbers devices relative to + the visible set, whereas pynvml and pyopencl enumerate all devices + on the machine, so the two only agree when the visible set starts + at zero and is contiguous. CUDA_VISIBLE_DEVICES entries may be + either an index or a UUID. + """ + + device = str(device).strip() + + # A UUID cannot be used to index into the OpenCL device list. + is_uuid = device.startswith("GPU-") or device.startswith("MIG-") + device_index = None if is_uuid else int(device) + + # Try to use pyopencl to detect the GPU vendor. + vendor = None + ocl_device = None + try: + import pyopencl as cl + + if device_index is None: + raise ValueError("Cannot index OpenCL devices by UUID") + + platforms = cl.get_platforms() + all_devices = [] + for platform in platforms: + try: + devices = platform.get_devices(device_type=cl.device_type.GPU) + all_devices.extend(devices) + except Exception: + continue + + if device_index < len(all_devices): + ocl_device = all_devices[device_index] + vendor = ocl_device.vendor + else: + msg = f"Device index {device_index} out of range. Found {len(all_devices)} GPU(s)." + _logger.error(msg) + raise IndexError(msg) + except IndexError: + raise + except Exception: + _logger.warning( + "Could not query GPU platform via OpenCL; falling back to pynvml for NVIDIA detection." + ) + + # NVIDIA: Use pynvml (also used as fallback when OpenCL is unavailable). + if vendor is None or "NVIDIA" in vendor: + try: + import pynvml + + pynvml.nvmlInit() + if is_uuid: + handle = pynvml.nvmlDeviceGetHandleByUUID(device.encode()) + else: + handle = pynvml.nvmlDeviceGetHandleByIndex(device_index) + memory = pynvml.nvmlDeviceGetMemoryInfo(handle) + pynvml.nvmlShutdown() + return (memory.used, memory.free, memory.total) + except Exception as e: + if vendor is None: + msg = f"Could not get GPU memory info for device {device} via OpenCL or pynvml: {e}" + else: + msg = ( + f"Could not get NVIDIA GPU memory info for device {device}: {e}" + ) + _logger.error(msg) + raise RuntimeError(msg) from e + + # AMD: Use OpenCL extension. + elif "AMD" in vendor or "Advanced Micro Devices" in vendor: + try: + import pyopencl as cl + + total = ocl_device.global_mem_size + + # cl_amd_device_attribute_query reports the free memory in + # KBytes, as a list with the whole heap first. + free_memory_info = ocl_device.get_info( + cl.device_info.GLOBAL_FREE_MEMORY_AMD + ) + free_kb = ( + free_memory_info[0] + if isinstance(free_memory_info, (list, tuple)) + else free_memory_info + ) + free = free_kb * 1024 + used = total - free + return (used, free, total) + except Exception as e: + msg = f"Could not get AMD GPU memory info for device {device}: {e}" + _logger.error(msg) + raise RuntimeError(msg) from e + + # Any other vendor: we have no way of querying the free memory. + else: + msg = ( + f"Unable to query the memory of device {device}: " + f"unsupported GPU vendor '{vendor}'." + ) + _logger.error(msg) + raise RuntimeError(msg) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 7771462..6995157 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -35,8 +35,14 @@ from .._utils import _lam_sym +from ._adaptive import WindowRanker as _WindowRanker +from ._adaptive import cold_rate as _cold_rate from ._base import RunnerBase as _RunnerBase +# Imported under its own name, since checkpoints written before it moved +# pickle it as somd2.runner._repex.DynamicsCache. +from ._dynamics_cache import DynamicsCache + @_njit def _mix_subset(replicas, energy_matrix, proposed, accepted): @@ -109,1228 +115,9 @@ def _mix_subset(replicas, energy_matrix, proposed, accepted): 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. - """ - - def __init__( - self, - system, - lambdas, - rest2_scale_factors, - num_gpus, - dynamics_kwargs, - gcmc_kwargs=None, - output_directory=None, - perturbed_system=None, - xml_filenames=None, - num_slots=None, - update_constraints=True, - constraint_lambda_index=None, - gpu_devices=None, - strided=False, - ): - """ - Constructor. - - Parameters - ---------- - - system: :class: `System `, List[:class: `System `] - The perturbable system, or systems, to be simulated. - - lambdas: np.ndarray - The lambda value for each replica. - - rest2_scale_factors: np.ndarray - The REST2 scaling factor for each replica. - - num_gpus: int - The number of GPUs to use. - - dynamics_kwargs: dict - A dictionary of default dynamics keyword arguments. - - gcmc_kwargs: dict - GCMC specific keyword arguments. If None, then GCMC is not used. - - output_directory: pathlib.Path - The directory for simulation output. - - perturbed_system: :class: `System ` - The perturbed end-state system used to seed starting coordinates for - lambda > 0.5 replicas. If None, the perturbed state is not used. - - xml_filenames: list of str - A list of file paths for the OpenMM XML output, one per replica. - If None, XML files are not written. - - num_slots: int - The number of dynamics objects (slots) to create. If None, then one - is created per replica. If fewer, then each slot is re-used to - propagate several replicas per cycle, changing its lambda value as - it goes. - - update_constraints: bool - Whether to update the constraints when changing the lambda value of - a slot. - - constraint_lambda_index: int - The index of the lambda value to create every context at, so that - constrained bond lengths are the same for every replica. If None, - each context is created at the lambda value of the first replica it - hosts, which is only consistent between replicas when there is a - context each. Only meaningful when 'update_constraints' is False, - since the constraints are otherwise updated whenever lambda - changes. - - gpu_devices: list - 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) - - if num_slots is None: - num_slots = num_replicas - - # Initialise attributes. - self._lambdas = lambdas - self._rest2_scale_factors = rest2_scale_factors - self._num_replicas = num_replicas - self._num_slots = num_slots - 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. - self._virtual_sites = None - self._gcmc_states = [None] * num_replicas - # GCMC statistics for the whole simulation, keyed by lambda value. - self._gcmc_stats = None - self._energy_trajectories = [None] * num_replicas - self._ghost_files = [None] * num_replicas - # Waters in the GCMC volume, recorded per replica while it is resident - # in its slot. Derived, so not stored in a checkpoint. - self._gcmc_num_waters = [None] * num_replicas - # Whether the last mix moved a replica's state, and so whether it must - # be pushed into the context before the next block. - self._state_moved = [False] * num_replicas - self._terminal_flip_stats = [[0, 0] for _ in range(num_replicas)] - self._num_proposed = _np.matrix(_np.zeros((num_replicas, num_replicas))) - self._num_accepted = _np.matrix(_np.zeros((num_replicas, num_replicas))) - self._num_swaps = _np.matrix(_np.zeros((num_replicas, num_replicas))) - - # Build the slot layout and the per-slot attributes. - self._build_slot_layout() - - # Create the dynamics objects. - self._create_dynamics( - system, - lambdas, - rest2_scale_factors, - num_gpus, - dynamics_kwargs, - gcmc_kwargs=gcmc_kwargs, - output_directory=output_directory, - perturbed_system=perturbed_system, - xml_filenames=xml_filenames, - ) - - def _build_slot_layout(self): - """ - Assign replicas to slots. - - Each slot is given a contiguous group of replicas, so that it only ever - moves between neighbouring lambda values. That keeps the change in the - 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 - - 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 - for slot, group in enumerate(self._groups): - for replica in group: - self._replica_slot[replica] = slot - - def __setstate__(self, state): - """ - Set the state of the object. - """ - - # Checkpoints written before slots were introduced stored the states - # unpermuted, with self._states holding the mapping to apply on - # restart. They are detected by the absence of "_num_slots". - is_legacy = "_num_slots" not in state - - for key, value in state.items(): - setattr(self, key, value) - - # Provide defaults for attributes added after the initial release, - # so that old checkpoint files can still be loaded. - n = len(self._lambdas) - if not hasattr(self, "_gcmc_stats"): - self._gcmc_stats = None - if not hasattr(self, "_gcmc_states"): - self._gcmc_states = [None] * n - if not hasattr(self, "_terminal_flip_stats"): - self._terminal_flip_stats = [[0, 0] for _ in range(n)] - if not hasattr(self, "_time"): - self._time = None - if not hasattr(self, "_num_replicas"): - self._num_replicas = n - if not hasattr(self, "_energy_trajectories"): - 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 - # _create_dynamics() when the run is restarted. Older checkpoints - # predate slots entirely, in which case there was one per replica. - if not hasattr(self, "_num_slots"): - self._num_slots = n - if not hasattr(self, "_update_constraints"): - 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 - # applied. Legacy states hold the destination of each replica's - # configuration, so are inverted. - if is_legacy: - sources = _np.argsort(self._states) - self._openmm_states = [self._openmm_states[s] for s in sources] - self._gcmc_states = [self._gcmc_states[s] for s in sources] - - # Every replica is seeded from its stored state on a restart, since the - # contexts are created from the input system rather than the checkpoint. - self._state_moved = [True] * n - - # Rebuild the slot layout, which is derived rather than stored. - self._build_slot_layout() - - # Checkpoints written before a sampler could be re-used across lambda - # values stored the GCMC statistics as a list of counters per replica. - # Convert these to a single dictionary keyed by lambda value. - if isinstance(self._gcmc_stats, list): - converted = {} - for lam, stats in zip(self._lambdas, self._gcmc_stats): - if _RunnerBase._is_legacy_gcmc_stats(stats): - converted.update(_RunnerBase._convert_legacy_gcmc_stats(stats, lam)) - self._gcmc_stats = converted if converted else None - - def __getstate__(self): - """ - Get the state of the object. - """ - - # Create the state dict. - d = { - "_lambdas": self._lambdas, - "_rest2_scale_factors": self._rest2_scale_factors, - "_num_replicas": self._num_replicas, - "_num_slots": self._num_slots, - "_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. - "_gcmc_samplers": len(self._gcmc_samplers) * [None], - "_gcmc_states": self._gcmc_states, - "_gcmc_stats": self._gcmc_stats, - "_terminal_flip_stats": self._terminal_flip_stats, - "_num_proposed": self._num_proposed, - "_num_accepted": self._num_accepted, - "_num_swaps": self._num_swaps, - } - - return d - - def _create_dynamics( - self, - system, - lambdas, - rest2_scale_factors, - num_gpus, - dynamics_kwargs, - gcmc_kwargs=None, - output_directory=None, - perturbed_system=None, - xml_filenames=None, - ): - """ - Create the dynamics objects. - - Parameters - ---------- - - system: :class: `System `, List[:class: `System `] - The perturbable system, or systems, to be simulated. - - lambdas: np.ndarray - The lambda value for each replica. - - rest2_scale_factors: np.ndarray - The REST2 scaling factor for each replica. - - num_gpus: int - The number of GPUs to use. - - dynamics_kwargs: dict - A dictionary of default dynamics keyword arguments. - - gcmc_kwargs: dict - GCMC specific keyword arguments. If None, then GCMC is not used. - - output_directory: pathlib.Path - The directory for simulation output. - - perturbed_system: :class: `System ` - The perturbed end-state system used to seed starting coordinates for - lambda > 0.5 replicas. If None, the perturbed state is not used. - - xml_filenames: list of str - A list of file paths for the OpenMM XML output, one per replica. - If None, XML files are not written. - """ - - from math import floor - - # Copy the dynamics keyword arguments. - dynamics_kwargs = dynamics_kwargs.copy() - - # Store the number of replicas. - num_replicas = len(lambdas) - - # Copy the GCMC keyword arguments. - if gcmc_kwargs is not None: - gcmc_kwargs = gcmc_kwargs.copy() - - # Initialise the dynamics object list. - self._dynamics = [] - - # Per-device memory tracking for estimation. - device_mem = {} - - # 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. - if gcmc_kwargs is not None: - self._ghost_files = [ - str(output_directory / f"gcmc_ghosts_{lam:.5f}.txt") for lam in lambdas - ] - - # Create the dynamics objects in serial. Each slot is created at the - # lambda value of the first replica that it hosts. - for i in range(self._num_slots): - # The replica that seeds this slot. - seed = self._groups[i][0] - - # The replica in the middle of the group, used to choose which end - # state the starting coordinates come from. A slot's context is - # created from a single system and every replica it hosts starts - # from that context, so taking the middle rather than the first - # keeps any mismatch to at most half a group, next to the lambda - # value at which the end state switches. - middle = self._groups[i][len(self._groups[i]) // 2] - - lam = lambdas[seed] - scale = rest2_scale_factors[seed] - - # Create the context at a common lambda value, so that constrained - # bond lengths are the same for every replica. Only set when the - # constraints aren't updated as the slot changes lambda. - if self._constraint_lambda_index is None: - build_lam = lam - else: - build_lam = lambdas[self._constraint_lambda_index] - - # Work out the device index. - device = self._slot_devices[i] - - # Record baseline memory before the first slot on this device. - if device not in device_mem: - used_before, _, total_mem = self._check_device_memory( - self._physical_device(device) - ) - device_mem[device] = { - "before": used_before, - "total": total_mem, - "count": 0, - } - - # Whether to seed from the perturbed end state. Swapping the end - # states reverses the lambda schedule, so the perturbed end state - # is then at lambda = 0 rather than lambda = 1. - seed_perturbed = (lambdas[middle] > 0.5) != dynamics_kwargs[ - "swap_end_states" - ] - - # This is a restart, get the system for the seeding replica. - if isinstance(system, list): - mols = system[seed] - # This is a new simulation. Use the perturbed system to seed the - # starting coordinates and periodic space. - elif perturbed_system is not None and seed_perturbed: - mols = perturbed_system - else: - mols = system - - # Delete an existing trajectory frames. - mols.delete_all_frames() - - # Overload the device and lambda value. - dynamics_kwargs["device"] = device - dynamics_kwargs["lambda_value"] = build_lam - dynamics_kwargs["rest2_scale"] = scale - - if gcmc_kwargs is not None: - try: - from loch import GCMCSampler - except: - msg = "loch is not installed. GCMC sampling cannot be performed." - _logger.error(msg) - - # Create the GCMC sampler, telling it every lambda value that - # this slot will host so that switching between them doesn't - # need to build an OpenMM context. - gcmc_sampler = GCMCSampler( - mols, - device=device, - lambda_value=lam, - rest2_scale=scale, - lambda_values=[lambdas[r] for r in self._groups[i]], - rest2_scales=[rest2_scale_factors[r] for r in self._groups[i]], - ghost_file=self._ghost_files[seed], - **gcmc_kwargs, - ) - - # Get the modified GCMC system. - mols = gcmc_sampler.system() - - # Store the GCMC sampler. - self._gcmc_samplers[i] = gcmc_sampler - - _logger.info( - f"Created GCMC sampler for lambda {lam:.5f} on device {device}" - ) - - # Create the dynamics object. - try: - dynamics = mols.dynamics(**dynamics_kwargs) - except Exception as e: - # Contexts are appended as they are created, so this is the - # number that fit before the failure. - num_created = len(self._dynamics) - msg = ( - f"Could not create dynamics object for lambda {lam:.5f} on " - f"device {device} after creating {num_created} of " - f"{self._num_slots} contexts: {e}" - ) - # Failing partway through points at exhausted resources. - # Failing on the first context does not, so don't advise on - # the number of contexts in that case. - if num_created > 0: - msg += f". {self._max_contexts_advice(num_created)}" - _logger.error(msg) - raise RuntimeError(msg) from e - - # Bind the GCMC sampler to the dynamics object. This allows the - # dynamics object to reset the water state in its internal OpenMM - # context following a crash recovery. - if gcmc_kwargs is not None: - gcmc_sampler.bind_dynamics(dynamics) - - # Append the dynamics object. - self._dynamics.append(dynamics) - - # Write the OpenMM XML file to the output directory. This is - # indexed by replica, so use the replica that seeded the slot. - if xml_filenames is not None: - _logger.info( - f"Writing OpenMM XML for lambda {lam:.5f} on device {device}" - ) - dynamics.to_xml(xml_filenames[seed]) - - # Track memory footprint for this device. - info = device_mem[device] - info["count"] += 1 - num_contexts = contexts_per_device[device] - - # Estimate memory after the first or second replica. - if info["count"] == 1: - used_mem, _, _ = self._check_device_memory( - self._physical_device(device) - ) - info["after_first"] = used_mem - - if num_contexts == 1: - # Only one replica on this device, use actual measurement. - est_total = used_mem - else: - # Wait for the second replica to get the marginal cost. - est_total = None - - elif info["count"] == 2: - used_mem, _, _ = self._check_device_memory( - self._physical_device(device) - ) - # The first replica includes one-time context overhead. - # The marginal cost of subsequent replicas is the difference - # between the second and first. - first_cost = info["after_first"] - info["before"] - marginal_cost = used_mem - info["after_first"] - est_total = ( - info["before"] + first_cost + marginal_cost * (num_contexts - 1) - ) - _logger.info( - f"Memory per replica on device {device}: " - f"first = {first_cost / (1024**2):.0f} MiB, " - f"marginal = {marginal_cost / (1024**2):.0f} MiB" - ) - else: - est_total = None - - if est_total is not None: - total_mem = info["total"] - - # If this exceeds the total memory, raise an error. - if est_total > total_mem: - baseline = info["before"] - replica_cost = first_cost + marginal_cost * (num_contexts - 1) - - # How many contexts fit on this device, given the measured - # cost of the first and the marginal cost of the rest. The - # slots are spread evenly over the devices, so the total is - # capped by what the busiest device can hold. - if marginal_cost > 0: - per_device = 1 + floor( - (total_mem - baseline - first_cost) / marginal_cost - ) - else: - per_device = info["count"] - - msg = ( - f"Not enough memory on device {device} for all assigned replicas. " - f"Baseline usage before simulation: {baseline / (1024**3):.2f} GB " - f"Estimated replica memory: {replica_cost / (1024**3):.2f} GB, " - f"Total estimated: {est_total / (1024**3):.2f} GB, " - f"Available memory: {total_mem / (1024**3):.2f} GB. " - f"{self._max_contexts_advice(per_device * num_gpus)}" - ) - _logger.error(msg) - raise MemoryError(msg) - - # If there's less than 20% free memory, raise a warning. - elif ((total_mem - est_total) / total_mem) < 0.2: - _logger.warning( - f"Device {device} will have less than 20% free memory " - f"after creating all assigned replicas. " - f"{est_total / (1024**3):.2f} GB, " - f"Available memory: {total_mem / (1024**3):.2f} GB." - ) - - else: - _logger.info( - f"Estimated memory usage on device {device} after creating all replicas: " - f"{est_total / (1024**3):.2f} GB, " - f"Available memory: {total_mem / (1024**3):.2f} GB." - ) - - _logger.info( - f"Created dynamics object for lambda {lam:.5f} on device {device}" - ) - - # Leave the slot marked as holding no replica, so that the first - # call to load_replica() does the full setup (lambda value, GCMC - # parameters, ghost file and sampling statistics) rather than - # assuming the seeding replica is already fully installed. - self._slot_replica[i] = None - - # Give each replica its own energy trajectory. These are seeded from a - # slot's own so that the "ensemble" property is carried over. A slot - # accumulates into the trajectory of whichever replica it hosts. - for replica in range(self._num_replicas): - slot = self._replica_slot[replica] - self._energy_trajectories[replica] = self._dynamics[ - slot - ]._d.energy_trajectory() - - # Seed the starting state for every replica from the context of the - # slot that hosts it. The GCMC water state must be seeded too, since - # load_replica() diffs against it: a replica with no stored state would - # be skipped, leaving the sampler holding the water configuration of - # whichever replica used the slot last. - for replica in range(self._num_replicas): - slot = self._replica_slot[replica] - - if self._openmm_states[replica] is None: - self.save_openmm_state(slot, replica) - - if self._gcmc_samplers[slot] is not None and ( - self._gcmc_states[replica] is None - ): - self.save_gcmc_state(slot, replica) - - def slot_for(self, replica): - """ - Return the index of the slot that hosts a given replica. - - Parameters - ---------- - - replica: int - The index of the replica. - - Returns - ------- - - int - The index of the slot. - """ - 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. - - When there is one slot per replica the slot and replica indices are - the same. Otherwise use slot_for() to map a replica to its slot. - - Parameters - ---------- - - slot: int - The index of the slot. - - Returns - ------- - - tuple - The dynamics object for the slot and its GCMC sampler. - """ - return self._dynamics[slot], self._gcmc_samplers[slot] - - def set(self, slot, dynamics): - """ - Set the dynamics object for a given slot. - - Parameters - ---------- - - slot: int - The index of the slot. - - dynamics: sire.legacy.Convert.SOMMContext - The dynamics object. - """ - self._dynamics[slot] = dynamics - - def delete(self, slot): - """ - Delete the dynamics object for a given slot. - - Parameters - ---------- - - slot: int - The index of the slot. - """ - self._dynamics[slot] = None - - def save_openmm_state(self, slot, replica): - """ - Save the state of a slot's dynamics object as the state of a replica. - - Parameters - ---------- - - slot: int - The index of the slot. - - replica: int - The index of the replica whose state this is. - """ - - # Get the current OpenMM state. - state = ( - self._dynamics[slot] - .context() - .getState(getPositions=True, getVelocities=True) - ) - - # Store positions, velocities, and box vectors as compact numpy arrays - # rather than the OpenMM State object, which serialises to XML when - # pickled and is orders of magnitude larger. - self._openmm_states[replica] = { - "positions": state.getPositions(asNumpy=True), - "velocities": state.getVelocities(asNumpy=True), - "box": state.getPeriodicBoxVectors(asNumpy=True), - } - - @staticmethod - def _get_positions(state): - """ - Return the positions from a saved OpenMM state. - - Parameters - ---------- - - state: dict or openmm.State - The state to read. Dicts (new format) hold the positions directly. - A bare openmm.State is accepted for backwards compatibility with - old checkpoint files, as it is by _apply_openmm_state(). - - Returns - ------- - - openmm.unit.Quantity - The positions. - """ - if isinstance(state, dict): - return state["positions"] - - return state.getPositions(asNumpy=True) - - @staticmethod - def _apply_openmm_state(context, state): - """ - Apply a saved OpenMM state to a context. - - Parameters - ---------- - - context: openmm.Context - The OpenMM context to update. - - state: dict or openmm.State - The state to apply. Dicts (new format) contain "positions", - "velocities", and "box" numpy arrays. A bare openmm.State is - accepted for backwards compatibility with old checkpoint files. - - Note that the step count and simulation time carried by an openmm.State - are deliberately not restored for the dict format. They are held - separately and applied by Dynamics._set_clock(). - """ - if isinstance(state, dict): - # Set the box before the positions, since a barostat may have - # changed it between the state being saved and restored. - context.setPeriodicBoxVectors(*state["box"]) - context.setPositions(state["positions"]) - context.setVelocities(state["velocities"]) - else: - # Legacy openmm.State from checkpoint files written before this - # format change. - context.setState(state) - - def save_gcmc_state(self, slot, replica): - """ - Save the current GCMC water state of a slot as that of a replica. - - Parameters - ---------- - - slot: int - The index of the slot. - - replica: int - The index of the replica whose state this is. - """ - # Get the GCMC sampler. - gcmc_sampler = self._gcmc_samplers[slot] - - # Store the state. - self._gcmc_states[replica] = gcmc_sampler.water_state() - - def get_clock(self, replica=None): - """ - Get the simulation clock. - - 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(), 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): - """ - Get the states of the dynamics objects. - - Returns - ------- - - np.ndarray - The states. - """ - return self._states.copy() - - def set_states(self, states): - """ - Set the states of the dynamics objects. - - Parameters - ---------- - - states: np.ndarray - The new states. - """ - self._states = states - - def load_replica(self, replica, clock=None): - """ - Make a replica resident in its slot, ready for a dynamics block. - - This sets the slot's lambda value, pushes the replica's state into the - OpenMM context, and points the slot at the replica's clock, energy - trajectory, GCMC water state and ghost file. - - Work that isn't needed is skipped. When there is one slot per replica - and no swap has taken place, the slot already holds everything the - replica needs and this reduces to a handful of comparisons. - - Only touches the replica's own slot, so this is safe to call - concurrently for replicas in different slots. - - Parameters - ---------- - - replica: int - The index of the replica. - - clock: dict - The simulation clock to restore, as returned by - Dynamics._get_clock(). If None, the slot's clock is left alone. - """ - slot = self._replica_slot[replica] - dynamics = self._dynamics[slot] - gcmc_sampler = self._gcmc_samplers[slot] - - # The replica that the slot currently holds. - resident = self._slot_replica[slot] - - # The slot needs new positions and velocities if it is being handed a - # different replica, or if the last mix moved this replica's state. - if resident != replica or self._state_moved[replica]: - self._apply_openmm_state(dynamics.context(), self._openmm_states[replica]) - - # Positions have changed underneath the context, so any cached - # energies are stale. - dynamics.clear_energy_cache() - - self._state_moved[replica] = False - - # Set the lambda value and REST2 scaling factor. This is a no-op if - # the slot is already at this lambda value. - if resident != replica: - dynamics.set_lambda( - self._lambdas[replica], - rest2_scale=self._rest2_scale_factors[replica], - update_constraints=self._update_constraints, - ) - - # Restore the clock and point the slot at this replica's energy - # trajectory, so that energies are accumulated against the replica - # rather than the slot. - if clock is not None: - dynamics._set_clock(clock) - dynamics.set_energy_trajectory(self._energy_trajectories[replica]) - - if gcmc_sampler is not None: - gcmc_sampler.push() - try: - # Swap the water state into the sampler. Diff against what the - # sampler currently holds, which is the state of whichever - # replica was last resident. - target = self._gcmc_states[replica] - if target is not None: - current = gcmc_sampler.water_state() - water_idxs = _np.where(current != target)[0] - - if len(water_idxs) > 0: - gcmc_sampler._set_water_state( - dynamics.context(), - indices=water_idxs, - states=target[water_idxs], - ) - - if resident != replica: - # Update the lambda dependent non-bonded parameters used to - # evaluate insertion and deletion energies, and append ghost - # residues to this replica's file. The sampler keeps its - # statistics per lambda value, so switching also switches - # to this replica's. - gcmc_sampler.set_lambda( - self._lambdas[replica], self._rest2_scale_factors[replica] - ) - gcmc_sampler.set_ghost_file(self._ghost_files[replica]) - finally: - gcmc_sampler.pop() - - self._slot_replica[slot] = replica - - def store_replica(self, replica): - """ - Save the state of a replica back out of its slot, so that the slot can - be handed to another replica. - - Parameters - ---------- - - replica: int - The index of the 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) - - # Count the waters against the slot's context rather than whichever - # one the sampler happens to be bound to, which is unset after the - # sampler has been reset. - gcmc_sampler = self._gcmc_samplers[slot] - gcmc_sampler.push() - try: - self._gcmc_num_waters[replica] = gcmc_sampler.num_waters( - context=self._dynamics[slot].context() - ) - finally: - gcmc_sampler.pop() - - def mix_states(self): - """ - Apply the result of a replica mix. - - The states are permuted here, but not pushed into the OpenMM contexts. - They are applied lazily by load_replica(), which is the only point at - which a slot is known to be free. This does the same amount of work as - applying them eagerly, since load_replica() pushes a state exactly when - the mix moved it. - - The permutation must happen here rather than being resolved lazily - through self._states. A slot is re-used within a cycle, so a replica - may be loaded after another replica has already stored its post-run - state; reading through the indirection at that point would pick up the - new state rather than the pre-mix one. - """ - # Permute the travelling state. This is a reference shuffle, so it is - # cheap even for large systems. Statistics and output files stay with - # the lambda window, so are not permuted. - self._openmm_states = [self._openmm_states[state] for state in self._states] - self._gcmc_states = [self._gcmc_states[state] for state in self._states] - - # Flag the replicas whose state moved, so that load_replica() knows it - # has to push new positions and velocities into the context. The flags - # accumulate, and are only cleared once the state has been pushed, so - # that a replica mixed twice without being loaded in between isn't left - # starting from whatever its context happens to hold. - self._state_moved = [ - moved or bool(state != i) - for i, (state, moved) in enumerate(zip(self._states, self._state_moved)) - ] - - # Update the swap matrix with each configuration's move. - for i, state in enumerate(self._states): - self._num_swaps[state, i] += 1 - - def get_proposed(self): - """ - Return the number of proposed swaps between replicas. - """ - return self._num_proposed - - def get_accepted(self): - """ - Return the number of accepted swaps between replicas. - """ - return self._num_accepted - - def get_swaps(self): - """ - Return the swap matrix. - """ - return self._num_swaps - - def _physical_device(self, device): - """ - Map an OpenMM device index to the physical device backing it. - - Parameters - ---------- - - device: int - The OpenMM device index, which is relative to the visible set. - - Returns - ------- - - int, str - The physical device, i.e. the matching entry from - CUDA_VISIBLE_DEVICES. Falls back to the OpenMM index if the visible - set is unknown, which is correct whenever it starts at zero and is - contiguous. - """ - gpu_devices = getattr(self, "_gpu_devices", None) - - if gpu_devices is None or device >= len(gpu_devices): - return device - - return gpu_devices[device] - - @staticmethod - def _max_contexts_advice(num_contexts): - """ - Advice for a user whose replicas don't fit in GPU memory. - - Parameters - ---------- - - num_contexts: int - The number of OpenMM contexts that are known to fit. - """ - - if num_contexts < 1: - return ( - "A single OpenMM context does not fit on this device. Reduce " - "the size of the system, or use a device with more memory." - ) - - return ( - f"Set 'max_contexts' to {num_contexts} or fewer to re-use contexts " - "across lambda values, which removes the memory limit on the " - "number of replicas. This also requires 'frame_frequency' to equal " - "'checkpoint_frequency'." - ) - - @staticmethod - def _check_device_memory(device=0): - """ - Check the memory usage of the specified GPU device. - - Parameters - ---------- - - device: int, str - The device to query. This is the physical device, i.e. an entry - from CUDA_VISIBLE_DEVICES (or the equivalent for other platforms), - not the index used by OpenMM. OpenMM numbers devices relative to - the visible set, whereas pynvml and pyopencl enumerate all devices - on the machine, so the two only agree when the visible set starts - at zero and is contiguous. CUDA_VISIBLE_DEVICES entries may be - either an index or a UUID. - """ - - device = str(device).strip() - - # A UUID cannot be used to index into the OpenCL device list. - is_uuid = device.startswith("GPU-") or device.startswith("MIG-") - device_index = None if is_uuid else int(device) - - # Try to use pyopencl to detect the GPU vendor. - vendor = None - ocl_device = None - try: - import pyopencl as cl - - if device_index is None: - raise ValueError("Cannot index OpenCL devices by UUID") - - platforms = cl.get_platforms() - all_devices = [] - for platform in platforms: - try: - devices = platform.get_devices(device_type=cl.device_type.GPU) - all_devices.extend(devices) - except Exception: - continue - - if device_index < len(all_devices): - ocl_device = all_devices[device_index] - vendor = ocl_device.vendor - else: - msg = f"Device index {device_index} out of range. Found {len(all_devices)} GPU(s)." - _logger.error(msg) - raise IndexError(msg) - except IndexError: - raise - except Exception: - _logger.warning( - "Could not query GPU platform via OpenCL; falling back to pynvml for NVIDIA detection." - ) - - # NVIDIA: Use pynvml (also used as fallback when OpenCL is unavailable). - if vendor is None or "NVIDIA" in vendor: - try: - import pynvml - - pynvml.nvmlInit() - if is_uuid: - handle = pynvml.nvmlDeviceGetHandleByUUID(device.encode()) - else: - handle = pynvml.nvmlDeviceGetHandleByIndex(device_index) - memory = pynvml.nvmlDeviceGetMemoryInfo(handle) - pynvml.nvmlShutdown() - return (memory.used, memory.free, memory.total) - except Exception as e: - if vendor is None: - msg = f"Could not get GPU memory info for device {device} via OpenCL or pynvml: {e}" - else: - msg = ( - f"Could not get NVIDIA GPU memory info for device {device}: {e}" - ) - _logger.error(msg) - raise RuntimeError(msg) from e - - # AMD: Use OpenCL extension. - elif "AMD" in vendor or "Advanced Micro Devices" in vendor: - try: - import pyopencl as cl - - total = ocl_device.global_mem_size - - # cl_amd_device_attribute_query reports the free memory in - # KBytes, as a list with the whole heap first. - free_memory_info = ocl_device.get_info( - cl.device_info.GLOBAL_FREE_MEMORY_AMD - ) - free_kb = ( - free_memory_info[0] - if isinstance(free_memory_info, (list, tuple)) - else free_memory_info - ) - free = free_kb * 1024 - used = total - free - return (used, free, total) - except Exception as e: - msg = f"Could not get AMD GPU memory info for device {device}: {e}" - _logger.error(msg) - raise RuntimeError(msg) from e - - # Any other vendor: we have no way of querying the free memory. - else: - msg = ( - f"Unable to query the memory of device {device}: " - f"unsupported GPU vendor '{vendor}'." - ) - _logger.error(msg) - raise RuntimeError(msg) + # 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 RepexRunner(_RunnerBase): @@ -1516,8 +303,8 @@ def __init__(self, system, config): constraint_lambda_index=self._constraint_lambda_index, gpu_devices=self._gpu_devices, strided=self._is_strided, + virtual_sites=self._virtual_sites, ) - self._dynamics_cache._virtual_sites = self._virtual_sites else: _logger.debug("Restarting from file") @@ -1534,16 +321,13 @@ def __init__(self, system, config): ) raise e - # Set again, since states saved by older versions don't have it. - self._dynamics_cache._virtual_sites = self._virtual_sites - # Derive the simulation time: prefer the value stored in the - # pickle (_time is set by the new-format _write_checkpoint_system); - # fall back to the Sire system for old-format checkpoints. - if self._dynamics_cache._time is not None and not isinstance( + # pickle (set by the new-format _write_checkpoint_system); fall + # back to the Sire system for old-format checkpoints. + if self._dynamics_cache.time is not None and not isinstance( self._system, list ): - time = self._dynamics_cache._time + time = self._dynamics_cache.time else: time = self._system[0].time() @@ -1579,9 +363,9 @@ def __init__(self, system, config): ) # Make sure the number of replicas is the same. - if len(self._dynamics_cache._lambdas) != self._config.num_lambda: + if self._dynamics_cache.num_replicas != self._config.num_lambda: _logger.error( - f"The number of replicas in the dynamics cache ({len(self._dynamics_cache._lambdas)}) " + f"The number of replicas in the dynamics cache ({self._dynamics_cache.num_replicas}) " f"does not match the number of replicas in the configuration ({self._config.num_lambda})." ) @@ -1590,23 +374,11 @@ def __init__(self, system, config): if not isinstance(self._system, list): self._system.set_time(time) - # The physical device list is not pickled, since the run may be - # restarted against a different set of GPUs. - self._dynamics_cache._gpu_devices = self._gpu_devices - - # Rebuild the slot layout from the current config, so that - # 'max_contexts' can change on restart. Everything that is restored - # is per-replica, so it doesn't depend on the grouping. - self._dynamics_cache._num_slots = self._num_slots - self._dynamics_cache._update_constraints = self._config.update_constraints - 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. - self._dynamics_cache._create_dynamics( + # The settings that aren't stored are taken from the current + # configuration, so that the GPUs and 'max_contexts' can change on + # restart. Everything that is restored is per-replica, so it doesn't + # depend on the slot layout. + self._dynamics_cache.prepare_restart( self._system, self._lambda_values, self._rest2_scale_factors, @@ -1614,20 +386,14 @@ def __init__(self, system, config): self._dynamics_kwargs, gcmc_kwargs=self._gcmc_kwargs, output_directory=self._config.output_directory, + num_slots=self._num_slots, + update_constraints=self._config.update_constraints, + constraint_lambda_index=self._constraint_lambda_index, + gpu_devices=self._gpu_devices, + strided=self._is_strided, + virtual_sites=self._virtual_sites, ) - # The OpenMM contexts are not reset here. Each replica's state is - # pushed into its slot by load_replica() at the start of its first - # block, which is the only point at which the slot is known to be - # free. - - # Restore the sampling statistics. A sampler keeps only the lambda - # values it visits, so each can be handed the whole simulation's. - if self._dynamics_cache._gcmc_stats is not None: - for slot in range(self._dynamics_cache._num_slots): - _, gcmc_sampler = self._dynamics_cache.get(slot) - gcmc_sampler.restore_stats(self._dynamics_cache._gcmc_stats) - # Log the GCMC sphere centre for each replica. This uses the stored # state rather than the context, since a slot only holds the positions # of the replica it last hosted. @@ -1636,9 +402,7 @@ def __init__(self, system, config): for i, lam in enumerate(self._lambda_values): _, gcmc_sampler = self._dynamics_cache.get(self._dynamics_cache.slot_for(i)) if gcmc_sampler is not None and gcmc_sampler._reference is not None: - positions = DynamicsCache._get_positions( - self._dynamics_cache._openmm_states[i] - ) + positions = self._dynamics_cache.stored_positions(i) target = gcmc_sampler._get_target_position( positions.value_in_unit(_omm_unit.angstrom) ) @@ -1654,9 +418,9 @@ def __init__(self, system, config): # If restarting, subtract the time already run from the total runtime if self._config.restart: time = ( - self._dynamics_cache._time + self._dynamics_cache.time if ( - self._dynamics_cache._time is not None + self._dynamics_cache.time is not None and not isinstance(self._system, list) ) else self._system[0].time() @@ -1699,7 +463,7 @@ def __init__(self, system, config): # until here so that _terminal_flip_samplers is always initialised first). if self._is_restart and self._terminal_flip_samplers is not None: for i in range(len(self._lambda_values)): - attempted, accepted = self._dynamics_cache._terminal_flip_stats[i] + attempted, accepted = self._dynamics_cache.terminal_flip_stats[i] self._terminal_flip_samplers[i].reset(attempted, accepted) from threading import Lock @@ -1766,7 +530,7 @@ def _replica_passes(self, cycle, replicas=None): list of list of int The replicas to propagate in each pass. """ - groups = self._dynamics_cache._groups + groups = self._dynamics_cache.groups if replicas is not None: replicas = set(replicas) @@ -1995,15 +759,182 @@ def _run(self, dynamics_executor, checkpoint_executor): checkpoint_executor: concurrent.futures.ThreadPoolExecutor The thread pool for writing checkpoint files. """ - - from math import ceil from time import time - from itertools import repeat - # Record the start time. start = time() + cycles, cycles_per_checkpoint, total_blocks = self._cycle_counts() + + # 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() + + self._prepare_replicas( + dynamics_executor, checkpoint_executor, replica_devices, cycles + ) + + # Current block number. + block = self._start_block + + # Record the start time for the production block. + prod_start = time() + + cycles_per_frame, cycles_per_gcmc, cycles_per_flip = self._move_cycles(cycles) + + # 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()) + + 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 + + self._start_adaptive() + + # 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} " + 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. + # For example: + # States: [ 2 0 1 3 4 5 6 7 8 9 10 ] + # means that the final positions and velocities from replica 2 are + # used to seed replica 0, those from replica 0 are used to seed + # replica 1, and so on. + _logger.info(f"States: {self._dynamics_cache.get_states()}") + + # Whether to checkpoint. + is_checkpoint = (i + 1) in checkpoint_set + + # Whether the checkpoint files are written at the end of this cycle. + do_checkpoint = is_checkpoint or i == cycles - 1 + + # 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 + + if not is_full: + _logger.info(f"Propagating {len(active)} hot replicas: {active}") + + results = self._propagate_cycle( + dynamics_executor, + checkpoint_executor, + cycle=i, + replicas=None if is_full else active, + replica_devices=replica_devices, + last_cycle=last_cycle, + move_cycles=(cycles_per_frame, cycles_per_gcmc, cycles_per_flip), + do_checkpoint=do_checkpoint, + ) + + 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) + + cycle_time = self._start_time + (i + 1) * self._config.energy_frequency + is_burnt_in = cycle_time > self._config.adaptive_burn_in + if is_adaptive: + self._record_cycle( + i, cycle_time, energy_matrix, active, checkpoint_cycles + ) + + # Mix the replicas that were propagated. + _logger.info("Mixing replicas") + self._dynamics_cache.set_states( + _mix_subset( + _np.array(active, dtype=_np.int64), + energy_matrix, + self._dynamics_cache.get_proposed(), + self._dynamics_cache.get_accepted(), + ) + ) + + # This only permutes the stored states. They are pushed into the + # contexts by load_replica() at the start of the next block, which + # is also where the pre-run state for crash recovery is captured. + self._dynamics_cache.mix_states() + + if do_checkpoint: + # Rank the windows before the state is saved, so that a restart + # carries on with the same hot replicas. + if is_adaptive and is_checkpoint and is_burnt_in: + self._rank_windows( + replica_devices, + device_capacity, + self._window_samples(), + cycles - (i + 1), + _cold_rate( + i + 1, + cycles, + checkpoint_cycles, + self._config.adaptive_mix_interval, + ), + ) + + self._write_cycle_checkpoint( + checkpoint_executor, + block, + total_blocks, + i == cycles - 1, + num_samples, + cycle_time, + is_checkpoint, + ) + num_samples = [0] * num_replicas + + # This is a checkpoint cycle. + if is_checkpoint: + # Update the block number. + block += 1 + + self._finish(start, prod_start) + + def _cycle_counts(self): + """ + Work out the number of replica exchange cycles in the run, between + checkpoints, and the number of checkpoint blocks. + + Returns + ------- + + cycles: int + The number of cycles. + + cycles_per_checkpoint: float + The number of cycles between checkpoints, which may be fractional. + + total_blocks: int + The number of checkpoint blocks, including a final partial block. + """ + from math import ceil + # Work out the number of repex cycles. cycles = (self._config.runtime / self._config.energy_frequency).value() @@ -2035,7 +966,6 @@ def _run(self, dynamics_executor, checkpoint_executor): # Handle the case where the checkpoint frequency is less than the energy frequency. if frac < 1.0: frac = 1.0 - checkpoint_frequency = self._config.energy_frequency # Store the number of repex cycles per block (may be fractional). cycles_per_checkpoint = frac @@ -2046,13 +976,32 @@ def _run(self, dynamics_executor, checkpoint_executor): num_blocks = 1 rem = 0 - # Create the replica list. - replica_list = list(range(self._config.num_lambda)) + return cycles, cycles_per_checkpoint, num_blocks + int(rem > 0) - # 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() + def _prepare_replicas( + self, dynamics_executor, checkpoint_executor, replica_devices, cycles + ): + """ + Minimise and equilibrate the replicas, unless this is a restart, and + write a checkpoint after equilibration so that a restart after an + early production crash doesn't need to re-equilibrate. + + Parameters + ---------- + + dynamics_executor: concurrent.futures.ThreadPoolExecutor + The thread pool for dynamics. + + checkpoint_executor: concurrent.futures.ThreadPoolExecutor + The thread pool for writing checkpoint files. + + replica_devices: list of int + The GPU that propagates each replica. + + cycles: int + The number of cycles in the run. + """ + from itertools import repeat # Minimise at each lambda value. Don't minimise on restart. if self._config.minimise and not self._is_restart: @@ -2113,18 +1062,30 @@ def _run(self, dynamics_executor, checkpoint_executor): checkpoint_executor, "Post-equilibration checkpoint", self._checkpoint, - replica_list, + list(range(self._config.num_lambda)), repeat(self._lambda_values), repeat(-1), repeat(cycles), ) - # Current block number. - block = self._start_block + def _move_cycles(self, cycles): + """ + Work out the number of cycles between frames, GCMC moves and terminal + flip moves. - # Record the start time for the production block. - prod_start = time() + Parameters + ---------- + + cycles: int + The number of cycles in the run. + + Returns + ------- + cycles_per_frame, cycles_per_gcmc, cycles_per_flip: int + The number of cycles between each. A move that isn't used is never + due. + """ # Store the number of blocks per-frame. For GCMC, we need to write the # indices of the current ghost water residues each time a frame is saved. # For GCMC simulations, the frame frequency is guaranteed to be a multiple @@ -2158,274 +1119,301 @@ def _run(self, dynamics_executor, checkpoint_executor): else: cycles_per_flip = cycles + 1 - # 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 + return cycles_per_frame, cycles_per_gcmc, cycles_per_flip + def _start_adaptive(self): + """ + Set up the ranking of windows for adaptive replica exchange, restoring + the data from a restart, or make every replica hot if it isn't used. + """ 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() + if self._config.adaptive_repex: + self._ranker = _WindowRanker( + num_replicas, + self._config.adaptive_min_samples, + self._dynamics_cache.get_proposed(), + self._dynamics_cache.get_accepted(), + is_ring_break=self._ring_pair is not None, + is_cached=self._is_cached, + ) # A burn-in raised on restart starts the adaptation again. if self._start_time >= self._config.adaptive_burn_in: - self._restore_adaptive_data() + self._ranker.restore(self._dynamics_cache.adaptive_data) else: - self._dynamics_cache._hot_replicas = None - self._dynamics_cache._cold_votes = [0] * num_replicas + self._dynamics_cache.reset_hot_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 + self._dynamics_cache.reset_hot_replicas() - # Converted once, since unit conversions are slow. - start_ns = self._start_time.to("ns") - energy_frequency_ns = self._config.energy_frequency.to("ns") + def _propagate_cycle( + self, + dynamics_executor, + checkpoint_executor, + cycle, + replicas, + replica_devices, + last_cycle, + move_cycles, + do_checkpoint, + ): + """ + Propagate the replicas for a cycle, 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. - # Perform the replica exchange simulation. - for i in range(cycles): - _logger.info( - f"Running dynamics for cycle {i + 1} of {cycles} " - f"(time = {start_ns + i * energy_frequency_ns:.4f} ns)" - ) + Parameters + ---------- + + dynamics_executor: concurrent.futures.ThreadPoolExecutor + The thread pool for dynamics. + + checkpoint_executor: concurrent.futures.ThreadPoolExecutor + The thread pool for writing checkpoint files. + + cycle: int + The index of the cycle. + + replicas: list of int + The replicas to propagate. If None, every replica is propagated. + + replica_devices: list of int + The GPU that propagates each replica. + + last_cycle: list of int + The cycle on which each replica was last propagated, or -1. + + move_cycles: tuple of int + The number of cycles between frames, GCMC moves and terminal flip + moves. + + do_checkpoint: bool + Whether the checkpoint files are written at the end of the cycle. + + Returns + ------- + + list of tuple + The index and energies of each replica that was propagated. + """ + from itertools import repeat + + cycles_per_frame, cycles_per_gcmc, cycles_per_flip = move_cycles + is_adaptive = self._config.adaptive_repex + results = [] + + # 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() + + for batch in self._replica_passes(cycle, replicas=replicas): + # 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 batch_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], cycle, cycles_per_gcmc) + for r in batch_replicas + ] + is_terminal_flip = [ + self._is_due(last_cycle[r], cycle, cycles_per_flip) + for r in batch_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], cycle, cycles_per_frame) + for r in batch_replicas + ] + + try: + for result, index, energies in dynamics_executor.map( + self._run_block, + batch_replicas, + repeat(self._lambda_values), + 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( + f"Dynamics failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {energies}" + ) + raise energies + results.append((index, energies)) + except KeyboardInterrupt: + dynamics_executor.shutdown(wait=False, cancel_futures=True) + _logger.error("Dynamics cancelled. Exiting.") + _sys.exit(1) + + # When contexts are shared, commit the replicas of this pass while + # they are still resident in their slots. The committed systems are + # held until the end of the cycle, so that every checkpoint file is + # still written under a single lock. With a context per replica + # everything is still resident once the cycle finishes, so + # committing is left to _checkpoint(). + if do_checkpoint and self._is_cached: + self._map_replicas( + checkpoint_executor, "Commit", self._commit_replica, batch + ) + + return results + + def _record_cycle( + self, cycle, cycle_time, energy_matrix, replicas, checkpoint_cycles + ): + """ + Record the data used to rank the windows for adaptive replica exchange, + logging when the burn-in is complete and counting swaps from then on. + Must be called before the replicas are mixed. + + Parameters + ---------- - # Log the states. This is the replica index for the state (positions - # and velocities) used to seed each replica for the current cycle. - # For example: - # States: [ 2 0 1 3 4 5 6 7 8 9 10 ] - # means that the final positions and velocities from replica 2 are - # used to seed replica 0, those from replica 0 are used to seed - # replica 1, and so on. - _logger.info(f"States: {self._dynamics_cache.get_states()}") + cycle: int + The index of the cycle. - # Clear the results list. - results = [] + cycle_time: sire.units.GeneralUnit + The simulation time at the end of the cycle. - # Whether to checkpoint. - is_checkpoint = (i + 1) in checkpoint_set + energy_matrix: np.ndarray + The reduced energy of each replica at each window. - # Whether the checkpoint files are written at the end of this cycle. - do_checkpoint = is_checkpoint or i == cycles - 1 + replicas: list of int + The replicas that were propagated. - # 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 + checkpoint_cycles: np.ndarray + The checkpoint cycles, numbered from one, in increasing order. + """ + # Nothing is recorded until the burn-in is complete. + if cycle_time <= self._config.adaptive_burn_in: + return + + if cycle_time - self._config.energy_frequency <= self._config.adaptive_burn_in: + ranked = checkpoint_cycles[checkpoint_cycles >= cycle + 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._ranker.set_swap_baseline( + self._dynamics_cache.get_proposed(), + self._dynamics_cache.get_accepted(), ) - active = list(range(num_replicas)) if is_full else hot + self._ranker.record(energy_matrix, replicas, self._bond_distances(replicas)) - if not is_full: - _logger.info(f"Propagating {len(active)} hot replicas: {active}") + def _write_cycle_checkpoint( + self, + checkpoint_executor, + block, + num_blocks, + is_final_block, + num_samples, + cycle_time, + is_checkpoint, + ): + """ + Write the checkpoint files once a cycle is complete, every one under a + single lock, so that an external process reading the output directory + always sees a coherent set rather than a mixture of new and old files. - # 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, 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 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), - 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( - f"Dynamics failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {energies}" - ) - raise energies - results.append((index, energies)) - except KeyboardInterrupt: - dynamics_executor.shutdown(wait=False, cancel_futures=True) - _logger.error("Dynamics cancelled. Exiting.") - _sys.exit(1) - - # When contexts are shared, commit the replicas of this pass - # while they are still resident in their slots. The committed - # systems are held until the end of the cycle, so that every - # checkpoint file is still written under a single lock. With a - # context per replica everything is still resident once the - # cycle finishes, so committing is left to _checkpoint(). - if do_checkpoint and self._is_cached: - self._map_replicas( - checkpoint_executor, "Commit", self._commit_replica, batch - ) + Parameters + ---------- - for index in active: - last_cycle[index] = i - num_samples[index] += 1 + checkpoint_executor: concurrent.futures.ThreadPoolExecutor + The thread pool for writing checkpoint files. - # Assemble an energy matrix from the results. - _logger.info("Assembling energy matrix") - energy_matrix = self._assemble_results(results) + block: int + The index of the checkpoint block. - # 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) + num_blocks: int + The number of checkpoint blocks in the run. - # Mix the replicas that were propagated. - _logger.info("Mixing replicas") - self._dynamics_cache.set_states( - _mix_subset( - _np.array(active, dtype=_np.int64), - energy_matrix, - self._dynamics_cache.get_proposed(), - self._dynamics_cache.get_accepted(), - ) - ) + is_final_block: bool + Whether this is the last cycle of the run. - # This only permutes the stored states. They are pushed into the - # contexts by load_replica() at the start of the next block, which - # is also where the pre-run state for crash recovery is captured. - self._dynamics_cache.mix_states() + num_samples: list of int + The number of energy samples each replica has collected since the + last checkpoint. - # Checkpoint. This happens once the whole cycle is complete, with - # every checkpoint file written under a single lock, so that an - # 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, - ) + cycle_time: sire.units.GeneralUnit + The simulation time at the end of the cycle. - # Create the lock. - lock = _FileLock(self._lock_file) + is_checkpoint: bool + Whether this is a checkpoint cycle, rather than a final save. + """ + from itertools import repeat - # Acquire the file lock to ensure that the checkpoint files are - # in a consistent state if read by another process. - with lock.acquire(timeout=self._config.timeout.to("seconds")): - # First backup existing checkpoint files. - self._map_replicas( - checkpoint_executor, - "Backup", - self._backup_checkpoint, - replica_list, - ) + replica_list = list(range(self._config.num_lambda)) - # Now write the new checkpoint files. - 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, - ) + lock = _FileLock(self._lock_file) + with lock.acquire(timeout=self._config.timeout.to("seconds")): + # First backup existing checkpoint files. + self._map_replicas( + checkpoint_executor, + "Backup", + self._backup_checkpoint, + replica_list, + ) - num_samples = [0] * num_replicas + # Now write the new checkpoint files. + self._map_replicas( + checkpoint_executor, + "Checkpoint", + self._checkpoint, + replica_list, + repeat(self._lambda_values), + repeat(block), + repeat(num_blocks), + repeat(is_final_block), + num_samples, + ) - # Replicas' own clocks fall behind while they are cold, so - # the simulation time is that of the cycle. - self._dynamics_cache._time = cycle_time + # 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) + self._save_repex_state(final=not is_checkpoint) - # This is a checkpoint cycle. - if is_checkpoint: - # Update the block number. - block += 1 + def _finish(self, start, prod_start): + """ + Log the performance and run time, remove the backup files, and mark + the run as complete. - # Record the end time for the production block. - prod_end = time() + Parameters + ---------- - # Record the end time. - end = time() + start: float + The time at which the run started, in seconds. - # Work how many fractional days the production block took. - prod_time = (prod_end - prod_start) / 86400 + prod_start: float + The time at which production started, in seconds. + """ + from time import time - # Record the average production speed. (ns/day per replica) - prod_speed = self._config.runtime.to("ns") / prod_time + end = time() - # Record the average production speed. + # Record the average production speed. (ns/day per replica) + prod_speed = self._config.runtime.to("ns") / ((end - prod_start) / 86400) _logger.info(f"Overall performance: {prod_speed:.2f} ns day-1") # Log the run time in minutes. @@ -2877,59 +1865,6 @@ def _equilibrate(self, index): return True, index, None - def _compute_energies(self, index): - """ - Compute the energies for a given replica by updating the OpenMM state - within the context and re-evaluating the potential energy. - - Energies are currently computed internally by Sire at the end of each - dynamics block, but this approach incurs an overhead due to the cost of - updating the force field parameters within the context when changing - lambda. This alternaitve method is left here for performance testing. - - Parameters - ---------- - - index: int - The index of the replica. - - Returns - ------- - - index: int - The index of the replica. - - energies: np.ndarray - The energies of the replica and each state. - """ - _logger.info( - f"Computing energies for {_lam_sym} = {self._lambda_values[index]:.5f}" - ) - - # Get the dynamics object. - dynamics, _ = self._dynamics_cache.get(self._dynamics_cache.slot_for(index)) - - # Create an array to hold the energies. - energies = _np.zeros(self._config.num_lambda) - - # Loop over the states. - for i in range(self._config.num_lambda): - # Set the state. - DynamicsCache._apply_openmm_state( - dynamics.context(), self._dynamics_cache._openmm_states[i] - ) - dynamics._d._clear_state() - - # Compute and store the energy for this state. - energies[i] = dynamics.current_potential_energy().value() - - # Reset the state. - DynamicsCache._apply_openmm_state( - dynamics.context(), self._dynamics_cache._openmm_states[index] - ) - - return index, energies - def _assemble_results(self, results): """ Assemble the results into a matrix. @@ -2983,52 +1918,6 @@ def _checkpoint_cycles(num_cycles, 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. @@ -3134,13 +2023,17 @@ def _window_samples(self): The number of samples for each window. """ energy_frequency = self._config.energy_frequency.to("ps") + clocks = [ + self._dynamics_cache.get_clock(replica) + for replica in range(self._dynamics_cache.num_replicas) + ] return [ ( 0 if clock is None else round(clock["_current_time"].to("ps") / energy_frequency) ) - for clock in self._dynamics_cache._clocks + for clock in clocks ] def _check_adaptive_config(self): @@ -3255,225 +2148,46 @@ def _is_due(last_cycle, cycle, cycles_per_move): """ 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): + def _bond_distances(self, 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. + The distance of the bond that is broken or made in each replica, for + a ring-breaking perturbation. 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): + dict + The distance in Angstrom for each replica, or None if there is no + single bond that is broken or made. """ - 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. + if self._ring_pair is None: + return None - Returns - ------- + import openmm.unit as _omm_unit - 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)) - ) + i, j = self._ring_pair + distances = {} + for k in replicas: + positions = self._dynamics_cache.stored_positions(k) + pair = _np.asarray(positions[[i, j]].value_in_unit(_omm_unit.angstrom)) + distances[k] = float(_np.linalg.norm(pair[1] - pair[0])) - drift = abs(estimates[0] - estimates[1]) - return drift > 0.25 and drift > 2.0 * _np.hypot(*errors) + return distances 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. + Rank the windows and choose the hot replicas for adaptive replica + exchange, which are kept in the dynamics cache so that a restart + carries on with them. See WindowRanker.rank(). Parameters ---------- @@ -3494,194 +2208,19 @@ def _rank_windows( 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" + hot, votes = self._ranker.rank( + window_devices, + device_capacity, + window_samples, + cycles_left, + cold_rate, + self._dynamics_cache.get_proposed(), + self._dynamics_cache.get_accepted(), + self._dynamics_cache.hot_replicas, + self._dynamics_cache.cold_votes, ) + self._dynamics_cache.hot_replicas = hot + self._dynamics_cache.cold_votes = votes def _check_restart(self): """ @@ -3715,7 +2254,7 @@ def _write_checkpoint_system(self, system, index, context=None, gcmc_sampler=Non velocities are already stored as compact numpy arrays in the OpenMM state dict. """ - self._dynamics_cache._time = system.time() + self._dynamics_cache.time = system.time() def _load_and_commit(self, index): """ @@ -3891,7 +2430,7 @@ def _checkpoint( # the water count and the statistics are read from what was # recorded while this replica was resident in its slot, since by # the time the cycle finishes the slot may hold another one. - num_waters = self._dynamics_cache._gcmc_num_waters[index] + num_waters = self._dynamics_cache.gcmc_num_waters(index) if gcmc_sampler is not None and num_waters is not None: stats = gcmc_sampler.get_stats().get(gcmc_sampler.stats_key(lam)) n_moves = stats["num_moves"] if stats is not None else 0 @@ -3979,7 +2518,7 @@ def _merge_gcmc_stats(self): """ stats = {} - for slot in range(self._dynamics_cache._num_slots): + for slot in range(self._dynamics_cache.num_slots): _, gcmc_sampler = self._dynamics_cache.get(slot) if gcmc_sampler is not None: stats.update(gcmc_sampler.get_stats()) @@ -4012,9 +2551,9 @@ 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() + self._dynamics_cache.adaptive_data = self._ranker.to_dict() else: - self._dynamics_cache._adaptive_data = None + self._dynamics_cache.adaptive_data = None with open(self._repex_state, "wb") as f: _pickle.dump(self._dynamics_cache, f) @@ -4023,10 +2562,10 @@ def _save_sampler_stats(self): Save GCMC and terminal flip sampler statistics to the dynamics cache prior to pickling. """ - self._dynamics_cache._gcmc_stats = self._merge_gcmc_stats() + self._dynamics_cache.gcmc_stats = self._merge_gcmc_stats() if self._terminal_flip_samplers is not None: - self._dynamics_cache._terminal_flip_stats = [ + self._dynamics_cache.terminal_flip_stats = [ [s.num_attempted, s.num_accepted] for s in self._terminal_flip_samplers ] diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index def9492..03ccf9c 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -21,6 +21,7 @@ __all__ = ["Runner"] +from dataclasses import dataclass as _dataclass from filelock import FileLock as _FileLock from time import time as _timer @@ -35,6 +36,31 @@ from ._base import RunnerBase as _RunnerBase +@_dataclass +class _Window: + """ + The state of a lambda window that is shared by the stages of its + simulation. + """ + + index: int + lambda_value: float + rest2_scale: float + device: int = None + gcmc_sampler: object = None + terminal_flip_sampler: object = None + dynamics: object = None + lambda_energy: list = None + lambda_grad: list = None + lambda_array: list = None + rest2_scale_factors: list = None + num_energy_neighbours: int = None + # The integration step reached, counted from the start of the simulation. + step: int = 0 + # The number of energy samples recorded when the last checkpoint was written. + num_energy_samples: int = 0 + + class Runner(_RunnerBase): """ Standard simulation runner class. (Uncoupled simulations.) @@ -84,6 +110,28 @@ def __init__(self, system, config): _logger.error(msg) raise ValueError(msg) + # Production is run in whole integration steps. This is checked before + # the base class writes the config file, so a rejected restart doesn't + # leave one behind. + from math import isclose + + for option in ( + "runtime", + "checkpoint_frequency", + "energy_frequency", + "frame_frequency", + "gcmc_frequency", + "terminal_flip_frequency", + ): + value = getattr(config, option) + if value is None: + continue + ratio = (value / config.timestep).value() + if not isclose(ratio, round(ratio), abs_tol=1e-4): + msg = f"'{option}' must be a multiple of 'timestep'." + _logger.error(msg) + raise ValueError(msg) + # Call the base class constructor. super().__init__(system, config) @@ -476,56 +524,152 @@ def _run( The duration of the simulation. """ - # Get the lambda value. lambda_value = self._lambda_values[index] + window = _Window( + index=index, + lambda_value=lambda_value, + rest2_scale=self._rest2_scale_factors[ + self._lambda_energy.index(lambda_value) + ], + device=device, + ) - # Get the index in the lambda_energy array. - nrg_index = self._lambda_energy.index(lambda_value) + time = self._remaining_time(system, window, is_restart) + if time is None: + return _sr.u("0ps") - # Get the REST2 scaling factor. - rest2_scale = self._rest2_scale_factors[nrg_index] + system = self._create_samplers(system, window) - # Check for completion if this is a restart. - if is_restart: - if isinstance(self._system, list): - time = system.time() - else: - # New format: time stored in .npz, not in the Sire system. - time = _sr.u( - f"{float(_np.load(self._filenames[index]['checkpoint_state'])['time_ps'].item()):.6f} ps" - ) - system.set_time(time) - if time > self._config.runtime - self._config.timestep: - _logger.success( - f"{_lam_sym} = {lambda_value} already complete. Skipping." - ) - return _sr.u("0ps") + # Don't minimise on restart. + if self._config.minimise and not is_restart: + system = self._minimise_window(system, window) - # Work out the current block number. - if self._config.checkpoint_frequency.value() > 0.0: - self._start_block = int( - round(time.value() / self._config.checkpoint_frequency.value(), 12) - ) - else: - self._start_block = 0 + is_equilibrated = ( + not is_restart and self._config.equilibration_time.value() > 0.0 + ) + if is_equilibrated: + system = self._equilibrate_window(system, window) + + self._set_energy_lambdas(window) + self._create_production_dynamics(system, window, is_restart, is_equilibrated) + + # Write a checkpoint immediately after equilibration so that a restart + # after an early production crash doesn't need to re-equilibrate. + if is_equilibrated: + self._write_post_equilibration_checkpoint(system, window) + + window.step = self._num_steps(self._config.runtime - time) + window.num_energy_samples = self._num_energy_samples(window) + + if self._config.checkpoint_frequency.value() > 0.0: + self._run_checkpoint_blocks(window, time) + else: + self._run_without_checkpoints(window, time) + + return time + + def _remaining_time(self, system, window, is_restart): + """ + Work out the simulation time remaining for a window, and the block + that it starts from. + + Parameters + ---------- + + system: :class: `System ` + The system to be simulated. + + window: _Window + The window being simulated. + + is_restart: bool + Whether this is a restart simulation. + + Returns + ------- - # Subtract the current time from the runtime. - time = self._config.runtime - time + time: sire.units.GeneralUnit + The time remaining, or None if the window is already complete. + """ + if not is_restart: + self._start_block = 0 + return self._config.runtime + + if isinstance(self._system, list): + time = system.time() + else: + # New format: time stored in .npz, not in the Sire system. + time = _sr.u( + f"{float(_np.load(self._filenames[window.index]['checkpoint_state'])['time_ps'].item()):.6f} ps" + ) + system.set_time(time) + if time > self._config.runtime - self._config.timestep: + _logger.success( + f"{_lam_sym} = {window.lambda_value} already complete. Skipping." + ) + return None + + # Work out the current block number. + if self._config.checkpoint_frequency.value() > 0.0: + self._start_block = int( + round(time.value() / self._config.checkpoint_frequency.value(), 12) + ) else: self._start_block = 0 - time = self._config.runtime - - def generate_lam_vals(lambda_base, increment=0.001): - """Generate lambda values for a given lambda_base and increment""" - if lambda_base + increment > 1.0 and lambda_base - increment < 0.0: - raise ValueError("Increment too large") - if lambda_base + increment > 1.0: - lam_vals = [lambda_base - increment] - elif lambda_base - increment < 0.0: - lam_vals = [lambda_base + increment] - else: - lam_vals = [lambda_base - increment, lambda_base + increment] - return lam_vals + + return self._config.runtime - time + + @staticmethod + def _gradient_lambdas(lambda_base, increment=0.001): + """ + The lambda values either side of a window that are used for + finite-difference gradient analysis. + + Parameters + ---------- + + lambda_base: float + The lambda value of the window. + + increment: float + The distance to each neighbouring lambda value. + + Returns + ------- + + list of float + The lambda values. + """ + if lambda_base + increment > 1.0 and lambda_base - increment < 0.0: + raise ValueError("Increment too large") + if lambda_base + increment > 1.0: + lam_vals = [lambda_base - increment] + elif lambda_base - increment < 0.0: + lam_vals = [lambda_base + increment] + else: + lam_vals = [lambda_base - increment, lambda_base + increment] + return lam_vals + + def _create_samplers(self, system, window): + """ + Create the GCMC and terminal flip samplers for a window, if required. + + Parameters + ---------- + + system: :class: `System ` + The system to be simulated. + + window: _Window + The window being simulated, which holds the samplers. + + Returns + ------- + + system: :class: `System ` + The system to simulate, which includes the GCMC waters. + """ + lambda_value = window.lambda_value # Prepare the GCMC sampler. if self._config.gcmc: @@ -537,26 +681,23 @@ def generate_lam_vals(lambda_base, increment=0.001): msg = "loch is not installed. GCMC sampling cannot be performed." _logger.error(msg) - gcmc_sampler = GCMCSampler( + window.gcmc_sampler = GCMCSampler( system, - device=int(device), + device=int(window.device), lambda_value=lambda_value, - rest2_scale=rest2_scale, - ghost_file=self._filenames[index]["gcmc_ghosts"], + rest2_scale=window.rest2_scale, + ghost_file=self._filenames[window.index]["gcmc_ghosts"], **self._gcmc_kwargs, ) # Get the GCMC system. - system = gcmc_sampler.system() - - else: - gcmc_sampler = None + system = window.gcmc_sampler.system() # Create the terminal flip sampler (if terminal groups were detected). if self._terminal_groups: from ._samplers import TerminalFlipSampler - terminal_flip_sampler = TerminalFlipSampler( + window.terminal_flip_sampler = TerminalFlipSampler( self._terminal_groups, float(self._config.temperature.value()), ) @@ -564,159 +705,233 @@ def generate_lam_vals(lambda_base, increment=0.001): f"Terminal flip sampler ready at {_lam_sym} = {lambda_value:.5f} " f"(every {self._config.terminal_flip_frequency})" ) - else: - terminal_flip_sampler = None - # Minimisation. Don't minimise on restart. - if self._config.minimise and not is_restart: - constraint = self._config.constraint - perturbable_constraint = self._config.perturbable_constraint + return system + + def _minimise_window(self, system, window): + """ + Minimise a window before it is simulated. - # Don't use constraints during minimisation. - if not self._config.minimisation_constraints: + Parameters + ---------- + + system: :class: `System ` + The system to minimise. + + window: _Window + The window being simulated. + + Returns + ------- + + system: :class: `System ` + The minimised system, or the original if minimisation failed and + errors are not raised. + """ + constraint = self._config.constraint + perturbable_constraint = self._config.perturbable_constraint + + # Don't use constraints during minimisation. + if not self._config.minimisation_constraints: + constraint = "none" + perturbable_constraint = "none" + # We will be performing an equilibration stage. + elif self._config.equilibration_time.value() > 0.0: + # Don't use constraints during equilibration. + if not self._config.equilibration_constraints: constraint = "none" perturbable_constraint = "none" - # We will be performing an equilibration stage. - elif not is_restart and self._config.equilibration_time.value() > 0.0: - # Don't use constraints during equilibration. - if not self._config.equilibration_constraints: - constraint = "none" - perturbable_constraint = "none" - try: - system = self._minimisation( - system, - lambda_value, - rest2_scale=rest2_scale, - device=device, - constraint=constraint, - perturbable_constraint=perturbable_constraint, - gcmc_sampler=gcmc_sampler, - ) - except Exception as e: - msg = f"Minimisation failed for {_lam_sym} = {lambda_value:.5f}: {e}" - if self._config.minimisation_errors: - raise RuntimeError(msg) - else: - _logger.warning(msg) + try: + system = self._minimisation( + system, + window.lambda_value, + rest2_scale=window.rest2_scale, + device=window.device, + constraint=constraint, + perturbable_constraint=perturbable_constraint, + gcmc_sampler=window.gcmc_sampler, + ) + except Exception as e: + msg = f"Minimisation failed for {_lam_sym} = {window.lambda_value:.5f}: {e}" + if self._config.minimisation_errors: + raise RuntimeError(msg) + else: + _logger.warning(msg) - # Equilibration. - is_equilibrated = False - if not is_restart and self._config.equilibration_time.value() > 0.0: - is_equilibrated = True - try: - # Run without saving energies or frames. - _logger.info(f"Equilibrating at {_lam_sym} = {lambda_value:.5f}") - - # Copy the dynamics kwargs. - dynamics_kwargs = self._dynamics_kwargs.copy() - - # Overload the dynamics kwargs with the simulation options. - dynamics_kwargs.update( - { - "device": device, - "lambda_value": lambda_value, - "rest2_scale": rest2_scale, - "timestep": self._config.equilibration_timestep, - "constraint": ( - "none" - if not self._config.equilibration_constraints - else self._config.constraint - ), - "perturbable_constraint": ( - "none" - if not self._config.equilibration_constraints - else self._config.perturbable_constraint - ), - } - ) + return system - # Create the dynamics object. - dynamics = system.dynamics(**dynamics_kwargs) + def _equilibrate_window(self, system, window): + """ + Equilibrate a window without saving energies or frames. - # Write the OpenMM XML file to the output directory. - if self._config.save_xml and not is_restart: - _logger.info( - f"Writing OpenMM XML for {_lam_sym} = {lambda_value:.5f}" - ) - dynamics.to_xml(self._filenames[index]["xml"]) + Parameters + ---------- - # Equilibrate with GCMC moves. - if gcmc_sampler is not None: - # Bind the GCMC sampler to the dynamics object. - gcmc_sampler.bind_dynamics(dynamics) + system: :class: `System ` + The system to equilibrate. - # This context is built from the system as it was committed - # by minimisation, so it doesn't carry the water state that - # the moves performed there left behind. - gcmc_sampler.push() - try: - gcmc_sampler._set_water_state(dynamics.context(), force=True) - finally: - gcmc_sampler.pop() + window: _Window + The window being simulated. - _logger.info( - f"Equilibrating with GCMC moves at {_lam_sym} = {lambda_value:.5f}" - ) + Returns + ------- - gcmc_sampler.push() - try: - for i in range(100): - gcmc_sampler.move(dynamics.context()) - finally: - gcmc_sampler.pop() + system: :class: `System ` + The equilibrated system, with its time reset. + """ + lambda_value = window.lambda_value + gcmc_sampler = window.gcmc_sampler + + try: + _logger.info(f"Equilibrating at {_lam_sym} = {lambda_value:.5f}") + + # Copy the dynamics kwargs. + dynamics_kwargs = self._dynamics_kwargs.copy() + + # Overload the dynamics kwargs with the simulation options. + dynamics_kwargs.update( + { + "device": window.device, + "lambda_value": lambda_value, + "rest2_scale": window.rest2_scale, + "timestep": self._config.equilibration_timestep, + "constraint": ( + "none" + if not self._config.equilibration_constraints + else self._config.constraint + ), + "perturbable_constraint": ( + "none" + if not self._config.equilibration_constraints + else self._config.perturbable_constraint + ), + } + ) - # Run without saving energies or frames. - dynamics.run( - self._config.equilibration_time, - energy_frequency=0, - frame_frequency=0, - save_velocities=False, - auto_fix_minimise=self._config.auto_fix_minimise, - save_crash_report=self._config.save_crash_report, + # Create the dynamics object. + dynamics = system.dynamics(**dynamics_kwargs) + + # Write the OpenMM XML file to the output directory. + if self._config.save_xml: + _logger.info(f"Writing OpenMM XML for {_lam_sym} = {lambda_value:.5f}") + dynamics.to_xml(self._filenames[window.index]["xml"]) + + # Equilibrate with GCMC moves. + if gcmc_sampler is not None: + # Bind the GCMC sampler to the dynamics object. + gcmc_sampler.bind_dynamics(dynamics) + + # This context is built from the system as it was committed by + # minimisation, so it doesn't carry the water state that the + # moves performed there left behind. + gcmc_sampler.push() + try: + gcmc_sampler._set_water_state(dynamics.context(), force=True) + finally: + gcmc_sampler.pop() + + _logger.info( + f"Equilibrating with GCMC moves at {_lam_sym} = {lambda_value:.5f}" ) - # Commit the system. - system = dynamics.commit() + gcmc_sampler.push() + try: + for i in range(100): + gcmc_sampler.move(dynamics.context()) + finally: + gcmc_sampler.pop() - # Reset the timer. - if self._initial_time[index].value() != 0: - system.set_time(self._initial_time[index]) - else: - system.set_time(_sr.u("0ps")) + # Run without saving energies or frames. + dynamics.run( + self._config.equilibration_time, + energy_frequency=0, + frame_frequency=0, + save_velocities=False, + auto_fix_minimise=self._config.auto_fix_minimise, + save_crash_report=self._config.save_crash_report, + ) - except Exception as e: - raise RuntimeError(f"Equilibration failed: {e}") + # Commit the system. + system = dynamics.commit() - # Work out the lambda values for finite-difference gradient analysis. - lambda_grad = generate_lam_vals(lambda_value) + # Reset the timer. + if self._initial_time[window.index].value() != 0: + system.set_time(self._initial_time[window.index]) + else: + system.set_time(_sr.u("0ps")) + + except Exception as e: + raise RuntimeError(f"Equilibration failed: {e}") + + return system + + def _set_energy_lambdas(self, window): + """ + Work out the lambda values, and their REST2 scaling factors, at which + energies are sampled for a window. + + Parameters + ---------- - # Create the array of lambda values for energy sampling. - lambda_energy = self._lambda_energy.copy() + window: _Window + The window being simulated, which holds the lambda values. + """ + # Work out the lambda values for finite-difference gradient analysis. + window.lambda_grad = self._gradient_lambdas(window.lambda_value) - # Sort the lambda values. - lambda_energy = sorted(lambda_energy) + # Create the sorted array of lambda values for energy sampling. + window.lambda_energy = sorted(self._lambda_energy.copy()) # Create the lambda array. - lambda_array = lambda_energy.copy() + lambda_array = window.lambda_energy.copy() # Add the lambda values for finite-difference gradient analysis. - lambda_array.extend(lambda_grad) + lambda_array.extend(window.lambda_grad) # Add additional REST2 scaling factors. rest2_scale_factors = self._rest2_scale_factors.copy() - rest2_scale_factors.extend([rest2_scale] * len(lambda_grad)) + rest2_scale_factors.extend([window.rest2_scale] * len(window.lambda_grad)) # Get the indices of the lambda values in sorted order. sorted_indices = [ i for i, _ in sorted(enumerate(lambda_array), key=lambda x: x[1]) ] - # Sort the lambda values. - lambda_array = sorted(lambda_array) + # Sort the lambda values, and the scaling factors to match. + window.lambda_array = sorted(lambda_array) + window.rest2_scale_factors = [rest2_scale_factors[i] for i in sorted_indices] + + # Set the number of neighbours used for the energy calculation. If not + # None, then we add one to account for the extra windows used for + # finite-difference gradient analysis. + if self._config.num_energy_neighbours is not None: + window.num_energy_neighbours = self._config.num_energy_neighbours + 1 + + def _create_production_dynamics(self, system, window, is_restart, is_equilibrated): + """ + Create the dynamics object for production, restoring the state of a + restart and the state of the samplers. + + Parameters + ---------- + + system: :class: `System ` + The system to be simulated. + + window: _Window + The window being simulated, which holds the dynamics object. + + is_restart: bool + Whether this is a restart simulation. - # Now sort the scaling factors. - rest2_scale_factors = [rest2_scale_factors[i] for i in sorted_indices] + is_equilibrated: bool + Whether the window was equilibrated, in which case the OpenMM XML + file has already been written. + """ + index = window.index + lambda_value = window.lambda_value + gcmc_sampler = window.gcmc_sampler _logger.info(f"Running dynamics at {_lam_sym} = {lambda_value:.5f}") @@ -726,14 +941,15 @@ def generate_lam_vals(lambda_base, increment=0.001): # Overload the dynamics kwargs with the simulation options. dynamics_kwargs.update( { - "device": device, + "device": window.device, "lambda_value": lambda_value, - "rest2_scale": rest2_scale, + "rest2_scale": window.rest2_scale, } ) # Create the dynamics object. dynamics = system.dynamics(**dynamics_kwargs) + window.dynamics = dynamics # Write the OpenMM XML file to the output directory (only if not already # written during equilibration). @@ -851,6 +1067,7 @@ def generate_lam_vals(lambda_base, increment=0.001): gcmc_stats, self._lambda_values[index] ) gcmc_sampler.restore_stats(gcmc_stats) + terminal_flip_sampler = window.terminal_flip_sampler if terminal_flip_sampler is not None and "terminal_flip" in stats: attempted, accepted = stats["terminal_flip"] terminal_flip_sampler.reset(attempted, accepted) @@ -870,640 +1087,442 @@ def generate_lam_vals(lambda_base, increment=0.001): f"[{target[0]:.3f}, {target[1]:.3f}, {target[2]:.3f}] A" ) - # Set the number of neighbours used for the energy calculation. - # If not None, then we add one to account for the extra windows - # used for finite-difference gradient analysis. - if self._config.num_energy_neighbours is not None: - num_energy_neighbours = self._config.num_energy_neighbours + 1 - else: - num_energy_neighbours = None + def _write_post_equilibration_checkpoint(self, system, window): + """ + Write a checkpoint for a window straight after equilibration. - # Store the current checkpoint frequency. - checkpoint_frequency = self._config.checkpoint_frequency + Parameters + ---------- - # Store the checkpoint time in nanoseconds. - checkpoint_interval = checkpoint_frequency.to("ns") + system: :class: `System ` + The equilibrated system. - # Write a checkpoint immediately after equilibration so that a restart - # after an early production crash doesn't need to re-equilibrate. - if is_equilibrated: - lock = _FileLock(self._lock_file) - with lock.acquire(timeout=self._config.timeout.to("seconds")): - _, error = self._checkpoint( - system, - index, - block=-1, - speed=0.0, - lambda_energy=lambda_energy, - lambda_grad=lambda_grad, - context=dynamics.context(), - gcmc_sampler=gcmc_sampler, + window: _Window + The window being simulated. + """ + lock = _FileLock(self._lock_file) + with lock.acquire(timeout=self._config.timeout.to("seconds")): + _, error = self._checkpoint( + system, + window.index, + block=-1, + speed=0.0, + lambda_energy=window.lambda_energy, + lambda_grad=window.lambda_grad, + context=window.dynamics.context(), + gcmc_sampler=window.gcmc_sampler, + ) + if error is not None: + msg = ( + f"Post-equilibration checkpoint failed for {_lam_sym} = " + f"{window.lambda_value:.5f}:\n{error}" ) - if error is not None: - msg = ( - f"Post-equilibration checkpoint failed for {_lam_sym} = " - f"{lambda_value:.5f}:\n{error}" - ) - _logger.error(msg) - raise error + _logger.error(msg) + raise error + _logger.info( + f"Writing post-equilibration checkpoint " + f"for {_lam_sym} = {window.lambda_value:.5f}" + ) + + def _run_checkpoint_blocks(self, window, time): + """ + Run production dynamics in blocks, writing a checkpoint after each. + + Parameters + ---------- + + window: _Window + The window being simulated. + + time: sire.units.GeneralUnit + The simulation time remaining. + """ + lambda_value = window.lambda_value + dynamics = window.dynamics + + # The number of steps in each block. The last is shorter if the time + # remaining isn't a whole number of checkpoint intervals. + num_steps = self._num_steps(time) + checkpoint_steps = self._num_steps(self._config.checkpoint_frequency) + blocks = [checkpoint_steps] * (num_steps // checkpoint_steps) + if num_steps % checkpoint_steps > 0: + blocks.append(num_steps % checkpoint_steps) + total_blocks = self._start_block + len(blocks) + + for i, block_steps in enumerate(blocks): + block = self._start_block + i + is_final_block = i == len(blocks) - 1 + block_start = _timer() + + try: + self._run_block(window, block_steps) + except Exception as e: + raise RuntimeError( + f"Dynamics block {block + 1} for {_lam_sym} = {lambda_value:.5f} failed: {e}" + ) + + try: + # Commit the current system. + system = dynamics.commit() + + # Calculate the speed in nanoseconds per day. + duration = (block_steps * self._config.timestep).to("ns") + speed = duration / ((_timer() - block_start) / 86400) + + self._write_checkpoint(window, system, block, speed, is_final_block) + + # Delete all trajectory frames from the Sire system within the + # dynamics object. + dynamics._d._sire_mols.delete_all_frames() + _logger.info( - f"Writing post-equilibration checkpoint " + f"Finished block {block + 1} of {total_blocks} " f"for {_lam_sym} = {lambda_value:.5f}" ) - # Store the start time. - start = _timer() + self._log_samplers(window, system) - # Run the simulation, checkpointing in blocks. - if checkpoint_frequency.value() > 0.0: - # Calculate the number of blocks and the remainder time. - frac = (time / checkpoint_frequency).value() + if is_final_block: + _logger.success( + f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1" + ) + except Exception as e: + raise RuntimeError( + f"Checkpoint failed for {_lam_sym} = {lambda_value:.5f}: {e}" + ) - # Handle the case where the runtime is less than the checkpoint frequency. - if frac < 1.0: - frac = 1.0 - checkpoint_frequency = time - checkpoint_interval = checkpoint_frequency.to("ns") + def _run_without_checkpoints(self, window, time): + """ + Run production dynamics in a single block, writing one checkpoint at + the end. - num_blocks = int(frac) - rem = round(frac - num_blocks, 12) + Parameters + ---------- - # Run the dynamics in blocks. - for block in range(int(num_blocks)): - # Add the start block number. - block += self._start_block + window: _Window + The window being simulated. - # Record the start time. - block_start = _timer() + time: sire.units.GeneralUnit + The simulation time remaining. + """ + lambda_value = window.lambda_value - # Run the dynamics. - try: - # Run in sub-blocks when any MC sampler is active or energy - # components are being saved; otherwise run the full block. - needs_subblock = ( - gcmc_sampler is not None - or terminal_flip_sampler is not None - or self._config.save_energy_components - ) - if needs_subblock: - runtime = _sr.u("0ps") - ec_elapsed = _sr.u("0ps") - flip_counter = 0 - save_frames = ( - gcmc_sampler is not None - and self._config.frame_frequency > 0 - ) - next_frame = ( - self._config.frame_frequency if save_frames else None - ) - # Sub-block size: shortest active MC frequency, or - # energy_frequency when only saving energy components. - if ( - gcmc_sampler is not None - and terminal_flip_sampler is not None - ): - block_size = min( - self._config.gcmc_frequency, - self._config.terminal_flip_frequency, - ) - elif gcmc_sampler is not None: - block_size = self._config.gcmc_frequency - elif terminal_flip_sampler is not None: - block_size = self._config.terminal_flip_frequency - else: - block_size = self._config.energy_frequency - # How often to attempt each MC move (in sub-block units). - gcmc_every = ( - max( - 1, - round( - (self._config.gcmc_frequency / block_size).value() - ), - ) - if gcmc_sampler is not None - else None - ) - mc_flip_every = ( - max( - 1, - round( - ( - self._config.terminal_flip_frequency - / block_size - ).value() - ), - ) - if terminal_flip_sampler is not None - else None - ) + start = _timer() - while runtime < checkpoint_frequency: - needs_pre_run_snapshot = False - - # GCMC move. - if ( - gcmc_sampler is not None - and flip_counter % gcmc_every == 0 - ): - _logger.info( - f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}" - ) - gcmc_sampler.push() - try: - gcmc_sampler.move(dynamics.context()) - finally: - gcmc_sampler.pop() - needs_pre_run_snapshot = self._config.auto_fix_minimise - - # Terminal flip move. - if ( - terminal_flip_sampler is not None - and flip_counter % mc_flip_every == 0 - ): - _logger.info( - f"Performing terminal flip move at " - f"{_lam_sym} = {lambda_value:.5f}" - ) - flip_accepted = terminal_flip_sampler.move( - dynamics.context() - ) - if flip_accepted: - if self._config.auto_fix_minimise: - needs_pre_run_snapshot = True - if self._config.randomise_velocities: - dynamics.randomise_velocities() - - # Snapshot the context state for crash recovery if - # any MC move changed positions. - if needs_pre_run_snapshot: - dynamics._d._pre_run_state = ( - dynamics.context().getState( - getPositions=True, getVelocities=True - ) - ) - - # Record ghost residues immediately before the dynamics - # block if a frame will be saved within it. - if save_frames and runtime + block_size >= next_frame: - self._save_ghost_residues(index, gcmc_sampler) - next_frame += self._config.frame_frequency - - # Run the dynamics block. - dynamics.run( - block_size, - energy_frequency=self._config.energy_frequency, - frame_frequency=self._config.frame_frequency, - lambda_windows=lambda_array, - rest2_scale_factors=rest2_scale_factors, - save_velocities=self._config.save_velocities, - auto_fix_minimise=self._config.auto_fix_minimise, - num_energy_neighbours=num_energy_neighbours, - null_energy=self._config.null_energy, - save_crash_report=self._config.save_crash_report, - excess_chemical_potential=( - self._mu_ex if gcmc_sampler is not None else None - ), - num_waters=( - _np.sum(gcmc_sampler.water_state()) - if gcmc_sampler is not None - else None - ), - ) + try: + self._run_block(window, self._num_steps(time)) + except Exception as e: + raise RuntimeError( + f"Dynamics for {_lam_sym} = {lambda_value:.5f} failed: {e}" + ) - runtime += block_size - ec_elapsed += block_size - flip_counter += 1 - - # Save energy components. - if self._config.save_energy_components: - self._save_energy_components( - index, - dynamics.context(), - (block * checkpoint_frequency + ec_elapsed).to( - "ns" - ), - ) - - else: - dynamics.run( - checkpoint_frequency, - energy_frequency=self._config.energy_frequency, - frame_frequency=self._config.frame_frequency, - lambda_windows=lambda_array, - rest2_scale_factors=rest2_scale_factors, - save_velocities=self._config.save_velocities, - auto_fix_minimise=self._config.auto_fix_minimise, - num_energy_neighbours=num_energy_neighbours, - null_energy=self._config.null_energy, - save_crash_report=self._config.save_crash_report, - ) - except Exception as e: - raise RuntimeError( - f"Dynamics block {block + 1} for {_lam_sym} = {lambda_value:.5f} failed: {e}" - ) + # Commit the current system. + system = window.dynamics.commit() - # Checkpoint. - try: - # Commit the current system. - system = dynamics.commit() - - # Record the end time. - block_end = _timer() - - # Work how many fractional days the block took. - block_time = (block_end - block_start) / 86400 - - # Calculate the speed in nanoseconds per day. - speed = checkpoint_interval / block_time - - # Check if this is the final block. - is_final_block = ( - block - self._start_block - ) == num_blocks - 1 and rem == 0 - - # Create the lock. - lock = _FileLock(self._lock_file) - - # Acquire the file lock to ensure that the checkpoint files are - # in a consistent state if read by another process. - with lock.acquire(timeout=self._config.timeout.to("seconds")): - # Backup any existing checkpoint files. - index, error = self._backup_checkpoint(index) - - if error is not None: - raise error - - # Write the checkpoint files. - index, error = self._checkpoint( - system, - index, - block, - speed, - lambda_energy=lambda_energy, - lambda_grad=lambda_grad, - is_final_block=is_final_block, - context=dynamics.context(), - gcmc_sampler=gcmc_sampler, - ) + # Calculate the speed in nanoseconds per day. + speed = time.to("ns") / ((_timer() - start) / 86400) - if error is not None: - raise error + try: + self._write_checkpoint(window, system, 0, speed, is_final_block=True) + except Exception as e: + msg = f"Checkpoint failed for {_lam_sym} = {lambda_value:.5f}: {e}" + _logger.error(msg) + raise RuntimeError(msg) - # Save sampler statistics alongside the checkpoint. - self._save_sampler_stats( - index, gcmc_sampler, terminal_flip_sampler - ) + _logger.success( + f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1" + ) - # Delete all trajectory frames from the Sire system within the - # dynamics object. - dynamics._d._sire_mols.delete_all_frames() + def _run_block(self, window, num_steps): + """ + Run a block of production dynamics. This is split into sub-blocks + when Monte Carlo moves are performed or energy components are saved. - _logger.info( - f"Finished block {block + 1} of {self._start_block + num_blocks} " - f"for {_lam_sym} = {lambda_value:.5f}" - ) + Parameters + ---------- - # Log the number of waters within the GCMC sampling volume. - if gcmc_sampler is not None: - gcmc_sampler.push() - try: - n_moves = gcmc_sampler._num_moves - acc_str = ( - f", acceptance rate = {gcmc_sampler.move_acceptance_probability():.3f}" - f" (ins = {gcmc_sampler.num_insertions()}, del = {gcmc_sampler.num_deletions()})" - if n_moves > 0 - else "" - ) - # Count against the context, since dynamics have - # run since the last move. - num_waters = gcmc_sampler.num_waters( - context=dynamics.context() - ) - _logger.info( - f"Current number of waters in GCMC volume at {_lam_sym} = {lambda_value:.5f} " - f"(time = {system.time().to('ns'):.4f} ns) is {num_waters}{acc_str}" - ) - finally: - gcmc_sampler.pop() - - # Log terminal flip acceptance rate. - if terminal_flip_sampler is not None: - _logger.info( - f"Terminal flip acceptance rate at " - f"{_lam_sym} = {lambda_value:.5f}: " - f"{terminal_flip_sampler.acceptance_rate:.3f} " - f"({terminal_flip_sampler.num_accepted}/" - f"{terminal_flip_sampler.num_attempted})" - ) + window: _Window + The window being simulated. - if is_final_block: - _logger.success( - f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1" - ) - except Exception as e: - raise RuntimeError( - f"Checkpoint failed for {_lam_sym} = {lambda_value:.5f}: {e}" - ) + num_steps: int + The number of integration steps in the block. + """ + index = window.index + lambda_value = window.lambda_value + dynamics = window.dynamics + gcmc_sampler = window.gcmc_sampler + terminal_flip_sampler = window.terminal_flip_sampler + + if ( + gcmc_sampler is None + and terminal_flip_sampler is None + and not self._config.save_energy_components + ): + self._run_dynamics(window, num_steps) + window.step += num_steps + return - # Handle the remainder time. (There will be no remainer when GCMC sampling.) - if rem > 0: - block += 1 - block_start = _timer() - try: - # Perform one terminal flip at the start of the remainder block. - if terminal_flip_sampler is not None: - _logger.info( - f"Performing terminal flip move at " - f"{_lam_sym} = {lambda_value:.5f}" - ) - if ( - terminal_flip_sampler.move(dynamics.context()) - and self._config.randomise_velocities - ): - dynamics.randomise_velocities() - - dynamics.run( - rem, - energy_frequency=self._config.energy_frequency, - frame_frequency=self._config.frame_frequency, - lambda_windows=lambda_array, - rest2_scale_factors=rest2_scale_factors, - save_velocities=self._config.save_velocities, - auto_fix_minimise=self._config.auto_fix_minimise, - num_energy_neighbours=num_energy_neighbours, - null_energy=self._config.null_energy, - save_crash_report=self._config.save_crash_report, - ) + # The number of steps between each MC move, or between energy samples + # when only saving energy components. + gcmc_steps = ( + self._num_steps(self._config.gcmc_frequency) + if gcmc_sampler is not None + else None + ) + flip_steps = ( + self._num_steps(self._config.terminal_flip_frequency) + if terminal_flip_sampler is not None + else None + ) + intervals = [n for n in (gcmc_steps, flip_steps) if n is not None] + if not intervals: + intervals = [self._num_steps(self._config.energy_frequency)] - # Commit the current system. - system = dynamics.commit() + save_frames = gcmc_sampler is not None and self._config.frame_frequency > 0 + frame_steps = self._num_steps(self._config.frame_frequency) - # Save the energy contribution for each force. - if self._config.save_energy_components: - self._save_energy_components( - index, dynamics.context(), system.time().to("ns") - ) + end = window.step + num_steps - # Record the end time. - block_end = _timer() - - # Work how many fractional days the block took. - block_time = (block_end - block_start) / 86400 - - # Calculate the speed in nanoseconds per day. - speed = checkpoint_interval / block_time - - # Create the lock. - lock = _FileLock(self._lock_file) - - # Acquire the file lock to ensure that the checkpoint files are - # in a consistent state if read by another process. - with lock.acquire(timeout=self._config.timeout.to("seconds")): - # Backup any existing checkpoint files. - index, error = self._backup_checkpoint(index) - - if error is not None: - raise error - - # Write the checkpoint files. - index, error = self._checkpoint( - system, - index, - block, - speed, - lambda_energy=lambda_energy, - lambda_grad=lambda_grad, - is_final_block=True, - context=dynamics.context(), - gcmc_sampler=gcmc_sampler, - ) + while window.step < end: + # Run to the next move, counting from the start of the simulation so + # that the moves keep their frequency across blocks. + step = min(min(n - window.step % n for n in intervals), end - window.step) - if error is not None: - raise error + needs_pre_run_snapshot = False - # Save sampler statistics alongside the checkpoint. - self._save_sampler_stats( - index, gcmc_sampler, terminal_flip_sampler - ) + # GCMC move. + if gcmc_sampler is not None and window.step % gcmc_steps == 0: + _logger.info(f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}") + gcmc_sampler.push() + try: + gcmc_sampler.move(dynamics.context()) + finally: + gcmc_sampler.pop() + needs_pre_run_snapshot = self._config.auto_fix_minimise - # Delete all trajectory frames from the Sire system within the - # dynamics object. - dynamics._d._sire_mols.delete_all_frames() + # Terminal flip move. + if terminal_flip_sampler is not None and window.step % flip_steps == 0: + _logger.info( + f"Performing terminal flip move at {_lam_sym} = {lambda_value:.5f}" + ) + flip_accepted = terminal_flip_sampler.move(dynamics.context()) + if flip_accepted: + if self._config.auto_fix_minimise: + needs_pre_run_snapshot = True + if self._config.randomise_velocities: + dynamics.randomise_velocities() + + # Snapshot the context state for crash recovery if any MC move + # changed positions. + if needs_pre_run_snapshot: + dynamics._d._pre_run_state = dynamics.context().getState( + getPositions=True, getVelocities=True + ) - _logger.info( - f"Finished block {block + 1} of {self._start_block + num_blocks} " - f"for {_lam_sym} = {lambda_value:.5f}" - ) + # Record ghost residues immediately before the dynamics block if a + # frame will be saved within it. + if ( + save_frames + and (window.step + step) // frame_steps > window.step // frame_steps + ): + self._save_ghost_residues(index, gcmc_sampler) + + # Run the dynamics block. + self._run_dynamics( + window, + step, + excess_chemical_potential=( + self._mu_ex if gcmc_sampler is not None else None + ), + num_waters=( + _np.sum(gcmc_sampler.water_state()) + if gcmc_sampler is not None + else None + ), + ) - _logger.success( - f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1" - ) - except Exception as e: - raise RuntimeError( - f"Final dynamics block for {_lam_sym} = {lambda_value:.5f} failed: {e}" - ) - else: - try: - # Run in sub-blocks when any MC sampler is active or energy - # components are being saved; otherwise run a single block. - needs_subblock = ( - gcmc_sampler is not None - or terminal_flip_sampler is not None - or self._config.save_energy_components + window.step += step + + # Save energy components. + if self._config.save_energy_components: + self._save_energy_components( + index, dynamics.context(), dynamics.current_time().to("ns") ) - if needs_subblock: - runtime = _sr.u("0ps") - ec_elapsed = _sr.u("0ps") - flip_counter = 0 - save_frames = ( - gcmc_sampler is not None and self._config.frame_frequency > 0 - ) - next_frame = self._config.frame_frequency if save_frames else None - # Sub-block size: shortest active MC frequency, or - # energy_frequency when only saving energy components. - if gcmc_sampler is not None and terminal_flip_sampler is not None: - block_size = min( - self._config.gcmc_frequency, - self._config.terminal_flip_frequency, - ) - elif gcmc_sampler is not None: - block_size = self._config.gcmc_frequency - elif terminal_flip_sampler is not None: - block_size = self._config.terminal_flip_frequency - else: - block_size = self._config.energy_frequency - # How often to attempt each MC move (in sub-block units). - gcmc_every = ( - max( - 1, round((self._config.gcmc_frequency / block_size).value()) - ) - if gcmc_sampler is not None - else None - ) - mc_flip_every = ( - max( - 1, - round( - ( - self._config.terminal_flip_frequency / block_size - ).value() - ), - ) - if terminal_flip_sampler is not None - else None - ) - time_base = self._config.runtime - time - while runtime < time: - needs_pre_run_snapshot = False + def _num_steps(self, duration): + """ + The number of integration steps in a length of time. - # GCMC move. - if gcmc_sampler is not None and flip_counter % gcmc_every == 0: - _logger.info( - f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}" - ) - gcmc_sampler.push() - try: - gcmc_sampler.move(dynamics.context()) - finally: - gcmc_sampler.pop() - needs_pre_run_snapshot = self._config.auto_fix_minimise - - # Terminal flip move. - if ( - terminal_flip_sampler is not None - and flip_counter % mc_flip_every == 0 - ): - _logger.info( - f"Performing terminal flip move at " - f"{_lam_sym} = {lambda_value:.5f}" - ) - flip_accepted = terminal_flip_sampler.move( - dynamics.context() - ) - if flip_accepted: - if self._config.auto_fix_minimise: - needs_pre_run_snapshot = True - if self._config.randomise_velocities: - dynamics.randomise_velocities() - - # Snapshot the context state for crash recovery if - # any MC move changed positions. - if needs_pre_run_snapshot: - dynamics._d._pre_run_state = dynamics.context().getState( - getPositions=True, getVelocities=True - ) + Parameters + ---------- - # Record ghost residues immediately before the dynamics - # block if a frame will be saved within it. - if save_frames and runtime + block_size >= next_frame: - self._save_ghost_residues(index, gcmc_sampler) - next_frame += self._config.frame_frequency - - # Run the dynamics block. - dynamics.run( - block_size, - energy_frequency=self._config.energy_frequency, - frame_frequency=self._config.frame_frequency, - lambda_windows=lambda_array, - rest2_scale_factors=rest2_scale_factors, - save_velocities=self._config.save_velocities, - auto_fix_minimise=self._config.auto_fix_minimise, - num_energy_neighbours=num_energy_neighbours, - null_energy=self._config.null_energy, - save_crash_report=self._config.save_crash_report, - excess_chemical_potential=( - self._mu_ex if gcmc_sampler is not None else None - ), - num_waters=( - _np.sum(gcmc_sampler.water_state()) - if gcmc_sampler is not None - else None - ), - ) + duration: sire.units.GeneralUnit + The length of time. - runtime += block_size - ec_elapsed += block_size - flip_counter += 1 + Returns + ------- - # Save energy components. - if self._config.save_energy_components: - self._save_energy_components( - index, - dynamics.context(), - (time_base + ec_elapsed).to("ns"), - ) + int + The number of steps, to the nearest whole step. + """ + return round((duration / self._config.timestep).value()) - else: - dynamics.run( - time, - energy_frequency=self._config.energy_frequency, - frame_frequency=self._config.frame_frequency, - lambda_windows=lambda_array, - rest2_scale_factors=rest2_scale_factors, - save_velocities=self._config.save_velocities, - auto_fix_minimise=self._config.auto_fix_minimise, - num_energy_neighbours=num_energy_neighbours, - null_energy=self._config.null_energy, - save_crash_report=self._config.save_crash_report, - ) - except Exception as e: - raise RuntimeError( - f"Dynamics for {_lam_sym} = {lambda_value:.5f} failed: {e}" - ) + @staticmethod + def _num_energy_samples(window): + """ + The number of energy samples recorded by a window's dynamics. - # Commit the current system. - system = dynamics.commit() + Parameters + ---------- - # Record the end time. - end = _timer() + window: _Window + The window being simulated. - # Work how many fractional days the simulation took. - days = (end - start) / 86400 + Returns + ------- - # Calculate the speed in nanoseconds per day. - speed = time.to("ns") / days + int + The number of samples. + """ + energy_trajectory = window.dynamics._d._energy_trajectory + return 0 if energy_trajectory is None else len(energy_trajectory) - # Create the lock. - lock = _FileLock(self._lock_file) + def _run_dynamics(self, window, num_steps, **kwargs): + """ + Run production dynamics, sampling energies at every lambda value. - # Acquire the file lock to ensure that the checkpoint files are - # in a consistent state if read by another process. - with lock.acquire(timeout=self._config.timeout.to("seconds")): - # Backup any existing checkpoint files. - index, error = self._backup_checkpoint(index) + Parameters + ---------- - if error is not None: - msg = f"Checkpoint backup failed for {_lam_sym} = {lambda_value:.5f}: {error}" - _logger.error(msg) - raise RuntimeError(msg) + window: _Window + The window being simulated. - # Write the checkpoint files. - index, error = self._checkpoint( - system, - index, - 0, - speed, - lambda_energy=lambda_energy, - lambda_grad=lambda_grad, - is_final_block=True, - context=dynamics.context(), - gcmc_sampler=gcmc_sampler, - ) + num_steps: int + The number of integration steps to run. A number of steps is passed + to Sire, since it truncates a time that is a rounding error short + of a whole step. - if error is not None: - msg = f"Checkpoint failed for {_lam_sym} = {lambda_value:.5f}: {error}" - _logger.error(msg) - raise RuntimeError(msg) + kwargs: dict + Additional keyword arguments for the dynamics run. + """ + window.dynamics.run( + num_steps, + energy_frequency=self._config.energy_frequency, + frame_frequency=self._config.frame_frequency, + lambda_windows=window.lambda_array, + rest2_scale_factors=window.rest2_scale_factors, + save_velocities=self._config.save_velocities, + auto_fix_minimise=self._config.auto_fix_minimise, + num_energy_neighbours=window.num_energy_neighbours, + null_energy=self._config.null_energy, + save_crash_report=self._config.save_crash_report, + **kwargs, + ) - # Save sampler statistics alongside the final checkpoint. - self._save_sampler_stats(index, gcmc_sampler, terminal_flip_sampler) + def _write_checkpoint(self, window, system, block, speed, is_final_block): + """ + Back up the previous checkpoint of a window, then write the new one and + the sampler statistics, all under the file lock so that another process + always reads a consistent set of files. - _logger.success( - f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1" + Parameters + ---------- + + window: _Window + The window being simulated. + + system: :class: `System ` + The system to checkpoint. + + block: int + The index of the block that was run. + + speed: float + The speed of the block in ns per day. + + is_final_block: bool + Whether this is the last block of the simulation. + """ + # Only the energy samples recorded since the last checkpoint are new. + num_energy_samples = self._num_energy_samples(window) + + lock = _FileLock(self._lock_file) + with lock.acquire(timeout=self._config.timeout.to("seconds")): + # Backup any existing checkpoint files. + _, error = self._backup_checkpoint(window.index) + if error is not None: + raise error + + # Write the checkpoint files. + _, error = self._checkpoint( + system, + window.index, + block, + speed, + lambda_energy=window.lambda_energy, + lambda_grad=window.lambda_grad, + is_final_block=is_final_block, + context=window.dynamics.context(), + gcmc_sampler=window.gcmc_sampler, + num_energy_rows=num_energy_samples - window.num_energy_samples, ) + if error is not None: + raise error + window.num_energy_samples = num_energy_samples - return time + # Save sampler statistics alongside the checkpoint. + self._save_sampler_stats( + window.index, window.gcmc_sampler, window.terminal_flip_sampler + ) + + def _log_samplers(self, window, system): + """ + Log the number of waters in the GCMC sampling volume and the terminal + flip acceptance rate after a block. + + Parameters + ---------- + + window: _Window + The window being simulated. + + system: :class: `System ` + The system after the block. + """ + lambda_value = window.lambda_value + gcmc_sampler = window.gcmc_sampler + terminal_flip_sampler = window.terminal_flip_sampler + + # Log the number of waters within the GCMC sampling volume. + if gcmc_sampler is not None: + gcmc_sampler.push() + try: + n_moves = gcmc_sampler._num_moves + acc_str = ( + f", acceptance rate = {gcmc_sampler.move_acceptance_probability():.3f}" + f" (ins = {gcmc_sampler.num_insertions()}, del = {gcmc_sampler.num_deletions()})" + if n_moves > 0 + else "" + ) + # Count against the context, since dynamics have run since the + # last move. + num_waters = gcmc_sampler.num_waters(context=window.dynamics.context()) + _logger.info( + f"Current number of waters in GCMC volume at {_lam_sym} = {lambda_value:.5f} " + f"(time = {system.time().to('ns'):.4f} ns) is {num_waters}{acc_str}" + ) + finally: + gcmc_sampler.pop() + + # Log terminal flip acceptance rate. + if terminal_flip_sampler is not None: + _logger.info( + f"Terminal flip acceptance rate at " + f"{_lam_sym} = {lambda_value:.5f}: " + f"{terminal_flip_sampler.acceptance_rate:.3f} " + f"({terminal_flip_sampler.num_accepted}/" + f"{terminal_flip_sampler.num_attempted})" + ) def _save_sampler_stats(self, index, gcmc_sampler, terminal_flip_sampler): """ diff --git a/tests/runner/test_gcmc.py b/tests/runner/test_gcmc.py index b8228dd..2bada83 100644 --- a/tests/runner/test_gcmc.py +++ b/tests/runner/test_gcmc.py @@ -130,3 +130,46 @@ def release(self, *args, **kwargs): ) assert num_lines() > 0 assert traj.num_frames() == num_lines() + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_runner_gcmc_runtime_remainder(ethane_methanol): + """ + Validate that a runtime which isn't a multiple of the checkpoint frequency + finishes with GCMC moves, with a ghost residue record for every frame. + """ + pytest.importorskip("loch") + + import sire as sr + + with tempfile.TemporaryDirectory() as tmpdir: + config = Config( + runtime="20fs", + output_directory=tmpdir, + energy_frequency="4fs", + checkpoint_frequency="8fs", + frame_frequency="4fs", + platform="cuda", + max_threads=1, + num_lambda=2, + gcmc=True, + gcmc_selection="resname LIG", + gcmc_frequency="4fs", + ) + + runner = Runner(ethane_methanol, config) + runner.run() + + for lam in runner._lambda_values: + ghost_file = Path(tmpdir) / f"gcmc_ghosts_{lam:.5f}.txt" + traj = sr.load( + str(Path(tmpdir) / "system0.prm7"), + str(Path(tmpdir) / f"traj_{lam:.5f}.dcd"), + ) + assert traj.num_frames() == 5 + assert len(ghost_file.read_text().splitlines()) == 5 + + # A GCMC move at the start of every 4 fs interval, including the + # remainder, for each window. + log = (Path(tmpdir) / config.log_file).read_text() + assert log.count("Performing GCMC move at") == 10 diff --git a/tests/runner/test_gcmc_stats.py b/tests/runner/test_gcmc_stats.py index c6ff648..d537267 100644 --- a/tests/runner/test_gcmc_stats.py +++ b/tests/runner/test_gcmc_stats.py @@ -2,7 +2,7 @@ import pytest from somd2.runner._base import RunnerBase -from somd2.runner._repex import DynamicsCache +from somd2.runner._dynamics_cache import DynamicsCache def counters(num_moves=0): diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 57a93cb..f9d630b 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -5,6 +5,7 @@ import tempfile from somd2.runner import RepexRunner +from somd2.runner._adaptive import WindowRanker, cold_rate, statistical_inefficiency from somd2.runner._base import RunnerBase from somd2.runner._repex import _mix_subset from somd2.config import Config @@ -83,7 +84,7 @@ def test_repex_mixing_moves_configurations(): Validate that each configuration is moved to the state that the mixing accepted, for a mix that can only be a 3-cycle. """ - from somd2.runner._repex import DynamicsCache + from somd2.runner._dynamics_cache import DynamicsCache num_replicas = 3 @@ -232,7 +233,7 @@ def test_repex_mixing_samples_boltzmann(slot_schedule, inverted): import pickle from types import SimpleNamespace - from somd2.runner._repex import DynamicsCache + from somd2.runner._dynamics_cache import DynamicsCache n = 4 num_cycles = 20000 @@ -1044,7 +1045,7 @@ def test_physical_device_mapping(gpu_devices, expected): backing it, since OpenMM numbers devices relative to the visible set whereas pynvml and pyopencl enumerate all of them. """ - from somd2.runner._repex import DynamicsCache + from somd2.runner._dynamics_cache import DynamicsCache cache = object.__new__(DynamicsCache) cache._gpu_devices = gpu_devices @@ -1058,7 +1059,7 @@ def test_max_contexts_advice(): must name a number that the user can pass to 'max_contexts', and mention the frame frequency constraint that comes with it. """ - from somd2.runner._repex import DynamicsCache + from somd2.runner._dynamics_cache import DynamicsCache advice = DynamicsCache._max_contexts_advice(6) assert "'max_contexts' to 6 or fewer" in advice @@ -1084,7 +1085,7 @@ def test_gcmc_state_follows_replica(): The ghost files and the sampling statistics belong to the lambda window rather than the configuration, so they must not be permuted. """ - from somd2.runner._repex import DynamicsCache + from somd2.runner._dynamics_cache import DynamicsCache num_replicas = 4 @@ -1131,7 +1132,7 @@ def test_check_device_memory_queries_requested_device(monkeypatch, device, key, import sys import types - from somd2.runner._repex import DynamicsCache + from somd2.runner._dynamics_cache import DynamicsCache pynvml = pytest.importorskip("pynvml") @@ -1173,7 +1174,7 @@ def test_legacy_checkpoint_restore(): loaded. These stored the replica states unpermuted, with the states array holding the mapping to apply on restart. """ - from somd2.runner._repex import DynamicsCache + from somd2.runner._dynamics_cache import DynamicsCache n = 4 legacy = { @@ -1209,6 +1210,43 @@ def test_legacy_checkpoint_restore(): assert cache._ghost_files == [None] * n +def test_checkpoint_from_before_the_cache_moved(monkeypatch): + """ + Validate that a checkpoint pickled when the dynamics cache was defined in + _repex can still be loaded. + """ + import pickle + + from somd2.runner._dynamics_cache import DynamicsCache + + n = 2 + cache = object.__new__(DynamicsCache) + cache.__setstate__( + { + "_lambdas": [0.0, 1.0], + "_rest2_scale_factors": [1.0] * n, + "_num_slots": n, + "_states": np.arange(n), + "_time": None, + "_openmm_states": [None] * n, + "_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 monkeypatch.context() as m: + m.setattr(DynamicsCache, "__module__", "somd2.runner._repex") + data = pickle.dumps(cache) + + assert b"somd2.runner._repex" in data + assert isinstance(pickle.loads(data), DynamicsCache) + + @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") def test_repex_max_contexts_change_on_restart(ethane_methanol): """ @@ -1486,7 +1524,7 @@ 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 + from somd2.runner._dynamics_cache import DynamicsCache cache = object.__new__(DynamicsCache) cache._num_replicas = 10 @@ -1509,35 +1547,44 @@ def test_statistical_inefficiency(): rng = np.random.default_rng(42) x = rng.normal(size=20000) - assert RepexRunner._statistical_inefficiency(x) == pytest.approx(1.0, abs=0.1) + assert 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) + assert statistical_inefficiency(y) == pytest.approx(expected, rel=0.15) -def _ranking_runner(num_lambda, **config): +def _ranking_runner(num_lambda, is_ring_break=False, is_cached=False, **config): """ Create a runner holding only what is needed to rank windows. """ - from types import SimpleNamespace + from somd2.runner._dynamics_cache import DynamicsCache 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, + + cache = object.__new__(DynamicsCache) + cache._num_replicas = num_lambda + cache._num_proposed = zeros + cache._num_accepted = zeros + cache._hot_replicas = None + cache._cold_votes = [0] * num_lambda + cache._adaptive_data = None + runner._dynamics_cache = cache + + runner._ranker = WindowRanker( + num_lambda, + runner._config.adaptive_min_samples, + zeros, + zeros, + is_ring_break=is_ring_break, + is_cached=is_cached, ) - runner._reset_adaptive_data() return runner @@ -1579,8 +1626,8 @@ def _fill_ranking_data(runner, hard_pair, num_samples=200, seed=7): 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)) + runner._ranker.upper[k] = list(rng.normal(0.5, sigma, num_samples)) + runner._ranker.lower[k + 1] = list(rng.normal(-0.5, sigma, num_samples)) def test_rank_windows(): @@ -1687,8 +1734,8 @@ def test_rank_windows_minimises_variance_for_compute(): # Forward and reverse estimates that agree, so the pair isn't flagged. rng = np.random.default_rng(11) - runner._adaptive_upper[3] = list(rng.normal(1.0, 1.0, 200)) - runner._adaptive_lower[4] = list(rng.normal(0.0, 1.0, 200)) + runner._ranker.upper[3] = list(rng.normal(1.0, 1.0, 200)) + runner._ranker.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] @@ -1725,8 +1772,8 @@ def rank(samples): # 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)) + runner._ranker.upper[k] = list(rng.normal(1.0, 1.0, 200)) + runner._ranker.lower[k + 1] = list(rng.normal(0.0, 1.0, 200)) for _ in range(2): _rank(runner, 2, window_samples=samples, cycles_left=1000) @@ -1829,7 +1876,7 @@ def test_cold_rate(first, num_cycles, cycles_per_checkpoint, mix): run loop's schedule. """ checkpoints = RepexRunner._checkpoint_cycles(num_cycles, cycles_per_checkpoint) - assert RepexRunner._cold_rate(first, num_cycles, checkpoints, mix) == pytest.approx( + assert cold_rate(first, num_cycles, checkpoints, mix) == pytest.approx( _full_cycle_share(first, num_cycles, cycles_per_checkpoint, mix) ) @@ -1860,8 +1907,8 @@ def test_rank_windows_without_uncertainty(): runner = _ranking_runner(4, adaptive_min_samples=10) for k in range(3): - runner._adaptive_upper[k] = [0.5] * 100 - runner._adaptive_lower[k + 1] = [-0.5] * 100 + runner._ranker.upper[k] = [0.5] * 100 + runner._ranker.lower[k + 1] = [-0.5] * 100 _rank(runner, num_gpus=1) assert runner._dynamics_cache._hot_replicas is None @@ -1875,21 +1922,50 @@ def test_adaptive_data_is_restored(): 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] + runner._ranker.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)) + expected = (runner._ranker.lower, runner._ranker.upper) + data = pickle.loads(pickle.dumps(runner._ranker.to_dict())) - runner._reset_adaptive_data() - runner._dynamics_cache._adaptive_data = data - runner._restore_adaptive_data() + zeros = np.zeros((8, 8)) + runner._ranker.reset(zeros, zeros) + runner._ranker.restore(data) - assert runner._adaptive_lower == expected[0] - assert runner._adaptive_upper == expected[1] - assert runner._adaptive_distance[2] == [1.0, 2.0] + assert runner._ranker.lower == expected[0] + assert runner._ranker.upper == expected[1] + assert runner._ranker.distance[2] == [1.0, 2.0] + + +@pytest.mark.parametrize("burn_in, is_restored", [("0 ps", True), ("1 ns", False)]) +def test_start_adaptive_restores_data(burn_in, is_restored): + """ + Validate that a restart restores the data used to rank the windows from + the dynamics cache, unless the burn-in has been raised beyond the current + time, which makes every window hot again. + """ + import sire as sr + + saved = _ranking_runner(8, adaptive_min_samples=10) + _fill_ranking_data(saved, hard_pair=3) + + runner = _ranking_runner(8, adaptive_repex=True, adaptive_burn_in=burn_in) + runner._ring_pair = None + runner._is_cached = False + runner._start_time = sr.u("100 ps") + runner._dynamics_cache._adaptive_data = saved._ranker.to_dict() + runner._dynamics_cache._hot_replicas = [2, 3, 4] + runner._dynamics_cache._cold_votes = [1] * 8 + + runner._start_adaptive() + + if is_restored: + assert runner._ranker.upper == saved._ranker.upper + assert runner._dynamics_cache._hot_replicas == [2, 3, 4] + else: + assert runner._ranker.upper == [[]] * 8 + assert runner._dynamics_cache._hot_replicas is None + assert runner._dynamics_cache._cold_votes == [0] * 8 def test_rank_windows_needs_enough_samples(): @@ -1903,7 +1979,7 @@ def test_rank_windows_needs_enough_samples(): _rank(runner, num_gpus=2) assert runner._dynamics_cache._hot_replicas is None - assert len(runner._adaptive_upper[0]) == 20 + assert len(runner._ranker.upper[0]) == 20 def test_rank_windows_keeps_undersampled_windows_hot(): @@ -1920,7 +1996,7 @@ def test_rank_windows_keeps_undersampled_windows_hot(): y = np.zeros(200) for i in range(1, len(y)): y[i] = 0.95 * y[i - 1] + rng.normal(0.0, 0.01) - runner._adaptive_upper[6] = list(0.5 + y) + runner._ranker.upper[6] = list(0.5 + y) _rank_twice(runner, num_gpus=2) assert 6 in runner._dynamics_cache._hot_replicas @@ -1931,19 +2007,18 @@ 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) + runner = _ranking_runner(8, is_ring_break=True, adaptive_min_samples=50) _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)) + runner._ranker.distance[k] = list(rng.normal(1.5, 0.05, 200)) y = np.zeros(200) for i in range(1, len(y)): y[i] = 0.95 * y[i - 1] + rng.normal() - runner._adaptive_distance[7] = list(10.0 + y) + runner._ranker.distance[7] = list(10.0 + y) _rank_twice(runner, num_gpus=2) assert 7 in runner._dynamics_cache._hot_replicas @@ -1984,8 +2059,7 @@ 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 + runner = _ranking_runner(8, is_cached=True, adaptive_min_samples=10) # The flagged pair 0 gives windows 0, 1 and 2, one on each of three GPUs. _fill_ranking_data(runner, hard_pair=0) @@ -2055,7 +2129,7 @@ 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 + from somd2.runner._dynamics_cache import DynamicsCache cache = object.__new__(DynamicsCache) cache._replica_slot = [r % num_slots for r in range(num_replicas)] @@ -2183,7 +2257,7 @@ def test_adaptive_repex_ranking_arguments(ethane_methanol, monkeypatch): def rank_windows(*args): calls.append(args) - baselines.append(runner._adaptive_proposed.sum()) + baselines.append(runner._ranker.proposed.sum()) monkeypatch.setattr(runner, "_rank_windows", rank_windows) runner.run() @@ -2281,11 +2355,11 @@ def test_adaptive_repex_records_bond_distance(syk_ring_break_mols): assert runner._ring_pair == (1, 7) runner.run() - for distances in runner._adaptive_distance: + for distances in runner._ranker.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) + assert np.mean(runner._ranker.distance[0]) == pytest.approx(1.8, abs=0.2) def test_find_ring_pair_needs_a_single_bond(syk_ring_break_mols): @@ -2310,7 +2384,7 @@ def test_rank_windows_flags_unconverged_pairs(): _fill_ranking_data(runner, hard_pair=3) rng = np.random.default_rng(3) - runner._adaptive_lower[1] = list(rng.normal(-3.0, 0.1, 200)) + runner._ranker.lower[1] = list(rng.normal(-3.0, 0.1, 200)) _rank_twice(runner, num_gpus=2) hot = runner._dynamics_cache._hot_replicas @@ -2322,7 +2396,7 @@ 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 somd2.runner._dynamics_cache import DynamicsCache from tests.runner.test_gcmc_stats import TestLegacyRepexCheckpoint cache = object.__new__(DynamicsCache) diff --git a/tests/runner/test_runner.py b/tests/runner/test_runner.py new file mode 100644 index 0000000..6af940e --- /dev/null +++ b/tests/runner/test_runner.py @@ -0,0 +1,169 @@ +from pathlib import Path + +import numpy as np +import pytest +import tempfile + +from somd2.runner import Runner +from somd2.config import Config +from somd2.io import parquet_to_dataframe + + +def test_runtime_remainder(ethane_methanol): + """ + Validate that a runtime that isn't a multiple of the checkpoint frequency + finishes with a block for the remaining time. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "20fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "12fs", + "frame_frequency": "4fs", + "platform": "CPU", + "max_threads": 1, + "num_lambda": 2, + } + + runner = Runner(ethane_methanol, Config(**config)) + runner.run() + + for lam in ["0.00000", "1.00000"]: + energy_traj, _ = parquet_to_dataframe( + Path(tmpdir) / f"energy_traj_{lam}.parquet" + ) + times = energy_traj.index.get_level_values(0) + assert list(np.round(times, 6)) == [0.004, 0.008, 0.012, 0.016, 0.020] + + checkpoint_state = np.load(Path(tmpdir) / f"checkpoint_{lam}.npz") + assert checkpoint_state["time_ps"].item() == pytest.approx(0.020) + + +def test_runtime_remainder_with_terminal_flips(phenethyl_mols): + """ + Validate that blocks which aren't a whole number of terminal flip intervals + don't overrun, and that the remainder is run with terminal flips. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "20fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "12fs", + "frame_frequency": "4fs", + "terminal_flip_frequency": "8fs", + "platform": "CPU", + "max_threads": 1, + "num_lambda": 2, + } + + runner = Runner(phenethyl_mols, Config(**config)) + runner.run() + + energy_traj, _ = parquet_to_dataframe( + Path(tmpdir) / "energy_traj_0.00000.parquet" + ) + times = energy_traj.index.get_level_values(0) + assert list(np.round(times, 6)) == [0.004, 0.008, 0.012, 0.016, 0.020] + + # Flips are attempted every 8 fs, at 0, 8 and 16 fs, for each window. + log = (Path(tmpdir) / "log.txt").read_text() + assert log.count("Performing terminal flip move at") == 6 + + +def test_terminal_flip_frequency_across_blocks(phenethyl_mols): + """ + Validate that terminal flips keep their frequency when the checkpoint + frequency isn't a multiple of it. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "48fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "12fs", + "frame_frequency": "4fs", + "terminal_flip_frequency": "8fs", + "platform": "CPU", + "max_threads": 1, + "num_lambda": 1, + } + + runner = Runner(phenethyl_mols, Config(**config)) + runner.run() + + # Every 8 fs, at 0, 8, 16, 24, 32 and 40 fs. + log = (Path(tmpdir) / "log.txt").read_text() + assert log.count("Performing terminal flip move at") == 6 + + +@pytest.mark.parametrize( + "option, value", + [ + ("runtime", "18fs"), + ("checkpoint_frequency", "10fs"), + ("energy_frequency", "6fs"), + ("terminal_flip_frequency", "6fs"), + ], +) +def test_times_must_be_whole_steps(ethane_methanol, option, value): + """ + Validate that times which aren't a whole number of integration steps are + rejected. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "24fs", + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "12fs", + "frame_frequency": "12fs", + "timestep": "4fs", + "platform": "CPU", + "max_threads": 1, + "num_lambda": 2, + option: value, + } + + with pytest.raises(ValueError, match=f"'{option}' must be a multiple"): + Runner(ethane_methanol, Config(**config)) + + +def test_energy_components_after_restart(ethane_methanol): + """ + Validate that energy components are time stamped correctly when a restart + has less than a checkpoint interval left to run. + """ + import pandas as pd + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "12fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "12fs", + "frame_frequency": "4fs", + "save_energy_components": True, + "platform": "CPU", + "max_threads": 1, + "num_lambda": 2, + } + + runner = Runner(ethane_methanol, Config(**config)) + runner.run() + del runner + + config["runtime"] = "20fs" + config["restart"] = True + runner = Runner(ethane_methanol, Config(**config)) + runner.run() + + components = pd.read_parquet(Path(tmpdir) / "energy_components_0.00000.parquet") + # The times are in ns. + times = list(np.round(components["time"] * 1e6, 6)) + assert times == [4, 8, 12, 16, 20] diff --git a/tests/runner/test_terminal_flip.py b/tests/runner/test_terminal_flip.py index 7901121..c21af4b 100644 --- a/tests/runner/test_terminal_flip.py +++ b/tests/runner/test_terminal_flip.py @@ -397,8 +397,8 @@ def test_runner_validation_frequency_multiple(ethane_methanol): output_directory=tmpdir, platform="cpu", num_lambda=2, - energy_frequency="4fs", - terminal_flip_frequency="3fs", # not a multiple of 4fs + energy_frequency="8fs", + terminal_flip_frequency="12fs", # not a multiple of 8fs ) with pytest.raises( ValueError, match="must be a multiple of 'energy_frequency'" diff --git a/tests/viewer/test_progress.py b/tests/viewer/test_progress.py index 44e106b..8ff6703 100644 --- a/tests/viewer/test_progress.py +++ b/tests/viewer/test_progress.py @@ -4,8 +4,8 @@ import pytest import sire as sr -from somd2._viewer._data import Simulation -from somd2.runner._repex import DynamicsCache +from somd2._viewer._data import Simulation, _RepexUnpickler +from somd2.runner._dynamics_cache import DynamicsCache LAMBDAS = [0.0, 0.5, 1.0] @@ -64,6 +64,26 @@ def _config(adaptive): } +@pytest.mark.parametrize( + "module", ["somd2.runner._dynamics_cache", "somd2.runner._repex"] +) +def test_repex_state_from_either_module(tmp_path, monkeypatch, module): + """ + Validate that the replica exchange state is read whether the cache was + pickled from its own module or, before it moved, from _repex. + """ + monkeypatch.setattr(DynamicsCache, "__module__", module) + _simulation(tmp_path, "1 ns") + + with open(tmp_path / "repex_state.pkl", "rb") as f: + data = f.read() + f.seek(0) + state = _RepexUnpickler(f).load() + + assert module.encode() in data + assert state._time.to("ns") == pytest.approx(1.0) + + def test_adaptive_progress_uses_simulation_time(tmp_path): """ Validate that the progress of an adaptive replica exchange simulation is