From 2f2e50bf60af2e1f1d7ed811bab4abb5b502fe30 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 4 Mar 2026 13:53:47 +0000 Subject: [PATCH 1/9] Don't return entire energy trajectory after each cycle. --- src/somd2/runner/_repex.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index f6ba67b..85e67c0 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1283,11 +1283,8 @@ def _run_block( gcmc_sampler.write_ghost_residues() # Get the energy at each lambda value. - energies = ( - dynamics._d.energy_trajectory() - .to_pandas(to_alchemlyb=True, energy_unit="kcal/mol") - .iloc[-1, :] - .to_numpy() + energies = _np.array( + [e.value() for e in list(dynamics.current_energies().values())[2:]] ) except Exception as e: From 2230dacee1df4b0ec80632e93185e8c1ac4244dd Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 4 Mar 2026 19:05:10 +0000 Subject: [PATCH 2/9] Update OpenCL instructions. --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index cb55685..a7899d8 100644 --- a/README.md +++ b/README.md @@ -71,11 +71,12 @@ pip install -e . > [!Note] > Pixi does not run conda post-link scripts, so the `ocl-icd-system` > symlink needed for OpenCL won't be created automatically. After -> creating the environment, run the following once to fix this: +> creating the environment (or after a pixi update), run the following +> to fix this: > > ```bash > pixi shell -> ln -s /etc/OpenCL/vendors "${CONDA_PREFIX}/etc/OpenCL/vendors/ocl-icd-system" +> ln -sfn /etc/OpenCL/vendors "${CONDA_PREFIX}/etc/OpenCL/vendors/ocl-icd-system" > ``` ### Testing From 2918d7cc07ba718baa17c60c0210d9db2c34ff98 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 5 Mar 2026 09:25:54 +0000 Subject: [PATCH 3/9] Update pre-commit. --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0043ca4..61e8edf 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -14,7 +14,7 @@ repos: # Python formatting and linting - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.8.4 + rev: v0.15.4 hooks: # Run the formatter - id: ruff-format From b0165b6cf4c5f27082e08edc6125c61dbd03104b Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 5 Mar 2026 09:28:06 +0000 Subject: [PATCH 4/9] Autoformat. --- src/somd2/config/_config.py | 2 +- src/somd2/io/_io.py | 2 +- src/somd2/runner/_base.py | 4 ++-- src/somd2/runner/_repex.py | 4 ++-- src/somd2/runner/_runner.py | 6 +++--- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py index 0dbdb83..a799f63 100644 --- a/src/somd2/config/_config.py +++ b/src/somd2/config/_config.py @@ -1813,7 +1813,7 @@ def gcmc_radius(self, gcmc_radius): gcmc_r = _sr.u(gcmc_radius) except: raise ValueError( - "Unable to parse 'gcmc_radius' " f"as a Sire GeneralUnit: {gcmc_radius}" + f"Unable to parse 'gcmc_radius' as a Sire GeneralUnit: {gcmc_radius}" ) if not gcmc_r.has_same_units(angstrom): diff --git a/src/somd2/io/_io.py b/src/somd2/io/_io.py index 6b66497..e845e6a 100644 --- a/src/somd2/io/_io.py +++ b/src/somd2/io/_io.py @@ -74,7 +74,7 @@ def dataframe_to_parquet(df, metadata, filepath=None, filename=None): table = table.replace_schema_metadata(combined_meta) if filename is None: if "lambda" in metadata and "temperature" in metadata: - filename = f"Lam_{metadata['lambda'].replace('.','')[:5]}_T_{metadata['temperature']}.parquet" + filename = f"Lam_{metadata['lambda'].replace('.', '')[:5]}_T_{metadata['temperature']}.parquet" else: filename = "output.parquet" if not filename.endswith(".parquet"): diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index aabfaaf..ab23e6d 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -1985,8 +1985,8 @@ def _save_energy_components(self, index, context): state = new_context.getState(getEnergy=True, groups={i}) name = f.getName() name_len = len(name) - header += f"{f.getName():>{name_len+2}}" - record += f"{state.getPotentialEnergy().value_in_unit(openmm.unit.kilocalories_per_mole):>{name_len+2}.2f}" + header += f"{f.getName():>{name_len + 2}}" + record += f"{state.getPotentialEnergy().value_in_unit(openmm.unit.kilocalories_per_mole):>{name_len + 2}.2f}" # Write to file. if self._nrg_sample == 0: diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 85e67c0..cefce02 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -993,7 +993,7 @@ def run(self): # Perform the replica exchange simulation. for i in range(cycles): - _logger.info(f"Running dynamics for cycle {i+1} of {cycles}") + _logger.info(f"Running dynamics for cycle {i + 1} of {cycles}") # Log the states. This is the replica index for the state (positions # and velocities) used to seed each replica for the current cycle. @@ -1678,7 +1678,7 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): dynamics._d._sire_mols.delete_all_frames() _logger.info( - f"Finished block {block+1} of {self._start_block + num_blocks} " + f"Finished block {block + 1} of {self._start_block + num_blocks} " f"for {_lam_sym} = {lam:.5f}" ) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 931b534..f5fcf7f 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -748,7 +748,7 @@ def generate_lam_vals(lambda_base, increment=0.001): except: pass raise RuntimeError( - f"Dynamics block {block+1} for {_lam_sym} = {lambda_value:.5f} failed: {e}" + f"Dynamics block {block + 1} for {_lam_sym} = {lambda_value:.5f} failed: {e}" ) # Checkpoint. @@ -809,7 +809,7 @@ def generate_lam_vals(lambda_base, increment=0.001): dynamics._d._sire_mols.delete_all_frames() _logger.info( - f"Finished block {block+1} of {self._start_block + num_blocks} " + f"Finished block {block + 1} of {self._start_block + num_blocks} " f"for {_lam_sym} = {lambda_value:.5f}" ) @@ -884,7 +884,7 @@ def generate_lam_vals(lambda_base, increment=0.001): dynamics._d._sire_mols.delete_all_frames() _logger.info( - f"Finished block {block+1} of {self._start_block + num_blocks} " + f"Finished block {block + 1} of {self._start_block + num_blocks} " f"for {_lam_sym} = {lambda_value:.5f}" ) From c576a120891074f69d94fe685e096e13d50f6b11 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 5 Mar 2026 16:16:07 +0000 Subject: [PATCH 5/9] Update to use raw NumPy float array. --- src/somd2/runner/_repex.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index cefce02..99e86cb 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1283,9 +1283,7 @@ def _run_block( gcmc_sampler.write_ghost_residues() # Get the energy at each lambda value. - energies = _np.array( - [e.value() for e in list(dynamics.current_energies().values())[2:]] - ) + energies = dynamics._current_energy_array() except Exception as e: try: From 18e58670facbce71bc84ce092768ed10882f6518 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 6 Mar 2026 12:19:19 +0000 Subject: [PATCH 6/9] Fix off-by-one error. --- src/somd2/runner/_repex.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index f6ba67b..bf4d15a 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1008,13 +1008,13 @@ def run(self): results = [] # Whether to checkpoint. - is_checkpoint = i > 0 and i % cycles_per_checkpoint == 0 + is_checkpoint = (i + 1) % cycles_per_checkpoint == 0 # Whether to perform a GCMC move before the dynamics block. is_gcmc = i % cycles_per_gcmc == 0 # Whether a frame is saved at the end of the cycle. - write_gcmc_ghosts = i > 0 and i % cycles_per_frame == 0 + write_gcmc_ghosts = (i + 1) % cycles_per_frame == 0 # Run a dynamics block for each replica, making sure only each GPU is only # oversubscribed by a factor of self._config.oversubscription_factor. From 8134c434f62f17a6d4c658593b6b66f6f08e1682 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 6 Mar 2026 12:43:47 +0000 Subject: [PATCH 7/9] Make GCMC residue index write self-consistent. --- src/somd2/runner/_repex.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index bf4d15a..8e3c5a0 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1011,7 +1011,7 @@ def run(self): is_checkpoint = (i + 1) % cycles_per_checkpoint == 0 # Whether to perform a GCMC move before the dynamics block. - is_gcmc = i % cycles_per_gcmc == 0 + is_gcmc = (i + 1) % cycles_per_gcmc == 0 # Whether a frame is saved at the end of the cycle. write_gcmc_ghosts = (i + 1) % cycles_per_frame == 0 @@ -1248,6 +1248,14 @@ def _run_block( # 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, @@ -1277,10 +1285,6 @@ def _run_block( # Save the GCMC state. if gcmc_sampler is not None: self._dynamics_cache.save_gcmc_state(index) - # The frame frequency was hit, so write the indices of the - # current ghost water residues to file. - if write_gcmc_ghosts: - gcmc_sampler.write_ghost_residues() # Get the energy at each lambda value. energies = ( From a5dc8ec949c27798fb948c7387eb57cad67ab533 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 6 Mar 2026 13:16:43 +0000 Subject: [PATCH 8/9] Handle fractional frequencies correctly. --- src/somd2/runner/_repex.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 8e3c5a0..a07bb3c 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -904,12 +904,12 @@ def run(self): frac = 1.0 checkpoint_frequency = self._config.energy_frequency - # Store the number of repex cycles per block. - cycles_per_checkpoint = int(frac) + # Store the number of repex cycles per block (may be fractional). + cycles_per_checkpoint = frac # Otherwise, we don't checkpoint. else: - cycles_per_checkpoint = cycles + cycles_per_checkpoint = float(cycles) num_blocks = 1 rem = 0 @@ -991,6 +991,10 @@ def run(self): else: cycles_per_gcmc = cycles + 1 + # Initialise the threshold for the next checkpoint cycle. This is a float + # to handle non-integer ratios between the checkpoint and energy frequencies. + next_checkpoint = cycles_per_checkpoint + # Perform the replica exchange simulation. for i in range(cycles): _logger.info(f"Running dynamics for cycle {i+1} of {cycles}") @@ -1007,8 +1011,9 @@ def run(self): # Clear the results list. results = [] - # Whether to checkpoint. - is_checkpoint = (i + 1) % cycles_per_checkpoint == 0 + # Whether to checkpoint. Use a float threshold to correctly handle + # non-integer ratios between the checkpoint and energy frequencies. + is_checkpoint = (i + 1) >= next_checkpoint - 1e-10 # Whether to perform a GCMC move before the dynamics block. is_gcmc = (i + 1) % cycles_per_gcmc == 0 @@ -1119,6 +1124,9 @@ def run(self): # Update the block number. block += 1 + # Advance the checkpoint threshold. + next_checkpoint += cycles_per_checkpoint + # Guard the repex state and transition matrix saving with a file lock. lock = _FileLock(self._lock_file) with lock.acquire(timeout=self._config.timeout.to("seconds")): From 0e6a2d889162873ab15062c6ad6899f5b53c4c50 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 6 Mar 2026 17:22:53 +0000 Subject: [PATCH 9/9] Add rest2_selection to GCMC kwargs. --- src/somd2/runner/_base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index ab23e6d..eda7412 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -764,6 +764,7 @@ def __init__(self, system, config): "coulomb_power": self._config.coulomb_power, "shift_coulomb": str(self._config.shift_coulomb), "shift_delta": str(self._config.shift_delta), + "rest2_selection": self._config.rest2_selection, "swap_end_states": self._config.swap_end_states, "tolerance": self._config.gcmc_tolerance, "restart": self._is_restart,