Exemplo n.º 1
0
 def __init__(self, adc_east, adc_west):
     #input ADC values
     self.adc_east = adc_east
     self.adc_west = adc_west
     #self.adc_east = RooRealVar("adc_east", "adc_east", 10, 1300)
     #self.adc_west = RooRealVar("adc_west", "adc_west", 10, 1300)
     #east Gaussian
     self.gauss_east = Gauss(self.adc_east, "east")
     self.gauss_east.mean_1n.setVal(72.9)
     self.gauss_east.sigma_1n.setVal(21.4)
     #west Gaussian
     self.gauss_west = Gauss(self.adc_west, "west")
     self.gauss_west.mean_1n.setVal(87.7)
     self.gauss_west.sigma_1n.setVal(26.9)
     #east Crystal Ball
     self.cb_east = CrystalBall(self.adc_east, "east")
     self.cb_east.mean_2n.setVal(166.)
     self.cb_east.sigma_2n.setVal(42.1)
     self.cb_east.alpha_2xn.setVal(-0.7)
     self.cb_east.n_2xn.setVal(0.5)
     #west Crystal Ball
     self.cb_west = CrystalBall(self.adc_west, "west")
     self.cb_west.mean_2n.setVal(174.1)
     self.cb_west.sigma_2n.setVal(29.3)
     self.cb_west.alpha_2xn.setVal(-0.3)
     self.cb_west.n_2xn.setVal(0.8)
     # (g_e + c_e)*(g_w + c_w) = g_e*g_w + c_e*c_w + g_e*c_w + c_e*g_w
     #self.num_max = 3000
     self.num_max = 300000
     #1n1n 2D Gaussian
     self.pdf_1n1n = RooProdPdf("pdf_1n1n", "pdf_1n1n", RooArgList(self.gauss_east.gauss_1n, self.gauss_west.gauss_1n))
     self.num_1n1n = RooRealVar("num_1n1n", "num_1n1n", 200, 0, self.num_max) # 1
     #1n2xn Gaussian * Crystal Ball
     self.pdf_1n2xn = RooProdPdf("pdf_1n2xn", "pdf_1n2xn", RooArgList(self.gauss_east.gauss_1n, self.cb_west.cb_2xn))
     self.num_1n2xn = RooRealVar("num_1n2xn", "num_1n2xn", 100, 0, self.num_max) # 1
     #2xn1n Crystal Ball * Gaussian
     self.pdf_2xn1n = RooProdPdf("pdf_2xn1n", "pdf_2xn1n", RooArgList(self.cb_east.cb_2xn, self.gauss_west.gauss_1n))
     self.num_2xn1n = RooRealVar("num_2xn1n", "num_2xn1n", 100, 0, self.num_max) # 1
     #2xn2xn 2D Crystal Ball
     self.pdf_2xn2xn = RooProdPdf("pdf_2xn2xn", "pdf_2xn2xn", RooArgList(self.cb_east.cb_2xn, self.cb_west.cb_2xn))
     self.num_2xn2xn = RooRealVar("num_2xn2xn", "num_2xn2xn", 50, 0, self.num_max) # 1
     #fit model
     self.model = RooAddPdf("model", "model", RooArgList(self.pdf_1n1n, self.pdf_1n2xn, self.pdf_2xn1n, self.pdf_2xn2xn),
     RooArgList(self.num_1n1n, self.num_1n2xn, self.num_2xn1n, self.num_2xn2xn))
Exemplo n.º 2
0
 def fit(self):
     pdfs = RooArgList()
     obsvars = RooArgSet('set')
     for constraint in self.constraint.values():
         pdfs.add(constraint.pdf_constraint)
         if hasattr(constraint, 'var_obs'):
             obsvars.add(constraint.var_obs)
     self.model = RooProdPdf('model', 'model', pdfs)
     self.data = RooDataSet('data', 'data', obsvars)
     self.data.add(obsvars)
     self.data.Print()
     self.fit_result = self.model.fitTo(self.data,
                                        ROOT.RooFit.PrintLevel(3),
                                        ROOT.RooFit.Optimize(1),
                                        ROOT.RooFit.Hesse(1),
                                        ROOT.RooFit.Minos(1),
                                        ROOT.RooFit.Strategy(2),
                                        ROOT.RooFit.Save(1))
Exemplo n.º 3
0
    def fit(self):
        fit_variable = RooRealVar("fit_variable", "fit_variable",
                                  self.fit_boundaries[0],
                                  self.fit_boundaries[1])
        fit_variable.setBins(self.histograms[self.data_label].nbins())
        variables = RooArgList()
        variables.add(fit_variable)
        variable_set = RooArgSet()
        variable_set.add(fit_variable)

        roofit_histograms = {}
        roofit_pdfs = {}
        roofit_variables = {}

        N_min = 0.
        N_max = self.normalisation[self.data_label] * 2.
        pdf_arglist = RooArgList()
        variable_arglist = RooArgList()

        roofit_histograms[self.data_label] = RooDataHist(
            self.data_label, self.data_label, variables,
            self.histograms[self.data_label])
        for sample in self.samples:
            roofit_histogram = RooDataHist(sample, sample, variables,
                                           self.histograms[sample])
            roofit_histograms[sample] = roofit_histogram
            roofit_pdf = RooHistPdf('pdf' + sample, 'pdf' + sample,
                                    variable_set, roofit_histogram)
            roofit_pdfs[sample] = roofit_pdf
            roofit_variable = RooRealVar(sample, sample + " events",
                                         self.normalisation[sample], N_min,
                                         N_max)
            roofit_variables[sample] = roofit_variable
            pdf_arglist.add(roofit_pdf)
            variable_arglist.add(roofit_variable)

        model = RooAddPdf('model', 'sum of all known', pdf_arglist,
                          variable_arglist)
        use_model = model
        if self.constraints:
            arg_set = RooArgSet(model)
            constraints = self.get_fit_normalisation_constraints(
                model, roofit_variables)
            for constraint in constraints:
                arg_set.add(constraint)
            model_with_constraints = RooProdPdf(
                "model_with_constraints", "model  with gaussian constraints",
                arg_set, RooLinkedList())
            use_model = model_with_constraints

        if self.method == 'TMinuit':
            #WARNING: number of cores changes the results!!!
            self.saved_result = use_model.fitTo(
                roofit_histograms[self.data_label],
                RooFit.Minimizer("Minuit2", "Migrad"),
                RooFit.NumCPU(1),
                RooFit.Extended(),
                RooFit.Save(),
            )

        results = {}
        for sample in self.samples:
            results[sample] = (roofit_variables[sample].getVal(),
                               roofit_variables[sample].getError())
        self.results = results
