Esempio n. 1
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))
Esempio n. 2
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))
Esempio 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
Esempio 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
    }
Esempio n. 5
0
class CouplingsFitter2(object):
    def __init__(self):
        self.poiLabels = []
        self.BR = dict(b=0.577,
                       tau=0.063,
                       mu=2.2e-4,
                       c=2.91e-2,
                       g=8.57e-2,
                       gamma=3.82e-3,
                       W=0.215,
                       Z=0.0264,
                       t=0.0)

        self.poi = dict()
        self.poilabels = dict()
        self.constraint = dict()
        self.canvases = dict()
        self._keep = []

    def addPOI(self, poi, label='', minimum=-0.3, maximum=0.3):
        '''Add a parameter of interest.

        Example:
        
        addPOI('Z','Z',-0.05,0.05)
        ->
        Z[0,-0.05,0.05]
        # adds variable Z with value 0,
        # allow it to scale between -0.05 and 0.05
        '''
        self.poi[poi] = ROOT.RooRealVar(poi, poi, 0, minimum, maximum)
        if label == '':
            label = poi
        self.poilabels[poi] = label

    def createWidthDeviation(self):
        '''Compute the width deviation (denoted \kappa_H^2 by M.Peskin in arxiv 1312.4974).
        
        Note that we fit an additive modification of the coupling: (1 + dx)
        is therefore equal to kappa_x
        '''
        expr = '0'
        sumBR = sum(self.BR.values())
        pwidths = []
        for dcoupling, br in self.BR.iteritems():
            pwidth = None
            if dcoupling in self.poi:
                pwidth = str(
                    br /
                    sumBR) + "*(1+" + dcoupling + ")*(1+" + dcoupling + ")"
            else:
                # using sm partial width
                pwidth = str(br / sumBR)
            pwidths.append(pwidth)
        expr = '+'.join(pwidths)
        if 'inv' in self.poi:
            expr = '(' + expr + ')/(1.0-inv)'
        else:
            # setting invisible width to 0.
            expr = '(' + expr + ')'
        dependentlist = RooArgList()
        for dep in self.poi.values():
            dependentlist.add(dep)
        self.width = RooGenericPdf('width', 'width', expr, dependentlist)

    def addConstraint(self, name, expr, deplist, mean, sigma):
        '''Add a constraint on one of the observables
        
        For example, for ZH inclusive: 
        
        f.addConstraint('Zh','(1+Z)*(1+Z)','Z',1,0.004)  

        Z is an additive modification of the gZ coupling w/r to the standard model,
        so 1+Z = \kappa_Z
        
        Zh is the pdf of the ratio of the yield w/r to the one expected in the standard model.
        This pdf depends on Z, as (1+Z)*(1+Z).
        
        ZhObs is the measured value, here 1 so we assume that the SM yield is observed.
        
        The fit varies the parameter of interest Z, thus modifying the pdf,
        while ZhObs is fixed at 1. The likelihood of each value of Z is evaluated at ZhObs on the pdf.
        '''
        print 'constraint:', name
        print expr
        print deplist
        print mean, '+/-', sigma
        deps = self._getdeps(deplist)
        self.constraint[name] = GaussianConstraint(name, expr, deps, mean,
                                                   sigma)

    def _getdeps(self, deplist):
        depnames = deplist
        try:
            depnames = deplist.split(',')
        except:
            pass
        deps = []
        for dep in depnames:
            if dep == 'width':
                deps.append(self.width)
            else:
                deps.append(self.poi[dep])
        return deps

    def addChannel(self, name, mean, sigma, prod, decay=None):
        expr_prod = '(1+{prod})*(1+{prod})'.format(prod=prod)
        expr = expr_prod
        variables = [prod]
        if decay:
            if isinstance(decay, basestring):
                expr_decay = '(1+{decay})*(1+{decay})'.format(decay=decay)
                variables.append(decay)
            else:
                decay_exprs = []
                for decay, fraction in decay:
                    n = mean * fraction
                    decay_expr = '{n}*(1+{decay})*(1+{decay})'.format(
                        n=n, decay=decay)
                    decay_exprs.append(decay_expr)
                    variables.append(decay)
                expr_decay = ' + '.join(decay_exprs)
            expr = '{expr_prod}*({expr_decay})/width'.format(
                expr_prod=expr_prod,
                expr_decay=expr_decay,
            )
            variables.append('width')
        variables = list(set(variables))
        deplist = ','.join(variables)
        self.addConstraint(name, expr, deplist, mean, sigma)

    def addUniformConstraint(self, name, expr):
        '''Adds a uniform constraint with pdf name, and expression expr.
        
        For example:
        
        f.addPOI('inv','inv', 0, 0.01)
        f.addUniformConstraint('Zhinv','inv') ####->Means free floating

        inv (the invisible BR) is free to float between 0 and 1%
        '''
        deps = self._getdeps([expr])
        self.constraint[name] = UniformConstraint(name, expr, deps)

    def info(self):
        for name, constraint in sorted(self.constraint.iteritems()):
            constraint.info()

    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))

    def canvas(self, name, *args):
        canvas = self.canvases.setdefault(name,
                                          ROOT.TCanvas(name, name, *args))
        canvas.cd()
        return canvas

    def keep(self, obj):
        self._keep.append(obj)
        return obj

    def createSummary(self):
        #sample the covariance matrix for the width
        # ROOT.gStyle.SetOptTitle(0)
        ROOT.gStyle.SetStatW(0.4)
        ROOT.gStyle.SetStatH(0.4)

        self.graph_couplings = ROOT.TGraphAsymmErrors(len(self.poi) + 2)

        order_BR = ['Z', 'W', 'b', 'c', 'g', 'tau', 'mu', 'gamma', 'inv']

        for br in order_BR:
            if not self.poi.get(br, None):
                order_BR.remove(br)

        for i, poiname in enumerate(order_BR):
            poi = self.poi.get(poiname)
            self.graph_couplings.SetPoint(i, i + 0.5, poi.getVal())
            self.graph_couplings.SetPointError(i, 0.0, 0.0,
                                               -poi.getAsymErrorLo(),
                                               poi.getAsymErrorHi())

        print 'Sampling the covariance matrix to propagate error on width'

        self.h_width = ROOT.TH1F('h_width', 'width', 1000, 0.5, 1.5)
        ntoys = 10000
        for i in range(ntoys):
            randomizedPars = self.fit_result.randomizePars()
            for j in range(0, randomizedPars.getSize()):
                self.poi[randomizedPars.at(j).GetName()].setVal(
                    randomizedPars.at(j).getVal())
            self.h_width.Fill(self.width.getVal())
            for cstr in self.constraint.values():
                cstr.fill_pull()
        self.graph_couplings.SetMarkerStyle(20)
        self.graph_couplings.SetLineWidth(3)
        can_couplings = self.canvas('couplings')
        self.graph_couplings.Draw("AP")
        self.graph_couplings.GetYaxis().SetTitle("68% CL on d(A) ")
        self.graph_couplings.GetXaxis().SetNdivisions(0)
        l = self.keep(ROOT.TLine())
        l.SetLineColor(ROOT.kRed)
        l.SetLineWidth(3)
        l.DrawLine(0.0, 0.0, len(self.poi) + 1.5, 0)

        self.graph_couplings.SetPoint(len(self.poi), len(self.poi) + 0.5, 0.0)
        self.graph_couplings.SetPointError(
            len(self.poi), 0.0, 0.0,
            self.h_width.GetRMS() / self.h_width.GetMean(),
            self.h_width.GetRMS() / self.h_width.GetMean())

        for i, poiname in enumerate(order_BR + ['#Gamma_{T}']):
            label = self.poilabels.get(poiname, poiname)
            tex = self.keep(
                ROOT.TLatex(i + 0.5,
                            0.95 * self.graph_couplings.GetYaxis().GetXmin(),
                            label))
            tex.Draw()

        print """
###############################################################
###############################################################
###############################################################
                         RESULTS
###############################################################
###############################################################
############################################################### 
              """

        print 'RESULTS FOR THE CONFIDENCE INTERVALS------>'
        for name in order_BR:
            poi = self.poi[name]
            poiLabel = self.poilabels.get(name, name)
            mind = poi.getAsymErrorLo() * 100
            maxd = poi.getAsymErrorHi() * 100
            avd = abs(maxd - mind) / 2.
            # print poiLabel+':   ('+str(poi.getAsymErrorLo())+','+str(poi.getAsymErrorHi())+'), ' + str(avd)
            print '{label:10}:\t{mind:5.3f}%\t{maxd:5.3f}%\t{avd:5.3f}%'.format(
                label=poiLabel, mind=mind, maxd=maxd, avd=avd)

        can_gamma = self.canvas('gamma')
        self.h_width.GetXaxis().SetTitle("#Gamma_{T}")
        self.h_width.GetYaxis().SetTitle("N toys")
        self.h_width.Draw()
        print 'Relative error on the total width ', self.h_width.GetRMS(
        ) / self.h_width.GetMean()
        print 'Please check the histogram to see that the dist is Gaussian. If not the fit is biased'
        print 'The fit can be biased when floating the width sometimes.'

        can_pulls = self.canvas('pulls', 1000, 1000)
        npulls = len(self.constraint)
        nxy = int(math.ceil(math.sqrt(npulls)))
        can_pulls.Divide(nxy, nxy)
        for i, c in enumerate(self.constraint.values()):
            can_pulls.cd(i + 1)
            c.pulls.Draw()
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, 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 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
Esempio n. 9
0
class CouplingsFitterTest(object):
    def __init__(self):
        self.poiLabels = []
        self.poi = dict()
        self.poilabels = dict()
        self.constraint = dict()
        self.canvases = dict()
        self._keep = []

    def addPOI(self, poi, label='', minimum=-0.3, maximum=0.3):
        '''Add a parameter of interest.

        Example:
        
        addPOI('Z','Z',-0.05,0.05)
        ->
        Z[0,-0.05,0.05]
        # adds variable Z with value 0,
        # allow it to scale between -0.05 and 0.05
        '''
        self.poi[poi] = ROOT.RooRealVar(poi, poi, 0, minimum, maximum)
        if label == '':
            label = poi
        self.poilabels[poi] = label

    def addConstraint(self, name, expr, deplist, mean, sigma):
        '''Add a constraint on one of the observables
        
        For example, for ZH inclusive: 
        
        f.addConstraint('Zh','(1+Z)*(1+Z)','Z',1,0.004)  

        Z is an additive modification of the gZ coupling w/r to the standard model,
        so 1+Z = \kappa_Z
        
        Zh is the pdf of the ratio of the yield w/r to the one expected in the standard model.
        This pdf depends on Z, as (1+Z)*(1+Z).
        
        ZhObs is the measured value, here 1 so we assume that the SM yield is observed.
        
        The fit varies the parameter of interest Z, thus modifying the pdf,
        while ZhObs is fixed at 1. The likelihood of each value of Z is evaluated at ZhObs on the pdf.
        '''
        deps = self._getdeps(deplist)
        self.constraint[name] = GaussianConstraint(name, expr, deps, mean,
                                                   sigma)

    def _getdeps(self, deplist):
        depnames = deplist
        try:
            depnames = deplist.split(',')
        except:
            pass
        deps = []
        for dep in depnames:
            if dep == 'width':
                deps.append(self.width)
            else:
                deps.append(self.poi[dep])
        return deps

    def addUniformConstraint(self, name, expr):
        '''Adds a uniform constraint with pdf name, and expression expr.
        
        For example:
        
        f.addPOI('inv','inv', 0, 0.01)
        f.addUniformConstraint('Zhinv','inv') ####->Means free floating

        inv (the invisible BR) is free to float between 0 and 1%
        '''
        deps = self._getdeps([expr])
        self.constraint[name] = UniformConstraint(name, expr, deps)

    def info(self):
        for name, constraint in sorted(self.constraint.iteritems()):
            constraint.info()

    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))

    def canvas(self, name, *args):
        canvas = self.canvases.setdefault(name,
                                          ROOT.TCanvas(name, name, *args))
        canvas.cd()
        return canvas

    def keep(self, obj):
        self._keep.append(obj)
        return obj

    def createSummary(self):
        #sample the covariance matrix for the width
        # ROOT.gStyle.SetOptTitle(0)
        ##        ROOT.gStyle.SetStatW(0.4);
        ##        ROOT.gStyle.SetStatH(0.4);
        ##
        ##        self.graph_couplings = ROOT.TGraphAsymmErrors(len(self.poi)+2)
        ##
        ##        order_BR = ['Z', 'W', 'b', 'c', 'g', 'tau', 'mu', 'gamma', 'inv']
        ##
        ##        for br in order_BR:
        ##            if not self.poi.get(br, None):
        ##                order_BR.remove(br)
        ##
        ##        for i, poiname in enumerate(order_BR):
        ##            poi = self.poi.get(poiname)
        ##            self.graph_couplings.SetPoint(i, i+0.5, poi.getVal())
        ##            self.graph_couplings.SetPointError(i, 0.0, 0.0, -poi.getAsymErrorLo(), poi.getAsymErrorHi())

        print 'Sampling the covariance matrix to propagate error on width'

        ##        self.h_width = ROOT.TH1F('h_width','width',1000,0.5,1.5)
        ntoys = 10000
        for i in range(ntoys):
            randomizedPars = self.fit_result.randomizePars()
            for j in range(0, randomizedPars.getSize()):
                self.poi[randomizedPars.at(j).GetName()].setVal(
                    randomizedPars.at(j).getVal())
##            self.h_width.Fill(self.width.getVal())
##            for cstr in self.constraint.values():
##                cstr.fill_pull()
##        self.graph_couplings.SetMarkerStyle(20)
##        self.graph_couplings.SetLineWidth(3)
##        can_couplings = self.canvas('couplings')
##        self.graph_couplings.Draw("AP")
##        self.graph_couplings.GetYaxis().SetTitle("68% CL on d(A) ")
##        self.graph_couplings.GetXaxis().SetNdivisions(0)
##        l= self.keep(ROOT.TLine())
##        l.SetLineColor(ROOT.kRed)
##        l.SetLineWidth(3)
##        l.DrawLine(0.0,0.0,len(self.poi)+1.5,0)

##        self.graph_couplings.SetPoint(len(self.poi),len(self.poi)+0.5,0.0)
##        self.graph_couplings.SetPointError(len(self.poi),0.0,0.0,self.h_width.GetRMS()/self.h_width.GetMean(),
##                            self.h_width.GetRMS()/self.h_width.GetMean())
##
##        for i, poiname in enumerate(order_BR+['#Gamma_{T}']):
##            label = self.poilabels.get(poiname, poiname)
##            tex = self.keep(ROOT.TLatex(i+0.5,0.95*self.graph_couplings.GetYaxis().GetXmin(),label))
##            tex.Draw()
##
        print """
###############################################################
###############################################################
###############################################################
                         RESULTS
###############################################################
###############################################################
############################################################### 
              """

        print 'RESULTS FOR THE CONFIDENCE INTERVALS------>'
        for name in self.poi:
            poi = self.poi[name]
            poiLabel = self.poilabels.get(name, name)
            mind = poi.getAsymErrorLo() * 100
            maxd = poi.getAsymErrorHi() * 100
            avd = abs(maxd - mind) / 2.
            # print poiLabel+':   ('+str(poi.getAsymErrorLo())+','+str(poi.getAsymErrorHi())+'), ' + str(avd)
            print '{label:10}:\t{mind:5.3f}%\t{maxd:5.3f}%\t{avd:5.3f}%'.format(
                label=poiLabel, mind=mind, maxd=maxd, avd=avd)
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
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
Esempio n. 12
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)
Esempio n. 13
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])