ewoksdata 2.1.0__py3-none-any.whl

ewoksdata/data/hdf5/config.py

@@ -0,0 +1,186 @@
+"""HDF5 configuration of optimal data storage (IO speed, compression, ...)"""
+
+from collections.abc import Mapping
+from numbers import Integral
+from typing import Optional
+
+import numpy
+from numpy.typing import DTypeLike
+
+try:
+    import hdf5plugin
+except ImportError:
+    hdf5plugin = None
+
+from .types import ShapeType
+from .types import VarH5pyShapeType
+from .types import VarShapeType
+
+DEFAULT_CHUNK_NBYTES = 1 << 20
+DEFAULT_COMPRESSION_LIMIT_NBYTES = 1 << 20
+DEFAULT_CHUNK_SPLIT = 4
+DEFAULT_COMPRESSION_SCHEME = "gzip-byteshuffle"
+
+# Default data size
+#  0D detector: 2 KB
+#  1D detector: 2 MB
+#  2D detector: 2 GB
+DEFAULT_SCAN_DIM_SIZE = 512
+DEFAULT_DETECTOR_DIM_SIZE = 1024
+DEFAULT_DTYPE = numpy.int32
+
+
+def dtype_nbytes(dtype: DTypeLike) -> int:
+    return numpy.dtype(dtype).itemsize
+
+
+def shape_to_size(shape: ShapeType) -> int:
+    # numpy.prod gives problems on windows
+    n = 1
+    for x in shape:
+        n *= x
+    return n
+
+
+def shape_to_nbytes(shape: ShapeType, dtype: DTypeLike) -> int:
+    return shape_to_size(shape) * dtype_nbytes(dtype)
+
+
+def guess_data_shape(
+    scan_shape: VarShapeType,
+    detector_shape: VarShapeType,
+    max_shape: Optional[VarH5pyShapeType],
+) -> ShapeType:
+    scan_shape = tuple(n if n else DEFAULT_SCAN_DIM_SIZE for n in scan_shape)
+    detector_shape = tuple(n if n else DEFAULT_SCAN_DIM_SIZE for n in detector_shape)
+    data_shape = scan_shape + detector_shape
+    if max_shape:
+        assert len(max_shape) == len(
+            data_shape
+        ), "HDF5 dataset shape must have the same dimensions as maxshape"
+        data_shape = tuple(
+            n1 if not n2 else max(n1, n2) for n1, n2 in zip(data_shape, max_shape)
+        )
+    return data_shape
+
+
+def guess_chunk_shape(
+    data_shape: ShapeType,
+    dtype: DTypeLike,
+    chunk_split: Optional[Integral] = None,
+    chunk_nbytes: Optional[Integral] = None,
+) -> Optional[ShapeType]:
+    """Try to guess the optimal chunk shape with these constraints:
+    * Split any dimension for partial access
+    * Below the maximal chunk size (1 MB by default, uncompressed)
+
+    The inner-most dimensions are split in `chunk_split` parts
+    until chunk_nbytes is reached. The chunk size in the outer
+    dimensions will be 1, unless the data size is too small.
+    """
+    if chunk_nbytes is None:
+        chunk_nbytes = DEFAULT_CHUNK_NBYTES
+    if chunk_split is None:
+        chunk_split = DEFAULT_CHUNK_SPLIT
+    itemsize = dtype_nbytes(dtype)
+    size = shape_to_size(data_shape)
+    nbytes = size * itemsize
+    if nbytes <= chunk_nbytes:
+        return None
+
+    max_size = chunk_nbytes // itemsize
+    current_size = 1
+    chunk_shape = []
+    for n_i in data_shape[-1::-1]:
+        if current_size >= max_size:
+            c_i = 1
+        else:
+            a = int(numpy.ceil(n_i / chunk_split))
+            b = int(numpy.ceil(max_size / current_size))
+            c_i = min(a, b)
+        chunk_shape.append(c_i)
+        current_size *= c_i
+    chunk_shape = tuple(chunk_shape[::-1])
+    if chunk_shape == data_shape:
+        return None
+    return chunk_shape
+
+
+def guess_compression(
+    data_shape: ShapeType,
+    dtype: DTypeLike,
+    compression_limit_nbytes: Optional[Integral] = None,
+) -> bool:
+    """Compression is needed when the total data size exceeds the limits (1 MB by default)."""
+    if compression_limit_nbytes is None:
+        compression_limit_nbytes = DEFAULT_COMPRESSION_LIMIT_NBYTES
+    nbytes = shape_to_nbytes(data_shape, dtype)
+    return nbytes > compression_limit_nbytes
+
+
+def get_compression_arguments(compression_scheme: Optional[str] = None) -> Mapping:
+    if compression_scheme:
+        compression_scheme = compression_scheme.lower()
+    if compression_scheme is None:
+        compression_scheme = DEFAULT_COMPRESSION_SCHEME
+    if compression_scheme == "none":
+        return dict()
+    elif compression_scheme == "gzip":
+        return {"compression": "gzip"}
+    elif compression_scheme == "byteshuffle":
+        return {"shuffle": True}
+    elif compression_scheme == "gzip-byteshuffle":
+        return {"compression": "gzip", "shuffle": True}
+    elif compression_scheme == "bitshuffle":
+        if hdf5plugin is None:
+            raise RuntimeError(
+                "Writer does not support HDF5 'bitshuffle' compression. Install the hdf5plugin library"
+            )
+        return hdf5plugin.Bitshuffle(nelems=0, lz4=False)
+    elif compression_scheme == "lz4-bitshuffle":
+        if hdf5plugin is None:
+            raise RuntimeError(
+                "Writer does not support HDF5 'bitshuffle' compression. Install the hdf5plugin library"
+            )
+        return hdf5plugin.Bitshuffle(nelems=0, lz4=True)
+    else:
+        raise ValueError(f"Unknown HDF5 compression '{compression_scheme}'")
+
+
+def guess_dataset_config(
+    scan_shape: VarShapeType,
+    detector_shape: VarShapeType,
+    dtype: Optional[DTypeLike] = None,
+    chunk_split: Optional[Integral] = None,
+    chunk_nbytes: Optional[Integral] = None,
+    compression_limit_nbytes: Optional[Integral] = None,
+    compression_scheme: Optional[str] = None,
+    max_shape: Optional[VarH5pyShapeType] = None,
+) -> dict:
+    """Dataset configuration passed the `h5py.Group.create_dataset` for optimal
+    storage (IO speed, compression, ...)
+    """
+    data_shape = guess_data_shape(
+        scan_shape=scan_shape, detector_shape=detector_shape, max_shape=max_shape
+    )
+    if dtype is None:
+        dtype = DEFAULT_DTYPE
+    chunk_shape = guess_chunk_shape(
+        data_shape=data_shape,
+        dtype=dtype,
+        chunk_split=chunk_split,
+        chunk_nbytes=chunk_nbytes,
+    )
+    config = {"chunks": chunk_shape}
+    compression = guess_compression(
+        data_shape=data_shape,
+        dtype=dtype,
+        compression_limit_nbytes=compression_limit_nbytes,
+    )
+    if compression:
+        config.update(get_compression_arguments(compression_scheme=compression_scheme))
+    chunking_required = compression or max_shape is not None
+    if chunking_required and chunk_shape is None:
+        # Do not let h5py guess the chunk size
+        config["chunks"] = data_shape
+    return config