Exemplo n.º 4
0
def buildTimePdf(config):
    """
    build time pdf, return pdf and associated data in dictionary
    """
    from B2DXFitters.WS import WS
    print 'CONFIGURATION'
    for k in sorted(config.keys()):
        print '    %32s: %32s' % (k, config[k])

    # start building the fit
    ws = RooWorkspace('ws_%s' % config['Context'])
    one = WS(ws, RooConstVar('one', '1', 1.0))
    zero = WS(ws, RooConstVar('zero', '0', 0.0))

    # start by defining observables
    time = WS(ws, RooRealVar('time', 'time [ps]', 0.2, 15.0))
    qf = WS(ws, RooCategory('qf', 'final state charge'))
    qf.defineType('h+', +1)
    qf.defineType('h-', -1)
    qt = WS(ws, RooCategory('qt', 'tagging decision'))
    qt.defineType('B+', +1)
    qt.defineType('Untagged', 0)
    qt.defineType('B-', -1)

    # now other settings
    Gamma = WS(ws, RooRealVar('Gamma', 'Gamma', 0.661))  # ps^-1
    DGamma = WS(ws, RooRealVar('DGamma', 'DGamma', 0.106))  # ps^-1
    Dm = WS(ws, RooRealVar('Dm', 'Dm', 17.719))  # ps^-1

    # HACK (1/2): be careful about lower bound on eta, since mistagpdf below
    # is zero below a certain value - generation in accept/reject would get
    # stuck
    eta = WS(
        ws,
        RooRealVar(
            'eta', 'eta', 0.35, 0.0 if 'FIT' in config['Context'] else
            (1. + 1e-5) * max(0.0, config['TrivialMistagParams']['omega0']),
            0.5))
    tageff = WS(ws, RooRealVar('tageff', 'tageff', 0.60, 0.0, 1.0))
    timeerr = WS(ws, RooRealVar('timeerr', 'timeerr', 0.040, 0.001, 0.100))

    # now build the PDF
    from B2DXFitters.timepdfutils import buildBDecayTimePdf
    from B2DXFitters.resmodelutils import getResolutionModel
    from B2DXFitters.acceptanceutils import buildSplineAcceptance

    obs = [qt, qf, time, eta, timeerr]

    acc, accnorm = buildSplineAcceptance(
        ws, time, 'Bs2DsPi_accpetance',
        config['SplineAcceptance']['KnotPositions'],
        config['SplineAcceptance']['KnotCoefficients'][config['Context']],
        'FIT' in config['Context'])  # float for fitting
    if 'GEN' in config['Context']:
        acc = accnorm  # use normalised acceptance for generation
    # get resolution model
    resmodel, acc = getResolutionModel(ws, config, time, timeerr, acc)

    # build a (mock) mistag distribution
    mistagpdfparams = {}  # start with parameters of mock distribution
    for sfx in ('omega0', 'omegaavg', 'f'):
        mistagpdfparams[sfx] = WS(
            ws,
            RooRealVar('Bs2DsPi_mistagpdf_%s' % sfx,
                       'Bs2DsPi_mistagpdf_%s' % sfx,
                       config['TrivialMistagParams'][sfx]))
    # build mistag pdf itself
    mistagpdf = WS(
        ws,
        MistagDistribution('Bs2DsPi_mistagpdf', 'Bs2DsPi_mistagpdf', eta,
                           mistagpdfparams['omega0'],
                           mistagpdfparams['omegaavg'], mistagpdfparams['f']))
    # build mistag calibration
    mistagcalibparams = {}  # start with parameters of calibration
    for sfx in ('p0', 'p1', 'etaavg'):
        mistagcalibparams[sfx] = WS(
            ws,
            RooRealVar('Bs2DsPi_mistagcalib_%s' % sfx,
                       'Bs2DsPi_mistagpdf_%s' % sfx,
                       config['MistagCalibParams'][sfx]))
    for sfx in ('p0', 'p1'):  # float calibration paramters
        mistagcalibparams[sfx].setConstant(False)
        mistagcalibparams[sfx].setError(0.1)
    # build mistag pdf itself
    omega = WS(
        ws,
        MistagCalibration('Bs2DsPi_mistagcalib', 'Bs2DsPi_mistagcalib', eta,
                          mistagcalibparams['p0'], mistagcalibparams['p1'],
                          mistagcalibparams['etaavg']))
    # build mock decay time error distribution (~ timeerr^6 * exp(-timerr /
    # (timerr_av / 7))
    terrpdf_shape = WS(
        ws,
        RooConstVar('timeerr_ac', 'timeerr_ac',
                    config['DecayTimeResolutionAvg'] / 7.))
    terrpdf_truth = WS(
        ws, RooTruthModel('terrpdf_truth', 'terrpdf_truth', timeerr))
    terrpdf_i0 = WS(
        ws,
        RooDecay('terrpdf_i0', 'terrpdf_i0', timeerr, terrpdf_shape,
                 terrpdf_truth, RooDecay.SingleSided))
    terrpdf_i1 = WS(
        ws,
        RooPolynomial('terrpdf_i1', 'terrpdf_i1', timeerr,
                      RooArgList(zero, zero, zero, zero, zero, zero, one), 0))
    terrpdf = WS(ws, RooProdPdf('terrpdf', 'terrpdf', terrpdf_i0, terrpdf_i1))

    # build the time pdf
    pdf = buildBDecayTimePdf(config,
                             'Bs2DsPi',
                             ws,
                             time,
                             timeerr,
                             qt,
                             qf, [[omega]], [tageff],
                             Gamma,
                             DGamma,
                             Dm,
                             C=one,
                             D=zero,
                             Dbar=zero,
                             S=zero,
                             Sbar=zero,
                             timeresmodel=resmodel,
                             acceptance=acc,
                             timeerrpdf=terrpdf,
                             mistagpdf=[mistagpdf],
                             mistagobs=eta)
    return {  # return things
        'ws': ws,
        'pdf': pdf,
        'obs': obs
    }
Exemplo n.º 5
0
def fitData(fulldata, ibin):

    cut = cut_base + '&& (mumuMass*mumuMass > %s && mumuMass*mumuMass < %s)' % (
        q2binning[ibin], q2binning[ibin + 1])
    data = fulldata.reduce(RooArgSet(tagged_mass, mumuMass, mumuMassE), cut)

    fraction = dict_s_rt[ibin] / (dict_s_rt[ibin] + dict_s_wt[ibin])
    print 'mistag fraction on MC for bin ', ibin, ' : ', fraction.n, '+/-', fraction.s

    ### creating RT component
    w.loadSnapshot("reference_fit_RT_%s" % ibin)
    sigmart1 = w.var("#sigma_{1}^{RT%s}" % ibin)
    sigmart2 = w.var("#sigma_{2}^{RT%s}" % ibin)
    massrt = w.var("mean^{RT%s}" % ibin)
    f1rt = w.var("f^{RT%s}" % ibin)

    theRTgauss = w.pdf("doublegaus_RT%s" % ibin)
    c_sigma_rt1 = _constrainVar(sigmart1)
    c_sigma_rt2 = _constrainVar(sigmart2)
    c_mean_rt = _constrainVar(massrt)
    c_f1rt = _constrainVar(f1rt)

    ### creating WT component
    w.loadSnapshot("reference_fit_WT_%s" % ibin)
    meanwt = w.var("mean^{WT%s}" % ibin)
    sigmawt = w.var("#sigma_{CB}^{WT%s}" % ibin)
    alphawt1 = w.var("#alpha_{1}^{WT%s}" % ibin)
    alphawt2 = w.var("#alpha_{2}^{WT%s}" % ibin)
    nwt1 = w.var("n_{1}^{WT%s}" % ibin)
    nwt2 = w.var("n_{2}^{WT%s}" % ibin)

    theWTgauss = w.pdf("doublecb_%s" % ibin)
    c_mean_wt = _constrainVar(meanwt)
    c_sigma_wt = _constrainVar(sigmawt)
    c_alpha_wt1 = _constrainVar(alphawt1)
    c_alpha_wt2 = _constrainVar(alphawt2)
    c_n_wt1 = _constrainVar(nwt1)
    c_n_wt2 = _constrainVar(nwt2)

    ### creating constraints for the RT component
    c_RTgauss = RooProdPdf(
        "c_RTgauss", "c_RTgauss",
        RooArgList(theRTgauss, c_sigma_rt1, c_sigma_rt2, c_mean_rt, c_f1rt))

    c_vars = RooArgSet(c_sigma_rt1, c_sigma_rt2, c_f1rt, c_mean_rt)
    c_vars.add(c_sigma_wt)
    c_vars.add(c_mean_wt)
    c_vars.add(c_alpha_wt1)
    c_vars.add(c_alpha_wt2)
    c_vars.add(c_n_wt1)
    c_vars.add(c_n_wt2)

    ### creating constraints for the WT component
    c_WTgauss = RooProdPdf(
        "c_WTgauss", "c_WTgauss",
        RooArgList(theWTgauss, c_alpha_wt1, c_n_wt1, c_sigma_wt, c_mean_wt,
                   c_alpha_wt2, c_n_wt2))

    frt = RooRealVar("F_{RT}", "frt", fraction.n, 0, 1)
    signalFunction = RooAddPdf("sumgaus", "rt+wt",
                               RooArgList(c_RTgauss, c_WTgauss),
                               RooArgList(frt))
    c_frt = RooGaussian("c_frt", "c_frt", frt,
                        ROOT.RooFit.RooConst(fraction.n),
                        ROOT.RooFit.RooConst(fraction.s))
    c_signalFunction = RooProdPdf("c_signalFunction", "c_signalFunction",
                                  RooArgList(signalFunction, c_frt))
    c_vars.add(frt)

    ### now create background parametrization
    slope = RooRealVar("slope", "slope", 0.5, -10, 10)
    bkg_exp = RooExponential("bkg_exp", "exponential", slope, tagged_mass)
    pol_c1 = RooRealVar("p1", "coeff x^0 term", 0.5, -10, 10)
    pol_c2 = RooRealVar("p2", "coeff x^1 term", 0.5, -10, 10)
    bkg_pol = RooChebychev("bkg_pol", "2nd order pol", tagged_mass,
                           RooArgList(pol_c1, pol_c2))

    nsig = RooRealVar("Yield", "signal frac", 4000, 0, 1000000)
    nbkg = RooRealVar("nbkg", "bkg fraction", 1000, 0, 550000)

    #     fitFunction = RooAddPdf ("fitfunction" , "fit function"  ,  RooArgList(c_signalFunction, bkg_pol), RooArgList(nsig, nbkg))
    fitFunction = RooAddPdf("fitfunction", "fit function",
                            RooArgList(c_signalFunction, bkg_exp),
                            RooArgList(nsig, nbkg))

    r = fitFunction.fitTo(data, RooFit.Extended(True), RooFit.Save(),
                          RooFit.Range("full"), RooFit.Verbose(False),
                          ROOT.RooFit.Constrain(c_vars))

    frame = tagged_mass.frame(RooFit.Range("full"))
    data.plotOn(frame, RooFit.Binning(35), RooFit.MarkerSize(.7))
    fitFunction.plotOn(frame)
    drawPdfComponents(fitFunction,
                      frame,
                      ROOT.kAzure,
                      RooFit.NormRange("full"),
                      RooFit.Range("full"),
                      isData=True)

    parList = RooArgSet(nsig, massrt, sigmart1, sigmart2, f1rt, meanwt,
                        sigmawt, alphawt1)
    parList.add(alphawt2)
    parList.add(nwt1)
    parList.add(nwt2)
    parList.add(frt)
    fitFunction.paramOn(frame, RooFit.Parameters(parList),
                        RooFit.Layout(0.62, 0.86, 0.89))
    frame.Draw()
    niceFrame(frame, '')
    frame.addObject(_writeFitStatus(r))

    if not args.year == 'test':
        writeCMS(frame, args.year, [q2binning[ibin], q2binning[ibin + 1]])
    frame.Draw()
    c1.SaveAs('fit_results_mass/save_fit_data_%s_%s_LMNR.pdf' %
              (ibin, args.year))
