ewoksdata 2.1.0__py3-none-any.whl

ewoksdata/data/hdf5/dataset_writer/_sequential_stack.py

@@ -0,0 +1,195 @@
+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 StackDatasetWriter(DatasetWriterBase):
+    """Append arrays of the same shape to each item of a new HDF5 dataset
+    in a 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 StackDatasetWriter(h5group, "mydataset") as writer:
+            writer.add_point([1,2,3], 0)
+            writer.add_point([10,11,12], 1)
+            writer.add_points([[13,14,15], [16,17,18]], 1)
+            writer.add_points([[4,5,6], [7,8,9]], 0)
+
+    Chunk size determination, chunk-aligned writing, 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._chunked: bool = False
+        self._nstack = nstack
+
+        self._buffers: List[List[ArrayLike]] = list()
+        self._chunk_size: ArrayLike = numpy.zeros(2, dtype=int)
+        self._flushed_size_dim1: ArrayLike = numpy.array(list(), dtype=int)
+
+    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 = (stack_index + 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
+
+        if options["chunks"]:
+            self._chunked = True
+            self._chunk_size = numpy.array(options["chunks"][:2], dtype=int)
+
+        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) -> bool:
+        """Append one array to one stack of the dataset."""
+        if self._dataset is None:
+            self._dataset = self._create_dataset(data, stack_index)
+
+        buffer = self._get_buffer(stack_index)
+        buffer.append(data)
+        return self.flush_buffer(align=True)
+
+    def add_points(self, data: ArrayLike, stack_index: int) -> bool:
+        """Append several arrays at once to one stack of the dataset."""
+        if self._dataset is None:
+            self._dataset = self._create_dataset(data[0], stack_index)
+
+        buffer = self._get_buffer(stack_index)
+        buffer.extend(data)
+        return self.flush_buffer(align=True)
+
+    def flush_buffer(self, align: bool = False) -> bool:
+        # Determine how many points to flush for each buffer in the stack
+        chunk_sizes = numpy.array([len(buffer) for buffer in self._buffers])
+        flushed_size_dim1 = self._flushed_size_dim1
+        size_dim0 = len(chunk_sizes)
+        assert size_dim0 == len(
+            flushed_size_dim1
+        ), "Number of buffers and number of flushed dim1 points must be the same"
+
+        chunk_size_dim0, chunk_size_dim1 = self._chunk_size[:2]
+        if chunk_size_dim0 == 0:
+            chunk_size_dim0 = size_dim0
+
+        if self._flush_time_expired():
+            flush_sizes = chunk_sizes
+        elif align and self._chunked:
+            size_dim0 = size_dim0 // chunk_size_dim0 * chunk_size_dim0
+            chunk_sizes = chunk_sizes[:size_dim0]
+            flushed_size_dim1 = flushed_size_dim1[:size_dim0]
+            if size_dim0:
+                n1 = chunk_sizes + (flushed_size_dim1 % chunk_size_dim1)
+                flush_sizes = n1 // chunk_size_dim1 * chunk_size_dim1
+                flush_sizes = numpy.minimum(flush_sizes, chunk_sizes)
+                for i0_chunk0 in range(0, size_dim0, chunk_size_dim0):
+                    flush_sizes[i0_chunk0 : i0_chunk0 + chunk_size_dim0] = min(
+                        flush_sizes[i0_chunk0 : i0_chunk0 + chunk_size_dim0]
+                    )
+            else:
+                flush_sizes = list()
+        else:
+            flush_sizes = chunk_sizes
+
+        if not any(flush_sizes):
+            return False
+
+        # Enlarge the dataset when needed
+        nalloc = self._dataset.shape[:2]
+        istart_dim1 = flushed_size_dim1
+        flushed_size_dim1 = istart_dim1 + flush_sizes
+        nalloc_new = numpy.array([size_dim0, max(flushed_size_dim1)])
+        if self._chunked and any(nalloc_new > nalloc):
+            for axis, n in enumerate(nalloc_new):
+                self._dataset.resize(n, axis=axis)
+
+        # Copy data from buffer to HDF5
+        for i0_chunk0 in range(0, size_dim0, chunk_size_dim0):
+            idx_dim0 = slice(i0_chunk0, i0_chunk0 + chunk_size_dim0)
+            buffers = self._buffers[idx_dim0]
+
+            flush_sizes_dim1 = flush_sizes[idx_dim0]
+            non_ragged_buffers = len(set(flush_sizes_dim1)) == 1
+
+            istart0_dim1 = istart_dim1[idx_dim0]
+            non_ragged_destination = len(set(istart0_dim1)) == 1
+
+            if non_ragged_destination and non_ragged_buffers:
+                data = [buffer[: flush_sizes_dim1[0]] for buffer in buffers]
+                idx_dim1 = slice(istart0_dim1[0], istart0_dim1[0] + flush_sizes_dim1[0])
+                self._dataset[idx_dim0, idx_dim1] = data
+            else:
+                for buffer, i_dim0, istart_dim1, i_flush_size_dim1 in zip(
+                    buffers,
+                    range(i0_chunk0, i0_chunk0 + chunk_size_dim0),
+                    istart0_dim1,
+                    flush_sizes_dim1,
+                ):
+                    self._dataset[
+                        i_dim0, istart_dim1 : istart_dim1 + i_flush_size_dim1, ...
+                    ] = buffer[:i_flush_size_dim1]
+
+        # Remove copied data from buffer
+        for i0 in range(size_dim0):
+            self._buffers[i0] = self._buffers[i0][flush_sizes[i0] :]
+            self._flushed_size_dim1[i0] = flushed_size_dim1[i0]
+
+        self._flush_hdf5()
+        self._last_flush = time.time()
+        return True
+
+    def _get_buffer(self, stack_index: int) -> List[ArrayLike]:
+        # Add stack buffers when needed
+        for _ in range(max(stack_index - len(self._buffers) + 1, 0)):
+            self._buffers.append(list())
+            self._flushed_size_dim1 = numpy.append(self._flushed_size_dim1, 0)
+        return self._buffers[stack_index]

ewoksdata/data/hdf5/types.py

@@ -0,0 +1,14 @@
+from numbers import Integral
+from typing import NewType
+from typing import Optional
+from typing import Tuple
+
+PositiveIntegral = NewType("PositiveIntegral", Integral)  # >= 0
+StrictPositiveIntegral = NewType("StrictPositiveIntegral", Integral)  # > 0
+ShapeType = Tuple[StrictPositiveIntegral]
+SizeType = PositiveIntegral
+
+VarShapeType = Tuple[Optional[PositiveIntegral]]  # 0 or None mark a variable dimension
+VarH5pyShapeType = Tuple[
+    Optional[StrictPositiveIntegral]
+]  # None marks a variable dimension

ewoksdata/data/nexus.py

@@ -0,0 +1,187 @@
+import os
+from contextlib import contextmanager
+from typing import Iterator
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
+from typing import Union
+
+import h5py
+from silx.io import h5py_utils
+
+from .url import DataUrl
+from .url import as_dataurl
+
+
+def ensure_nxclass(group: h5py.Group) -> None:
+    """
+    Ensure the HDF5 group has an appropriate ``NX_class`` attribute.
+
+    The attribute is assigned based on the group's depth in the hierarchy:
+      * Root group (``/``): ``NXroot``
+      * First-level group: ``NXentry``
+      * Deeper groups: ``NXcollection``
+
+    :param group: The HDF5 group to update.
+    :type group: h5py.Group
+    """
+    if group.attrs.get("NX_class"):
+        return
+    groups = [s for s in group.name.split("/") if s]
+    n = len(groups)
+    if n == 0:
+        group.attrs["NX_class"] = "NXroot"
+    elif n == 1:
+        group.attrs["NX_class"] = "NXentry"
+    else:
+        group.attrs["NX_class"] = "NXcollection"
+
+
+def select_default_plot(nxdata: h5py.Group) -> None:
+    """
+    Set the default NXdata group for visualization in the parent hierarchy.
+
+    The function walks up the group tree and sets the ``default`` attribute
+    in each parent to point to the given NXdata group.
+
+    :param nxdata: An HDF5 group representing an NXdata node.
+    :type nxdata: h5py.Group
+    """
+    parent = nxdata.parent
+    for name in nxdata.name.split("/")[::-1]:
+        if not name:
+            continue
+        parent.attrs["default"] = name
+        parent = parent.parent
+
+
+def create_url(url: str, overwrite: bool = False, **open_options) -> DataUrl:
+    """
+    Create (or ensure the existence of) an HDF5 group specified by a URL.
+
+    :param url: The HDF5 URL (e.g. ``path/to/file.h5::/entry/data``).
+    :param overwrite: If True, overwrite the final group if it exists.
+                      If False, reuse existing groups. Default is False.
+    :param open_options: Additional options passed to ``h5py.File``.
+    :return: A DataUrl object pointing to the created or existing group.
+    """
+    url = as_dataurl(url)
+    filename = url.file_path()
+    data_path = url.data_path() or "/"
+    dirname = os.path.dirname(filename)
+    if dirname:
+        os.makedirs(dirname, exist_ok=True)
+    _ = open_options.setdefault("mode", "a")
+    with h5py_utils.open_item(filename, "/", **open_options) as parent:
+        h5item, _ = _create_h5group(parent, data_path, overwrite=overwrite)
+        return as_dataurl(f"{filename}::{h5item.name}")
+
+
+def get_nxentry(h5item: Union[h5py.Dataset, h5py.Group]) -> h5py.Group:
+    """
+    Return the top-level NXentry group containing the given HDF5 item.
+
+    :param h5item: An HDF5 dataset or group.
+    :return: The NXentry group containing ``h5item``.
+    :raises ValueError: If the item does not belong to an NXentry hierarchy.
+    """
+    parts = [s for s in h5item.name.split("/") if s]
+    if parts:
+        return h5item.file[parts[0]]
+    raise ValueError("HDF5 item must be part of an NXentry")
+
+
+@contextmanager
+def create_nexus_group(
+    url: Union[str, DataUrl],
+    retry_timeout=None,
+    retry_period=None,
+    default_levels: Optional[Sequence[str]] = None,
+    **open_options,
+) -> Iterator[Tuple[h5py.Group, bool]]:
+    """
+    Context manager to create or access an HDF5 group specified by a URL.
+
+    The data path inside HDF5 is ensured to have at least ``default_levels`` levels.
+
+    :param url: The HDF5 URL (e.g. ``path/to/file.h5::/entry/data``).
+    :param retry_timeout: Maximum time in seconds to retry opening the file.
+    :param retry_period: Time in seconds between retries.
+    :param default_levels: Minimum path hierarchy touse when missing. Default is ``["/", "results"]``.
+    :param open_options: Additional options passed to ``h5py.File``.
+    :yields: Tuple of (HDF5 group, already_existed flag).
+    """
+    url = as_dataurl(url)
+    filename = url.file_path()
+    itemname = url.data_path() or "/"
+    dirname = os.path.dirname(filename)
+    if dirname:
+        os.makedirs(dirname, exist_ok=True)
+    open_options.setdefault("mode", "a")
+    with h5py_utils.open_item(
+        filename,
+        "/",
+        retry_timeout=retry_timeout,
+        retry_period=retry_period,
+        **open_options,
+    ) as root:
+        yield _create_h5group(root, itemname, default_levels=default_levels)
+
+
+def _create_h5group(
+    parent: h5py.Group,
+    data_path: str,
+    default_levels: Optional[Sequence[str]] = None,
+    overwrite: bool = False,
+) -> Tuple[h5py.Group, bool]:
+    """
+    Create or retrieve an HDF5 group under a given parent.
+
+    Ensures that intermediate groups exist, optionally enforcing a minimum
+    depth using ``default_levels``.
+
+    Missing ``NX_class`` attributes are automatically created on all levels.
+    The final group can be replaced if ``overwrite=True``.
+
+    :param parent: Parent group under which the path should be created.
+    :param data_path: Group path to create relative to ``parent``.
+    :param default_levels: Minimum hierarchy to enforce.
+                           If ``data_path`` is shorter, extra levels are added.
+    :param overwrite: If True, delete and recreate the final group if it exists.
+                      Default is False.
+    :return: Tuple of (HDF5 group, already_existed flag).
+    """
+    # Ensure the data path has at least N levels as defined by ``default_levels``
+    groups = [s for s in data_path.split("/") if s]
+    if default_levels:
+        default_levels = list(default_levels)
+    else:
+        default_levels = ["results"]
+    if len(groups) < len(default_levels):
+        groups += default_levels[len(groups) :]
+
+    # Ensure parent groups exist + ensure NX_class
+    create = False
+    ensure_nxclass(parent)
+    for group in groups[:-1]:
+        if group in parent:
+            parent = parent[group]
+        else:
+            parent = parent.create_group(group)
+            create = True
+        ensure_nxclass(parent)
+
+    # Final group: create, overwrite or confirm existence + ensure NX_class
+    group = groups[-1]
+    if group in parent:
+        if overwrite:
+            del parent[group]
+            parent = parent.create_group(group)
+        else:
+            parent = parent[group]
+    else:
+        parent = parent.create_group(group)
+        create = True
+    ensure_nxclass(parent)
+
+    return parent, not create

ewoksdata/data/url.py

@@ -0,0 +1,30 @@
+import sys
+from typing import Tuple
+from typing import Union
+
+from silx.io.url import DataUrl
+
+
+def as_dataurl(url: Union[str, DataUrl]) -> DataUrl:
+    if isinstance(url, str) and sys.platform == "win32":
+        url = url.replace("\\", "/")
+    if not isinstance(url, DataUrl):
+        url = DataUrl(url)
+    return url
+
+
+def as_h5url(url: Union[str, DataUrl]) -> DataUrl:
+    url = as_dataurl(url)
+    if not url.data_path():
+        return DataUrl(f"{url.file_path()}::/")
+    return url
+
+
+def h5dataset_url_parse(url: Union[str, DataUrl]) -> Tuple[str, str, Tuple]:
+    url = as_h5url(url)
+    filename = str(url.file_path())
+    h5path = url.data_path()
+    idx = url.data_slice()
+    if idx is None:
+        idx = tuple()
+    return filename, h5path, idx

ewoksdata/data/utils.py

@@ -0,0 +1,53 @@
+from numbers import Number
+from typing import List
+from typing import Mapping
+from typing import Optional
+from typing import Tuple
+
+import numpy
+
+
+def is_data(data):
+    if isinstance(data, (numpy.ndarray, Number)):
+        return True
+    if isinstance(data, (str, list)) and data:
+        return True
+    return False
+
+
+def data_from_storage(data, remove_numpy=True):
+    if isinstance(data, numpy.ndarray):
+        if not remove_numpy:
+            return data
+        elif data.ndim == 0:
+            return data.item()
+        else:
+            return data.tolist()
+    elif isinstance(data, Mapping):
+        return {
+            k: data_from_storage(v, remove_numpy=remove_numpy)
+            for k, v in data.items()
+            if not k.startswith("@")
+        }
+    else:
+        return data
+
+
+def _validate_names(names: Optional[List[str]], name_type: str) -> None:
+    if names is None:
+        return
+    for n in names:
+        if not isinstance(n, str) or not n.strip():
+            raise ValueError(f"Invalid {name_type} name: {n!r}")
+
+
+def validate_and_init_names(
+    lima_names: Optional[List[str]],
+    counter_names: Optional[List[str]],
+) -> Tuple[List[str], List[str]]:
+    _validate_names(lima_names, "LIMA")
+    _validate_names(counter_names, "Counter")
+    return (
+        lima_names if lima_names is not None else [],
+        counter_names if counter_names is not None else [],
+    )

ewoksdata/tests/conftest.py

@@ -0,0 +1,30 @@
+import numpy
+import pytest
+
+from . import data
+
+
+@pytest.fixture(scope="session")
+def bliss_scan(tmpdir_factory):
+    tmpdir = tmpdir_factory.mktemp("sample_dataset")
+    npoints_per_file = 3
+    npoints = 31
+    scannr = 2
+    image = numpy.zeros((10, 12))
+    return (
+        str(
+            data.save_bliss_scan(
+                tmpdir,
+                image,
+                npoints_per_file,
+                npoints,
+                scannr,
+                subscannr=1,
+                lima_names=("p3", "p4"),
+                counter_names=("diode1", "diode2"),
+                sequence="add",
+            )
+        ),
+        npoints,
+        scannr,
+    )

ewoksdata/tests/data/__init__.py

@@ -0,0 +1 @@
+from .bliss_scans import *  # noqa F403

ewoksdata/tests/data/bliss_scans.py

@@ -0,0 +1,137 @@
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+from typing import Tuple
+
+import h5py
+import numpy
+
+__all__ = ["save_bliss_scan"]
+
+
+def save_bliss_scan(
+    tmpdir: Path,
+    image: numpy.ndarray,
+    npoints_per_file: int,
+    npoints: int,
+    scannr: int,
+    subscannr: int = 1,
+    sequence: str = "add",
+    lima_names: Tuple[str, ...] = ("p3",),
+    counter_names: Tuple[str, ...] = ("diode1", "diode2"),
+    positioner_names: Tuple[str, ...] = ("diffty", "diffrz"),
+    title: Optional[str] = None,
+) -> Path:
+    dataset_filename = tmpdir / "sample_dataset.h5"
+    lima_dirname = tmpdir / f"scan{scannr:04d}"
+    lima_dirname.mkdir()
+    with h5py.File(str(dataset_filename), mode="w") as root:
+        root.attrs["NX_class"] = "NXroot"
+        root.attrs["creator"] = "bliss"
+        scan = root.create_group(f"{scannr}.{subscannr}")
+        scan.attrs["NX_class"] = "NXentry"
+        if title is not None:
+            scan["title"] = title
+
+        instrument = scan.create_group("instrument")
+        instrument.attrs["NX_class"] = "NXinstrument"
+        measurement = scan.create_group("measurement")
+        measurement.attrs["NX_class"] = "NXcollection"
+
+        if positioner_names:
+            positioners = instrument.create_group("positioners")
+            positioners.attrs["NX_class"] = "NXcollection"
+
+        for name in counter_names:
+            _save_counter(scan, name, npoints)
+        for name in lima_names:
+            _save_lima(
+                scan,
+                name,
+                image,
+                npoints,
+                npoints_per_file,
+                lima_dirname,
+                scannr,
+                sequence,
+            )
+        for name in positioner_names:
+            _save_positioner(scan, name, npoints)
+
+        scan["end_time"] = datetime.now().astimezone().isoformat()
+    return dataset_filename
+
+
+def _save_counter(scan: h5py.Group, name: str, npoints: int) -> None:
+    diode = _add_detector(scan, name)
+    diode["data"] = numpy.arange(npoints)
+
+
+def _save_lima(
+    scan: h5py.Group,
+    name: str,
+    image: numpy.ndarray,
+    npoints: int,
+    npoints_per_file: int,
+    lima_dirname: Path,
+    scannr: int,
+    sequence: str,
+) -> None:
+    tot_shape = (npoints,) + image.shape
+    lima_shape = (npoints_per_file,) + image.shape
+    nlima = npoints // npoints_per_file + bool(npoints % npoints_per_file)
+    lima_dtype = float
+    layout = h5py.VirtualLayout(tot_shape, dtype=lima_dtype)
+    for i in range(nlima):
+        with h5py.File(str(lima_dirname / f"{name}_{i}.h5"), mode="w") as limaroot:
+            limaroot.attrs["NX_class"] = "NXroot"
+            entry = limaroot.create_group("entry_0000")
+            entry.attrs["NX_class"] = "NXentry"
+            lmeasurement = entry.create_group("measurement")
+            lmeasurement.attrs["NX_class"] = "NXcollection"
+            lmeasurement["data"] = numpy.full(lima_shape, numpy.nan)
+
+            i0 = i * npoints_per_file
+            i1 = min(i0 + npoints_per_file, npoints)
+            for itot in range(i0, i1):
+                if sequence == "add":
+                    lmeasurement["data"][itot - i0, ...] = image + itot
+                else:
+                    m = itot / (npoints - 1) + 1
+                    if itot == 0:
+                        monitor = _add_detector(scan, "monitor")
+                        monitor["data"] = numpy.zeros(npoints)
+                    monitor["data"][itot] = m
+                    lmeasurement["data"][itot - i0, ...] = image * m
+
+            spath = os.path.join(f"scan{scannr:04d}", f"{name}_{i}.h5")
+            sdset = "/entry_0000/measurement/data"
+
+            sshape = (i1 - i0,) + lima_shape[1:]
+            vsource = h5py.VirtualSource(spath, sdset, shape=sshape, dtype=lima_dtype)
+
+            layout[i0:i1] = vsource[:]
+    detector = _add_detector(scan, name)
+    detector.create_virtual_dataset("data", layout, fillvalue=numpy.nan)
+
+
+def _add_detector(scan: h5py.Group, name: str) -> h5py.Group:
+    grp = scan["instrument"].create_group(name)
+    grp.attrs["NX_class"] = "NXdetector"
+    scan["measurement"][name] = h5py.SoftLink(f"{grp.name}/data")
+    return grp
+
+
+def _save_positioner(scan: h5py.Group, name: str, npoints: int) -> None:
+    positioner = _add_positioner(scan, name)
+    positioner["value"] = numpy.arange(npoints)
+
+
+def _add_positioner(scan: h5py.Group, name: str) -> h5py.Group:
+    instrument = scan["instrument"]
+    grp = instrument.create_group(name)
+    grp.attrs["NX_class"] = "Nxpositioner"
+    scan["measurement"][name] = h5py.SoftLink(f"{grp.name}/value")
+    instrument["positioners"][name] = h5py.SoftLink(f"{grp.name}/value")
+    return grp