Esempio n. 1
0
def generate_testfiles():
    h = TH1D("gaussian_hist", "Gaussian histgram", 100, -3, 3)
    h.FillRandom("gaus", 1000)

    file = TFile(
        os.path.dirname(__file__) + ('/../testfiles/root_testfiles.root'),
        'RECREATE')
    file.cd()
    h.Write()

    x = RooRealVar("D0_M", "m(K_{S}^{0}K^{+}K^{-})", 1860, 1800, 1930,
                   "\\mathrm{MeV}/c^{2}")
    x.setBins(130)
    m1 = RooRealVar("m1", "mean 1", 1864, 1860, 1870)
    s1 = RooRealVar("s1", "sigma 1", 2, 0, 5)
    g1 = RooGaussian("g1", "Gaussian 1", x, m1, s1)

    m2 = RooRealVar("m2", "mean 2", 1864, 1860, 1870)
    s2 = RooRealVar("s2", "sigma 2", 4, 0, 5)
    g2 = RooGaussian("g2", "Gaussian 2", x, m2, s2)

    f1 = RooRealVar("f", "f", 0.5, 0, 1)
    m = RooAddPdf("model", "model", RooArgList(g1, g2), f1)

    data = m.generate(x, 1e6)

    x.Write("x")
    m.Write("model")
    data.Write("data")

    file.Close()

    return
Esempio n. 2
0
def rooFit203():

    print ">>> setup model..."
    x = RooRealVar("x", "x", -10, 10)
    mean = RooRealVar("mean", "mean of gaussian", 0, -10, 10)
    gauss = RooGaussian("gauss", "gaussian PDF", x, mean, RooConst(1))

    # Construct px = 1 (flat in x)
    px = RooPolynomial("px", "px", x)

    # Construct model = f*gx + (1-f)px
    f = RooRealVar("f", "f", 0., 1.)
    model = RooAddPdf("model", "model", RooArgList(gauss, px), RooArgList(f))
    data = model.generate(RooArgSet(x), 10000)  # RooDataSet

    print ">>> fit to full data range..."
    result_full = model.fitTo(data, Save(kTRUE))  # RooFitResult

    print "\n>>> fit \"signal\" range..."
    # Define "signal" range in x as [-3,3]
    x.setRange("signal", -3, 3)
    result_sig = model.fitTo(data, Save(kTRUE),
                             Range("signal"))  # RooFitResult

    print "\n>>> plot and print results..."
    # Make plot frame in x and add data and fitted model
    frame1 = x.frame(Title("Fitting a sub range"))  # RooPlot
    data.plotOn(frame1, Name("data"))
    model.plotOn(frame1, Range("Full"), LineColor(kBlue),
                 Name("model"))  # Add shape in full ranged dashed
    model.plotOn(frame1, LineStyle(kDashed), LineColor(kRed),
                 Name("model2"))  # By default only fitted range is shown

    print "\n>>> result of fit on all data:"
    result_full.Print()

    print "\n>>> result of fit in in signal region (note increased error on signal fraction):"
    result_sig.Print()

    print ">>> draw on canvas..."
    canvas = TCanvas("canvas", "canvas", 100, 100, 800, 600)
    legend = TLegend(0.2, 0.85, 0.4, 0.65)
    legend.SetTextSize(0.032)
    legend.SetBorderSize(0)
    legend.SetFillStyle(0)
    gPad.SetLeftMargin(0.14)
    gPad.SetRightMargin(0.02)
    frame1.GetYaxis().SetLabelOffset(0.008)
    frame1.GetYaxis().SetTitleOffset(1.4)
    frame1.GetYaxis().SetTitleSize(0.045)
    frame1.GetXaxis().SetTitleSize(0.045)
    frame1.Draw()
    legend.AddEntry("data", "data", 'LEP')
    legend.AddEntry("model", "fit (full range)", 'L')
    legend.AddEntry("model2", "fit (signal range)", 'L')
    legend.Draw()
    canvas.SaveAs("rooFit203.png")
Esempio n. 3
0
def test_correlated_values():

    try:
        import uncertainties
    except ImportError:
        raise SkipTest("uncertainties package is not installed")
    from rootpy.stats.correlated_values import correlated_values

    # construct pdf and toy data following example at
    # http://root.cern.ch/drupal/content/roofit

    # --- Observable ---
    mes = RooRealVar("mes", "m_{ES} (GeV)", 5.20, 5.30)

    # --- Parameters ---
    sigmean = RooRealVar("sigmean", "B^{#pm} mass", 5.28, 5.20, 5.30)
    sigwidth = RooRealVar("sigwidth", "B^{#pm} width", 0.0027, 0.001, 1.)

    # --- Build Gaussian PDF ---
    signal = RooGaussian("signal", "signal PDF", mes, sigmean, sigwidth)

    # --- Build Argus background PDF ---
    argpar = RooRealVar("argpar", "argus shape parameter", -20.0, -100., -1.)
    background = RooArgusBG("background", "Argus PDF",
                            mes, RooFit.RooConst(5.291), argpar)

    # --- Construct signal+background PDF ---
    nsig = RooRealVar("nsig", "#signal events", 200, 0., 10000)
    nbkg = RooRealVar("nbkg", "#background events", 800, 0., 10000)
    model = RooAddPdf("model", "g+a",
                      RooArgList(signal,background),
                      RooArgList(nsig,nbkg))

    # --- Generate a toyMC sample from composite PDF ---
    data = model.generate(RooArgSet(mes), 2000)

    # --- Perform extended ML fit of composite PDF to toy data ---
    fitresult = model.fitTo(data, RooFit.Save(), RooFit.PrintLevel(-1))

    nsig, nbkg = correlated_values(["nsig", "nbkg"], fitresult)

    # Arbitrary math expression according to what the `uncertainties`
    # package supports, automatically computes correct error propagation
    sum_value = nsig + nbkg
    value, error = sum_value.nominal_value, sum_value.std_dev

    workspace = Workspace(name='workspace')
    # import the data
    assert_false(workspace(data))
    with TemporaryFile():
        workspace.Write()
def test_correlated_values():

    try:
        import uncertainties
    except ImportError:
        raise SkipTest("uncertainties package is not installed")
    from rootpy.stats.correlated_values import correlated_values

    # construct pdf and toy data following example at
    # http://root.cern.ch/drupal/content/roofit

    # --- Observable ---
    mes = RooRealVar("mes", "m_{ES} (GeV)", 5.20, 5.30)

    # --- Parameters ---
    sigmean = RooRealVar("sigmean", "B^{#pm} mass", 5.28, 5.20, 5.30)
    sigwidth = RooRealVar("sigwidth", "B^{#pm} width", 0.0027, 0.001, 1.)

    # --- Build Gaussian PDF ---
    signal = RooGaussian("signal", "signal PDF", mes, sigmean, sigwidth)

    # --- Build Argus background PDF ---
    argpar = RooRealVar("argpar", "argus shape parameter", -20.0, -100., -1.)
    background = RooArgusBG("background", "Argus PDF",
                            mes, RooFit.RooConst(5.291), argpar)

    # --- Construct signal+background PDF ---
    nsig = RooRealVar("nsig", "#signal events", 200, 0., 10000)
    nbkg = RooRealVar("nbkg", "#background events", 800, 0., 10000)
    model = RooAddPdf("model", "g+a",
                      RooArgList(signal,background),
                      RooArgList(nsig,nbkg))

    # --- Generate a toyMC sample from composite PDF ---
    data = model.generate(RooArgSet(mes), 2000)

    # --- Perform extended ML fit of composite PDF to toy data ---
    fitresult = model.fitTo(data, RooFit.Save(), RooFit.PrintLevel(-1))

    nsig, nbkg = correlated_values(["nsig", "nbkg"], fitresult)

    # Arbitrary math expression according to what the `uncertainties`
    # package supports, automatically computes correct error propagation
    sum_value = nsig + nbkg
    value, error = sum_value.nominal_value, sum_value.std_dev

    workspace = Workspace(name='workspace')
    # import the data
    assert_false(workspace(data))
    with TemporaryFile():
        workspace.Write()
Esempio n. 5
0
def toy_run(nevents):
    lower = -1
    upper = 1
    # create observables
    obs = RooRealVar("obs", "obs1", lower, upper)
    # create parameters
    mean1 = RooRealVar("mean1", "mean of gaussian", 0, -1, 1)
    sigma1 = RooRealVar("sigma1", "sigma of gaussian", 0.1, -1, 1)
    gauss1 = RooGaussian("gauss1", "gaussian PDF", obs, mean1, sigma1)

    mean2 = RooRealVar("mean2", "mean of gaussian", 0.5, -1, 1)
    sigma2 = RooRealVar("sigma2", "sigma of gaussian", 0.2, -1, 1)
    gauss2 = RooGaussian("gauss2", "gaussian PDF", obs, mean2, sigma2)
    frac = RooRealVar("frac", "Fraction of a gauss", 0.5, 0, 1)
    arg_list = RooArgList(
        gauss1,
        gauss2,
        gauss2,
        gauss2,
        gauss2,
        # gauss2,
        gauss2,
        gauss2,
        gauss1)
    arg_list.addOwned(gauss2)
    pdf = RooAddPdf(
        "sum_pdf",
        "sum of pdfs",
        arg_list,
        RooArgList(
            frac,
            frac,
            frac,
            # frac,
            # frac,
            frac,
            frac,
            frac,
            frac,
            frac))

    # obs, pdf = build_pdf()

    timer = zfit_benchmark.timer.Timer(f"Toys {nevents}")
    with timer:
        data = pdf.generate(RooArgSet(obs), nevents)
        pdf.fitTo(data)
        # mgr.generateAndFit(n_toys, nevents)

    return float(timer.elapsed)
