Skip to content

Add support for caching and re-using dynamics contexts - #191

Merged
lohedges merged 24 commits into
develfrom
feature_dynamics_cache
Aug 5, 2026
Merged

Add support for caching and re-using dynamics contexts#191
lohedges merged 24 commits into
develfrom
feature_dynamics_cache

Conversation

@lohedges

@lohedges lohedges commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

This PR extends the DynamicsCache used for replica exchange with proper caching, allowing contexts to be re-used for simulations at lambda values other than which they were created. This allows us to run replica exchange on more modest hardware, where there isn't sufficient memory to create all contexts up front, as well as offering the ability to run simulations with a very large number of replicas, e.g. large ABFE simulations.

The user can now specify a max_contexts option which limits the total number of contexts that are created. When max_contexts is None, then the existing behaviour is preserved. Setting max_contexts to less than the number of replicas means that contexts will be re-used. In this situation, we have added book-keeping to store the energy trajectory for each replica so that we can swap it on re-use. We then call set_lambda(), delete_all_frames(), and re-set the integrator clock to the start of the repex cycle. In situations where max_contexts is equal to the num_gpus * oversubscription_factor, there is no performance loss for using the new approach, i.e. you only pay a hit when there are GPU resources that are under utilised. (We warn when this is the case.) When using all replicas, we now suggest an appropriate max_contexts if GPU memory is exhausted on setup. (We don't automatically fall back to the max_contexts pathway due to the two caveats below.)

Running with a restricted number of contexts comes with two caveats:

  1. It is not possible (without effort and performance loss) to cache trajectory frames, so this approach is limited to a frame_frequency equal to the checkpoint_frequency, which is a common production use case anyway.
  2. Re-using a context for a simulation at another lambda involves adjusting force field parameters, which could potentially involve a change in constraints. When this happens, the context needs a rebuild, which massively impacts performance. To mitigate this, I have added an update_constraints option that specifies whether constraints should be updated when max_contexts is used. (By default they are.) If False, then the user can additionally supply a constraint_lambda_index option that specifies the replica that should be used for the constraint value. I intend to see if there is a good default for this, e.g. should we use the lambda value where the constraint is at its mid-point, or one of the end-states? In previous work I found that, when running simulations in forward/reverse, you'd get the correct result when pinning the constraint to one end-state, i.e. lambda=0 for forwards and lambda=1 for reverse, then averaging the ddG values.

The PR also makes the GCMC equilibration pathway consistent for the Runner and RepexRunner, as well as removing minimisation on restart, which was inconsistent with the other restart guards.

@lohedges lohedges added the enhancement New feature or request label Aug 4, 2026
@lohedges
lohedges merged commit 8504cc5 into devel Aug 5, 2026
6 of 8 checks passed
@lohedges
lohedges deleted the feature_dynamics_cache branch August 5, 2026 11:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant