[Python][UHI] Start introducing UHI serialization

Author: siliataider

bindings/pyroot/pythonizations/CMakeLists.txt

@@ -127,10 +127,11 @@ set(py_sources
   ROOT/_pythonization/_ttree.py
   ROOT/_pythonization/_tvector3.py
   ROOT/_pythonization/_tvectort.py
-  ROOT/_pythonization/_uhi/main.py
+  ROOT/_pythonization/_uhi/__init__.py
   ROOT/_pythonization/_uhi/tags.py
   ROOT/_pythonization/_uhi/indexing.py
   ROOT/_pythonization/_uhi/plotting.py
+  ROOT/_pythonization/_uhi/serialization.py
   ${PYROOT_EXTRA_PYTHON_SOURCES}
 )
 

bindings/pyroot/pythonizations/python/ROOT/_facade.py

@@ -479,7 +479,7 @@ def uhi(self):
         uhi_module.__file__ = "<module ROOT>"
         uhi_module.__package__ = self
         try:
-            from ._pythonization._uhi.main import _add_module_level_uhi_helpers
+            from ._pythonization._uhi import _add_module_level_uhi_helpers
 
             _add_module_level_uhi_helpers(uhi_module)
         except ImportError:

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_th1.py

@@ -241,7 +241,7 @@ def _FillWithArrayTH1(self, *args):
 for klass in _th1_derived_classes_to_pythonize:
     pythonization(klass)(inject_constructor_releasing_ownership)
 
-    from ROOT._pythonization._uhi.main import _add_plotting_features
+    from ROOT._pythonization._uhi import _add_plotting_features
 
     # Add UHI plotting features
     pythonization(klass)(_add_plotting_features)
@@ -252,20 +252,24 @@ def _enable_numpy_fill(klass):
         klass._Fill = klass.Fill
         klass.Fill = _FillWithArrayTH1
 
+    # Add serialization features
+    from ROOT._pythonization._uhi import _add_serialization_features
+    pythonization(klass)(_add_serialization_features)
+
 
 @pythonization("TH1")
 def pythonize_th1(klass):
     # Parameters:
     # klass: class to be pythonized
-    from ROOT._pythonization._uhi.main import _add_indexing_features
+    from ROOT._pythonization._uhi import _add_indexing_features
 
     # Support hist *= scalar
     klass.__imul__ = _imul
 
     klass._Original_SetDirectory = klass.SetDirectory
     klass.SetDirectory = _SetDirectory_SetOwnership
 
-    # Add UHI indexing features
+    # Add UHI indexing and serialization features
     _add_indexing_features(klass)
 
     inject_clone_releasing_ownership(klass)

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_th2.py

@@ -78,7 +78,7 @@ def _FillWithArrayTH2(self, *args):
 for klass in _th2_derived_classes_to_pythonize:
     pythonization(klass)(inject_constructor_releasing_ownership)
 
-    from ROOT._pythonization._uhi.main import _add_plotting_features
+    from ROOT._pythonization._uhi import _add_plotting_features
 
     # Add UHI plotting features
     pythonization(klass)(_add_plotting_features)

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_th3.py

@@ -27,7 +27,7 @@
 for klass in _th3_derived_classes_to_pythonize:
     pythonization(klass)(inject_constructor_releasing_ownership)
 
-    from ROOT._pythonization._uhi.main import _add_plotting_features
+    from ROOT._pythonization._uhi import _add_plotting_features
 
     # Add UHI plotting features
     pythonization(klass)(_add_plotting_features)

…/python/ROOT/_pythonization/_uhi/main.py …hon/ROOT/_pythonization/_uhi/__init__.pybindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi/main.py renamed to bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi/__init__.py bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi/main.py renamed to bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi/__init__.py

@@ -64,3 +64,28 @@ def _add_plotting_features(klass: Any) -> None:
     klass.counts = _counts
     klass.axes = property(_axes)
     klass.values = values_func_dict.get(klass.__name__, _values_default)