Esempio n. 6
0
def test_plottable():

    # construct pdf and toy data following example at
    # http://root.cern.ch/drupal/content/roofit

    # Observable
    mes = RooRealVar("mes", "m_{ES} (GeV)", 5.20, 5.30)

    # Parameters
    sigmean = RooRealVar("sigmean", "B^{#pm} mass", 5.28, 5.20, 5.30)
    sigwidth = RooRealVar("sigwidth", "B^{#pm} width", 0.0027, 0.001, 1.)

    # Build Gaussian PDF
    signal = RooGaussian("signal", "signal PDF", mes, sigmean, sigwidth)

    # Build Argus background PDF
    argpar = RooRealVar("argpar", "argus shape parameter", -20.0, -100., -1.)
    background = RooArgusBG("background", "Argus PDF",
                            mes, RooFit.RooConst(5.291), argpar)

    # Construct signal+background PDF
    nsig = RooRealVar("nsig", "#signal events", 200, 0., 10000)
    nbkg = RooRealVar("nbkg", "#background events", 800, 0., 10000)
    model = RooAddPdf("model", "g+a",
                      RooArgList(signal, background),
                      RooArgList(nsig, nbkg))

    # Generate a toyMC sample from composite PDF
    data = model.generate(RooArgSet(mes), 2000)

    # Perform extended ML fit of composite PDF to toy data
    fitresult = model.fitTo(data, RooFit.Save(), RooFit.PrintLevel(-1))

    # Plot toy data and composite PDF overlaid
    mesframe = asrootpy(mes.frame())
    data.plotOn(mesframe)
    model.plotOn(mesframe)

    for obj in mesframe.objects:
        assert_true(obj)
    for curve in mesframe.curves:
        assert_true(curve)
    for hist in mesframe.data_hists:
        assert_true(hist)
    assert_true(mesframe.plotvar)
    with TemporaryFile():
        mesframe.Write()
Esempio n. 7
0
def test_plottable():

    # construct pdf and toy data following example at
    # http://root.cern.ch/drupal/content/roofit

    # Observable
    mes = RooRealVar("mes", "m_{ES} (GeV)", 5.20, 5.30)

    # Parameters
    sigmean = RooRealVar("sigmean", "B^{#pm} mass", 5.28, 5.20, 5.30)
    sigwidth = RooRealVar("sigwidth", "B^{#pm} width", 0.0027, 0.001, 1.)

    # Build Gaussian PDF
    signal = RooGaussian("signal", "signal PDF", mes, sigmean, sigwidth)

    # Build Argus background PDF
    argpar = RooRealVar("argpar", "argus shape parameter", -20.0, -100., -1.)
    background = RooArgusBG("background", "Argus PDF",
                            mes, RooFit.RooConst(5.291), argpar)

    # Construct signal+background PDF
    nsig = RooRealVar("nsig", "#signal events", 200, 0., 10000)
    nbkg = RooRealVar("nbkg", "#background events", 800, 0., 10000)
    model = RooAddPdf("model", "g+a",
                      RooArgList(signal, background),
                      RooArgList(nsig, nbkg))

    # Generate a toyMC sample from composite PDF
    data = model.generate(RooArgSet(mes), 2000)

    # Perform extended ML fit of composite PDF to toy data
    fitresult = model.fitTo(data, RooFit.Save(), RooFit.PrintLevel(-1))

    # Plot toy data and composite PDF overlaid
    mesframe = asrootpy(mes.frame())
    data.plotOn(mesframe)
    model.plotOn(mesframe)

    for obj in mesframe.objects:
        assert_true(obj)
    for curve in mesframe.curves:
        assert_true(curve)
    for hist in mesframe.data_hists:
        assert_true(hist)
    assert_true(mesframe.plotvar)
    with TemporaryFile():
        mesframe.Write()
Esempio n. 8
0
def rooFit502():

    print ">>> setup model components..."
    x = RooRealVar("x", "x", 0, 10)
    mean = RooRealVar("mean", "mean of gaussians", 5, 0, 10)
    sigma1 = RooRealVar("sigma1", "width of gaussians", 0.5)
    sigma2 = RooRealVar("sigma2", "width of gaussians", 1)
    sig1 = RooGaussian("sig1", "Signal component 1", x, mean, sigma1)
    sig2 = RooGaussian("sig2", "Signal component 2", x, mean, sigma2)
    a0 = RooRealVar("a0", "a0", 0.5, 0., 1.)
    a1 = RooRealVar("a1", "a1", -0.2, 0., 1.)
    bkg = RooChebychev("bkg", "Background", x, RooArgList(a0, a1))

    print ">>> sum model components..."
    sig1frac = RooRealVar("sig1frac", "fraction of component 1 in signal", 0.8,
                          0., 1.)
    sig = RooAddPdf("sig", "Signal", RooArgList(sig1, sig2),
                    RooArgList(sig1frac))
    bkgfrac = RooRealVar("bkgfrac", "fraction of background", 0.5, 0., 1.)
    model = RooAddPdf("model", "g1+g2+a", RooArgList(bkg, sig),
                      RooArgList(bkgfrac))

    print ">>> generate data..."
    data = model.generate(RooArgSet(x), 1000)  # RooDataSet

    print ">>> create workspace, import data and model..."
    workspace = RooWorkspace("workspace", "workspace")  # empty RooWorkspace
    getattr(workspace, 'import')(model)  # import model and all its components
    getattr(workspace, 'import')(data)  # import data
    #workspace.import(model) # causes synthax error in python
    #workspace.import(data)  # causes synthax error in python

    print "\n>>> print workspace contents:"
    workspace.Print()

    print "\n>>> save workspace in file..."
    workspace.writeToFile("rooFit502_workspace.root")

    print ">>> save workspace in memory (gDirectory)..."
    gDirectory.Add(workspace)
Esempio n. 9
0
def rooFit602():
    
    print ">>> setup model..."
    x      = RooRealVar("x","x",0,10)
    mean   = RooRealVar("mean","mean of gaussian",5)
    sigma1 = RooRealVar("sigma1","width of gaussian",0.5)
    sigma2 = RooRealVar("sigma2","width of gaussian",1)
    sig1   = RooGaussian("sig1","Signal component 1",x,mean,sigma1)
    sig2   = RooGaussian("sig2","Signal component 2",x,mean,sigma2)
    a0     = RooRealVar("a0","a0",0.5,0.,1.)
    a1     = RooRealVar("a1","a1",-0.2,0.,1.)
    bkg    = RooChebychev("bkg","Background",x,RooArgSet(a0,a1))
    sig1frac = RooRealVar("sig1frac","fraction of component 1 in signal",0.8,0.,1.)
    sig      = RooAddPdf("sig","Signal",RooArgList(sig1,sig2),sig1frac)
    bkgfrac = RooRealVar("bkgfrac","fraction of background",0.5,0.,1.)
    model   = RooAddPdf("model","g1+g2+a",RooArgList(bkg,sig),bkgfrac)
    
    print ">>> create binned dataset..."
    data = model.generate(RooArgSet(x),10000) # RooDataSet
    hist = data.binnedClone() # RooDataHist
    
    # Construct a chi^2 of the data and the model.
    # When a p.d.f. is used in a chi^2 fit, the probability density scaled
    # by the number of events in the dataset to obtain the fit function
    # If model is an extended p.d.f, the expected number events is used
    # instead of the observed number of events.
    model.chi2FitTo(hist)

    # NB: It is also possible to fit a RooAbsReal function to a RooDataHist
    # using chi2FitTo(). 

    # Note that entries with zero bins are _not_ allowed 
    # for a proper chi^2 calculation and will give error
    # messages
    data_small = date.reduce(EventRange(1,100)) # RooDataSet
    hist_small = data_small.binnedClone() # RooDataHist
    chi2_lowstat("chi2_lowstat","chi2",model,hist_small)
    print ">>> chi2_lowstat.getVal() = %s" % chi2_lowstat.getVal()
Esempio n. 10
0
def rooFit601():

    print ">>> setup pdf and likelihood..."
    x = RooRealVar("x", "x", -20, 20)
    mean = RooRealVar("mean", "mean of g1 and g2", 0)
    sigma1 = RooRealVar("sigma1", "width of g1", 3)
    sigma2 = RooRealVar("sigma2", "width of g2", 4, 3.0,
                        6.0)  # intentional strong correlations
    gauss1 = RooGaussian("gauss1", "gauss1", x, mean, sigma1)
    gauss2 = RooGaussian("gauss2", "gauss2", x, mean, sigma2)
    frac = RooRealVar("frac", "frac", 0.5, 0.0, 1.0)
    model = RooAddPdf("model", "model", RooArgList(gauss1, gauss2),
                      RooArgList(frac))

    print ">>> generate to data..."
    data = model.generate(RooArgSet(x), 1000)  # RooDataSet

    print ">>> construct unbinned likelihood of model wrt data..."
    nll = model.createNLL(data)  # RooAbsReal

    print ">>> interactive minimization and error analysis with MINUIT interface object..."
    minuit = RooMinuit(nll)

    print ">>> set avtive verbosity for logging of MINUIT parameter space stepping..."
    minuit.setVerbose(kTRUE)

    print ">>> call MIGRAD to minimize the likelihood..."
    minuit.migrad()

    print "\n>>> parameter values and error estimates that are back propagated from MINUIT:"
    model.getParameters(RooArgSet(x)).Print("s")

    print "\n>>> disable verbose logging..."
    minuit.setVerbose(kFALSE)

    print ">>> run HESSE to calculate errors from d2L/dp2..."
    minuit.hesse()

    print ">>> value of and error on sigma2 (back propagated from MINUIT):"
    sigma2.Print()

    print "\n>>> run MINOS on sigma2 parameter only..."
    minuit.minos(RooArgSet(sigma2))

    print "\n>>> value of and error on sigma2 (back propagated from MINUIT after running MINOS):"
    sigma2.Print()

    print "\n>>> saving results, contour plots..."
    # Save a snapshot of the fit result. This object contains the initial
    # fit parameters, the final fit parameters, the complete correlation
    # matrix, the EDM, the minimized FCN , the last MINUIT status code and
    # the number of times the RooFit function object has indicated evaluation
    # problems (e.g. zero probabilities during likelihood evaluation)
    result = minuit.save()  # RooFitResult

    # Make contour plot of mx vs sx at 1,2,3 sigma
    frame1 = minuit.contour(frac, sigma2, 1, 2, 3)  # RooPlot
    frame1.SetTitle("RooMinuit contour plot")

    # Print the fit result snapshot
    result.Print("v")

    print "\n>>> change value of \"mean\" parameter..."
    mean.setVal(0.3)

    # Rerun MIGRAD,HESSE
    print ">>> rerun MIGRAD, HESSE..."
    minuit.migrad()
    minuit.hesse()

    print ">>> value on and error of frac:"
    frac.Print()

    print "\n>>> fix value of \"sigma\" parameter (setConstant)..."
    sigma2.setConstant(kTRUE)

    print ">>> rerun MIGRAD, HESSE..."
    minuit.migrad()
    minuit.hesse()
    frac.Print()