def fitData(fulldata, ibin, n_bkg, w):

    cut = cut_base + '&& (mumuMass*mumuMass > %s && mumuMass*mumuMass < %s)' % (
        q2binning[ibin], q2binning[ibin + 1])
    fulldata_v2 = fulldata.reduce(
        RooArgSet(tagged_mass, mumuMass, mumuMassE, randVar), cut)

    ## reduce to data-like statistics
    nDataEntries = fulldata_v2.sumEntries()
    nDesired = n_bin[ibin] / nDataEntries
    cut = 'rand < %f' % nDesired
    signaldata = fulldata_v2.reduce(
        RooArgSet(tagged_mass, mumuMass, mumuMassE), cut)
    n_realsignal = signaldata.sumEntries()

    nrt_mc = _getFittedVar("nRT_%s" % ibin, w)
    nwt_mc = _getFittedVar("nWT_%s" % ibin, w)
    fraction = nrt_mc / (nrt_mc + nwt_mc)

    ### creating RT component
    w.loadSnapshot("reference_fit_RT_%s" % ibin)
    meanrt = w.var("mean^{RT%s}" % ibin)
    sigmart = RooRealVar()
    sigmart1 = RooRealVar()
    sigmart2 = RooRealVar()
    alphart1 = RooRealVar()
    alphart2 = RooRealVar()
    nrt1 = RooRealVar()
    nrt2 = RooRealVar()

    ## double cb fast
    if ibin < 5:
        sigmart = w.var("#sigma_{CB}^{RT%s}" % ibin)
        alphart1 = w.var("#alpha_{1}^{RT%s}" % ibin)
        alphart2 = w.var("#alpha_{2}^{RT%s}" % ibin)
        nrt1 = w.var("n_{1}^{RT%s}" % ibin)
        nrt2 = w.var("n_{2}^{RT%s}" % ibin)

    ## double cb old
    else:
        sigmart1 = w.var("#sigma_{CBRT0}^{%s}" % ibin)
        sigmart2 = w.var("#sigma_{CBRT1}^{%s}" % ibin)
        alphart1 = w.var("#alpha_{RT0}^{%s}" % ibin)
        alphart2 = w.var("#alpha_{RT1}^{%s}" % ibin)
        nrt1 = w.var("n_{RT0}^{%s}" % ibin)
        nrt2 = w.var("n_{RT1}^{%s}" % ibin)

    pars_init_vals = {}
    theRTgauss = w.pdf("doublecb_RT%s" % ibin)
    if ibin < 5:
        c_sigma_rt = _constrainVar(sigmart, 1, pars_init_vals)
    else:
        c_sigma_rt1 = _constrainVar(sigmart1, 1, pars_init_vals)
        c_sigma_rt2 = _constrainVar(sigmart2, 1, pars_init_vals)

    c_alpha_rt1 = _constrainVar(alphart1, 1, pars_init_vals)
    c_alpha_rt2 = _constrainVar(alphart2, 1, pars_init_vals)
    c_n_rt1 = _constrainVar(nrt1, 1, pars_init_vals)
    c_n_rt2 = _constrainVar(nrt2, 1, pars_init_vals)

    ### creating WT component
    w.loadSnapshot("reference_fit_WT_%s" % ibin)
    meanwt = w.var("mean^{WT%s}" % ibin)
    sigmawt = w.var("#sigma_{CB}^{WT%s}" % ibin)
    alphawt1 = w.var("#alpha_{1}^{WT%s}" % ibin)
    alphawt2 = w.var("#alpha_{2}^{WT%s}" % ibin)
    nwt1 = w.var("n_{1}^{WT%s}" % ibin)
    nwt2 = w.var("n_{2}^{WT%s}" % ibin)

    theWTgauss = w.pdf("doublecb_%s" % ibin)
    c_sigma_wt = _constrainVar(sigmawt, 1, pars_init_vals)
    c_alpha_wt1 = _constrainVar(alphawt1, 1, pars_init_vals)
    c_alpha_wt2 = _constrainVar(alphawt2, 1, pars_init_vals)
    c_n_wt1 = _constrainVar(nwt1, 1, pars_init_vals)
    c_n_wt2 = _constrainVar(nwt2, 1, pars_init_vals)

    ### creating constraints for the RT component
    c_vars = RooArgSet()
    if ibin < 5:
        c_RTgauss = RooProdPdf(
            "c_RTgauss", "c_RTgauss",
            RooArgList(theRTgauss, c_alpha_rt1, c_n_rt1, c_sigma_rt,
                       c_alpha_rt2, c_n_rt2))
        c_vars = RooArgSet(c_sigma_rt, c_alpha_rt1, c_alpha_rt2, c_n_rt1,
                           c_n_rt2)
    else:
        c_RTgauss = RooProdPdf(
            "c_RTgauss", "c_RTgauss",
            RooArgList(theRTgauss, c_alpha_rt1, c_n_rt1, c_sigma_rt1,
                       c_sigma_rt2, c_alpha_rt2, c_n_rt2))
        c_vars = RooArgSet(c_sigma_rt1, c_sigma_rt2, c_alpha_rt1, c_alpha_rt2,
                           c_n_rt1, c_n_rt2)

    ### creating constraints for the WT component
    c_WTgauss = RooProdPdf(
        "c_WTgauss", "c_WTgauss",
        RooArgList(theWTgauss, c_alpha_wt1, c_n_wt1, c_sigma_wt, c_alpha_wt2,
                   c_n_wt2))
    c_vars.add(c_sigma_wt)
    c_vars.add(c_alpha_wt1)
    c_vars.add(c_alpha_wt2)
    c_vars.add(c_n_wt1)
    c_vars.add(c_n_wt2)

    frt = RooRealVar("F_{RT}", "frt", fraction.n, 0, 1)
    signalFunction = RooAddPdf("sumgaus", "rt+wt",
                               RooArgList(c_RTgauss, c_WTgauss),
                               RooArgList(frt))
    c_frt = RooGaussian("c_frt", "c_frt", frt,
                        ROOT.RooFit.RooConst(fraction.n),
                        ROOT.RooFit.RooConst(fraction.s))

    ### creating constraints for the difference between the two peaks
    deltaPeaks = RooFormulaVar("deltaPeaks", "@0 - @1",
                               RooArgList(meanrt, meanwt))
    c_deltaPeaks = RooGaussian(
        "c_deltaPeaks",
        "c_deltaPeaks",
        deltaPeaks,
        ROOT.RooFit.RooConst(deltaPeaks.getVal()),
        ROOT.RooFit.RooConst(0.0005)  ## value to be checked
    )

    c_signalFunction = RooProdPdf(
        "c_signalFunction", "c_signalFunction",
        RooArgList(signalFunction, c_frt, c_deltaPeaks))
    c_vars.add(frt)
    c_vars.add(deltaPeaks)

    ### now create background parametrization
    slope = RooRealVar("slope", "slope", 0.5, -10, 10)
    bkg_exp = RooExponential("bkg_exp", "exponential", slope, tagged_mass)
    pol_c1 = RooRealVar("p1", "coeff x^0 term", 0.5, -10, 10)
    pol_c2 = RooRealVar("p2", "coeff x^1 term", 0.5, -10, 10)
    bkg_pol = RooChebychev("bkg_pol", "2nd order pol", tagged_mass,
                           RooArgList(pol_c1, pol_c2))

    nsig = RooRealVar("Yield", "signal frac", nrt_mc.n + nwt_mc.n, 0, 1000000)
    nbkg = RooRealVar("nbkg", "bkg fraction", 1000, 0, 550000)

    print nsig.getVal()
    fitFunction = RooAddPdf("fitfunction", "fit function",
                            RooArgList(c_signalFunction, bkg_pol),
                            RooArgList(nsig, nbkg))

    pars_to_tune = [
        sigmawt, alphawt1, alphawt2, nwt1, nwt2, alphart1, alphart2, nrt1, nrt2
    ]
    if ibin < 5:
        pars_to_tune.append(sigmart)
    else:
        pars_to_tune.append(sigmart1)
        pars_to_tune.append(sigmart2)

    ## add toy bkg
    for itoy in range(args.ntoys):
        data = deepcopy(signaldata)
        toy_bkg = generateBkg(tagged_mass, ibin, n_bkg)
        data.append(toy_bkg)
        print 'toy number', itoy

        for ipar in pars_to_tune:
            ipar.setVal(pars_init_vals[ipar.GetName()])


