Up edges creation

Author: Descanonge

src/xarray_histogram/core.py

@@ -12,7 +12,6 @@
 from collections import abc
 from copy import copy
 from functools import partial, reduce
-from typing import cast
 
 import boost_histogram as bh
 import numpy as np
@@ -44,9 +43,6 @@
 RangeType = tuple[float | None, float | None]
 
 
-# TODO add flow
-
-
 def histogram(
     x: xr.DataArray,
     /,
@@ -535,30 +531,51 @@ def get_coord(name: str, ax: bh.axis.Axis, dtype: np.dtype, flow: bool) -> xr.Da
     attrs = dict(bin_type=type(ax).__name__, underflow=underflow, overflow=overflow)
 
     if isinstance(ax, bh.axis.Integer):
-        lefts = ax.edges[:-1].astype(dtype)
+        if dtype.kind not in "uib":
+            raise TypeError(f"Cannot use Integer axis for dtype {dtype}")
+
+        lefts = ax.edges[:-1].astype("int")
+
+        # deal with bool variables
+        if dtype.kind == "b" and not (underflow or overflow):
+            lefts = lefts.astype("bool")
+
         # use min/max possible encoded values to indicate flow
+        bins_dtype = lefts.dtype
         if underflow:
-            if dtype.kind == "u":
-                dtype = np.dtype(f"i{min(dtype.itemsize * 2, 8)}")
-            vmin = np.iinfo(dtype).min
-            lefts = np.concatenate(([vmin], lefts), dtype=dtype)
+            vmin = np.iinfo(bins_dtype).min
+            lefts = np.concatenate(([vmin], lefts), dtype=bins_dtype)
         if overflow:
-            vmax = np.iinfo(dtype).max
-            lefts = np.concatenate((lefts, [vmax]), dtype=dtype)
+            vmax = np.iinfo(bins_dtype).max
+            lefts = np.concatenate((lefts, [vmax]), dtype=bins_dtype)
 
     elif isinstance(ax, bh.axis.IntCategory):
-        lefts = np.asarray([cast(int, ax.bin(i)) for i in range(ax.size)])
-        lefts = lefts.astype(dtype)
+        if dtype.kind not in "uib":
+            raise TypeError(f"Cannot use Integer axis for dtype {dtype}")
+
+        lefts = np.asarray([ax.bin(i) for i in range(ax.size)], dtype="int")
+
+        # deal with bool variables
+        if dtype.kind == "b" and not (underflow or overflow):
+            lefts = lefts.astype("bool")
+
+        bins_dtype = lefts.dtype
         if overflow:
-            lefts = np.concatenate((lefts, [np.iinfo(dtype).max]), dtype=dtype)
+            lefts = np.concatenate(
+                (lefts, [np.iinfo(bins_dtype).max]), dtype=bins_dtype
+            )
 
     elif isinstance(ax, bh.axis.StrCategory):
+        if dtype.kind not in "SU":
+            raise TypeError(f"Cannot use StrCategory axis for dtype {dtype}")
         lefts = np.asarray([ax.bin(i) for i in range(ax.size)])
         if overflow:
             lefts = np.concatenate((lefts, ["_flow_bin"]))
 
     else:
-        lefts = ax.edges[:-1].astype(dtype, casting="safe")
+        if dtype.kind not in "biuf":
+            raise TypeError(f"Cannot use {type(ax).__name__} axis for dtype {dtype}")
+        lefts = ax.edges[:-1]
         attrs["right_edge"] = ax.edges[-1]
         if underflow:
             lefts = np.concatenate(([-np.inf], lefts))
@@ -575,13 +592,19 @@ def _bins_name(variable: str) -> str:
 def get_edges(coord: xr.DataArray) -> xr.DataArray:
     """Return edges positions."""
     name = coord.name
-    if coord.attrs["bin_type"] in ["Integer", "IntCategory", "StrCategory"]:
-        return xr.DataArray(coord.values, dims=[name], name=name)
+    bin_type = coord.attrs["bin_type"]
+    if bin_type in ["IntCategory", "StrCategory"]:
+        raise TypeError(f"Edges not available for {bin_type} bins type.")
 
-    # insert right_edge
+    overflow = coord.attrs.get("overflow", False)
+
+    if bin_type == "Integer":
+        right_edge = coord[-2 if overflow else -1] + 1
+    else:
+        right_edge = coord.attrs["right_edge"]
     values = coord.values
-    insert = values.size - 1 if coord.attrs.get("overflow", False) else values.size
-    values = np.insert(values, insert, [coord.attrs["right_edge"]])
+    insert = values.size - 1 if overflow else values.size
+    values = np.insert(values, insert, [right_edge])
 
     return xr.DataArray(values, dims=[name], name=name)
 

tests/test_accessor.py

@@ -29,6 +29,18 @@ def get_blank_histogram(n_var: int = 1) -> xr.DataArray:
     return h
 
 
+def get_hist(*axes: bh.axis.Axis, flow=True) -> xr.DataArray:
+    data = []
+    for ax in axes:
+        x = get_array([2])
+        if isinstance(ax, bh.axis.Integer | bh.axis.IntCategory):
+            x = x.astype("int")
+        if isinstance(ax, bh.axis.StrCategory):
+            x = x.astype("U")
+        data.append(x)
+    return xh.histogramdd(*data, bins=axes, flow=flow)
+
+
 class TestAccessibility:
     @pytest.mark.parametrize("x", [get_array([20]), get_array([20, 5])], ids=id_x)
     @bool_param("density")
@@ -65,8 +77,42 @@ def test_variable_argument() -> None:
 
 
 class TestEdges:
+    def test_edges(self) -> None:
+        # Regular
+        h = get_hist(bh.axis.Regular(3, 0.0, 0.3), flow=False)
+        assert_allclose(h.hist.edges(), [0.0, 0.1, 0.2, 0.3])
+        h = get_hist(bh.axis.Regular(3, 0.0, 0.3, underflow=False))
+        assert_allclose(h.hist.edges(), [0.0, 0.1, 0.2, 0.3, np.inf])
+        h = get_hist(bh.axis.Regular(3, 0.0, 0.3))
+        assert_allclose(h.hist.edges(), [-np.inf, 0.0, 0.1, 0.2, 0.3, np.inf])
+
+        # Integer
+        h = get_hist(bh.axis.Integer(0, 3), flow=False)
+        assert_allclose(h.hist.edges(), [0, 1, 2, 3])
+        h = get_hist(bh.axis.Integer(0, 3, underflow=False))
+        dtype = h.var1_bins.dtype
+        vmin = np.iinfo(dtype).min
+        vmax = np.iinfo(dtype).max
+        assert_allclose(h.hist.edges(), [0, 1, 2, 3, vmax])
+        h = get_hist(bh.axis.Integer(0, 3))
+        assert_allclose(h.hist.edges(), [vmin, 0, 1, 2, 3, vmax])
+
+        # Variable
+        h = get_hist(bh.axis.Variable([0, 1, 3, 10]), flow=False)
+        assert_allclose(h.hist.edges(), [0, 1, 3, 10])
+        h = get_hist(bh.axis.Variable([0, 1, 3, 10], underflow=False))
+        assert_allclose(h.hist.edges(), [0, 1, 3, 10, np.inf])
+        h = get_hist(bh.axis.Variable([0, 1, 3, 10]))
+        assert_allclose(h.hist.edges(), [-np.inf, 0, 1, 3, 10, np.inf])
+
+        # Not supported
+        for ax in [bh.axis.IntCategory([0, 1, 2]), bh.axis.StrCategory(["a", "b"])]:
+            h = get_hist(ax)
+            with pytest.raises(TypeError):
+                h.hist.edges()
+
     def test_infer_right_edge(self) -> None:
-        h = get_blank_histogram()
+        h = get_hist(bh.axis.Regular(10, 0.0, 1.0), flow=False)
         # this reset right edge
         h = h.assign_coords(var1_bins=np.arange(0, 10))
 
@@ -78,10 +124,6 @@ def test_infer_right_edge(self) -> None:
         with pytest.raises(ValueError):
             _ = h_wrong.hist
 
-    def test_basic(self) -> None:
-        h = get_blank_histogram()
-        assert_allclose(h.hist.edges("var1"), np.arange(0, 11))
-
     def test_centers(self):
         h = get_blank_histogram()
         assert_allclose(h.hist.centers(), np.arange(0, 10) + 0.5)