Esempio n. 11
0
pdf_combine1 = RooAddPdf("pdf_combine1"," gauss4 + gauss5 ", RooArgList(gauss4 , gauss5 ), RooArgList( frac_combine1 ))


#pdf_combine1.plotOn(xframe3,RooFit.LineColor(RooFit.kBlue))

gauss4.plotOn(xframe3, RooFit.Normalization( frac_combine1.getVal()   ,RooAbsReal.Relative),RooFit.LineColor(RooFit.kOrange))
gauss5.plotOn(xframe3, RooFit.Normalization( 1-frac_combine1.getVal() ,RooAbsReal.Relative),RooFit.LineColor(RooFit.kCyan))
pdf_combine1.plotOn(xframe3, RooFit.Normalization(1.0,RooAbsReal.Relative) ,RooFit.LineColor(RooFit.kBlue))

# -------------------------------------------
# 4. use combine PDF to generate MC

xframe4 = x.frame(RooFit.Title("4. use combine PDF to generate MC"))

data2 = pdf_combine1.generate(RooArgSet(x),500)

#pdf_combine1.plotOn(xframe4, RooFit.Normalization(500, RooAbsReal.NumEvent) , RooFit.LineColor(RooFit.kOrange))
data2.plotOn(xframe4)
pdf_combine1.plotOn(xframe4,RooFit.LineColor(RooFit.kBlue))

# -------------------------------------------
# 5. use combine PDF to fit the toy MC

xframe5 = x.frame(RooFit.Title("5. use combine PDF to fit the toy MC"))

# gauss 6
mean6 = RooRealVar("mean6","mean of gaussian",-1,-10,10)
sigma6 = RooRealVar("sigma6","width of gaussian",4,0.1,10)

gauss6 = RooGaussian("gauss6","gaussian PDF",x,mean6,sigma6)
Esempio n. 12
0
def rf501_simultaneouspdf():
    # C r e a t e   m o d e l   f o r   p h y s i c s   s a m p l e
    # -------------------------------------------------------------

    # Create observables
    x = RooRealVar( "x", "x", -8, 8 ) 

    # Construct signal pdf
    mean = RooRealVar( "mean", "mean", 0, -8, 8 ) 
    sigma = RooRealVar( "sigma", "sigma", 0.3, 0.1, 10 ) 
    gx = RooGaussian( "gx", "gx", x, mean, sigma ) 

    # Construct background pdf
    a0 = RooRealVar( "a0", "a0", -0.1, -1, 1 ) 
    a1 = RooRealVar( "a1", "a1", 0.004, -1, 1 ) 
    px = RooChebychev( "px", "px", x, RooArgList( a0, a1 ) ) 

    # Construct composite pdf
    f = RooRealVar( "f", "f", 0.2, 0., 1. ) 
    model = RooAddPdf( "model", "model", RooArgList( gx, px ), RooArgList( f ) ) 



    # C r e a t e   m o d e l   f o r   c o n t r o l   s a m p l e
    # --------------------------------------------------------------

    # Construct signal pdf. 
    # NOTE that sigma is shared with the signal sample model
    mean_ctl = RooRealVar( "mean_ctl", "mean_ctl", -3, -8, 8 ) 
    gx_ctl = RooGaussian( "gx_ctl", "gx_ctl", x, mean_ctl, sigma ) 

    # Construct the background pdf
    a0_ctl = RooRealVar( "a0_ctl", "a0_ctl", -0.1, -1, 1 ) 
    a1_ctl = RooRealVar( "a1_ctl", "a1_ctl", 0.5, -0.1, 1 ) 
    px_ctl = RooChebychev( "px_ctl", "px_ctl", x, RooArgList( a0_ctl, a1_ctl ) ) 

    # Construct the composite model
    f_ctl = RooRealVar( "f_ctl", "f_ctl", 0.5, 0., 1. ) 
    model_ctl = RooAddPdf( "model_ctl", "model_ctl", RooArgList( gx_ctl, px_ctl ),
                           RooArgList( f_ctl ) ) 
    


    # G e n e r a t e   e v e n t s   f o r   b o t h   s a m p l e s 
    # ---------------------------------------------------------------

    # Generate 1000 events in x and y from model
    data = model.generate( RooArgSet( x ), 100 ) 
    data_ctl = model_ctl.generate( RooArgSet( x ), 2000 ) 



    # C r e a t e   i n d e x   c a t e g o r y   a n d   j o i n   s a m p l e s 
    # ---------------------------------------------------------------------------
    # Define category to distinguish physics and control samples events
    sample = RooCategory( "sample", "sample" ) 
    sample.defineType( "physics" ) 
    sample.defineType( "control" ) 

    # Construct combined dataset in (x,sample)
    combData = RooDataSet( "combData", "combined data", RooArgSet(x), RooFit.Index( sample ),
                          RooFit.Import( "physics", data ),
                          RooFit.Import( "control", data_ctl ) ) 



    # C o n s t r u c t   a   s i m u l t a n e o u s   p d f   i n   ( x , s a m p l e )
    # -----------------------------------------------------------------------------------

    # Construct a simultaneous pdf using category sample as index
    simPdf = RooSimultaneous( "simPdf", "simultaneous pdf", sample ) 

    # Associate model with the physics state and model_ctl with the control state
    simPdf.addPdf( model, "physics" ) 
    simPdf.addPdf( model_ctl, "control" ) 



    # P e r f o r m   a   s i m u l t a n e o u s   f i t
    # ---------------------------------------------------

    # Perform simultaneous fit of model to data and model_ctl to data_ctl
    simPdf.fitTo( combData ) 



    # P l o t   m o d e l   s l i c e s   o n   d a t a    s l i c e s 
    # ----------------------------------------------------------------

    # Make a frame for the physics sample
    frame1 = x.frame( RooFit.Bins( 30 ), RooFit.Title( "Physics sample" ) ) 

    # Plot all data tagged as physics sample
    combData.plotOn( frame1, RooFit.Cut( "sample==sample::physics" ) ) 

    # Plot "physics" slice of simultaneous pdf. 
    # NBL You _must_ project the sample index category with data using ProjWData 
    # as a RooSimultaneous makes no prediction on the shape in the index category 
    # and can thus not be integrated
    simPdf.plotOn( frame1, RooFit.Slice( sample, "physics" ),
                   RooFit.ProjWData( RooArgSet(sample), combData ) ) 
    simPdf.plotOn( frame1, RooFit.Slice( sample, "physics" ),
                   RooFit.Components( "px" ),
                   RooFit.ProjWData( RooArgSet(sample), combData ),
                   RooFit.LineStyle( kDashed ) ) 

    # The same plot for the control sample slice
    frame2 = x.frame( RooFit.Bins( 30 ), RooFit.Title( "Control sample" ) ) 
    combData.plotOn( frame2, RooFit.Cut( "sample==sample::control" ) ) 
    simPdf.plotOn( frame2, RooFit.Slice( sample, "control" ),
                  RooFit.ProjWData( RooArgSet(sample), combData ) ) 
    simPdf.plotOn( frame2, RooFit.Slice( sample, "control" ),
                  RooFit.Components( "px_ctl" ),
                  RooFit.ProjWData( RooArgSet(sample), combData ),
                  RooFit.LineStyle( kDashed ) ) 



    c = TCanvas( "rf501_simultaneouspdf", "rf403_simultaneouspdf", 800, 400 ) 
    c.Divide( 2 ) 
    c.cd( 1 )
    gPad.SetLeftMargin( 0.15 )
    frame1.GetYaxis().SetTitleOffset( 1.4 )
    frame1.Draw() 
    c.cd( 2 )
    gPad.SetLeftMargin( 0.15 )
    frame2.GetYaxis().SetTitleOffset( 1.4 )
    frame2.Draw() 
    raw_input()
sigmean = RooRealVar("sigmean", "B^{#pm} mass", 5.28, 5.20, 5.30)
sigwidth = RooRealVar("sigwidth", "B^{#pm} width", 0.0027, 0.001, 1.)
 
# --- Build Gaussian PDF ---
signal = RooGaussian("signal", "signal PDF", mes, sigmean, sigwidth)

argpar = RooRealVar("argpar", "argus shape parameter", -20.0, -100., -1.)
background = RooArgusBG("background", "Argus PDF", mes, RooFit.RooConst(5.291), argpar)
 
# --- Construct signal+background PDF ---
nsig = RooRealVar("nsig", "#signal events", 200, 0., 10000)
nbkg = RooRealVar("nbkg", "#background events", 800, 0., 10000)
model = RooAddPdf("model", "g+a", RooArgList(signal, background), RooArgList(nsig, nbkg))

# --- Generate a toyMC sample from composite PDF ---
data = model.generate(RooArgSet(mes), 2000)
 
# --- Perform extended ML fit of composite PDF to toy data ---
model.fitTo(data)
 
# --- Plot toy data and composite PDF overlaid ---
mesframe = mes.frame()
data.plotOn(mesframe)
model.plotOn(mesframe)
model.plotOn(mesframe, RooFit.Components('background'), RooFit.LineStyle(kDashed))

mesframe.Draw()

