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 new file mode 100644 index 00000000..0564de2a --- /dev/null +++ b/src/hist/hist2root.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import math + +import ROOT +from coffea import processor + +import hist.basehist + + +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 + 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) + + raise TypeError("Not a known ROOT type") + + +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()): + out.SetBinContent(i + 1, counter[key]) + out.GetXaxis().SetBinLabel(i + 1, key) + return out + + +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] + 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: hist.basehist.BaseHist, hname: str) -> ROOT.TH2D: + 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