ewoksdata/data/hdf5/dataset_writer/__init__.py

@@ -0,0 +1,4 @@
+from ._non_sequential import IndexedDatasetWriter  # noqa F401
+from ._non_sequential_stack import IndexedStackDatasetWriter  # noqa F401
+from ._sequential import DatasetWriter  # noqa F401
+from ._sequential_stack import StackDatasetWriter  # noqa F401

ewoksdata/data/hdf5/dataset_writer/_base.py

@@ -0,0 +1,55 @@
+import time
+from typing import Optional
+
+import h5py
+import numpy
+
+
+class DatasetWriterBase:
+    def __init__(
+        self,
+        parent: h5py.Group,
+        name: str,
+        attrs: Optional[dict] = None,
+        flush_period: Optional[float] = None,
+    ) -> None:
+        self._file = parent.file
+        self._parent = parent
+        self._name = name
+        self._attrs = attrs
+        self._dataset_name = f"{parent.name}/{name}"
+        self._dataset = None
+        self._flush_period = flush_period
+        self._last_flush = None
+
+    @property
+    def dataset_name(self) -> str:
+        return self._dataset_name
+
+    def __enter__(self) -> "DatasetWriterBase":
+        return self
+
+    def __exit__(self, *args) -> None:
+        self.flush_buffer()
+
+    @property
+    def dataset(self) -> Optional[h5py.Dataset]:
+        return self._dataset
+
+    def _create_dataset(self, first_data_point: numpy.ndarray) -> h5py.Dataset:
+        raise NotImplementedError
+
+    def flush_buffer(self, align: bool = False) -> bool:
+        raise NotImplementedError
+
+    def _flush_time_expired(self) -> bool:
+        if self._flush_period is None:
+            return False
+        if self._last_flush is None:
+            self._last_flush = time.time()
+            return False
+        return (time.time() - self._last_flush) >= self._flush_period
+
+    def _flush_hdf5(self) -> None:
+        """Explicit HDF5 flush for non-locking readers."""
+        self._file.flush()