print 'nsig:',nsig.getValV(), '+-', nsig.getError()
print 'nbkg:', nbkg.getValV(), '+-', nbkg.getError()
print 'mes:', mes.getValV(), '+-', mes.getError()
def rooFit501():

    print ">>> setup model for physics sample..."
    x = RooRealVar("x", "x", -8, 8)
    mean = RooRealVar("mean", "mean", 0, -8, 8)
    sigma = RooRealVar("sigma", "sigma", 0.3, 0.1, 10)
    gauss = RooGaussian("gx", "gx", x, mean, sigma)
    a0 = RooRealVar("a0", "a0", -0.1, -1, 1)
    a1 = RooRealVar("a1", "a1", 0.004, -1, 1)
    px = RooChebychev("px", "px", x, RooArgList(a0, a1))
    f = RooRealVar("f", "f", 0.2, 0., 1.)
    model = RooAddPdf("model", "model", RooArgList(gauss, px), RooArgList(f))

    print ">>> setup model for control sample..."
    # NOTE: sigma is shared with the signal sample model
    mean_ctrl = RooRealVar("mean_ctrl", "mean_ctrl", -3, -8, 8)
    gauss_ctrl = RooGaussian("gauss_ctrl", "gauss_ctrl", x, mean_ctrl, sigma)
    a0_ctrl = RooRealVar("a0_ctrl", "a0_ctrl", -0.1, -1, 1)
    a1_ctrl = RooRealVar("a1_ctrl", "a1_ctrl", 0.5, -0.1, 1)
    px_ctrl = RooChebychev("px_ctrl", "px_ctrl", x,
                           RooArgList(a0_ctrl, a1_ctrl))
    f_ctrl = RooRealVar("f_ctrl", "f_ctrl", 0.5, 0., 1.)
    model_ctrl = RooAddPdf("model_ctrl", "model_ctrl",
                           RooArgList(gauss_ctrl, px_ctrl), RooArgList(f_ctrl))

    print ">>> generate events for both samples..."
    data = model.generate(RooArgSet(x), 100)  # RooDataSet
    data_ctrl = model_ctrl.generate(RooArgSet(x), 2000)  # RooDataSet

    print ">>> create index category and join samples..."
    # Define category to distinguish physics and control samples events
    sample = RooCategory("sample", "sample")
    sample.defineType("physics")
    sample.defineType("control")

    print ">>> construct combined dataset in (x,sample)..."
    combData = RooDataSet("combData", "combined data", RooArgSet(x),
                          Index(sample), Import("physics", data),
                          Import("control", data_ctrl))

    print ">>> construct a simultaneous pdf in (x,sample)..."
    # Construct a simultaneous pdf using category sample as index
    simPdf = RooSimultaneous("simPdf", "simultaneous pdf", sample)

    # Associate model with the physics state and model_ctrl with the control state
    simPdf.addPdf(model, "physics")
    simPdf.addPdf(model_ctrl, "control")

    print ">>> perform a simultaneous fit..."
    # Perform simultaneous fit of model to data and model_ctrl to data_ctrl
    simPdf.fitTo(combData)

    print "\n>>> plot model slices on data slices..."
    frame1 = x.frame(Bins(30), Title("Physics sample"))  # RooPlot
    combData.plotOn(frame1, Cut("sample==sample::physics"))

    # Plot "physics" slice of simultaneous pdf.
    # NBL You _must_ project the sample index category with data using ProjWData
    # as a RooSimultaneous makes no prediction on the shape in the index category
    # and can thus not be integrated
    simPdf.plotOn(frame1, Slice(sample, "physics"),
                  ProjWData(RooArgSet(sample), combData))
    simPdf.plotOn(frame1, Slice(sample, "physics"), Components("px"),
                  ProjWData(RooArgSet(sample), combData), LineStyle(kDashed))

    print "\n>>> plot control sample slices..."
    frame2 = x.frame(Bins(30), Title("Control sample"))  # RooPlot
    combData.plotOn(frame2, Cut("sample==sample::control"))
    simPdf.plotOn(frame2, Slice(sample, "control"),
                  ProjWData(RooArgSet(sample), combData))
    simPdf.plotOn(frame2, Slice(sample, "control"), Components("px_ctrl"),
                  ProjWData(RooArgSet(sample), combData), LineStyle(kDashed))

    print "\n>>> draw on canvas..."
    canvas = TCanvas("canvas", "canvas", 100, 100, 1400, 600)
    canvas.Divide(2)
    canvas.cd(1)
    gPad.SetLeftMargin(0.15)
    gPad.SetRightMargin(0.02)
    frame1.GetYaxis().SetLabelOffset(0.008)
    frame1.GetYaxis().SetTitleOffset(1.6)
    frame1.GetYaxis().SetTitleSize(0.045)
    frame1.GetXaxis().SetTitleSize(0.045)
    frame1.Draw()
    canvas.cd(2)
    gPad.SetLeftMargin(0.15)
    gPad.SetRightMargin(0.02)
    frame2.GetYaxis().SetLabelOffset(0.008)
    frame2.GetYaxis().SetTitleOffset(1.6)
    frame2.GetYaxis().SetTitleSize(0.045)
    frame2.GetXaxis().SetTitleSize(0.045)
    frame2.Draw()
    canvas.SaveAs("rooFit501.png")
Esempio n. 15
0
def rooFit204():
    
    print ">>> setup model signal components: gaussians..."
    x      = RooRealVar("x","x",0,10)
    mean   = RooRealVar("mean","mean of gaussian",5)
    sigma1 = RooRealVar("sigma1","width of gaussians",0.5)
    sigma2 = RooRealVar("sigma2","width of gaussians",1)
    sig1   = RooGaussian("sig1","Signal component 1",x,mean,sigma1)
    sig2   = RooGaussian("sig2","Signal component 2",x,mean,sigma2)
    sig1frac = RooRealVar("sig1frac","fraction of component 1 in signal",0.8,0.,1.)
    sig      = RooAddPdf("sig","Signal",RooArgList(sig1,sig2),RooArgList(sig1frac))
    
    print ">>> setup model background components: Chebychev polynomial..."
    a0  = RooRealVar("a0","a0",0.5,0.,1.)
    a1  = RooRealVar("a1","a1",-0.2,0.,1.)
    bkg = RooChebychev("bkg","Background",x,RooArgList(a0,a1))
    
    print ">>> construct extended components with specified range..."
    # Define signal range in which events counts are to be defined
    x.setRange("signalRange",5,6)
    
    # Associated nsig/nbkg as expected number of events with sig/bkg _in_the_range_ "signalRange"
    nsig = RooRealVar("nsig","number of signal events in signalRange",    500,0.,10000) 
    nbkg = RooRealVar("nbkg","number of background events in signalRange",500,0,10000) 
    esig = RooExtendPdf("esig","extended signal pdf",    sig,nsig,"signalRange") 
    ebkg = RooExtendPdf("ebkg","extended background pdf",bkg,nbkg,"signalRange") 
    
    print ">>> sum extended components..."
    # Construct sum of two extended p.d.f. (no coefficients required)
    model = RooAddPdf("model","(g1+g2)+a",RooArgList(ebkg,esig))
    
    print ">>> sample data, fit model..."
    data = model.generate(RooArgSet(x),1000) # RooDataSet 
    result = model.fitTo(data,Extended(kTRUE),Save()) # RooFitResult
    
    print "\n>>> fit result:"
    result.Print()
    
    
    
    print "\n>>> plot everything..."
    frame1 = x.frame(Title("Fitting a sub range")) # RooPlot
    data.plotOn(frame1,Binning(50),Name("data"))
    model.plotOn(frame1,LineColor(kBlue),Name("model"))
    argset1 = RooArgSet(bkg)
    model.plotOn(frame1,Components(argset1),LineStyle(kDashed),LineColor(kBlue),Name("bkg"),Normalization(1.0,RooAbsReal.RelativeExpected))
    #model.plotOn(frame1,Components(argset1),LineStyle(kDashed),LineColor(kRed),Name("bkg2"))
    
    print ">>> draw on canvas..."
    canvas = TCanvas("canvas","canvas",100,100,800,600)
    legend = TLegend(0.2,0.85,0.4,0.7)
    legend.SetTextSize(0.032)
    legend.SetBorderSize(0)
    legend.SetFillStyle(0)
    gPad.SetLeftMargin(0.14); gPad.SetRightMargin(0.02)
    frame1.GetYaxis().SetLabelOffset(0.008)
    frame1.GetYaxis().SetTitleOffset(1.4)
    frame1.GetYaxis().SetTitleSize(0.045)
    frame1.GetXaxis().SetTitleSize(0.045)
    frame1.Draw()
    legend.AddEntry("data",   "data",           'LEP')
    legend.AddEntry("model",  "fit",            'L')
    legend.AddEntry("bkg",    "background only",'L')
    #legend.AddEntry("bkg2",   "background only (no extended norm)",'L')
    legend.Draw()
    canvas.SaveAs("rooFit204.png")
