legend-pydataobj 1.8.1__py3-none-any.whl → 1.10.0__py3-none-any.whl

lgdo/lh5/store.py

@@ -5,13 +5,16 @@ HDF5 files.
 
 from __future__ import annotations
 
+import bisect
 import logging
 import os
 import sys
 from collections.abc import Mapping, Sequence
+from inspect import signature
 from typing import Any
 
 import h5py
+import numpy as np
 from numpy.typing import ArrayLike
 
 from .. import types
@@ -34,7 +37,9 @@ class LH5Store:
     lgdo.waveformtable.WaveformTable
     """
 
-    def __init__(self, base_path: str = "", keep_open: bool = False) -> None:
+    def __init__(
+        self, base_path: str = "", keep_open: bool = False, locking: bool = False
+    ) -> None:
         """
         Parameters
         ----------
@@ -43,12 +48,21 @@ class LH5Store:
         keep_open
             whether to keep files open by storing the :mod:`h5py` objects as
             class attributes.
+        locking
+            whether to lock files when reading
         """
         self.base_path = "" if base_path == "" else utils.expand_path(base_path)
         self.keep_open = keep_open
+        self.locking = locking
         self.files = {}
 
-    def gimme_file(self, lh5_file: str | h5py.File, mode: str = "r") -> h5py.File:
+    def gimme_file(
+        self,
+        lh5_file: str | h5py.File,
+        mode: str = "r",
+        page_buffer: int = 0,
+        **file_kwargs,
+    ) -> h5py.File:
         """Returns a :mod:`h5py` file object from the store or creates a new one.
 
         Parameters
@@ -57,12 +71,20 @@ class LH5Store:
             LH5 file name.
         mode
             mode in which to open file. See :class:`h5py.File` documentation.
+        page_buffer
+            enable paged aggregation with a buffer of this size in bytes
+            Only used when creating a new file. Useful when writing a file
+            with a large number of small datasets. This is a short-hand for
+            ``(fs_stragety="page", fs_pagesize=[page_buffer])``
+        file_kwargs
+            Keyword arguments for :class:`h5py.File`
         """
         if isinstance(lh5_file, h5py.File):
             return lh5_file
 
         if mode == "r":
             lh5_file = utils.expand_path(lh5_file, base_path=self.base_path)
+            file_kwargs["locking"] = self.locking
 
         if lh5_file in self.files:
             return self.files[lh5_file]
@@ -72,20 +94,30 @@ class LH5Store:
         else:
             full_path = lh5_file
 
+        file_exists = os.path.exists(full_path)
         if mode != "r":
             directory = os.path.dirname(full_path)
             if directory != "" and not os.path.exists(directory):
                 log.debug(f"making path {directory}")
                 os.makedirs(directory)
 
-        if mode == "r" and not os.path.exists(full_path):
+        if mode == "r" and not file_exists:
             msg = f"file {full_path} not found"
             raise FileNotFoundError(msg)
+        if not file_exists:
+            mode = "w"
 
-        if mode != "r" and os.path.exists(full_path):
+        if mode != "r" and file_exists:
             log.debug(f"opening existing file {full_path} in mode '{mode}'")
 
-        h5f = h5py.File(full_path, mode)
+        if mode == "w":
+            file_kwargs.update(
+                {
+                    "fs_strategy": "page",
+                    "fs_page_size": page_buffer,
+                }
+            )
+        h5f = h5py.File(full_path, mode, **file_kwargs)
 
         if self.keep_open:
             self.files[lh5_file] = h5f
@@ -135,6 +167,7 @@ class LH5Store:
         obj_buf: types.LGDO = None,
         obj_buf_start: int = 0,
         decompress: bool = True,
+        **file_kwargs,
     ) -> tuple[types.LGDO, int]:
         """Read LH5 object data from a file in the store.
 