ewoksdata/data/hdf5/dataset_writer/_non_sequential.py

@@ -0,0 +1,139 @@
+import logging
+import time
+from typing import List
+from typing import Optional
+
+import h5py
+import numpy
+from numpy.typing import ArrayLike
+
+from ..config import guess_dataset_config
+from ..types import StrictPositiveIntegral
+from ._base import DatasetWriterBase
+
+logger = logging.getLogger(__name__)
+
+
+class IndexedDatasetWriter(DatasetWriterBase):
+    """Append arrays of the same shape to a new HDF5 dataset in a potentially non-sequential manner.
+
+    Instead of creating a dataset with the :code:`h5py` API
+
+    .. code-block::python
+
+        h5group["mydataset"] = [[1,2,3], [4,5,6], [7,8,9]]
+
+    it can be done like this
+
+    .. code-block::python
+
+        with IndexedDatasetWriter(h5group, "mydataset") as writer:
+            writer.add_point(1, [4,5,6])
+            writer.add_points([0, 2], [[1,2,3], [7,8,9]])
+
+    Chunk size determination, compression and flushing is handled.
+    """
+
+    def __init__(
+        self,
+        parent: h5py.Group,
+        name: str,
+        npoints: Optional[StrictPositiveIntegral] = None,
+        attrs: Optional[dict] = None,
+        flush_period: Optional[float] = None,
+        overwrite: bool = False,
+    ) -> None:
+        super().__init__(parent, name, attrs=attrs, flush_period=flush_period)
+        self._npoints = npoints
+        self._overwrite = overwrite
+        self._npoints_added: int = 0
+
+        self._buffer: List[tuple[int, ArrayLike]] = []
+
+    def _create_dataset(self, first_data_point: numpy.ndarray) -> h5py.Dataset:
+        scan_shape = (self._npoints,)
+        detector_shape = first_data_point.shape
+        dtype = first_data_point.dtype
+
+        if self._npoints is None:
+            shape = (1,) + detector_shape
+            max_shape = scan_shape + detector_shape
+        else:
+            shape = scan_shape + detector_shape
+            max_shape = None
+
+        options = guess_dataset_config(
+            scan_shape, detector_shape, dtype=dtype, max_shape=max_shape
+        )
+        options.update(
+            shape=shape,
+            dtype=dtype,
+            fillvalue=numpy.nan,  # converts to 0 for integers
+        )
+        if max_shape:
+            options["maxshape"] = max_shape
+
+        if self._overwrite and self._name in self._parent:
+            del self._parent[self._name]
+
+        dset = self._parent.create_dataset(self._name, **options)
+        if self._attrs:
+            dset.attrs.update(self._attrs)
+        return dset
+
+    @property
+    def npoints_added(self) -> int:
+        return self._npoints_added
+
+    def add_point(
+        self,
+        data: ArrayLike,
+        index: int,
+    ) -> bool:
+        """Append one array to the dataset."""
+        if self._dataset is None:
+            self._dataset = self._create_dataset(numpy.asarray(data))
+
+        self._buffer.append((index, data))
+        self._npoints_added += 1
+        return self.flush_buffer()
+
+    def add_points(self, data: ArrayLike, indices: ArrayLike) -> bool:
+        """Append several arrays at once to the dataset."""
+        if self._dataset is None:
+            self._dataset = self._create_dataset(numpy.asarray(data[0]))
+
+        self._buffer.extend(zip(indices, data))
+        self._npoints_added += len(data)
+        return self.flush_buffer()
+
+    def flush_buffer(self, align: bool = False) -> bool:
+        if not self._buffer:
+            return False
+
+        if align:
+            logger.debug(
+                "Chunk-aligned flushing is only supported by DatasetWriter. "
+                "Only flush when `flush_period` has passed since the last flush."
+            )
+            if self._flush_time_expired():
+                return False
+
+        indices, values = zip(*self._buffer)
+        indices = numpy.asarray(indices)
+        values = numpy.asarray(values)
+
+        # Enlarge the dataset when needed
+        max_index = int(indices.max())
+        if max_index >= self._dataset.shape[0]:
+            self._dataset.resize(max_index + 1, axis=0)
+
+        # Copy data from buffer to HDF5
+        self._dataset[indices] = values
+
+        # Remove copied data from buffer
+        self._buffer.clear()
+
+        self._flush_hdf5()
+        self._last_flush = time.time()
+        return True

