diff --git a/CHANGELOG.md b/CHANGELOG.md index 1df7d6f..dab67f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,31 @@ Changelog ========= +[2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 +-------------------------------------------------------------------------------------- + +* Add support for generating Boresch restraints for absolute binding free energy calculations [#166](https://github.com/OpenBioSim/somd2/pull/166). +* Give alchemical ions their own plain morph lambda schedule so they interpolate correctly under non-standard lambda schedules [#169](https://github.com/OpenBioSim/somd2/pull/169). +* Persist alchemical ion identity across restarts so the same molecule is reused regardless of GCMC state [#172](https://github.com/OpenBioSim/somd2/pull/172). +* Use perisistent `ThreadPoolExector` objects within the main replica exchange dynamics block [#175](https://github.com/OpenBioSim/somd2/pull/175). +* Allow `oversubscription_factor` to change on restart [#177](https://github.com/OpenBioSim/somd2/pull/177). +* Restrict energy component decomposition to force groups that are used for integration [#180](https://github.com/OpenBioSim/somd2/pull/180). +* Parallelise replica mixing [#181](https://github.com/OpenBioSim/somd2/pull/181). +* 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 `single` [#191](https://github.com/OpenBioSim/somd2/pull/191). +* Add support for generating Morse restraints for ring-breaking perturbations [#194](https://github.com/OpenBioSim/somd2/pull/194). +* Remove the unused `kappa` lever equations from the ring-breaking/making lambda schedules [#195](https://github.com/OpenBioSim/somd2/pull/195). +* Accept stream file paths for the `restraints` and `lambda_schedule` configuration options, so they can be set from the command line [#198](https://github.com/OpenBioSim/somd2/pull/198). +* Account for off-site charges (virtual sites) when computing the charge difference between the end states. They are held as a molecule property rather than on the atoms, so a charge-preserving perturbation could appear to change charge and be given spurious alchemical ions [#200](https://github.com/OpenBioSim/somd2/pull/200). +* Handle `num_lambda=1`, which previously raised a `ZeroDivisionError` when generating the lambda values. The `RepexRunner` now rejects a single lambda window, since there is nothing to exchange with and the regular `Runner` is faster [#203](https://github.com/OpenBioSim/somd2/pull/203). +* Detect the available GPUs once in the base runner and re-use the list, rather than the `RepexRunner` querying `CUDA_VISIBLE_DEVICES` regardless of the chosen platform. Replica exchange is now also permitted on the HIP platform [#206](https://github.com/OpenBioSim/somd2/pull/206). +* Query the free memory of AMD GPUs with `CL_DEVICE_GLOBAL_FREE_MEMORY_AMD` rather than `CL_DEVICE_BOARD_NAME_AMD`, which returns the device name [#206](https://github.com/OpenBioSim/somd2/pull/206). + [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 6ec3b90..0f9cf96 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@

- SOMD + SOMD2

@@ -11,7 +11,25 @@ [![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://www.gnu.org/licenses/gpl-3.0) Open-source GPU accelerated molecular dynamics engine for alchemical free-energy -simulations. Built on top of [Sire](https://github.com/OpenBioSim/sire) and [OpenMM](https://github.com/openmm/openmm). The code is still under active development and is not yet ready for general use. +simulations. Built on top of [Sire](https://github.com/OpenBioSim/sire) and [OpenMM](https://github.com/openmm/openmm). + +## Features + +- **Perturbations**: relative binding free energies, + [absolute binding free energies](#absolute-binding-free-energies), + [ring-breaking](#ring-breaking-perturbations), + [charge-change](#charge-change-perturbations), and protein mutations. +- **[GCMC](#gcmc)**: grand canonical Monte Carlo water sampling. +- **[Replica exchange](#replica-exchange)**: Hamiltonian replica exchange + between lambda windows. +- **[REST2](#rest2)**: replica exchange with solute scaling. +- **[Terminal ring flips](#terminal-ring-flip-monte-carlo)**: Monte Carlo moves + to improve sampling of terminal aromatic rings. +- **[Ghost atom modifications](#ghost-atom-modifications)**: modification of + ghost atom bonded terms to avoid spurious coupling to the physical system. +- **[Multiple GPUs](#running-somd2-using-one-or-more-gpus)**: lambda windows are + distributed across the available devices, with optional + [oversubscription](#gpu-oversubscription). ## Installation @@ -74,7 +92,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 @@ -93,6 +111,31 @@ You should now have a `somd2` executable in your path. To test, run: somd2 --help ``` +### Keeping up to date + +During a development cycle the OpenBioSim packages are pinned only to a +`YYYY.N.0.dev` version, not to a specific build. `somd2` and its dependencies +therefore need to be kept in sync, so always update the whole stack together +rather than `somd2` alone. + +For a conda install, update everything in one go: + +``` +conda update -c conda-forge -c openbiosim/label/dev sire biosimspace ghostly loch somd2 +``` + +For a standalone pixi install, pull the latest `somd2` and refresh the +pre-built dependencies: + +``` +git pull +pixi update +``` + +For a full source install, `git pull` in *every* repository you have installed +(`sire`, `biosimspace`, `ghostly`, `loch` and `somd2`), not just `somd2`. Since +`sire` is compiled, you will also need to rebuild it. + ## Development Pre-commit hooks are used to ensure consistent code formatting and linting. @@ -114,10 +157,10 @@ pre-commit run --all-files ## Usage In order to run an alchemical free-energy simulation you will need to -first create a stream file containing the _perturbable_ system of interest. +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,48 +172,119 @@ 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 +A larger collection of input files and end-to-end tutorials, covering everything +from a simple charge-change validation system to full case studies, can be found +in the [somd2_examples](https://github.com/OpenBioSim/somd2_examples) repository. + +### 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). + +This is required in all cases, since SOMD2 enumerates the devices to run on from +the environment variable itself. -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 -and set. In the case that this detection fails, or if there are multiple platforms -available, the `--platform` option can be set (for example `--platform cuda`). +By default SOMD2 uses `--platform auto`, which selects the first platform +registered by OpenMM in order of preference: CUDA, OpenCL, HIP, Metal, Reference, +then CPU. If detection fails, or if you want a specific platform, use the +`--platform` option (for example `--platform cuda`). -By default, `SOMD2` will automatically manage the distribution of lambda windows +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 -`CUDA_VISIBLE_DEVICES` are set as above would restrict `SOMD2` to using only +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. +## Restarting + +A simulation can be continued from the files in its output directory using the +`--restart` option: + +``` +somd2 perturbable_system.bss --restart --output-directory output +``` + +Each λ window (or replica) resumes from its most recent checkpoint. The +configuration used for the original run is written to `config.yaml` in the +output directory, controlled by `--write-config`, which is enabled by default. +This file is required in order to restart, since the current configuration is +validated against it. + +Only a limited set of options may be changed on restart. Broadly, anything that +would change the perturbation or the Hamiltonian is fixed, whereas options +controlling how long to run for, what to write out, and which hardware to use +can be varied. The most useful of these is `--runtime`, which allows a completed +simulation to be extended. SOMD2 will tell you which option is at fault if you +change one that isn't allowed. + +> [!NOTE] +> If the most recent checkpoint files are incomplete or corrupt, for example +> when recovering from a crash, pass `--use-backup` to restart from the last +> but one checkpoint instead. + +## Hydrogen mass repartitioning + +By default SOMD2 applies hydrogen mass repartitioning (HMR), scaling hydrogen +masses by the factor given by `--h-mass-factor` (default 1.5). This is what +allows the default `--timestep` of 4 fs. + +If the masses of your input system have already been repartitioned, or you want +to use a different repartitioning scheme, pass `--no-hmr` so that the masses of +the input system are used as they are. + +> [!NOTE] +> A 4 fs timestep is not stable without repartitioning, so if you disable HMR +> you will need to reduce `--timestep` accordingly, or supply a system that has +> already been repartitioned. + ## 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, +SOMD2 supports Hamiltonian replica exchange (HREX) simulations, which can be +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,10 +293,10 @@ 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 +comprises *all* atoms in perturbable molecules, but can be controlled via the `--rest2-selection` option. This should be a `Sire` selection string that specifies additional atoms of interest, i.e. those in regular, non-perturbable molecules. If the selection does contain atoms within perturbable molecules, then only @@ -223,7 +337,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 @@ -242,9 +356,131 @@ geometry. To override this for all groups: somd2 perturbable_system.bss --terminal-flip-frequency "1 ps" --terminal-flip-angle "180 degrees" ``` +## Lambda schedules + +The way that the perturbation is applied across the lambda coordinate is +controlled by the `--lambda-schedule` option, which defaults to +`standard_morph`, which is intended for use with relative binding free +energy (RBFE) simulations. The available schedules are: + +| Schedule | Description | +| --- | --- | +| `standard_morph` | Linear interpolation between the two end states. | +| `charge_scaled_morph` | As above, but with charges scaled at intermediate lambda values. | +| `annihilate` | Absolute binding free energies, removing all non-bonded interactions. | +| `decouple` | Absolute binding free energies, removing only intermolecular interactions. | +| `ring_break_morph` | Ring-breaking perturbations. | +| `reverse_ring_break_morph` | Ring-making perturbations, i.e. the reverse of the above. | + +For the `annihilate`, `decouple`, and ring-breaking schedules, appropriate +restraints can be generated automatically. See the sections below. + +## Absolute binding free energies + +Absolute binding free energy (ABFE) calculations are supported using the +`annihilate` and `decouple` lambda schedules. Both first discharge the ligand, +then remove its Lennard-Jones interactions: `annihilate` removes all non-bonded +interactions, including those within the ligand, whereas `decouple` retains the +intramolecular terms. + +``` +somd2 perturbable_system.bss --lambda-schedule decouple +``` + +The ligand must be restrained within the binding site. If no restraints are +passed, a Boresch restraint is generated automatically for the bound leg, i.e. +when the system contains both a protein and water. This is done by minimising +the system, running a short trajectory at lambda = 0, then choosing the anchor +atoms and force constants from it. The length of this trajectory and the +frequency at which frames are saved can be controlled with the +`--restraint-search-time` and `--restraint-search-frequency` options. By +default the receptor anchor atoms are chosen from the protein backbone; use +`--restraint-search-receptor-selection` to pass a `Sire` selection string +instead. + +The restraint is written to `abfe_restraint.s3` in the output directory and is +reloaded on restart, since the accumulated free energy corresponds to that +particular restraint. The standard state correction is logged and written to +the metadata of the energy trajectory, so analysis code can apply it without +needing to scan the log. + +> [!NOTE] +> The Beutler soft-core form, enabled with `--softcore-form beutler`, is only +> supported with the ABFE schedules, or a custom schedule. + +## Ring-breaking perturbations + +Perturbations that break (or form) a ring are supported using the +`ring_break_morph` schedule, or `reverse_ring_break_morph` for the ring-making +direction. + +``` +somd2 perturbable_system.bss --lambda-schedule ring_break_morph +``` + +These perturbations require a pair of Morse restraints on the atoms of the bond +that is broken. If no restraints are passed, both are generated automatically. +A "hard" Morse potential replaces the harmonic bond, inheriting its force +constant and equilibrium length, and is switched off as a weaker "soft" Morse +restraint holds the fragment in place. Their well depths and the force constant +of the soft restraint can be controlled with the `--morse-hard-well-depth`, +`--morse-soft-well-depth`, and `--morse-soft-force-constant` options. + +Unlike the ABFE restraints, these are regenerated on each run rather than being +cached, since they are derived from the bond parameters alone and are therefore +identical every time. + +> [!NOTE] +> The defaults are a reasonable starting point, but ring-breaking +> perturbations are demanding. A non-uniform spacing of lambda values, set with +> `--lambda-values`, is typically needed to obtain good overlap around the point +> at which the bond is broken. The +> [alchemate](https://github.com/akalpokas/alchemate) package provides +> workflows for iteratively optimising the lambda schedule. + +## Charge-change perturbations + +Perturbations that change the net charge of the system are handled +automatically using the co-alchemical ion method. The charge difference between +the two end states is computed when the system is loaded, and, if it is +non-zero, a number of water molecules equal to the absolute charge difference +are perturbed into counter-ions alongside the main perturbation, keeping the +total charge constant at every lambda value. The waters furthest from the +perturbable molecule are chosen, and the ion type is picked to offset the +charge change, re-using the parameters of a free ion already present in the +system where possible. + +No options are needed to enable this. The automatically detected value can be +overridden with `--charge-difference`, which takes the perturbed charge minus +the reference charge: + +``` +somd2 perturbable_system.bss --charge-difference -1 +``` + +The molecules chosen as alchemical ions are written to `alchemical_ions.npz` in +the output directory and reused on restart, so that ion selection does not +depend on anything that might have changed between runs. + +Since a co-alchemical ion is only meaningful in the bulk, SOMD2 can restrain it +away from the perturbable region. Passing a distance to +`--coalchemical-restraint-dist` adds an inverse-distance restraint between each +ion and the atom closest to the centre of geometry of the perturbable molecule, +preventing the ion from drifting into the binding site and interacting with the +protein or ligand: + +``` +somd2 perturbable_system.bss --coalchemical-restraint-dist "10 A" +``` + +> [!NOTE] +> These restraints are *added* to any others in use. Restraints passed via the +> Python API, and those generated automatically for the ABFE and ring-breaking +> schedules described above, are all retained. + ## Debugging with energy components -To help diagnose simulation instabilities, `SOMD2` can record the potential +To help diagnose simulation instabilities, SOMD2 can record the potential energy contribution from each OpenMM force group. This is enabled with the `--save-energy-components` flag: @@ -271,7 +507,7 @@ The recording interval depends on the runner and active samplers: ## Copying output files during a simulation -When `SOMD2` writes checkpoint files it acquires an exclusive +When SOMD2 writes checkpoint files it acquires an exclusive [file lock](https://py-filelock.readthedocs.io) on `somd2.lock` inside the output directory. This guarantees that checkpoint files are always in a consistent state on disk. @@ -297,7 +533,7 @@ with FileLock("/path/to/output/somd2.lock"): ``` > [!NOTE] -> The `--timeout` option (default: `300 s`) controls how long `SOMD2` will +> The `--timeout` option (default: `300 s`) controls how long SOMD2 will > wait to re-acquire the lock after your copy completes. If you hold the lock > for longer than this, the simulation will raise a `Timeout` error. @@ -330,7 +566,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 @@ -344,23 +580,19 @@ be chosen as a trade off between accuracy and computational cost. A value of aro 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) +These are enabled by default, but can be disabled using the `--no-ghost-modifications` +option. Modifications are implemented using the [ghostly](https://github.com/OpenBioSim/ghostly) package. ## Note for SOMD1 users -For existing users of `somd1`, it's possible to generate input for `somd2` by passing -`--somd2 True` to the `prepareFEP.py` setup script. This will write a `somd2` compatible -stream file. - -Additionally, `somd2` can be run in `somd1` _compatibility_ mode by passing the -``--somd1-compatibility`` command-line option to the `somd2` executable. This ensures -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`. +SOMD2 can be run in SOMD1 *compatibility* mode by passing the +`--somd1-compatibility` command-line option to the `somd2` executable. This ensures +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 +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 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`.) @@ -376,7 +608,7 @@ BSS.Stream.save(system, "somd1") (This will write a stream file called `somd1.bss`.) -This can then be run with `somd2` using the following: +This can then be run with SOMD2 using the following: ``` somd2 somd1.bss --pert-file somd1.pert --somd1-compatibility @@ -385,13 +617,13 @@ somd2 somd1.bss --pert-file somd1.pert --somd1-compatibility (This only shows the limited options required. Others will take default values and can be set accordingly.) If you want to load an existing system from a perturbation file and use the -new `somd2` [ghost atom bonded-term modifications](https://github.com/OpenBioSim/ghostly), +new SOMD2 [ghost atom bonded-term modifications](https://github.com/OpenBioSim/ghostly), 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,13 +635,31 @@ 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 -`SOMD2` can also be used as a Python API, allowing it to be embedded +SOMD2 can also be used as a Python API, allowing it to be embedded within other Python scripts. +A few options take objects rather than values, so cannot be set directly on the +command line. A custom lambda schedule can be passed to `lambda_schedule` as a +`sire.cas.LambdaSchedule`, rather than one of the named schedules, and +user-defined restraints can be passed to `restraints`. + +Both options can also be set via a YAML configuration file, where they are +stored as a hex string of the serialised object. This is the form written to +`config.yaml`, so the simplest way to obtain one is to configure the option in +Python, run a simulation, and re-use the value from the resulting file. + +Alternatively, both accept a path to a [Sire](https://github.com/OpenBioSim/sire) +stream file containing the serialised object, which can be written with +`sire.stream.save`: + +``` +somd2 perturbable_system.bss --lambda-schedule my_schedule.s3 --restraints my_restraints.s3 +``` + ## Known issues If using the regular `Runner` class via the Python API, then you will need to @@ -423,3 +673,7 @@ some systems, causing the simulation to hang. This can be resolved by either reducing the frequency at which frames are stored, or checkpointing more 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 simulation output, +try setting the `PYMBAR_DISABLE_JAX` environment variable to `1`. diff --git a/pixi.toml b/pixi.toml index e77834c..612b388 100644 --- a/pixi.toml +++ b/pixi.toml @@ -6,9 +6,9 @@ platforms = ["linux-64", "osx-arm64"] [dependencies] python = ">=3.10" # main -biosimspace = ">=2026.1.0,<2026.2.0" +biosimspace = ">=2026.2.0,<2026.3.0" # devel -#biosimspace = "==2026.2.0.dev" +#biosimspace = "==2026.3.0.dev" filelock = "*" ghostly = "*" loch = "*" diff --git a/recipes/somd2/recipe.yaml b/recipes/somd2/recipe.yaml index 7303d8e..74073ae 100644 --- a/recipes/somd2/recipe.yaml +++ b/recipes/somd2/recipe.yaml @@ -20,9 +20,9 @@ requirements: - versioningit run: # main - - biosimspace >=2026.1.0,<2026.2.0 + - biosimspace >=2026.2.0,<2026.3.0 # devel - #- biosimspace ==2026.2.0.dev + #- biosimspace ==2026.3.0.dev - filelock - ghostly - loch @@ -37,6 +37,10 @@ tests: - python: imports: - somd2 + # AmberTools can be pulled in transitively, and installs tools into + # site-packages whose metadata still pins numpy <2, so 'pip check' + # fails even though nothing here uses them. + pip_check: false - script: - PYTHONPATH=. pytest -vvv --color=yes --import-mode=importlib ./tests files: diff --git a/src/somd2/__init__.py b/src/somd2/__init__.py index b3070df..7f0a24a 100644 --- a/src/somd2/__init__.py +++ b/src/somd2/__init__.py @@ -34,10 +34,20 @@ # Store the somd2 version. from ._version import __version__ -# Store the sire version. +# Store the sire version. Unlike somd2/BioSimSpace/ghostly/loch (which use +# versioningit and only append a "+g" local version segment +# for non-release (".dev") builds, omitting it entirely for a clean tagged +# release), sire exposes its version and revision id as separate attributes, +# with __revisionid__ always set regardless of release status. Build a +# composite string using the same "+g" convention as the other +# packages, only appending it for non-release (".dev") builds, so that all +# five version strings are formatted consistently. from sire import __version__ as _sire_version from sire import __revisionid__ as _sire_revisionid +if ".dev" in _sire_version: + _sire_version = f"{_sire_version}+g{_sire_revisionid}" + # Store the BioSimSpace version. from BioSimSpace import __version__ as _biosimspace_version @@ -61,7 +71,7 @@ def get_versions(): """ return { "somd2": __version__, - "sire": f"{_sire_version}+{_sire_revisionid}", + "sire": _sire_version, "biosimspace": _biosimspace_version, "ghostly": _ghostly_version, "loch": _loch_version, diff --git a/src/somd2/_utils/_schedules.py b/src/somd2/_utils/_schedules.py index 555e456..1f1f0e2 100644 --- a/src/somd2/_utils/_schedules.py +++ b/src/somd2/_utils/_schedules.py @@ -27,7 +27,44 @@ ] -def annihilate(fix_epsilon=True): +def _set_boresch_lever_equations(s, stage_dihedral, stage_distance_angle): + """ + Set the equations for a "split" restraint_lever Boresch restraint (see + sire.restraints.boresch's restraint_lever parameter), reproducing the + RXRX protocol's staged restraint turn-on (Table S1 of the RXRX paper's + SI): within each of the two named stages, the corresponding restraint + group ramps from ~0 to 1 following a geometric progression, while the + other group is held fixed. 'stage_dihedral' is the stage over which + the dihedral restraint group ramps on (with the distance/angle group + held at 0); 'stage_distance_angle' is the stage over which the + distance/angle group ramps on (with the dihedral group held at 1, + already fully on). + + Note: this aligns the restraint turn-on with SOMD2's own decharge/ + annihilate(or decouple) stage boundaries, rather than reproducing the + RXRX paper's exact global 50/50 window split (which falls partway + through the annihilate/decouple stage, since the paper's own decharge + stage is only 21 of 64 total bound-leg windows) - the relative sizes of + SOMD2's stages are controlled by lambda_values weighting, not fixed. + """ + from sire.legacy.CAS import Exp as _Exp + import math as _math + + # Geometric progression from ~0.01 (fully off) to 1.0 (fully on), matching + # the ratio observed in the RXRX paper's published lambda schedule. + ramp_on = _Exp((1 - s.lam()) * _math.log(0.01)) + + s.set_equation(stage=stage_dihedral, lever="restraint_dihedral", equation=ramp_on) + s.set_equation(stage=stage_dihedral, lever="restraint_distance_angle", equation=0) + s.set_equation(stage=stage_distance_angle, lever="restraint_dihedral", equation=1) + s.set_equation( + stage=stage_distance_angle, + lever="restraint_distance_angle", + equation=ramp_on, + ) + + +def annihilate(fix_epsilon=True, restraint_lever="split"): """ Build the ABFE lambda schedule using decharge → annihilate. @@ -44,12 +81,28 @@ def annihilate(fix_epsilon=True): If False, epsilon is scaled normally from initial to final and the LRC follows naturally. + restraint_lever : str, optional + How the Boresch restraint is controlled by this schedule, matching + sire.restraints.boresch's restraint_lever parameter. Either "split" + (default), where the dihedral restraint terms are turned on during + decharge and the distance/angle terms are turned on during + annihilate, reproducing the RXRX protocol's staged restraint + turn-on, or "combined", where the whole restraint is turned on + together during the decharge stage. The Boresch restraint object + passed to the simulation must have a matching restraint_lever value. + Returns ------- schedule : sire.legacy.CAS.LambdaSchedule The lambda schedule. """ + if restraint_lever not in ("combined", "split"): + raise ValueError( + "'restraint_lever' must be either 'combined' or 'split', " + f"got {restraint_lever!r}" + ) + from sire.cas import LambdaSchedule as _LambdaSchedule # Start with the standard decouple schedule and modify the stages and @@ -65,14 +118,22 @@ def annihilate(fix_epsilon=True): lever="charge", equation=s.lam() * s.final() + s.initial() * (1 - s.lam()), ) - s.set_equation(stage="decharge", lever="restraint", equation=s.lam() * s.final()) s.add_stage( "annihilate", equation=(-s.lam() + 1) * s.initial() + s.lam() * s.final(), ) s.set_equation(stage="annihilate", lever="charge", equation=s.final()) - s.set_equation(stage="annihilate", lever="restraint", equation=s.final()) + + if restraint_lever == "split": + _set_boresch_lever_equations( + s, stage_dihedral="decharge", stage_distance_angle="annihilate" + ) + else: + s.set_equation( + stage="decharge", lever="restraint", equation=s.lam() * s.final() + ) + s.set_equation(stage="annihilate", lever="restraint", equation=s.final()) if fix_epsilon: s.set_equation(stage="annihilate", lever="epsilon", equation=s.initial()) @@ -86,7 +147,7 @@ def annihilate(fix_epsilon=True): return s -def decouple(fix_epsilon=True): +def decouple(fix_epsilon=True, restraint_lever="split"): """ Build the ABFE lambda schedule using decharge → decouple. @@ -101,12 +162,28 @@ def decouple(fix_epsilon=True): ghost-LRC force is then explicitly scaled to zero over the stage. If False, epsilon is scaled normally and the LRC follows naturally. + restraint_lever : str, optional + How the Boresch restraint is controlled by this schedule, matching + sire.restraints.boresch's restraint_lever parameter. Either "split" + (default), where the dihedral restraint terms are turned on during + decharge and the distance/angle terms are turned on during decouple, + reproducing the RXRX protocol's staged restraint turn-on, or + "combined", where the whole restraint is turned on together during + the decharge stage. The Boresch restraint object passed to the + simulation must have a matching restraint_lever value. + Returns ------- schedule : sire.legacy.CAS.LambdaSchedule The lambda schedule. """ + if restraint_lever not in ("combined", "split"): + raise ValueError( + "'restraint_lever' must be either 'combined' or 'split', " + f"got {restraint_lever!r}" + ) + from sire.cas import LambdaSchedule as _LambdaSchedule # Start with the standard decouple schedule and modify the stages and @@ -114,7 +191,6 @@ def decouple(fix_epsilon=True): # we will use this approach for prototyping. s = _LambdaSchedule.standard_decouple() - s.set_equation(stage="decouple", lever="restraint", equation=s.final()) s.set_equation(stage="decouple", lever="kappa", force="ghost/ghost", equation=0) s.set_equation(stage="decouple", lever="kappa", force="ghost-14", equation=0) s.set_equation(stage="decouple", lever="charge", equation=s.final()) @@ -142,7 +218,16 @@ def decouple(fix_epsilon=True): s.set_equation( stage="decharge", lever="kappa", force="ghost-14", equation=-s.lam() + 1 ) - s.set_equation(stage="decharge", lever="restraint", equation=s.initial() * s.lam()) + + if restraint_lever == "split": + _set_boresch_lever_equations( + s, stage_dihedral="decharge", stage_distance_angle="decouple" + ) + else: + s.set_equation(stage="decouple", lever="restraint", equation=s.final()) + s.set_equation( + stage="decharge", lever="restraint", equation=s.initial() * s.lam() + ) return s diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index d4482be..b2b6598 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -75,13 +75,20 @@ class Config: ], "log_level": [level.lower() for level in _logger._core.levels], "softcore_form": ["zacharias", "taylor", "beutler"], + "precision": ["single", "mixed", "double"], } + # Options that advertise a set of choices, but which also accept other + # forms, e.g. the path to a stream file. These are validated by the setter, + # rather than by argparse. + _open_choices = ["lambda_schedule"] + # A dictionary of nargs for the various options. _nargs = { "lambda_values": "+", "lambda_energy": "+", "rest2_scale": "+", + "restraints": "+", } def __init__( @@ -132,12 +139,16 @@ def __init__( num_energy_neighbours=None, null_energy="1e6 kcal/mol", platform="auto", + precision="single", 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, @@ -172,6 +183,12 @@ def __init__( save_xml=False, page_size=None, timeout="300 s", + restraint_search_time="1 ns", + restraint_search_frequency="10 ps", + restraint_search_receptor_selection=None, + morse_hard_well_depth="150 kcal mol-1", + morse_soft_well_depth="50 kcal mol-1", + morse_soft_force_constant="125 kcal mol-1 A-2", ): """ Constructor. @@ -228,8 +245,10 @@ def __init__( then this will be set to the same as 'lambda_values', or the values defined by 'num_lambda' if 'lambda_values' is not set. - lambda_schedule: str - Lambda schedule to use for alchemical free energy simulations. + lambda_schedule: str, sire.cas.LambdaSchedule + Lambda schedule to use for alchemical free energy simulations. This + can be the name of one of the standard schedules, or the path to a + Sire stream file containing a custom LambdaSchedule. charge_scale_factor: float Factor by which to scale charges for charge scaled morph. @@ -245,9 +264,10 @@ def __init__( The soft-core shift-delta parameter. This is used to soften the Lennard-Jones interaction. - restraints: sire.mm._MM.Restraints - A single set of restraints, or a list of sets of restraints that - will be applied to the atoms during the simulation. + restraints: str, sire.mm._MM.Restraints + One or more paths to Sire stream files containing the sets of + restraints that will be applied to the atoms during the simulation. + A stream file may hold a single set, or a list of sets. constraint: str Constraint type to use for non-perturbable molecules. @@ -350,6 +370,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. @@ -374,14 +400,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. + 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 @@ -556,6 +619,35 @@ def __init__( null_energy: str The energy value to use for lambda windows that are not being computed as part of the energy trajectory. + + restraint_search_time: str + Length of the short pre-production trajectory used to auto-generate + a Boresch restraint when running an ABFE simulation without a + user-supplied restraint. Defaults to "1 ns". + + restraint_search_frequency: str + Frame-saving frequency during the restraint-search trajectory. + Defaults to "10 ps". Should be small enough to yield at least 50 + frames over ``restraint_search_time``. + + restraint_search_receptor_selection: str + Sire selection string for receptor anchor atom candidates used + during automatic Boresch restraint generation. If None, the default + backbone selection is used (CA, C, N atoms in non-water molecules). + + morse_hard_well_depth: str + The well depth of the "hard" Morse potential that replaces the + broken bond when auto-generating restraints for a ring-breaking + simulation. + + morse_soft_well_depth: str + The well depth of the "soft" Morse restraint that holds the broken + fragment in place when auto-generating restraints for a + ring-breaking simulation. + + morse_soft_force_constant: str + The force constant of the "soft" Morse restraint used when + auto-generating restraints for a ring-breaking simulation. """ # Setup logger before doing anything else @@ -605,12 +697,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 @@ -644,9 +740,13 @@ def __init__( self.num_energy_neighbours = num_energy_neighbours self.null_energy = null_energy self.page_size = page_size - + self.restraint_search_time = restraint_search_time + self.restraint_search_frequency = restraint_search_frequency + self.restraint_search_receptor_selection = restraint_search_receptor_selection + self.morse_hard_well_depth = morse_hard_well_depth + self.morse_soft_well_depth = morse_soft_well_depth + self.morse_soft_force_constant = morse_soft_force_constant self.write_config = write_config - self.overwrite = overwrite def __str__(self): @@ -1089,44 +1189,51 @@ def lambda_schedule(self, lambda_schedule): "'lambda_schedule' must be of type 'str' or 'LambdaSchedule' object" ) if isinstance(lambda_schedule, str): - # Strip whitespace and convert to lower case. - lambda_schedule = lambda_schedule.strip().lower() - if lambda_schedule == "standard_morph": + # Strip whitespace. The keyword comparison is made against a + # lower case copy, since the string may also be a path, which + # is case sensitive. + lambda_schedule = lambda_schedule.strip() + keyword = lambda_schedule.lower() + if keyword == "standard_morph": self._lambda_schedule = _LambdaSchedule.standard_morph() self._lambda_schedule_name = "standard_morph" - elif lambda_schedule == "charge_scaled_morph": + elif keyword == "charge_scaled_morph": self._lambda_schedule = _LambdaSchedule.charge_scaled_morph(0.2) self._lambda_schedule_name = "charge_scaled_morph" - elif lambda_schedule == "ring_break_morph": + elif keyword == "ring_break_morph": from .._utils._schedules import ( ring_break_morph as _ring_break_morph, ) self._lambda_schedule = _ring_break_morph() self._lambda_schedule_name = "ring_break_morph" - elif lambda_schedule == "reverse_ring_break_morph": + elif keyword == "reverse_ring_break_morph": from .._utils._schedules import ( reverse_ring_break_morph as _reverse_ring_break_morph, ) self._lambda_schedule = _reverse_ring_break_morph() self._lambda_schedule_name = "reverse_ring_break_morph" - elif lambda_schedule == "annihilate": + elif keyword == "annihilate": self._lambda_schedule = None self._lambda_schedule_name = "annihilate" - elif lambda_schedule == "decouple": + elif keyword == "decouple": self._lambda_schedule = None self._lambda_schedule_name = "decouple" else: - try: - self._lambda_schedule = self._from_hex(lambda_schedule) - self._lambda_schedule_name = None - except Exception: + schedule = self._from_string( + lambda_schedule, + "lambda_schedule", + hint=", or one of the following strings: " + f"{', '.join(self._choices['lambda_schedule'])}", + ) + if not isinstance(schedule, _LambdaSchedule): raise ValueError( - "Unable to deserialise 'lambda_schedule'. Ensure that this is a " - "hex string representation of a valid LambdaSchedule object, or " - f"one of the following strings: {', '.join(self._choices['lambda_schedule'])}" + f"'lambda_schedule' deserialised to a " + f"'{type(schedule).__name__}', not a 'LambdaSchedule'." ) + self._lambda_schedule = schedule + self._lambda_schedule_name = None else: self._lambda_schedule = lambda_schedule self._lambda_schedule_name = None @@ -1212,32 +1319,34 @@ def restraints(self): @restraints.setter def restraints(self, restraints): - # If not supplied as a list, convert to a list. + # If not supplied as a list, convert to a list. Note that a string is + # itself iterable, so must be wrapped explicitly. if restraints is not None: - if not isinstance(restraints, _Iterable): + if isinstance(restraints, str) or not isinstance(restraints, _Iterable): restraints = [restraints] - # Check that all restraints are of the correct type. - deserialised_restraints = [] + # Resolve each entry, keeping objects and deserialised strings in + # the order they were given. + resolved_restraints = [] for restraint in restraints: - if isinstance(restraint, _sr.mm._MM.Restraints): - continue - elif isinstance(restraint, str): - try: - restraint = self._from_hex(restraint) - except Exception: - raise ValueError( - "Unable to deserialise restraint. Ensure that this " - "is a hex string representation of a valid sire.mm._MM.Restraints object." - ) - deserialised_restraints.append(restraint) + if isinstance(restraint, str): + restraint = self._from_string(restraint.strip(), "restraints") + + # A stream file may hold a list of sets of restraints, e.g. the + # pair used for a ring-breaking perturbation. + if isinstance(restraint, _Iterable): + resolved_restraints.extend(restraint) else: + resolved_restraints.append(restraint) + + # Check that all restraints are of the correct type. + for restraint in resolved_restraints: + if not isinstance(restraint, _sr.mm._MM.Restraints): raise ValueError( "'restraints' must be a sire.mm._MM.Restraints object, or a list of these objects." ) - if len(deserialised_restraints) > 0: - restraints = deserialised_restraints + restraints = resolved_restraints self._restraints = restraints @@ -1647,6 +1756,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 @@ -1758,6 +1882,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 @@ -2467,6 +2640,89 @@ def _from_hex(hex): return obj + @classmethod + def _from_string(cls, string, name, hint=""): + """ + Internal method to deserialise a Sire object from a string, which can + either be the path to a stream file, or the hex string representation + of the serialised object. + + Parameters + ---------- + + string: str + The path to a stream file, or a hex string representation of the + Sire object. + + name: str + The name of the option being set, used for error messages. + + hint: str + An additional clause appended to the error message, e.g. listing + the keywords that the option also accepts. + + Returns + ------- + + obj: + The deserialised Sire object. + """ + from pathlib import Path as _Path + + # Work out whether this is a path to an existing file. A hex string can + # exceed the maximum filename length, which raises rather than simply + # returning False on some platforms. + try: + is_file = _Path(string).is_file() + except Exception: + is_file = False + + if is_file: + from sire.stream import load + + try: + return load(string) + except Exception as e: + raise ValueError( + f"Unable to load '{name}' from stream file '{string}': {e}" + ) + else: + try: + return cls._from_hex(string) + except Exception: + raise ValueError( + f"Unable to interpret '{name}'. Expected the path to a Sire " + f"stream file, or a hex string of a serialised object{hint}." + ) + + def __getstate__(self): + """ + Hex-encode the same fields that to_yaml()/from_yaml() already + hex-encode (currently 'restraints' and 'lambda_schedule'), since + these legacy Sire objects are not guaranteed to have native pickle + support. This is needed so that a Config holding these can be sent + to a spawned worker process, e.g. via + concurrent.futures.ProcessPoolExecutor. + """ + state = self.__dict__.copy() + if state.get("_restraints") is not None: + state["_restraints"] = [ + self._to_hex(restraint) for restraint in state["_restraints"] + ] + if state.get("_lambda_schedule") is not None: + state["_lambda_schedule"] = self._to_hex(state["_lambda_schedule"]) + return state + + def __setstate__(self, state): + """Reverse the hex-encoding performed in __getstate__.""" + if state.get("_restraints") is not None: + state["_restraints"] = [ + self._from_hex(restraint) for restraint in state["_restraints"] + ] + if state.get("_lambda_schedule") is not None: + state["_lambda_schedule"] = self._from_hex(state["_lambda_schedule"]) + self.__dict__.update(state) + @classmethod def _create_parser(cls): """ @@ -2475,6 +2731,7 @@ def _create_parser(cls): import argparse import inspect + import re # Inspect the signature to get the parameters. sig = inspect.signature(Config.__init__) @@ -2482,7 +2739,7 @@ def _create_parser(cls): params = { key: value for key, value in params.items() - if key not in ["self", "args", "kwargs", "restraints"] + if key not in ["self", "args", "kwargs"] } # Get the docstring. @@ -2499,7 +2756,7 @@ def _create_parser(cls): # Loop over all lines in the docstring until we find the parameter. for line in doc: line = line.strip() - if line.startswith(param): + if re.match(rf"{re.escape(param)}\s*:", line): found_param = True elif found_param: if line == "": @@ -2541,14 +2798,26 @@ def _create_parser(cls): # This parameter has choices. if param in cls._choices: - parser.add_argument( - f"--{cli_param}", - type=typ, - default=params[param].default, - choices=cls._choices[param], - help=help[param], - required=False, - ) + # Other forms are also accepted, so advertise the choices in the + # help text, but leave the validation to the setter. + if param in cls._open_choices: + parser.add_argument( + f"--{cli_param}", + type=typ, + default=params[param].default, + metavar="{" + ",".join(cls._choices[param]) + "}", + help=help[param], + required=False, + ) + else: + parser.add_argument( + f"--{cli_param}", + type=typ, + default=params[param].default, + choices=cls._choices[param], + help=help[param], + required=False, + ) # This is a standard parameter. else: if typ == bool: @@ -2571,6 +2840,144 @@ def _create_parser(cls): return parser + @property + def restraint_search_time(self): + return self._restraint_search_time + + @restraint_search_time.setter + def restraint_search_time(self, restraint_search_time): + if not isinstance(restraint_search_time, str): + raise TypeError("'restraint_search_time' must be of type 'str'") + + from sire.units import picosecond + + try: + t = _sr.u(restraint_search_time) + except: + raise ValueError( + f"Unable to parse 'restraint_search_time' as a Sire GeneralUnit: {restraint_search_time}" + ) + + if not t.has_same_units(picosecond): + raise ValueError("'restraint_search_time' units are invalid.") + + self._restraint_search_time = t + + @property + def restraint_search_frequency(self): + return self._restraint_search_frequency + + @restraint_search_frequency.setter + def restraint_search_frequency(self, restraint_search_frequency): + if not isinstance(restraint_search_frequency, str): + raise TypeError("'restraint_search_frequency' must be of type 'str'") + + from sire.units import picosecond + + try: + t = _sr.u(restraint_search_frequency) + except: + raise ValueError( + f"Unable to parse 'restraint_search_frequency' as a Sire GeneralUnit: {restraint_search_frequency}" + ) + + if not t.has_same_units(picosecond): + raise ValueError("'restraint_search_frequency' units are invalid.") + + self._restraint_search_frequency = t + + @property + def restraint_search_receptor_selection(self): + return self._restraint_search_receptor_selection + + @restraint_search_receptor_selection.setter + def restraint_search_receptor_selection(self, restraint_search_receptor_selection): + if restraint_search_receptor_selection is not None: + if not isinstance(restraint_search_receptor_selection, str): + raise TypeError( + "'restraint_search_receptor_selection' must be of type 'str'" + ) + self._restraint_search_receptor_selection = restraint_search_receptor_selection + + @property + def morse_hard_well_depth(self): + return self._morse_hard_well_depth + + @morse_hard_well_depth.setter + def morse_hard_well_depth(self, morse_hard_well_depth): + self._morse_hard_well_depth = self._parse_well_depth( + morse_hard_well_depth, "morse_hard_well_depth" + ) + + @property + def morse_soft_well_depth(self): + return self._morse_soft_well_depth + + @morse_soft_well_depth.setter + def morse_soft_well_depth(self, morse_soft_well_depth): + self._morse_soft_well_depth = self._parse_well_depth( + morse_soft_well_depth, "morse_soft_well_depth" + ) + + @property + def morse_soft_force_constant(self): + return self._morse_soft_force_constant + + @morse_soft_force_constant.setter + def morse_soft_force_constant(self, morse_soft_force_constant): + if not isinstance(morse_soft_force_constant, str): + raise TypeError("'morse_soft_force_constant' must be of type 'str'") + + from sire.units import angstrom, kcal_per_mol + + try: + k = _sr.u(morse_soft_force_constant) + except: + raise ValueError( + "Unable to parse 'morse_soft_force_constant' as a Sire " + f"GeneralUnit: {morse_soft_force_constant}" + ) + + if not k.has_same_units(kcal_per_mol / (angstrom * angstrom)): + raise ValueError("'morse_soft_force_constant' units are invalid.") + + self._morse_soft_force_constant = k + + @staticmethod + def _parse_well_depth(value, name): + """ + Internal helper to validate a Morse potential well depth. + + Parameters + ---------- + + value: str + The well depth as a string, e.g. "150 kcal mol-1". + + name: str + The name of the option, used in error messages. + + Returns + ------- + + well_depth: sire.units.GeneralUnit + The parsed well depth. + """ + if not isinstance(value, str): + raise TypeError(f"'{name}' must be of type 'str'") + + from sire.units import kcal_per_mol + + try: + de = _sr.u(value) + except: + raise ValueError(f"Unable to parse '{name}' as a Sire GeneralUnit: {value}") + + if not de.has_same_units(kcal_per_mol): + raise ValueError(f"'{name}' units are invalid.") + + return de + def _reset_logger(self, logger): """ Internal method to reset the logger. diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 5e49b4d..8e8890b 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -207,7 +207,8 @@ def __init__(self, system, config): if c0 != c1: msg = ( "End-state connectivities are different. If this is a ring-breaking " - "perturbation, please set 'lambda_schedule_name' to 'ring_breaking'." + "perturbation, please set 'lambda_schedule_name' to 'ring_break_morph' " + "(or 'reverse_ring_break_morph' for the reverse perturbation)." ) _logger.warning(msg) break @@ -223,6 +224,20 @@ def __init__(self, system, config): except: self._has_water = False + # Check for protein (distinguishes ABFE from AHFE). A protein+ligand + # system has more than one non-water molecule with at least 3 atoms; + # a solvated ligand alone has only one. + try: + non_water_mols = self._system["(not water) and (atomidx > 1)"].molecules() + self._has_protein = non_water_mols.num_molecules() > 1 + except: + self._has_protein = False + + # Set by _generate_boresch_restraint() when a Boresch restraint is + # auto-generated, then written into the energy trajectory parquet + # metadata (see _checkpoint). + self._standard_state_correction = None + # Warn if dispersion correction is requested but can't be applied. if self._config.use_dispersion_correction and not self._has_water: msg = "Cannot use dispersion correction for vacuum simulations. Disabling!" @@ -253,6 +268,24 @@ def __init__(self, system, config): self._config._extra_args["use_gcmc_lrc"] = True self._config._extra_args["num_gcmc_waters"] = self._config.gcmc_num_waters + # Auto-generate Morse restraints for ring-breaking perturbations with no + # user-supplied restraint. This is done before any modification of the + # bonded terms, and before the reference system is stored and the restart + # checks are performed, since the hard restraint replaces a bond in the + # system, which must match the checkpoints. + if self._is_ring_break and self._config.restraints is None: + try: + self._config.restraints = self._generate_morse_restraints() + except Exception as e: + msg = ( + "Unable to generate Morse restraints for ring-breaking " + f"perturbation: {e}. If the Morse potential has already been " + "applied to the input system, then pass the corresponding " + "restraints using the 'restraints' option." + ) + _logger.error(msg) + raise RuntimeError(msg) + # We're running in SOMD1 compatibility mode. if self._config.somd1_compatibility: from .._utils._somd1 import make_compatible @@ -274,7 +307,7 @@ def __init__(self, system, config): elif self._config.ghost_modifications: from ghostly import modify - _logger.info("Applying modifications to ghost atom bonded terms") + _logger.info("Applying modifications to ghost atom bonded terms.") try: self._system, self._modifications = modify(self._system) # Angle optimisation can sometimes fail. @@ -433,6 +466,8 @@ def __init__(self, system, config): # Set the lambda values. if self._config.lambda_values: self._lambda_values = self._config.lambda_values + elif self._config.num_lambda == 1: + self._lambda_values = [0.0] else: self._lambda_values = [ round(i / (self._config.num_lambda - 1), 5) @@ -601,14 +636,18 @@ def __init__(self, system, config): # used to skip duplicate rows on restart. self._last_ec_time = {} + # Per-window cache of the integrator's integration force groups bitmask. + self._integration_groups = {} + # Store the current system as a reference. self._reference_system = self._system.clone() # 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 @@ -875,6 +914,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, @@ -977,6 +1017,249 @@ def __init__(self, system, config): # Update the maximum number of threads. _sr.legacy.Base.set_max_num_threads(sire_threads) + @property + def _is_abfe_bound(self): + """ + Whether this is the bound leg of an ABFE simulation: annihilate/decouple + schedule with a solvated protein present. False for the free leg (ligand + in solvent, no protein), even though it may use the same lambda schedule. + """ + return ( + self._config._lambda_schedule_name in ("annihilate", "decouple") + and self._has_protein + and self._has_water + ) + + @property + def _is_ring_break(self): + """ + Whether this is a ring-breaking (or ring-making) simulation, i.e. one + using the 'ring_break_morph' lambda schedule, or its reverse. + """ + return self._config._lambda_schedule_name in ( + "ring_break_morph", + "reverse_ring_break_morph", + ) + + def _generate_morse_restraints(self): + """ + Return the pair of Morse restraints required by the 'ring_break_morph' + lambda schedule, or its reverse. Called automatically when running a + ring-breaking simulation with no user-supplied restraint. + + The "hard" restraint directly replaces the harmonic bond that is broken + (or formed) by the perturbation, inheriting its force constant and + equilibrium length. The "soft" restraint acts on the same pair of atoms + and holds the broken fragment in place while the hard restraint is + switched off. + + Returns + ------- + + restraints: [sire.mm.MorsePotentialRestraints] + The hard and soft Morse restraints, in the order expected by the + schedule's 'morse_hard' and 'morse_soft' levers. + + Notes + ----- + + As a side effect, ``self._system`` is updated with the replacement of + the broken bond by the hard Morse potential. + """ + from sire.restraints import morse_potential as _morse_potential + + _logger.info( + "No restraints supplied for ring-breaking perturbation. " + "Generating default Morse restraints." + ) + + hard_restraints, self._system = _morse_potential( + self._system, + de=self._config.morse_hard_well_depth, + auto_parametrise=True, + direct_morse_replacement=True, + name="morse_hard", + ) + + # Restrain the same pair of atoms as the hard restraint, at the same + # equilibrium distance. + soft_restraints, _ = _morse_potential( + self._system, + atoms0=hard_restraints[0].atom0(), + atoms1=hard_restraints[0].atom1(), + r0=hard_restraints[0].r0(), + k=self._config.morse_soft_force_constant, + de=self._config.morse_soft_well_depth, + auto_parametrise=False, + direct_morse_replacement=False, + name="morse_soft", + ) + + _logger.info(f"Hard Morse restraint: {hard_restraints[0]}") + _logger.info(f"Soft Morse restraint: {soft_restraints[0]}") + + return [hard_restraints, soft_restraints] + + def _generate_boresch_restraint(self, device=None): + """ + Return a Boresch restraint for the ABFE simulation, either by loading + one saved from a previous run or by running a short lambda=0 trajectory. + Called automatically before minimisation/equilibration of the production + windows/replicas when the simulation is ABFE and no restraint has been + supplied. + + The input system is assumed to already be equilibrated (SOMD2 does not + run a separate equilibration stage for the restraint search); it is + minimised using the production minimisation settings, then a short + trajectory is run at lambda=0, matching the dynamics settings used for + production, to derive the restraint geometry and force constants. + + Parameters + ---------- + + device : int, optional + GPU device number to use for the restraint-search run. + + Returns + ------- + + restraints : sire.mm.BoreschRestraints + + Notes + ----- + + As a side effect, when the restraint is generated (rather than reused + from a checkpointed restart), ``self._system`` is re-seeded from the + least-strained trajectory frame returned by ``boresch_search``. The + restraint equilibrium values are trajectory averages, so the input + structure is generally not consistent with them; starting production + from it instead would leave the restraint badly strained at t=0 and can + blow the simulation up as the restraint is switched on. + + The cached restraint is only reused on a genuine (checkpointed) restart, + where it must be kept identical to the one the accumulated free energy + was computed with. When there is no checkpoint (a fresh run, or a crash + before any progress was checkpointed) the restraint is regenerated. This + avoids getting pinned to a restraint that is itself the cause of the + crash: reusing it and re-seeding the same frame would just reproduce the + crash on every restart, whereas a fresh search may pick a different frame + or anchor, and re-seeds ``self._system`` naturally. + """ + from sire.restraints import boresch_search + + restraint_file = str(self._config.output_directory / "abfe_restraint.s3") + + # Only reuse a saved restraint when continuing an actual checkpointed + # restart: the restraint must match the one the accumulated free energy + # was computed with, and the coordinates come from the checkpoint (so no + # re-seeding is needed). Deliberately do NOT reuse it when there is no + # checkpoint - see the Notes above. + if self._is_restart: + if not _Path(restraint_file).exists(): + raise FileNotFoundError( + "Restarting an ABFE bound-leg simulation, but no saved " + f"restraint was found at {restraint_file}. The restraint " + "cannot be regenerated mid-simulation without invalidating " + "the accumulated free energy." + ) + + _logger.info(f"Loading existing Boresch restraint from {restraint_file}") + restraints = _sr.stream.load(restraint_file) + + from sire.restraints import get_standard_state_correction + + correction = get_standard_state_correction( + restraints[0], temperature=self._config.temperature + ) + self._standard_state_correction = float( + correction.to(_sr.units.kcal_per_mol) + ) + + return restraints + + _logger.info( + "No restraint supplied for ABFE. Running Boresch restraint search." + ) + + search_system = self._system + + if self._config.minimise: + constraint = self._config.constraint + perturbable_constraint = self._config.perturbable_constraint + + # Don't use constraints during minimisation. + if not self._config.minimisation_constraints: + constraint = "none" + perturbable_constraint = "none" + + min_dynamics_kwargs = self._dynamics_kwargs.copy() + min_dynamics_kwargs.update( + { + "device": device, + "lambda_value": 0.0, + "constraint": constraint, + "perturbable_constraint": perturbable_constraint, + } + ) + + min_dynamics = search_system.dynamics(**min_dynamics_kwargs) + min_dynamics.minimise(timeout=self._config.timeout) + search_system = min_dynamics.commit() + + dynamics_kwargs = self._dynamics_kwargs.copy() + dynamics_kwargs.update( + { + "device": device, + "lambda_value": 0.0, + } + ) + + dynamics = search_system.dynamics(**dynamics_kwargs) + dynamics.run( + self._config.restraint_search_time, + energy_frequency=0, + frame_frequency=self._config.restraint_search_frequency, + save_velocities=False, + auto_fix_minimise=self._config.auto_fix_minimise, + save_crash_report=self._config.save_crash_report, + ) + search_system = dynamics.commit() + + search_kwargs = {"temperature": self._config.temperature} + if self._config.restraint_search_receptor_selection is not None: + search_kwargs["receptor_selection"] = ( + self._config.restraint_search_receptor_selection + ) + + restraints, correction, starting_structure = boresch_search( + search_system, **search_kwargs + ) + + # Cache so it can be written into the energy trajectory parquet + # metadata (see _checkpoint), letting analysis code automatically + # apply the correction without needing to scan the logs. + self._standard_state_correction = float(correction.to(_sr.units.kcal_per_mol)) + _logger.info( + f"Boresch restraint generated. Standard state correction: " + f"{self._standard_state_correction:.4f} kcal mol-1" + ) + + # Re-seed production from the least-strained frame so the restraint is + # essentially relaxed at t=0 (see the docstring Notes). The frame comes + # from the perturbable search system, so link its properties back to the + # reference (lambda=0) end state, matching how the seed systems are + # handled elsewhere (see _perturbed_system), and drop the search + # trajectory frames so only the single starting snapshot is retained. + starting_structure = _sr.morph.link_to_reference(starting_structure) + starting_structure.delete_all_frames() + self._system = starting_structure + + # Save so that a genuine (checkpointed) restart can reuse the exact same + # restraint without re-running the search. + _sr.stream.save(restraints, restraint_file) + + return restraints + def _check_space(self): """ Check if the system has a periodic space. @@ -1005,8 +1288,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") @@ -1031,12 +1319,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." ) @@ -1353,10 +1643,10 @@ def _create_alchemical_ions( restraint_distance, ) - try: - restraints.add(restraint) - except: + if restraints is None: restraints = restraint + else: + restraints.add(restraint) # Update the system. system.update(merged) @@ -1629,6 +1919,10 @@ def _compare_configs(config1, config2): "log_file", "overwrite", "timeout", + "oversubscription_factor", + "max_contexts", + "restraint_search_time", + "restraint_search_frequency", ] for key in config1.keys(): if key not in allowed_diffs: @@ -2067,8 +2361,7 @@ def _checkpoint( if not is_post_equilibration: metadata = { "attrs": df.attrs, - "somd2 version": versions["somd2"], - "sire version": versions["sire"], + "versions": versions, "lambda": f"{lam:.5f}", "speed": speed, "temperature": str(self._config.temperature.value()), @@ -2078,6 +2371,13 @@ def _checkpoint( if lambda_grad is not None: metadata["lambda_grad"] = [f"{v:.5f}" for v in lambda_grad] + # Add the standard state correction, if a Boresch restraint + # was auto-generated for this ABFE run. + if self._standard_state_correction is not None: + metadata["standard_state_correction"] = ( + f"{self._standard_state_correction:.6f}" + ) + if is_final_block: # Save the end-state GCMC topologies for trajectory analysis and visualisation. # This topology contains additional water molecules that are used for GCMC @@ -2234,6 +2534,53 @@ def _write_checkpoint_system(self, system, index, context=None, gcmc_sampler=Non system.delete_all_frames() _sr.stream.save(system, self._filenames[index]["checkpoint"]) + @staticmethod + def _is_legacy_gcmc_stats(stats): + """ + Whether GCMC statistics are in the format used before a sampler could + be re-used across lambda values. + + Those were a flat dictionary of counters for a single lambda value, + rather than a dictionary of counters keyed by lambda value. + + Parameters + ---------- + + stats: dict + The GCMC sampling statistics. + + Returns + ------- + + bool + Whether the statistics are in the old format. + """ + return isinstance(stats, dict) and "num_moves" in stats + + @staticmethod + def _convert_legacy_gcmc_stats(stats, lambda_value): + """ + Convert GCMC statistics from the old format to the current one. + + Parameters + ---------- + + stats: dict + A flat dictionary of counters, for a single lambda value. + + lambda_value: float + The lambda value that the statistics belong to. + + Returns + ------- + + dict + The statistics, keyed by lambda value. + """ + from loch import GCMCSampler as _GCMCSampler + + return {_GCMCSampler.stats_key(lambda_value): dict(stats)} + def _backup_checkpoint(self, index): """ Create a backup of the previous checkpoint files. @@ -2336,10 +2683,19 @@ def _save_energy_components(self, index, context, time_ns): if time_ns <= self._last_ec_time[index]: return + if index not in self._integration_groups: + self._integration_groups[index] = ( + context.getIntegrator().getIntegrationForceGroups() + ) + integration_groups = self._integration_groups[index] + # Use the named force groups already assigned by sire_to_openmm_system, # sorted alphabetically for a consistent column order across runs. + # Skip any group not actually used for integration. energies = {} for name, grp in sorted(context._force_group_map.items()): + if not integration_groups & (1 << grp): + continue state = context.getState(getEnergy=True, groups=(1 << grp)) energies[name] = state.getPotentialEnergy().value_in_unit( openmm.unit.kilocalories_per_mole diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index e112aea..5dc4d66 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,33 +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._gcmc_stats = [None] * len(lambdas) - 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._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._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( @@ -132,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) @@ -143,13 +212,52 @@ def __setstate__(self, state): # so that old checkpoint files can still be loaded. n = len(self._lambdas) if not hasattr(self, "_gcmc_stats"): - self._gcmc_stats = [None] * n + self._gcmc_stats = None if not hasattr(self, "_gcmc_states"): self._gcmc_states = [None] * n if not hasattr(self, "_terminal_flip_stats"): - self._terminal_flip_stats = [[0, 0]] * 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. + # Convert these to a single dictionary keyed by lambda value. + if isinstance(self._gcmc_stats, list): + converted = {} + for lam, stats in zip(self._lambdas, self._gcmc_stats): + if _RunnerBase._is_legacy_gcmc_stats(stats): + converted.update(_RunnerBase._convert_legacy_gcmc_stats(stats, lam)) + self._gcmc_stats = converted if converted else None def __getstate__(self): """ @@ -160,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, @@ -241,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) @@ -266,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 @@ -281,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: @@ -291,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, ) @@ -317,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 @@ -330,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] @@ -383,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) @@ -413,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) ) @@ -479,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): """ @@ -500,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): """ @@ -551,43 +829,168 @@ def set_states(self, states): """ self._states = states - def mix_states(self, old_states): + def load_replica(self, replica, clock=None): """ - Mix the states of the dynamics objects. + 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 ---------- - old_states : numpy.ndarray - The state indices from before the last replica mix. + + 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. """ - # Mix the states. - for i, state in enumerate(self._states): - # The state has changed. - if i != state: - _logger.debug(f"Replica {i} seeded from state {state}") - self._apply_openmm_state( - self._dynamics[i].context(), self._openmm_states[state] - ) + slot = self._replica_slot[replica] + dynamics = self._dynamics[slot] + gcmc_sampler = self._gcmc_samplers[slot] - # 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] + # The replica that the slot currently holds. + resident = self._slot_replica[slot] - # Update the water state in the GCMCSampler. - self._gcmc_samplers[i].push() - try: - self._gcmc_samplers[i]._set_water_state( - self._dynamics[i].context(), + # The slot needs new positions and velocities if it is being handed a + # different replica, or if the last mix moved this replica's state. + if resident != replica or self._state_moved[replica]: + self._apply_openmm_state(dynamics.context(), self._openmm_states[replica]) + + # Positions have changed underneath the context, so any cached + # energies are stale. + dynamics.clear_energy_cache() + + self._state_moved[replica] = False + + # Set the lambda value and REST2 scaling factor. This is a no-op if + # the slot is already at this lambda value. + if resident != replica: + dynamics.set_lambda( + self._lambdas[replica], + rest2_scale=self._rest2_scale_factors[replica], + update_constraints=self._update_constraints, + ) + + # Restore the clock and point the slot at this replica's energy + # trajectory, so that energies are accumulated against the replica + # rather than the slot. + if clock is not None: + dynamics._set_clock(clock) + dynamics.set_energy_trajectory(self._energy_trajectories[replica]) + + if gcmc_sampler is not None: + gcmc_sampler.push() + try: + # Swap the water state into the sampler. Diff against what the + # sampler currently holds, which is the state of whichever + # replica was last resident. + target = self._gcmc_states[replica] + if target is not None: + current = gcmc_sampler.water_state() + water_idxs = _np.where(current != target)[0] + + if len(water_idxs) > 0: + gcmc_sampler._set_water_state( + dynamics.context(), indices=water_idxs, - states=self._gcmc_states[state][water_idxs], + states=target[water_idxs], ) - finally: - self._gcmc_samplers[i].pop() - # Update the swap matrix. + 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_num_waters[replica] = gcmc_sampler.num_waters( + context=self._dynamics[slot].context() + ) + finally: + gcmc_sampler.pop() + + def mix_states(self, old_states): + """ + 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)) + ] + + # Update the swap matrix. + for i, state in enumerate(self._states): self._num_swaps[old_states[i], state] += 1 def get_proposed(self): @@ -634,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): """ @@ -811,6 +1239,36 @@ 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) + + # There is nothing to exchange with a single replica. + if self._num_replicas < 2: + msg = ( + "Replica exchange requires at least two " + f"{_lam_sym} values. Please use the Runner class instead." + ) + _logger.error(msg) + raise ValueError(msg) + + 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 + # construction time and won't pick up a restraint added afterwards. + if self._is_abfe_bound and self._config.restraints is None: + try: + restraints = self._generate_boresch_restraint(device=0) + except Exception as e: + msg = f"Unable to generate Boresch restraint for ABFE simulation: {e}" + _logger.error(msg) + raise RuntimeError(msg) + self._config.restraints = restraints + self._dynamics_kwargs["restraints"] = restraints + # Store the name of the dynamics cache pickle file. self._repex_state = self._config.output_directory / "repex_state.pkl" @@ -831,7 +1289,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" @@ -884,6 +1346,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, ) @@ -959,6 +1424,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, @@ -970,43 +1445,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() - if self._dynamics_cache._gcmc_stats[i] is not None: - gcmc_sampler.restore_stats(self._dynamics_cache._gcmc_stats[i]) - - # 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" @@ -1067,6 +1531,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})" @@ -1143,20 +1810,23 @@ 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. + dynamics_executor = ThreadPoolExecutor(max_workers=num_workers) + checkpoint_executor = ThreadPoolExecutor(max_workers=num_checkpoint_workers) + # Create the replica list. replica_list = list(range(self._config.num_lambda)) - # 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}" @@ -1171,12 +1841,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( @@ -1190,35 +1860,55 @@ 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): replicas = replica_list[ j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] - with ThreadPoolExecutor( - max_workers=num_checkpoint_workers - ) as executor: - try: - for index, error in executor.map( - self._checkpoint, - replicas, - repeat(self._lambda_values), - repeat(-1), - repeat(cycles), - ): - if error is not None: - msg = ( - f"Post-equilibration checkpoint failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}:\n{error}" - ) - _logger.error(msg) - raise error - except KeyboardInterrupt: - _logger.error( - "Post-equilibration checkpoint cancelled. Exiting." - ) - _sys.exit(1) + try: + for index, error in checkpoint_executor.map( + self._checkpoint, + replicas, + repeat(self._lambda_values), + repeat(-1), + repeat(cycles), + ): + if error is not None: + msg = ( + f"Post-equilibration checkpoint failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}:\n{error}" + ) + _logger.error(msg) + raise error + except KeyboardInterrupt: + checkpoint_executor.shutdown(wait=False, cancel_futures=True) + _logger.error( + "Post-equilibration checkpoint cancelled. Exiting." + ) + _sys.exit(1) # Current block number. block = self._start_block @@ -1302,13 +1992,25 @@ 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] - with ThreadPoolExecutor(max_workers=num_workers) as executor: + # 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 executor.map( + for result, index, energies in dynamics_executor.map( self._run_block, replicas, repeat(self._lambda_values), @@ -1316,6 +2018,7 @@ def run(self): repeat(write_gcmc_ghosts), repeat(is_terminal_flip), repeat(time_ns), + repeat(clock), ): if not result: _logger.error( @@ -1324,11 +2027,45 @@ def run(self): 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: + # 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 + ) + _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) @@ -1342,21 +2079,23 @@ def run(self): j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] - with ThreadPoolExecutor(max_workers=num_workers) as executor: - try: - for index, error in executor.map( - self._backup_checkpoint, - replicas, - ): - if not result: - _logger.error( - f"Backup failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error - except KeyboardInterrupt: - _logger.error("Backup cancelled. Exiting.") - _sys.exit(1) + try: + for index, error in checkpoint_executor.map( + self._backup_checkpoint, + replicas, + ): + if error: + _logger.error( + f"Backup failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}: {error}" + ) + raise error + except KeyboardInterrupt: + checkpoint_executor.shutdown( + wait=False, cancel_futures=True + ) + _logger.error("Backup cancelled. Exiting.") + _sys.exit(1) # Now write the new checkpoint files. for j in range(num_checkpoint_batches): @@ -1365,25 +2104,27 @@ def run(self): j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] - with ThreadPoolExecutor(max_workers=num_workers) as executor: - try: - for index, error in executor.map( - self._checkpoint, - replicas, - repeat(self._lambda_values), - repeat(block), - repeat(num_blocks + int(rem > 0)), - repeat(i == cycles - 1), - ): - if error: - _logger.error( - f"Checkpoint failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error - except KeyboardInterrupt: - _logger.error("Checkpoint cancelled. Exiting.") - _sys.exit(1) + try: + for index, error in checkpoint_executor.map( + self._checkpoint, + replicas, + repeat(self._lambda_values), + repeat(block), + repeat(num_blocks + int(rem > 0)), + repeat(i == cycles - 1), + ): + if error: + _logger.error( + f"Checkpoint failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}: {error}" + ) + raise error + except KeyboardInterrupt: + checkpoint_executor.shutdown( + wait=False, cancel_futures=True + ) + _logger.error("Checkpoint cancelled. Exiting.") + _sys.exit(1) # Assemble an energy matrix from the results. _logger.info("Assembling energy matrix") @@ -1400,16 +2141,11 @@ def run(self): self._dynamics_cache.get_accepted(), ) ) - self._dynamics_cache.mix_states(old_states) - # 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: @@ -1439,6 +2175,9 @@ def run(self): with open(self._repex_state, "wb") as f: _pickle.dump(self._dynamics_cache, f) + dynamics_executor.shutdown(wait=True) + checkpoint_executor.shutdown(wait=True) + # Record the end time for the production block. prod_end = time() @@ -1497,10 +2236,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 ---------- @@ -1528,6 +2273,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 ------- @@ -1546,13 +2297,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 @@ -1566,9 +2318,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: @@ -1577,13 +2326,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 ) @@ -1617,12 +2365,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: @@ -1666,8 +2411,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() @@ -1698,17 +2447,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. @@ -1720,13 +2469,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 @@ -1757,8 +2513,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() @@ -1797,17 +2557,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 @@ -1821,7 +2581,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( @@ -1843,17 +2603,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. @@ -1865,21 +2625,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 @@ -1915,7 +2670,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) @@ -1993,6 +2748,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. @@ -2029,13 +2858,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: @@ -2047,7 +2886,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: @@ -2060,23 +2900,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: @@ -2195,15 +3036,36 @@ def _mix_replicas(num_replicas, energy_matrix, proposed, accepted): return states + def _merge_gcmc_stats(self): + """ + Merge the GCMC sampling statistics from every sampler. + + A sampler accumulates statistics for each lambda value it visits, so + the results are gathered into a single dictionary keyed by lambda + value. Samplers only report the lambda values they visit, so the keys + are disjoint and the merge order doesn't matter. + + Returns + ------- + + dict + The statistics for each lambda value, or None if not using GCMC. + """ + stats = {} + + 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()) + + return stats if stats else None + def _save_sampler_stats(self): """ Save GCMC and terminal flip sampler statistics to the dynamics cache prior to pickling. """ - for i in range(len(self._lambda_values)): - _, gcmc_sampler = self._dynamics_cache.get(i) - if gcmc_sampler is not None: - self._dynamics_cache._gcmc_stats[i] = gcmc_sampler.get_stats() + self._dynamics_cache._gcmc_stats = self._merge_gcmc_stats() if self._terminal_flip_samplers is not None: self._dynamics_cache._terminal_flip_stats = [ @@ -2226,7 +3088,7 @@ def _save_transition_matrix(self): swaps[i_state, j_state] + swaps[j_state, i_state] ) / denom else: - t[i_state, i_state] = 1.0 + t_ij[i_state, i_state] = 1.0 # Backup the existing transition matrix, if it exists. if self._repex_matrix.exists(): diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 0200658..5b6283e 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -215,6 +215,18 @@ def run(self): else: self._max_workers = 1 + # Auto-generate a Boresch restraint for ABFE runs with no user-supplied restraint. + if self._is_abfe_bound and self._config.restraints is None: + device = self._gpu_pool[0] if self._is_gpu else None + try: + restraints = self._generate_boresch_restraint(device=device) + except Exception as e: + msg = f"Unable to generate Boresch restraint for ABFE simulation: {e}" + _logger.error(msg) + raise RuntimeError(msg) + self._config.restraints = restraints + self._dynamics_kwargs["restraints"] = restraints + import concurrent.futures as _futures import multiprocessing as _mp @@ -546,8 +558,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 @@ -570,6 +582,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}" @@ -624,6 +637,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}" ) @@ -716,12 +738,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 @@ -732,6 +751,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: @@ -790,9 +816,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: @@ -800,7 +828,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() @@ -809,7 +836,12 @@ def generate_lam_vals(lambda_base, increment=0.001): stats = self._load_sampler_stats(index) if stats is not None: if gcmc_sampler is not None and "gcmc" in stats: - gcmc_sampler.restore_stats(stats["gcmc"]) + gcmc_stats = stats["gcmc"] + if self._is_legacy_gcmc_stats(gcmc_stats): + gcmc_stats = self._convert_legacy_gcmc_stats( + gcmc_stats, self._lambda_values[index] + ) + gcmc_sampler.restore_stats(gcmc_stats) if terminal_flip_sampler is not None and "terminal_flip" in stats: attempted, accepted = stats["terminal_flip"] terminal_flip_sampler.reset(attempted, accepted) @@ -1133,9 +1165,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() @@ -1510,6 +1547,7 @@ def _minimisation( device=None, constraint="none", perturbable_constraint="none", + gcmc_sampler=None, ): """ Minimise a system. @@ -1535,6 +1573,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 ------- @@ -1562,6 +1604,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_config.py b/tests/runner/test_config.py index 39bd948..ce0168e 100644 --- a/tests/runner/test_config.py +++ b/tests/runner/test_config.py @@ -83,3 +83,146 @@ def test_logfile_creation(): assert Path.exists(runner._config.output_directory / runner._config.log_file) somd2._logger.remove() + + +def test_morse_restraint_options(): + """Validate that the Morse restraint options are parsed correctly.""" + import math + + import pytest + + # The defaults are parsed as Sire units. + config = Config() + assert config.morse_hard_well_depth == sr.u("150 kcal mol-1") + assert config.morse_soft_well_depth == sr.u("50 kcal mol-1") + assert config.morse_soft_force_constant == sr.u("125 kcal mol-1 A-2") + + # Equivalent units are accepted, and converted. + config = Config(morse_hard_well_depth="418.4 kJ mol-1") + assert math.isclose( + config.morse_hard_well_depth.to(sr.units.kcal_per_mol), 100.0, rel_tol=1e-6 + ) + + # Well depths must be energies. + for option in ("morse_hard_well_depth", "morse_soft_well_depth"): + with pytest.raises(TypeError): + Config(**{option: 150}) + + with pytest.raises(ValueError, match="Unable to parse"): + Config(**{option: "not a unit"}) + + with pytest.raises(ValueError, match="units are invalid"): + Config(**{option: "150 kcal mol-1 A-2"}) + + # The force constant must be an energy per unit area. + with pytest.raises(TypeError): + Config(morse_soft_force_constant=125) + + with pytest.raises(ValueError, match="Unable to parse"): + Config(morse_soft_force_constant="not a unit") + + with pytest.raises(ValueError, match="units are invalid"): + Config(morse_soft_force_constant="125 kcal mol-1") + + +def test_lambda_schedule_input_forms(): + """Validate that all supported lambda schedule input forms are accepted.""" + import os + + import pytest + + schedule = sr.cas.LambdaSchedule.standard_morph() + + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "schedule.s3") + sr.stream.save(schedule, path) + + # A named schedule, which is case insensitive. + config = Config(lambda_schedule="DECOUPLE") + assert config._lambda_schedule_name == "decouple" + + # The path to a stream file. + config = Config(lambda_schedule=path) + assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule) + assert config._lambda_schedule_name is None + + # A hex string of the serialised object. + config = Config(lambda_schedule=Config._to_hex(schedule)) + assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule) + + # The object itself. + config = Config(lambda_schedule=schedule) + assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule) + + # Anything else is rejected. + with pytest.raises(ValueError, match="Unable to interpret"): + Config(lambda_schedule="not_a_schedule") + + # A stream file holding the wrong type of object. + wrong_path = os.path.join(tmpdir, "wrong.s3") + sr.stream.save(sr.cas.Symbol("x"), wrong_path) + with pytest.raises(ValueError, match="not a 'LambdaSchedule'"): + Config(lambda_schedule=wrong_path) + + +def test_restraints_input_forms(): + """Validate that all supported restraint input forms are accepted.""" + import os + + import pytest + + mols = sr.load_test_files("ala.top", "ala.crd") + restraint0 = sr.restraints.positional(mols, atoms="atomidx 0") + restraint1 = sr.restraints.positional(mols, atoms="atomidx 1") + + with tempfile.TemporaryDirectory() as tmpdir: + path0 = os.path.join(tmpdir, "restraint0.s3") + both_path = os.path.join(tmpdir, "both.s3") + sr.stream.save(restraint0, path0) + sr.stream.save([restraint0, restraint1], both_path) + + # A single object, or a list of objects. + assert len(Config(restraints=restraint0).restraints) == 1 + assert len(Config(restraints=[restraint0, restraint1]).restraints) == 2 + + # The path to a stream file, or a list of paths. + assert len(Config(restraints=path0).restraints) == 1 + assert len(Config(restraints=[path0, path0]).restraints) == 2 + + # A stream file holding a list of sets of restraints. + assert len(Config(restraints=both_path).restraints) == 2 + + # A hex string of the serialised object. + assert len(Config(restraints=Config._to_hex(restraint0)).restraints) == 1 + + # Objects and paths can be mixed, and all are retained. + config = Config(restraints=[restraint0, path0]) + assert len(config.restraints) == 2 + assert all( + isinstance(restraint, sr.mm._MM.Restraints) + for restraint in config.restraints + ) + + # Anything else is rejected. + with pytest.raises(ValueError, match="Unable to interpret"): + Config(restraints="not_a_restraint") + + # A stream file holding the wrong type of object. + wrong_path = os.path.join(tmpdir, "wrong.s3") + sr.stream.save(sr.cas.LambdaSchedule.standard_morph(), wrong_path) + with pytest.raises(ValueError, match="must be a sire.mm._MM.Restraints"): + Config(restraints=wrong_path) + + +def test_help_text_scraping(): + """Validate that help text isn't truncated by the parameter name.""" + parser = Config._create_parser() + + for action in parser._actions: + if action.dest == "restraints": + break + + # The description wraps onto a line starting with the parameter name, which + # must not be mistaken for the start of the next parameter. + assert "applied to the atoms" in action.help + assert "a list of sets" in action.help 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_gcmc_stats.py b/tests/runner/test_gcmc_stats.py new file mode 100644 index 0000000..c6ff648 --- /dev/null +++ b/tests/runner/test_gcmc_stats.py @@ -0,0 +1,97 @@ +import numpy as np +import pytest + +from somd2.runner._base import RunnerBase +from somd2.runner._repex import DynamicsCache + + +def counters(num_moves=0): + """A flat dictionary of counters, as written before the format changed.""" + return { + "num_moves": num_moves, + "num_accepted": num_moves, + "num_insertions": 0, + "num_deletions": 0, + "num_accepted_attempts": 0, + } + + +class TestLegacyDetection: + """Tests for detecting the format of GCMC statistics.""" + + def test_flat_counters_are_legacy(self): + assert RunnerBase._is_legacy_gcmc_stats(counters(5)) + + def test_keyed_by_lambda_is_current(self): + assert not RunnerBase._is_legacy_gcmc_stats({"0.00000": counters(5)}) + + @pytest.mark.parametrize("stats", [None, {}, []]) + def test_other_values_are_not_legacy(self, stats): + assert not RunnerBase._is_legacy_gcmc_stats(stats) + + def test_conversion_keys_by_lambda(self): + converted = RunnerBase._convert_legacy_gcmc_stats(counters(5), 0.33333) + assert converted == {"0.33333": counters(5)} + + def test_conversion_is_a_copy(self): + """The original must not be aliased into the converted result.""" + original = counters(5) + converted = RunnerBase._convert_legacy_gcmc_stats(original, 0.0) + original["num_moves"] = 99 + assert converted["0.00000"]["num_moves"] == 5 + + +class TestLegacyRepexCheckpoint: + """Tests for restoring a replica exchange checkpoint.""" + + @staticmethod + def make_state(gcmc_stats, lambdas=(0.0, 0.5, 1.0)): + n = len(lambdas) + return { + "_lambdas": list(lambdas), + "_rest2_scale_factors": [1.0] * n, + "_states": np.arange(n), + "_time": None, + "_openmm_states": [None] * n, + "_gcmc_samplers": [None] * n, + "_gcmc_states": [None] * n, + "_gcmc_stats": gcmc_stats, + "_terminal_flip_stats": [[0, 0]] * n, + "_num_proposed": np.zeros((n, n)), + "_num_accepted": np.zeros((n, n)), + "_num_swaps": np.zeros((n, n)), + } + + def restore(self, gcmc_stats, **kwargs): + cache = object.__new__(DynamicsCache) + cache.__setstate__(self.make_state(gcmc_stats, **kwargs)) + return cache + + def test_per_replica_list_is_converted(self): + """A list of counters per replica becomes a map keyed by lambda.""" + cache = self.restore([counters(i) for i in range(3)]) + + assert cache._gcmc_stats == { + "0.00000": counters(0), + "0.50000": counters(1), + "1.00000": counters(2), + } + + def test_no_gcmc_gives_none(self): + """A checkpoint from a run without GCMC has no statistics.""" + assert self.restore([None, None, None])._gcmc_stats is None + + def test_current_format_is_untouched(self): + """A checkpoint already in the current format is left alone.""" + stats = {"0.00000": counters(4), "1.00000": counters(9)} + assert self.restore(dict(stats))._gcmc_stats == stats + + def test_missing_attribute_defaults_to_none(self): + """A checkpoint predating GCMC statistics has none.""" + state = self.make_state(None) + del state["_gcmc_stats"] + + cache = object.__new__(DynamicsCache) + cache.__setstate__(state) + + assert cache._gcmc_stats is None 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}" + ) diff --git a/tests/runner/test_ring_break_restraints.py b/tests/runner/test_ring_break_restraints.py new file mode 100644 index 0000000..255931a --- /dev/null +++ b/tests/runner/test_ring_break_restraints.py @@ -0,0 +1,232 @@ +import tempfile + +import pytest + +from somd2.config import Config +from somd2.runner import Runner + + +def _config(tmpdir, **kwargs): + """Return a minimal ring-breaking config rooted at 'tmpdir'.""" + options = { + "output_directory": tmpdir, + "lambda_schedule": "ring_break_morph", + "num_lambda": 3, + "runtime": "12fs", + "energy_frequency": "4fs", + "frame_frequency": "4fs", + "checkpoint_frequency": "4fs", + "equilibration_time": "0fs", + "minimise": False, + "platform": "CPU", + "max_threads": 1, + } + options.update(kwargs) + return Config(**options) + + +def _restraint(restraints, name): + """Return the single restraint from the set called 'name'.""" + for restraint_set in restraints: + if str(restraint_set.name()) == name: + assert len(restraint_set) == 1 + return restraint_set[0] + raise AssertionError(f"No restraint set named {name!r} in {restraints}") + + +def test_restraints_are_generated(syk_ring_break_mols): + """ + Ensure that a pair of Morse restraints is automatically generated for a + ring-breaking perturbation when no restraint is supplied, and that they act + on the same pair of atoms, at the same equilibrium distance. + """ + with tempfile.TemporaryDirectory() as tmpdir: + runner = Runner(syk_ring_break_mols.clone(), _config(tmpdir)) + + restraints = runner._config.restraints + assert restraints is not None + assert len(restraints) == 2 + + hard = _restraint(restraints, "morse_hard") + soft = _restraint(restraints, "morse_soft") + + # Both restraints act on the bond that is broken. + assert hard.atom0() == soft.atom0() + assert hard.atom1() == soft.atom1() + assert hard.r0() == soft.r0() + + # The restraints are passed through to the dynamics. + assert runner._dynamics_kwargs["restraints"] is restraints + + +def test_restraints_match_config(syk_ring_break_mols): + """ + Ensure that the generated restraints use the well depths and force constant + from the config, and that the hard restraint inherits the force constant of + the bond that it replaces. + """ + import sire as sr + + with tempfile.TemporaryDirectory() as tmpdir: + config = _config( + tmpdir, + morse_hard_well_depth="123 kcal mol-1", + morse_soft_well_depth="45 kcal mol-1", + morse_soft_force_constant="67 kcal mol-1 A-2", + ) + runner = Runner(syk_ring_break_mols.clone(), config) + + hard = _restraint(runner._config.restraints, "morse_hard") + soft = _restraint(runner._config.restraints, "morse_soft") + + assert hard.de() == sr.u("123 kcal mol-1") + assert soft.de() == sr.u("45 kcal mol-1") + assert soft.k() == sr.u("67 kcal mol-1 A-2") + + # The hard restraint is auto-parametrised from the broken bond, so its + # force constant comes from the bond, not the config. + assert hard.k() != soft.k() + assert hard.k().value() > 0 + + +def test_broken_bond_is_replaced(syk_ring_break_mols): + """ + Ensure that the hard restraint replaces the harmonic bond that is broken by + the perturbation, i.e. that the bond is removed from the runner's system. + Leaving both in place would double count the interaction. + """ + with tempfile.TemporaryDirectory() as tmpdir: + mols = syk_ring_break_mols.clone() + runner = Runner(mols, _config(tmpdir)) + + hard = _restraint(runner._config.restraints, "morse_hard") + + def num_bonds(system, idx0, idx1): + """Count the bond potentials between a pair of atom indices.""" + atoms = system.atoms() + atom0 = atoms[idx0] + atom1 = atoms[idx1] + mol = system[atom0.molecule().number()] + info = mol.info() + expected = {atom0.index().value(), atom1.index().value()} + + count = 0 + for bond_prop in ("bond0", "bond1"): + for potential in mol.property(bond_prop).potentials(): + idxs = { + info.atom_idx(potential.atom0()).value(), + info.atom_idx(potential.atom1()).value(), + } + if idxs == expected: + count += 1 + return count + + # The unmodified input still has the bond, in the reference end state + # only, since it is broken by the perturbation. + assert num_bonds(syk_ring_break_mols, hard.atom0(), hard.atom1()) == 1 + + # The runner's system has it removed, replaced by the Morse restraint. + assert num_bonds(runner._system, hard.atom0(), hard.atom1()) == 0 + + +def test_restraints_are_deterministic(syk_ring_break_mols): + """ + Ensure that generating the restraints twice from the same input gives + identical restraints. A restart regenerates them from the input system + rather than reloading them, so they must not drift between runs, otherwise + the accumulated free energy would be invalidated. + """ + with tempfile.TemporaryDirectory() as tmpdir0: + runner0 = Runner(syk_ring_break_mols.clone(), _config(tmpdir0)) + + with tempfile.TemporaryDirectory() as tmpdir1: + runner1 = Runner(syk_ring_break_mols.clone(), _config(tmpdir1)) + + for name in ("morse_hard", "morse_soft"): + assert _restraint(runner0._config.restraints, name) == _restraint( + runner1._config.restraints, name + ) + + +def test_reverse_schedule_generates_restraints(syk_ring_break_mols): + """ + Ensure that restraints are also generated for the ring-making direction, + which uses the reversed schedule. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = _config(tmpdir, lambda_schedule="reverse_ring_break_morph") + runner = Runner(syk_ring_break_mols.clone(), config) + + assert runner._config.restraints is not None + assert len(runner._config.restraints) == 2 + + +def test_user_restraints_are_not_overridden(syk_ring_break_mols): + """ + Ensure that a user-supplied restraint is left alone, and that the system is + not modified behind their back. + """ + import sire as sr + + mols = syk_ring_break_mols.clone() + + restraints = sr.restraints.distance( + mols, + atoms0=0, + atoms1=1, + k="10 kcal mol-1 A-2", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + runner = Runner(mols, _config(tmpdir, restraints=restraints)) + + assert len(runner._config.restraints) == 1 + assert runner._config.restraints[0] == restraints + + +@pytest.mark.parametrize("schedule", ["standard_morph", "charge_scaled_morph"]) +def test_no_restraints_for_other_schedules(schedule, ethane_methanol): + """ + Ensure that Morse restraints are only generated for ring-breaking + schedules. + """ + with tempfile.TemporaryDirectory() as tmpdir: + config = _config(tmpdir, lambda_schedule=schedule) + runner = Runner(ethane_methanol.clone(), config) + + assert runner._config.restraints is None + + +def test_no_broken_bond_raises(ethane_methanol): + """ + Ensure that a clear error is raised when a ring-breaking schedule is used + for a perturbation that doesn't break (or form) a bond. + """ + with tempfile.TemporaryDirectory() as tmpdir: + with pytest.raises(RuntimeError, match="Unable to generate Morse restraints"): + Runner(ethane_methanol.clone(), _config(tmpdir)) + + +def test_already_applied_raises(syk_ring_break_mols): + """ + Ensure that a helpful error is raised if the Morse potential has already + been applied to the input system, but the corresponding restraints were not + passed via the config. The replacement must not be applied twice. + """ + import sire as sr + + mols = syk_ring_break_mols.clone() + + # Apply the Morse replacement, as a user following the existing workflow + # would, but don't pass the restraints to the config. + _, mols = sr.restraints.morse_potential( + mols, + de="150 kcal mol-1", + auto_parametrise=True, + direct_morse_replacement=True, + name="morse_hard", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + with pytest.raises(RuntimeError, match="already been applied"): + Runner(mols, _config(tmpdir)) diff --git a/tests/schedules/test_abfe.py b/tests/schedules/test_abfe.py new file mode 100644 index 0000000..e44a25a --- /dev/null +++ b/tests/schedules/test_abfe.py @@ -0,0 +1,153 @@ +import pytest + +from somd2._utils._schedules import annihilate, decouple + +# Lambda schedules are always symmetric between the two stages (decharge is +# the first stage, annihilate/decouple is the second), so both builders share +# identical expected lever values. +BUILDERS = [annihilate, decouple] + +_LAMBDA_VALUES = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] + +# Expected "restraint" lever values (restraint_lever="combined"): ramps +# linearly 0 -> 1 across the whole first stage (decharge), then held at 1. +_COMBINED_RESTRAINT = [0.0, 0.4, 0.8, 1.0, 1.0, 1.0] + +# Expected "restraint_dihedral"/"restraint_distance_angle" lever values +# (restraint_lever="split"): a geometric progression from ~0.01 to 1.0 across +# each stage in turn, reproducing the ratio in Table S1 of the RXRX paper's +# SI. dihedral ramps during the first stage then holds at 1; distance_angle +# holds at 0 during the first stage then ramps during the second. +_SPLIT_DIHEDRAL = [0.01, 0.06309573444801934, 0.39810717055349737, 1.0, 1.0, 1.0] +_SPLIT_DISTANCE_ANGLE = [0.0, 0.0, 0.0, 0.025118864315095805, 0.15848931924611143, 1.0] + + +def _morph(schedule, lever, lambda_value): + """ + Query a lever's value exactly as SireOpenMM's LambdaLever does at + runtime: lambda_schedule.morph("*", restraint_name, 1.0, 1.0, lambda_value). + """ + return schedule.morph("*", lever, 1.0, 1.0, lambda_value) + + +@pytest.mark.parametrize("builder", BUILDERS) +def test_restraint_lever_defaults_to_split(builder): + """ + annihilate()/decouple() default to restraint_lever="split", matching + boresch_search()'s own default of restraint_lever="split" for its + default protocol="rxrx" + """ + schedule = builder() + assert "restraint_dihedral" in schedule.get_levers() + assert "restraint_distance_angle" in schedule.get_levers() + assert "restraint" not in schedule.get_levers() + + +@pytest.mark.parametrize("builder", BUILDERS) +def test_restraint_lever_invalid_raises(builder): + with pytest.raises(ValueError, match="restraint_lever"): + builder(restraint_lever="not_a_real_lever") + + +@pytest.mark.parametrize("builder", BUILDERS) +def test_restraint_lever_combined(builder): + """ + restraint_lever="combined" sets a single "restraint" lever that ramps + 0 -> 1 across the first stage (decharge), then holds at 1. + """ + schedule = builder(restraint_lever="combined") + assert "restraint" in schedule.get_levers() + assert "restraint_dihedral" not in schedule.get_levers() + assert "restraint_distance_angle" not in schedule.get_levers() + + values = [_morph(schedule, "restraint", lv) for lv in _LAMBDA_VALUES] + assert values == pytest.approx(_COMBINED_RESTRAINT, abs=1e-6) + + +@pytest.mark.parametrize("builder", BUILDERS) +def test_restraint_lever_split(builder): + """ + restraint_lever="split" sets two independent levers ("restraint_dihedral" + and "restraint_distance_angle"), each following a geometric progression + across its own stage while the other is held fixed, reproducing the RXRX + protocol's staged restraint turn-on. + """ + schedule = builder(restraint_lever="split") + + dihedral_values = [ + _morph(schedule, "restraint_dihedral", lv) for lv in _LAMBDA_VALUES + ] + distance_angle_values = [ + _morph(schedule, "restraint_distance_angle", lv) for lv in _LAMBDA_VALUES + ] + + assert dihedral_values == pytest.approx(_SPLIT_DIHEDRAL, abs=1e-6) + assert distance_angle_values == pytest.approx(_SPLIT_DISTANCE_ANGLE, abs=1e-6) + + +@pytest.mark.skipif( + "openmm" not in __import__("sire").convert.supported_formats(), + reason="openmm support is not available", +) +def test_restraint_lever_split_openmm_system(): + """ + End-to-end regression test: build a real dynamics object with a "split" + Boresch restraint and confirm the two independently-lambda-addressable + OpenMM Forces (distance/angle and dihedral) exist, with 'rho' values + that follow the expected geometric progression as lambda changes, + matching what test_restraint_lever_split checks at the schedule level. + """ + import xml.etree.ElementTree as ET + + import sire as sr + + mols = sr.load_test_files("boresch_restraints.prm7", "boresch_restraints.dcd") + mols.update(sr.morph.decouple(mols.molecule(1), as_new_molecule=False)) + + restraints = sr.restraints.boresch( + mols, + receptor=[692, 702, 704], + ligand=[1496, 1498, 1499], + kr="1 kcal mol-1 A-2", + ktheta=["80 kcal mol-1 rad-2"] * 2, + kphi=["80 kcal mol-1 rad-2"] * 3, + r0="4.56908 A", + theta0=["82.5581 degrees", "94.9595 degrees"], + phi0=["27.9429 degrees", "125.68 degrees", "-107.008 degrees"], + angle_potential="restricted_bending", + restraint_lever="split", + ) + + schedule = decouple(restraint_lever="split") + + d = mols.dynamics( + timestep="2fs", + temperature="298 K", + schedule=schedule, + lambda_value=0.0, + map={"restraints": restraints}, + ) + + expected_rho = { + 0.0: {"distance_angle": 0.0, "dihedral": 0.01}, + 0.5: {"distance_angle": 0.01, "dihedral": 1.0}, + 1.0: {"distance_angle": 1.0, "dihedral": 1.0}, + } + + for lam, expected in expected_rho.items(): + d.set_lambda(lam) + root = ET.fromstring(d.to_xml()) + + rhos = {} + for force in root.iter("Force"): + if force.get("name") == "BoreschRestraintForce": + kind = ( + "distance_angle" if "e_bond" in force.get("energy") else "dihedral" + ) + rhos[kind] = float(force.find("Bonds")[0].get("param1")) + + assert set(rhos.keys()) == {"distance_angle", "dihedral"} + assert rhos["distance_angle"] == pytest.approx( + expected["distance_angle"], abs=1e-6 + ) + assert rhos["dihedral"] == pytest.approx(expected["dihedral"], abs=1e-6)