Esempio n. 16
0
class BaseFitter(object):
    def __init__(self, histos, uncertainty):
        self._histos = histos
        self._uncertainty = uncertainty
        self._make_underlying_model()
        self._make_dataset()
        self._make_fit_model()
        self._fit()

    def _make_underlying_model(self):
        self.pdfs = {}
        self.yields = {}  # yields are plain floats
        self.ryields = {}  # keep track of roofit objects for memory management
        h = self._histos[0]
        nbins = h.GetXaxis().GetNbins()
        xmin = h.GetXaxis().GetXmin()
        xmax = h.GetXaxis().GetXmax()
        self.xvar = RooRealVar("x", "x", xmin, xmax)
        self.xvar.setBins(nbins)
        self.pdfs = {}
        self.hists = []
        pdfs = RooArgList()
        yields = RooArgList()
        for histo in self._histos:
            if histo.Integral() == 0:
                continue
            compname = histo.GetName()
            hist = RooDataHist(compname, compname, RooArgList(self.xvar),
                               histo)
            SetOwnership(hist, False)
            # self.hists.append(hist)
            pdf = RooHistPdf(compname, compname, RooArgSet(self.xvar), hist)
            self.pdfs[compname] = pdf
            # self.pdfs[compname].Print()
            pdfs.add(pdf)
            nevts = histo.Integral()
            uncertainty = self._uncertainty
            nmin = min(0, nevts * (1 - uncertainty))
            nmax = nevts * (1 + uncertainty)
            theyield = RooRealVar('n{}'.format(compname),
                                  'n{}'.format(compname), nevts, nmin, nmax)
            self.ryields[compname] = theyield
            self.yields[compname] = nevts
            yields.add(theyield)

        self.underlying_model = RooAddPdf('model', 'model', pdfs, yields)

    def _make_fit_model(self):
        pass

    def _make_dataset(self):
        nevents = sum(self.yields.values())
        self.data = self.underlying_model.generate(RooArgSet(self.xvar),
                                                   nevents)

    def _fit(self):
        self.tresult = self.underlying_model.fitTo(self.data,
                                                   RooFit.Extended(),
                                                   RooFit.Save(),
                                                   RooFit.PrintEvalErrors(-1))

    def print_result(self, comps):
        print 'input background uncertainty:', self._uncertainty
        self.tresult.Print()
        signal_percent_unc = None
        for comp in comps:
            yzh = self.ryields[comp]
            zh_val = yzh.getVal()
            zh_err = yzh.getError()
            percent_unc = zh_err / zh_val * 100.
            if signal_percent_unc is None:
                signal_percent_unc = percent_unc
            print '{} yield  = {:8.2f}'.format(comp, zh_val)
            print '{} uncert = {:8.2f}%'.format(comp, percent_unc)
            print '{} abs uncert = {:8.2f}'.format(comp, zh_err)
        return percent_unc

    def draw_pdfs(self):
        self.pframe = self.xvar.frame()
        for pdf in self.pdfs.values():
            pdf.plotOn(self.pframe)
        self.pframe.Draw()

    def draw_data(self):
        self.mframe = self.xvar.frame()
        self.data.plotOn(self.mframe)
        self.underlying_model.plotOn(self.mframe)
        for icomp, compname in enumerate(self.pdfs):
            self.underlying_model.plotOn(self.mframe,
                                         RooFit.Components(compname),
                                         RooFit.LineColor(icomp + 1))
        self.mframe.Draw()
        gPad.Modified()
        gPad.Update()
Esempio n. 17
0
frac_combine4   = RooRealVar("frac_combine4",   "fraction of gauss4 wrt gauss5", 0.7, 0.,   1.)

pdf_combine4 = RooAddPdf("pdf_combine4"," gauss12 + gauss13 ", RooArgList(gauss12 , gauss13 ), RooArgList( frac_combine4 ))


#pdf_combine1.plotOn(xframe3,RooFit.LineColor(RooFit.kBlue))

#gauss12.plotOn(xframe8, RooFit.Normalization( frac_combine4.getVal()   ,RooAbsReal.Relative),RooFit.LineColor(RooFit.kOrange))
#gauss13.plotOn(xframe8, RooFit.Normalization( 1-frac_combine4.getVal() ,RooAbsReal.Relative),RooFit.LineColor(RooFit.kCyan))
#pdf_combine2.plotOn(xframe8, RooFit.Normalization(1.0,RooAbsReal.Relative) ,RooFit.LineColor(RooFit.kBlue))

# generate toy MC

n_generate = 5000

data4 = pdf_combine4.generate(RooArgSet(x), n_generate )


data4_SB  = RooDataSet("data4_SB", "data4 SB", RooArgSet(x), RooFit.Import( data4 ), RooFit.Cut("x<-2||x>4") )

# set range

x.setRange("signal_region",-2 ,4 )
x.setRange("left_side_band_region",-10 ,-2 )
x.setRange("right_side_band_region",4 ,10 )

# use gauss 13 to generate toy MC and cut right SB 

data_gauss13 = gauss13.generate(RooArgSet(x), n_generate*(1-frac_combine4.getVal() ) )
#data_gauss13.plotOn(xframe8,RooFit.LineColor(RooFit.kGreen))
data_gauss13_right_SB = RooDataSet("data_gauss13_right_SB", "data_gauss13_right_SB", RooArgSet(x), RooFit.Import( data_gauss13 ), RooFit.Cut("x>4") )
Esempio n. 18
0
from ROOT import RooRealVar, RooGaussian, RooChebychev, RooAddPdf, RooArgList, RooArgSet, RooFit

x = RooRealVar("x", "x", -1, 1)

# Use RooGaussian in the generation
mean = RooRealVar("mean", "mean of gaussian", 0, -1, 1)
sigma = RooRealVar("sigma", "sigma of gaussian", 0.1, -1, 1)
sig = RooGaussian("gauss", "gaussian PDF", x, mean, sigma)

# Background
a0 = RooRealVar("a0", "a0", 0.5, 0., 1.)
a1 = RooRealVar("a1", "a1", -0.2, 0., 1.)
bkg = RooChebychev("bkg", "Background", x, RooArgList(a0, a1))

bkgfrac = RooRealVar("bkgfrac", "fraction of background", 0.5, 0., 1.)
model = RooAddPdf("model", "g+a", RooArgList(bkg, sig), RooArgList(bkgfrac))

data = model.generate(RooArgSet(x), 10000)

model.fitTo(data)
Esempio n. 19
0
vybgd = RooRealVar("nbgd", 'nbgd', ybgd, 0, 100000)
model = RooAddPdf("model", "model", RooArgList(sig, bgd),
                  RooArgList(vysig, vybgd))
#model = RooAddPdf("model", "model", RooArgList(sig),
#                  RooArgList(vysig))

c2 = TCanvas()

mframe = x.frame()
model.plotOn(mframe)
mframe.Draw()


######### data generation, parametrized

data = model.generate(RooArgSet(x), ndata)

######### fit, parametrized

result = model.fitTo(data, RooFit.Extended(), RooFit.Save())

cf_p = TCanvas('cf_p', 'fit, parametrized')

cf_p_frame = x.frame()
data.plotOn(cf_p_frame)
model.plotOn(cf_p_frame)
cf_p_frame.Draw()


######### model definition, template  
# --- Build Gaussian PDF ---
signal = RooGaussian("signal", "signal PDF", mes, sigmean, sigwidth)

argpar = RooRealVar("argpar", "argus shape parameter", -20.0, -100., -1.)
background = RooArgusBG("background", "Argus PDF", mes, RooFit.RooConst(5.291),
                        argpar)

# --- Construct signal+background PDF ---
nsig = RooRealVar("nsig", "#signal events", 200, 0., 10000)
nbkg = RooRealVar("nbkg", "#background events", 800, 0., 10000)
model = RooAddPdf("model", "g+a", RooArgList(signal, background),
                  RooArgList(nsig, nbkg))

# --- Generate a toyMC sample from composite PDF ---
data = model.generate(RooArgSet(mes), 2000)

# --- Perform extended ML fit of composite PDF to toy data ---
model.fitTo(data)

# --- Plot toy data and composite PDF overlaid ---
mesframe = mes.frame()
data.plotOn(mesframe)
model.plotOn(mesframe)
model.plotOn(mesframe, RooFit.Components('background'),
             RooFit.LineStyle(kDashed))

mesframe.Draw()

