Skip to content

Wrap GCMC push/pop pairs in try/finally to guard against crashes #131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lohedges merged 1 commit into from
Mar 24, 2026
Merged

Wrap GCMC push/pop pairs in try/finally to guard against crashes #131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 55 additions & 51 deletions src/somd2/runner/_repex.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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)
74 changes: 49 additions & 25 deletions src/somd2/runner/_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down
Loading