diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 1120d53..99a1286 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -2373,6 +2373,7 @@ def _checkpoint( is_final_block=False, context=None, gcmc_sampler=None, + num_energy_rows=None, ): """ Save a checkpoint file. @@ -2401,6 +2402,10 @@ def _checkpoint( is_final_block: bool Whether this is the final block of the simulation. + num_energy_rows: int + The number of energy samples collected since the last checkpoint. + If None, then a full block is assumed. + Returns ------- @@ -2427,6 +2432,13 @@ def _checkpoint( # parquet-related work in this case. is_post_equilibration = block == -1 + def new_rows(df): + # The rows sampled since the last checkpoint. Sliced from the + # start, since a count of zero must select no rows. + if num_energy_rows is None: + return df.iloc[-self._energy_per_block :] + return df.iloc[len(df) - num_energy_rows :] + # Get the energy trajectory. if not is_post_equilibration: df = system.energy_trajectory(to_alchemlyb=True, energy_unit="kT") @@ -2518,7 +2530,7 @@ def _checkpoint( # yet (e.g. checkpoint_frequency=0), create one from scratch. _energy_traj = self._filenames[index]["energy_traj"] if _Path(_energy_traj).exists(): - _parquet_append(_energy_traj, df.iloc[-self._energy_per_block :]) + _parquet_append(_energy_traj, new_rows(df)) else: _dataframe_to_parquet(df, metadata=metadata, filename=_energy_traj) @@ -2558,10 +2570,7 @@ def _checkpoint( ): _dataframe_to_parquet(df, metadata=metadata, filename=filename) else: - _parquet_append( - filename, - df.iloc[-self._energy_per_block :], - ) + _parquet_append(filename, new_rows(df)) if not is_post_equilibration: self._flush_energy_components(index) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 98ea1cf..d094086 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -79,6 +79,61 @@ def __init__(self, system, config): _logger.error(msg) raise ValueError(msg) + from math import isclose + + def is_multiple(value, interval): + ratio = (value / interval).value() + return isclose(ratio, round(ratio), abs_tol=1e-4) + + # Dynamics are run in whole integration steps. + for option in ( + "runtime", + "checkpoint_frequency", + "energy_frequency", + "frame_frequency", + "gcmc_frequency", + "terminal_flip_frequency", + ): + value = getattr(config, option) + if value is not None and not is_multiple(value, config.timestep): + msg = f"'{option}' must be a multiple of 'timestep'." + _logger.error(msg) + raise ValueError(msg) + + # Blocks are split into sub-blocks between Monte Carlo moves, or energy + # component samples, so they must hold a whole number of them. + is_checkpointing = config.checkpoint_frequency.value() > 0.0 + intervals = {} + if config.gcmc: + intervals["gcmc_frequency"] = ( + config.gcmc_frequency or config.energy_frequency + ) + if config.terminal_flip_frequency is not None: + intervals["terminal_flip_frequency"] = config.terminal_flip_frequency + if config.save_energy_components: + intervals["energy_frequency"] = config.energy_frequency + options = ( + ("runtime", "checkpoint_frequency") if is_checkpointing else ("runtime",) + ) + for name, interval in intervals.items(): + for option in options: + if not is_multiple(getattr(config, option), interval): + msg = f"'{option}' must be a multiple of '{name}'." + _logger.error(msg) + raise ValueError(msg) + + if is_checkpointing: + # A final partial block doesn't perform GCMC moves. + if config.gcmc and not is_multiple( + config.runtime, config.checkpoint_frequency + ): + msg = ( + "'runtime' must be a multiple of 'checkpoint_frequency' " + "for GCMC simulations." + ) + _logger.error(msg) + raise ValueError(msg) + # Call the base class constructor. super().__init__(system, config) @@ -906,6 +961,10 @@ def generate_lam_vals(lambda_base, increment=0.001): # Store the start time. start = _timer() + # The number of energy samples recorded when the last checkpoint was + # written, so that each checkpoint appends only the new ones. + num_samples = len(dynamics._d._energy_trajectory) + # Run the simulation, checkpointing in blocks. if checkpoint_frequency.value() > 0.0: # Calculate the number of blocks and the remainder time. @@ -989,7 +1048,12 @@ def generate_lam_vals(lambda_base, increment=0.001): else None ) - while runtime < checkpoint_frequency: + # Count sub-blocks rather than compare times, since a + # rounding error can otherwise add one. + num_subblocks = round( + (checkpoint_frequency / block_size).value() + ) + while flip_counter < num_subblocks: needs_pre_run_snapshot = False # GCMC move. @@ -1071,9 +1135,7 @@ def generate_lam_vals(lambda_base, increment=0.001): self._save_energy_components( index, dynamics.context(), - (block * checkpoint_frequency + ec_elapsed).to( - "ns" - ), + dynamics.current_time().to("ns"), ) else: @@ -1136,10 +1198,13 @@ def generate_lam_vals(lambda_base, increment=0.001): is_final_block=is_final_block, context=dynamics.context(), gcmc_sampler=gcmc_sampler, + num_energy_rows=len(dynamics._d._energy_trajectory) + - num_samples, ) if error is not None: raise error + num_samples = len(dynamics._d._energy_trajectory) # Save sampler statistics alongside the checkpoint. self._save_sampler_stats( @@ -1214,8 +1279,12 @@ def generate_lam_vals(lambda_base, increment=0.001): ): dynamics.randomise_velocities() + # Run the remaining time as a whole number of steps, since + # Sire truncates a time a rounding error short of a step. dynamics.run( - rem, + round( + (rem * checkpoint_frequency / self._config.timestep).value() + ), energy_frequency=self._config.energy_frequency, frame_frequency=self._config.frame_frequency, lambda_windows=lambda_array, @@ -1268,6 +1337,8 @@ def generate_lam_vals(lambda_base, increment=0.001): is_final_block=True, context=dynamics.context(), gcmc_sampler=gcmc_sampler, + num_energy_rows=len(dynamics._d._energy_trajectory) + - num_samples, ) if error is not None: @@ -1346,7 +1417,10 @@ def generate_lam_vals(lambda_base, increment=0.001): ) time_base = self._config.runtime - time - while runtime < time: + # Count sub-blocks rather than compare times, since a + # rounding error can otherwise add one. + num_subblocks = round((time / block_size).value()) + while flip_counter < num_subblocks: needs_pre_run_snapshot = False # GCMC move. @@ -1481,6 +1555,7 @@ def generate_lam_vals(lambda_base, increment=0.001): is_final_block=True, context=dynamics.context(), gcmc_sampler=gcmc_sampler, + num_energy_rows=len(dynamics._d._energy_trajectory) - num_samples, ) if error is not None: diff --git a/tests/runner/test_terminal_flip.py b/tests/runner/test_terminal_flip.py index 7901121..c04f8a3 100644 --- a/tests/runner/test_terminal_flip.py +++ b/tests/runner/test_terminal_flip.py @@ -397,8 +397,8 @@ def test_runner_validation_frequency_multiple(ethane_methanol): output_directory=tmpdir, platform="cpu", num_lambda=2, - energy_frequency="4fs", - terminal_flip_frequency="3fs", # not a multiple of 4fs + energy_frequency="8fs", + terminal_flip_frequency="20fs", # not a multiple of 8fs ) with pytest.raises( ValueError, match="must be a multiple of 'energy_frequency'"