print 'nsig:', nsig.getValV(), '+-', nsig.getError()
print 'nbkg:', nbkg.getValV(), '+-', nbkg.getError()
Esempio n. 21
0
def rooFit205():

    print ">>> setup model signal components: gaussians..."
    x = RooRealVar("x", "x", 0, 10)
    mean = RooRealVar("mean", "mean of gaussians", 5)
    sigma1 = RooRealVar("sigma1", "width of gaussians", 0.5)
    sigma2 = RooRealVar("sigma2", "width of gaussians", 1)
    sig1 = RooGaussian("sig1", "Signal component 1", x, mean, sigma1)
    sig2 = RooGaussian("sig2", "Signal component 2", x, mean, sigma2)
    sig1frac = RooRealVar("sig1frac", "fraction of component 1 in signal", 0.8,
                          0., 1.)
    sig = RooAddPdf("sig", "Signal", RooArgList(sig1, sig2),
                    RooArgList(sig1frac))

    print ">>> setup model background components: Chebychev polynomial plus exponential..."
    a0 = RooRealVar("a0", "a0", 0.5, 0., 1.)
    a1 = RooRealVar("a1", "a1", -0.2, 0., 1.)
    bkg1 = RooChebychev("bkg1", "Background 1", x, RooArgList(a0, a1))
    alpha = RooRealVar("alpha", "alpha", -1)
    bkg2 = RooExponential("bkg2", "Background 2", x, alpha)
    bkg1frac = RooRealVar("bkg1frac", "fraction of component 1 in background",
                          0.2, 0., 1.)
    bkg = RooAddPdf("bkg", "Signal", RooArgList(bkg1, bkg2),
                    RooArgList(bkg1frac))

    print ">>> sum signal and background component..."
    bkgfrac = RooRealVar("bkgfrac", "fraction of background", 0.5, 0., 1.)
    model = RooAddPdf("model", "g1+g2+a", RooArgList(bkg, sig),
                      RooArgList(bkgfrac))

    print ">>> setup basic plot with data and full pdf..."
    data = model.generate(RooArgSet(x), 1000)  # RooDataSet
    frame1 = x.frame(
        Title("Component plotting of pdf=(sig1+sig2)+(bkg1+bkg2)"))  # RooPlot
    data.plotOn(frame1, Name("data"))
    model.plotOn(frame1, Name("model"))

    print ">>> clone frame for reuse..."
    frame2 = frame1.Clone("frame2")  # RooPlot
    frame2.SetTitle("Get components with regular expressions")

    print ">>> make omponent by object reference..."
    # Plot multiple background components specified by object reference
    # Note that specified components may occur at any level in object tree
    # (e.g bkg is component of 'model' and 'sig2' is component 'sig')
    argset1 = RooArgSet(bkg)
    argset2 = RooArgSet(bkg2)
    argset3 = RooArgSet(bkg, sig2)
    model.plotOn(frame1, Components(argset1), LineColor(kRed), Name("bkgs1"))
    model.plotOn(frame1, Components(argset2), LineStyle(kDashed),
                 LineColor(kRed), Name("bkg2"))
    model.plotOn(frame1, Components(argset3), LineStyle(kDotted),
                 Name("bkgssig21"))

    print "\n>>> make component by name / regular expressions..."
    model.plotOn(frame2, Components("bkg"), LineColor(kAzure - 4),
                 Name("bkgs2"))  # by name
    model.plotOn(frame2, Components("bkg1,sig2"), LineColor(kAzure - 4),
                 LineStyle(kDotted), Name("bkg1sig22"))  # by name
    model.plotOn(frame2, Components("sig*"), LineColor(kAzure - 4),
                 LineStyle(kDashed), Name("sigs2"))  # with regexp (wildcard *)
    model.plotOn(frame2, Components("bkg1,sig*"), LineColor(kYellow),
                 LineStyle(kDashed),
                 Name("bkg1sigs2"))  # with regexp (,) #Invisible()

    print "\n>>> draw pfds and fits on canvas..."
    canvas = TCanvas("canvas", "canvas", 100, 100, 1400, 600)
    legend1 = TLegend(0.22, 0.85, 0.4, 0.65)
    legend2 = TLegend(0.22, 0.85, 0.4, 0.65)
    for legend in [legend1, legend2]:
        legend.SetTextSize(0.032)
        legend.SetBorderSize(0)
        legend.SetFillStyle(0)
    canvas.Divide(2)
    canvas.cd(1)
    gPad.SetLeftMargin(0.15)
    gPad.SetRightMargin(0.02)
    frame1.GetYaxis().SetLabelOffset(0.008)
    frame1.GetYaxis().SetTitleOffset(1.6)
    frame1.GetYaxis().SetTitleSize(0.045)
    frame1.GetXaxis().SetTitleSize(0.045)
    frame1.Draw()
    legend1.AddEntry("data", "data", 'LEP')
    legend1.AddEntry("model", "model", 'L')
    legend1.AddEntry("bkgs1", "bkg", 'L')
    legend1.AddEntry("bkg2", "bkg2", 'L')
    legend1.AddEntry("bkgssig21", "bkg,sig2", 'L')
    legend1.Draw()
    canvas.cd(2)
    gPad.SetLeftMargin(0.15)
    gPad.SetRightMargin(0.02)
    frame2.GetYaxis().SetLabelOffset(0.008)
    frame2.GetYaxis().SetTitleOffset(1.6)
    frame2.GetYaxis().SetTitleSize(0.045)
    frame2.GetXaxis().SetTitleSize(0.045)
    frame2.Draw()
    legend2.AddEntry("data", "data", 'LEP')
    legend2.AddEntry("model", "model", 'L')
    legend2.AddEntry("bkgs2", "\"bkg\"", 'L')
    legend2.AddEntry("bkg1sig22", "\"bkg1,sig2\"", 'L')
    legend2.AddEntry("sigs2", "\"sig*\"", 'L')
    legend2.AddEntry("bkg1sigs2", "\"bkg1,sig*\"", 'L')
    legend2.Draw()
    canvas.SaveAs("rooFit205.png")
def rooFit206():

    print ">>> setup model signal components: gaussians..."
    x = RooRealVar("x", "x", 0, 10)
    mean = RooRealVar("mean", "mean of gaussians", 5)
    sigma1 = RooRealVar("sigma1", "width of gaussians", 0.5)
    sigma2 = RooRealVar("sigma2", "width of gaussians", 1)
    sig1 = RooGaussian("sig1", "Signal component 1", x, mean, sigma1)
    sig2 = RooGaussian("sig2", "Signal component 2", x, mean, sigma2)
    sig1frac = RooRealVar("sig1frac", "fraction of component 1 in signal", 0.8,
                          0., 1.)
    sig = RooAddPdf("sig", "Signal", RooArgList(sig1, sig2),
                    RooArgList(sig1frac))

    print ">>> setup model background components: Chebychev polynomial plus exponential..."
    a0 = RooRealVar("a0", "a0", 0.5, 0., 1.)
    a1 = RooRealVar("a1", "a1", -0.2, 0., 1.)
    bkg1 = RooChebychev("bkg1", "Background 1", x, RooArgList(a0, a1))
    alpha = RooRealVar("alpha", "alpha", -1)
    bkg2 = RooExponential("bkg2", "Background 2", x, alpha)
    bkg1frac = RooRealVar("bkg1frac", "fraction of component 1 in background",
                          0.2, 0., 1.)
    bkg = RooAddPdf("bkg", "Signal", RooArgList(bkg1, bkg2),
                    RooArgList(bkg1frac))

    print ">>> sum signal and background component..."
    bkgfrac = RooRealVar("bkgfrac", "fraction of background", 0.5, 0., 1.)
    model = RooAddPdf("model", "g1+g2+a", RooArgList(bkg, sig),
                      RooArgList(bkgfrac))

    print ">>> composite tree in ASCII:"
    model.Print("t")

    print "\n>>> write to txt file"
    model.printCompactTree("", "rooFit206_asciitree.txt")

    print ">>> draw composite tree graphically..."
    # Print GraphViz DOT file with representation of tree
    model.graphVizTree("rooFit206_model.dot")

    # Make graphic output file with one of the GraphViz tools
    # (freely available from www.graphviz.org)
    #
    # 'Top-to-bottom graph'
    # unix> dot -Tgif -o rooFit206_model_dot.gif rooFit206_model.dot
    #
    # 'Spring-model graph'
    # unix> fdp -Tgif -o rooFit206_model_fdp.gif rooFit206_model.dot

    print ">>> plot everything..."
    data = model.generate(RooArgSet(x), 1000)  # RooDataSet
    frame1 = x.frame(
        Title("Component plotting of pdf=(sig1+sig2)+(bkg1+bkg2)"))  # RooPlot
    data.plotOn(frame1, Name("data"), Binning(40))
    model.plotOn(frame1, Name("model"))
    argset1 = RooArgSet(bkg)
    argset2 = RooArgSet(bkg2)
    argset3 = RooArgSet(bkg, sig2)
    model.plotOn(frame1, Components(argset1), LineColor(kRed), Name("bkg"))
    model.plotOn(frame1, Components(argset2), LineStyle(kDashed),
                 LineColor(kRed), Name("bkg2"))
    model.plotOn(frame1, Components(argset3), LineStyle(kDotted),
                 Name("bkgsig2"))

    print "\n>>> draw pfds and fits on canvas..."
    canvas = TCanvas("canvas", "canvas", 100, 100, 800, 600)
    legend = TLegend(0.22, 0.85, 0.4, 0.65)
    legend.SetTextSize(0.032)
    legend.SetBorderSize(0)
    legend.SetFillStyle(0)
    gPad.SetLeftMargin(0.15)
    gPad.SetRightMargin(0.02)
    frame1.GetYaxis().SetLabelOffset(0.008)
    frame1.GetYaxis().SetTitleOffset(1.6)
    frame1.GetYaxis().SetTitleSize(0.045)
    frame1.GetXaxis().SetTitleSize(0.045)
    frame1.Draw()
    legend.AddEntry("data", "data", 'LEP')
    legend.AddEntry("model", "model", 'L')
    legend.AddEntry("bkg", "bkg", 'L')
    legend.AddEntry("bkg2", "bkg2", 'L')
    legend.AddEntry("bkgsig2", "bkg,sig2", 'L')
    legend.Draw()
    canvas.SaveAs("rooFit206.png")
