Fixes for root2numpy

mkShapesRDF/lib/utils.py

@@ -1,3 +1,6 @@
+import os
+
+
 def getFrameworkPath():
     r"""Utility function to get the absolute path to the mkShapesRDF framework
 
@@ -6,15 +9,13 @@ def getFrameworkPath():
         string
             absolute path to the mkShapesRDF framework (ends with ``/``)
     """
-    import os
-    import inspect
 
-    fwPath = os.path.realpath(inspect.getfile(inspect.currentframe()))  # this file
-    fwPath = os.path.dirname(fwPath)  # lib
-    fwPath = os.path.dirname(fwPath)  # mkShapesRDF (source code)
-    fwPath = os.path.dirname(fwPath)  # mkShapesRDF
-    fwPath = os.path.abspath(fwPath)  # abs path to mkShapesRDF
+    try:
+        fwPath = os.environ["STARTPATH"]
+        fwPath = fwPath[: -len("start.sh")]
+    except Exception as _:
+        raise Exception(
+            "STARTPATH is not set! Please be sure you've activated the environment"
+        )
 
-    if not fwPath.endswith("/"):
-        fwPath += "/"
     return fwPath

mkShapesRDF/shapeAnalysis/histo_utils.py

@@ -72,7 +72,7 @@ def fold(h, ifrom, ito):
         sumw2[:, :, ifrom] = 0.0
 
 
-def postPlot(hTotal, doFold, unroll=True):
+def postPlot(hTotal, doFold, unroll=True, delete_oldHisto=True):
     """
     Takes TH1, TH2 or TH3, folds if necessary and unrolls
 
@@ -140,7 +140,8 @@ def postPlot(hTotal, doFold, unroll=True):
         stats1d[3] = stats2d[3] + stats2d[5]
         hTotal.PutStats(stats1d.GetArray())
 
-        hTotal_rolled.Delete()
+        if delete_oldHisto:
+            hTotal_rolled.Delete()
 
     if unroll and isinstance(hTotal, ROOT.TH3) and hTotal.GetDimension() == 3:
         # --- transform 2D into 1D
@@ -199,7 +200,8 @@ def postPlot(hTotal, doFold, unroll=True):
         stats1d[3] = stats3d[3] + stats3d[5] + stats3d[8]
         hTotal.PutStats(stats1d.GetArray())
 
-        hTotal_rolled.Delete()
+        if delete_oldHisto:
+            hTotal_rolled.Delete()
 
     return hTotal
 

mkShapesRDF/shapeAnalysis/rnp.py

@@ -30,7 +30,7 @@ def rnp_array(array, copy=True):
     return arr
 
 
-def rnp_hist2array(h, include_overflow=False, copy=True):
+def rnp_hist2array(h, include_overflow=False, copy=True, return_errors=False):
     """
     Converts histogram into a numpy array
 
@@ -55,16 +55,28 @@ def rnp_hist2array(h, include_overflow=False, copy=True):
         shape = (h.GetNbinsY() + 2, h.GetNbinsX() + 2)
     elif isinstance(h, ROOT.TH1):
         shape = (h.GetNbinsX() + 2,)
+
+    if return_errors:
+        errors = np.sqrt(rnp_array(h.GetSumw2(), copy=copy))
+    else:
+        errors = np.zeros_like(arr)
+
     arr = arr.reshape(shape)
+    errors = errors.reshape(shape)
     if not include_overflow:
         slices = []
         for axis, bins in enumerate(shape):
             slices.append(slice(1, -1))
         arr = arr[tuple(slices)]
+        errors = errors[tuple(slices)]
+    arr = np.transpose(arr)
+    errors = np.transpose(errors)
+    if return_errors:
+        return arr, errors
     return arr
 
 