#         r = fitFunction.fitTo(data,
#                               RooFit.Extended(True),
#                               RooFit.Range("full"),
#                               ROOT.RooFit.Constrain(c_vars),
# #                               ROOT.RooFit.Minimizer("Minuit2","migrad"),
#                               ROOT.RooFit.Hesse(True),
#                               ROOT.RooFit.Strategy(2),
#                               ROOT.RooFit.Minos(False),
#                              )
# #         print 'fit with Hesse strategy 2 done, now Minos'
        r = fitFunction.fitTo(
            data,
            RooFit.Extended(True),
            RooFit.Save(),
            RooFit.Range("full"),
            RooFit.Verbose(False),
            ROOT.RooFit.Constrain(c_vars),
            #                           ROOT.RooFit.Minimizer("Minuit2","migrad"),
            #                           ROOT.RooFit.Hesse(True),
            ROOT.RooFit.Strategy(2),
            ROOT.RooFit.Minos(True),
        )
        # #         r.Print()
        # #     r.correlationMatrix().Print()
        fitStats['data%s_itoy%s' % (ibin, itoy)] = r.status()
        covStats['data%s_itoy%s' % (ibin, itoy)] = r.covQual()
        frame = tagged_mass.frame(RooFit.Range("full"))
        data.plotOn(frame, RooFit.Binning(nbins), RooFit.MarkerSize(.7))
        fitFunction.plotOn(frame, RooFit.NormRange("full"),
                           RooFit.Range("full"))

        ## evaluate sort of chi2 and save number of RT/WT events
        observables = RooArgSet(tagged_mass)
        flparams = fitFunction.getParameters(observables)
        nparam = int(
            flparams.selectByAttrib("Constant", ROOT.kFALSE).getSize())
        pdfstring = "fitfunction_Norm[tagged_mass]_Range[full]_NormRange[full]"
        chi2s['data%s_itoy%s' % (ibin, itoy)] = frame.chiSquare(
            pdfstring, "h_fulldata", nparam)
        frame.addObject(_writeChi2(chi2s['data%s_itoy%s' % (ibin, itoy)]))

        ## save plot only if 1 toy is run
        if args.ntoys == 1:

            drawPdfComponents(fitFunction,
                              frame,
                              ROOT.kAzure,
                              RooFit.NormRange("full"),
                              RooFit.Range("full"),
                              isData=True)
            fitFunction.paramOn(frame, RooFit.Layout(0.62, 0.86, 0.89))
            #             parList = RooArgSet (nsig, nbkg, meanrt, meanwt, alphart1, alphart2, meanwt, sigmawt)
            #             if ibin < 5 :
            #                 parList.add(sigmart)
            #             else:
            #                 parList.add(sigmart1)
            #                 parList.add(sigmart2)
            #             parList.add(alphawt1)
            #             parList.add(alphawt2)
            #             parList.add(nwt1)
            #             parList.add(nwt2)
            #             parList.add(frt)
            #             fitFunction.paramOn(frame, RooFit.Parameters(parList), RooFit.Layout(0.62,0.86,0.89))

            frame.Draw()
            niceFrame(frame, '')
            frame.addObject(_writeFitStatus(r))

            c1 = ROOT.TCanvas()
            upperPad = ROOT.TPad('upperPad', 'upperPad', 0., 0.35, 1., 1.)
            lowerPad = ROOT.TPad('lowerPad', 'lowerPad', 0., 0.0, 1., 0.345)
            upperPad.SetBottomMargin(0.012)
            lowerPad.SetTopMargin(0)
            lowerPad.SetBottomMargin(0.2)

            upperPad.Draw()
            lowerPad.Draw()
            upperPad.cd()
            frame.Draw()
            if not args.year == 'test':
                writeCMS(frame, args.year,
                         [q2binning[ibin], q2binning[ibin + 1]], 1)
            frame.Draw()

            ## add plot of pulls
            lowerPad.cd()
            hpull = frame.pullHist("h_fulldata", pdfstring)
            frame2 = tagged_mass.frame(RooFit.Range("full"), RooFit.Title(''))
            frame2.addPlotable(hpull, "P")
            niceFrameLowerPad(frame2, 'pull')
            frame2.Draw()
            line = ROOT.TLine(5.0, 1, 5.6, 1)
            line.SetLineColor(ROOT.kGreen + 3)
            line.Draw()

            for ilog in [True, False]:
                upperPad.SetLogy(ilog)
                c1.SaveAs(
                    'fit_results_mass_checkOnMC/toybkg/save_fit_data_%s_%s_nbkg%s_LMNR_Final%s_%s_update_pol2bkg.pdf'
                    % (ibin, args.year, n_bkg, '_logScale' * ilog, itoy))

        out_f.cd()
        #         r.Write('results_data_%s_ntoy%s'%(ibin,itoy))

        ## compare nkbg fitted w/ original value
        nbkgs['data%s_itoy%s' %
              (ibin, itoy)] = (toy_bkg.sumEntries() -
                               nbkg.getVal()) / toy_bkg.sumEntries()
        nsigs['data%s_itoy%s' %
              (ibin, itoy)] = (n_realsignal - nsig.getVal()) / n_realsignal
