diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 6444811..523e7d7 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -532,12 +532,14 @@ def mix_states(self): # Update the water state in the GCMCSampler. self._gcmc_samplers[i].push() - self._gcmc_samplers[i]._set_water_state( - self._dynamics[i].context(), - indices=water_idxs, - states=self._gcmc_states[state][water_idxs], - ) - self._gcmc_samplers[i].pop() + try: + self._gcmc_samplers[i]._set_water_state( + self._dynamics[i].context(), + indices=water_idxs, + states=self._gcmc_states[state][water_idxs], + ) + finally: + self._gcmc_samplers[i].pop() # Update the swap matrix. old_state = self._old_states[i] @@ -820,12 +822,14 @@ def __init__(self, system, config): # 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[state], - force=True, - ) - gcmc_sampler.pop() + try: + gcmc_sampler._set_water_state( + dynamics.context(), + states=self._dynamics_cache._gcmc_states[state], + force=True, + ) + finally: + gcmc_sampler.pop() # Conversion factor for reduced potential. kT = (_sr.units.k_boltz * self._config.temperature).to(_sr.units.kcal_per_mol) @@ -1277,13 +1281,13 @@ def _run_block( 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() + try: + # Perform the GCMC move. + _logger.info(f"Performing GCMC move at {_lam_sym} = {lam:.5f}") + gcmc_sampler.move(dynamics.context()) + finally: + # Remove the PyCUDA context from the stack. + gcmc_sampler.pop() # Save the GCMC state. self._dynamics_cache.save_gcmc_state(index) @@ -1340,15 +1344,15 @@ def _minimise(self, index): if gcmc_sampler is not None: # Push the PyCUDA context on top of the stack. gcmc_sampler.push() - - _logger.info( - f"Pre-equilibrating with GCMC moves at {_lam_sym} = {self._lambda_values[index]:.5f}" - ) - for i in range(100): - gcmc_sampler.move(dynamics.context()) - - # Remove the PyCUDA context from the stack. - gcmc_sampler.pop() + try: + _logger.info( + f"Pre-equilibrating with GCMC moves at {_lam_sym} = {self._lambda_values[index]:.5f}" + ) + for i in range(100): + gcmc_sampler.move(dynamics.context()) + finally: + # Remove the PyCUDA context from the stack. + gcmc_sampler.pop() # Minimise. dynamics.minimise(timeout=self._config.timeout) @@ -1433,15 +1437,15 @@ def _equilibrate(self, index): if gcmc_sampler is not None: # Push the PyCUDA context on top of the stack. gcmc_sampler.push() - - _logger.info( - f"Equilibrating with GCMC moves at {_lam_sym} = {self._lambda_values[index]:.5f}" - ) - for i in range(100): - gcmc_sampler.move(dynamics.context()) - - # Remove the PyCUDA context from the stack. - gcmc_sampler.pop() + try: + _logger.info( + f"Equilibrating with GCMC moves at {_lam_sym} = {self._lambda_values[index]:.5f}" + ) + for i in range(100): + gcmc_sampler.move(dynamics.context()) + finally: + # Remove the PyCUDA context from the stack. + gcmc_sampler.pop() # Store the current water state. water_state = gcmc_sampler.water_state() @@ -1693,14 +1697,14 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): if gcmc_sampler is not None: # Push the PyCUDA context on top of the stack. gcmc_sampler.push() - - _logger.info( - f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} " - f"is {gcmc_sampler.num_waters()}" - ) - - # Remove the PyCUDA context from the stack. - gcmc_sampler.pop() + try: + _logger.info( + f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} " + f"is {gcmc_sampler.num_waters()}" + ) + finally: + # Remove the PyCUDA context from the stack. + gcmc_sampler.pop() if is_final_block: _logger.success(f"{_lam_sym} = {lam:.5f} complete") @@ -1823,12 +1827,12 @@ def _reset_gcmc_sampler(gcmc_sampler, dynamics): # Push the PyCUDA context on top of the stack. gcmc_sampler.push() - - # Set the water state. - gcmc_sampler._set_water_state(dynamics.context(), force=True) - - # Remove the PyCUDA context from the stack. - gcmc_sampler.pop() + try: + # Set the water state. + gcmc_sampler._set_water_state(dynamics.context(), force=True) + finally: + # Remove the PyCUDA context from the stack. + gcmc_sampler.pop() # Re-bind the GCMC sampler to the dynamics object. gcmc_sampler.bind_dynamics(dynamics) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index fb8b625..14ad6d9 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -521,8 +521,12 @@ def generate_lam_vals(lambda_base, increment=0.001): f"Equilibrating with GCMC moves at {_lam_sym} = {lambda_value:.5f}" ) - for i in range(100): - gcmc_sampler.move(dynamics.context()) + gcmc_sampler.push() + try: + for i in range(100): + gcmc_sampler.move(dynamics.context()) + finally: + gcmc_sampler.pop() # Run without saving energies or frames. dynamics.run( @@ -610,20 +614,24 @@ def generate_lam_vals(lambda_base, increment=0.001): if self._is_restart: from openmm.unit import angstrom - # First set all waters to non-ghosts. - gcmc_sampler._set_water_state( - dynamics.context(), - states=_np.ones(len(gcmc_sampler._water_indices)), - force=True, - ) + gcmc_sampler.push() + try: + # First set all waters to non-ghosts. + gcmc_sampler._set_water_state( + dynamics.context(), + states=_np.ones(len(gcmc_sampler._water_indices)), + force=True, + ) - # Now set the ghost waters. - gcmc_sampler._set_water_state( - dynamics.context(), - self._restart_ghost_waters[index], - states=_np.zeros(len(gcmc_sampler._water_indices)), - force=True, - ) + # Now set the ghost waters. + gcmc_sampler._set_water_state( + dynamics.context(), + self._restart_ghost_waters[index], + states=_np.zeros(len(gcmc_sampler._water_indices)), + force=True, + ) + finally: + gcmc_sampler.pop() # Finally, reset the context positions to match the restart system. dynamics.context().setPositions( @@ -634,10 +642,14 @@ def generate_lam_vals(lambda_base, increment=0.001): # the water state in the new context to match the equilibrated system. elif is_equilibrated: # Reset the water state. - gcmc_sampler._set_water_state( - dynamics.context(), - force=True, - ) + gcmc_sampler.push() + try: + gcmc_sampler._set_water_state( + dynamics.context(), + force=True, + ) + finally: + gcmc_sampler.pop() # Set the number of neighbours used for the energy calculation. # If not None, then we add one to account for the extra windows @@ -728,7 +740,11 @@ def generate_lam_vals(lambda_base, increment=0.001): _logger.info( f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}" ) - gcmc_sampler.move(dynamics.context()) + gcmc_sampler.push() + try: + gcmc_sampler.move(dynamics.context()) + finally: + gcmc_sampler.pop() else: dynamics.run( @@ -816,10 +832,14 @@ def generate_lam_vals(lambda_base, increment=0.001): # Log the number of waters within the GCMC sampling volume. if gcmc_sampler is not None: - _logger.info( - f"Current number of waters in GCMC volume at {_lam_sym} = {lambda_value:.5f} " - f"is {gcmc_sampler.num_waters()}" - ) + gcmc_sampler.push() + try: + _logger.info( + f"Current number of waters in GCMC volume at {_lam_sym} = {lambda_value:.5f} " + f"is {gcmc_sampler.num_waters()}" + ) + finally: + gcmc_sampler.pop() if is_final_block: _logger.success( @@ -928,7 +948,11 @@ def generate_lam_vals(lambda_base, increment=0.001): _logger.info( f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}" ) - gcmc_sampler.move(dynamics.context()) + gcmc_sampler.push() + try: + gcmc_sampler.move(dynamics.context()) + finally: + gcmc_sampler.pop() # Update the runtime. runtime += self._config.energy_frequency