Merge pull request #290 from ndawe/hist_bin_edges

hist2array: return_edges returns the bin edges along each axis

Author: ndawe

root_numpy/_hist.py

@@ -129,7 +129,7 @@ def fill_profile(profile, array, weights=None, return_indices=False):
         "ROOT.TProfile, ROOT.TProfile2D, or ROOT.TProfile3D")
 
 
-def hist2array(hist, include_overflow=False, copy=True):
+def hist2array(hist, include_overflow=False, copy=True, return_edges=False):
     """Convert a ROOT histogram into a NumPy array
 
     Parameters
@@ -143,11 +143,17 @@ def hist2array(hist, include_overflow=False, copy=True):
         If True (the default) then copy the underlying array, otherwise the
         NumPy array will view (and not own) the same memory as the ROOT
         histogram's array.
+    return_edges : bool, optional (default=False)
+        If True, also return the bin edges along each axis.
 
     Returns
     -------
     array : numpy array
         A NumPy array containing the histogram bin values
+    edges : list of numpy arrays
+        A list of numpy arrays where each array contains the bin edges along
+        the corresponding axis of ``hist``. Overflow and underflow bins are not
+        included.
 
     Raises
     ------
@@ -220,6 +226,25 @@ def hist2array(hist, include_overflow=False, copy=True):
             array = _librootnumpy.thn2array(ROOT.AsCObject(hist),
                                             shape, dtype)
 
+    if return_edges:
+        if simple_hist:
+            ndims = hist.GetDimension()
+            axis_getters = ['GetXaxis', 'GetYaxis', 'GetZaxis'][:ndims]
+        else:
+            ndims = hist.GetNdimensions()
+            axis_getters = ['GetAxis'] * ndims
+
+        edges = []
+        for idim, axis_getter in zip(range(ndims), axis_getters):
+            # GetXaxis expects 0 parameters while we need the axis in GetAxis
+            ax = getattr(hist, axis_getter)(*(() if simple_hist else (idim,)))
+            # `edges` is Nbins + 1 in order to have the last bin's upper edge as well
+            edges.append(np.empty(ax.GetNbins() + 1, dtype=np.double))
+            # load the lower edges into `edges`
+            ax.GetLowEdge(edges[-1])
+            # Get the upper edge of the last bin
+            edges[-1][-1] = ax.GetBinUpEdge(ax.GetNbins())
+
     if not include_overflow:
         # Remove overflow and underflow bins
         array = array[tuple([slice(1, -1) for idim in range(array.ndim)])]
@@ -228,7 +253,10 @@ def hist2array(hist, include_overflow=False, copy=True):
         # Preserve x, y, z -> axis 0, 1, 2 order
         array = np.transpose(array)
         if copy:
-            return np.copy(array)
+            array = np.copy(array)
+
+    if return_edges:
+        return array, edges
     return array
 
 

root_numpy/tests.py

@@ -1103,15 +1103,18 @@ def make_histogram(hist_type, shape, fill=True):
     hist_cls = getattr(ROOT, 'TH{0}{1}'.format(ndim, hist_type))
     if ndim == 1:
         hist = hist_cls(hist_cls.__name__, '',
-                        shape[0], 0, 1)
+                        shape[0], 0, shape[0])
         func = ROOT.TF1('func', 'x')
     elif ndim == 2:
         hist = hist_cls(hist_cls.__name__, '',
-                        shape[1], 0, 1, shape[0], 0, 1)
+                        shape[1], 0, shape[1],
+                        shape[0], 0, shape[0])
         func = ROOT.TF2('func', 'x*y')
     elif ndim == 3:
         hist = hist_cls(hist_cls.__name__, '',
-                        shape[2], 0, 1, shape[1], 0, 1, shape[0], 0, 1)
+                        shape[2], 0, shape[2],
+                        shape[1], 0, shape[1],
+                        shape[0], 0, shape[0])
         func = ROOT.TF3('func', 'x*y*z')
     else:
         raise ValueError("ndim must be 1, 2, or 3")  # pragma: no cover
@@ -1130,16 +1133,17 @@ def check_hist2array(hist, include_overflow, copy):
         else:
             assert_equal(array.shape[iaxis],
                          getattr(hist, 'GetNbins{0}'.format(axis))())
-    # non-zero elements
-    assert_true(np.any(array))
+    hist_sum = hist.Integral()
+    assert_true(hist_sum > 0)
+    assert_equal(hist_sum, np.sum(array))
 
 
 def check_hist2array_THn(hist):
     hist_thn = ROOT.THn.CreateHn("", "", hist)
     array = rnp.hist2array(hist)
     array_thn = rnp.hist2array(hist_thn)
     # non-zero elements
-    assert_true(np.any(array))
+    assert_true(np.any(array_thn))
     # arrays should be identical
     assert_array_equal(array, array_thn)
 
@@ -1149,7 +1153,7 @@ def check_hist2array_THnSparse(hist):
     array = rnp.hist2array(hist)
     array_thnsparse = rnp.hist2array(hist_thnsparse)
     # non-zero elements
-    assert_true(np.any(array))
+    assert_true(np.any(array_thnsparse))
     # arrays should be identical
     assert_array_equal(array, array_thnsparse)
 
@@ -1164,6 +1168,20 @@ def test_hist2array():
             yield check_hist2array, hist, True, False
             yield check_hist2array, hist, True, True
             yield check_hist2array_THn, hist
+            # check that the memory was copied
+            arr = rnp.hist2array(hist, copy=True)
+            hist_sum = hist.Integral()
+            assert_true(hist_sum > 0)
+            hist.Reset()
+            assert_equal(np.sum(arr), hist_sum)
+            # check that the memory is shared
+            hist = make_histogram(hist_type, shape=(5,) * ndim)
+            arr = rnp.hist2array(hist, copy=False)
+            hist_sum = hist.Integral()
+            assert_true(hist_sum > 0)
+            assert_true(np.sum(arr) == hist_sum)
+            hist.Reset()
+            assert_true(np.sum(arr) == 0)
 
 
 def test_hist2array_THn():
@@ -1182,6 +1200,24 @@ def test_hist2array_THnSparse():
             yield check_hist2array_THnSparse, hist
 
 
+def check_hist2array_edges(hist, ndim, bins):
+    _, edges = rnp.hist2array(hist, return_edges=True)
+    assert_equal(len(edges), ndim)
+    for axis_edges in edges:
+        assert_array_equal(axis_edges, np.arange(bins + 1, dtype=np.double))
+
+
+def test_hist2array_edges():
+    for ndim in (1, 2, 3):
+        for bins in (1, 2, 5):
+            hist = make_histogram('D', shape=(bins,) * ndim)
+            yield check_hist2array_edges, hist, ndim, bins
+            hist = ROOT.THn.CreateHn("", "", make_histogram('D', shape=(bins,) * ndim))
+            yield check_hist2array_edges, hist, ndim, bins
+            hist = ROOT.THnSparse.CreateSparse("", "", make_histogram('D', shape=(bins,) * ndim))
+            yield check_hist2array_edges, hist, ndim, bins
+
+
 def check_array2hist(hist):
     shape = np.array([hist.GetNbinsX(), hist.GetNbinsY(), hist.GetNbinsZ()])
     shape = shape[:hist.GetDimension()]