feat: add read/write support

Co-authored-by: Aryaman Jeendgar <jeendgararyaman@gmail.com>
Signed-off-by: Henry Schreiner <henryschreineriii@gmail.com>

Author: henryiii

pyproject.toml

@@ -47,6 +47,9 @@ schema = [
   "fastjsonschema",
   "importlib-resources; python_version<'3.9'",
 ]
+hdf5 = [
+  "h5py",
+]
 
 [dependency-groups]
 docs = [
@@ -62,6 +65,7 @@ test = [
   "boost-histogram>=1.0",
   "fastjsonschema",
   "importlib-resources; python_version<'3.9'",
+  "h5py",
 ]
 dev = [{ include-group = "test"}]
 
@@ -89,7 +93,7 @@ warn_unreachable = true
 enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
 
 [[tool.mypy.overrides]]
-module = ["fastjsonschema"]
+module = ["fastjsonschema", "h5py"]
 ignore_missing_imports = true
 
 

src/uhi/io/__init__.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import h5py
+import numpy as np
+
+from ..typing.serialization import AnyAxis, AnyHistogram, AnyStorage, Histogram
+
+__all__ = ["read", "write"]
+
+
+def __dir__() -> list[str]:
+    return __all__
+
+
+def write(grp: h5py.Group, /, histogram: AnyHistogram) -> None:
+    """
+    Write a histogram to an HDF5 group.
+    """
+    # All referenced objects will be stored inside of /{name}/ref_axes
+    hist_folder_storage = grp.create_group("ref_axes")
+
+    # Metadata
+
+    if "metadata" in histogram:
+        metadata_grp = grp.create_group("metadata")
+        for key, val1 in histogram["metadata"].items():
+            metadata_grp.attrs[key] = val1
+
+    # Axes
+    axes_dataset = grp.create_dataset(
+        "axes", len(histogram["axes"]), dtype=h5py.special_dtype(ref=h5py.Reference)
+    )
+    for i, axis in enumerate(histogram["axes"]):
+        # Iterating through the axes, calling `create_axes_object` for each of them,
+        # creating references to new groups and appending it to the `items` dataset defined above
+        ax_group = hist_folder_storage.create_group(f"axis_{i}")
+        ax_info = axis.copy()
+        ax_metadata = ax_info.pop("metadata", None)
+        ax_edges = np.asarray(ax_info.pop("edges", None))
+        ax_cats: list[int] | list[str] | None = ax_info.pop("categories", None)
+        for key, val2 in ax_info.items():
+            ax_group.attrs[key] = val2
+        if ax_metadata is not None:
+            ax_metadata_grp = ax_group.create_group("metadata")
+            for k, v in ax_metadata.items():
+                ax_metadata_grp.attrs[k] = v
+        if ax_edges is not None:
+            ax_group.create_dataset("edges", shape=ax_edges.shape, data=ax_edges)
+        if ax_cats is not None:
+            ax_group.create_dataset("items", shape=len(ax_cats), data=ax_cats)
+        axes_dataset[i] = ax_group.ref
+
+    # Storage
+    storage_grp = grp.create_group("storage")
+    storage_type = histogram["storage"]["type"]
+
+    storage_grp.attrs["type"] = storage_type
+
+    for key, val3 in histogram["storage"].items():
+        if key == "type":
+            continue
+        npvalue = np.asarray(val3)
+        storage_grp.create_dataset(key, shape=npvalue.shape, data=npvalue)
+
+
+def read(grp: h5py.Group, /) -> Histogram:
+    """
+    Read a histogram from an HDF5 group.
+    """
+    axes_grp = grp["axes"]
+    axes_ref = grp["ref_axes"]
+    assert isinstance(axes_ref, h5py.Group)
+    assert isinstance(axes_grp, h5py.Dataset)
+
+    axes = [AnyAxis(**axes_ref[unref_axis_ref].attrs) for unref_axis_ref in axes_ref]  # type: ignore[typeddict-item]
+
+    storage_grp = grp["storage"]
+    assert isinstance(storage_grp, h5py.Group)
+    storage = AnyStorage(type=storage_grp.attrs["type"])
+    for key in storage_grp:
+        storage[key] = np.asarray(storage_grp[key])  # type: ignore[literal-required]
+
+    histogram_dict = AnyHistogram(axes=axes, storage=storage)
+    if "metadata" in grp:
+        histogram_dict["metadata"] = dict(grp["metadata"].attrs.items())
+
+    return histogram_dict  # type: ignore[return-value]

src/uhi/io/hdf5.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+__all__ = ["default", "object_hook"]
+
+
+CONVERT_KEYS = [
+    "values",
+    "variances",
+    "edges",
+    "counts",
+    "sum_of_weights",
+    "sum_of_weights_squared",
+]
+
+
+def __dir__() -> list[str]:
+    return __all__
+
+
+def default(obj: Any, /) -> Any:
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()  # Convert ndarray to list
+    msg = f"Object of type {type(obj)} is not JSON serializable"
+    raise TypeError(msg)
+
+
+def object_hook(dct: dict[str, Any], /) -> dict[str, Any]:
+    """
+    Decode a histogram from a dictionary.
+    """
+
+    for item in CONVERT_KEYS:
+        if item in dct and isinstance(dct[item], list):
+            dct[item] = np.asarray(dct[item])
+
+    return dct

src/uhi/io/json.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import json
+import zipfile
+from typing import Any
+
+import numpy as np
+
+from ..typing.serialization import AnyHistogram, Histogram
+from .json import CONVERT_KEYS
+
+__all__ = ["read", "write"]
+
+
+def write(
+    zip_file: zipfile.ZipFile,
+    /,
+    name: str,
+    histogram: AnyHistogram,
+) -> None:
+    """
+    Write a histogram to a zip file.
+    """
+    # Write out numpy arrays to files in the zipfile
+    for storage_key in list(histogram["storage"]):
+        if storage_key in CONVERT_KEYS:
+            path = f"{name}_storage_{storage_key}.npy"
+            np.save(zip_file.open(path, "w"), histogram["storage"][storage_key])  # type: ignore[literal-required]
+            histogram["storage"][storage_key] = path  # type: ignore[literal-required]
+
+    for axis in histogram["axes"]:
+        for key in list(axis):
+            if key in CONVERT_KEYS:
+                path = f"{name}_axis_{key}.npy"
+                np.save(zip_file.open(path, "w"), axis[key])  # type: ignore[literal-required]
+                axis[key] = path  # type: ignore[literal-required]
+
+    hist_json = json.dumps(histogram)
+    zip_file.writestr(f"{name}.json", hist_json)
+
+
+def read(zip_file: zipfile.ZipFile, /, name: str) -> Histogram:
+    """
+    Read histograms from a zip file.
+    """
+
+    def object_hook(dct: dict[str, Any], /) -> dict[str, Any]:
+        for item in CONVERT_KEYS:
+            if item in dct and isinstance(dct[item], str):
+                dct[item] = np.load(zip_file.open(dct[item]))
+        return dct
+
+    with zip_file.open(f"{name}.json") as f:
+        return json.load(f, object_hook=object_hook)  # type: ignore[no-any-return]

src/uhi/io/zip.py

@@ -52,15 +52,21 @@
         }
       }
     },