def fitData(fulldata, ibin, w):

    cut  = cut_base + '&& (mumuMass*mumuMass > %s && mumuMass*mumuMass < %s)'%(q2binning[ibin], q2binning[ibin+1])
    data = fulldata.reduce(RooArgSet(tagged_mass,mumuMass,mumuMassE), cut)

    nrt_mc = _getFittedVar("nRT_%s"%ibin, w)
    nwt_mc = _getFittedVar("nWT_%s"%ibin, w)
    fraction = nrt_mc / (nrt_mc + nwt_mc)
    print 'mistag fraction on MC for bin ', ibin , ' : ' , fraction.n , '+/-', fraction.s 
    
    ### creating RT component
    w.loadSnapshot("reference_fit_RT_%s"%ibin)
    mean_rt      = w.var("mean^{RT%s}"%ibin)
    sigma_rt1    = w.var("#sigma_{RT1}^{%s}"%ibin)

    sigma_rt2 = RooRealVar()
    alpha_rt1 = RooRealVar()
    alpha_rt2 = RooRealVar()
    n_rt1     = RooRealVar()
    n_rt2     = RooRealVar()
    f1rt     = RooRealVar()

    ## double cb fast
    if ibin < 4:
        alpha_rt1    = w.var("#alpha_{RT1}^{%s}"%ibin)
        alpha_rt2    = w.var("#alpha_{RT2}^{%s}"%ibin)
        n_rt1        = w.var("n_{RT1}^{%s}"%ibin)
        n_rt2        = w.var("n_{RT2}^{%s}"%ibin)

    ## double cb old
    else:
        sigma_rt2    = w.var("#sigma_{RT2}^{%s}"%ibin)
        alpha_rt1    = w.var("#alpha_{RT1}^{%s}"%ibin)
        alpha_rt2    = w.var("#alpha_{RT2}^{%s}"%ibin)
        n_rt1        = w.var("n_{RT1}^{%s}"%ibin)
        n_rt2        = w.var("n_{RT2}^{%s}"%ibin)
        f1rt        = w.var("f^{RT%s}"%ibin)

    theRTgauss  = w.pdf("doublecb_RT%s"%ibin)   

    ### creating WT component
    w.loadSnapshot("reference_fit_WT_%s"%ibin)
    mean_wt      = w.var("mean_{WT}^{%s}"%ibin)
    sigma_wt     = w.var("#sigma_{WT1}^{%s}"%ibin)
    alpha_wt1    = w.var("#alpha_{WT1}^{%s}"%ibin)
    alpha_wt2    = w.var("#alpha_{WT2}^{%s}"%ibin)
    n_wt1        = w.var("n_{WT1}^{%s}"%ibin)
    n_wt2        = w.var("n_{WT2}^{%s}"%ibin)
    theWTgauss  = w.pdf("doublecb_%s"%ibin)   

    ### creating variable for the difference between the two peaks
    deltaPeaks = RooFormulaVar("deltaPeaks%s"%ibin, "@0 - @1", RooArgList(mean_rt, mean_wt))  
    frt              = RooRealVar ("F_{RT}%s"%ibin   , "frt"             , fraction.n , 0, 1)
    signalFunction   = RooAddPdf  ("sumgaus%s"%ibin  , "rt+wt"           , RooArgList(theRTgauss,theWTgauss), RooArgList(frt))

    ### now create background parametrization
    slope         = RooRealVar    ("slope_%s"%ibin   , "slope"           ,    0.5,   -10, 10);
    bkg_exp       = RooExponential("bkg_exp%s"%ibin , "exponential"     ,  slope,   tagged_mass  );
    pol_c1        = RooRealVar    ("p1_%s"%ibin      , "coeff x^0 term"  ,    0.5,   -10, 10);
    pol_c2        = RooRealVar    ("p2_%s"%ibin      , "coeff x^1 term"  ,    0.5,   -10, 10);
    bkg_pol       = RooPolynomial ("bkg_pol%s"%ibin  , "2nd order pol"   ,  tagged_mass, RooArgList(pol_c1, pol_c2));

    fsig          = RooRealVar("fsig%s"%ibin   , "fsig"    ,     0.9,     0,   1);
#     nsig          = RooRealVar("Yield%s"%ibin  , "signal frac"    ,     1000,     0,   10000);
#     nbkg          = RooRealVar("nbkg%s"%ibin   , "bkg fraction"   ,     1000,     0,   500000);
    nsig        = RooRealVar("Yield"         , "signal frac"   ,   600000,     0,   5000000);
    nbkg        = RooRealVar("nbkg"          , "bkg fraction"  ,   100000,     0,   2000000);

#     if ibin==4:  
#         nsig.setRange(500000,1500000)
#         nsig.setVal(900000)
#         nbkg.setRange(80000,1000000)
#         nbkg.setVal(100000)


    ### creating constraints
    c_vars = RooArgSet()
    c_pdfs = RooArgSet()

    c_sigma_rt1   = _constrainVar(sigma_rt1, 1)
    c_alpha_rt1   = _constrainVar(alpha_rt1, 1)
    c_alpha_rt2   = _constrainVar(alpha_rt2, 1)
    c_n_rt1       = _constrainVar(n_rt1, 1)
    c_n_rt2       = _constrainVar(n_rt2, 1)

    c_sigma_wt    = _constrainVar(sigma_wt,  1)
    c_alpha_wt1   = _constrainVar(alpha_wt1, 1)
    c_alpha_wt2   = _constrainVar(alpha_wt2, 1)
    c_n_wt1       = _constrainVar(n_wt1, 1)
    c_n_wt2       = _constrainVar(n_wt2, 1)


    if ibin < 4:
        c_pdfs = RooArgSet(c_sigma_rt1, c_alpha_rt1, c_alpha_rt2, c_n_rt1, c_n_rt2)
        c_vars = RooArgSet(sigma_rt1,     alpha_rt1,   alpha_rt2,   n_rt1,   n_rt2)
    else:
        c_sigma_rt2   = _constrainVar(sigma_rt2, 1)
        c_pdfs = RooArgSet(c_sigma_rt1, c_sigma_rt2, c_alpha_rt1, c_alpha_rt2, c_n_rt1, c_n_rt2)
        c_vars = RooArgSet(  sigma_rt1,   sigma_rt2,   alpha_rt1,   alpha_rt2,   n_rt1,   n_rt2)
    
    c_pdfs.add(c_sigma_wt);  c_vars.add(sigma_wt)
    c_pdfs.add(c_alpha_wt1); c_vars.add(alpha_wt1)
    c_pdfs.add(c_alpha_wt2); c_vars.add(alpha_wt2)
    c_pdfs.add(c_n_wt1);     c_vars.add(n_wt1)
    c_pdfs.add(c_n_wt2);     c_vars.add(n_wt2)

    c_deltaPeaks = RooGaussian("c_deltaPeaks%s"%ibin , "c_deltaPeaks", deltaPeaks, ROOT.RooFit.RooConst( deltaPeaks.getVal() ), 
                                ROOT.RooFit.RooConst( 0.0005 )  ## value to be checked
                                ) 
    c_pdfs.add(c_deltaPeaks)
    c_vars.add(deltaPeaks)
    c_frt            = RooGaussian("c_frt%s"%ibin    , "c_frt"           , frt,  ROOT.RooFit.RooConst(fraction.n) , ROOT.RooFit.RooConst(frt_sigma[ibin]) )
    c_pdfs.add(c_frt)
    c_vars.add(frt)
    
    constr_list = RooArgList(c_pdfs)
    constr_list.add(signalFunction)
    c_signalFunction = RooProdPdf ("c_signalFunction", "c_signalFunction", constr_list)     

#     mean        = RooRealVar ("mass"         , "mean"          ,  B0Mass_,   3,    7, "GeV")
#     sigma       = RooRealVar ("#sigma_{1}"   , "sigma"         ,  0.028,     0,   10, "GeV")
#     signalGauss = RooGaussian("signalGauss"  , "signal gauss"  ,  tagged_mass,  mean,sigma)
#     
#     sigma2       = RooRealVar ("#sigma_{2}"    , "sigma2"         ,  0.048,     0,   0.07, "GeV")
#     signalGauss2 = RooGaussian("signalGauss2"  , "signal gauss2"  ,  tagged_mass,  mean,sigma2)
#     f1           = RooRealVar ("f1"            , "f1"             ,  0.8  ,     0.,   1.)
#     gaus         = RooAddPdf  ("gaus"          , "gaus1+gaus2"    , RooArgList(signalGauss,signalGauss2), RooArgList(f1))
    
#     pol_c1      = RooRealVar ("p1"           , "coeff x^0 term",    0.5,   -10, 10);
#     pol_c2      = RooRealVar ("p2"           , "coeff x^1 term",    0.5,   -10, 10);
#     pol_c3      = RooRealVar ("p3"           , "coeff x^2 term",    0.5,   -10, 10);
#     slope       = RooRealVar ("slope"        , "slope"         ,    0.5,   -10, 10);
#     bkg_exp     = RooExponential("bkg_exp"   , "exponential"   ,  slope,   tagged_mass  );
#     bkg_pol     = RooChebychev("bkg_pol"     , "2nd order pol" ,  tagged_mass, RooArgList(pol_c1,pol_c2));

    fitFunction = RooAddPdf ("fitfunction%s"%ibin , "fit function"  ,  RooArgList(c_signalFunction, bkg_exp), RooArgList(nsig, nbkg))
    
