fix: support transform axes, more tests

Signed-off-by: Henry Schreiner <[email protected]>

Author: henryiii

pyproject.toml

@@ -227,6 +227,8 @@ messages_control.disable = [
   "too-many-statements",
   "too-many-positional-arguments",
   "wrong-import-position",
+  "unused-argument", # Covered by ruff
+  "unsubscriptable-object", # Wrongly triggered
 ]
 
 [tool.ruff.lint]

src/boost_histogram/serialization/_axis.py

@@ -22,15 +22,26 @@ def _axis_to_dict(ax: Any, /) -> dict[str, Any]:
 @_axis_to_dict.register(axis.Integer)
 def _(ax: axis.Regular | axis.Integer, /) -> dict[str, Any]:
     """Convert a Regular axis to a dictionary."""
-    data = {
-        "type": "regular",
-        "lower": ax.edges[0],
-        "upper": ax.edges[-1],
-        "bins": ax.size,
-        "underflow": ax.traits.underflow,
-        "overflow": ax.traits.overflow,
-        "circular": ax.traits.circular,
-    }
+
+    # Special handling if the axis has a transform
+    if isinstance(ax, axis.Regular) and ax.transform is not None:
+        data = {
+            "type": "variable",
+            "edges": ax.edges,
+            "underflow": ax.traits.underflow,
+            "overflow": ax.traits.overflow,
+            "circular": ax.traits.circular,
+        }
+    else:
+        data = {
+            "type": "regular",
+            "lower": ax.edges[0],
+            "upper": ax.edges[-1],
+            "bins": ax.size,
+            "underflow": ax.traits.underflow,
+            "overflow": ax.traits.overflow,
+            "circular": ax.traits.circular,
+        }
     if ax.metadata is not None:
         data["metadata"] = ax.metadata
 

tests/test_hdf5.py

@@ -15,19 +15,25 @@
 @pytest.mark.parametrize(
     ("storage_type", "fill_args", "fill_kwargs"),
     [
-        pytest.param(bh.storage.Weight(), [0.3, 0.3, 0.4, 1.2], {}, id="weight"),
+        pytest.param(bh.storage.Double(), [0.3, 0.3, 0.4, 1.2], {}, id="double"),
+        pytest.param(bh.storage.Int64(), [0.3, 0.3, 0.4, 1.2], {}, id="int64"),
         pytest.param(
-            bh.storage.WeightedMean(),
-            [0.3, 0.3, 0.4, 1.2, 1.6],
-            {"sample": [1, 2, 3, 4, 4], "weight": [1, 1, 1, 1, 2]},
-            id="weighted_mean",
+            bh.storage.AtomicInt64(), [0.3, 0.3, 0.4, 1.2], {}, id="atomicint"
         ),
+        pytest.param(bh.storage.Unlimited(), [0.3, 0.3, 0.4, 1.2], {}, id="unlimited"),
+        pytest.param(bh.storage.Weight(), [0.3, 0.3, 0.4, 1.2], {}, id="weight"),
         pytest.param(
             bh.storage.Mean(),
             [0.3, 0.3, 0.4, 1.2, 1.6],
             {"sample": [1, 2, 3, 4, 4]},
             id="mean",
         ),
+        pytest.param(
+            bh.storage.WeightedMean(),
+            [0.3, 0.3, 0.4, 1.2, 1.6],
+            {"sample": [1, 2, 3, 4, 4], "weight": [1, 1, 1, 1, 2]},
+            id="weighted_mean",
+        ),
     ],
 )
 def test_hdf5_storage(
@@ -51,7 +57,14 @@ def test_hdf5_storage(
 
     # checking types of the reconstructed axes
     assert type(actual_hist.axes[0]) is type(re_constructed_hist.axes[0])
-    assert actual_hist.storage_type == re_constructed_hist.storage_type
+
+    if isinstance(storage_type, bh.storage.Unlimited):
+        actual_hist_storage = bh.storage.Double()
+    elif isinstance(storage_type, bh.storage.AtomicInt64):
+        actual_hist_storage = bh.storage.Int64()
+    else:
+        actual_hist_storage = storage_type
+    assert isinstance(actual_hist_storage, re_constructed_hist.storage_type)
 
     # checking values of the essential inputs of the axes
     assert actual_hist.axes[0].traits == re_constructed_hist.axes[0].traits
@@ -76,3 +89,43 @@ def test_hdf5_storage(
         assert actual_hist.counts() == pytest.approx(
             re_constructed_hist.counts(), abs=1e-4, rel=1e-9
         )
+
+
+def test_hdf5_2d(tmp_path: Path) -> None:
+    h = bh.Histogram(bh.axis.Integer(0, 4), bh.axis.StrCategory(["a", "b", "c"]))
+    h.fill([0, 1, 1, 1], ["a", "b", "b", "c"])
+
+    filepath = tmp_path / "hist.h5"
+    with h5py.File(filepath, "x") as f:
+        grp = f.create_group("test_hist")
+        s.write_hdf5_schema(grp, h)
+
+    with h5py.File(filepath) as f:
+        re_constructed_hist = s.read_hdf5_schema(f["test_hist"])
+
+    actual_hist = h.copy()
+
+    assert isinstance(re_constructed_hist.axes[0], bh.axis.Regular)
+    assert type(actual_hist.axes[1]) is type(re_constructed_hist.axes[1])
+
+    assert (
+        actual_hist.axes[0].traits.underflow
+        == re_constructed_hist.axes[0].traits.underflow
+    )
+    assert (
+        actual_hist.axes[0].traits.overflow
+        == re_constructed_hist.axes[0].traits.overflow
+    )
+    assert (
+        actual_hist.axes[0].traits.circular
+        == re_constructed_hist.axes[0].traits.circular
+    )
+    assert actual_hist.axes[1].traits == re_constructed_hist.axes[1].traits
+
+    assert np.asarray(actual_hist.axes[0].edges) == pytest.approx(
+        np.asarray(re_constructed_hist.axes[0].edges)
+    )
+    assert list(actual_hist.axes[1]) == list(re_constructed_hist.axes[1])
+
+    assert h.values() == pytest.approx(re_constructed_hist.values())
+    assert h.values(flow=True) == pytest.approx(re_constructed_hist.values(flow=True))

tests/test_serialization_generic.py

@@ -99,6 +99,26 @@ def test_weighted_mean_to_dict() -> None:
     assert data["storage"]["data"]["variances"] == pytest.approx(np.zeros(4))
 
 
+def test_transform_log_axis_to_dict() -> None:
+    h = bh.Histogram(bh.axis.Regular(10, 1, 10, transform=bh.axis.transform.log))
+    data = generic.to_dict(h)
+
+    assert data["axes"][0]["type"] == "variable"
+    assert data["axes"][0]["edges"] == pytest.approx(
+        np.exp(np.linspace(0, np.log(10), 11))
+    )
+
+
+def test_transform_sqrt_axis_to_dict() -> None:
+    h = bh.Histogram(bh.axis.Regular(10, 0, 10, transform=bh.axis.transform.sqrt))
+    data = generic.to_dict(h)
+
+    assert data["axes"][0]["type"] == "variable"
+    assert data["axes"][0]["edges"] == pytest.approx(
+        (np.linspace(0, np.sqrt(10), 11)) ** 2
+    )
+
+
 @pytest.mark.parametrize(
     "storage_type",
     [