diff --git a/.img/somd2.png b/.img/somd2.png index df14e91..b2e9c40 100644 Binary files a/.img/somd2.png and b/.img/somd2.png differ diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 84e01f9..b0d5663 100644 --- a/src/somd2/runner/_repex.py +++ b/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() diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index f5fcf7f..fb8b625 100644 --- a/src/somd2/runner/_runner.py +++ b/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,