"""HDF5 data loader."""
from __future__ import annotations
from pathlib import Path
from typing import Any
import h5py
import numpy as np
from ..base import BaseDataLoader, register_loader
from ..bundle import DataBundle
[docs]
@register_loader("hdf5")
@register_loader("default")
class HDF5Loader(BaseDataLoader):
"""Load readout traces from HDF5 files."""
supported_formats = (".h5", ".hdf5")
def __init__(
self,
traces_key: str = "DD",
labels_key: str = "y",
time_key: str = "t_bin",
):
self.traces_key = traces_key
self.labels_key = labels_key
self.time_key = time_key
[docs]
def load(self, path: str | Path, **kwargs: Any) -> DataBundle:
traces_key = kwargs.get("traces_key", self.traces_key)
labels_key = kwargs.get("labels_key", self.labels_key)
time_key = kwargs.get("time_key", self.time_key)
path = Path(path)
meta: dict[str, Any] = {"source": str(path), "format": "hdf5"}
with h5py.File(path, "r") as f:
if traces_key not in f:
available = list(f.keys())
raise KeyError(
f"Traces key '{traces_key}' not found. Available: {available}"
)
traces = np.array(f[traces_key])
labels = np.array(f[labels_key]) if labels_key in f else None
time = np.array(f[time_key]) if time_key in f else None
return DataBundle(traces=traces, labels=labels, time=time, metadata=meta)
[docs]
@classmethod
def from_config(cls, config: dict[str, Any]) -> "HDF5Loader":
keys = config.get("hdf5_keys", {})
return cls(
traces_key=keys.get("traces", "DD"),
labels_key=keys.get("labels", "y"),
time_key=keys.get("time", "t_bin"),
)