-def rnp_array2hist(array, h):
+def rnp_array2hist(array, h, errors=None):
     """
     Sets bin contents with a numpy array, modifying it in place
 
@@ -82,6 +94,19 @@ def rnp_array2hist(array, h):
         shape = (h.GetNbinsX() + 2, h.GetNbinsY() + 2)
     elif isinstance(h, ROOT.TH1):
         shape = (h.GetNbinsX() + 2,)
+
+    if errors is not None:
+        if errors.shape != array.shape:
+            raise ValueError(
+                "Bin contents and Bin errors have different shapes",
+                array.shape,
+                errors.shape,
+            )
+        _errors = errors
+    else:
+        # dummy vector, it won't be set at the end
+        _errors = np.zeros_like(array)
+
     if array.shape != shape:
         slices = []
         for axis, bins in enumerate(shape):
@@ -92,11 +117,20 @@ def rnp_array2hist(array, h):
             else:
                 raise ValueError("array and histogram are not compatible")
         array_overflow = np.zeros(shape, dtype=dtype)
+        errors_overflow = np.zeros(shape, dtype=dtype)
         array_overflow[tuple(slices)] = array
+        errors_overflow[tuple(slices)] = _errors
         array = array_overflow
+        _errors = errors_overflow
+
     if array.shape != shape:
-        raise "array2hist: Different shape between array and h"
+        raise ValueError("array2hist: Different shape between array and h")
+
     array = np.ravel(np.transpose(array))
+    _errors = np.ravel(np.transpose(_errors))
+
     nx = len(array)
     arr = memoryview(array)
     h.Set(nx, arr)
+    if errors is not None:
+        h.SetError(_errors)

tests/test_rnp.py

@@ -0,0 +1,99 @@
+import ROOT
+import numpy as np
+from mkShapesRDF.shapeAnalysis import rnp
+from mkShapesRDF.shapeAnalysis import histo_utils
+
+ROOT.gROOT.SetBatch(True)
+
+rng = np.random.Generator(np.random.PCG64(1281))
+
+
+def test_array2hist():
+    nbins_x = 10
+    nbins_y = 15
+
+    nbins_x_with_fold = nbins_x + 2
+    nbins_y_with_fold = nbins_y + 2
+
+    bin_content = rng.normal(120, 8, size=(nbins_x_with_fold, nbins_y_with_fold))
+    bin_error = rng.normal(48, 2, size=(nbins_x_with_fold, nbins_y_with_fold))
+
+    h = ROOT.TH2D("", "", nbins_x, 0, 1, nbins_y, 0, 1)
+
+    # just for test
+    rnp.rnp_array2hist(bin_content, h)
+    _ = rnp.rnp_hist2array(
+        h, include_overflow=True, copy=True
+    )
+
+    # now the actual set and get method
+    rnp.rnp_array2hist(bin_content, h, bin_error)
+
+    bin_content2, bin_error2 = rnp.rnp_hist2array(
+        h, include_overflow=True, copy=True, return_errors=True
+    )
+
+    assert bin_content.shape == bin_content2.shape
+    assert np.all(bin_content == bin_content2)
+    assert bin_error.shape == bin_error2.shape
+    assert np.all(bin_error == bin_error2)
+
+    # No fold
+    h_unroll = histo_utils.postPlot(h, doFold=0, unroll=True, delete_oldHisto=False)
+
+    # We are good people -> we work x-major
+    bin_content3, bin_error3 = rnp.rnp_hist2array(
+        h, include_overflow=False, copy=True, return_errors=True
+    )
+    bin_content3 = bin_content3.flatten()
+    bin_error3 = bin_error3.flatten()
+    bin_content4, bin_error4 = rnp.rnp_hist2array(
+        h_unroll, include_overflow=False, copy=True, return_errors=True
+    )
+
+    assert bin_content3.shape == bin_content4.shape
+    assert np.all(bin_content3 == bin_content4)
+    assert bin_error3.shape == bin_error4.shape
+    assert np.all(bin_error3 == bin_error4)
+
+    # Do fold
+    h_for_unroll = h.Clone()
+    h_for_unroll.SetName("tmp2") # just to suppress warning of memoryleak
+    h_unroll = histo_utils.postPlot(h_for_unroll, doFold=3, unroll=True, delete_oldHisto=False)
+
+    # We're passing below h which has 2 dimension -> bin_content3 will be 2D
+    # We are good people -> we work x-major
+    bin_content3, bin_error3 = rnp.rnp_hist2array(
+        h, include_overflow=True, copy=True, return_errors=True
+    )
+
+
+    bin_content3 = bin_content3[1:-1, 1:-1]
+    bin_error3 = bin_error3[1:-1, 1:-1]
+
+    edges = np.zeros_like(bin_content3) != 0.0
+    edges[0, :] = True
+    edges[-1, :] = True
+    edges[:, 0] = True
+    edges[:, -1] = True
+
+    bin_content3 = bin_content3.flatten()
+    bin_error3 = bin_error3.flatten()
+    edges = edges.flatten()
+    bin_content4, bin_error4 = rnp.rnp_hist2array(
+        h_unroll, include_overflow=False, copy=True, return_errors=True
+    )
+
+    # Here we know that the bin content of under and overflow was folded
+    # we therefore expect the bin content to be different in the edges
+
+    assert bin_content3.shape == bin_content4.shape
+    assert np.all(bin_content3[~edges] == bin_content4[~edges])
+    assert np.all(bin_content3[edges] != bin_content4[edges])
+    assert bin_error3.shape == bin_error4.shape
+    assert np.all(bin_error3[~edges] == bin_error4[~edges])
+    assert np.all(bin_error3[edges] != bin_error4[edges])
+
+
+test_array2hist()
+print('Good, no errors for now ;)')