#     r = fitFunction.fitTo(data, 
# #                           RooFit.Extended(True), 
#                           RooFit.Range("full"), 
#                           ROOT.RooFit.Constrain(c_vars),
#                           ROOT.RooFit.Minimizer("Minuit2","migrad"),
#                           ROOT.RooFit.Hesse(True),
#                           ROOT.RooFit.Strategy(2),
#                           ROOT.RooFit.Minos(False),
#                          )
    print 'fit with Hesse strategy 2 done, now Minos'    
    r = fitFunction.fitTo(data, 
                          RooFit.Extended(True), 
                          RooFit.Save(), 
                          RooFit.Range("full"), 
                          RooFit.Verbose(False),
                          ROOT.RooFit.Constrain(c_vars),
#                           ROOT.RooFit.Minimizer("Minuit2","migrad"),
#                           ROOT.RooFit.Hesse(True),
#                           ROOT.RooFit.Strategy(2),
#                           ROOT.RooFit.Minos(False),
                         )

    r.Print()
    r.correlationMatrix().Print()
    fitStats['data%s'%(ibin)] = r.status()
    covStats['data%s'%(ibin)] = r.covQual()
    frame = tagged_mass.frame( RooFit.Range("full") )
    data.plotOn(frame, RooFit.Binning(nbins), RooFit.MarkerSize(.7))
    fitFunction.plotOn(frame, RooFit.NormRange("full"), RooFit.Range("full"))

    ## evaluate sort of chi2 and save number of RT/WT events
    observables = RooArgSet(tagged_mass)
    flparams = fitFunction.getParameters(observables)
    nparam = int(flparams.selectByAttrib("Constant",ROOT.kFALSE).getSize())
    pdfstring = "fitfunction%s_Norm[tagged_mass]_Range[full]_NormRange[full]"%ibin
    chi2s['data%s'%ibin] = frame.chiSquare(pdfstring, "h_fulldata",  nparam)
    frame. addObject(_writeChi2( chi2s['data%s'%ibin] ))

    drawPdfComponents(fitFunction, frame, ROOT.kAzure, RooFit.NormRange("full"), RooFit.Range("full"), isData = True)
#     fitFunction.paramOn(frame, RooFit.Layout(0.62,0.86,0.89))

    parList = RooArgSet (nsig, mean_rt, sigma_rt, alpha_rt1, alpha_rt2, n_rt1, n_rt2, mean_wt, sigma_wt)
#     parList.add(alphawt1)
#     parList.add(alphawt2)
#     parList.add(nwt1)
#     parList.add(nwt2)
    parList.add(frt)
    fitFunction.paramOn(frame, RooFit.Parameters(parList), RooFit.Layout(0.62,0.86,0.89))

    frame.Draw()
    niceFrame(frame, '')
    frame. addObject(_writeFitStatus(r))

    c1 = ROOT.TCanvas() 
    upperPad = ROOT.TPad('upperPad' , 'upperPad' , 0., 0.35 , 1.,  1.    )  
    lowerPad = ROOT.TPad('lowerPad' , 'lowerPad' , 0., 0.0  , 1.,  0.345 )  
    upperPad.SetBottomMargin(0.012)
    lowerPad.SetTopMargin(0)
    lowerPad.SetBottomMargin(0.2)

    upperPad.Draw(); lowerPad.Draw()
    upperPad.cd()
    frame.Draw()
    if not args.year=='test':  writeCMS(frame, args.year, [ q2binning[ibin], q2binning[ibin+1] ], 0)
    frame.Draw()

    ## add plot of pulls
    lowerPad.cd()
    hpull = frame.pullHist("h_fulldata", pdfstring)
    frame2 =  tagged_mass.frame(RooFit.Range("full"), RooFit.Title(''))
    frame2.addPlotable(hpull,"P") 
    niceFrameLowerPad(frame2, 'pull')
    frame2.Draw()
    line = ROOT.TLine(5.0,1,5.6,1)
    line.SetLineColor(ROOT.kGreen+3)
    line.Draw()

    for ilog in [True,False]:
        upperPad.SetLogy(ilog)
        c1.SaveAs('fit_results_mass/save_fit_data_%s_%s_LMNR_Update%s_newSigmaFRT_pars_Jpsi.pdf'%(ibin, args.year, '_logScale'*ilog))


    out_f.cd()
    r.Write('results_data_%s'%(ibin))

    params = fitFunction.getParameters(RooArgSet(tagged_mass)) 
    out_w.saveSnapshot("reference_fit_data_%s"%(ibin),params,ROOT.kTRUE) 
    getattr(out_w, 'import')(fitFunction)
def get_fitted_normalisation_from_ROOT(channel, input_files, variable,
                                       met_type, b_tag_bin):
    results = {}
    initial_values = {}
    templates = {}

    for variable_bin in variable_bins_ROOT[variable]:
        histograms = get_histograms(channel,
                                    input_files,
                                    variable=variable,
                                    met_type=met_type,
                                    variable_bin=variable_bin,
                                    b_tag_bin=b_tag_bin,
                                    rebin=measurement_config.rebin)
        # create signal histograms
        h_eta_signal = histograms['TTJet'] + histograms['SingleTop']

        N_ttbar_before_fit = histograms['TTJet'].Integral()
        N_SingleTop_before_fit = histograms['SingleTop'].Integral()
        N_vjets_before_fit = histograms['V+Jets'].Integral()
        N_qcd_before_fit = histograms['QCD'].Integral()
        N_signal_before_fit = N_ttbar_before_fit + N_SingleTop_before_fit

        N_ttbar_error_before_fit = sum(histograms['TTJet'].errors())
        N_SingleTop_error_before_fit = sum(histograms['SingleTop'].errors())
        N_vjets_error_before_fit = sum(histograms['V+Jets'].errors())
        N_QCD_error_before_fit = sum(histograms['QCD'].errors())

        if (N_SingleTop_before_fit != 0):
            TTJet_SingleTop_ratio = N_ttbar_before_fit / N_SingleTop_before_fit
        else:
            print 'Bin ', variable_bin, ': ttbar/singleTop ratio undefined for %s channel! Setting to 0.' % channel
            TTJet_SingleTop_ratio = 0

        leptonAbsEta = RooRealVar("leptonAbsEta", "leptonAbsEta", 0., 2.4)
        # this has to move to dps.utils.Fitting.py
        vars = RooArgList()
        vars.add(leptonAbsEta)
        vars_set = RooArgSet()
        vars_set.add(leptonAbsEta)
        n_event_obs = histograms['data'].Integral()

        lowerBound = 0.
        upperBound = n_event_obs + 10 * sqrt(n_event_obs)
        n_init = n_event_obs / 2.

        data = RooDataHist("data", "dataset with leptonAbsEta", vars,
                           histograms['data'])
        rh_vj = RooDataHist("rh_vj", "vj", vars, histograms['V+Jets'])
        rh_qcd = RooDataHist("rh_qcd", "qcd", vars, histograms['QCD'])
        rh_signal = RooDataHist("rh_signal", "signal", vars, h_eta_signal)

        pdf_vj = RooHistPdf("pdf_vj", "V+Jets pdf", vars_set, rh_vj, 0)
        pdf_qcd = RooHistPdf("pdf_qcd", "QCD pdf ", vars_set, rh_qcd, 0)
        pdf_signal = RooHistPdf("pdf_signal", "single top pdf", vars_set,
                                rh_signal, 0)

        # RooRealVar(const char *name, const char *title, Double_t value, Double_t minValue, Double_t maxValue, const char *unit) :
        nSignal = RooRealVar("nSignal", "number of single top + ttbar events",
                             N_signal_before_fit, lowerBound, upperBound,
                             "event")
        nvj = RooRealVar("nvj", "number of V+Jets bgnd events",
                         N_vjets_before_fit, lowerBound, upperBound, "event")
        nqcd = RooRealVar("nqcd", "number of QCD bgnd events",
                          N_QCD_error_before_fit, lowerBound, upperBound,
                          "event")

        model = RooAddPdf("model", "sig+vj+qcd",
                          RooArgList(pdf_signal, pdf_vj, pdf_qcd),
                          RooArgList(nSignal, nvj, nqcd))
        vj_constraint = RooGaussian("nvj_constraint", "nvj_constraint", nvj,
                                    RooFit.RooConst(N_vjets_before_fit),
                                    RooFit.RooConst(0.5 * N_vjets_before_fit))
        qcd_constraint = RooGaussian("nqcd_constraint", "nqcd_constraint",
                                     nqcd, RooFit.RooConst(N_qcd_before_fit),
                                     RooFit.RooConst(2 * N_qcd_before_fit))
        model_with_constraints = RooProdPdf(
            "model_with_constraints", "model with gaussian constraints",
            RooArgSet(model, vj_constraint, qcd_constraint), RooLinkedList())
        model_with_constraints.fitTo(data, RooFit.Minimizer(
            "Minuit2",
            "Migrad"))  #WARNING: number of cores changes the results!!!
        #         nll = model.createNLL(data, RooFit.NumCPU(2))
        #         RooMinuit(nll).migrad()
        #         frame1 = nSignal.frame(RooFit.Bins(100), RooFit.Range(lowerBound, n_event_obs), RooFit.Title("LL and profileLL in nSignal"))
        #         nll.plotOn(frame1, RooFit.ShiftToZero())
        #         frame2 = nvj.frame(RooFit.Bins(100), RooFit.Range(lowerBound, n_event_obs), RooFit.Title("LL and profileLL in nvj"))
        #         nll.plotOn(frame2, RooFit.ShiftToZero())
        #         frame3 = nqcd.frame(RooFit.Bins(100), RooFit.Range(lowerBound, n_event_obs), RooFit.Title("LL and profileLL in nqcd"))
        #         nll.plotOn(frame3, RooFit.ShiftToZero())
        #
        #         pll_nSignal = nll.createProfile(nSignal)
        #         pll_nSignal.plotOn(frame1, RooFit.LineColor(2))
        #         frame1.SetMinimum(0)
        #         frame1.SetMaximum(3)
        #
        #         pll_nvj = nll.createProfile(nvj)
        #         pll_nvj.plotOn(frame2, RooFit.LineColor(2))
        #         frame2.SetMinimum(0)
        #         frame2.SetMaximum(3)
        #
        #         pll_nqcd = nll.createProfile(nqcd)
        #         pll_nqcd.plotOn(frame3, RooFit.LineColor(2))
        #         frame3.SetMinimum(0)
        #         frame3.SetMaximum(3)
        #         c = TCanvas("profilell","profilell",1200, 400)
        #         c.Divide(3)
        #         c.cd(1)
        #         frame1.Draw()
        #         c.cd(2)
        #         frame2.Draw()
        #         c.cd(3)
        #         frame3.Draw()
        #         c.SaveAs('profileLL.png')
        #         model.fitTo(data, RooFit.Minimizer("Minuit2", "Migrad"), RooFit.NumCPU(1))#WARNING: number of cores changes the results!!!
        fit_results = {}
        fit_results['signal'] = (nSignal.getVal(), nSignal.getError())
        fit_results['QCD'] = ufloat(nqcd.getVal(), nqcd.getError())
        fit_results['V+Jets'] = ufloat(nvj.getVal(), nvj.getError())

        N_ttbar, N_SingleTop = decombine_result(fit_results['signal'],
                                                TTJet_SingleTop_ratio)
        fit_results['signal'] = ufloat(nSignal.getVal(), nSignal.getError())
        fit_results['TTJet'] = ufloat(N_ttbar)
        fit_results['SingleTop'] = ufloat(N_SingleTop)

        if results == {}:  # empty
            for sample in fit_results.keys():
                results[sample] = [fit_results[sample]]
        else:
            for sample in fit_results.keys():
                results[sample].append(fit_results[sample])

    return results, None, None