Esempio n. 23
0
def rooFit201():
    
    print ">>> setup model component: gaussian signals and Chebychev polynomial background..."
    x      = RooRealVar("x","x",0,11)
    mean   = RooRealVar("mean","mean of gaussians",5)
    sigma1 = RooRealVar("sigma1","width of gaussians",0.5)
    sigma2 = RooRealVar("sigma2","width of gaussians",1)
    sig1   = RooGaussian("sig1","Signal component 1",x,mean,sigma1)
    sig2   = RooGaussian("sig2","Signal component 2",x,mean,sigma2)
    
    a0  = RooRealVar("a0","a0",0.5,0.,1.)
    a1  = RooRealVar("a1","a1",-0.2,0.,1.)
    bkg = RooChebychev("bkg","Background",x,RooArgList(a0,a1))
    
    
    
    print ">>>\n>>> METHOD 1 - Two RooAddPdfs"
    print ">>> add signal components..."
    # Sum the signal components into a composite signal p.d.f.
    sig1frac = RooRealVar("sig1frac","fraction of component 1 in signal",0.8,0.,1.)
    sig      = RooAddPdf("sig","Signal",RooArgList(sig1,sig2),RooArgList(sig1frac))
    
    print ">>> add signal and background..."
    # Sum the composite signal and background
    bkgfrac = RooRealVar("bkgfrac","fraction of background",0.5,0.,1.)
    model   = RooAddPdf("model","g1+g2+a",RooArgList(bkg,sig),RooArgList(bkgfrac))
    
    print ">>> sample, fit and plot model..."
    data = model.generate(RooArgSet(x),1000) # RooDataSet
    model.fitTo(data)
    frame1 = x.frame(Title("Example of composite pdf=(sig1+sig2)+bkg")) # RooPlot
    data.plotOn(frame1,Binning(50),Name("data"))
    model.plotOn(frame1,Name("model"))
    
    # Overlay the background component of model with a dashed line
    argset1 = RooArgSet(bkg)
    model.plotOn(frame1,Components(argset1),LineWidth(2),Name("bkg")) #,LineStyle(kDashed)
    
    # Overlay the background+sig2 components of model with a dotted line
    argset2 = RooArgSet(bkg,sig2)
    model.plotOn(frame1,Components(argset2),LineWidth(2),LineStyle(kDashed),LineColor(kAzure-4),Name("bkgsig2")) #,LineStyle(kDotted)
    
    print "\n>>> structure of composite pdf:"
    model.Print("t") # "tree" mode
    
    print "\n>>> parameters:"
    params = model.getVariables() # RooArgSet
    params.Print("v")
    params.Print()
    
    print "\n>>> params.find(\"...\").getVal():"
    print ">>>   sigma1  = %.2f"   % params.find("sigma1").getVal()
    print ">>>   sigma2  = %.2f"   % params.find("sigma2").getVal()
    print ">>>   bkgfrac = %5.2f"  % params.find("bkgfrac").getVal()
    print ">>>   sig1frac = %5.2f" % params.find("sig1frac").getVal()
    
    print ">>>\n>>> components:"
    comps   = model.getComponents() # RooArgSet
    sig     = comps.find("sig")     # RooAbsArg
    sigVars = sig.getVariables()    # RooArgSet
    sigVars.Print()
    
    
    
    print ">>>\n>>> METHOD 2 - One RooAddPdf with recursive fractions"
    # Construct sum of models on one go using recursive fraction interpretations
    #   model2 = bkg + (sig1 + sig2)
    model2 = RooAddPdf("model","g1+g2+a",RooArgList(bkg,sig1,sig2),RooArgList(bkgfrac,sig1frac),kTRUE)
    
    # NB: Each coefficient is interpreted as the fraction of the
    # left-hand component of the i-th recursive sum, i.e.
    #   sum4 = A + ( B + ( C + D ) )
    # with fraction fA, fB and fC expands to
    #   sum4 = fA*A + (1-fA)*(fB*B + (1-fB)*(fC*C + (1-fC)*D))
    
    print ">>> plot recursive addition model..."
    argset3 = RooArgSet(bkg,sig2)
    model2.plotOn(frame1,LineColor(kRed),LineStyle(kDashDotted),LineWidth(3),Name("model2"))
    model2.plotOn(frame1,Components(argset3),LineColor(kMagenta),LineStyle(kDashDotted),LineWidth(3),Name("bkgsig22"))
    model2.Print("t")
    
    
    
    print ">>> draw pdfs and fits on canvas..."
    canvas = TCanvas("canvas","canvas",100,100,800,600)
    legend = TLegend(0.57,0.87,0.95,0.65)
    legend.SetTextSize(0.030)
    legend.SetBorderSize(0)
    legend.SetFillStyle(0)
    gPad.SetLeftMargin(0.14); gPad.SetRightMargin(0.02)
    frame1.GetYaxis().SetLabelOffset(0.008)
    frame1.GetYaxis().SetTitleOffset(1.4)
    frame1.GetYaxis().SetTitleSize(0.045)
    frame1.GetXaxis().SetTitleSize(0.045)
    frame1.Draw()
    legend.AddEntry("data",    "data",                            'LEP')
    legend.AddEntry("model",   "composite model",                 'L')
    legend.AddEntry("model2",  "composite model (method 2)",      'L')
    legend.AddEntry("bkg",     "background only",                 'L')
    legend.AddEntry("bkgsig2", "background + signal 2",           'L')
    legend.AddEntry("bkgsig22","background + signal 2 (method 2)",'L')
    legend.Draw()
    canvas.SaveAs("rooFit201.png")
def rooFit202():

    print ">>> setup model component: gaussian signals and Chebychev polynomial background..."
    x = RooRealVar("x", "x", 0, 10)
    mean = RooRealVar("mean", "mean of gaussian", 5)
    sigma1 = RooRealVar("sigma1", "width of gaussian", 0.5)
    sigma2 = RooRealVar("sigma2", "width of gaussian", 1.0)
    sig1 = RooGaussian("sig1", "Signal component 1", x, mean, sigma1)
    sig2 = RooGaussian("sig2", "Signal component 2", x, mean, sigma2)

    a0 = RooRealVar("a0", "a0", 0.5, 0., 1.)
    a1 = RooRealVar("a1", "a1", -0.2, 0., 1.)
    bkg = RooChebychev("bkg", "Background", x, RooArgList(a0, a1))

    # Sum the signal components into a composite signal p.d.f.
    sig1frac = RooRealVar("sig1frac", "fraction of component 1 in signal", 0.8,
                          0., 1.)
    sig = RooAddPdf("sig", "Signal", RooArgList(sig1, sig2),
                    RooArgList(sig1frac))

    print ">>>\n>>> METHOD 1"
    print ">>> construct extended composite model..."
    # Sum the composite signal and background into an extended pdf nsig*sig+nbkg*bkg
    nsig = RooRealVar("nsig", "number of signal events", 500, 0., 10000)
    nbkg = RooRealVar("nbkg", "number of background events", 500, 0, 10000)
    model = RooAddPdf("model", "(g1+g2)+a", RooArgList(bkg, sig),
                      RooArgList(nbkg, nsig))

    print ">>> sample, fit and plot extended model...\n"
    # Generate a data sample of expected number events in x from model
    #   nsig + nbkg = model.expectedEvents()
    # NOTE: since the model predicts a specific number events, one can
    #       omit the requested number of events to be generated
    # Introduce Poisson fluctuation with Extended(kTRUE)
    data = model.generate(RooArgSet(x), Extended(kTRUE))  # RooDataSet

    # Fit model to data, extended ML term automatically included
    # NOTE: Composite extended pdfs can only be successfully fit if the extended likelihood
    #       term -log(Poisson(Nobs,Nexp)) is included in the minimization because they have
    #       one extra degree of freedom in their parameterization that is constrained by
    #       this extended term. If a pdf is capable of calculating an extended term (i.e.
    #       any extended RooAddPdf), the extended term is AUTOMATICALLY included in the
    #       likelihood calculation. Override this behaviour with Extended():
    #           Extended(kTRUE)  ADD extended likelihood term
    #           Extended(kFALSE) DO NOT ADD extended likelihood
    #model.fitTo(data,Extended(kTRUE))
    model.fitTo(data)

    print "\n>>> plot data, model and model components..."
    # Plot data and PDF overlaid, use expected number of events for pdf projection
    # normalization, rather than observed number of events, data.numEntries()
    frame1 = x.frame(Title("extended ML fit example"))  # RooPlot
    data.plotOn(frame1, Binning(30), Name("data"))
    model.plotOn(frame1, Normalization(1.0, RooAbsReal.RelativeExpected),
                 Name("model"))

    # Overlay the background components of model
    # NOTE: By default, the pdf is normalized to event count of the last dataset added
    #       to the plot frame. Use "RelativeExpected" to normalize to the expected
    #       event count of the pdf instead
    argset1 = RooArgSet(bkg)
    argset2 = RooArgSet(sig1)
    argset3 = RooArgSet(sig2)
    argset4 = RooArgSet(bkg, sig2)
    model.plotOn(frame1, Components(argset1), LineStyle(kDashed),
                 LineColor(kBlue),
                 Normalization(1.0, RooAbsReal.RelativeExpected), Name("bkg"))
    #model.plotOn(frame1,Components(argset1),LineStyle(kDashed),LineColor(kBlue),  Name("bkg2"))
    model.plotOn(frame1, Components(argset2), LineStyle(kDotted),
                 LineColor(kMagenta),
                 Normalization(1.0, RooAbsReal.RelativeExpected), Name("sig1"))
    model.plotOn(frame1, Components(argset3), LineStyle(kDotted),
                 LineColor(kPink),
                 Normalization(1.0, RooAbsReal.RelativeExpected), Name("sig2"))
    model.plotOn(frame1, Components(argset4), LineStyle(kDashed),
                 LineColor(kAzure - 4),
                 Normalization(1.0, RooAbsReal.RelativeExpected),
                 Name("bkgsig2"))

    print "\n>>> structure of composite pdf:"
    model.Print("t")  # "tree" mode

    print "\n>>> parameters:"
    params = model.getVariables()  # RooArgSet
    params.Print("v")
    params.Print()

    print "\n>>> params.find(\"...\").getVal():"
    print ">>>   sigma1   = %.2f" % params.find("sigma1").getVal()
    print ">>>   sigma2   = %.2f" % params.find("sigma2").getVal()
    print ">>>   nsig     = %6.2f,  sig1frac = %5.2f" % (
        params.find("nsig").getVal(), params.find("sig1frac").getVal())
    print ">>>   nbkg     = %6.2f" % params.find("nbkg").getVal()

    print ">>>\n>>> components:"
    comps = model.getComponents()  # RooArgSet
    sig = comps.find("sig")  # RooAbsArg
    sigVars = sig.getVariables()  # RooArgSet
    sigVars.Print()

    print ">>>\n>>> METHOD 2"
    print ">>> construct extended components first..."
    # Associated nsig/nbkg as expected number of events with sig/bkg
    nsig = RooRealVar("nsig", "number of signal events", 500, 0., 10000)
    nbkg = RooRealVar("nbkg", "number of background events", 500, 0, 10000)
    esig = RooExtendPdf("esig", "extended signal pdf", sig, nsig)
    ebkg = RooExtendPdf("ebkg", "extended background pdf", bkg, nbkg)

    print ">>> sum extended components without coefficients..."
    # Construct sum of two extended p.d.f. (no coefficients required)
    model2 = RooAddPdf("model2", "(g1+g2)+a", RooArgList(ebkg, esig))

    # METHOD 2 is functionally completely equivalent to METHOD 1.
    # Its advantage is that the yield parameter is associated to the shape pdf
    # directly, while in METHOD 1 the association is made after constructing
    # a RooAddPdf. Also, class RooExtendPdf offers extra functionality to
    # interpret event counts in a different range.

    print ">>> plot model..."
    model2.plotOn(frame1, LineStyle(kDashed), LineColor(kRed),
                  Normalization(1.0, RooAbsReal.RelativeExpected),
                  Name("model2"))

    print ">>> draw on canvas..."
    canvas = TCanvas("canvas", "canvas", 100, 100, 800, 600)
    legend = TLegend(0.2, 0.85, 0.4, 0.65)
    legend.SetTextSize(0.032)
    legend.SetBorderSize(0)
    legend.SetFillStyle(0)
    gPad.SetLeftMargin(0.14)
    gPad.SetRightMargin(0.02)
    frame1.GetYaxis().SetLabelOffset(0.008)
    frame1.GetYaxis().SetTitleOffset(1.4)
    frame1.GetYaxis().SetTitleSize(0.045)
    frame1.GetXaxis().SetTitleSize(0.045)
    frame1.Draw()
    legend.AddEntry("data", "data", 'LEP')
    legend.AddEntry("model", "composite model", 'L')
    legend.AddEntry("model2", "composite model (method 2)", 'L')
    legend.AddEntry("bkg", "background only", 'L')
    #legend.AddEntry("bkg2",   "background only (no extended norm)", 'L')
    legend.AddEntry("sig1", "signal 1", 'L')
    legend.AddEntry("sig2", "signal 2", 'L')
    legend.AddEntry("bkgsig2", "background + signal 2", 'L')
    legend.Draw()
    canvas.SaveAs("rooFit202.png")
