diff --git a/mkShapesRDF/lib/utils.py b/mkShapesRDF/lib/utils.py index e50d3dda..c3642370 100644 --- a/mkShapesRDF/lib/utils.py +++ b/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 diff --git a/mkShapesRDF/shapeAnalysis/histo_utils.py b/mkShapesRDF/shapeAnalysis/histo_utils.py index 35a09284..4df1df78 100644 --- a/mkShapesRDF/shapeAnalysis/histo_utils.py +++ b/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 diff --git a/mkShapesRDF/shapeAnalysis/rnp.py b/mkShapesRDF/shapeAnalysis/rnp.py index b8a33593..63c93581 100644 --- a/mkShapesRDF/shapeAnalysis/rnp.py +++ b/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) diff --git a/tests/test_rnp.py b/tests/test_rnp.py new file mode 100644 index 00000000..955a16d4 --- /dev/null +++ b/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 ;)') \ No newline at end of file