fix: NumPy syntax now handles dd

Author: henryiii

src/boost_histogram/_internal/hist.py

@@ -454,18 +454,36 @@ def _compute_commonindex(self, index):
 
         return indexes
 
-    def to_numpy(self, flow=False):
+    @inject_signature("self, flow=False, *, dd=False")
+    def to_numpy(self, flow=False, **kwargs):
         """
         Convert to a Numpy style tuple of return arrays.
 
+        Parameters
+        ----------
+
+        flow : bool = False
+            Include the flow bins.
+        dd : bool = False
+            Use the histogramdd return syntax, where the edges are in a tuple
+
         Return
         ------
         contents : Array[Any]
             The bin contents
         *edges : Array[float]
             The edges for each dimension
         """
-        return self._hist.to_numpy(flow)
+
+        with KWArgs(kwargs) as kw:
+            dd = kw.optional("dd", False)
+
+        return_tuple = self._hist.to_numpy(flow)
+
+        if dd:
+            return return_tuple[0], return_tuple[1:]
+        else:
+            return return_tuple
 
     @inject_signature("self, *, deep=True")
     def copy(self, **kwargs):

src/boost_histogram/numpy.py

@@ -78,9 +78,9 @@ def histogramdd(
     if density:
         areas = _reduce(_mul, hist.axes.widths)
         density = hist.view() / hist.sum() / areas
-        return (density,) + hist.to_numpy()[1:]
+        return (density, hist.to_numpy()[1:])
 
-    return hist if bh_cls is not None else hist.to_numpy()
+    return hist if bh_cls is not None else hist.to_numpy(dd=True)
 
 
 @_inject_signature(
@@ -89,7 +89,13 @@ def histogramdd(
 def histogram2d(
     x, y, bins=10, range=None, normed=None, weights=None, density=None, **kwargs
 ):
-    return histogramdd((x, y), bins, range, normed, weights, density, **kwargs)
+    result = histogramdd((x, y), bins, range, normed, weights, density, **kwargs)
+
+    if isinstance(result, tuple):
+        data, (edgesx, edgesy) = result
+        return data, edgesx, edgesy
+    else:
+        return result
 
 
 @_inject_signature(
@@ -110,7 +116,13 @@ def histogram(
                 "Upgrade numpy to 1.13+ to use string arguments to boost-histogram's histogram function"
             )
         bins = np.histogram_bin_edges(a, bins, range, weights)
-    return histogramdd((a,), (bins,), (range,), normed, weights, density, **kwargs)
+
+    result = histogramdd((a,), (bins,), (range,), normed, weights, density, **kwargs)
+    if isinstance(result, tuple):
+        data, (edges,) = result
+        return data, edges
+    else:
+        return result
 
 
 # Process docstrings

tests/test_internal_histogram.py

@@ -117,6 +117,25 @@ def test_growing_histogram():
     assert hist.size == 17
 
 
+def test_numpy_dd():
+    h = bh.Histogram(
+        bh.axis.Regular(10, 0, 1), bh.axis.Regular(5, 0, 1), storage=bh.storage.Int64()
+    )
+
+    for i in range(10):
+        for j in range(5):
+            x, y = h.axes[0].centers[i], h.axes[1].centers[j]
+            v = i + j * 10 + 1
+            h.fill([x] * v, [y] * v)
+
+    h2, x2, y2 = h.to_numpy()
+    h1, (x1, y1) = h.to_numpy(dd=True)
+
+    assert_array_equal(h1, h2)
+    assert_array_equal(x1, x2)
+    assert_array_equal(y1, y2)
+
+
 def test_numpy_flow():
     h = bh.Histogram(
         bh.axis.Regular(10, 0, 1), bh.axis.Regular(5, 0, 1), storage=bh.storage.Int64()

tests/test_numpy_interface.py

@@ -116,4 +116,44 @@ def test_histogram2d_object():
     np.testing.assert_array_equal(h1, h2)
 
     with pytest.raises(KeyError):
-        bh_h2 = bh.numpy.histogram2d(x, y, density=True, histogram=bh.Histogram)
+        bh.numpy.histogram2d(x, y, density=True, histogram=bh.Histogram)
+
+
+def test_histogramdd():
+    x = np.array([0.3, 0.3, 0.1, 0.8, 0.34, 0.03, 0.32, 0.65])
+    y = np.array([0.4, 0.5, 0.22, 0.65, 0.32, 0.01, 0.23, 1.98])
+    z = np.array([0.5, 0.7, 0.0, 0.65, 0.72, 0.01, 0.3, 1.4])
+
+    h1, (e1x, e1y, e1z) = np.histogramdd([x, y, z])
+    h2, (e2x, e2y, e2z) = bh.numpy.histogramdd([x, y, z])
+
+    np.testing.assert_array_almost_equal(e1x, e2x)
+    np.testing.assert_array_almost_equal(e1y, e2y)
+    np.testing.assert_array_almost_equal(e1z, e2z)
+    np.testing.assert_array_equal(h1, h2)
+
+    h1, (e1x, e1y, e1z) = np.histogramdd([x, y, z], density=True)
+    h2, (e2x, e2y, e2z) = bh.numpy.histogramdd([x, y, z], density=True)
+
+    np.testing.assert_array_almost_equal(e1x, e2x)
+    np.testing.assert_array_almost_equal(e1y, e2y)
+    np.testing.assert_array_almost_equal(e1z, e2z)
+    np.testing.assert_array_almost_equal(h1, h2)
+
+
+def test_histogramdd_object():
+    x = np.array([0.3, 0.3, 0.1, 0.8, 0.34, 0.03, 0.32, 0.65])
+    y = np.array([0.4, 0.5, 0.22, 0.65, 0.32, 0.01, 0.23, 1.98])
+    z = np.array([0.5, 0.7, 0.0, 0.65, 0.72, 0.01, 0.3, 1.4])
+
+    h1, (e1x, e1y, e1z) = np.histogramdd([x, y, z])
+    bh_h2 = bh.numpy.histogramdd([x, y, z], histogram=bh.Histogram)
+    h2, (e2x, e2y, e2z) = bh_h2.to_numpy(dd=True)
+
+    np.testing.assert_array_almost_equal(e1x, e2x)
+    np.testing.assert_array_almost_equal(e1y, e2y)
+    np.testing.assert_array_almost_equal(e1z, e2z)
+    np.testing.assert_array_equal(h1, h2)
+
+    with pytest.raises(KeyError):
+        bh.numpy.histogramdd([x, y, z], density=True, histogram=bh.Histogram)