Skip to content

Perform GCMC before dynamics so energies correspond to correct state #133

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 27, 2026
Merged

Perform GCMC before dynamics so energies correspond to correct state #133

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
35 changes: 17 additions & 18 deletions src/somd2/runner/_repex.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1242,6 +1242,22 @@ def _run_block(
# Get the dynamics object (and GCMC sampler).
dynamics, gcmc_sampler = self._dynamics_cache.get(index)

# Perform the GCMC move before dynamics so that the energies
# computed during dynamics are consistent with the state used
# for replica exchange mixing.
if gcmc_sampler is not None and is_gcmc:
gcmc_sampler.push()
try:
_logger.info(f"Performing GCMC move at {_lam_sym} = {lam:.5f}")
gcmc_sampler.move(dynamics.context())
finally:
gcmc_sampler.pop()

# Write ghost residues immediately after the GCMC move so the
# ghost state and frame (saved during dynamics) are consistent.
if write_gcmc_ghosts:
gcmc_sampler.write_ghost_residues()

_logger.info(f"Running dynamics at {_lam_sym} = {lam:.5f}")

# Draw new velocities from the Maxwell-Boltzmann distribution.
Expand Down Expand Up @@ -1272,27 +1288,10 @@ def _run_block(
)

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()
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)

# Save the OpenMM state after any GCMC move so the context is consistent.
# Save the OpenMM state.
self._dynamics_cache.save_openmm_state(index)

# Get the energy at each lambda value.
Expand Down
80 changes: 46 additions & 34 deletions src/somd2/runner/_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -722,6 +722,29 @@ def generate_lam_vals(lambda_base, increment=0.001):

# Loop until we reach the runtime.
while runtime < checkpoint_frequency:
# Perform a GCMC move before dynamics so the ghost
# state is consistent with the energies computed
# during dynamics.
_logger.info(
f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}"
)
gcmc_sampler.push()
try:
gcmc_sampler.move(dynamics.context())
finally:
gcmc_sampler.pop()

# Write ghost residues immediately after the GCMC
# move if a frame will be saved in the upcoming
# dynamics block.
if (
save_frames
and runtime + self._config.energy_frequency
>= next_frame
):
gcmc_sampler.write_ghost_residues()
next_frame += self._config.frame_frequency

# Run the dynamics in blocks of the GCMC frequency.
dynamics.run(
self._config.gcmc_frequency,
Expand All @@ -748,23 +771,6 @@ def generate_lam_vals(lambda_base, increment=0.001):
# Update the runtime.
runtime += self._config.energy_frequency

# 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.push()
try:
gcmc_sampler.move(dynamics.context())
finally:
gcmc_sampler.pop()

else:
dynamics.run(
checkpoint_frequency,
Expand Down Expand Up @@ -948,7 +954,29 @@ def generate_lam_vals(lambda_base, increment=0.001):
next_frame = self._config.frame_frequency

# Loop until we reach the runtime.
while runtime <= time:
while runtime < time:
# Perform a GCMC move before dynamics so the ghost
# state is consistent with the energies computed
# during dynamics.
_logger.info(
f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}"
)
gcmc_sampler.push()
try:
gcmc_sampler.move(dynamics.context())
finally:
gcmc_sampler.pop()

# Write ghost residues immediately after the GCMC
# move if a frame will be saved in the upcoming
# dynamics block.
if (
save_frames
and runtime + self._config.energy_frequency >= next_frame
):
gcmc_sampler.write_ghost_residues()
next_frame += self._config.frame_frequency

# Run the dynamics in blocks of the GCMC frequency.
dynamics.run(
self._config.gcmc_frequency,
Expand All @@ -963,24 +991,8 @@ def generate_lam_vals(lambda_base, increment=0.001):
save_crash_report=self._config.save_crash_report,
)

# Perform a GCMC move.
_logger.info(
f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}"
)
gcmc_sampler.push()
try:
gcmc_sampler.move(dynamics.context())
finally:
gcmc_sampler.pop()

# 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 save_frames and runtime >= next_frame:
gcmc_sampler.write_ghost_residues()
next_frame += self._config.frame_frequency
else:
dynamics.run(
time,
Expand Down
Loading