# -*- coding: utf-8 -*-
"""
HDF5 save/load helpers for scan outputs.
The scan output is a *nested* dict of numpy arrays mirroring Result (produced by
results_to_arrays), optionally with a scan_meta group carrying the swept
axes/labels/flow; it is written to nested HDF5 groups and read back into a
ResultArray. The on-disk schema is therefore generated from the dataclasses —
there is no hand-maintained key list. Optimisation scans now use this same
format (the optimised psi map lives in scan_meta/extra).
"""
from __future__ import annotations
import h5py as hp
import numpy as np
[docs]def _write_group(grp, data: dict) -> None:
"""Recursively write a (possibly nested) dict of arrays into an HDF5 group."""
for key, value in data.items():
if isinstance(value, dict):
_write_group(grp.create_group(key), value)
else:
grp[key] = value
[docs]def _read_group(grp) -> dict:
"""Recursively read an HDF5 group into a (possibly nested) dict of arrays."""
out = {}
for key, item in grp.items():
out[key] = _read_group(item) if isinstance(item, hp.Group) else np.array(item)
return out
[docs]def save_hdf5(filename: str, data_dict: dict) -> None:
"""Write a (possibly nested) dict of arrays to filename (mode "w")."""
with hp.File(filename, "w") as f:
_write_group(f, data_dict)
[docs]def load_hdf5(filename: str) -> dict:
"""Read filename into a (possibly nested) dict of arrays."""
with hp.File(filename, "r") as f:
return _read_group(f)
# %% Save/Load Functions
[docs]def save_out(name: str, out) -> None:
"""Save a scan result (a ResultArray or its nested dict) to <name>.hdf5.
Prefer result.save(name), which also writes the scan metadata; this
helper persists the raw data dict only.
"""
data = out.to_dict() if hasattr(out, "to_dict") else out
save_hdf5(f"{name}.hdf5", data)
[docs]def load_out(file: str):
"""Load a scan result from an HDF5 file as a ResultArray (with metadata)."""
from albums.result_io import ResultArray
return ResultArray.load(file)