From cae3bf875b46ee9d50c28c959a49e86c659f5c6e Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 20 Mar 2026 10:22:38 +0000 Subject: [PATCH] Apply saved replica exchange mixing on restart to restore post-mix states. --- src/somd2/runner/_repex.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 33ef339..6444811 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -812,15 +812,17 @@ def __init__(self, system, config): for i in range(len(self._lambda_values)): dynamics, gcmc_sampler = self._dynamics_cache.get(i) - # Reset the OpenMM state. - dynamics.context().setState(self._dynamics_cache._openmm_states[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] + dynamics.context().setState(self._dynamics_cache._openmm_states[state]) # Reset the GCMC water state. if gcmc_sampler is not None: gcmc_sampler.push() gcmc_sampler._set_water_state( dynamics.context(), - states=self._dynamics_cache._gcmc_states[i], + states=self._dynamics_cache._gcmc_states[state], force=True, ) gcmc_sampler.pop()