Simultaneous fit

simul.py

import ROOT as R
from ROOT import RooFit as RF

# 'Unblind' or 'Blind'
# BLINDING = 'Blind'
BLINDING = 'Unblind'
BLIND_STR = '6LJgnoyA4Zr95gec'

def roovar_value(rv):
    '''Return a pretty-printed central value and uncertainty of rv.

    The string returned is of the form 'XX +/- YY'.
    Keyword arguments:
    rv -- Instance of a RooRealVar
    '''
    return '{0:.5f} +/- {1:.5f}'.format(rv.getVal(), rv.getError())


def add_blinded_araw(w):
    '''Add a blinded ARaw parameter to the workspace w.'''
    blind = R.RooUnblindPrecision(
        'ARaw_unblind', 'ARaw_unblind', BLIND_STR,
        w.var('ARaw').getVal(), 1, w.var('ARaw'), w.cat('blinding')
    )
    getattr(w, 'import')(blind, R.RooCmdArg())


def simultaneous_fit(toy=False, entries=10000, split=0.8, asym=0.):
    '''Perform a simultaneous fit to the +/- tagged samples.

    Keyword arguments:
    toy -- Use generated toy MC if True, else use pK-pi+ sample (default: False)
    count -- Number of toy MC event to generate (default: 10,000)
    split -- The ratio of generated signal-to-background events (default: 0.8)
    asym -- Initial value of ARaw within (-1., 1.) exclusively (default: 0)
    '''
    if not toy:
        f = R.TFile('selected-pKpi-2011-20r1-MagUp-sample.root')
        t = f.Get('DecayTree')
        entries = t.GetEntries()

    w = R.RooWorkspace('w')
    w_import = getattr(w, 'import')
    # Fit variable
    w.factory('Lambdac_M[2220, 2360]')
    # Categories
    w.factory('Lambdac_ID[positive=4122, negative=-4122]')
    w.cat('Lambdac_ID').SetTitle('Lambda_c charge')
    # Yield asymmetries
    w.factory('ARaw[{0}, -10, 10]'.format(asym))
    w.factory('blinding[Unblind=0, Blind=1]')
    w.cat('blinding').setLabel(BLINDING)
    add_blinded_araw(w)
    w.factory('ABkg[0, -1, 1]')
    # Yields
    w.factory('y_sig[{0}, {1}, {2}]'.format(entries*split, 0, entries))
    w.factory('y_bkg[{0}, {1}, {2}]'.format(entries*(1 - split), 0, entries))
    # PDFs
    w.factory('RooGaussian::pdf_g1(Lambdac_M, mu[2290, 2285, 2295], s1[4, 0, 15])')
    w.factory('RooGaussian::pdf_g2(Lambdac_M, mu, s2[7, 0, 15])')
    w.factory('SUM::pdf_sig(gfrac[0.5, 0, 1]*pdf_g1, pdf_g2)')
    w.factory('RooExponential::pdf_bkg(Lambdac_M, lambda[-0.001, -0.1, 0.1])')
    w.factory('SUM::pdf_tot(y_sig*pdf_sig, y_bkg*pdf_bkg)')
    # Create one PDF per Lambdac_ID category (+/-), giving each PDF
    # separate signal means and widths, and separate background decay constants
    w.factory('SIMCLONE::pdf_tot_temp('
        'pdf_tot, $SplitParam({mu, s1, s2, gfrac, lambda, y_sig, y_bkg}, Lambdac_ID)'
    ')')
    w.Print()
    # Edit the PDFs to use the yield asymmetry parameters
    w.factory('EDIT::pdf_tot2_positive('
        'pdf_tot_positive,'
        "y_sig_positive=expr('y_sig*(1 + ARaw_unblind)/2.0', {y_sig, ARaw_unblind}),"
        "y_bkg_positive=expr('y_bkg*(1 + ABkg)/2.0', {y_bkg, ABkg})"
    ')')
    w.factory('EDIT::pdf_tot2_negative('
        'pdf_tot_negative,'
        "y_sig_negative=expr('y_sig*(1 - ARaw_unblind)/2.0', {y_sig, ARaw_unblind}),"
        "y_bkg_negative=expr('y_bkg*(1 - ABkg)/2.0', {y_bkg, ABkg})"
    ')')
    w.factory('SIMUL::pdf_tot_sim(Lambdac_ID,positive=pdf_tot2_positive,negative=pdf_tot2_negative)')

    # Data
    w.defineSet('fit_vars', 'Lambdac_M,Lambdac_ID')
    fit_vars = R.RooArgList(w.var('Lambdac_M'), w.cat('Lambdac_ID'))
    if toy:
        w.factory('EDIT::pdf_tot_gen(pdf_tot_sim, ARaw_unblind=ARaw_true[{0}])'.format(asym))
        data = w.pdf('pdf_tot_gen').generate(
            w.set('fit_vars'),
            entries,
            RF.Name('data')
        )
    else:
        data = R.RooDataSet('data', 'data', t, w.set('fit_vars'))
        f.Close()
    w_import(data, R.RooCmdArg())
    w.data('data').table(w.cat('Lambdac_ID')).Print('v')

    # Do it
    fit_result = w.pdf('pdf_tot_sim').fitTo(w.data('data'), RF.Save(True))
    fit_result.Print()

    # Plot each tag separately
    v = w.var('Lambdac_M')
    frames = []
    for charge in ['positive', 'negative']:
        frame = v.frame(RF.Title('Lambdac_ID::{0}'.format(charge)), RF.Bins(70))
        frame.GetYaxis().SetTitleOffset(1.4)
        w.data('data').plotOn(
            frame,
            RF.Cut('Lambdac_ID==Lambdac_ID::{0}'.format(charge)),
            RF.MarkerSize(0.3)
        )
        w.pdf('pdf_tot_sim').plotOn(
            frame,
            RF.Slice(w.cat('Lambdac_ID'), charge),
            RF.ProjWData(R.RooArgSet(w.cat('Lambdac_ID')), w.data('data')),
            RF.Components('pdf_bkg_{0}'.format(charge)),
            RF.LineColor(R.kRed),
            RF.LineStyle(R.kDashed)
        )
        w.pdf('pdf_tot_sim').plotOn(
            frame,
            RF.Slice(w.cat('Lambdac_ID'), charge),
            RF.ProjWData(R.RooArgSet(w.cat('Lambdac_ID')), w.data('data')),
            RF.Components('pdf_tot2_{0}'.format(charge)),
            RF.LineColor(R.kBlue + 1)
        )
        w.data('data').plotOn(
            frame,
            RF.Cut('Lambdac_ID==Lambdac_ID::{0}'.format(charge)),
            RF.MarkerSize(0.3)
        )
        frames.append(frame)
    num_frames = len(frames)
    y_max = max([frame.GetMaximum() for frame in frames])

    # Draw it
    c = R.TCanvas('c', 'c', 600, 400)
    c.Divide(num_frames, 2)
    x_step = 1.0/num_frames
    curr_pad = 0
    for idx, frame in enumerate(frames):
        # Draw the data and the fit
        curr_pad += 1
        c.cd(curr_pad)
        R.gPad.SetPad(x_step*idx, 0.25, x_step*(idx + 1), 1)
        # Equalise the y-axes for side-by-side comparison
        frame.SetMaximum(y_max)
        frame.Draw()
        # Draw the pulls
        curr_pad += 1
        c.cd(curr_pad)
        R.gPad.SetPad(x_step*idx, 0, x_step*(idx + 1), 0.25)
        pull = v.frame(RF.Title(' '))
        pull.addPlotable(frame.pullHist(), 'BX0')
        pull.GetXaxis().SetTitle('')
        pull.SetMaximum(5)
        pull.SetMinimum(-5)
        pull.GetYaxis().SetNdivisions(503)
        pull.GetYaxis().SetTitle('#Delta/#sigma')
        pull.Draw()
        x_min = frame.GetXaxis().GetXmin()
        x_max = frame.GetXaxis().GetXmax()
        l1 = R.TLine(x_min, 2, x_max, 2)
        l2 = R.TLine(x_min, -2, x_max, -2)
        l1.SetLineColor(R.kRed)
        l2.SetLineColor(R.kRed)
        l1.Draw()
        l2.Draw()
        # Prevent garbage collection
        frame.l1 = l1
        frame.l2 = l2
    c.SaveAs('simultaneous_fit.pdf')

    print '='*40
    print 'Initial values'
    print '-'*40
    print 'Initial events:', entries
    if toy:
        print 'Initial value of ARaw:', asym
        print 'Initial split value:', split
        print 'Expected number of signal entries:', entries*split
        print 'Expected number of background entries:', entries*(1 - split)
    print '='*40
    print 'Results of toy fit' if toy else 'Results of data fit'
    print '-'*40
    print 'Fitted value of signal entries:', roovar_value(w.var('y_sig'))
    print 'Fitted value of background entries:', roovar_value(w.var('y_bkg'))
    print 'Fitted value of ARaw:', roovar_value(w.var('ARaw'))
    print '='*40

if __name__ == '__main__':
    # Suppress RooFit INFO messages
    R.RooMsgService.instance().setGlobalKillBelow(RF.WARNING)
    # Different seed on each run
    # R.RooRandom.randomGenerator().SetSeed(0)
    # simultaneous_fit()
    simultaneous_fit(True, entries=10000, split=0.5, asym=0.2)