[Python][UHI] Add serialization tests

Author: siliataider

bindings/pyroot/pythonizations/test/uhi_serialization_json.py

@@ -0,0 +1,145 @@
+"""
+Tests to verify that TH1 and derived histograms conform to the UHI Serialization interfaces.
+"""
+"""
+Tests to verify that ROOT histograms can be serialized to and from
+the UHI JSON IR (round-trip) for 1D, 2D, and 3D histograms.
+"""
+
+import json
+
+import numpy as np
+import pytest
+import ROOT
+import uhi.io.json
+
+
+# Helpers
+# TODO import from plotting
+def _iterate_bins(hist, flow=False):
+    ranges = [
+        range(0 if flow else 1, hist.GetNbinsX() + (2 if flow else 1)),
+        range(0 if flow else 1, hist.GetNbinsY() + (2 if flow else 1)) if hist.GetDimension() > 1 else [0],
+        range(0 if flow else 1, hist.GetNbinsZ() + (2 if flow else 1)) if hist.GetDimension() > 2 else [0],
+    ]
+    yield from ((i, j, k)[: hist.GetDimension()] for i in ranges[0] for j in ranges[1] for k in ranges[2])
+
+
+def _bin_contents(hist, flow=False):
+    return np.array(
+        [hist.GetBinContent(*idx) for idx in _iterate_bins(hist, flow=flow)]
+    )
+
+
+def _roundtrip(hist):
+    """
+    ROOT histogram -> JSON -> UHI IR -> ROOT histogram
+    """
+    json_str = json.dumps(hist, default=uhi.io.json.default)
+    ir = json.loads(json_str, object_hook=uhi.io.json.object_hook)
+    return hist.__class__(ir)
+
+
+# Tests
+class TestTH1Serialization:
+    def test_roundtrip_1d(self):
+        h = ROOT.TH1D("h1", "h1", 10, -5, 5)
+        h[...] = np.arange(10)
+
+        # test constructor based serialization roundtrip
+        h_loaded = _roundtrip(h)
+
+        assert np.array_equal(
+            _bin_contents(h, flow=False),
+            _bin_contents(h_loaded, flow=False),
+        )
+
+        # test classmethod based serialization roundtrip
+        ir = h._to_uhi_()
+        h_loaded_cls = ROOT.TH1D._from_uhi_(ir)
+        assert np.array_equal(
+            _bin_contents(h, flow=False),
+            _bin_contents(h_loaded_cls, flow=False),
+        )
+
+    def test_roundtrip_1d_flow(self):
+        h = ROOT.TH1D("h1f", "h1f", 5, 0, 5)
+        h.Fill(-1)
+        h.Fill(6)
+        h[...] = np.arange(5)
+
+        # test constructor based serialization roundtrip
+        h_loaded = _roundtrip(h)
+
+        assert np.array_equal(
+            _bin_contents(h, flow=True),
+            _bin_contents(h_loaded, flow=True),
+        )
+
+        # test classmethod based serialization roundtrip
+        ir = h._to_uhi_()
+        h_loaded_cls = ROOT.TH1D._from_uhi_(ir)
+        assert np.array_equal(
+            _bin_contents(h, flow=True),
+            _bin_contents(h_loaded_cls, flow=True),
+        )
+
+
+class TestTH2Serialization:
+    def test_roundtrip_2d(self):
+        h = ROOT.TH2D("h2", "h2", 4, 0, 4, 3, 0, 3)
+
+        values = np.arange(12).reshape(4, 3)
+        h[...] = values
+
+        h_loaded = _roundtrip(h)
+
+        assert np.array_equal(
+            _bin_contents(h, flow=False),
+            _bin_contents(h_loaded, flow=False),
+        )
+
+    def test_roundtrip_2d_flow(self):
+        h = ROOT.TH2D("h2f", "h2f", 3, 0, 3, 2, 0, 2)
+        h.Fill(-1, 1)
+        h.Fill(4, 3)
+        h[...] = np.arange(6).reshape(3, 2)
+
+        h_loaded = _roundtrip(h)
+
+        assert np.array_equal(
+            _bin_contents(h, flow=True),
+            _bin_contents(h_loaded, flow=True),
+        )
+
+
+class TestTH3Serialization:
+    def test_roundtrip_3d(self):
+        h = ROOT.TH3D("h3", "h3", 3, 0, 3, 2, 0, 2, 2, 0, 2)
+
+        values = np.arange(12).reshape(3, 2, 2)
+        h[...] = values
+
+        h_loaded = _roundtrip(h)
+
+        assert np.array_equal(
+            _bin_contents(h, flow=False),
+            _bin_contents(h_loaded, flow=False),
+        )
+
+    def test_roundtrip_3d_flow(self):
+        h = ROOT.TH3D("h3f", "h3f", 2, 0, 2, 2, 0, 2, 2, 0, 2)
+        h.Fill(-1, 0, 0)
+        h.Fill(3, 3, 3)
+        h[...] = np.arange(8).reshape(2, 2, 2)
+
+        h_loaded = _roundtrip(h)
+
+        assert np.array_equal(
+            _bin_contents(h, flow=True),
+            _bin_contents(h_loaded, flow=True),
+        )
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main(args=[__file__]))