Exemplo n.º 9
0
def fitData(fulldata, ibin, nRT_fromMC, nWT_fromMC):

    cut = cut_base + '&& (mumuMass*mumuMass > %s && mumuMass*mumuMass < %s)' % (
        q2binning[ibin], q2binning[ibin + 1])
    fulldata_v2 = fulldata.reduce(
        RooArgSet(tagged_mass, mumuMass, mumuMassE, randVar), cut)

    ## reduce to data-like statistics
    nDataEntries = fulldata_v2.sumEntries()
    nDesired = n_bin[ibin] / nDataEntries
    cut = 'rand < %f' % nDesired
    data = fulldata_v2.reduce(RooArgSet(tagged_mass, mumuMass, mumuMassE), cut)

    fraction = dict_s_rt[ibin] / (dict_s_rt[ibin] + dict_s_wt[ibin])
    print 'mistag fraction on MC for bin ', ibin, ' : ', fraction.n, '+/-', fraction.s

    ### creating RT component
    w.loadSnapshot("reference_fit_RT_%s" % ibin)
    mean_rt = w.var("mean_{RT}^{%s}" % ibin)
    sigma_rt1 = w.var("#sigma_{RT1}^{%s}" % ibin)

    sigma_rt2 = RooRealVar()
    alpha_rt1 = RooRealVar()
    alpha_rt2 = RooRealVar()
    n_rt1 = RooRealVar()
    n_rt2 = RooRealVar()
    f1rt = RooRealVar()

    ## double cb fast
    if ibin < 4:
        alpha_rt1 = w.var("#alpha_{RT1}^{%s}" % ibin)
        alpha_rt2 = w.var("#alpha_{RT2}^{%s}" % ibin)
        n_rt1 = w.var("n_{RT1}^{%s}" % ibin)
        n_rt2 = w.var("n_{RT2}^{%s}" % ibin)

    ## double cb old
    else:
        sigma_rt2 = w.var("#sigma_{RT2}^{%s}" % ibin)
        alpha_rt1 = w.var("#alpha_{RT1}^{%s}" % ibin)
        alpha_rt2 = w.var("#alpha_{RT2}^{%s}" % ibin)
        n_rt1 = w.var("n_{RT1}^{%s}" % ibin)
        n_rt2 = w.var("n_{RT2}^{%s}" % ibin)
        f1rt = w.var("f^{RT%s}" % ibin)

    theRTgauss = w.pdf("doublecb_RT%s" % ibin)

    ### creating WT component
    w.loadSnapshot("reference_fit_WT_%s" % ibin)
    mean_wt = w.var("mean_{WT}^{%s}" % ibin)
    sigma_wt = w.var("#sigma_{WT1}^{%s}" % ibin)
    alpha_wt1 = w.var("#alpha_{WT1}^{%s}" % ibin)
    alpha_wt2 = w.var("#alpha_{WT2}^{%s}" % ibin)
    n_wt1 = w.var("n_{WT1}^{%s}" % ibin)
    n_wt2 = w.var("n_{WT2}^{%s}" % ibin)
    theWTgauss = w.pdf("doublecb_%s" % ibin)

    ### creating variable for the difference between the two peaks
    deltaPeaks = RooFormulaVar("deltaPeaks", "@0 - @1",
                               RooArgList(mean_rt, mean_wt))
    frt = RooRealVar("F_{RT}", "frt", fraction.n, 0, 1)
    signalFunction = RooAddPdf("sumgaus", "rt+wt",
                               RooArgList(theRTgauss, theWTgauss),
                               RooArgList(frt))

    ### now create background parametrization
    slope = RooRealVar("slope", "slope", 0.5, -10, 10)
    bkg_exp = RooExponential("bkg_exp", "exponential", slope, tagged_mass)
    pol_c1 = RooRealVar("p1", "coeff x^0 term", 0.5, -10, 10)
    pol_c2 = RooRealVar("p2", "coeff x^1 term", 0.5, -10, 10)
    bkg_pol = RooChebychev("bkg_pol", "2nd order pol", tagged_mass,
                           RooArgList(pol_c1))

    nsig = RooRealVar("Yield", "signal frac", nRT_fromMC + nWT_fromMC, 0,
                      1000000)
    nbkg = RooRealVar("nbkg", "bkg fraction", 1000, 0, 550000)

    print nsig.getVal()

    ### creating constraints
    c_vars = RooArgSet()
    c_pdfs = RooArgSet()

    c_sigma_rt1 = _constrainVar(sigma_rt1, 1)
    c_alpha_rt1 = _constrainVar(alpha_rt1, 1)
    c_alpha_rt2 = _constrainVar(alpha_rt2, 1)
    c_n_rt1 = _constrainVar(n_rt1, 1)
    c_n_rt2 = _constrainVar(n_rt2, 1)

    c_sigma_wt = _constrainVar(sigma_wt, 1)
    c_alpha_wt1 = _constrainVar(alpha_wt1, 1)
    c_alpha_wt2 = _constrainVar(alpha_wt2, 1)
    c_n_wt1 = _constrainVar(n_wt1, 1)
    c_n_wt2 = _constrainVar(n_wt2, 1)

    if ibin < 4:
        c_pdfs = RooArgSet(c_sigma_rt1, c_alpha_rt1, c_alpha_rt2, c_n_rt1,
                           c_n_rt2)
        c_vars = RooArgSet(sigma_rt1, alpha_rt1, alpha_rt2, n_rt1, n_rt2)
    else:
        c_sigma_rt2 = _constrainVar(sigma_rt2, 1)
        c_f1rt = _constrainVar(f1rt, 1)
        c_pdfs = RooArgSet(c_sigma_rt1, c_sigma_rt2, c_alpha_rt1, c_alpha_rt2,
                           c_n_rt1, c_n_rt2, c_f1rt)
        c_vars = RooArgSet(sigma_rt1, sigma_rt2, alpha_rt1, alpha_rt2, n_rt1,
                           n_rt2, f1rt)

    c_pdfs.add(c_sigma_wt)
    c_vars.add(sigma_wt)
    c_pdfs.add(c_alpha_wt1)
    c_vars.add(alpha_wt1)
    c_pdfs.add(c_alpha_wt2)
    c_vars.add(alpha_wt2)
    c_pdfs.add(c_n_wt1)
    c_vars.add(n_wt1)
    c_pdfs.add(c_n_wt2)
    c_vars.add(n_wt2)

    ### creating constraints for the difference between the two peaks
    c_deltaPeaks = RooGaussian(
        "c_deltaPeaks",
        "c_deltaPeaks",
        deltaPeaks,
        ROOT.RooFit.RooConst(deltaPeaks.getVal()),
        ROOT.RooFit.RooConst(0.0005)  ## value to be checked
    )
    c_pdfs.add(c_deltaPeaks)
    c_vars.add(deltaPeaks)

    c_frt = RooGaussian("c_frt%s" % ibin, "c_frt", frt,
                        ROOT.RooFit.RooConst(fraction.n),
                        ROOT.RooFit.RooConst(frt_sigma[ibin]))
    c_pdfs.add(c_frt)
    c_vars.add(frt)

    constr_list = RooArgList(c_pdfs)
    constr_list.add(signalFunction)
    c_signalFunction = RooProdPdf("c_signalFunction", "c_signalFunction",
                                  constr_list)

    fitFunction = c_signalFunction

    r = fitFunction.fitTo(
        data,
        #                           RooFit.Extended(True),
        RooFit.Range("full"),
        ROOT.RooFit.Constrain(c_vars),
        ROOT.RooFit.Minimizer("Minuit2", "migrad"),
        ROOT.RooFit.Hesse(True),
        ROOT.RooFit.Strategy(2),
        ROOT.RooFit.Minos(False),
    )
    print 'fit with Hesse strategy 2 done, now Minos'
    r = fitFunction.fitTo(
        data,
        #                           RooFit.Extended(True),
        RooFit.Save(),
        RooFit.Range("full"),
        RooFit.Verbose(False),
        ROOT.RooFit.Constrain(c_vars),
        #                           ROOT.RooFit.Minimizer("Minuit2","migrad"),
        #                           ROOT.RooFit.Hesse(True),
        ROOT.RooFit.Strategy(2),
        ROOT.RooFit.Minos(True),
    )

    r.Print()
    r.correlationMatrix().Print()
    fitStats['data%s' % (ibin)] = r.status()
    covStats['data%s' % (ibin)] = r.covQual()
    frame = tagged_mass.frame(RooFit.Range("full"))
    data.plotOn(frame, RooFit.Binning(nbins), RooFit.MarkerSize(.7))
    fitFunction.plotOn(frame, RooFit.NormRange("full"), RooFit.Range("full"))

    ## evaluate sort of chi2 and save number of RT/WT events
    observables = RooArgSet(tagged_mass)
    flparams = fitFunction.getParameters(observables)
    nparam = int(flparams.selectByAttrib("Constant", ROOT.kFALSE).getSize())
    pdfstring = "c_signalFunction_Norm[tagged_mass]_Range[full]_NormRange[full]"
    chi2s['data%s' % ibin] = frame.chiSquare(pdfstring, "h_fulldata", nparam)
    frame.addObject(_writeChi2(chi2s['data%s' % ibin]))

    drawPdfComponents(fitFunction,
                      frame,
                      ROOT.kAzure,
                      RooFit.NormRange("full"),
                      RooFit.Range("full"),
                      isData=True)

    parList = RooArgSet(nsig, mean_rt, mean_wt, sigma_rt1, sigma_rt2,
                        alpha_rt1, alpha_rt2, mean_wt, sigma_wt)
    parList.add(alpha_wt1)
    parList.add(alpha_wt2)
    parList.add(n_wt1)
    parList.add(n_wt2)
    parList.add(frt)
    fitFunction.paramOn(frame, RooFit.Parameters(parList),
                        RooFit.Layout(0.62, 0.86, 0.89))

    frame.Draw()
    niceFrame(frame, '')
    frame.addObject(_writeFitStatus(r))

    c1 = ROOT.TCanvas()
    upperPad = ROOT.TPad('upperPad', 'upperPad', 0., 0.35, 1., 1.)
    lowerPad = ROOT.TPad('lowerPad', 'lowerPad', 0., 0.0, 1., 0.345)
    upperPad.SetBottomMargin(0.012)
    lowerPad.SetTopMargin(0)
    lowerPad.SetBottomMargin(0.2)

    upperPad.Draw()
    lowerPad.Draw()
    upperPad.cd()
    frame.Draw()
    if not args.year == 'test':
        writeCMS(frame, args.year, [q2binning[ibin], q2binning[ibin + 1]], 1)
    frame.Draw()

    ## add plot of pulls
    lowerPad.cd()
    hpull = frame.pullHist("h_fulldata", pdfstring)
    frame2 = tagged_mass.frame(RooFit.Range("full"), RooFit.Title(''))
    frame2.addPlotable(hpull, "P")
    niceFrameLowerPad(frame2, 'pull')
    frame2.Draw()
    line = ROOT.TLine(5.0, 1, 5.6, 1)
    line.SetLineColor(ROOT.kGreen + 3)
    line.Draw()

    for ilog in [True, False]:
        upperPad.SetLogy(ilog)
        c1.SaveAs(
            'fit_results_mass_checkOnMC/save_fit_data_%s_%s_LMNR_newSigmaFRT_%s.pdf'
            % (ibin, args.year, '_logScale' * ilog))

    out_f.cd()
    r.Write('results_data_%s' % (ibin))

    params = fitFunction.getParameters(RooArgSet(tagged_mass))
    w.saveSnapshot("reference_fit_data_%s" % (ibin), params, ROOT.kTRUE)
    getattr(w, 'import')(fitFunction)