ewoksdata/data/hdf5/dataset_writer/_non_sequential_stack.py

@@ -0,0 +1,158 @@
+import logging
+import time
+from typing import List
+from typing import Optional
+
+import h5py
+import numpy
+from numpy.typing import ArrayLike
+
+from ..config import guess_dataset_config
+from ..types import StrictPositiveIntegral
+from ._base import DatasetWriterBase
+
+logger = logging.getLogger(__name__)
+
+
+class IndexedStackDatasetWriter(DatasetWriterBase):
+    """Append arrays of the same shape to each item of a new HDF5 dataset
+    in a potentially non-sequential manner per item. So each item of the HDF5 dataset is a
+    stack to which we can append data in a sequential manner.
+
+    Instead of creating a dataset with the :code:`h5py` API
+
+    .. code-block::python
+
+        stack0 = [[1,2,3], [4,5,6], [7,8,9]]
+        stack1 = [[10,11,12], [13,14,15], [16,17,18]]
+        h5group["mydataset"] = [stack0, stack1]
+
+    it can be done like this
+
+    .. code-block::python
+
+        with IndexedStackDatasetWriter(h5group, "mydataset") as writer:
+            writer.add_point([4,5,6], 0, 1)
+            writer.add_point([13,14,15], 1, 1)
+            writer.add_points([[10,11,12], [16,17,18]], [1, 1], [0, 2])
+            writer.add_points([[1,2,3], [7,8,9]], [0, 0], [0, 2])
+
+    Chunk size determination, compression and flushing is handled.
+    """
+
+    def __init__(
+        self,
+        parent: h5py.Group,
+        name: str,
+        npoints: Optional[StrictPositiveIntegral] = None,
+        nstack: Optional[StrictPositiveIntegral] = None,
+        attrs: Optional[dict] = None,
+        flush_period: Optional[float] = None,
+    ) -> None:
+        super().__init__(parent, name, attrs=attrs, flush_period=flush_period)
+        self._npoints = npoints
+        self._nstack = nstack
+        self._buffer: List[tuple[int, int, ArrayLike]] = []
+
+    def _create_dataset(
+        self, first_data_point: ArrayLike, stack_index: int
+    ) -> h5py.Dataset:
+        scan_shape = (self._nstack, self._npoints)
+        first_data_point = numpy.asarray(first_data_point)
+        detector_shape = first_data_point.shape
+        dtype = first_data_point.dtype
+
+        if self._nstack is None or self._npoints is None:
+            shape = (1, 1) + detector_shape
+            max_shape = scan_shape + detector_shape
+        else:
+            shape = scan_shape + detector_shape
+            max_shape = None
+
+        options = guess_dataset_config(
+            scan_shape, detector_shape, dtype=dtype, max_shape=max_shape
+        )
+        options.update(
+            shape=shape,
+            dtype=dtype,
+            fillvalue=numpy.nan,  # converts to 0 for integers
+        )
+        if max_shape:
+            options["maxshape"] = max_shape
+
+        dset = self._parent.create_dataset(self._name, **options)
+        if self._attrs:
+            dset.attrs.update(self._attrs)
+        return dset
+
+    def add_point(self, data: ArrayLike, stack_index: int, point_index: int) -> bool:
+        """Append one array to one stack of the dataset."""
+        if self._dataset is None:
+            self._dataset = self._create_dataset(data, stack_index)
+
+        self._buffer.append((stack_index, point_index, data))
+        return self.flush_buffer()
+
+    def add_points(
+        self,
+        data: ArrayLike,
+        stack_indices: ArrayLike,
+        point_indices: ArrayLike,
+    ) -> bool:
+        """Append several arrays at once to one stack of the dataset."""
+        if not (len(stack_indices) == len(point_indices) == len(data)):
+            raise ValueError(
+                "stack_indices, point_indices, and data must have same length"
+            )
+
+        if self._dataset is None:
+            self._dataset = self._create_dataset(data[0], stack_indices[0])
+
+        self._buffer.extend(
+            zip(
+                map(int, stack_indices),
+                map(int, point_indices),
+                data,
+            )
+        )
+        return self.flush_buffer()
+
+    def flush_buffer(self, align: bool = False) -> bool:
+        if not self._buffer:
+            return False
+
+        if align:
+            logger.debug(
+                "Chunk-aligned flushing is only supported by StackDatasetWriter. "
+                "Only flush when `flush_period` has passed since the last flush."
+            )
+            if self._flush_time_expired():
+                return False
+
+        stack_indices, point_indices, values = zip(*self._buffer)
+        stack_indices = numpy.asarray(stack_indices)
+        point_indices = numpy.asarray(point_indices)
+        values = numpy.asarray(values)
+
+        # Enlarge the dataset when needed
+        max_stack = int(stack_indices.max())
+        max_point = int(point_indices.max())
+
+        shape0, shape1 = self._dataset.shape[:2]
+
+        if max_stack >= shape0:
+            self._dataset.resize(max_stack + 1, axis=0)
+
+        if max_point >= shape1:
+            self._dataset.resize(max_point + 1, axis=1)
+
+        # Copy data from buffer to HDF5
+        for i0, i1, value in zip(stack_indices, point_indices, values):
+            self._dataset[i0, i1] = value
+
+        # Remove copied data from buffer
+        self._buffer.clear()
+
+        self._flush_hdf5()
+        self._last_flush = time.time()
+        return True