@@ -143,13 +176,62 @@ class LH5Store:
         .lh5.core.read
         """
         # grab files from store
-        if not isinstance(lh5_file, (str, h5py.File)):
-            lh5_obj = [self.gimme_file(f, "r")[name] for f in list(lh5_file)]
+        if isinstance(lh5_file, (str, h5py.File)):
+            lh5_obj = self.gimme_file(lh5_file, "r", **file_kwargs)[name]
         else:
-            lh5_obj = self.gimme_file(lh5_file, "r")[name]
-
+            lh5_files = list(lh5_file)
+            n_rows_read = 0
+
+            for i, h5f in enumerate(lh5_files):
+                if (
+                    isinstance(idx, (list, tuple))
+                    and len(idx) > 0
+                    and not np.isscalar(idx[0])
+                ):
+                    # a list of lists: must be one per file
+                    idx_i = idx[i]
+                elif idx is not None:
+                    # make idx a proper tuple if it's not one already
+                    if not (isinstance(idx, tuple) and len(idx) == 1):
+                        idx = (idx,)
+                    # idx is a long continuous array
+                    n_rows_i = utils.read_n_rows(name, h5f)
+                    # find the length of the subset of idx that contains indices
+                    # that are less than n_rows_i
+                    n_rows_to_read_i = bisect.bisect_left(idx[0], n_rows_i)
+                    # now split idx into idx_i and the remainder
+                    idx_i = np.array(idx[0])[:n_rows_to_read_i]
+                    idx = np.array(idx[0])[n_rows_to_read_i:] - n_rows_i
+                else:
+                    idx_i = None
+                n_rows_i = n_rows - n_rows_read
+
+                obj_buf, n_rows_read_i = self.read(
+                    name,
+                    h5f,
+                    start_row,
+                    n_rows_i,
+                    idx_i,
+                    use_h5idx,
+                    field_mask,
+                    obj_buf,
+                    obj_buf_start,
+                    decompress,
+                )
+
+                n_rows_read += n_rows_read_i
+                if n_rows_read >= n_rows or obj_buf is None:
+                    return obj_buf, n_rows_read
+                start_row = 0
+                obj_buf_start += n_rows_read_i
+            return obj_buf, n_rows_read
+
+        if isinstance(idx, (list, tuple)) and len(idx) > 0 and not np.isscalar(idx[0]):
+            idx = idx[0]
         return _serializers._h5_read_lgdo(
-            lh5_obj,
+            lh5_obj.id,
+            lh5_obj.file.filename,
+            lh5_obj.name,
             start_row=start_row,
             n_rows=n_rows,
             idx=idx,
@@ -170,6 +252,7 @@ class LH5Store:
         n_rows: int | None = None,
         wo_mode: str = "append",
         write_start: int = 0,
+        page_buffer: int = 0,
         **h5py_kwargs,
     ) -> None:
         """Write an LGDO into an LH5 file.
@@ -199,10 +282,17 @@ class LH5Store:
         # write_object:overwrite.
         mode = "w" if wo_mode == "of" else "a"
 
