From 40c0d20a50bab8b32f2c2d880ce7daf86c0f027a Mon Sep 17 00:00:00 2001 From: PythonFZ Date: Thu, 5 Mar 2026 15:15:42 +0100 Subject: [PATCH 1/3] use lazy `znh5md.IO` with `file_factory` for ASEMD frames property Replace eager frame loading with a lazy `znh5md.IO` handle using a context-manager-based file factory. Write all frames to a single `md.h5` file instead of per-index files. Co-Authored-By: Claude Opus 4.6 --- ipsuite/dynamics/md_nodes.py | 48 +++++++++++++++--------------------- 1 file changed, 20 insertions(+), 28 deletions(-) diff --git a/ipsuite/dynamics/md_nodes.py b/ipsuite/dynamics/md_nodes.py index b5daefda..5ed84678 100644 --- a/ipsuite/dynamics/md_nodes.py +++ b/ipsuite/dynamics/md_nodes.py @@ -1,3 +1,4 @@ +import contextlib import logging import sys import typing as t @@ -125,10 +126,8 @@ class ASEMD(zntrack.Node): Output path for model-specific output files. laufband_path : Path Path to the job queue database file. - frames : list[ase.Atoms] - Property that returns all trajectory frames from saved files. - structures : list[list[ase.Atoms]] - Property that returns structures organized by simulation run. + frames : znh5md.IO + Property that returns a lazy IO handle for trajectory frames. Examples -------- @@ -168,25 +167,20 @@ class ASEMD(zntrack.Node): laufband_path: Path = zntrack.outs_path(zntrack.nwd / "laufband.sqlite") @property - def frames(self) -> list[ase.Atoms]: - files = list(self.state.fs.glob((self.frames_path / "*.h5").as_posix())) - frames = [] - for file in files: - with self.state.fs.open(file, "rb") as f: - with h5py.File(f) as file: - frames.extend(znh5md.IO(file_handle=file)[:]) - return frames - - @property - def structures(self) -> list[list[ase.Atoms]]: - """Return the structures as a list of lists of Atoms.""" - files = list(self.state.fs.glob((self.frames_path / "*.h5").as_posix())) - structures = [] - for file in files: - with self.state.fs.open(file, "rb") as f: - with h5py.File(f) as file: - structures.append(znh5md.IO(file_handle=file)[:]) - return structures + def frames(self) -> znh5md.IO: + @contextlib.contextmanager + def _factory() -> t.Generator[h5py.File, None, None]: + if self.state.rev is None and self.state.remote is None: + with h5py.File(self.frames_path / "md.h5", "r") as file: + yield file + else: + with self.state.fs.open( + (self.frames_path / "md.h5").as_posix(), "rb" + ) as f: + with h5py.File(f) as file: + yield file + + return znh5md.IO(file_factory=_factory) def initialize_md(self): self.model_outs.mkdir(parents=True, exist_ok=True) @@ -263,7 +257,7 @@ def run_md(self, idx: int, atoms: ase.Atoms) -> int: # noqa: C901 ncols=120, ) io = znh5md.IO( - self.frames_path / f"{idx}.h5", + self.frames_path / "md.h5", ) # We do not save the starting configuration. E.g. step 0 is not saved! with Live(console=progress.console, refresh_per_second=10) as live: @@ -388,10 +382,8 @@ class ASEMDSafeSampling(ASEMD): Output path for model-specific output files. laufband_path : Path Path to the job queue database file. - frames : list[ase.Atoms] - Property that returns all trajectory frames from saved files. - structures : list[list[ase.Atoms]] - Property that returns structures organized by simulation run. + frames : znh5md.IO + Property that returns a lazy IO handle for trajectory frames. Examples -------- From e355eeb18ea8c3de549a9e3f08df780ab9dad195 Mon Sep 17 00:00:00 2001 From: PythonFZ Date: Thu, 5 Mar 2026 15:25:32 +0100 Subject: [PATCH 2/3] factory / support multiple files via materialization --- ipsuite/dynamics/md_nodes.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/ipsuite/dynamics/md_nodes.py b/ipsuite/dynamics/md_nodes.py index 5ed84678..51100257 100644 --- a/ipsuite/dynamics/md_nodes.py +++ b/ipsuite/dynamics/md_nodes.py @@ -166,21 +166,34 @@ class ASEMD(zntrack.Node): model_outs: Path = zntrack.outs_path(zntrack.nwd / "model") laufband_path: Path = zntrack.outs_path(zntrack.nwd / "laufband.sqlite") - @property - def frames(self) -> znh5md.IO: + def _make_file_factory(self, file_path: str) -> t.Callable: @contextlib.contextmanager def _factory() -> t.Generator[h5py.File, None, None]: if self.state.rev is None and self.state.remote is None: - with h5py.File(self.frames_path / "md.h5", "r") as file: + with h5py.File(file_path, "r") as file: yield file else: - with self.state.fs.open( - (self.frames_path / "md.h5").as_posix(), "rb" - ) as f: + with self.state.fs.open(file_path, "rb") as f: with h5py.File(f) as file: yield file - return znh5md.IO(file_factory=_factory) + return _factory + + @property + def frames(self) -> znh5md.IO | list[ase.Atoms]: + files = sorted( + self.state.fs.glob((self.frames_path / "*.h5").as_posix()) + ) + if not files: + raise FileNotFoundError( + f"No HDF5 files found in {self.frames_path}" + ) + if len(files) == 1: + return znh5md.IO(file_factory=self._make_file_factory(files[0])) + return sum( + [znh5md.IO(file_factory=self._make_file_factory(f))[:] for f in files], + [], + ) def initialize_md(self): self.model_outs.mkdir(parents=True, exist_ok=True) From 1f5f64b0c0be661ba9a8c399cd12b46dd955cf75 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 5 Mar 2026 14:25:45 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- ipsuite/dynamics/md_nodes.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/ipsuite/dynamics/md_nodes.py b/ipsuite/dynamics/md_nodes.py index 51100257..2882b695 100644 --- a/ipsuite/dynamics/md_nodes.py +++ b/ipsuite/dynamics/md_nodes.py @@ -181,13 +181,9 @@ def _factory() -> t.Generator[h5py.File, None, None]: @property def frames(self) -> znh5md.IO | list[ase.Atoms]: - files = sorted( - self.state.fs.glob((self.frames_path / "*.h5").as_posix()) - ) + files = sorted(self.state.fs.glob((self.frames_path / "*.h5").as_posix())) if not files: - raise FileNotFoundError( - f"No HDF5 files found in {self.frames_path}" - ) + raise FileNotFoundError(f"No HDF5 files found in {self.frames_path}") if len(files) == 1: return znh5md.IO(file_factory=self._make_file_factory(files[0])) return sum(