Test simultaneous fitting in kinematic bins.

test_simultaneous_fit.py

import ROOT as R
from ROOT import RooFit as RF

from charmproduction import config, utilities


def create_binning(var, edges):
    """Return RooBinning for var with bin edges.

    Lowest and highest bin edge are automatically added from
    RooRealVar.getMin and RooRealVar.getMax.
    """
    binning = R.RooBinning(var.getMin(), var.getMax())
    for edge in edges:
        binning.addBoundary(edge)
    return binning


def test_simultaneous_fit(mode, polarity, year):
    """Try get a simultaneous fit to invariant mass."""
    w = R.RooWorkspace('w')
    tfile = R.TFile((
        '/Users/apearce/Physics/Data/CharmProduction2010/Reprocessed'
        '/DVntuple.2010.12.MagUp.root'
    ))
    n = tfile.Get('TupleD0ToKpipipi/DecayTree')
    nentries = n.GetEntries()

    # Create observables
    w.factory('D0_M[1800, 1930]')
    w.factory('D0_IPCHI2_OWNPV[0, 1e5]')
    w.factory('D0_PT[0, 1e4]')
    w.factory('D0_Y[2, 5]')
    observables = R.RooArgSet(
        w.var('D0_M'),
        w.var('D0_IPCHI2_OWNPV'),
        w.var('D0_PT'),
        w.var('D0_Y')
    )

    # Create yield variables
    w.factory('yield_sig[0, {0}]'.format(nentries))
    w.factory('yield_bkg[0, {0}]'.format(nentries))

    # Load data
    data = R.RooDataSet('data', 'data', n, observables)

    # Add the log(IPCHI2) transform
    w.factory((
        'expr::D0_LOG10_IPCHI2_OWNPV('
        '"log10(D0_IPCHI2_OWNPV)",'
        'D0_IPCHI2_OWNPV'
        ')'
    ))
    data.addColumn(w.function('D0_LOG10_IPCHI2_OWNPV'))

    # Create binning categories and add columns to data
    pT_binning = create_binning(w.var('D0_PT'), [2e3, 3e3])
    y_binning = create_binning(w.var('D0_Y'), [3.5])
    w.var('D0_PT').setBinning(pT_binning, 'pT_binning')
    w.var('D0_Y').setBinning(y_binning, 'y_binning')
    pT_bins = R.RooBinningCategory('pT_bins', '', w.var('D0_PT'), 'pT_binning')
    y_bins = R.RooBinningCategory('y_bins', '', w.var('D0_Y'), 'y_binning')
    data.addColumn(pT_bins)
    data.addColumn(y_bins)

    getattr(w, 'import')(data, R.RooCmdArg())

    # Create mass shape PDFs
    # Signal shape
    w.factory('mu_m_sig[1850, 1880]')
    w.factory('sigma_m_sig[2, 10]')
    w.factory((
        'RooGaussian::pdf_m_sig('
        'D0_M, mu_m_sig, sigma_m_sig'
        ')'
    ))
    # Background shape
    w.factory('a0_m_bkg[0, -0.2, 0.2]')
    w.factory('a0_delta_m_bkg[0]')
    w.factory('expr::a0_what_m_bkg("@0 + @1", a0_m_bkg, a0_delta_m_bkg)')
    w.factory((
        'RooChebychev::pdf_m_bkg('
        'D0_M, {a0_what_m_bkg}'
        ')'
    ))

    # Create extended PDF
    w.factory((
        'SUM::pdf_m_tot('
        'yield_sig*pdf_m_sig,'
        'yield_bkg*pdf_m_bkg'
        ')'
    ))

    # Create simultaneous PDF
    w.factory((
        'SIMCLONE::pdf_sim_m_tot('
        'pdf_m_tot,'
        '$SplitParam({a0_delta_m_bkg,yield_sig,yield_bkg}, {pT_bins, y_bins})'
        ')'
    ))

    # Fit
    w.pdf('pdf_sim_m_tot').fitTo(w.data('data'))

    # Plot
    pT_nbins = w.var('D0_PT').numBins('pT_binning')
    y_nbins = w.var('D0_Y').numBins('y_binning')
    projwdata = RF.ProjWData(
        R.RooArgSet(w.cat('pT_bins'), w.cat('y_bins')),
        w.data('data')
    )
    mosaic = R.TCanvas('mosiac', 'mosiac', pT_nbins*200, y_nbins*200)
    mosaic.Divide(pT_nbins, y_nbins)
    pad = 1
    for y_bin in range(y_nbins):
        w.cat('y_bins').setIndex(y_bin)
        for pT_bin in range(pT_nbins):
            w.cat('pT_bins').setIndex(pT_bin)

            pT_label = w.cat('pT_bins').getLabel()
            y_label = w.cat('y_bins').getLabel()

            f = w.var('D0_M').frame(
                RF.Bins(65),
                RF.Title('{0} {1}'.format(pT_label, y_label))
            )
            cut = 'pT_bins == pT_bins::{0} && y_bins == y_bins::{1}'.format(
                pT_label, y_label
            )
            pT_slice = RF.Slice(w.cat('pT_bins'), pT_label)
            y_slice = RF.Slice(w.cat('y_bins'), y_label)
            print '-'*40
            print 'Cuts:', cut
            print '-'*40
            w.data('data').plotOn(f, RF.Cut(cut), RF.MarkerSize(0.5))
            w.pdf('pdf_sim_m_tot').plotOn(
                f,
                projwdata,
                pT_slice,
                y_slice,
                RF.Components('pdf_m_sig*'),
                RF.LineStyle(2)
            )
            w.pdf('pdf_sim_m_tot').plotOn(
                f,
                projwdata,
                pT_slice,
                y_slice,
                RF.Components('pdf_m_bkg*'),
                RF.LineColor(R.kRed),
                RF.LineStyle(3)
            )
            w.pdf('pdf_sim_m_tot').plotOn(
                f,
                projwdata,
                pT_slice,
                y_slice,
            )

            # Split the pad again and resize them
            currpad = mosaic.cd(pad)
            currpad.Divide(1, 2)
            currpad.cd(1)
            R.gPad.SetPad(0, 0.25, 1, 1)
            currpad.cd(2)
            R.gPad.SetPad(0, 0, 1, 0.25)

            # Draw the fit in the first
            currpad.cd(1)
            f.Draw()

            # Draw the pulls in the second
            currpad.cd(2)
            pulls = w.var('D0_M').frame(RF.Title(' '))
            pulls.addPlotable(f.pullHist(), 'BX0')
            pulls.SetMaximum(5)
            pulls.SetMinimum(-5)
            pulls.GetYaxis().SetTitle('#Delta/#sigma')
            pulls.Draw()

            pad += 1
    mosaic.SaveAs('test_simultaneous_fit_mosaic.pdf')

    w.Print('v')


if __name__ == '__main__':
    utilities.quiet_mode()
    test_simultaneous_fit(config.D0ToKpipipi, config.magup, 2010)

Stick this in analysis/, edit the (hardcoded, oops) path to the data and run it.