+        file_kwargs = {
+            k: h5py_kwargs[k]
+            for k in h5py_kwargs & signature(h5py.File).parameters.keys()
+        }
+
         return _serializers._h5_write_lgdo(
             obj,
             name,
-            self.gimme_file(lh5_file, mode=mode),
+            self.gimme_file(
+                lh5_file, mode=mode, page_buffer=page_buffer, **file_kwargs
+            ),
             group=group,
             start_row=start_row,
             n_rows=n_rows,

lgdo/lh5/tools.py

@@ -128,7 +128,7 @@ def show(
 
     # open file
     if isinstance(lh5_file, str):
-        lh5_file = h5py.File(utils.expand_path(lh5_file), "r")
+        lh5_file = h5py.File(utils.expand_path(lh5_file), "r", locking=False)
 
     # go to group
     if lh5_group != "/":

lgdo/lh5/utils.py

@@ -125,7 +125,12 @@ def get_h5_group(
         else:
             group = base_group.create_group(group)
             if grp_attrs is not None:
-                group.attrs.update(grp_attrs)
+                group.attrs.update(
+                    {
+                        k: v.encode("utf-8") if isinstance(v, str) else v
+                        for k, v in grp_attrs.items()
+                    }
+                )
             return group
     if (
         grp_attrs is not None
@@ -141,7 +146,13 @@ def get_h5_group(
         log.debug(f"overwriting {group}.attrs...")
         for key in group.attrs:
             group.attrs.pop(key)
-        group.attrs.update(grp_attrs)
+
+        group.attrs.update(
+            {
+                k: v.encode("utf-8") if isinstance(v, str) else v
+                for k, v in grp_attrs.items()
+            }
+        )
 
     return group
 

lgdo/lh5_store.py

@@ -20,6 +20,7 @@ from .types import (
     ArrayOfEncodedEqualSizedArrays,  # noqa: F401
     ArrayOfEqualSizedArrays,  # noqa: F401
     FixedSizeArray,  # noqa: F401
+    Histogram,  # noqa: F401
     Scalar,
     Struct,
     Table,  # noqa: F401

lgdo/types/__init__.py

@@ -6,6 +6,7 @@ from .array import Array
 from .arrayofequalsizedarrays import ArrayOfEqualSizedArrays
 from .encoded import ArrayOfEncodedEqualSizedArrays, VectorOfEncodedVectors
 from .fixedsizearray import FixedSizeArray
+from .histogram import Histogram
 from .lgdo import LGDO
 from .scalar import Scalar
 from .struct import Struct
@@ -18,6 +19,7 @@ __all__ = [
     "ArrayOfEqualSizedArrays",
     "ArrayOfEncodedEqualSizedArrays",
     "FixedSizeArray",
+    "Histogram",
     "LGDO",
     "Scalar",
     "Struct",

lgdo/types/histogram.py

@@ -0,0 +1,419 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable
+from typing import Any
+
+import hist
+import numpy as np
+from numpy.typing import NDArray
+
+from .array import Array
+from .lgdo import LGDO
+from .scalar import Scalar
+from .struct import Struct
+
+log = logging.getLogger(__name__)
+
+
+class Histogram(Struct):
+    class Axis(Struct):
+        def __init__(
+            self,
+            edges: NDArray | Array | None,
+            first: float | None,
+            last: float | None,
+            step: float | None,
+            closedleft: bool = True,
+            binedge_attrs: dict[str, Any] | None = None,
+        ) -> None:
+            """
+            A special struct to group axis parameters for use in a :class:`Histogram`.
+
+            Depending on the parameters, an axis either can have
+
+            * a binning described by a range object, if ``first``, ``last`` and ``step``
+              are passed, or
+            * a variable binning described by the ``edges`` array.
+
+            Parameters
+            ----------
+            edges
+                an array of edges that describe the binning of this axis.
+            first
+                left edge of the leftmost bin
+            last
+                right edge of the rightmost bin
+            step
+                step size (width of each bin)
+            closedleft
+                if True, the bin intervals are left-closed :math:`[a,b)`;
+                if False, intervals are right-closed :math:`(a,b]`.
+            binedge_attrs
+                attributes that will be added to the ``binedges`` LGDO that
+                is part of the axis struct.
+            """
+            if edges is not None and (
+                first is not None or last is not None or step is not None
+            ):
+                msg = "can only construct Axis either from edges or from range"
+                raise ValueError(msg)
+            if edges is None and (first is None or last is None or step is None):
+                msg = "did not pass all range parameters"
+                raise ValueError(msg)
+
+            if edges is None:
+                edges = Struct(
+                    {
+                        "first": Scalar(first),
+                        "last": Scalar(last),
+                        "step": Scalar(step),
+                    },
+                    binedge_attrs,
+                )
+            else:
+                if not isinstance(edges, Array):
+                    edges = Array(edges, attrs=binedge_attrs)
+                elif binedge_attrs is not None:
+                    msg = "passed both binedge as Array LGDO instance and binedge_attrs"
+                    raise ValueError(msg)
+
+                if len(edges.nda.shape) != 1:
+                    msg = "must pass an array<1>{real} as edges vector"
+                    raise ValueError(msg)
+
+            super().__init__({"binedges": edges, "closedleft": Scalar(closedleft)})
+
+        @classmethod
+        def from_edges(
+            cls,
+            edges: NDArray | Iterable[float],
+            binedge_attrs: dict[str, Any] | None = None,
+        ) -> Histogram.Axis:
+            """Create a new axis with variable binning described by ``edges``."""
+            edges = np.array(edges)
+            return cls(edges, None, None, None, True, binedge_attrs)
+
+        @classmethod
+        def from_range_edges(
+            cls,
+            edges: NDArray | Iterable[float],
+            binedge_attrs: dict[str, Any] | None = None,
+        ) -> Histogram.Axis:
+            """Create a new axis from the binning described by ``edges``, but try to convert it to
+            a evenly-spaced range object first.
+
+            .. warning ::
+
+                This function might return a wrong binning, especially in the case of very small
+                magnitudes of the spacing. See the documentation of :func:`numpy.isclose` for
+                details. Use this function only with caution, if you know the binning's order of
+                magniutude.
+            """
+            edges = np.array(edges)
+            edge_diff = np.diff(edges)
+            if np.any(~np.isclose(edge_diff, edge_diff[0])):
+                return cls(edges, None, None, None, True, binedge_attrs)
+            return cls(None, edges[0], edges[-1], edge_diff[0], True, binedge_attrs)
+
+        @property
+        def is_range(self) -> bool:
+            return isinstance(self["binedges"], Struct)
+
+        @property
+        def first(self) -> float:
+            if not self.is_range:
+                msg = "Axis is not a range"
+                raise TypeError(msg)
+            return self["binedges"]["first"].value
+
+        @property
+        def last(self) -> float:
+            if not self.is_range:
+                msg = "Axis is not a range"
+                raise TypeError(msg)
+            return self["binedges"]["last"].value
+
+        @property
+        def step(self) -> float:
+            if not self.is_range:
+                msg = "Axis is not a range"
+                raise TypeError(msg)
+            return self["binedges"]["step"].value
+
+        @property
+        def closedleft(self) -> bool:
+            return self["closedleft"].value
+
+        @property
+        def nbins(self) -> int:
+            """Return the number of bins, both for variable and range binning."""
+            if self.is_range:
+                bins = (self.last - self.first) / self.step
+                bins_int = int(np.rint(bins))
+                assert np.isclose(bins, bins_int)
+                return bins_int
+            return len(self["binedges"].nda) - 1
+
+        @property
+        def edges(self) -> NDArray:
+            """Return all binedges, both for variable and range binning."""
+            if self.is_range:
+                return np.linspace(self.first, self.last, self.nbins + 1)
+            return self["binedges"].nda
+
+        def __str__(self) -> str:
+            thr_orig = np.get_printoptions()["threshold"]
+            np.set_printoptions(threshold=8)
+
+            if self.is_range:
+                string = f"first={self.first}, last={self.last}, step={self.step}"
+            else:
+                string = f"edges={self.edges}"
+            string += f", closedleft={self.closedleft}"
+
+            attrs = self.get_binedgeattrs()
+            if attrs:
+                string += f" with attrs={attrs}"
+
+            np.set_printoptions(threshold=thr_orig)
+            return string
+
+        def get_binedgeattrs(self, datatype: bool = False) -> dict:
+            """Return a copy of the LGDO attributes dictionary of the binedges
+
+            Parameters
+            ----------
+            datatype
+                if ``False``, remove ``datatype`` attribute from the output
+                dictionary.
+            """
+            return self["binedges"].getattrs(datatype)
+
+    def __init__(
+        self,
+        weights: hist.Hist | NDArray | Array,
+        binning: None
+        | Iterable[Histogram.Axis]
+        | Iterable[NDArray]
+        | Iterable[tuple[float, float, float]] = None,
+        isdensity: bool = False,
+        attrs: dict[str, Any] | None = None,
+        binedge_attrs: dict[str, Any] | None = None,
+        flow: bool = True,
+    ) -> None:
+        """A special struct to contain histogrammed data.
+
+        Parameters
+        ----------
+        weights
+            An :class:`numpy.ndarray` to be used for this object's internal
+            array, or a :class:`hist.Hist` object, whose data view is used
+            for this object's internal array.
+            Note: the array/histogram view is used directly, not copied
+        binning
+            * has to by None if a :class:`hist.Hist` has been passed as ``weights``
+            * can be a list of pre-initialized :class:`Histogram.Axis`
+            * can be a list of tuples, each representing a range, ``(first, last, step)``
+            * can be a list of numpy arrays, as returned by :func:`numpy.histogramdd`.
+        isdensity
+            If True, all bin contents represent a density (amount per volume), and not
+            an absolute amount.
+        binedge_attrs
+            attributes that will be added to the all ``binedges`` of all axes.
+            This does not work if :class:`Histogram.Axis` instances are directly passed
+            as binning.
+        attrs
+            a set of user attributes to be carried along with this LGDO.
+        flow
+            If ``False``, discard counts in over-/underflow bins of the passed
+            :class:`hist.Hist` instance. If ``True``, this data will also be discarded,
+            but a warning is emitted.
+
+            .. note ::
+
+                :class:`Histogram` does not support storing counts in overflow or
+                underflow bins. This parameter just controls, whether a warning will
+                be emitted.
+        """
+        if isinstance(weights, hist.Hist):
+            if binning is not None:
+                msg = "not allowed to pass custom binning if constructing from hist.Hist instance"
+                raise ValueError(msg)
+            if isdensity:
+                msg = "not allowed to pass isdensity=True if constructing from hist.Hist instance"
+                raise ValueError(msg)
+
+            if weights.sum(flow=True) != weights.sum(flow=False) and flow:
+                log.warning(
+                    "flow bins of hist.Hist cannot be represented, their counts are discarded"
+                )
+            weights_view = weights.view(flow=False)
+            if type(weights_view) is not np.ndarray:
+                msg = "only simple numpy-backed storages can be used in a hist.Hist"
+                raise ValueError(msg)
+            w = Array(weights_view)
+
+            b = []
+            for ax in weights.axes:
+                if not isinstance(ax, (hist.axis.Regular, hist.axis.Variable)):
+                    msg = "only regular or variable axes of hist.Hist can be converted"
+                    raise ValueError(msg)
+                if isinstance(ax, hist.axis.Regular):
+                    step = (ax.edges[-1] - ax.edges[0]) / ax.size
+                    bax = Histogram.Axis(
+                        None, ax.edges[0], ax.edges[-1], step, True, binedge_attrs
+                    )
+                    b.append(bax)
+                else:
+                    b.append(Histogram.Axis.from_edges(ax.edges, binedge_attrs))
+        else:
+            if binning is None:
+                msg = "need to also pass binning if passing histogram as array"
+                raise ValueError(msg)
+            w = weights if isinstance(weights, Array) else Array(weights)
+
+            if all(isinstance(ax, Histogram.Axis) for ax in binning):
+                if binedge_attrs is not None:
+                    msg = "passed both binedges as Axis instances and binedge_attrs"
+                    raise ValueError(msg)
+                b = binning
+            elif all(isinstance(ax, np.ndarray) for ax in binning):
+                b = [Histogram.Axis.from_edges(ax, binedge_attrs) for ax in binning]
+            elif all(isinstance(ax, tuple) for ax in binning):
+                b = [Histogram.Axis(None, *ax, True, binedge_attrs) for ax in binning]
+            else:
+                msg = "invalid binning object passed"
+                raise ValueError(msg)
+
+            if len(binning) != len(w.nda.shape):
+                msg = "binning and weight dimensions do not match"
+                raise ValueError(msg)
+            for i, ax in enumerate(b):
+                if ax.nbins != w.nda.shape[i]:
+                    msg = f"bin count does not match weight count along axis {i}"
+                    raise ValueError(msg)
+
+        b = Struct({f"axis_{i}": a for i, a in enumerate(b)})
+
+        super().__init__(
+            {"binning": b, "weights": w, "isdensity": Scalar(isdensity)},
+            attrs,
+        )
+
+    @property
+    def isdensity(self) -> bool:
+        return self["isdensity"].value
+
+    @property
+    def weights(self) -> Array:
+        return self["weights"]
+
+    @property
+    def binning(self) -> tuple[Histogram.Axis, ...]:
+        bins = sorted(self["binning"].items())
+        assert all(isinstance(v, Histogram.Axis) for k, v in bins)
+        return tuple(v for _, v in bins)
+
+    def __setitem__(self, name: str, obj: LGDO) -> None:
+        # do not allow for new attributes on this
+        msg = "histogram fields cannot be mutated"
+        raise TypeError(msg)
+
+    def __getattr__(self, name: str) -> None:
+        # do not allow for new attributes on this
+        msg = "histogram fields cannot be mutated"
+        raise TypeError(msg)
+
+    def add_field(self, name: str | int, obj: LGDO) -> None:  # noqa: ARG002
+        """
+        .. error ::
+
+            Not applicable: A histogram cannot be used as a struct
+        """
+        msg = "histogram fields cannot be mutated"
+        raise TypeError(msg)
+
+    def remove_field(self, name: str | int, delete: bool = False) -> None:  # noqa: ARG002
+        """
+        .. error ::
+
+            Not applicable: A histogram cannot be used as a struct
+        """
+        msg = "histogram fields cannot be mutated"
+        raise TypeError(msg)
+
+    def __str__(self) -> str:
+        string = "{\n"
+        for k, v in enumerate(self.binning):
+            string += f" 'axis_{k}': {v},\n"
+        string += "}"
+
+        attrs = self.getattrs()
+        if attrs:
+            string += f" with attrs={attrs}"
+
+        return string
+
+    def view_as(
+        self,
+        library: str,
+    ) -> tuple[NDArray] | hist.Hist:
+        r"""View the histogram data as a third-party format data structure.
+
+        This is typically a zero-copy or nearly zero-copy operation.
+
+        Supported third-party formats are:
+
+        - ``np``: returns a tuple of binning and an :class:`np.ndarray`, similar
+          to the return value of :func:`numpy.histogramdd`.
+        - ``hist``: returns an :class:`hist.Hist` that holds **a copy** of this
+          histogram's data.
+
+        Warning
+        -------
+        Viewing as ``hist`` will perform a copy of the stored histogram data.
+
+        Parameters
+        ----------
+        library
+            format of the returned data view.
+
+        See Also
+        --------
+        .LGDO.view_as
+        """
+        if library == "hist":
+            if self.isdensity:
+                msg = "hist.Hist cannot represent density histograms"
+                raise ValueError(msg)
+
+            hist_axes = []
+            for a in self.binning:
+                if not a.closedleft:
+                    msg = "hist.Hist cannot represent right-closed intervals"
+                    raise ValueError(msg)
+                if a.is_range:
+                    hist_ax = hist.axis.Regular(
+                        bins=a.nbins,
+                        start=a.first,
+                        stop=a.last,
+                        underflow=False,
+                        overflow=False,
+                    )
+                else:
+                    hist_ax = hist.axis.Variable(
+                        a.edges,
+                        underflow=False,
+                        overflow=False,
+                    )
+                hist_axes.append(hist_ax)
+
+            return hist.Hist(*hist_axes, data=self.weights.view_as("np"))
+
+        if library == "np":
+            edges = tuple([a.edges for a in self.binning])
+            return self.weights.view_as("np"), edges
+
+        msg = f"{library!r} is not a supported third-party format."
+        raise TypeError(msg)

lgdo/types/table.py

@@ -168,7 +168,7 @@ class Table(Struct):
         self.add_field(name, obj, use_obj_size=use_obj_size)
 
     def remove_column(self, name: str, delete: bool = False) -> None:
-        """Alias for :meth:`.remove_field` using table terminology 'column'."""
+        """Alias for :meth:`Struct.remove_field` using table terminology 'column'."""
         super().remove_field(name, delete)
 
     def join(