Merge pull request #98 from ManuelHu/histogram

add histogram LGDO

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
     "colorlog",
     "h5py>=3.2",
     "hdf5plugin",
+    "hist",
     "numba!=0.53.*,!=0.54.*",
     "numexpr",
     "numpy>=1.21",

src/lgdo/__init__.py

@@ -33,6 +33,8 @@
   :class:`dict`
 * :class:`.Table`: a :class:`.Struct` whose elements ("columns") are all array
   types with the same length (number of rows)
+* :class:`.Histogram`: holds an array of histogrammed data, and the associated
+  binning of arbitrary dimensionality.
 
 Currently the primary on-disk format for LGDO object is LEGEND HDF5 (LH5) files. IO
 is done via the class :class:`.lh5_store.LH5Store`. LH5 files can also be
@@ -50,6 +52,7 @@
     ArrayOfEncodedEqualSizedArrays,
     ArrayOfEqualSizedArrays,
     FixedSizeArray,
+    Histogram,
     Scalar,
     Struct,
     Table,
@@ -63,6 +66,7 @@
     "ArrayOfEqualSizedArrays",
     "ArrayOfEncodedEqualSizedArrays",
     "FixedSizeArray",
+    "Histogram",
     "LGDO",
     "Scalar",
     "Struct",

src/lgdo/lh5/_serializers/__init__.py

@@ -7,6 +7,7 @@
     _h5_read_ndarray,
 )
 from .read.composite import (
+    _h5_read_histogram,
     _h5_read_lgdo,
     _h5_read_struct,
     _h5_read_table,
@@ -32,6 +33,7 @@
     "_h5_read_array_of_equalsized_arrays",
     "_h5_read_struct",
     "_h5_read_table",
+    "_h5_read_histogram",
     "_h5_read_scalar",
     "_h5_read_array_of_encoded_equalsized_arrays",
     "_h5_read_vector_of_encoded_vectors",

src/lgdo/lh5/_serializers/read/composite.py

@@ -13,6 +13,7 @@
     ArrayOfEncodedEqualSizedArrays,
     ArrayOfEqualSizedArrays,
     FixedSizeArray,
+    Histogram,
     Scalar,
     Struct,
     Table,
@@ -168,6 +169,19 @@ def _h5_read_lgdo(
             decompress=decompress,
         )
 
+    if lgdotype is Histogram:
+        return _h5_read_histogram(
+            h5o,
+            start_row=start_row,
+            n_rows=n_rows,
+            idx=idx,
+            use_h5idx=use_h5idx,
+            field_mask=field_mask,
+            obj_buf=obj_buf,
+            obj_buf_start=obj_buf_start,
+            decompress=decompress,
+        )
+
     if lgdotype is ArrayOfEncodedEqualSizedArrays:
         return _h5_read_array_of_encoded_equalsized_arrays(
             h5o,
@@ -385,3 +399,49 @@ def _h5_read_table(
     utils.check_obj_buf_attrs(obj_buf.attrs, attrs, h5g)
 
     return obj_buf, n_rows_read
+
+
+def _h5_read_histogram(
+    h5g,
+    start_row=0,
+    n_rows=sys.maxsize,
+    idx=None,
+    use_h5idx=False,
+    field_mask=None,
+    obj_buf=None,
+    obj_buf_start=0,
+    decompress=True,
+):
+    if obj_buf is not None or obj_buf_start != 0:
+        msg = "reading a histogram into an existing object buffer is not supported"
+        raise LH5DecodeError(msg, h5g)
+
+    struct, n_rows_read = _h5_read_struct(
+        h5g,
+        start_row,
+        n_rows,
+        idx,
+        use_h5idx,
+        field_mask,
+        decompress,
+    )
+    isdensity = struct.isdensity.value
+    binning = []
+    for _, a in struct.binning.items():
+        be = a.binedges
+        if isinstance(be, Struct):
+            b = (None, be.first.value, be.last.value, be.step.value, a.closedleft.value)
+        elif isinstance(be, Array):
+            b = (be, None, None, None, a.closedleft.value)
+        else:
+            msg = "unexpected binning of histogram"
+            raise LH5DecodeError(msg, h5g)
+        ax = Histogram.Axis(*b)
+        # copy attrs to "clone" the "whole" struct.
+        ax.attrs = a.getattrs(datatype=True)
+        ax["binedges"].attrs = be.getattrs(datatype=True)
+        binning.append(ax)
+    weights = struct.weights.view_as("np")
+    histogram = Histogram(weights, binning, isdensity)
+
+    return histogram, n_rows_read

src/lgdo/lh5/_serializers/write/composite.py

@@ -63,8 +63,15 @@ def _h5_write_lgdo(
         msg = f"can't overwrite '{name}' in wo_mode 'write_safe'"
         raise LH5EncodeError(msg, lh5_file, group, name)
 
-    # struct or table or waveform table
+    # struct, table, waveform table or histogram.
     if isinstance(obj, types.Struct):
+        if isinstance(obj, types.Histogram) and wo_mode not in ["w", "o", "of"]:
+            msg = f"can't append-write histogram in wo_mode '{wo_mode}'"
+            raise LH5EncodeError(msg, lh5_file, group, name)
+        if isinstance(obj, types.Histogram) and write_start != 0:
+            msg = f"can't write histogram in wo_mode '{wo_mode}' with write_start != 0"
+            raise LH5EncodeError(msg, lh5_file, group, name)
+
         return _h5_write_struct(
             obj,
             name,

src/lgdo/lh5/datatype.py

@@ -14,6 +14,7 @@
             lgdo.ArrayOfEncodedEqualSizedArrays,
             r"^array_of_encoded_equalsized_arrays<1,1>\{.+\}$",
         ),
+        (lgdo.Histogram, r"^struct\{binning,weights,isdensity\}$"),
         (lgdo.Struct, r"^struct\{.*\}$"),
         (lgdo.Table, r"^table\{.*\}$"),
         (lgdo.FixedSizeArray, r"^fixedsize_array<\d+>\{.+\}$"),

src/lgdo/lh5_store.py

@@ -20,6 +20,7 @@
     ArrayOfEncodedEqualSizedArrays,  # noqa: F401
     ArrayOfEqualSizedArrays,  # noqa: F401
     FixedSizeArray,  # noqa: F401
+    Histogram,  # noqa: F401
     Scalar,
     Struct,
     Table,  # noqa: F401

src/lgdo/types/__init__.py

@@ -6,6 +6,7 @@
 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 @@
     "ArrayOfEqualSizedArrays",
     "ArrayOfEncodedEqualSizedArrays",
     "FixedSizeArray",
+    "Histogram",
     "LGDO",
     "Scalar",
     "Struct",