+
+
+"""
+Implementation of the serialization component of the UHI
+"""
+
+def _TH1_Constructor(self, *args, **kwargs):
+    """
+    If UHI IR is detected, use the UHI deserialization constructor
+    else forward to the original
+    """
+    if len(args) == 1 and isinstance(args[0], dict):
+        from .serialization import _from_uhi_
+
+        _from_uhi_(self, args[0])
+    else:
+        self._original_init_(*args, **kwargs)
+
+def _add_serialization_features(klass: Any) -> None:
+    from .serialization import _to_uhi_
+
+    klass._to_uhi_ = _to_uhi_
+
+    klass._original_init_ = klass.__init__
+    klass.__init__ = _TH1_Constructor

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi/plotting.py

@@ -38,6 +38,14 @@ def circular(self) -> bool:
     def discrete(self) -> bool:
         return self._discrete
 
+    @property
+    def underflow(self) -> bool:
+        return True
+
+    @property
+    def overflow(self) -> bool:
+        return True
+
 
 class PlottableAxisBase(ABC):
     def __init__(self, tAxis: Any) -> None:
@@ -104,7 +112,7 @@ def _hasWeights(hist: Any) -> bool:
 def _axes(self) -> Tuple[Union[PlottableAxisContinuous, PlottableAxisDiscrete], ...]:
     return tuple(PlottableAxisFactory.create(_get_axis(self, i)) for i in range(self.GetDimension()))
 
-
+# TODO this is not correct?
 def _kind(self) -> Kind:
     # TProfile -> MEAN, everything else -> COUNT
     if self.__class__.__name__.startswith("TProfile"):
@@ -182,6 +190,8 @@ def _counts(self, flow=False) -> np.typing.NDArray[Any]:  # noqa: F821
         where=sum_of_weights_squared != 0,
     )
 
+def _get_sum_of_weights(self) -> np.typing.NDArray[Any]:  # noqa: F821
+    return self.values()
 
 def _get_sum_of_weights(self) -> np.typing.NDArray[Any]:  # noqa: F821
     return self.values()

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi/serialization.py

