"""Contain utility functions for HDF5 file handling.
.. autoclass:: H5File
.. autofunction:: print_hdf5_structure
.. autofunction:: assert_datasets_almost_equal
"""
import logging
from pathlib import Path
from typing import Literal, Self, Sequence, TypeAlias, cast
import numpy as np
from h5py import Dataset as H5Dataset
from h5py import File, Group
from numpy.typing import DTypeLike
from mojito.dataset import Dataset
logger = logging.getLogger(__name__)
FileMode: TypeAlias = Literal["r", "r+", "a", "w", "w-"]
"""Type for file open modes, compatible with :class:`h5py.File` modes."""
def is_write_mode(mode: FileMode) -> bool:
"""Check if the given file mode allows writing to the HDF5 file.
Parameters
----------
mode
File open mode to check.
Returns
-------
True if the mode allows writing, False otherwise.
"""
return mode in ("r+", "a", "w", "w-")
[docs]
class H5File:
"""Mixin for HDF5 file handling in context-manager scenarios.
This mixin provides the context-manager methods :meth:`__enter__` and
:meth:`__exit__` for classes that handle HDF5 files, ensuring proper resource
management.
Parameters
----------
path
Path to the HDF5 file.
mode
File open mode (default "r" for read-only).
"""
def __init__(self, path: str | Path, *, mode: FileMode = "r") -> None:
self.path = str(path)
self.mode: FileMode = mode
self._file: File | None = None
def __enter__(self) -> Self:
"""Open the HDF5 file and return the context manager."""
self._file = File(self.path, self.mode)
return self
def __exit__(self, exc_type, exc_value, traceback) -> None:
"""Close the HDF5 file."""
if self._file is not None:
self._file.close()
self._file = None
@property
def _opened_file(self) -> File:
"""Ensure that the HDF5 file is open and return the file object.
Raises
------
RuntimeError
If the HDF5 file is not open.
Returns
-------
The open HDF5 file object.
"""
if self._file is None:
raise RuntimeError("HDF5 file is not open")
return self._file
def _get_dataset(self, name: str) -> Dataset:
"""Retrieve a dataset from the HDF5 file.
Parameters
----------
name
Name of the dataset to retrieve.
Returns
-------
The requested dataset.
Raises
------
KeyError
If the dataset does not exist in the HDF5 file, or if the object with the
specified name is not a dataset.
"""
dataset = self._opened_file.get(name, None)
if dataset is None or not isinstance(dataset, H5Dataset):
raise KeyError(f"Dataset '{name}' not found in the HDF5 file")
# h5py.Dataset implements __array__ at runtime, but this is not reflected in
# its type stubs, so it is not statically recognized as a Dataset protocol
# member. Cast explicitly to silence the resulting false-positive type error.
return cast(Dataset, dataset)
def _require_dataset(
self,
name: str,
*,
shape: tuple[int, ...],
dtype: DTypeLike = np.float64,
) -> Dataset:
"""Retrieve or create a dataset from the HDF5 file.
This function attempts to retrieve a dataset with the given name from the HDF5
file. If the dataset does not exist, it will be created.
Parameters
----------
name
Name of the dataset to retrieve or create.
shape
Shape of the dataset to create if it does not exist.
dtype
Data type of the dataset to create if it does not exist.
Returns
-------
The requested or newly created dataset.
"""
# Only create the dataset if it does not exist and the file is writable
if name not in self._opened_file and is_write_mode(self.mode):
self._opened_file.create_dataset(name, shape=shape, dtype=dtype)
try:
dataset = self._get_dataset(name)
except KeyError as error:
raise KeyError(
f"Dataset '{name}' not found in the HDF5 file; "
"open in write mode to create it when accessed"
) from error
return dataset
[docs]
def print_hdf5_structure(path: str | Path, *, print_attrs: bool = False) -> None:
"""Print the structure of an HDF5 file with visual hierarchy.
Recursively prints the hierarchy of groups and datasets in an HDF5 file,
displaying their names, types, and shapes. Optionally prints attributes for
groups and datasets.
Parameters
----------
path
Path to the HDF5 file.
print_attrs
Whether to print attributes for groups and datasets.
Example
-------
>>> print_hdf5_structure("data.h5", print_attrs=True)
HDF5 Structure of data.h5
└── [GROUP] /
├── @version: 1.0
├── [DATASET] /dataset1
│ dtype: float64, shape: (100,)
└── [GROUP] /group1
└── [DATASET] /group1/dataset2
dtype: int32, shape: (50, 50)
@units: meters
"""
def print_structure(
obj: Group | H5Dataset,
prefix: str = "",
is_last: bool = True,
) -> None:
"""Recursively print HDF5 structure with tree-like visualization."""
current_prefix = "└── " if is_last else "├── "
next_prefix = " " if is_last else "│ "
if isinstance(obj, Group):
print(f"{prefix}{current_prefix}[GROUP] {obj.name}")
if print_attrs and obj.attrs:
attr_keys = list(obj.attrs.keys())
for i, key in enumerate(attr_keys):
is_last_attr = i == len(attr_keys) - 1
attr_prefix = (
prefix + next_prefix + ("└── " if is_last_attr else "├── ")
)
print(f"{attr_prefix}@{key}: {obj.attrs[key]}")
keys = list(obj.keys())
for i, key in enumerate(keys):
next_obj = obj[key]
if isinstance(next_obj, (Group, H5Dataset)):
is_last_child = i == len(keys) - 1
print_structure(next_obj, prefix + next_prefix, is_last_child)
elif isinstance(obj, H5Dataset):
name = obj.name
print(f"{prefix}{current_prefix}[DATASET] {name}")
print(f"{prefix}{next_prefix} dtype: {obj.dtype}, shape: {obj.shape}")
if print_attrs and obj.attrs:
attr_keys = list(obj.attrs.keys())
for i, key in enumerate(attr_keys):
is_last_attr = i == len(attr_keys) - 1
attr_prefix = (
prefix + next_prefix + ("└── " if is_last_attr else "├── ")
)
print(f"{attr_prefix}@{key}: {obj.attrs[key]}")
with File(str(path), "r") as f:
print(f"HDF5 Structure of {path}")
print_structure(f)
[docs]
def assert_datasets_almost_equal(
datasets: Sequence[Dataset],
*,
chunk: int = 100_000,
) -> None:
"""Assert that the given datasets are almost equal, reading them in chunks.
This is a helper function to check that datasets are identical across groups
without loading them entirely into memory.
Note that chunking is only done along the first dimension of the datasets,
so this function is most effective for datasets where the first dimension is
the largest.
Parameters
----------
datasets
Sequence of datasets to compare. All datasets must have the same shape.
chunk
Chunk size to use on the first axis when comparing the datasets. This
limits memory usage when comparing large datasets.
Raises
------
ValueError
If the chunk size is not positive.
AssertionError
If the datasets have different shapes.
AssertionError
If the datasets are not almost equal within the given chunk size.
"""
# Check chunk size is positive
if chunk <= 0:
raise ValueError("Chunk size must be positive")
# Return if there are no datasets to compare
if len(datasets) < 2:
return
# Take first dataset as reference
reference = datasets[0]
# Check that all datasets have the same shape
reference_shape = reference.shape
for dataset in datasets[1:]:
if dataset.shape != reference_shape:
raise AssertionError("Datasets have different shapes")
# For 0-D (scalar) datasets, compare directly without chunking
if len(reference_shape) == 0:
for dataset in datasets[1:]:
np.testing.assert_array_almost_equal(
dataset[()],
reference[()],
err_msg="Inconsistent datasets across files",
)
return
# For higher-dimensional datasets, compare in chunks along the first axis
# and raise a warning if the number of dimensions is greater than 1
if len(reference_shape) > 1:
logger.warning(
"Comparing datasets with more than 1 dimension using chunking "
"along the first axis. This may be memory-intensive if the "
"datasets dare large along other axes."
)
# Check that all datasets are almost equal in chunks
for i in range(0, reference_shape[0], chunk):
reference_chunk = reference[i : i + chunk]
for dataset in datasets[1:]:
np.testing.assert_array_almost_equal(
dataset[i : i + chunk],
reference_chunk,
err_msg="Inconsistent datasets across files",
)