diff --git a/CHANGELOG.md b/CHANGELOG.md index d699ec4..d9dbe38 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,10 @@ Changelog * Fixed the replica exchange GPU memory check querying the wrong device when `CUDA_VISIBLE_DEVICES` does not start at zero, since OpenMM numbers devices relative to the visible set whereas `pynvml` enumerates all of them [#183](https://github.com/OpenBioSim/somd2/issues/183). * Store GCMC sampling statistics per lambda value, converting those from earlier checkpoints on restart [#184](https://github.com/OpenBioSim/somd2/pull/184). * Link restart systems to the reference end state rather than the perturbed one, since that is the coordinate set that dynamics maintains. Perturbable molecules were otherwise resumed from the coordinates they were built with [#189](https://github.com/OpenBioSim/somd2/pull/189). +* Add `max_contexts` to cap the number of OpenMM contexts used for replica exchange, re-using each across lambda values so that GPU memory no longer limits the number of replicas [#191](https://github.com/OpenBioSim/somd2/pull/191). +* Skip minimisation on restart [#191](https://github.com/OpenBioSim/somd2/pull/191). +* Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket [#191](https://github.com/OpenBioSim/somd2/pull/191). +* Add a `precision` option for GPU platforms, defaulting to `mixed` [#191](https://github.com/OpenBioSim/somd2/pull/191). [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026 -------------------------------------------------------------------------------------- diff --git a/README.md b/README.md index cad485e..f9b8e63 100644 --- a/README.md +++ b/README.md @@ -74,7 +74,7 @@ Then install `somd2` into the environment: pip install -e . ``` -> [!Note] +> [!NOTE] > Pixi does not run conda post-link scripts, so the `ocl-icd-system` > symlink needed for OpenCL won't be created automatically. After > creating the environment (or after a pixi update), run the following @@ -117,7 +117,7 @@ In order to run an alchemical free-energy simulation you will need to first create a stream file containing the _perturbable_ system of interest. This can be created using [BioSimSpace](https://github.com/OpenBioSim/biosimspace). For example, following the tutorial -[here](https://biosimspace.openbiosim.org/versions/2023.4.0/tutorials/hydration_freenrg.html). +[here](https://biosimspace.openbiosim.org/tutorials/hydration_freenrg.html). Once the system is created, it can be streamed to file using, e.g.: ```python @@ -129,23 +129,23 @@ BSS.Stream.save(system, "perturbable_system") You can then run a simulation with: ``` -somd2 perturtbable_system.bss +somd2 perturbable_system.bss ``` The help message provides information on all of the supported options, along with their default values. Options can be specified on the command line, or using a YAML configuration file, passed with the `--config` option. Any options -explicity set on the command line will override those set via the config file. +explicitly set on the command line will override those set via the config file. An example perturbable system for a methane to ethanol perturbation in solvent can be found [here](https://sire.openbiosim.org/m/merged_molecule.s3.bz2). This is a `bzip2` compressed file that will need to be extracted before use. -#### Running SOMD2 using one or more GPUs +### Running SOMD2 using one or more GPUs In order to run using GPUs you will first need to set the relevant environment -variable. For example, to run using 4 CUDA enabled GPUS set `CUDA_VISIBLE_DEVICES=0,1,2,3` -(for openCL and HIP use `OPENCL_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` respectively). +variable. For example, to run using 4 CUDA enabled GPUs set `CUDA_VISIBLE_DEVICES=0,1,2,3` +(for OpenCL and HIP use `OPENCL_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` respectively). By default `SOMD2` will run using the CPU platform, however if the relevant environment variable has been set (as above) the new platform will be detected @@ -154,23 +154,45 @@ available, the `--platform` option can be set (for example `--platform cuda`). By default, `SOMD2` will automatically manage the distribution of lambda windows across all listed devices. In order to restrict the number of devices used -the `--max_gpus` option can be set, for example setting `max_gpus=2` while +the `--max-gpus` option can be set, for example setting `--max-gpus 2` while `CUDA_VISIBLE_DEVICES` are set as above would restrict `SOMD2` to using only GPUs 0 and 1. ## Replica exchange `SOMD2` supports Hamiltonian replica exchange (HREX) simulations, which can be -enabled using the `--replica-exchange` option. Note that dynamics contexts will -be created up-front for all replicas, so this can be memory intensive. As such, +enabled using the `--replica-exchange` option. By default, dynamics contexts are +created up-front for all replicas, so this can be memory intensive. As such, replica exchange is intended for use on multi-GPU nodes with a large amount of -memory. For optimal performance, it is recommended that the number of replicas -be a multiple of the number of GPUs. It is also possible to oversubscribe the -GPUs, i.e. have more than one replica running on a GPU at a time. This can be -controlled via the `--oversubscription-factor` option, e.g. a value of 2 would -allow 2 replicas to run on each GPU at a time. - -The swap frequency for replica exchange is controlled by the `energy-frequency` +memory. It is also possible to oversubscribe the GPUs, i.e. have more than one +replica running on a GPU at a time. This can be controlled via the +`--oversubscription-factor` option, e.g. a value of 2 would allow 2 replicas to +run on each GPU at a time. This requires the NVIDIA multi-process service (MPS) +to be enabled, see [GPU oversubscription](#gpu-oversubscription) below. + +If the number of replicas you want doesn't fit in GPU memory, use the +`--max-contexts` option to cap the number of contexts that are created. Each +context is then re-used to propagate several replicas per cycle, changing its +lambda value as it goes, so the number of replicas is no longer limited by +memory. For example, `--num-lambda 24 --max-contexts 4` runs 24 replicas using +the memory of 4. This costs some performance, since the replicas sharing a +context run one after another rather than at the same time, so only use it when +one context per replica won't fit. When contexts are re-used, `--frame-frequency` +must equal `--checkpoint-frequency`. + +For optimal performance, it is recommended that the number of contexts, i.e. the +number of replicas, or `--max-contexts` if it is set, be a multiple of the number +of GPUs, and no smaller than the number of GPUs multiplied by the +oversubscription factor. `SOMD2` will warn you if this isn't the case. + +Changing the lambda value of a context requires it to be reinitialised whenever a +constrained bond length actually perturbs with lambda, which is slow. If this +overhead is significant, pass `--no-update-constraints` to freeze the +constrained bond lengths at those of a single lambda value, chosen with +`--constraint-lambda-index`. Both options are ignored unless contexts are being +re-used. + +The swap frequency for replica exchange is controlled by the `--energy-frequency` option, i.e. we compute the energies for all replicas at this frequency, then attempt to mix the replicas. A larger value will improve performance, but may reduce the efficiency of the exchange. @@ -179,7 +201,7 @@ reduce the efficiency of the exchange. We also support Replica Exchange with Solute Scaling ([REST2](https://pubs.acs.org/doi/10.1021/jp204407d)) simulations to facilitate sampling for perturbations -involving conformational changes, e.g. ring flips. This can be enabled +involving conformational changes, e.g. ring flips. This can be enabled using the `--rest2-scale` option, which specifies the "temperature" of the REST2 region relative to the rest of the system. By default, the REST2 region comprises _all_ atoms in perturbable molecules, but can be controlled via the @@ -223,7 +245,7 @@ from `conda-forge`. SOMD2 supports terminal ring flip Monte Carlo (MC) moves to improve sampling of terminal aromatic rings in perturbable ligands, as described in -[this paper](https://chemrxiv.org/doi/full/10.26434/chemrxiv-2025-2zkx5). +[this paper](https://doi.org/10.26434/chemrxiv-2025-2zkx5). Each move attempts a discrete rotation of a terminal ring around the bond connecting it to the rest of the molecule, accepted or rejected via the Metropolis criterion. Terminal ring groups are detected automatically from @@ -330,7 +352,7 @@ free_nrg = BSS.FreeEnergy.Relative.difference(pmf1, pmf2) When running HREX with a large number of replicas it can become computationally expensive to compute energies. (We need the energies of each replica at each -lamdba value.) As a shortcut, it's possible to truncate the neighbourhood of +lambda value.) As a shortcut, it's possible to truncate the neighbourhood of windows for which we compute energies, then use a large null energy for the remaining windows. This can be controlled via the `--num-energy-neighbours` option. For example, setting this to 2 would compute energies for the current window and @@ -345,7 +367,7 @@ We support modification of ghost atom bonded terms to avoid spurious coupling to the physical system using the approach described in [this](https://pubs.acs.org/doi/10.1021/acs.jctc.0c01328) paper. These are enabled by default, but can be disabled using the ``--no-ghost-modifications`` -option. Modifications are implemented using the [ghostly](https://gitbub.com/OpenBioSim/ghostly) +option. Modifications are implemented using the [ghostly](https://github.com/OpenBioSim/ghostly) package. ## Note for SOMD1 users @@ -360,7 +382,7 @@ that the perturbation used is consistent with the approach from `somd1`, i.e. it uses the same modifications for bonded-terms involving dummy atoms as `somd1`. Finally, it is also possible to run `somd2` using an existing `somd1` perturbation -file. To do so, you will also need to create a stream file representating the +file. To do so, you will also need to create a stream file representing the λ = 0 state. For existing input generated by `prepareFEP.py`, this can be done as follows. (This assumes that the output has a prefix `somd1`.) @@ -391,7 +413,7 @@ then simply omit the `--somd1-compatibility` option. ## GPU oversubscription If you have an NVIDIA GPU that supports the multi-process service (MPS), you can -oversubscibe the GPU to run multiple OpenMM contexts on the same GPU at once, +oversubscribe the GPU to run multiple OpenMM contexts on the same GPU at once, increasing the throughput of your simulation. To do this, you will need to first enable MPS by running the following command: @@ -403,7 +425,7 @@ The number of contexts that can be run in parallel is then controlled by the `--oversubscription-factor` option, which defaults to 1. More details on MPS, including tuning options, can be found in the following -[techical blog](https://developer.nvidia.com/blog/maximizing-openmm-molecular-dynamics-throughput-with-nvidia-multi-process-service/). +[technical blog](https://developer.nvidia.com/blog/maximizing-openmm-molecular-dynamics-throughput-with-nvidia-multi-process-service/). ## Python API @@ -425,5 +447,5 @@ frequently. (Frames are written to disk and cleared from memory at each checkpoint.) PyMBAR uses JAX by default for GPU acceleration, which can cause issues in -some environments. If you encounter issues when analysing simlation output, +some environments. If you encounter issues when analysing simulation output, try setting the `PYMBAR_DISABLE_JAX` environment variable to `1`. diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 9efd79c..adaa0e6 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -75,6 +75,7 @@ class Config: ], "log_level": [level.lower() for level in _logger._core.levels], "softcore_form": ["zacharias", "taylor", "beutler"], + "precision": ["single", "mixed", "double"], } # A dictionary of nargs for the various options. @@ -132,12 +133,16 @@ def __init__( num_energy_neighbours=None, null_energy="1e6 kcal/mol", platform="auto", + precision="mixed", max_threads=None, max_gpus=None, max_sire_threads=None, opencl_platform_index=0, oversubscription_factor=1, replica_exchange=False, + max_contexts=None, + update_constraints=True, + constraint_lambda_index=0, randomise_velocities=False, perturbed_system=None, terminal_flip_frequency=None, @@ -353,6 +358,12 @@ def __init__( platform: str Platform to run simulation on. + precision: str + The floating point precision to use on GPU platforms. 'single' is fastest, + 'double' is slowest, and 'mixed' computes forces in single precision but + accumulates and integrates in double. Ignored by platforms that do not + support it, such as CPU. + max_threads: int Maximum number of CPU threads to use for simulation. (Default None, uses all available) Does nothing if platform is set to CUDA. @@ -377,15 +388,51 @@ def __init__( Whether to run replica exchange simulation. Currently this can only be used when GPU resources are available. + max_contexts: int + The maximum number of OpenMM contexts to create for a replica exchange + simulation. If None, then one context is created per replica, which is + fastest, but limits the number of replicas to those that fit in GPU memory. + If fewer contexts than replicas are requested, then each context is re-used + to propagate several replicas per cycle, changing its lambda value as it + goes. This lifts the memory limit at the cost of some performance. When + re-using contexts, 'frame_frequency' must equal 'checkpoint_frequency'. + + update_constraints: bool + Whether the constraints are updated when the lambda value of a context is + changed, i.e. whether constrained bond lengths are allowed to perturb with + lambda. This is only used when contexts are re-used across lambda values, + i.e. when 'max_contexts' is less than the number of replicas. Updating the + constraints is correct, but requires the OpenMM context to be reinitialised + whenever a constrained bond length actually changes, which is slow. Set this + to False if that overhead is significant; the constrained bond lengths are + then frozen at those of the lambda value given by + 'constraint_lambda_index'. Note that this is distinct from + 'dynamic_constraints', which controls where the constraint lengths are + taken from rather than whether they track lambda. + + constraint_lambda_index: int + The index of the lambda value at which to fix the constrained bond lengths + when 'update_constraints' is False. Every context is created at this lambda + value, so that the constraints are the same for all replicas rather than + depending on which context a replica is assigned to. The default of zero is + arbitrary but consistent; a lambda schedule that perturbs bonds away from + the end states may warrant a different choice. This is only used for + replica exchange simulations, and only when 'max_contexts' is less than the + number of replicas, 'update_constraints' is False, and a constrained bond + length actually perturbs with lambda. + randomise_velocities: bool Whether to randomise velocities at the start of each replica exchange cycle or following a terminal flip Monte Carlo move. perturbed_system: str The path to a stream file containing a Sire system for the equilibrated perturbed - end state (lambda = 1). This will be used as the starting conformation all lambda - windows > 0.5 when performing a replica exchange simulation. (Note that this assumes - that the "coordinates1" property specifies the coordinates for perturbable molecules.) + end state (lambda = 1). This is the same system as the input, but with the + "coordinates1" property of any perturbable molecules holding the equilibrated + coordinates for the lambda = 1 state. It is used as the starting conformation for + the lambda windows closest to the perturbed end state when performing a replica + exchange simulation, i.e. those with lambda > 0.5, or lambda < 0.5 when + 'swap_end_states' is True. terminal_flip_frequency: str Frequency at which to attempt terminal ring flip Monte Carlo moves. If None @@ -624,12 +671,16 @@ def __init__( self.checkpoint_frequency = checkpoint_frequency self.num_checkpoint_workers = num_checkpoint_workers self.platform = platform + self.precision = precision self.max_threads = max_threads self.max_gpus = max_gpus self.max_sire_threads = max_sire_threads self.opencl_platform_index = opencl_platform_index self.oversubscription_factor = oversubscription_factor self.replica_exchange = replica_exchange + self.max_contexts = max_contexts + self.update_constraints = update_constraints + self.constraint_lambda_index = constraint_lambda_index self.randomise_velocities = randomise_velocities self.perturbed_system = perturbed_system self.terminal_flip_frequency = terminal_flip_frequency @@ -1667,6 +1718,21 @@ def platform(self, platform): else: self._platform = "cpu" + @property + def precision(self): + return self._precision + + @precision.setter + def precision(self, precision): + if not isinstance(precision, str): + raise TypeError("'precision' must be of type 'str'") + precision = precision.lower().replace(" ", "") + if precision not in self._choices["precision"]: + raise ValueError( + f"'precision' not recognised. Valid options are: {', '.join(self._choices['precision'])}" + ) + self._precision = precision + @property def max_threads(self): return self._max_threads @@ -1778,6 +1844,55 @@ def replica_exchange(self, replica_exchange): raise ValueError("'replica_exchange' must be of type 'bool'") self._replica_exchange = replica_exchange + @property + def max_contexts(self): + return self._max_contexts + + @max_contexts.setter + def max_contexts(self, max_contexts): + if max_contexts is None or ( + isinstance(max_contexts, str) + and max_contexts.lower().replace(" ", "") == "none" + ): + self._max_contexts = None + return + + if not isinstance(max_contexts, int): + try: + max_contexts = int(max_contexts) + except Exception: + raise ValueError("'max_contexts' must be of type 'int'") + if max_contexts < 1: + raise ValueError("'max_contexts' must be greater than 0") + self._max_contexts = max_contexts + + @property + def update_constraints(self): + return self._update_constraints + + @update_constraints.setter + def update_constraints(self, update_constraints): + if not isinstance(update_constraints, bool): + raise ValueError("'update_constraints' must be of type 'bool'") + self._update_constraints = update_constraints + + @property + def constraint_lambda_index(self): + return self._constraint_lambda_index + + @constraint_lambda_index.setter + def constraint_lambda_index(self, constraint_lambda_index): + if not isinstance(constraint_lambda_index, int): + try: + constraint_lambda_index = int(constraint_lambda_index) + except Exception: + raise ValueError("'constraint_lambda_index' must be of type 'int'") + if constraint_lambda_index < 0: + raise ValueError( + "'constraint_lambda_index' must be greater than or equal to 0" + ) + self._constraint_lambda_index = constraint_lambda_index + @property def randomise_velocities(self): return self._randomise_velocities diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 158f176..044e3a9 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -624,8 +624,9 @@ def __init__(self, system, config): # Create a clone of the fully-prepared reference system with the # perturbed end-state coordinates and periodic space. This is done # after all system preparation so that the clone inherits the same - # topology and properties. It is used to seed starting coordinates - # for lambda > 0.5 replicas. + # topology and properties. It is used to seed starting coordinates for + # the replicas closest to the perturbed end state, i.e. lambda > 0.5, + # or lambda < 0.5 when the end states are swapped. if self._config.replica_exchange and self._config.perturbed_system is not None: from sire.legacy.IO import setCoordinates as _setCoordinates @@ -892,6 +893,7 @@ def __init__(self, system, config): "cutoff": self._config.cutoff, "cutoff_type": self._config.cutoff_type, "platform": self._config.platform, + "precision": self._config.precision, "rest2_selection": self._config.rest2_selection, "shift_coulomb": self._config.shift_coulomb, "shift_delta": self._config.shift_delta, @@ -1188,8 +1190,13 @@ def _check_end_state_constraints(self): """ Internal function to check whether the constraints are the same at the two end states. + + Sets self._end_state_constraints_differ, which records whether any + constrained bond length changes with lambda. """ + self._end_state_constraints_differ = False + # Find all perturbable molecules in the system.. pert_mols = self._system.molecules("property is_perturbable") @@ -1214,12 +1221,14 @@ def _check_end_state_constraints(self): # Check for equivalence. if len(constraints0) != len(constraints1): + self._end_state_constraints_differ = True _logger.info( f"Constraints are at not the same at {_lam_sym} = 0 and {_lam_sym} = 1." ) else: for c0, c1 in zip(constraints0, constraints1): if c0 != c1: + self._end_state_constraints_differ = True _logger.info( f"Constraints are at not the same at {_lam_sym} = 0 and {_lam_sym} = 1." ) @@ -1802,6 +1811,7 @@ def _compare_configs(config1, config2): "overwrite", "timeout", "oversubscription_factor", + "max_contexts", "restraint_search_time", "restraint_search_frequency", ] diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 57f8a32..fe6c4a6 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -54,6 +54,9 @@ def __init__( output_directory=None, perturbed_system=None, xml_filenames=None, + num_slots=None, + update_constraints=True, + constraint_lambda_index=None, gpu_devices=None, ): """ @@ -91,34 +94,65 @@ def __init__( 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. """ - # Warn if the number of replicas is not a multiple of the number of GPUs. - if len(lambdas) > num_gpus and len(lambdas) % num_gpus != 0: - _logger.warning( - "The number of replicas is not a multiple of the number of GPUs. " - "This may result in suboptimal performance." - ) + 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._states = _np.array(range(len(lambdas))) + self._states = _np.array(range(num_replicas)) self._time = None - self._openmm_states = [None] * len(lambdas) - self._gcmc_samplers = [None] * len(lambdas) - self._gcmc_states = [None] * len(lambdas) + self._openmm_states = [None] * num_replicas + self._gcmc_states = [None] * num_replicas # GCMC statistics for the whole simulation, keyed by lambda value. self._gcmc_stats = None - self._terminal_flip_stats = [[0, 0]] * len(lambdas) - self._num_proposed = _np.matrix(_np.zeros((len(lambdas), len(lambdas)))) - self._num_accepted = _np.matrix(_np.zeros((len(lambdas), len(lambdas)))) - self._num_swaps = _np.matrix(_np.zeros((len(lambdas), len(lambdas)))) + 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( @@ -133,10 +167,44 @@ def __init__( 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. + + The layout is derived from the number of replicas and slots, so it is + rebuilt rather than stored in a checkpoint. + """ + self._gcmc_samplers = [None] * self._num_slots + self._slot_replica = [None] * self._num_slots + + self._groups = [ + [int(r) for r in group] + for group in _np.array_split( + _np.arange(self._num_replicas), self._num_slots + ) + ] + + # 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) @@ -148,9 +216,38 @@ def __setstate__(self, state): if not hasattr(self, "_gcmc_states"): self._gcmc_states = [None] * n if not hasattr(self, "_terminal_flip_stats"): - self._terminal_flip_stats = [[0, 0]] * n + 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 + 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 + + # 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. + if is_legacy: + self._openmm_states = [self._openmm_states[s] for s in self._states] + self._gcmc_states = [self._gcmc_states[s] for s in self._states] + + # 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. @@ -171,6 +268,9 @@ def __getstate__(self): 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, "_openmm_states": self._openmm_states, @@ -252,21 +352,51 @@ def _create_dynamics( # Per-device memory tracking for estimation. device_mem = {} - # Work out how many replicas are assigned to each device. - # Replicas are assigned round-robin, so the first (num_replicas % num_gpus) - # devices get one extra replica. - base = floor(num_replicas / num_gpus) - remainder = num_replicas % num_gpus + # Work out how many slots are assigned to each device. + # Slots are assigned round-robin, so the first (num_slots % num_gpus) + # devices get one extra slot. + base = floor(self._num_slots / num_gpus) + remainder = self._num_slots % num_gpus contexts_per_device = [ base + (1 if i < remainder else 0) for i in range(num_gpus) ] - # Create the dynamics objects in serial. - for i, (lam, scale) in enumerate(zip(lambdas, rest2_scale_factors)): + # 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 = i % num_gpus - # Record baseline memory before the first replica on this device. + # 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) @@ -277,12 +407,19 @@ def _create_dynamics( "count": 0, } - # This is a restart, get the system for this replica. + # 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[i] - # This is a new simulation. For lambda > 0.5, use the perturbed - # system to seed the starting coordinates and periodic space. - elif perturbed_system is not None and lam > 0.5: + 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 @@ -292,7 +429,7 @@ def _create_dynamics( # Overload the device and lambda value. dynamics_kwargs["device"] = device - dynamics_kwargs["lambda_value"] = lam + dynamics_kwargs["lambda_value"] = build_lam dynamics_kwargs["rest2_scale"] = scale if gcmc_kwargs is not None: @@ -302,15 +439,17 @@ def _create_dynamics( msg = "loch is not installed. GCMC sampling cannot be performed." _logger.error(msg) - ghost_file = str(output_directory / f"gcmc_ghosts_{lam:.5f}.txt") - - # Create the GCMC sampler. + # 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, - ghost_file=ghost_file, + 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, ) @@ -328,7 +467,19 @@ def _create_dynamics( try: dynamics = mols.dynamics(**dynamics_kwargs) except Exception as e: - msg = f"Could not create dynamics object for lambda {lam:.5f} on device {device}: {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 @@ -341,12 +492,13 @@ def _create_dynamics( # Append the dynamics object. self._dynamics.append(dynamics) - # Write the OpenMM XML file to the output directory. + # 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[i]) + dynamics.to_xml(xml_filenames[seed]) # Track memory footprint for this device. info = device_mem[device] @@ -394,12 +546,25 @@ def _create_dynamics( 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"Available memory: {total_mem / (1024**3):.2f} GB. " + f"{self._max_contexts_advice(per_device * num_gpus)}" ) _logger.error(msg) raise MemoryError(msg) @@ -424,65 +589,120 @@ def _create_dynamics( f"Created dynamics object for lambda {lam:.5f} on device {device}" ) - def get(self, index): + # 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): """ - Get the dynamics object (and GCMC sampler) for a given index. + Return the index of the slot that hosts a given replica. Parameters ---------- - index: int + replica: int The index of the replica. Returns ------- + int + The index of the slot. + """ + return self._replica_slot[replica] + + 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 replica and its GCMC sampler. + The dynamics object for the slot and its GCMC sampler. """ - return self._dynamics[index], self._gcmc_samplers[index] + return self._dynamics[slot], self._gcmc_samplers[slot] - def set(self, index, dynamics): + def set(self, slot, dynamics): """ - Set the dynamics object for a given index. + Set the dynamics object for a given slot. Parameters ---------- - index: int - The index of the replica. + slot: int + The index of the slot. dynamics: sire.legacy.Convert.SOMMContext The dynamics object. """ - self._dynamics[index] = dynamics + self._dynamics[slot] = dynamics - def delete(self, index): + def delete(self, slot): """ - Delete the dynamics object for a given index. + Delete the dynamics object for a given slot. Parameters ---------- - index: int - The index of the replica. + slot: int + The index of the slot. """ - self._dynamics[index] = None + self._dynamics[slot] = None - def save_openmm_state(self, index): + def save_openmm_state(self, slot, replica): """ - Save the state of the dynamics object. + Save the state of a slot's dynamics object as the state of a replica. Parameters ---------- - index: int - The index of the replica. + 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[index] + self._dynamics[slot] .context() .getState(getPositions=True, getVelocities=True) ) @@ -490,12 +710,36 @@ def save_openmm_state(self, index): # 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[index] = { + 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): """ @@ -511,32 +755,55 @@ def _apply_openmm_state(context, 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"]) - if state["box"] is not None: - context.setPeriodicBoxVectors(*state["box"]) else: # Legacy openmm.State from checkpoint files written before this # format change. context.setState(state) - def save_gcmc_state(self, index): + def save_gcmc_state(self, slot, replica): """ - Save the current GCMC water state for the replica. + Save the current GCMC water state of a slot as that of a replica. Parameters ---------- - index: int - The index of the replica. + 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[index] + gcmc_sampler = self._gcmc_samplers[slot] # Store the state. - self._gcmc_states[index] = gcmc_sampler.water_state() + self._gcmc_states[replica] = gcmc_sampler.water_state() + + def get_clock(self): + """ + Get the simulation clock. + + Every replica advances by the same amount each cycle, so the clock is + common to all of them and is read from the first slot. + + Returns + ------- + + dict + The clock, as returned by Dynamics._get_clock(). + """ + return self._dynamics[0]._get_clock() def get_states(self): """ @@ -562,66 +829,167 @@ def set_states(self, states): """ self._states = states - def _seed_replica(self, i): + def load_replica(self, replica, clock=None): """ - Apply the (possibly new, post-mix) state to replica i's context, - including any GCMC water-state swap. Only touches replica i's own - context/sampler, so this is safe to run concurrently with other - replicas' calls from a thread pool. + 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. """ - state = self._states[i] + slot = self._replica_slot[replica] + dynamics = self._dynamics[slot] + gcmc_sampler = self._gcmc_samplers[slot] - _logger.debug(f"Replica {i} seeded from state {state}") - self._apply_openmm_state( - self._dynamics[i].context(), self._openmm_states[state] - ) + # 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]) - # Swap the water state in the GCMCSamplers. - if self._gcmc_samplers[i] is not None: - # Find the indices of the water states that differ. - water_idxs = _np.where(self._gcmc_states[i] != self._gcmc_states[state])[0] + # Positions have changed underneath the context, so any cached + # energies are stale. + dynamics.clear_energy_cache() - # Update the water state in the GCMCSampler. - self._gcmc_samplers[i].push() + 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) + + 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_samplers[i]._set_water_state( - self._dynamics[i].context(), - indices=water_idxs, - states=self._gcmc_states[state][water_idxs], + self._gcmc_num_waters[replica] = gcmc_sampler.num_waters( + context=self._dynamics[slot].context() ) finally: - self._gcmc_samplers[i].pop() + gcmc_sampler.pop() - def mix_states(self, old_states, executor=None): + def mix_states(self, old_states): """ - Mix the states of the dynamics objects. + 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. Parameters ---------- + old_states : numpy.ndarray The state indices from before the last replica mix. + """ + # 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)) + ] - executor : concurrent.futures.ThreadPoolExecutor, optional - Executor used to apply the per-replica state changes (an - OpenMM setPositions/setVelocities/setPeriodicBoxVectors call - per changed replica, each against a different context) in - parallel. Each replica's context is independent of every - other, so this is safe. Falls back to a serial loop if not - provided. - """ - # Replicas whose state actually changed. - changed = [i for i, state in enumerate(self._states) if i != state] - - if executor is not None and len(changed) > 1: - # Consume the map so we block until every replica is seeded - # and any exception raised in a worker thread propagates here. - list(executor.map(self._seed_replica, changed)) - else: - for i in changed: - self._seed_replica(i) - - # Update the swap matrix. Cheap, CPU-only bookkeeping - kept out of - # the parallel section above to avoid any shared-array races. + # Update the swap matrix. for i, state in enumerate(self._states): self._num_swaps[old_states[i], state] += 1 @@ -669,6 +1037,31 @@ def _physical_device(self, 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): """ @@ -836,6 +1229,12 @@ def __init__(self, system, config): # devices relative to the visible set, so index i is gpu_devices[i]. self._gpu_devices = list(gpu_devices)[: self._num_gpus] + # Work out how many OpenMM contexts (slots) to create. When there are + # fewer slots than replicas, each slot is re-used to propagate several + # replicas per cycle, changing its lambda value as it goes. + self._num_replicas = len(self._lambda_values) + self._set_num_slots() + # Auto-generate a Boresch restraint for ABFE runs with no user-supplied # restraint. This must happen before the dynamics cache is built below, # since the per-replica OpenMM contexts it creates are fixed at @@ -870,7 +1269,11 @@ def __init__(self, system, config): perturbable_constraint = self._config.perturbable_constraint # Don't use constraints during minimisation. - if self._config.minimise and not self._config.minimisation_constraints: + if ( + self._config.minimise + and not self._is_restart + and not self._config.minimisation_constraints + ): constraint = "none" perturbable_constraint = "none" @@ -923,6 +1326,9 @@ def __init__(self, system, config): perturbed_system=self._perturbed_system, output_directory=self._config.output_directory, xml_filenames=xml_filenames, + num_slots=self._num_slots, + update_constraints=self._config.update_constraints, + constraint_lambda_index=self._constraint_lambda_index, gpu_devices=self._gpu_devices, ) @@ -998,6 +1404,16 @@ def __init__(self, system, config): # 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._build_slot_layout() + # Create the dynamics objects. self._dynamics_cache._create_dynamics( self._system, @@ -1009,46 +1425,32 @@ def __init__(self, system, config): output_directory=self._config.output_directory, ) - # Reset the state of the OpenMM contexts and GCMC samplers. - for i in range(len(self._lambda_values)): - dynamics, gcmc_sampler = self._dynamics_cache.get(i) - - # Reset the OpenMM state, applying the last replica exchange - # mixing so the correct post-mix state is restored. - state = self._dynamics_cache._states[i] - DynamicsCache._apply_openmm_state( - dynamics.context(), self._dynamics_cache._openmm_states[state] - ) - - # Reset the GCMC water state and restore statistics. - if gcmc_sampler is not None: - gcmc_sampler.push() - try: - gcmc_sampler._set_water_state( - dynamics.context(), - states=self._dynamics_cache._gcmc_states[state], - force=True, - ) - finally: - gcmc_sampler.pop() - - # Samplers keep only the lambda values they visit, so it's - # safe to hand each of them the whole simulation's stats. - if self._dynamics_cache._gcmc_stats is not None: - gcmc_sampler.restore_stats(self._dynamics_cache._gcmc_stats) - - # Log the GCMC sphere centre for each replica using the actual context - # positions (accurate for both fresh runs and restarts). + # 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. import openmm.unit as _omm_unit for i, lam in enumerate(self._lambda_values): - dynamics, gcmc_sampler = self._dynamics_cache.get(i) + _, 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: - state = dynamics.context().getState(getPositions=True) - positions = state.getPositions(asNumpy=True).value_in_unit( - _omm_unit.angstrom + positions = DynamicsCache._get_positions( + self._dynamics_cache._openmm_states[i] + ) + target = gcmc_sampler._get_target_position( + positions.value_in_unit(_omm_unit.angstrom) ) - target = gcmc_sampler._get_target_position(positions) _logger.info( f"Initial GCMC sphere centre for lambda {lam:.5f}: " f"[{target[0]:.3f}, {target[1]:.3f}, {target[2]:.3f}] A" @@ -1109,6 +1511,209 @@ def __init__(self, system, config): # Create a lock to guard the dynamics cache. self._lock = Lock() + # Systems committed while a replica was resident in its slot, awaiting + # the end of cycle checkpoint. Only used when contexts are shared, + # keyed by replica index and emptied by _checkpoint(). + self._committed = {} + + def _build_lambda(self, replica): + """ + Return the lambda value to create a replica's context at. + + This is the replica's own lambda value, unless the constraints are + fixed at a common one, in which case a rebuilt context has to use that + too or it would pick up the constrained bond lengths of its own lambda + value instead. + + Parameters + ---------- + + replica: int + The index of the replica. + + Returns + ------- + + float + The lambda value to create the context at. + """ + if self._constraint_lambda_index is None: + return self._lambda_values[replica] + + return self._lambda_values[self._constraint_lambda_index] + + def _replica_passes(self, cycle): + """ + Work out which replicas to propagate in each pass of a cycle. + + A slot can only host one replica at a time, so each pass takes at most + one replica from each slot. Groups are traversed in alternating + directions on successive cycles, so that a slot always moves to a + neighbouring lambda window, including across the cycle boundary. That + keeps the change in force field parameters as small as possible. + + When there is one slot per replica there is a single pass containing + every replica, which is the same as propagating them all together. + + Parameters + ---------- + + cycle: int + The index of the current cycle. + + Returns + ------- + + list of list of int + The replicas to propagate in each pass. + """ + groups = self._dynamics_cache._groups + + # Traverse the groups backwards on odd cycles. + is_reversed = cycle % 2 == 1 + + passes = [] + for i in range(max(len(group) for group in groups)): + batch = [] + for group in groups: + if i < len(group): + batch.append(group[len(group) - 1 - i] if is_reversed else group[i]) + passes.append(batch) + + return passes + + def _safe_batches(self, num_workers, cycle=0): + """ + Yield batches of replicas that can be processed concurrently. + + Two replicas that share a slot must never be processed at the same + time, since they would be using the same dynamics object and OpenMM + context. Batches are therefore taken from within a single pass, which + holds at most one replica per slot, and then split by the number of + workers. + + With one slot per replica there is a single pass containing every + replica, so this is just a split by the number of workers. + + Parameters + ---------- + + num_workers: int + The maximum number of replicas in a batch. + + cycle: int + The index of the current cycle, which sets the traversal order. + + Yields + ------ + + list of int + A batch of replicas that is safe to process concurrently. + """ + from math import ceil + + for batch in self._replica_passes(cycle): + for i in range(ceil(len(batch) / num_workers)): + yield batch[i * num_workers : (i + 1) * num_workers] + + def _set_num_slots(self): + """ + Work out the number of OpenMM contexts (slots) to create, validating + the configuration options that only apply when contexts are re-used + across lambda values. + + Sets self._num_slots and self._is_cached. + """ + + num_replicas = self._num_replicas + + if self._config.max_contexts is None: + self._num_slots = num_replicas + else: + self._num_slots = min(self._config.max_contexts, num_replicas) + + # There is a context per replica, so nothing is re-used and all of the + # constraints below are irrelevant. A context then keeps the lambda + # value it was created at, so there is no need to fix the constraints + # at a common one. + self._is_cached = self._num_slots < num_replicas + self._constraint_lambda_index = None + + # Contexts run concurrently, so it's the number of them that determines + # how evenly the GPUs are loaded, not the number of replicas. + num_workers = self._num_gpus * self._config.oversubscription_factor + + if self._num_slots < num_workers: + _logger.warning( + f"The number of contexts ({self._num_slots}) is less than the " + f"number of workers ({num_workers}). Some GPUs will be left idle." + ) + elif self._num_slots % self._num_gpus != 0: + _logger.warning( + f"The number of contexts ({self._num_slots}) is not a multiple " + f"of the number of GPUs ({self._num_gpus}). This may result in " + "suboptimal performance." + ) + + if not self._is_cached: + if self._config.max_contexts is not None: + _logger.info( + f"Creating one OpenMM context per replica ({num_replicas})" + ) + return + + # Frames can only be saved on checkpoint cycles when contexts are + # re-used. Within a cycle a context propagates several replicas in + # turn, so frames from different replicas would otherwise accumulate + # in the same internal trajectory. Tying frames to checkpoints means + # each one is written out and cleared before the context is handed to + # the next replica. + if ( + self._save_frames + and self._config.frame_frequency != self._config.checkpoint_frequency + ): + msg = ( + "'frame_frequency' must equal 'checkpoint_frequency' when " + "'max_contexts' is less than the number of replicas." + ) + _logger.error(msg) + raise ValueError(msg) + + # When the constraints aren't updated as a slot changes lambda, they stay + # as they were when its context was created. Create every context at the + # same lambda value, so that the constrained bond lengths are uniform + # across replicas rather than depending on which slot a replica happens + # to be assigned to. Only needed if they actually perturb. + if not self._config.update_constraints and self._end_state_constraints_differ: + # Which lambda value is used matters less than every replica using + # the same one, since where the bonds actually perturb depends on + # the lambda schedule. + if self._config.constraint_lambda_index >= num_replicas: + msg = ( + f"'constraint_lambda_index' " + f"({self._config.constraint_lambda_index}) is out of range " + f"for {num_replicas} {_lam_sym} values." + ) + _logger.error(msg) + raise ValueError(msg) + + self._constraint_lambda_index = self._config.constraint_lambda_index + _logger.warning( + f"'update_constraints' is False. Constrained bond lengths will not " + f"perturb with lambda, and are fixed at those of " + f"{_lam_sym} = " + f"{self._lambda_values[self._constraint_lambda_index]:.5f} " + f"for every replica." + ) + + from math import ceil + + _logger.info( + f"Re-using {self._num_slots} OpenMM context(s) across " + f"{num_replicas} replicas: {ceil(num_replicas / self._num_slots)} " + "pass(es) per cycle" + ) + def __str__(self): """Return a string representation of the object.""" return f"RepexRunner(system={self._system}, config={self._config})" @@ -1185,7 +1790,6 @@ def run(self): ) # Work out the required number of batches. - num_batches = ceil(self._config.num_lambda / num_workers) num_checkpoint_batches = ceil(self._config.num_lambda / num_checkpoint_workers) # Persistent thread pools, reused across every batch and cycle. @@ -1195,14 +1799,14 @@ def run(self): # Create the replica list. replica_list = list(range(self._config.num_lambda)) - # Minimise at each lambda value. - if self._config.minimise: - for i in range(num_batches): + # Minimise at each lambda value. Don't minimise on restart. + if self._config.minimise and not self._is_restart: + for batch in self._safe_batches(num_workers): with ThreadPoolExecutor(max_workers=num_workers) as executor: try: for success, index, e in executor.map( self._minimise, - replica_list[i * num_workers : (i + 1) * num_workers], + batch, ): if not success: msg = f"Minimisation failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" @@ -1217,12 +1821,12 @@ def run(self): # Equilibrate the system. if self._is_equilibration and not self._is_restart: - for i in range(num_batches): + for batch in self._safe_batches(num_workers): with ThreadPoolExecutor(max_workers=num_workers) as executor: try: for success, index, e in executor.map( self._equilibrate, - replica_list[i * num_workers : (i + 1) * num_workers], + batch, ): if not success: _logger.error( @@ -1236,6 +1840,28 @@ def run(self): # Write a checkpoint immediately after equilibration so that a restart # after an early production crash doesn't need to re-equilibrate. if self._is_equilibration and not self._is_restart: + # When contexts are shared, commit each replica while it is + # resident so that every checkpoint file can still be written + # under a single lock below. + if self._is_cached: + for batch in self._safe_batches(num_checkpoint_workers): + try: + for index, error in checkpoint_executor.map( + self._load_and_commit, + batch, + ): + if error is not None: + msg = ( + f"Post-equilibration commit failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}:\n{error}" + ) + _logger.error(msg) + raise error + except KeyboardInterrupt: + checkpoint_executor.shutdown(wait=False, cancel_futures=True) + _logger.error("Post-equilibration commit cancelled. Exiting.") + _sys.exit(1) + lock = _FileLock(self._lock_file) with lock.acquire(timeout=self._config.timeout.to("seconds")): for j in range(num_checkpoint_batches): @@ -1346,33 +1972,80 @@ def run(self): else None ) - # Run a dynamics block for each replica, making sure only each GPU is only - # oversubscribed by a factor of self._config.oversubscription_factor. - for j in range(num_batches): - replicas = replica_list[j * num_workers : (j + 1) * num_workers] - try: - for result, index, energies in dynamics_executor.map( - self._run_block, - replicas, - repeat(self._lambda_values), - repeat(is_gcmc), - repeat(write_gcmc_ghosts), - repeat(is_terminal_flip), - repeat(time_ns), - ): - if not result: - _logger.error( - f"Dynamics failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {energies}" + # Whether the checkpoint files are written at the end of this cycle. + do_checkpoint = is_checkpoint or i == cycles - 1 + + # Capture the simulation clock at the start of the cycle. Every + # replica advances by the same amount each cycle, so a slot has to + # be rewound to this point before it propagates the next replica. + 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): + # Run a dynamics block for each replica in the pass, making sure + # each GPU is only oversubscribed by a factor of + # self._config.oversubscription_factor. + for j in range(ceil(len(batch) / num_workers)): + replicas = batch[j * num_workers : (j + 1) * num_workers] + try: + for result, index, energies in dynamics_executor.map( + self._run_block, + replicas, + repeat(self._lambda_values), + repeat(is_gcmc), + repeat(write_gcmc_ghosts), + repeat(is_terminal_flip), + repeat(time_ns), + repeat(clock), + ): + 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: + for j in range(ceil(len(batch) / num_checkpoint_workers)): + replicas = batch[ + j * num_checkpoint_workers : (j + 1) + * num_checkpoint_workers + ] + try: + for index, error in checkpoint_executor.map( + self._commit_replica, + replicas, + ): + if error: + _logger.error( + f"Commit failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}: {error}" + ) + raise error + except KeyboardInterrupt: + checkpoint_executor.shutdown( + wait=False, cancel_futures=True ) - raise energies - results.append((index, energies)) - except KeyboardInterrupt: - dynamics_executor.shutdown(wait=False, cancel_futures=True) - _logger.error("Dynamics cancelled. Exiting.") - _sys.exit(1) - - # Checkpoint. - if is_checkpoint or i == cycles - 1: + _logger.error("Commit cancelled. Exiting.") + _sys.exit(1) + + # 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: # Create the lock. lock = _FileLock(self._lock_file) @@ -1391,7 +2064,7 @@ def run(self): self._backup_checkpoint, replicas, ): - if not result: + if error: _logger.error( f"Backup failed for {_lam_sym} = " f"{self._lambda_values[index]:.5f}: {error}" @@ -1448,16 +2121,11 @@ def run(self): self._dynamics_cache.get_accepted(), ) ) - self._dynamics_cache.mix_states(old_states, executor=dynamics_executor) - - # Snapshot the pre-run state for crash recovery. - if self._config.auto_fix_minimise: - for i, state in enumerate(self._dynamics_cache.get_states()): - self._dynamics_cache._dynamics[i]._d._pre_run_state = ( - self._dynamics_cache._dynamics[i] - .context() - .getState(getPositions=True, getVelocities=True) - ) + + # 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(old_states) # This is a checkpoint cycle. if is_checkpoint: @@ -1548,10 +2216,16 @@ def _run_block( write_gcmc_ghosts=False, is_terminal_flip=False, time_ns=None, + clock=None, ): """ Run a dynamics block for a given replica. + The replica is made resident in its slot, propagated, then stored back + out again so that the slot can be handed to the next replica. When + there is one slot per replica, loading and storing reduce to the state + bookkeeping that replica exchange does anyway. + Parameters ---------- @@ -1579,6 +2253,12 @@ def _run_block( The current simulation time in nanoseconds, used when saving energy components. If None, energy components are not saved. + clock: dict + The simulation clock at the start of the cycle, as returned by + Dynamics._get_clock(). Every replica advances by the same amount + each cycle, so a slot must be rewound to the start of the cycle + before propagating the next replica. + Returns ------- @@ -1597,13 +2277,14 @@ def _run_block( lam = lambdas[index] try: + # Make the replica resident in its slot. + self._dynamics_cache.load_replica(index, clock=clock) + # Get the dynamics object (and GCMC sampler). - dynamics, gcmc_sampler = self._dynamics_cache.get(index) + dynamics, gcmc_sampler = self._dynamics_cache.get( + self._dynamics_cache.slot_for(index) + ) - # Track whether any MC move changed the context positions so we - # can update _pre_run_state once at the end. Only needed when - # crash recovery is enabled. - needs_pre_run_snapshot = False auto_fix_minimise = self._config.auto_fix_minimise # Perform the GCMC move before dynamics so that the energies @@ -1617,9 +2298,6 @@ def _run_block( finally: gcmc_sampler.pop() - if auto_fix_minimise: - needs_pre_run_snapshot = True - # Write ghost residues immediately after the GCMC move so the # ghost state and frame (saved during dynamics) are consistent. if write_gcmc_ghosts: @@ -1628,13 +2306,12 @@ def _run_block( # Perform a terminal flip move before dynamics if requested. if self._terminal_flip_samplers is not None and is_terminal_flip: _logger.info(f"Performing terminal flip move at {_lam_sym} = {lam:.5f}") - if self._terminal_flip_samplers[index].move(dynamics.context()): - if auto_fix_minimise: - needs_pre_run_snapshot = True + self._terminal_flip_samplers[index].move(dynamics.context()) - # Snapshot the context state for crash recovery if any MC move - # changed positions. - if needs_pre_run_snapshot: + # Snapshot the context state for crash recovery. The slot was just + # seeded with this replica's state, and any MC move above may have + # changed it again, so this is always required. + if auto_fix_minimise: dynamics._d._pre_run_state = dynamics.context().getState( getPositions=True, getVelocities=True ) @@ -1668,12 +2345,9 @@ def _run_block( ), ) - if gcmc_sampler is not None: - # Save the GCMC state. - self._dynamics_cache.save_gcmc_state(index) - - # Save the OpenMM state. - self._dynamics_cache.save_openmm_state(index) + # Save the replica's state back out of the slot, so that the slot + # can be handed to the next replica. + self._dynamics_cache.store_replica(index) # Save the energy contribution for each force. if self._config.save_energy_components and time_ns is not None: @@ -1717,8 +2391,12 @@ def _minimise(self, index): _logger.info(f"Minimising at {_lam_sym} = {self._lambda_values[index]:.5f}") try: + # Make the replica resident in its slot. + self._dynamics_cache.load_replica(index) + # Get the dynamics object (and GCMC sampler). - dynamics, gcmc_sampler = self._dynamics_cache.get(index) + slot = self._dynamics_cache.slot_for(index) + dynamics, gcmc_sampler = self._dynamics_cache.get(slot) if gcmc_sampler is not None and not self._is_restart: gcmc_sampler.push() @@ -1749,17 +2427,17 @@ def _minimise(self, index): system = dynamics.commit() # Delete the dynamics object. - self._dynamics_cache.delete(index) + self._dynamics_cache.delete(slot) # Work out the device index. - device = index % self._num_gpus + device = slot % self._num_gpus # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() # Overload the device and lambda value. dynamics_kwargs["device"] = device - dynamics_kwargs["lambda_value"] = self._lambda_values[index] + dynamics_kwargs["lambda_value"] = self._build_lambda(index) dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index] # Create the production dynamics object. @@ -1771,13 +2449,20 @@ def _minimise(self, index): if gcmc_sampler is not None: self._reset_gcmc_sampler(gcmc_sampler, dynamics) - # Set the new dynamics object. - self._dynamics_cache.set(index, dynamics) + # Set the new dynamics object. The rebuilt object has a + # fresh, empty energy trajectory, but no energies are + # recorded during minimisation or equilibration, and + # load_replica() re-attaches the replica's trajectory + # before the first production block. + self._dynamics_cache.set(slot, dynamics) _logger.info( f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}" ) + # Save the minimised state back out of the slot. + self._dynamics_cache.store_replica(index) + except Exception as e: return False, index, e @@ -1808,8 +2493,12 @@ def _equilibrate(self, index): _logger.info(f"Equilibrating at {_lam_sym} = {self._lambda_values[index]:.5f}") try: + # Make the replica resident in its slot. + self._dynamics_cache.load_replica(index) + # Get the dynamics object (and GCMC sampler). - dynamics, gcmc_sampler = self._dynamics_cache.get(index) + slot = self._dynamics_cache.slot_for(index) + dynamics, gcmc_sampler = self._dynamics_cache.get(slot) if gcmc_sampler is not None: gcmc_sampler.push() @@ -1848,17 +2537,17 @@ def _equilibrate(self, index): system = dynamics.commit() # Delete the current dynamics object. - self._dynamics_cache.delete(index) + self._dynamics_cache.delete(slot) # Work out the device index. - device = index % self._num_gpus + device = slot % self._num_gpus # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() # Overload the device and lambda value. dynamics_kwargs["device"] = device - dynamics_kwargs["lambda_value"] = self._lambda_values[index] + dynamics_kwargs["lambda_value"] = self._build_lambda(index) dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index] dynamics_kwargs["timestep"] = self._config._equilibration_timestep dynamics_kwargs["constraint"] = constraint @@ -1872,7 +2561,7 @@ def _equilibrate(self, index): self._reset_gcmc_sampler(gcmc_sampler, dynamics) # Update the dynamics object in the cache. - self._dynamics_cache.set(index, dynamics) + self._dynamics_cache.set(slot, dynamics) # Equilibrate. dynamics.run( @@ -1894,17 +2583,17 @@ def _equilibrate(self, index): system.set_time(_sr.u("0ps")) # Delete the dynamics object. - self._dynamics_cache.delete(index) + self._dynamics_cache.delete(slot) # Work out the device index. - device = index % self._num_gpus + device = slot % self._num_gpus # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() # Overload the device and lambda value. dynamics_kwargs["device"] = device - dynamics_kwargs["lambda_value"] = self._lambda_values[index] + dynamics_kwargs["lambda_value"] = self._build_lambda(index) dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index] # Create the production dynamics object. @@ -1916,21 +2605,16 @@ def _equilibrate(self, index): if gcmc_sampler is not None: self._reset_gcmc_sampler(gcmc_sampler, dynamics) - # Compute the current number of waters in the GCMC sampling - # volume after equilibration. - gcmc_sampler.push() - try: - gcmc_sampler.num_waters(context=dynamics.context()) - finally: - gcmc_sampler.pop() - # Set the new dynamics object. - self._dynamics_cache.set(index, dynamics) + self._dynamics_cache.set(slot, dynamics) _logger.info( f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}" ) + # Save the equilibrated state back out of the slot. + self._dynamics_cache.store_replica(index) + except Exception as e: return False, index, e @@ -1966,7 +2650,7 @@ def _compute_energies(self, index): ) # Get the dynamics object. - dynamics, _ = self._dynamics_cache.get(index) + 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) @@ -2044,6 +2728,80 @@ def _write_checkpoint_system(self, system, index, context=None, gcmc_sampler=Non """ self._dynamics_cache._time = system.time() + def _load_and_commit(self, index): + """ + Make a replica resident in its slot and commit it, for the + post-equilibration checkpoint. + + Equilibration stores every replica's state back out of its slot, so a + replica has to be loaded again before it can be committed. + + Parameters + ---------- + + index: int + The index of the replica. + + Returns + ------- + + index: int + The index of the replica. + + exception: Exception + The exception if the commit failed. + """ + try: + self._dynamics_cache.load_replica(index) + except Exception as e: + return index, e + + return self._commit_replica(index) + + def _commit_replica(self, index): + """ + Commit a replica while it is still resident in its slot, holding the + result until the checkpoint files are written at the end of the cycle. + + This is only used when contexts are shared between replicas. A slot + holds the positions and trajectory frames of the replica it last + hosted, so a replica has to be committed before its slot is handed on. + Deferring only the file writes keeps every checkpoint file within a + single lock, so that an external reader never sees a partial set. + + Parameters + ---------- + + index: int + The index of the replica. + + Returns + ------- + + index: int + The index of the replica. + + exception: Exception + The exception if the commit failed. + """ + try: + slot = self._dynamics_cache.slot_for(index) + dynamics, _ = self._dynamics_cache.get(slot) + + # commit() returns a clone, so the frames can be cleared straight + # away, ready for the next replica to use the slot. + system = dynamics.commit() + speed = dynamics.time_speed() + dynamics._d._sire_mols.delete_all_frames() + + with self._lock: + self._committed[index] = (system, speed) + + except Exception as e: + return index, e + + return index, None + def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): """ Checkpoint the simulation. @@ -2080,13 +2838,23 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): lam = lambdas[index] # Get the dynamics object (and GCMC sampler). - dynamics, gcmc_sampler = self._dynamics_cache.get(index) + slot = self._dynamics_cache.slot_for(index) + dynamics, gcmc_sampler = self._dynamics_cache.get(slot) - # Commit the current system. - system = dynamics.commit() + # Use the system committed by _commit_replica() if there is one. + # When contexts are shared the replica is no longer resident in its + # slot by the time the cycle finishes, so it was committed earlier, + # while it still was. + committed = self._committed.pop(index, None) - # Get the simulation speed. - speed = dynamics.time_speed() + if committed is None: + # Commit the current system. + system = dynamics.commit() + + # Get the simulation speed. + speed = dynamics.time_speed() + else: + system, speed = committed # Call the base class checkpoint method to save the system state. with self._lock: @@ -2098,7 +2866,8 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): return index, error # Delete all trajectory frames from the Sire system within the - # dynamics object. + # dynamics object. This is a no-op when the replica was committed + # earlier, since the frames were cleared then to free the slot. dynamics._d._sire_mols.delete_all_frames() if block == -1: @@ -2111,23 +2880,24 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): f"for {_lam_sym} = {lam:.5f}" ) - # 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 "" - ) - _logger.info( - f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} " - f"is {gcmc_sampler.num_waters()}{acc_str}" - ) - finally: - gcmc_sampler.pop() + # Log the number of waters within the GCMC sampling volume. Both + # 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] + 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 + acc_str = ( + f", acceptance rate = {stats['num_accepted'] / n_moves:.3f}" + f" (ins = {stats['num_insertions']}, del = {stats['num_deletions']})" + if n_moves > 0 + else "" + ) + _logger.info( + f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} " + f"is {num_waters}{acc_str}" + ) # Log terminal flip acceptance rate for this replica. if self._terminal_flip_samplers is not None: @@ -2263,8 +3033,8 @@ def _merge_gcmc_stats(self): """ stats = {} - for i in range(len(self._lambda_values)): - _, gcmc_sampler = self._dynamics_cache.get(i) + 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()) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index fb0dac7..0017b44 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -560,8 +560,8 @@ def generate_lam_vals(lambda_base, increment=0.001): else: terminal_flip_sampler = None - # Minimisation. - if self._config.minimise: + # Minimisation. Don't minimise on restart. + if self._config.minimise and not is_restart: constraint = self._config.constraint perturbable_constraint = self._config.perturbable_constraint @@ -584,6 +584,7 @@ def generate_lam_vals(lambda_base, increment=0.001): 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}" @@ -638,6 +639,15 @@ def generate_lam_vals(lambda_base, increment=0.001): # 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}" ) @@ -730,12 +740,9 @@ def generate_lam_vals(lambda_base, increment=0.001): import openmm.unit as _omm_unit _npz_state = _np.load(self._filenames[index]["checkpoint_state"]) - dynamics.context().setPositions( - _npz_state["positions"] * _omm_unit.nanometer - ) - dynamics.context().setVelocities( - _npz_state["velocities"] * _omm_unit.nanometer / _omm_unit.picosecond - ) + + # Set the box before the positions, since a barostat may have + # changed it between the state being saved and restored. if "box" in _npz_state: from openmm import Vec3 as _Vec3 @@ -746,6 +753,13 @@ def generate_lam_vals(lambda_base, increment=0.001): _Vec3(*_box[2]) * _omm_unit.nanometer, ) + dynamics.context().setPositions( + _npz_state["positions"] * _omm_unit.nanometer + ) + dynamics.context().setVelocities( + _npz_state["velocities"] * _omm_unit.nanometer / _omm_unit.picosecond + ) + # Reset the GCMC sampler. This resets the sampling statistics and clears # the associated OpenMM forces. if gcmc_sampler is not None: @@ -804,9 +818,11 @@ def generate_lam_vals(lambda_base, increment=0.001): finally: gcmc_sampler.pop() - # Otherwise, if we've performed equilibration, then we need to reset - # the water state in the new context to match the equilibrated system. - elif is_equilibrated: + # Otherwise, reset the water state in the new context to match the + # system that the preparation stages left behind. Both minimisation + # and equilibration perform GCMC moves, so the state held by the + # sampler is ahead of the one the context was built with. + else: # Reset the water state. gcmc_sampler.push() try: @@ -814,7 +830,6 @@ def generate_lam_vals(lambda_base, increment=0.001): dynamics.context(), force=True, ) - gcmc_sampler.num_waters(context=dynamics.context()) finally: gcmc_sampler.pop() @@ -1152,9 +1167,14 @@ def generate_lam_vals(lambda_base, increment=0.001): 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"is {gcmc_sampler.num_waters()}{acc_str}" + f"is {num_waters}{acc_str}" ) finally: gcmc_sampler.pop() @@ -1529,6 +1549,7 @@ def _minimisation( device=None, constraint="none", perturbable_constraint="none", + gcmc_sampler=None, ): """ Minimise a system. @@ -1554,6 +1575,10 @@ def _minimisation( perturbable_constraint: str The constraint for perturbable molecules. + gcmc_sampler: :class: `GCMCSampler ` + A GCMC sampler to pre-equilibrate the water with before minimising. + If None, then no GCMC moves are performed. + Returns ------- @@ -1581,6 +1606,24 @@ def _minimisation( # Create a dynamics object. dynamics = system.dynamics(**dynamics_kwargs) + # Pre-equilibrate the water before minimising, so that a dry pocket + # is filled before the geometry relaxes into it. The context is + # created from the sampler's own system, so its water state already + # matches and only needs binding. + if gcmc_sampler is not None: + gcmc_sampler.bind_dynamics(dynamics) + + _logger.info( + f"Pre-equilibrating with GCMC moves at {_lam_sym} = {lambda_value:.5f}" + ) + + gcmc_sampler.push() + try: + for i in range(100): + gcmc_sampler.move(dynamics.context()) + finally: + gcmc_sampler.pop() + # Run the minimisation. dynamics.minimise(timeout=self._config.timeout) diff --git a/tests/runner/test_gcmc.py b/tests/runner/test_gcmc.py new file mode 100644 index 0000000..c952966 --- /dev/null +++ b/tests/runner/test_gcmc.py @@ -0,0 +1,63 @@ +""" +Tests for GCMC sampling via the regular (non replica exchange) runner. +""" + +import pytest +import re +import tempfile + +from pathlib import Path + +from somd2.config import Config +from somd2.runner import Runner + +from tests.conftest import has_cuda + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_runner_gcmc_without_a_selection(ethane_methanol): + """ + Validate GCMC sampling with no 'gcmc_selection', where moves are attempted + within the entire simulation box rather than a region around a selection. + + The sampler then has no reference, so it cannot count the waters within a + region. Reporting the water count has to account for that, which is what + this exercises: counting it raised before there was a path for the + reference-free case. + """ + pytest.importorskip("loch") + + with tempfile.TemporaryDirectory() as tmpdir: + config = Config( + runtime="8fs", + output_directory=tmpdir, + energy_frequency="4fs", + checkpoint_frequency="4fs", + frame_frequency="4fs", + platform="cuda", + max_threads=1, + num_lambda=2, + gcmc=True, + gcmc_frequency="4fs", + ) + + # The bulk-only path is the point of this test, so make sure it can't + # stop being exercised without the test failing. + assert config.gcmc_selection is None + + runner = Runner(ethane_methanol, config) + runner.run() + + # GCMC ran, so the ghost residues were written. + for lam in runner._lambda_values: + assert (Path(tmpdir) / f"gcmc_ghosts_{lam:.5f}.txt").exists() + + # With no region the count is the number of non-ghost waters in the + # box, which is never zero for a solvated system. + log = (Path(tmpdir) / config.log_file).read_text() + counts = [ + int(x) + for x in re.findall(r"number of waters in GCMC volume.*? is (\d+)", log) + ] + assert counts, "no water count was logged" + assert all(count > 0 for count in counts), f"zero water count logged: {counts}" diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index e7b970d..d0a7bea 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -150,6 +150,462 @@ def test_rest2_selection(ethane_methanol, rest2_selection, is_valid): runner = RunnerBase(ethane_methanol, Config(**config)) +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +@pytest.mark.parametrize("max_contexts", [1, 2, 3, 4]) +def test_repex_bounded_contexts(ethane_methanol, max_contexts): + """ + Validate that a replica exchange simulation runs when there are fewer + OpenMM contexts than replicas, so that each context is re-used to + propagate several replicas per cycle. + """ + num_lambda = 4 + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "12fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": max_contexts, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + # Only the requested number of contexts should have been created. + assert len(runner._dynamics_cache._dynamics) == max_contexts + + # Every replica must be assigned to exactly one slot. + groups = runner._dynamics_cache._groups + assert sorted(r for group in groups for r in group) == list(range(num_lambda)) + + runner.run() + + # Output is per replica, regardless of how many contexts were used. + assert (Path(tmpdir) / "repex_matrix.txt").exists() + for i in range(num_lambda): + assert Path(runner._filenames[i]["energy_traj"]).exists() + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_frame_frequency_constraint(ethane_methanol): + """ + Validate that frames can only be saved on checkpoint cycles when contexts + are re-used across lambda values. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "12fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "8fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": 4, + "replica_exchange": True, + "max_contexts": 2, + } + + with pytest.raises(ValueError, match="frame_frequency"): + RepexRunner(ethane_methanol, Config(**config)) + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +@pytest.mark.parametrize("update_constraints", [True, False]) +def test_repex_update_constraints(ethane_methanol, update_constraints): + """ + Validate both constraint modes. Ethane to methanol does perturb a + constrained bond length, so update_constraints=True forces the context to + be reinitialised on every lambda change. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "12fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": 4, + "replica_exchange": True, + "max_contexts": 2, + "update_constraints": update_constraints, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + runner.run() + + assert (Path(tmpdir) / "repex_matrix.txt").exists() + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_bounded_contexts_output_equivalence(ethane_methanol): + """ + Validate that re-using contexts produces the same output structure as one + context per replica. Energies are not compared: a shared context consumes + the integrator's random number stream in a different order, so the + trajectories legitimately differ. + """ + import pandas as pd + + num_lambda = 4 + + def run(max_contexts, tmpdir): + config = { + "runtime": "16fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": max_contexts, + } + runner = RepexRunner(ethane_methanol, Config(**config)) + runner.run() + return [ + pd.read_parquet(runner._filenames[i]["energy_traj"]) + for i in range(num_lambda) + ], list(runner._lambda_values) + + with tempfile.TemporaryDirectory() as tmpdir: + full, lambdas = run(num_lambda, tmpdir) + + with tempfile.TemporaryDirectory() as tmpdir: + cached, _ = run(1, tmpdir) + + # One record per energy_frequency interval, starting at the first. Pinning + # the values rather than only comparing the two runs catches a clock that + # is wrong the same way in both. + expected_times = [0.004, 0.008, 0.012, 0.016] + + for i in range(num_lambda): + assert len(cached[i]) == len(full[i]) + assert list(cached[i].columns) == list(full[i].columns) + assert cached[i].index.equals(full[i].index) + + for records in (full[i], cached[i]): + times = [round(t, 6) for t in records.index.get_level_values(0)] + assert times == expected_times, f"replica {i} recorded times {times}" + + # Each replica records energies at every lambda value. + assert list(records.columns) == lambdas + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_bounded_contexts_restart(ethane_methanol): + """ + Validate that a replica exchange simulation using fewer contexts than + replicas can be restarted, that each replica resumes from the state it + stopped at, and that the energy trajectory is extended rather than + restarted. + """ + import pandas as pd + + num_lambda = 4 + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "8fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": 2, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + runner.run() + + num_rows = [ + len(pd.read_parquet(runner._filenames[i]["energy_traj"])) + for i in range(num_lambda) + ] + + # The state each replica finished at, which the checkpoint holds. + import openmm.unit as omm_unit + + stopped = [ + { + "positions": state["positions"].value_in_unit(omm_unit.nanometer), + "velocities": state["velocities"].value_in_unit( + omm_unit.nanometer / omm_unit.picosecond + ), + "box": state["box"].value_in_unit(omm_unit.nanometer), + } + for state in runner._dynamics_cache._openmm_states + ] + + # Restart, extending the runtime. + config["runtime"] = "16fs" + config["restart"] = True + + runner = RepexRunner(ethane_methanol, Config(**config)) + + # Every replica must resume from where it stopped. The contexts are + # created from the input system, so the only thing carrying the + # simulated state across a restart is the checkpoint. + for i in range(num_lambda): + state = runner._dynamics_cache._openmm_states[i] + for key, unit in ( + ("positions", omm_unit.nanometer), + ("velocities", omm_unit.nanometer / omm_unit.picosecond), + ("box", omm_unit.nanometer), + ): + assert np.allclose( + state[key].value_in_unit(unit), stopped[i][key], atol=1e-6 + ), f"replica {i} {key} not restored" + + # The input coordinates must not be what was restored, otherwise the + # checks above would pass even if the checkpoint were ignored. + import sire as sr + + inputs = sr.io.get_coords_array(runner._system) + restored = runner._dynamics_cache._openmm_states[0]["positions"].value_in_unit( + omm_unit.angstrom + ) + assert not np.allclose(restored, inputs, atol=1e-3) + + # Restoring the checkpoint into the cache is not enough: the contexts + # are created from the input system, so the state has to reach them + # too. Loading a replica is what pushes it. + for i in range(num_lambda): + runner._dynamics_cache.load_replica(i) + dynamics, _ = runner._dynamics_cache.get(runner._dynamics_cache.slot_for(i)) + positions = ( + dynamics.context() + .getState(getPositions=True) + .getPositions(asNumpy=True) + .value_in_unit(omm_unit.nanometer) + ) + assert np.allclose(positions, stopped[i]["positions"], atol=1e-5), ( + f"replica {i} positions not pushed into its context" + ) + + runner.run() + + # 8 fs at 4 fs intervals, extended to 16 fs. The clock has to continue + # from where it stopped rather than restarting at zero, so the records + # must run to the new runtime with no repeats or gaps. + expected_times = [0.004, 0.008, 0.012, 0.016] + + for i in range(num_lambda): + extended = pd.read_parquet(runner._filenames[i]["energy_traj"]) + assert len(extended) > num_rows[i] + + times = [round(t, 6) for t in extended.index.get_level_values(0)] + assert times == expected_times, f"replica {i} recorded times {times}" + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +@pytest.mark.parametrize("max_contexts", [1, 4]) +def test_repex_checkpoint_single_lock(ethane_methanol, max_contexts): + """ + Validate that every checkpoint file is written within a single acquisition + of the file lock, so that a process streaming the output off the machine + always sees a coherent set rather than a mixture of new and old files. + """ + import somd2.runner._repex as repex_module + + num_lambda = 4 + acquisitions = [] + + real_filelock = repex_module._FileLock + + class CountingFileLock(real_filelock): + def acquire(self, *args, **kwargs): + acquisitions.append(1) + return super().acquire(*args, **kwargs) + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "8fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": max_contexts, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + repex_module._FileLock = CountingFileLock + try: + runner.run() + finally: + repex_module._FileLock = real_filelock + + # Two cycles, each taking the lock once for the checkpoint files and once + # for the repex state, plus a final acquisition. This must not scale with + # the number of passes. + assert len(acquisitions) == 5 + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +@pytest.mark.parametrize("max_contexts", [1, 4]) +def test_repex_gcmc_bounded_contexts(ethane_methanol, max_contexts): + """ + Validate that GCMC sampling works when contexts are re-used across lambda + values, so a slot's single sampler is re-parameterised and re-pointed at + the ghost file of whichever replica it hosts. + """ + pytest.importorskip("loch") + + num_lambda = 4 + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "8fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": max_contexts, + "gcmc": True, + "gcmc_selection": "resname LIG", + "gcmc_frequency": "4fs", + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + runner.run() + + assert (Path(tmpdir) / "repex_matrix.txt").exists() + + # One ghost file per lambda, each with a line per saved frame. A slot + # writing to the wrong file would leave these unbalanced. + counts = [] + for lam in runner._lambda_values: + ghost_file = Path(tmpdir) / f"gcmc_ghosts_{lam:.5f}.txt" + assert ghost_file.exists() + counts.append(len(ghost_file.read_text().strip().splitlines())) + + assert len(set(counts)) == 1, f"unbalanced ghost files: {counts}" + assert counts[0] > 0 + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_gcmc_without_a_selection(ethane_methanol): + """ + Validate GCMC sampling with no 'gcmc_selection', where moves are attempted + within the entire simulation box rather than a region around a selection. + + The sampler then has no reference, so it cannot count the waters within a + region, and every move samples the whole box. Counting the waters on each + replica handover has to account for that. + """ + pytest.importorskip("loch") + + num_lambda = 4 + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "8fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": 2, + "gcmc": True, + "gcmc_frequency": "4fs", + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + # The bulk-only path is the point of this test, so make sure it can't + # stop being exercised without the test failing. + assert runner._dynamics_cache._gcmc_samplers[0]._reference is None + + runner.run() + + assert (Path(tmpdir) / "repex_matrix.txt").exists() + for lam in runner._lambda_values: + assert (Path(tmpdir) / f"gcmc_ghosts_{lam:.5f}.txt").exists() + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_concurrent_slots(ethane_methanol): + """ + Validate that replicas sharing a slot are never propagated concurrently. + Oversubscribing exercises this on a single GPU, since the worker count is + the number of GPUs times the oversubscription factor. + + This is also the only test that equilibrates, so it covers moving replicas + in and out of their slots during equilibration, the context rebuild when + the constraints change, and the post-equilibration checkpoint. + """ + num_lambda = 4 + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "12fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "equilibration_time": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": 2, + "oversubscription_factor": 2, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + # Guard against the equilibration coverage being lost silently. + assert runner._is_equilibration + + # Minimising without constraints and equilibrating with them means the + # contexts are rebuilt part way through, which is the path being + # covered here. + assert not runner._config.minimisation_constraints + assert runner._config.equilibration_constraints + + # Every batch must contain at most one replica per slot. + num_workers = runner._num_gpus * config["oversubscription_factor"] + for batch in runner._safe_batches(num_workers): + slots = [runner._dynamics_cache.slot_for(r) for r in batch] + assert len(slots) == len(set(slots)), f"batch {batch} shares a slot" + + runner.run() + + assert (Path(tmpdir) / "repex_matrix.txt").exists() + + @pytest.mark.parametrize( "gpu_devices, expected", [ @@ -169,12 +625,8 @@ def test_rest2_selection(ethane_methanol, rest2_selection, is_valid): def test_physical_device_mapping(gpu_devices, expected): """ Validate that an OpenMM device index is mapped to the physical device - backing it. - - OpenMM numbers devices relative to the visible set, whereas pynvml and - pyopencl enumerate every device on the machine. Querying the memory of a - device by its OpenMM index therefore reports the wrong GPU whenever the - visible set does not start at zero. + 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 @@ -184,6 +636,70 @@ def test_physical_device_mapping(gpu_devices, expected): assert [cache._physical_device(i) for i in range(len(expected))] == expected +def test_max_contexts_advice(): + """ + Validate the advice given when the replicas don't fit in GPU memory. It + 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 + + advice = DynamicsCache._max_contexts_advice(6) + assert "'max_contexts' to 6 or fewer" in advice + assert "frame_frequency" in advice and "checkpoint_frequency" in advice + + # Nothing fits, so there is no number to suggest. + advice = DynamicsCache._max_contexts_advice(0) + assert "max_contexts" not in advice + assert "does not fit" in advice + + +def test_gcmc_state_follows_replica(): + """ + Validate that the GCMC water occupancy travels with the configuration it + belongs to when replicas are mixed. + + A slot holds one GCMC sampler but hosts several replicas, and load_replica() + installs a replica's water state by diffing it against whatever the last + resident left in the sampler. If the occupancy did not follow the positions + through a mix, a replica would run with another replica's waters, which + gives plausible numbers rather than an obviously wrong output file. + + 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 + + num_replicas = 4 + + cache = object.__new__(DynamicsCache) + + # Label both states with the replica they came from, so that a replica + # holding mismatched positions and waters is detectable. + cache._openmm_states = list(range(num_replicas)) + cache._gcmc_states = list(range(num_replicas)) + cache._ghost_files = [f"ghosts_{i}.txt" for i in range(num_replicas)] + cache._state_moved = [False] * num_replicas + cache._num_swaps = np.zeros((num_replicas, num_replicas)) + + # Mix twice, since a slot is re-used within a cycle. + for states in ([2, 0, 3, 1], [1, 3, 0, 2]): + old_states = list(range(num_replicas)) + expected = [cache._gcmc_states[state] for state in states] + + cache._states = states + cache.mix_states(old_states) + + # The water occupancy follows the same permutation as the positions. + assert cache._gcmc_states == expected + + # Every replica holds the positions and waters of the same origin. + assert cache._openmm_states == cache._gcmc_states + + # The ghost files stay with the lambda window. + assert cache._ghost_files == [f"ghosts_{i}.txt" for i in range(num_replicas)] + + @pytest.mark.parametrize( "device, key, value", [ @@ -204,8 +720,7 @@ def test_check_device_memory_queries_requested_device(monkeypatch, device, key, pynvml = pytest.importorskip("pynvml") - # Force the OpenCL branch to fail so that the pynvml path is always taken, - # regardless of what the machine running the tests has installed. + # Force the OpenCL branch to fail so the pynvml path is always taken. broken = types.SimpleNamespace() def get_platforms(): @@ -235,3 +750,275 @@ def by_uuid(uuid): assert DynamicsCache._check_device_memory(device) == (1, 2, 3) assert requested == {key: value} + + +def test_legacy_checkpoint_restore(): + """ + Validate that a checkpoint written before slots existed can still be + loaded. These stored the replica states unpermuted, with the states array + holding the mapping to apply on restart. + """ + from somd2.runner._repex import DynamicsCache + + n = 4 + legacy = { + "_lambdas": [0.0, 0.33, 0.67, 1.0], + "_rest2_scale_factors": [1.0] * n, + "_states": np.array([2, 0, 1, 3]), + "_time": None, + "_openmm_states": [f"state{i}" for i in range(n)], + "_gcmc_samplers": [None] * n, + "_gcmc_states": [f"water{i}" for i in range(n)], + "_gcmc_stats": [None] * n, + "_terminal_flip_stats": [[0, 0]] * n, + "_num_proposed": np.zeros((n, n)), + "_num_accepted": np.zeros((n, n)), + "_num_swaps": np.zeros((n, n)), + } + + cache = object.__new__(DynamicsCache) + cache.__setstate__(dict(legacy)) + + # Converted to the current convention: each replica's own state, with the + # last mix applied. + assert cache._openmm_states == ["state2", "state0", "state1", "state3"] + assert cache._gcmc_states == ["water2", "water0", "water1", "water3"] + + # Every replica is seeded from its stored state on a restart. + assert cache._state_moved == [True] * n + + # Attributes postdating the checkpoint are defaulted, one slot per replica. + assert cache._num_slots == n + assert cache._groups == [[0], [1], [2], [3]] + assert cache._energy_trajectories == [None] * n + assert cache._ghost_files == [None] * n + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_max_contexts_change_on_restart(ethane_methanol): + """ + Validate that the number of contexts can change on restart. The slot + layout is rebuilt from the configuration rather than restored. + """ + num_lambda = 4 + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "8fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": 2, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + assert len(runner._dynamics_cache._dynamics) == 2 + runner.run() + + # Restart with a different number of contexts. + config["runtime"] = "16fs" + config["restart"] = True + config["max_contexts"] = 4 + + runner = RepexRunner(ethane_methanol, Config(**config)) + assert len(runner._dynamics_cache._dynamics) == 4 + assert runner._dynamics_cache._groups == [[0], [1], [2], [3]] + runner.run() + + assert (Path(tmpdir) / "repex_matrix.txt").exists() + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_gcmc_lambda_cache_warm(ethane_methanol, monkeypatch): + """ + Validate that a GCMC sampler builds exactly one OpenMM context, scanning + it over every lambda it will host, and none once running. A mismatch + between the cached lambdas and those passed to set_lambda would show up + here as an extra build. + """ + loch = pytest.importorskip("loch") + + num_lambda = 4 + calls = [] + + real_precompute = loch.GCMCSampler._precompute_lambdas + + def counting_precompute(self, lambda_values, rest2_scales): + # Only record calls with work to do. Deduplicated, since the caller + # may name the same lambda twice but the scan extracts it once. + missing = sorted( + { + (float(lam), float(scale)) + for lam, scale in zip(lambda_values, rest2_scales) + if (float(lam), float(scale)) not in self._lambda_params + } + ) + if missing: + calls.append(missing) + return real_precompute(self, lambda_values, rest2_scales) + + monkeypatch.setattr(loch.GCMCSampler, "_precompute_lambdas", counting_precompute) + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "8fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": 1, + "gcmc": True, + "gcmc_selection": "resname LIG", + "gcmc_frequency": "4fs", + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + # A single context build, covering every lambda value in one scan. + assert len(calls) == 1 + assert sorted(lam for lam, _ in calls[0]) == sorted(runner._lambda_values) + + runner.run() + + # No further context may be built once the simulation is running. + assert len(calls) == 1 + + # An uncached lambda still works, building and caching on demand. + _, sampler = runner._dynamics_cache.get(0) + uncached = 0.123456 + assert uncached not in runner._lambda_values + + sampler.push() + try: + sampler.set_lambda(uncached) + finally: + sampler.pop() + + assert len(calls) == 2 + assert (uncached, sampler._rest2_scale) in sampler._lambda_params + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +@pytest.mark.parametrize("swap_end_states", [False, True]) +def test_repex_perturbed_system_seeding(ethane_methanol, swap_end_states): + """ + Validate that the end states are seeded from the right coordinates when + contexts are shared. + + A context is created from a single system and every replica it hosts starts + from that context, so the end state a replica starts from is chosen from + the middle of the group rather than its first replica. That keeps any + mismatch next to the lambda value at which the end state switches, instead + of it depending on where the groups happen to fall. + + Swapping the end states reverses the lambda schedule, so the perturbed end + state moves to lambda = 0 and the seeding must mirror with it. + + Only the perturbable molecule is displaced, so the assertions also cover + the property path that its coordinates travel along: read from + 'coordinates1' via link_to_perturbed, written to 'coordinates0', then read + back via link_to_reference. + """ + import sire as sr + + # A perturbed end state, with the perturbable molecule displaced so that + # its coordinates are distinct. + perturbed = ethane_methanol.clone() + perturbed.set_property("space", ethane_methanol.property("space")) + coords = sr.io.get_coords_array(ethane_methanol) + from sire.legacy.IO import setCoordinates + + # Flag the atoms of the perturbable molecules. get_coords_array returns the + # atoms in molecule order, so the offset tracks the array index. + is_perturbable = np.zeros(len(coords), dtype=bool) + offset = 0 + for mol in ethane_methanol.molecules(): + num_atoms = mol.num_atoms() + if mol.has_property("is_perturbable"): + is_perturbable[offset : offset + num_atoms] = True + offset += num_atoms + assert is_perturbable.any(), "no perturbable molecules in the test system" + + coords[is_perturbable] += 1.0 + + # Write to 'coordinates1', which is the property that the runner reads via + # link_to_perturbed. + perturbed = sr.system.System( + setCoordinates(perturbed._system, coords.tolist(), True) + ) + + # Ten replicas across three contexts, a layout in which the switch falls + # inside a group. + num_lambda = 10 + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "4fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": num_lambda, + "replica_exchange": True, + "max_contexts": 3, + "perturbed_system": perturbed, + "swap_end_states": swap_end_states, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + reference = sr.io.get_coords_array(runner._system) + target = sr.io.get_coords_array(runner._perturbed_system) + + import openmm.unit as omm_unit + + seeded = [] + for i in range(num_lambda): + positions = runner._dynamics_cache._openmm_states[i][ + "positions" + ].value_in_unit(omm_unit.angstrom) + from_reference = np.allclose(positions, reference, atol=1e-3) + from_target = np.allclose(positions, target, atol=1e-3) + assert from_reference != from_target, f"replica {i} matches neither" + seeded.append("perturbed" if from_target else "reference") + + # The end states themselves must always be right. Swapping the end + # states puts the perturbed one at lambda = 0. + if swap_end_states: + assert seeded[0] == "perturbed" + assert seeded[-1] == "reference" + else: + assert seeded[0] == "reference" + assert seeded[-1] == "perturbed" + + # Both systems must be used, otherwise the option does nothing. + assert set(seeded) == {"reference", "perturbed"} + + # Only the group containing the switch can be seeded from the wrong end + # state, and then for no more than half of it. Choosing the end state + # from the first replica of a group rather than its middle breaks this. + lambdas = runner._lambda_values + for group in runner._dynamics_cache._groups: + wrong = [ + i + for i in group + if (seeded[i] == "perturbed") != ((lambdas[i] > 0.5) != swap_end_states) + ] + assert len(wrong) <= len(group) // 2, ( + f"group {group} has {len(wrong)} replicas seeded from the " + f"wrong end state: {wrong}" + )