Merge pull request #301 from cbourjau/array2hist_with_errors

Optionally pass bin errors in array2hist

root_numpy/_hist.py

@@ -260,7 +260,7 @@ def hist2array(hist, include_overflow=False, copy=True, return_edges=False):
     return array
 
 
-def array2hist(array, hist):
+def array2hist(array, hist, errors=None):
     """Convert a NumPy array into a ROOT histogram
 
     Parameters
@@ -270,6 +270,9 @@ def array2hist(array, hist):
         ROOT histogram.
     hist : ROOT TH1, TH2, or TH3
         A ROOT histogram.
+    errors : numpy array
+        A numpy array of errors with matching dimensionality as the
+        bin contents array. If not given, no errors are set
 
     Returns
     -------
@@ -381,6 +384,14 @@ def array2hist(array, hist):
     dtype = np.dtype(DTYPE_ROOT2NUMPY[hist_type])
     # No copy is made if the dtype is the same as input
     _array = np.ascontiguousarray(array, dtype=dtype)
+    if errors is not None:
+        if errors.shape != array.shape:
+            raise ValueError("Contents and errors are not compatible")
+        # errors are specified as doubles in SetError function
+        _errors = np.ascontiguousarray(errors, dtype=np.float64)
+    else:
+        _errors = None
+
     if _array.ndim != len(shape):
         raise ValueError(
             "array and histogram do not have "
@@ -400,8 +411,16 @@ def array2hist(array, hist):
         array_overflow = np.zeros(shape, dtype=dtype)
         array_overflow[tuple(slices)] = _array
         _array = array_overflow
+
+        if _errors is not None:
+            errors_overflow = np.zeros(shape, dtype=np.float64)
+            errors_overflow[tuple(slices)] = _errors
+            _errors = errors_overflow
+
     ARRAY_NUMPY2ROOT[len(shape)][hist_type](
         ROOT.AsCObject(hist), np.ravel(np.transpose(_array)))
     # Set the number of entries to the number of array elements
     hist.SetEntries(_array.size)
+    if _errors is not None:
+        hist.SetError(np.ravel(_errors.T))
     return hist

root_numpy/tests/test_hist.py

@@ -136,6 +136,27 @@ def check_array2hist(hist):
     arr_hist = rnp.hist2array(hist)
     assert_array_equal(arr_hist, arr)
 
+    # Check behaviour if errors are supplied
+    errors = arr * 0.1
+    _hist = hist.Clone()
+    _hist.Reset()
+    rnp.array2hist(arr, _hist, errors=errors)
+    arr_hist = rnp.hist2array(_hist)
+    assert_array_equal(arr_hist, arr)
+    if hist.GetDimension() == 1:
+        errors_from_hist = np.array([_hist.GetBinError(ix)
+                                     for ix in range(1, _hist.GetNbinsX() + 1)])
+    if hist.GetDimension() == 2:
+        errors_from_hist = np.array([[_hist.GetBinError(ix, iy)
+                                      for iy in range(1, _hist.GetNbinsY() + 1)]
+                                     for ix in range(1, _hist.GetNbinsX() + 1)])
+    if hist.GetDimension() == 3:
+        errors_from_hist = np.array([[[_hist.GetBinError(ix, iy, iz)
+                                       for iz in range(1, _hist.GetNbinsZ() + 1)]
+                                      for iy in range(1, _hist.GetNbinsY() + 1)]
+                                     for ix in range(1, _hist.GetNbinsX() + 1)])
+    assert_array_equal(errors, errors_from_hist)
+
     shape_overflow = shape + 2
     arr_overflow = RNG.randint(0, 10, size=shape_overflow)
     hist_overflow = hist.Clone()
@@ -170,12 +191,21 @@ def test_array2hist():
     assert_raises(ValueError, rnp.array2hist,
                   np.arange(4).reshape(2, 2),
                   ROOT.TH2D('test', '', 4, 0, 1, 3, 0, 1))
+    # shape of errors and content array does not match
+    assert_raises(ValueError, rnp.array2hist,
+                  np.arange(4).reshape(2, 2),
+                  ROOT.TH2D('test', '', 4, 0, 1, 4, 0, 1),
+                  np.arange(6).reshape(2, 3))
 
     for ndim in (1, 2, 3):
         for hist_type in 'DFISC':
             hist = make_histogram(hist_type, shape=(5,) * ndim, fill=False)
             yield check_array2hist, hist
 
+    # Check for histograms with unequal dimensions (reveals issues with transposing)
+    hist = make_histogram(hist_type, shape=(5, 6, 7), fill=False)
+    check_array2hist(hist)
+
 
 def test_fill_hist():
     n_samples = 1000