From f742b59800b34e73dca4e014744f15aff2493fba Mon Sep 17 00:00:00 2001 From: "S. Lehti" Date: Mon, 16 Aug 2021 10:29:12 +0300 Subject: [PATCH 1/3] First commit of hist2root.py --- src/hist/hist2root.py | 94 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) create mode 100644 src/hist/hist2root.py diff --git a/src/hist/hist2root.py b/src/hist/hist2root.py new file mode 100644 index 00000000..6170ff92 --- /dev/null +++ b/src/hist/hist2root.py @@ -0,0 +1,94 @@ +from coffea import processor +import ROOT +import math + +def convert(hist,hname=""): + """Convert a 1-dimensional or 2-dimensional`Hist` object or a counter to ROOT.TH1D or ROOT.TH2D + Parameters + ---------- + hist : Hist + A 1-dimensional or 2-dimensional histogram object or a counter + hname: new name for the output histogram + Returns + ------- + out + A ``ROOT.TH1D`` or ``ROOT.TH2D`` object + Examples + -------- + Saving histograms in the coffea result into a ROOT file + + result = processor.run_uproot_job( + samples, + "Events", + Analysis(), + processor.iterative_executor, + {"schema": NanoAODSchema}, + ) + + import ROOT + import hist2root + + fOUT = ROOT.TFile.Open('output.root','RECREATE') + for k in result.keys(): + hist2root.convert(result[k],k).Write() + fOUT.Close() + """ + if isinstance(hist,processor.defaultdict_accumulator): + return convertCounter(hist,hname) + + if len(hist.axes()) == 1: + return convert2TH1D(hist,hname) + + if len(hist.axes()) == 2: + return convert2TH2D(hist,hname) + +def convertCounter(counter,hname): + n = len(counter.keys()) + out = ROOT.TH1D(hname,"",n,0,n) + for i,key in enumerate(counter.keys()): + out.SetBinContent(i+1,counter[key]) + out.GetXaxis().SetBinLabel(i+1,key) + return out + +def convert2TH1D(hist,hname): + name = hist.label + if len(name) > 0: + name = hname + title = hist.label + axis = hist.axes()[0] + edges = axis.edges(overflow="none") + edgesOFlow = axis.edges(overflow="all") + + out = ROOT.TH1D(name,title,len(edges)-1,edges) + for i in range(0,len(edgesOFlow)-1): + bincenter = edgesOFlow[i] + 0.5*(edgesOFlow[i+1]-edgesOFlow[i]) + x = bincenter + w,wsq = hist.values(sumw2=True,overflow="all")[()] + out.Fill(x,w[i]) + out.SetBinError(i,math.sqrt(wsq[i])) + return out + +def convert2TH2D(hist,hname): + name = hist.label + if len(name) > 0: + name = hname + title = hist.label + axis1 = hist.axes()[0] + xedges = axis1.edges(overflow="none") + xedgesOFlow = axis1.edges(overflow="all") + axis2 = hist.axes()[1] + yedges = axis2.edges(overflow="none") + yedgesOFlow = axis2.edges(overflow="all") + + out = ROOT.TH2D(name,title,len(xedges)-1,xedges,len(yedges)-1,yedges) + + for i in range(0,len(xedgesOFlow)-1): + xbincenter = xedgesOFlow[i] + 0.5*(xedgesOFlow[i+1]-xedgesOFlow[i]) + x = xbincenter + for j in range(0,len(yedgesOFlow)-1): + ybincenter = yedgesOFlow[j] + 0.5*(yedgesOFlow[j+1]-yedgesOFlow[j]) + y = ybincenter + w,wsq = hist.values(sumw2=True,overflow="all")[()] + out.Fill(x,y,w[j][i]) + out.SetBinError(i,j,math.sqrt(wsq[j][i])) + return out From fab49bc99cc2bb5da230d976163b56e67e6d1505 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Aug 2021 07:40:51 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/hist/hist2root.py | 73 +++++++++++++++++++++++-------------------- 1 file changed, 39 insertions(+), 34 deletions(-) diff --git a/src/hist/hist2root.py b/src/hist/hist2root.py index 6170ff92..0c392f80 100644 --- a/src/hist/hist2root.py +++ b/src/hist/hist2root.py @@ -1,8 +1,10 @@ -from coffea import processor -import ROOT import math -def convert(hist,hname=""): +import ROOT +from coffea import processor + + +def convert(hist, hname=""): """Convert a 1-dimensional or 2-dimensional`Hist` object or a counter to ROOT.TH1D or ROOT.TH2D Parameters ---------- @@ -33,45 +35,48 @@ def convert(hist,hname=""): hist2root.convert(result[k],k).Write() fOUT.Close() """ - if isinstance(hist,processor.defaultdict_accumulator): - return convertCounter(hist,hname) - + if isinstance(hist, processor.defaultdict_accumulator): + return convertCounter(hist, hname) + if len(hist.axes()) == 1: - return convert2TH1D(hist,hname) + return convert2TH1D(hist, hname) if len(hist.axes()) == 2: - return convert2TH2D(hist,hname) + return convert2TH2D(hist, hname) + -def convertCounter(counter,hname): +def convertCounter(counter, hname): n = len(counter.keys()) - out = ROOT.TH1D(hname,"",n,0,n) - for i,key in enumerate(counter.keys()): - out.SetBinContent(i+1,counter[key]) - out.GetXaxis().SetBinLabel(i+1,key) + out = ROOT.TH1D(hname, "", n, 0, n) + for i, key in enumerate(counter.keys()): + out.SetBinContent(i + 1, counter[key]) + out.GetXaxis().SetBinLabel(i + 1, key) return out - -def convert2TH1D(hist,hname): + + +def convert2TH1D(hist, hname): name = hist.label if len(name) > 0: - name = hname + name = hname title = hist.label - axis = hist.axes()[0] + axis = hist.axes()[0] edges = axis.edges(overflow="none") edgesOFlow = axis.edges(overflow="all") - out = ROOT.TH1D(name,title,len(edges)-1,edges) - for i in range(0,len(edgesOFlow)-1): - bincenter = edgesOFlow[i] + 0.5*(edgesOFlow[i+1]-edgesOFlow[i]) + out = ROOT.TH1D(name, title, len(edges) - 1, edges) + for i in range(0, len(edgesOFlow) - 1): + bincenter = edgesOFlow[i] + 0.5 * (edgesOFlow[i + 1] - edgesOFlow[i]) x = bincenter - w,wsq = hist.values(sumw2=True,overflow="all")[()] - out.Fill(x,w[i]) - out.SetBinError(i,math.sqrt(wsq[i])) + w, wsq = hist.values(sumw2=True, overflow="all")[()] + out.Fill(x, w[i]) + out.SetBinError(i, math.sqrt(wsq[i])) return out -def convert2TH2D(hist,hname): + +def convert2TH2D(hist, hname): name = hist.label if len(name) > 0: - name = hname + name = hname title = hist.label axis1 = hist.axes()[0] xedges = axis1.edges(overflow="none") @@ -79,16 +84,16 @@ def convert2TH2D(hist,hname): axis2 = hist.axes()[1] yedges = axis2.edges(overflow="none") yedgesOFlow = axis2.edges(overflow="all") - - out = ROOT.TH2D(name,title,len(xedges)-1,xedges,len(yedges)-1,yedges) - for i in range(0,len(xedgesOFlow)-1): - xbincenter = xedgesOFlow[i] + 0.5*(xedgesOFlow[i+1]-xedgesOFlow[i]) + out = ROOT.TH2D(name, title, len(xedges) - 1, xedges, len(yedges) - 1, yedges) + + for i in range(0, len(xedgesOFlow) - 1): + xbincenter = xedgesOFlow[i] + 0.5 * (xedgesOFlow[i + 1] - xedgesOFlow[i]) x = xbincenter - for j in range(0,len(yedgesOFlow)-1): - ybincenter = yedgesOFlow[j] + 0.5*(yedgesOFlow[j+1]-yedgesOFlow[j]) + for j in range(0, len(yedgesOFlow) - 1): + ybincenter = yedgesOFlow[j] + 0.5 * (yedgesOFlow[j + 1] - yedgesOFlow[j]) y = ybincenter - w,wsq = hist.values(sumw2=True,overflow="all")[()] - out.Fill(x,y,w[j][i]) - out.SetBinError(i,j,math.sqrt(wsq[j][i])) + w, wsq = hist.values(sumw2=True, overflow="all")[()] + out.Fill(x, y, w[j][i]) + out.SetBinError(i, j, math.sqrt(wsq[j][i])) return out From f4abf4433a20381f143fccd8d55f977184ff6866 Mon Sep 17 00:00:00 2001 From: Henry Schreiner Date: Sat, 21 Aug 2021 10:17:54 -0400 Subject: [PATCH 3/3] WIP: start adding static typing --- pyproject.toml | 2 ++ src/hist/hist2root.py | 26 ++++++++++++++++---------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a2637fb7..7c6434d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,5 +69,7 @@ module = [ "iminuit.*", "mplhep.*", "hist.version", + "ROOT.*", + "coffea.*", ] ignore_missing_imports = true diff --git a/src/hist/hist2root.py b/src/hist/hist2root.py index 0c392f80..0564de2a 100644 --- a/src/hist/hist2root.py +++ b/src/hist/hist2root.py @@ -1,11 +1,15 @@ +from __future__ import annotations + import math import ROOT from coffea import processor +import hist.basehist + -def convert(hist, hname=""): - """Convert a 1-dimensional or 2-dimensional`Hist` object or a counter to ROOT.TH1D or ROOT.TH2D +def convert(hist: hist.basehist.BaseHist, hname: str = "") -> ROOT.TH1D | ROOT.TH2D: + """Convert a 1-dimensional or 2-dimensional `Hist` object or a counter to ROOT.TH1D or ROOT.TH2D Parameters ---------- hist : Hist @@ -38,14 +42,16 @@ def convert(hist, hname=""): if isinstance(hist, processor.defaultdict_accumulator): return convertCounter(hist, hname) - if len(hist.axes()) == 1: + if len(hist.axes) == 1: return convert2TH1D(hist, hname) - if len(hist.axes()) == 2: + if len(hist.axes) == 2: return convert2TH2D(hist, hname) + raise TypeError("Not a known ROOT type") + -def convertCounter(counter, hname): +def convertCounter(counter: hist.basehist.BaseHist, hname: str): n = len(counter.keys()) out = ROOT.TH1D(hname, "", n, 0, n) for i, key in enumerate(counter.keys()): @@ -54,12 +60,12 @@ def convertCounter(counter, hname): return out -def convert2TH1D(hist, hname): +def convert2TH1D(hist: hist.basehist.BaseHist, hname: str) -> ROOT.TH1D: name = hist.label if len(name) > 0: name = hname title = hist.label - axis = hist.axes()[0] + axis = hist.axes[0] edges = axis.edges(overflow="none") edgesOFlow = axis.edges(overflow="all") @@ -73,15 +79,15 @@ def convert2TH1D(hist, hname): return out -def convert2TH2D(hist, hname): +def convert2TH2D(hist: hist.basehist.BaseHist, hname: str) -> ROOT.TH2D: name = hist.label if len(name) > 0: name = hname title = hist.label - axis1 = hist.axes()[0] + axis1 = hist.axes[0] xedges = axis1.edges(overflow="none") xedgesOFlow = axis1.edges(overflow="all") - axis2 = hist.axes()[1] + axis2 = hist.axes[1] yedges = axis2.edges(overflow="none") yedgesOFlow = axis2.edges(overflow="all")