ewoksdata/data/hdf5/dataset_writer/_sequential.py

@@ -0,0 +1,143 @@
+import time
+from typing import List
+from typing import Optional
+
+import h5py
+import numpy
+from numpy.typing import ArrayLike
+
+from ..config import guess_dataset_config
+from ..types import StrictPositiveIntegral
+from ._base import DatasetWriterBase
+
+
+class DatasetWriter(DatasetWriterBase):
+    """Append arrays of the same shape to a new HDF5 dataset in a sequential manner.
+
+    Instead of creating a dataset with the :code:`h5py` API
+
+    .. code-block::python
+
+        h5group["mydataset"] = [[1,2,3], [4,5,6], [7,8,9]]
+
+    it can be done like this
+
+    .. code-block::python
+
+        with DatasetWriter(h5group, "mydataset") as writer:
+            writer.add_point([1,2,3])
+            writer.add_points([[4,5,6], [7,8,9]])
+
+    Chunk size determination, chunk-aligned writing, compression and flushing is handled.
+    """
+
+    def __init__(
+        self,
+        parent: h5py.Group,
+        name: str,
+        npoints: Optional[StrictPositiveIntegral] = None,
+        attrs: Optional[dict] = None,
+        flush_period: Optional[float] = None,
+        overwrite: bool = False,
+    ) -> None:
+        super().__init__(parent, name, attrs=attrs, flush_period=flush_period)
+        self._npoints = npoints
+        self._overwrite = overwrite
+        self._chunked: bool = False
+        self._npoints_added: int = 0
+
+        self._buffer: List[ArrayLike] = list()
+        self._chunk_size: int = 0
+        self._flushed_size: int = 0
+
+    def _create_dataset(self, first_data_point: numpy.ndarray) -> h5py.Dataset:
+        scan_shape = (self._npoints,)
+        detector_shape = first_data_point.shape
+        dtype = first_data_point.dtype
+
+        if self._npoints is None:
+            shape = (1,) + detector_shape
+            max_shape = scan_shape + detector_shape
+        else:
+            shape = scan_shape + detector_shape
+            max_shape = None
+
+        options = guess_dataset_config(
+            scan_shape, detector_shape, dtype=dtype, max_shape=max_shape
+        )
+        options.update(
+            shape=shape,
+            dtype=dtype,
+            fillvalue=numpy.nan,  # converts to 0 for integers
+        )
+
+        if max_shape:
+            options["maxshape"] = max_shape
+
+        if options["chunks"]:
+            self._chunked = True
+            self._chunk_size = options["chunks"][0]
+
+        if self._overwrite and self._name in self._parent:
+            del self._parent[self._name]
+
+        dset = self._parent.create_dataset(self._name, **options)
+        if self._attrs:
+            dset.attrs.update(self._attrs)
+        return dset
+
+    @property
+    def npoints_added(self) -> int:
+        return self._npoints_added
+
+    def add_point(self, data: ArrayLike) -> bool:
+        """Append one array to the dataset."""
+        if self._dataset is None:
+            self._dataset = self._create_dataset(data)
+
+        self._buffer.append(data)
+        self._npoints_added += 1
+        return self.flush_buffer(align=True)
+
+    def add_points(self, data: ArrayLike) -> bool:
+        """Append several arrays at once to the dataset."""
+        if self._dataset is None:
+            self._dataset = self._create_dataset(data[0])
+
+        self._buffer.extend(data)
+        self._npoints_added += len(data)
+        return self.flush_buffer(align=True)
+
+    def flush_buffer(self, align: bool = False) -> bool:
+        # Determine how many points to flush
+        chunk_size = len(self._buffer)
+
+        if self._flush_time_expired():
+            flush_size = chunk_size
+        elif align and self._chunked:
+            n = chunk_size + (self._flushed_size % self._chunk_size)
+            flush_size = n // self._chunk_size * self._chunk_size
+            flush_size = min(flush_size, chunk_size)
+        else:
+            flush_size = chunk_size
+
+        if flush_size == 0:
+            return False
+
+        # Enlarge the dataset when needed
+        nalloc = self._dataset.shape[0]
+        istart = self._flushed_size
+        flushed_size = istart + flush_size
+        if self._chunked and flushed_size > nalloc:
+            self._dataset.resize(flushed_size, axis=0)
+
+        # Copy data from buffer to HDF5
+        self._dataset[istart : istart + flush_size] = self._buffer[:flush_size]
+
+        # Remove copied data from buffer
+        self._buffer = self._buffer[flush_size:]
+        self._flushed_size = flushed_size
+
+        self._flush_hdf5()
+        self._last_flush = time.time()
+        return True