Fix GCMC trajectory synchronisation

src/somd2/runner/_repex.py

@@ -1242,28 +1242,6 @@ def _run_block(
             # Draw new velocities from the Maxwell-Boltzmann distribution.
             dynamics.randomise_velocities()
 
-            # Perform a GCMC move. For repex this needs to be done before the
-            # dynamics block so that the final energies, which are used in the
-            # repex acceptance criteria, are correct.
-            if is_gcmc and gcmc_sampler is not None:
-                # Push the PyCUDA context on top of the stack.
-                gcmc_sampler.push()
-
-                # Perform the GCMC move.
-                _logger.info(f"Performing GCMC move at {_lam_sym} = {lam:.5f}")
-                gcmc_sampler.move(dynamics.context())
-
-                # Remove the PyCUDA context from the stack.
-                gcmc_sampler.pop()
-
-                # A frame was saved at the end of the last cycle, so write
-                # the current ghost water residue indices to file. This is
-                # done here, immediately after the GCMC move, since the
-                # sampler state is only updated during GCMC moves and waters
-                # may have moved in/out of the GCMC sphere during dynamics.
-                if write_gcmc_ghosts:
-                    gcmc_sampler.write_ghost_residues()
-
             # Run the dynamics.
             dynamics.run(
                 self._config.energy_frequency,
@@ -1287,13 +1265,30 @@ def _run_block(
                 ),
             )
 
-            # Set the state.
-            self._dynamics_cache.save_openmm_state(index)
-
-            # Save the GCMC state.
             if gcmc_sampler is not None:
+                # Write ghost residues before the GCMC move so the ghost state
+                # is consistent with the saved frame (which is also captured
+                # before the GCMC move).
+                if write_gcmc_ghosts:
+                    gcmc_sampler.write_ghost_residues()
+
+                if is_gcmc:
+                    # Push the PyCUDA context on top of the stack.
+                    gcmc_sampler.push()
+
+                    # Perform the GCMC move.
+                    _logger.info(f"Performing GCMC move at {_lam_sym} = {lam:.5f}")
+                    gcmc_sampler.move(dynamics.context())
+
+                    # Remove the PyCUDA context from the stack.
+                    gcmc_sampler.pop()
+
+                # Save the GCMC state.
                 self._dynamics_cache.save_gcmc_state(index)
 
+            # Save the OpenMM state after any GCMC move so the context is consistent.
+            self._dynamics_cache.save_openmm_state(index)
+
             # Get the energy at each lambda value.
             energies = dynamics._current_energy_array()
 

src/somd2/runner/_runner.py

@@ -690,7 +690,7 @@ def generate_lam_vals(lambda_base, increment=0.001):
                         next_frame = self._config.frame_frequency
 
                         # Loop until we reach the runtime.
-                        while runtime <= checkpoint_frequency:
+                        while runtime < checkpoint_frequency:
                             # Run the dynamics in blocks of the GCMC frequency.
                             dynamics.run(
                                 self._config.gcmc_frequency,
@@ -714,21 +714,22 @@ def generate_lam_vals(lambda_base, increment=0.001):
                                 ),
                             )
 
-                            # Perform a GCMC move.
-                            _logger.info(
-                                f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}"
-                            )
-                            gcmc_sampler.move(dynamics.context())
-
                             # Update the runtime.
                             runtime += self._config.energy_frequency
 
-                            # If a frame is saved, then we need to save current indices
-                            # of the ghost water residues.
+                            # If a frame is saved, write the ghost residue indices
+                            # before the GCMC move so the ghost state is consistent
+                            # with the saved frame.
                             if save_frames and runtime >= next_frame:
                                 gcmc_sampler.write_ghost_residues()
                                 next_frame += self._config.frame_frequency
 
+                            # Perform a GCMC move.
+                            _logger.info(
+                                f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}"
+                            )
+                            gcmc_sampler.move(dynamics.context())
+
                     else:
                         dynamics.run(
                             checkpoint_frequency,