ewoksdata 2.2.0__tar.gz → 2.3.0rc1__tar.gz

{ewoksdata-2.2.0/src/ewoksdata.egg-info → ewoksdata-2.3.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ewoksdata
-Version: 2.2.0
+Version: 2.3.0rc1
 Summary: Ewoks tasks for data access
 Author-email: ESRF <dau-pydev@esrf.fr>
 License: # MIT License

{ewoksdata-2.2.0 → ewoksdata-2.3.0rc1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ewoksdata"
-version = "2.2.0"
+version = "2.3.0rc1"
 authors = [{ name = "ESRF", email = "dau-pydev@esrf.fr" }]
 description = "Ewoks tasks for data access"
 readme = { file = "README.md", content-type = "text/markdown" }

{ewoksdata-2.2.0 → ewoksdata-2.3.0rc1}/src/ewoksdata/data/blissdata/blissdatav2.py

@@ -26,12 +26,14 @@ if _BLISSDATA_VERSION >= Version("2.3.0"):
     from blissdata.exceptions import IndexNoMoreThereError
     from blissdata.exceptions import IndexNotYetThereError
     from blissdata.exceptions import IndexWontBeThereError
+    from blissdata.streams import BaseStream
     from blissdata.streams import CursorGroup
 else:
     from blissdata.redis_engine.exceptions import EndOfStream
     from blissdata.redis_engine.exceptions import IndexNoMoreThereError
     from blissdata.redis_engine.exceptions import IndexNotYetThereError
     from blissdata.redis_engine.exceptions import IndexWontBeThereError
+    from blissdata.streams.base import BaseStream
     from blissdata.streams.base import CursorGroup
 
 logger = logging.getLogger(__name__)
@@ -178,10 +180,8 @@ def last_lima_image(scan_key: str, lima_name: str) -> ArrayLike:
 
 
 def _get_streams(
-    scan_key: str,
-    lima_names: List[str],
-    counter_names: List[str],
-):
+    scan_key: str, lima_names: List[str], counter_names: List[str]
+) -> Dict[str, BaseStream]:
     data_store = _get_data_store()
     scan = data_store.load_scan(scan_key)
 
@@ -191,20 +191,17 @@ def _get_streams(
     streams = dict()
 
     for name, stream in scan.streams.items():
-        if (
-            stream.event_stream.encoding["type"] == "json"
-            and "lima" in stream.info["format"]
-        ):
-            if name.split(":")[-2] in lima_names:
-                streams[name.split(":")[-2]] = stream
-
-        elif stream.info["plugin"] == "lima2":
-            for lima_name in lima_names:
-                if name == f"{lima_name}:frame":
-                    streams[lima_name] = stream
-
-        elif name.split(":")[-1] in counter_names:
-            streams[name.split(":")[-1]] = stream
+        name_parts = name.split(":")
+        if _is_lima_image_stream(stream, name_parts):
+            input_name = name_parts[0]
+            if input_name not in lima_names:
+                continue
+        elif name_parts[-1] in counter_names:
+            input_name = name_parts[-1]
+        else:
+            continue
+
+        streams[input_name] = stream
 
     nnames = len(lima_names) + len(counter_names)
     nstreams = len(streams)
@@ -215,6 +212,22 @@ def _get_streams(
     return streams
 
 
+def _is_lima_image_stream(stream: BaseStream, name_parts: Tuple[str]) -> bool:
+    if "plugin" in stream.info:
+        if stream.info["plugin"] == "lima2":
+            return (
+                name_parts[-1] == "frame" and stream.info.get("master_file") is not None
+            )
+        else:
+            return stream.info["plugin"] == "lima"
+
+    if stream.event_stream.encoding["type"] == "json":
+        # lima_v1, lima_v2
+        return "lima" in stream.info.get("format", "")
+
+    return False
+
+
 def iter_bliss_scan_data_from_memory_slice(
     scan_key: str,
     lima_names: List[str],

ewoksdata-2.3.0rc1/src/ewoksdata/data/hdf5/dataset_writer/_base.py

@@ -0,0 +1,112 @@
+import time
+from typing import Iterator
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
+
+import h5py
+import numpy
+
+
+class DatasetWriterBase:
+    """Write an HDF5 dataset incrementally with support for
+
+    - chunk size determination
+    - compression algorithm selection
+    - chunk-aligned flushing of the buffer to libhdf5 memory
+    - opt-out explicit flushing of libhdf5 memory to disk
+
+    Flushing happens in two stages, both controlled by the `flush_period`
+    argument:
+
+    - Buffer -> libhdf5: adding data writes the buffered data that covers
+      entire chunks to HDF5. The remaining buffered data is written on an
+      explicit `flush_buffer()` call, when the writer exits and depending
+      on `flush_period`:
+
+      - `None`: never when adding data.
+      - `0`: every time data is added.
+      - `> 0`: when data is added after the period expired.
+
+      Datasets without chunking never buffer.
+
+    - libhdf5 -> disk: an explicit HDF5 flush makes updates visible to
+      non-locking readers. It happens when the writer exits and depending
+      on `flush_period`:
+
+      - `None` (default): never before the writer exits.
+      - `0`: every time data is written to HDF5.
+      - `> 0`: when data is written to HDF5 after the period expired.
+    """
+
+    def __init__(
+        self,
+        parent: h5py.Group,
+        name: str,
+        attrs: Optional[dict] = None,
+        flush_period: Optional[float] = 1.0,
+    ) -> 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()
+        self._flush_hdf5()  # in case throttled
+
+    @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:
+        """Never expires when no flush period is defined."""
+        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 so non-locking readers to see updates on retry."""
+        # NOTE: flushing a file with compressed datasets is expensive
+        self._file.flush()
+        self._last_flush = time.time()
+
+    def _flush_hdf5_throttled(self) -> bool:
+        """HDF5 flush when the flush period expired."""
+        if self._flush_time_expired():
+            self._flush_hdf5()
+            return True
+        return False
+
+
+def iter_contiguous_slices(
+    indices: Sequence[int],
+) -> Iterator[Tuple[numpy.ndarray, slice]]:
+    """Split indices in runs of consecutive values."""
+    indices = numpy.asarray(indices)
+    order = numpy.argsort(indices, kind="stable")
+    sorted_indices = indices[order]
+    breaks = numpy.nonzero(numpy.diff(sorted_indices) != 1)[0] + 1
+    for positions in numpy.split(order, breaks):
+        start = int(indices[positions[0]])
+        yield positions, slice(start, start + len(positions))

{ewoksdata-2.2.0 → ewoksdata-2.3.0rc1}/src/ewoksdata/data/hdf5/dataset_writer/_non_sequential.py

@@ -1,5 +1,5 @@
 import logging
-import time
+from typing import Dict
 from typing import List
 from typing import Optional
 
@@ -10,6 +10,7 @@ from numpy.typing import ArrayLike
 from ..config import guess_dataset_config
 from ..types import StrictPositiveIntegral
 from ._base import DatasetWriterBase
+from ._base import iter_contiguous_slices
 
 logger = logging.getLogger(__name__)
 
@@ -30,8 +31,6 @@ class IndexedDatasetWriter(DatasetWriterBase):
         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__(
@@ -40,15 +39,19 @@ class IndexedDatasetWriter(DatasetWriterBase):
         name: str,
         npoints: Optional[StrictPositiveIntegral] = None,
         attrs: Optional[dict] = None,
-        flush_period: Optional[float] = None,
+        flush_period: Optional[float] = 1.0,
         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._chunk_size: int = 0
         self._npoints_added: int = 0
 
-        self._buffer: List[tuple[int, ArrayLike]] = []
+        # Points are buffered by index
+        self._buffer: Dict[int, ArrayLike] = dict()
+        self._nbuffered_per_chunk: Dict[int, int] = dict()
 
     def _create_dataset(self, first_data_point: numpy.ndarray) -> h5py.Dataset:
         scan_shape = (self._npoints,)
@@ -73,6 +76,10 @@ class IndexedDatasetWriter(DatasetWriterBase):
         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]
 
@@ -94,46 +101,88 @@ class IndexedDatasetWriter(DatasetWriterBase):
         if self._dataset is None:
             self._dataset = self._create_dataset(numpy.asarray(data))
 
-        self._buffer.append((index, data))
+        self._buffer_point(index, data)
         self._npoints_added += 1
-        return self.flush_buffer()
+        return self.flush_buffer(align=True)
 
     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))
+        for index, value in zip(indices, data):
+            self._buffer_point(index, value)
         self._npoints_added += len(data)
-        return self.flush_buffer()
+        return self.flush_buffer(align=True)
+
+    def _buffer_point(self, index: int, data: ArrayLike) -> None:
+        index = int(index)
+        if self._chunked and index not in self._buffer:
+            chunk_index = index // self._chunk_size
+            nbuffered = self._nbuffered_per_chunk.get(chunk_index, 0)
+            self._nbuffered_per_chunk[chunk_index] = nbuffered + 1
+        self._buffer[index] = data
 
     def flush_buffer(self, align: bool = False) -> bool:
+        # Determine which buffered points to flush
         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
+        if self._flush_time_expired():
+            flush_indices = list(self._buffer)
+        elif align and self._chunked:
+            flush_indices = self._complete_chunk_indices()
+        else:
+            flush_indices = list(self._buffer)
 
-        indices, values = zip(*self._buffer)
-        indices = numpy.asarray(indices)
-        values = numpy.asarray(values)
+        if not flush_indices:
+            return False
 
         # Enlarge the dataset when needed
-        max_index = int(indices.max())
+        max_index = max(flush_indices)
         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
+        values = numpy.asarray([self._buffer[index] for index in flush_indices])
+        for source, destination in iter_contiguous_slices(flush_indices):
+            self._dataset[destination] = values[source]
 
         # Remove copied data from buffer
-        self._buffer.clear()
+        if len(flush_indices) == len(self._buffer):
+            self._buffer.clear()
+            self._nbuffered_per_chunk.clear()
+        else:
+            for index in flush_indices:
+                del self._buffer[index]
+                chunk_index = index // self._chunk_size
+                nbuffered = self._nbuffered_per_chunk[chunk_index] - 1
+                if nbuffered:
+                    self._nbuffered_per_chunk[chunk_index] = nbuffered
+                else:
+                    del self._nbuffered_per_chunk[chunk_index]
+
+        return self._flush_hdf5_throttled()
+
+    def _complete_chunk_indices(self) -> List[int]:
+        """Indices of the buffered points that cover entire chunks."""
+        chunk_indices = sorted(
+            chunk_index
+            for chunk_index, nbuffered in self._nbuffered_per_chunk.items()
+            if nbuffered == self._chunk_npoints(chunk_index)
+        )
+        return [
+            index
+            for chunk_index in chunk_indices
+            for index in range(
+                chunk_index * self._chunk_size,
+                chunk_index * self._chunk_size + self._chunk_npoints(chunk_index),
+            )
+        ]
 
-        self._flush_hdf5()
-        self._last_flush = time.time()
-        return True
+    def _chunk_npoints(self, chunk_index: int) -> int:
+        """Number of points in a chunk. The last chunk can be smaller when
+        the number of points is fixed."""
+        if self._npoints is None:
+            return self._chunk_size
+        return min(self._chunk_size, self._npoints - chunk_index * self._chunk_size)

ewoksdata-2.3.0rc1/src/ewoksdata/data/hdf5/dataset_writer/_non_sequential_stack.py

@@ -0,0 +1,238 @@
+import logging
+from typing import Dict
+from typing import Optional
+from typing import Tuple
+
+import h5py
+import numpy
+from numpy.typing import ArrayLike
+
+from ..config import guess_dataset_config
+from ..types import StrictPositiveIntegral
+from ._base import DatasetWriterBase
+from ._base import iter_contiguous_slices
+
+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])
+    """
+
+    def __init__(
+        self,
+        parent: h5py.Group,
+        name: str,
+        npoints: Optional[StrictPositiveIntegral] = None,
+        nstack: Optional[StrictPositiveIntegral] = None,
+        attrs: Optional[dict] = None,
+        flush_period: Optional[float] = 1.0,
+    ) -> None:
+        super().__init__(parent, name, attrs=attrs, flush_period=flush_period)
+        self._npoints = npoints
+        self._nstack = nstack
+        self._chunked: bool = False
+        self._chunk_shape: Tuple[int, int] = (0, 0)
+
+        # Points are buffered by (stack index, point index).
+        self._buffer: Dict[Tuple[int, int], ArrayLike] = dict()
+        self._nbuffered_per_chunk: Dict[Tuple[int, int], int] = dict()
+
+    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
+
+        if options["chunks"]:
+            self._chunked = True
+            self._chunk_shape = tuple(options["chunks"][:2])
+
+        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_point(stack_index, point_index, data)
+        return self.flush_buffer(align=True)
+
+    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])
+
+        for stack_index, point_index, value in zip(stack_indices, point_indices, data):
+            self._buffer_point(stack_index, point_index, value)
+        return self.flush_buffer(align=True)
+
+    def _buffer_point(
+        self, stack_index: int, point_index: int, data: ArrayLike
+    ) -> None:
+        position = int(stack_index), int(point_index)
+        if self._chunked and position not in self._buffer:
+            chunk_index = (
+                position[0] // self._chunk_shape[0],
+                position[1] // self._chunk_shape[1],
+            )
+            self._nbuffered_per_chunk[chunk_index] = (
+                self._nbuffered_per_chunk.get(chunk_index, 0) + 1
+            )
+        self._buffer[position] = data
+
+    def flush_buffer(self, align: bool = False) -> bool:
+        if not self._buffer:
+            return False
+
+        if self._flush_time_expired():
+            flushed = self._flush_all()
+        elif align and self._chunked:
+            flushed = self._flush_complete_chunks()
+        else:
+            flushed = self._flush_all()
+
+        if not flushed:
+            return False
+        return self._flush_hdf5_throttled()
+
+    def _flush_all(self) -> bool:
+        """Write all buffered points to HDF5."""
+        positions = list(self._buffer)
+        stack_indices = numpy.array([stack_index for stack_index, _ in positions])
+        point_indices = numpy.array([point_index for _, point_index in positions])
+        values = numpy.asarray(list(self._buffer.values()))
+
+        self._enlarge_dataset(int(stack_indices.max()), int(point_indices.max()))
+
+        for stack in numpy.unique(stack_indices):
+            mask = stack_indices == stack
+            stack_values = values[mask]
+            for source, destination in iter_contiguous_slices(point_indices[mask]):
+                self._dataset[int(stack), destination] = stack_values[source]
+
+        self._buffer.clear()
+        self._nbuffered_per_chunk.clear()
+        return True
+
+    def _flush_complete_chunks(self) -> bool:
+        """Write the buffered points that cover entire chunks to HDF5."""
+        chunk_indices = sorted(
+            chunk_index
+            for chunk_index, nbuffered in self._nbuffered_per_chunk.items()
+            if nbuffered == self._chunk_npoints(chunk_index)
+        )
+        if not chunk_indices:
+            return False
+
+        chunk_extents = [
+            (
+                chunk_index,
+                self._chunk_slice(0, chunk_index),
+                self._chunk_slice(1, chunk_index),
+            )
+            for chunk_index in chunk_indices
+        ]
+
+        self._enlarge_dataset(
+            max(stacks.stop for _, stacks, _ in chunk_extents) - 1,
+            max(points.stop for _, _, points in chunk_extents) - 1,
+        )
+
+        for chunk_index, stacks, points in chunk_extents:
+            block = numpy.asarray(
+                [
+                    [
+                        self._buffer[(stack_index, point_index)]
+                        for point_index in range(points.start, points.stop)
+                    ]
+                    for stack_index in range(stacks.start, stacks.stop)
+                ]
+            )
+            self._dataset[stacks, points] = block
+
+            for stack_index in range(stacks.start, stacks.stop):
+                for point_index in range(points.start, points.stop):
+                    del self._buffer[(stack_index, point_index)]
+            del self._nbuffered_per_chunk[chunk_index]
+
+        return True
+
+    def _enlarge_dataset(self, max_stack: int, max_point: int) -> None:
+        """Enlarge the dataset when needed."""
+        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)
+
+    def _chunk_slice(self, axis: int, chunk_index: Tuple[int, int]) -> slice:
+        """Dataset range covered by a chunk along the stack axis (0) or
+        point axis (1). Edge chunks can be smaller when the number of
+        stacks or points is fixed."""
+        nmax = self._nstack if axis == 0 else self._npoints
+        start = chunk_index[axis] * self._chunk_shape[axis]
+        stop = start + self._chunk_shape[axis]
+        if nmax is not None:
+            stop = min(stop, nmax)
+        return slice(start, stop)
+
+    def _chunk_npoints(self, chunk_index: Tuple[int, int]) -> int:
+        """Number of points in a chunk. Edge chunks can be smaller when the
+        number of stacks or points is fixed."""
+        stacks = self._chunk_slice(0, chunk_index)
+        points = self._chunk_slice(1, chunk_index)
+        return (stacks.stop - stacks.start) * (points.stop - points.start)

{ewoksdata-2.2.0 → ewoksdata-2.3.0rc1}/src/ewoksdata/data/hdf5/dataset_writer/_sequential.py

@@ -1,4 +1,3 @@
-import time
 from typing import List
 from typing import Optional
 
@@ -27,8 +26,6 @@ class DatasetWriter(DatasetWriterBase):
         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__(
@@ -37,7 +34,7 @@ class DatasetWriter(DatasetWriterBase):
         name: str,
         npoints: Optional[StrictPositiveIntegral] = None,
         attrs: Optional[dict] = None,
-        flush_period: Optional[float] = None,
+        flush_period: Optional[float] = 1.0,
         overwrite: bool = False,
     ) -> None:
         super().__init__(parent, name, attrs=attrs, flush_period=flush_period)
@@ -110,16 +107,16 @@ class DatasetWriter(DatasetWriterBase):
 
     def flush_buffer(self, align: bool = False) -> bool:
         # Determine how many points to flush
-        chunk_size = len(self._buffer)
+        buffer_size = len(self._buffer)
 
         if self._flush_time_expired():
-            flush_size = chunk_size
+            flush_size = buffer_size
         elif align and self._chunked:
-            n = chunk_size + (self._flushed_size % self._chunk_size)
+            n = buffer_size + (self._flushed_size % self._chunk_size)
             flush_size = n // self._chunk_size * self._chunk_size
-            flush_size = min(flush_size, chunk_size)
+            flush_size = min(flush_size, buffer_size)
         else:
-            flush_size = chunk_size
+            flush_size = buffer_size
 
         if flush_size == 0:
             return False
@@ -138,6 +135,4 @@ class DatasetWriter(DatasetWriterBase):
         self._buffer = self._buffer[flush_size:]
         self._flushed_size = flushed_size
 
-        self._flush_hdf5()
-        self._last_flush = time.time()
-        return True
+        return self._flush_hdf5_throttled()