from __future__ import print_function
import ROOT
x = ROOT.RooRealVar("x", "x", -10, 10)
x.setBins(40)
p0 = ROOT.RooPolynomial("px", "px", x)
data = p0.generate({x}, 1000)
wFunc = ROOT.RooFormulaVar("w", "event weight", "(x*x+10)", [x])
w = data.addColumn(wFunc)
data.Print()
wdata = ROOT.RooDataSet(data.GetName(), data.GetTitle(), data, data.get(), "", w.GetName())
wdata.Print()
a0 = ROOT.RooRealVar("a0", "a0", 1)
a1 = ROOT.RooRealVar("a1", "a1", 0, -1, 1)
a2 = ROOT.RooRealVar("a2", "a2", 1, 0, 10)
p2 = ROOT.RooPolynomial("p2", "p2", x, [a0, a1, a2], 0)
r_ml_wgt = p2.fitTo(wdata, Save=True)
r_ml_wgt_corr = p2.fitTo(wdata, Save=True, SumW2Error=True)
frame = x.frame(Title="Unbinned ML fit, chi^2 fit to weighted data")
wdata.plotOn(frame, DataError="SumW2")
p2.plotOn(frame)
genPdf = ROOT.RooGenericPdf("genPdf", "x*x+10", [x])
data2 = genPdf.generate({x}, 1000)
data3 = genPdf.generate({x}, 43000)
r_ml_unw10 = p2.fitTo(data2, Save=True)
r_ml_unw43 = p2.fitTo(data3, Save=True)
binnedData = wdata.binnedClone()
binnedData.Print("v")
chi2 = p2.createChi2(binnedData, ROOT.RooFit.DataError("SumW2"))
m = ROOT.RooMinimizer(chi2)
m.migrad()
m.hesse()
r_chi2_wgt = m.save()
p2.plotOn(frame, LineStyle="--", LineColor="r")
print("==> ML Fit results on 1K unweighted events")
r_ml_unw10.Print()
print("==> ML Fit results on 43K unweighted events")
r_ml_unw43.Print()
print("==> ML Fit results on 1K weighted events with a summed weight of 43K")
r_ml_wgt.Print()
print("==> Corrected ML Fit results on 1K weighted events with a summed weight of 43K")
r_ml_wgt_corr.Print()
print("==> Chi2 Fit results on 1K weighted events with a summed weight of 43K")
r_chi2_wgt.Print()
c = ROOT.TCanvas("rf403_weightedevts", "rf403_weightedevts", 600, 600)
ROOT.gPad.SetLeftMargin(0.15)
frame.GetYaxis().SetTitleOffset(1.8)
frame.Draw()
c.SaveAs("rf403_weightedevts.png")