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 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 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..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, @@ -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: diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index f6ba67b..84e01f9 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,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}" ) 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}" )