Exemplo n.º 10
0
# Get tree
rfile = get_file('data/smalltree-new-MC.root', 'read')
ftree = get_object('ftree', rfile)

# dt dataset for denom
trigger = 'HLT2Topo3BodyTOS'
triggerVar = RooRealVar(trigger, trigger, 0, 2)
cut = trigger+'>0'
tmpdtdataset = RooDataSet('dtdataset', 'dt dataset', RooArgSet(dt, triggerVar),
                            RooFit.Import(ftree), RooFit.Cut(cut))
dtdatahist = tmpdtdataset.binnedClone('dtdatahist', 'Binned dt')
dtdatahist = dtdatahist.reduce(dtargset)
errorPdf = RooHistPdf('errorPdf', 'Time error Hist PDF', dtargset, dtdatahist)

decaywdt = RooProdPdf('decaywdt', 'Decay function with dt distribution',
                      RooArgSet(errorPdf),
                      RooFit.Conditional(decayargset, timeargset))

del tmpdtdataset, dtdatahist

# Variable width binning
nbins1 = 50
nbins2 = 24
binedges1 = numpy.zeros(nbins1+1 , dtype=float)
binedges2 = numpy.zeros(nbins2+1 , dtype=float)
logwidth = log10((1E-2 + 2E-4) / 2E-4)

for i in range(nbins1+1):
    # binedges1[i] = 10**(log10(2E-4) + i*logwidth/nbins1)
    binedges1[i] = epsilon + i*1E-2/nbins1
    # print '%2d: %e' % (i, binedges1[i])