import ROOT
from array import array
def makeTH1(name, mean, sigma):
"""Create ROOT TH1 filled with a Gaussian distribution."""
hh = ROOT.TH1D(name, name, 100, -10, 10)
for i in range(1000):
hh.Fill(ROOT.gRandom.Gaus(mean, sigma))
return hh
def makeTTree():
"""Create ROOT ROOT.TTree filled with a Gaussian distribution in x and a uniform distribution in y."""
tree = ROOT.TTree("tree", "tree")
px = array("d", [0])
py = array("d", [0])
pz = array("d", [0])
pi = array("i", [0])
tree.Branch("x", px, "x/D")
tree.Branch("y", py, "y/D")
tree.Branch("z", pz, "z/D")
tree.Branch("i", pi, "i/I")
for i in range(100):
px[0] = ROOT.gRandom.Gaus(0, 3)
py[0] = ROOT.gRandom.Uniform() * 30 - 15
pz[0] = ROOT.gRandom.Gaus(0, 5)
pi[0] = i % 3
tree.Fill()
return tree
hh_1 = makeTH1("hh1", 0, 3)
hh_2 = makeTH1("hh2", -3, 1)
hh_3 = makeTH1("hh3", +3, 4)
x = ROOT.RooRealVar("x", "x", -10, 10)
c = ROOT.RooCategory("c", "c")
c.defineType("SampleA")
c.defineType("SampleB")
c.defineType("SampleC")
dh = ROOT.RooDataHist("dh", "dh", [x], Index=c, Import={"SampleA": hh_1, "SampleB": hh_2, "SampleC": hh_3})
dh.Print()
dh2 = ROOT.RooDataHist("dh", "dh", [x], Index=c, Import={"SampleA": hh_1, "SampleB": hh_2, "SampleC": hh_3})
dh2.Print()
tree = makeTTree()
y = ROOT.RooRealVar("y", "y", -10, 10)
z = ROOT.RooRealVar("z", "z", -10, 10)
ds = ROOT.RooDataSet("ds", "ds", {x, y}, Import=tree)
ds.Print()
ds2 = ROOT.RooDataSet("ds2", "ds2", {x, y, z}, Import=tree, Cut="y+z<0")
ds2.Print()
i = ROOT.RooRealVar("i", "i", 0, 5)
ds3 = ROOT.RooDataSet("ds3", "ds3", {i, x}, Import=tree)
ds3.Print()
icat = ROOT.RooCategory("i", "i", {"State0": 0, "State1": 1})
ds4 = ROOT.RooDataSet("ds4", "ds4", {icat, x}, Import=tree)
ds4.Print()
dsA = ds2.reduce({x, y}, "z<-5")
dsB = ds2.reduce({x, y}, "abs(z)<5")
dsC = ds2.reduce({x, y}, "z>5")
dsABC = ROOT.RooDataSet(
"dsABC",
"dsABC",
{x, y},
ROOT.RooFit.Import("SampleA", dsA),
ROOT.RooFit.Import("SampleB", dsB),
Index=c,
Import=("SampleC", dsC),
)
dsABC.Print()