@@ -0,0 +1,129 @@
+# Author: Silia Taider CERN  10/2025
+
+################################################################################
+# Copyright (C) 1995-2025, Rene Brun and Fons Rademakers.                      #
+# All rights reserved.                                                         #
+#                                                                              #
+# For the licensing terms see $ROOTSYS/LICENSE.                                #
+# For the list of contributors see $ROOTSYS/README/CREDITS.                    #
+################################################################################
+from __future__ import annotations
+
+from typing import Any
+
+import ROOT
+
+from .plotting import PlottableAxisBase, _get_sum_of_weights, _get_sum_of_weights_squared, _hasWeights
+from .tags import _get_axis
+
+"""
+Implementation of the serialization component of the UHI
+"""
+
+
+def _axis_to_dict(root_axis: ROOT.TAxis, uhi_axis: PlottableAxisBase) -> dict[str, Any]:
+    """
+    Return a dictionary representation of the given ROOT axis.
+    """
+    return {
+        "type": "regular",
+        "lower": root_axis.GetBinLowEdge(root_axis.GetFirst()),
+        "upper": root_axis.GetBinUpEdge(root_axis.GetLast()),
+        "bins": root_axis.GetNbins(),
+        "underflow": uhi_axis.traits.underflow,
+        "overflow": uhi_axis.traits.overflow,
+        "circular": uhi_axis.traits.circular,
+    }
+
+
+def _axis_from_dict(axis_dict: dict[str, Any]) -> list[Any]:
+    """
+    Return the arguments needed to construct the corresponding ROOT histogram axis.
+    For now only supports regular axes.
+    """
+
+    axis_type = axis_dict["type"]
+
+    if axis_type == "regular":
+        nbins = axis_dict["bins"]
+        lower = axis_dict["lower"]
+        upper = axis_dict["upper"]
+        return [nbins, lower, upper]
+
+    raise ValueError(f"Unsupported axis type for conversion to ROOT: {axis_type}")
+
+
+def _storage_to_dict(hist: Any) -> dict[str, Any]:
+    """
+    Logic:
+    - If histogram is a profile (TProfile*) --> Kind="MEAN":
+        - if histogram has Sumw2: type is weighted_mean_storage (if _hasWeights(hist))
+        - else: storage type is mean_storage
+    - Else (TH1*/TH2*/TH3*) --> Kind="COUNT":
+        - if histogram has Sumw2: type is weighted_storage
+        - else if histogram is TH*I: type is int_storage
+        - else: type is double_storage
+    """
+    storage_dict = {
+        "values": hist.values(),
+    }
+
+    if hist.kind == "MEAN":
+        storage_dict["variances"] = hist.variances()
+
+        if _hasWeights(hist):
+            storage_dict["type"] = "weighted_mean"
+            storage_dict["sum_of_weights"] = _get_sum_of_weights(hist)
+            storage_dict["sum_of_weights_squared"] = _get_sum_of_weights_squared(hist)
+        else:
+            storage_dict["type"] = "mean"
+            storage_dict["counts"] = hist.counts()
+
+    else:  # COUNT
+        if _hasWeights(hist):
+            storage_dict["type"] = "weighted"
+            storage_dict["variances"] = hist.variances()
+        else:
+            if hist.ClassName().endswith("I"):
+                storage_dict["type"] = "int"
+            else:
+                storage_dict["type"] = "double"
+
+    return storage_dict
+
+
+def _set_histogram_storage_from_dict(hist: Any, storage_dict: dict[str, Any]) -> None:
+    """
+    Set the histogram storage (values and statistics) from the given storage dictionary.
+    """
+    hist_values = storage_dict["values"]
+    hist[...] = hist_values
+
+    stype = storage_dict.get("type")
+    if stype in ["weighted_mean", "mean", "weighted"]:
+        hist.variances()[:] = storage_dict["variances"]
+
+
+def _to_uhi_(self) -> dict[str, Any]:
+    return {
+        "uhi_schema": 1,
+        "writer_info": {"ROOT": {"version": ROOT.__version__, "class": self.ClassName(), "name": self.GetName()}},
+        "axes": [_axis_to_dict(_get_axis(self, i), self.axes[i]) for i in range(self.GetDimension())],
+        "storage": _storage_to_dict(self),
+    }
+
+
+def _from_uhi_(self, uhi_dict: dict[str, Any]) -> ROOT.TH1:
+    # rebuild axes
+    axes = uhi_dict["axes"]
+    axes_specs = [_axis_from_dict(axis_dict) for axis_dict in axes]
+
+    # construct the histogram
+    ctor_args = ["h_uhi", "h_uhi"]
+    for axis_spec in axes_specs:
+        ctor_args.extend(axis_spec)
+
+    self._original_init_(*ctor_args)
+
+    # set storage
+    _set_histogram_storage_from_dict(self, uhi_dict["storage"])

tutorials/hist/repro_uhi_serialization.py

@@ -0,0 +1,53 @@
+import json
+
+import numpy as np
+import ROOT
+import uhi.io.json
+
+import hist
+
+## to_uhi
+
+h = ROOT.TH1D("h", "h", 10, -5, 5)
+h[...] = range(10)
+print("\nh =", h)
+print("values=", h.values())
+
+ob = json.dumps(h, default=uhi.io.json.default)
+ir = json.loads(ob, object_hook=uhi.io.json.object_hook)
+
+h_loaded = hist.Hist(ir)
+print("\nh_loaded =")
+print("values=", h_loaded.values())
+
+
+## from_uhi
+
+h2 = hist.Hist(hist.axis.Regular(10, -5, 5, name="x"))
+h2[...] = range(10)
+print("\nh2 =")
+print("values=", h2.values())
+
+ob2 = json.dumps(h2, default=uhi.io.json.default)
+ir2 = json.loads(ob2, object_hook=uhi.io.json.object_hook)
+
+h2_loaded = ROOT.TH1D(ir2)
+print("\nh2_loaded =", h2_loaded)
+print("values=", h2_loaded.values())
+
+
+## 2D histogram
+h2d = hist.Hist(
+    hist.axis.Regular(10, -5, 5, name="x"),
+    hist.axis.Regular(10, -5, 5, name="y"),
+)
+h2d[...] = np.arange(100).reshape(10, 10)
+print("\nh2d =")
+print("values=", h2d.values())
+
+ob3 = json.dumps(h2d, default=uhi.io.json.default)
+ir3 = json.loads(ob3, object_hook=uhi.io.json.object_hook)
+
+h2d_loaded = hist.Hist(ir3)
+print("\nh2d_loaded =")
+print("values=", h2d_loaded.values())