Example of sWeight generation and sWeight distributions in RooFit

sweights_example.png

sweights_example.py

import ROOT
from ROOT import RooFit as RF

# Create the observables (x), parameters, and PDFs
w = ROOT.RooWorkspace('w')
w.factory('RooGaussian::sig_pdf(x[0, 10], mean[5], sigma[1])')
w.factory('RooExponential::bkg_pdf(x, lambda[0.2])')
tot_pdf = w.factory(
    'SUM::tot_pdf(nsig[1000, 0, 2000]*sig_pdf, nbkg[1200, 0, 2000]*bkg_pdf)'
)
x = w.var('x')

# Generate some Monte Carlo data and fit to it
data = tot_pdf.generate(ROOT.RooArgSet(x), 2200)
tot_pdf.fitTo(data, RF.Extended(True))

# Compute sWeights
ROOT.RooStats.SPlot(
    'sData', 'sData', data, tot_pdf,
    ROOT.RooArgList(w.var('nsig'), w.var('nbkg'))
)
# Import the dataset so we can easily access everything
getattr(w, 'import')(data, ROOT.RooCmdArg())

# Plot the fit
frame_fit = x.frame(RF.Title('Gaussian + polynomial fit'))
data.plotOn(frame_fit)
tot_pdf.plotOn(frame_fit, RF.Components('bkg_pdf'), RF.LineColor(ROOT.kRed))
tot_pdf.plotOn(frame_fit, RF.Components('sig_pdf'), RF.LineStyle(ROOT.kDashed))
tot_pdf.plotOn(frame_fit)

# Plot the signal sWeights
nsig_sw = w.var('nsig_sw')
nsig_sw.SetTitle('Signal sWeight')
nsig_sw.setRange(-1.0, 1.5)
frame_sigsw = nsig_sw.frame(RF.Title('Signal sWeights'))
data.plotOn(frame_sigsw)
# Plot the background sWeights
nbkg_sw = w.var('nbkg_sw')
nbkg_sw.SetTitle('Background sWeight')
nbkg_sw.setRange(-0.5, 2.0)
frame_bkgsw = w.var('nbkg_sw').frame(RF.Title('Background sWeights'))
data.plotOn(frame_bkgsw)
# Plot the signal sWeight against the observable x
frame_sigsw_x = w.var('x').frame(RF.Title('sWeights vs. x'))
data.plotOnXY(frame_sigsw_x, RF.YVar(nsig_sw),
              RF.MarkerColor(ROOT.kBlue), RF.Name('sig'))
data.plotOnXY(frame_sigsw_x, RF.YVar(nbkg_sw),
              RF.MarkerColor(ROOT.kRed), RF.Name('bkg'))
legend = ROOT.TLegend(0.89, 0.89, 0.5, 0.7)
legend.AddEntry(frame_sigsw_x.findObject('sig'), 'Signal sWeights', 'p')
legend.AddEntry(frame_sigsw_x.findObject('bkg'), 'Background sWeights', 'p')
legend.SetBorderSize(0)

# Save the plot and sWeights distribution as a PDF
c = ROOT.TCanvas('c', 'c', 1600, 1600)
c.Divide(2, 2)
c.cd(1)
frame_fit.Draw()
c.cd(2)
frame_sigsw.Draw()
c.cd(3)
frame_bkgsw.Draw()
c.cd(4)
frame_sigsw_x.Draw()
legend.Draw()
c.SaveAs('sweights_example.pdf')