Merge pull request #126 from OpenBioSim/fix_124

Fix issue #124

Author: lohedges

.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

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

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

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"):

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,
@@ -1985,8 +1986,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:

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,9 +991,13 @@ 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}")
+            _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.
@@ -1007,14 +1011,15 @@ def run(self):
             # Clear the results list.
             results = []
 
-            # Whether to checkpoint.
-            is_checkpoint = i > 0 and i % 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 % 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 > 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.
@@ -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")):
@@ -1248,6 +1256,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,18 +1293,9 @@ 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 = (
-                dynamics._d.energy_trajectory()
-                .to_pandas(to_alchemlyb=True, energy_unit="kcal/mol")
-                .iloc[-1, :]
-                .to_numpy()
-            )
+            energies = dynamics._current_energy_array()
 
         except Exception as e:
             try:
@@ -1681,7 +1688,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}"
             )
 

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}"
                     )