From 1d8c6018c3a47dad8b59a9c04adabf4df1c8efc2 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 21 Sep 2026 12:51:33 +0100 Subject: [PATCH] Buffer energy components and write them to parquet at checkpoint time. --- CHANGELOG.md | 1 + src/somd2/runner/_base.py | 53 ++++++++-- tests/runner/test_energy_components.py | 131 +++++++++++++++++++++++++ 3 files changed, 176 insertions(+), 9 deletions(-) create mode 100644 tests/runner/test_energy_components.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9352933..6f419f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ Changelog * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR. * Increase the default `cutoff` from 7.5 Å to 9 Å, matching common practice and giving faster PME on current GPUs, since the shorter cutoff shifts too much work onto the reciprocal space grid [#209](https://github.com/OpenBioSim/somd2/pull/209). +* Buffer energy components and write them at checkpoint time, rather than rewriting the parquet file on every energy save, which cost a few milliseconds per replica per cycle and grew with the length of the run [#212](https://github.com/OpenBioSim/somd2/pull/212). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 8e8890b..023e67d 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -636,6 +636,11 @@ def __init__(self, system, config): # used to skip duplicate rows on restart. self._last_ec_time = {} + # Per-window energy-components rows collected since the last checkpoint, + # and the number of rows at which they are written out regardless. + self._ec_rows = {} + self._max_ec_rows = 10000 + # Per-window cache of the integrator's integration force groups bitmask. self._integration_groups = {} @@ -2500,6 +2505,9 @@ def _checkpoint( df.iloc[-self._energy_per_block :], ) + if not is_post_equilibration: + self._flush_energy_components(index) + except Exception as e: return index, e @@ -2644,8 +2652,11 @@ def _backup_checkpoint(self, index): def _save_energy_components(self, index, context, time_ns): """ - Internal function to save the energy components for each force group to a - Parquet file. + Internal function to record the energy components for each force group. + Rows are buffered and written to the Parquet file by + _flush_energy_components() at checkpoint time, so that the file is only + rewritten once per checkpoint and stays consistent with the other + checkpoint files. Parameters ---------- @@ -2660,10 +2671,7 @@ def _save_energy_components(self, index, context, time_ns): The current simulation time in nanoseconds. """ - import json as _json import openmm - import pandas as _pd - import pyarrow as _pa import pyarrow.parquet as _pq_local filepath = self._filenames[index]["energy_components"] @@ -2701,8 +2709,37 @@ def _save_energy_components(self, index, context, time_ns): openmm.unit.kilocalories_per_mole ) - row = {"time": round(time_ns, 6)} | energies - df = _pd.DataFrame([row]) + rows = self._ec_rows.setdefault(index, []) + rows.append({"time": round(time_ns, 6)} | energies) + self._last_ec_time[index] = time_ns + + # Bound the memory used when checkpoints are rare or disabled. + if len(rows) >= self._max_ec_rows: + self._flush_energy_components(index) + + def _flush_energy_components(self, index): + """ + Write the energy components buffered by _save_energy_components() to + the Parquet file for a window. + + Parameters + ---------- + + index : int + The index of the window or replica. + """ + + import json as _json + import pandas as _pd + import pyarrow as _pa + import pyarrow.parquet as _pq_local + + rows = self._ec_rows.pop(index, []) + if not rows: + return + + filepath = self._filenames[index]["energy_components"] + df = _pd.DataFrame(rows) path = _Path(filepath) if path.exists() and path.stat().st_size > 0: @@ -2719,8 +2756,6 @@ def _save_energy_components(self, index, context, time_ns): ) _pq_local.write_table(table, filepath) - self._last_ec_time[index] = time_ns - def _restore_backup_files(self): """ Restore backup files in the working directory. diff --git a/tests/runner/test_energy_components.py b/tests/runner/test_energy_components.py new file mode 100644 index 0000000..ef9826a --- /dev/null +++ b/tests/runner/test_energy_components.py @@ -0,0 +1,131 @@ +from pathlib import Path +import tempfile + +import pyarrow.parquet as pq +import pytest + +from somd2.runner import Runner, RepexRunner +from somd2.config import Config + +from tests.conftest import has_cuda + + +def _config(tmpdir, runtime, platform, repex, restart=False): + return Config( + runtime=runtime, + restart=restart, + output_directory=tmpdir, + energy_frequency="4fs", + checkpoint_frequency="12fs", + frame_frequency="12fs", + platform=platform, + max_threads=1, + num_lambda=2, + replica_exchange=repex, + save_energy_components=True, + ) + + +def _times(tmpdir, lam): + table = pq.read_table(Path(tmpdir) / f"energy_components_{lam}.parquet") + return table.column("time").to_pylist() + + +@pytest.mark.parametrize( + "runner_class, platform", + [ + (Runner, "CPU"), + pytest.param( + RepexRunner, + "cuda", + marks=pytest.mark.skipif(not has_cuda, reason="CUDA not available."), + ), + ], +) +def test_energy_components_rows(ethane_methanol, runner_class, platform): + """ + Validate that the energy components file holds one row per energy save, + including across a restart. + """ + repex = runner_class is RepexRunner + with tempfile.TemporaryDirectory() as tmpdir: + runner = runner_class(ethane_methanol, _config(tmpdir, "12fs", platform, repex)) + runner.run() + + for lam in ("0.00000", "1.00000"): + times = _times(tmpdir, lam) + assert len(times) == 3 + assert times == sorted(set(times)) + + runner = runner_class( + ethane_methanol, _config(tmpdir, "24fs", platform, repex, restart=True) + ) + runner.run() + + for lam in ("0.00000", "1.00000"): + times = _times(tmpdir, lam) + assert len(times) == 6 + assert times == sorted(set(times)) + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_energy_components_written_at_checkpoint(ethane_methanol): + """ + Validate that energy components are buffered between checkpoints rather + than written on every energy save. + """ + with tempfile.TemporaryDirectory() as tmpdir: + runner = RepexRunner(ethane_methanol, _config(tmpdir, "24fs", "cuda", True)) + + saves = [] + flushes = [] + save = runner._save_energy_components + flush = runner._flush_energy_components + + def counting_save(index, context, time_ns): + saves.append(index) + return save(index, context, time_ns) + + def counting_flush(index): + flushes.append(index) + return flush(index) + + runner._save_energy_components = counting_save + runner._flush_energy_components = counting_flush + + runner.run() + + # Six energy saves and two checkpoints per replica. + assert saves.count(0) == 6 + assert flushes.count(0) == 2 + assert len(_times(tmpdir, "0.00000")) == 6 + + +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_energy_components_buffer_limit(ethane_methanol): + """ + Validate that the buffer is written out when it reaches its size limit, + so that memory is bounded when checkpoints are rare. + """ + with tempfile.TemporaryDirectory() as tmpdir: + runner = RepexRunner(ethane_methanol, _config(tmpdir, "24fs", "cuda", True)) + runner._max_ec_rows = 2 + + flushes = [] + flush = runner._flush_energy_components + + def counting_flush(index): + flushes.append(index) + return flush(index) + + runner._flush_energy_components = counting_flush + + runner.run() + + # The limit is reached after the second and fifth saves, since the + # checkpoint after the third empties the buffer, plus one flush at each + # of the two checkpoints. + assert flushes.count(0) == 4 + times = _times(tmpdir, "0.00000") + assert len(times) == 6 + assert times == sorted(set(times))