+    "ndarray": {
+      "type": "array",
+      "items": {
+        "oneOf": [{ "type": "number" }, { "$ref": "#/$defs/ndarray" }]
+      },
+      "description": "A ND (nested) array of numbers."
+    },
     "data_array": {
       "oneOf": [
         {
           "type": "string",
           "description": "A path (similar to URI) to the floating point bin data"
         },
         {
-          "type": "array",
-          "items": { "type": "number" }
+          "$ref": "#/$defs/ndarray"
         }
       ]
     },

src/uhi/resources/histogram.schema.json

@@ -1,9 +1,24 @@
+"""Serialization types for UHI.
+
+Two types of dictionaries are defined here:
+
+1. ``AnyAxis``, ``AnyStorage``, and ``AnyHistogram`` are used for inputs. They represent
+   the merger of all possible types.
+2. ``Axis``, ``Storage``, and ``histogram`` are used for outputs. These have precise entries
+   defined for each Literal type.
+"""
+
 from __future__ import annotations
 
-from collections.abc import Sequence
 from typing import Literal, TypedDict, Union
 
+from numpy.typing import ArrayLike
+
 __all__ = [
+    "AnyAxis",
+    "AnyHistogram",
+    "AnyStorage",
+    "Axis",
     "BooleanAxis",
     "CategoryIntAxis",
     "CategoryStrAxis",
@@ -12,6 +27,7 @@
     "IntStorage",
     "MeanStorage",
     "RegularAxis",
+    "Storage",
     "VariableAxis",
     "WeightedMeanStorage",
     "WeightedStorage",
@@ -40,7 +56,7 @@ class RegularAxis(_RequiredRegularAxis, total=False):
 
 class _RequiredVariableAxis(TypedDict):
     type: Literal["variable"]
-    edges: list[float] | str
+    edges: ArrayLike | str
     underflow: bool
     overflow: bool
     circular: bool
@@ -80,43 +96,84 @@ class BooleanAxis(_RequiredBooleanAxis, total=False):
 
 class IntStorage(TypedDict):
     type: Literal["int"]
-    values: Sequence[int] | str
+    values: ArrayLike | str
 
 
 class DoubleStorage(TypedDict):
     type: Literal["double"]
-    values: Sequence[float] | str
+    values: ArrayLike | str
 
 
 class WeightedStorage(TypedDict):
     type: Literal["weighted"]
-    values: Sequence[float] | str
-    variances: Sequence[float] | str
+    values: ArrayLike | str
+    variances: ArrayLike | str
 
 
 class MeanStorage(TypedDict):
     type: Literal["mean"]
-    counts: Sequence[float] | str
-    values: Sequence[float] | str
-    variances: Sequence[float] | str
+    counts: ArrayLike | str
+    values: ArrayLike | str
+    variances: ArrayLike | str
 
 
 class WeightedMeanStorage(TypedDict):
     type: Literal["weighted_mean"]
-    sum_of_weights: Sequence[float] | str
-    sum_of_weights_squared: Sequence[float] | str
-    values: Sequence[float] | str
-    variances: Sequence[float] | str
+    sum_of_weights: ArrayLike | str
+    sum_of_weights_squared: ArrayLike | str
+    values: ArrayLike | str
+    variances: ArrayLike | str
+
+
+Storage = Union[
+    IntStorage, DoubleStorage, WeightedStorage, MeanStorage, WeightedMeanStorage
+]
+
+Axis = Union[RegularAxis, VariableAxis, CategoryStrAxis, CategoryIntAxis, BooleanAxis]
+
+
+class _RequiredAnyStorage(TypedDict):
+    type: Literal["int", "double", "weighted", "mean", "weighted_mean"]
+
+
+class AnyStorage(_RequiredAnyStorage, total=False):
+    values: ArrayLike | str
+    variances: ArrayLike | str
+    sum_of_weights: ArrayLike | str
+    sum_of_weights_squared: ArrayLike | str
+    counts: ArrayLike | str
+
+
+class _RequiredAnyAxis(TypedDict):
+    type: Literal["regular", "variable", "category_str", "category_int", "boolean"]
+
+
+class AnyAxis(_RequiredAnyAxis, total=False):
+    metadata: dict[str, SupportedMetadata]
+    lower: float
+    upper: float
+    bins: int
+    edges: ArrayLike | str
+    categories: list[str] | list[int]
+    underflow: bool
+    overflow: bool
+    flow: bool
+    circular: bool
 
 
 class _RequiredHistogram(TypedDict):
-    axes: list[
-        RegularAxis | VariableAxis | CategoryStrAxis | CategoryIntAxis | BooleanAxis
-    ]
-    storage: (
-        IntStorage | DoubleStorage | WeightedStorage | MeanStorage | WeightedMeanStorage
-    )
+    axes: list[Axis]
+    storage: Storage
 
 
 class Histogram(_RequiredHistogram, total=False):
     metadata: dict[str, SupportedMetadata]
+
+
+class _RequiredAnyHistogram(TypedDict):
+    axes: list[AnyAxis]
+    storage: AnyStorage
+
+
+class AnyHistogram(_RequiredAnyHistogram, total=False):
+    metadata: dict[str, SupportedMetadata]