Esempio n. 25
0
class BaseFitter(object):
    def __init__(self, plot):
        assert (isinstance(plot, DataMCPlot))
        self.plot = plot
        self._make_underlying_model()
        self._make_dataset()
        self._make_fit_model()
        self._fit()

    def _make_underlying_model(self):
        self.pdfs = {}
        self.yields = {}  # yields are plain floats
        self.ryields = {}  # keep track of roofit objects for memory management
        nbins, xmin, xmax = self.plot.histos[0].GetBinning()
        self.xvar = RooRealVar("x", "x", xmin, xmax)
        self.xvar.setBins(nbins)
        self.pdfs = {}
        self.hists = []
        pdfs = RooArgList()
        yields = RooArgList()
        for compname, comp in self.plot.histosDict.iteritems():
            print compname
            assert (isinstance(comp, Histogram))
            hist = RooDataHist(compname, compname, RooArgList(self.xvar),
                               comp.weighted)
            SetOwnership(hist, False)
            # self.hists.append(hist)
            pdf = RooHistPdf(compname, compname, RooArgSet(self.xvar), hist)
            self.pdfs[compname] = pdf
            self.pdfs[compname].Print()
            pdfs.add(pdf)

            nevts = comp.Integral(xmin=xmin, xmax=xmax)
            theyield = RooRealVar('n{}'.format(compname),
                                  'n{}'.format(compname), nevts, 0, 50000)
            self.ryields[compname] = theyield
            self.yields[compname] = nevts
            yields.add(theyield)

        self.underlying_model = RooAddPdf('model', 'model', pdfs, yields)

    def _make_fit_model(self):
        pass

    def _make_dataset(self):
        nevents = sum(self.yields.values())
        self.data = self.underlying_model.generate(RooArgSet(self.xvar),
                                                   nevents)

    def _fit(self):
        self.tresult = self.underlying_model.fitTo(self.data,
                                                   RooFit.Extended(),
                                                   RooFit.Save())
        self.tresult.Print()

    def draw_pdfs(self):
        self.pframe = self.xvar.frame()
        for pdf in self.pdfs.values():
            pdf.plotOn(self.pframe)
        self.pframe.Draw()

    def draw_data(self):
        self.canvas_data = TCanvas()
        self.mframe = self.xvar.frame()
        self.data.plotOn(self.mframe)
        self.underlying_model.plotOn(self.mframe)
        for icomp, compname in enumerate(self.pdfs):
            self.underlying_model.plotOn(self.mframe,
                                         RooFit.Components(compname),
                                         RooFit.LineColor(icomp + 1))
        self.mframe.Draw()
Esempio n. 26
0
class BaseFitter(object):
    def __init__(self, plot):
        assert (isinstance(plot, DataMCPlot))
        self.plot = plot
        self._make_underlying_model()
        self._make_dataset()
        self._make_fit_model()
        self._fit()

    def _make_underlying_model(self):
        self.pdfs = {}
        self.yields = {}  # yields are plain floats
        self.ryields = {}  # keep track of roofit objects for memory management
        nbins, xmin, xmax = self.plot.histos[0].GetBinning()
        self.xvar = RooRealVar("x", "x", xmin, xmax)
        self.xvar.setBins(nbins)
        self.pdfs = {}
        self.hists = []
        pdfs = RooArgList()
        yields = RooArgList()
        for compname, comp in self.plot.histosDict.iteritems():
            if comp.weighted.Integral() == 0:
                continue
            assert (isinstance(comp, Histogram))
            hist = RooDataHist(compname, compname, RooArgList(self.xvar),
                               comp.weighted)
            SetOwnership(hist, False)
            # self.hists.append(hist)
            pdf = RooHistPdf(compname, compname, RooArgSet(self.xvar), hist)
            self.pdfs[compname] = pdf
            # self.pdfs[compname].Print()
            pdfs.add(pdf)
            nevts = comp.Integral(xmin=xmin, xmax=xmax)
            nmin = min(0, nevts * (1 - comp.uncertainty))
            nmax = nevts * (1 + comp.uncertainty)
            theyield = RooRealVar('n{}'.format(compname),
                                  'n{}'.format(compname), nevts, nmin, nmax)
            self.ryields[compname] = theyield
            self.yields[compname] = nevts
            yields.add(theyield)

        self.underlying_model = RooAddPdf('model', 'model', pdfs, yields)

    def _make_fit_model(self):
        pass

    def _make_dataset(self):
        nevents = sum(self.yields.values())
        self.data = self.underlying_model.generate(RooArgSet(self.xvar),
                                                   nevents)

    def _fit(self):
        self.tresult = self.underlying_model.fitTo(self.data,
                                                   RooFit.Extended(),
                                                   RooFit.Save(),
                                                   RooFit.PrintEvalErrors(-1))

    def print_result(self):
        self.tresult.Print()
        yzh = self.ryields['ZH']
        zh_val = yzh.getVal()
        zh_err = yzh.getError()
        percent_unc = zh_err / zh_val * 100.
        print 'ZH yield  = {:8.2f}'.format(zh_val)
        print 'ZH uncert = {:8.2f}%'.format(percent_unc)
        return percent_unc

    def draw_pdfs(self):
        self.pframe = self.xvar.frame()
        for pdf in self.pdfs.values():
            pdf.plotOn(self.pframe)
        self.pframe.Draw()

    def draw_data(self):
        self.mframe = self.xvar.frame()
        self.data.plotOn(self.mframe)
        self.underlying_model.plotOn(self.mframe)
        for icomp, compname in enumerate(self.pdfs):
            self.underlying_model.plotOn(self.mframe,
                                         RooFit.Components(compname),
                                         RooFit.LineColor(icomp + 1))
        self.mframe.Draw()
def fillWorkspace(workspace):

    print ">>> setup model components..."
    x = RooRealVar("x", "x", 0, 10)
    mean = RooRealVar("mean", "mean of gaussians", 5, 0, 10)
    sigma1 = RooRealVar("sigma1", "width of gaussians", 0.5)
    sigma2 = RooRealVar("sigma2", "width of gaussians", 1)
    sig1 = RooGaussian("sig1", "Signal component 1", x, mean, sigma1)
    sig2 = RooGaussian("sig2", "Signal component 2", x, mean, sigma2)
    a0 = RooRealVar("a0", "a0", 0.5, 0., 1.)
    a1 = RooRealVar("a1", "a1", -0.2, 0., 1.)
    bkg = RooChebychev("bkg", "Background", x, RooArgList(a0, a1))

    print ">>> sum model components..."
    sig1frac = RooRealVar("sig1frac", "fraction of component 1 in signal", 0.8,
                          0., 1.)
    sig = RooAddPdf("sig", "Signal", RooArgList(sig1, sig2),
                    RooArgList(sig1frac))
    bkgfrac = RooRealVar("bkgfrac", "fraction of background", 0.5, 0., 1.)
    model = RooAddPdf("model", "g1+g2+a", RooArgList(bkg, sig),
                      RooArgList(bkgfrac))

    print ">>> import model into workspace..."
    getattr(workspace, 'import')(model)  # import model and all its components
    #workspace.import(model) # causes synthax error in python

    print "\n>>> encode definition of parameters and observables in workspace..."
    # Define named sets "parameters" and "observables", which list which variables should
    # be considered parameters and observables by the users convention
    #
    # Variables appearing in sets _must_ live in the workspace already, or the autoImport
    # flag of defineSet must be set to import them on the fly. Named sets contain only
    # references to the original variables, therefore the value of observables in named
    # sets already reflect their 'current' value
    params = model.getParameters(RooArgSet(x))  # RooArgSet
    workspace.defineSet("parameters", RooArgSet(params))
    workspace.defineSet("observables", RooArgSet(x))

    # Encode reference value for parameters in workspace:
    # Define a parameter 'snapshot' in the pdf
    # Unlike a named set, a parameter snapshot stores an independent set of values for
    # a given set of variables in the workspace. The values can be stored and reloaded
    # into the workspace variable objects using the loadSnapshot() and saveSnapshot()
    # methods. A snapshot saves the value of each variable, any errors that are stored
    # with it as well as the 'Constant' flag that is used in fits to determine if a
    # parameter is kept fixed or not.

    print ">>> generate and fit data..."
    # Do a dummy fit to a (supposedly) reference dataset here and store the results
    # of that fit into a snapshot
    refData = model.generate(RooArgSet(x), 10000)  # RooDataSet
    model.fitTo(refData, PrintLevel(-1))

    print "\n>>> save fit results into a snapshot in the workspace..."
    # The kTRUE flag imports the values of the objects in (*params) into the workspace
    # If not set, the present values of the workspace parameters objects are stored
    workspace.saveSnapshot("reference_fit", params, kTRUE)

    print ">>> make another fit with the signal component forced to zero..."
    bkgfrac.setVal(1)
    bkgfrac.setConstant(kTRUE)
    bkgfrac.removeError()
    model.fitTo(refData, PrintLevel(-1))

    print "\n>>> save new fit parameters in different snapshot..."
    workspace.saveSnapshot("reference_fit_bkgonly", params, kTRUE)