Python ROOT.RooArgList类代码示例

def get_arglist(args):
    """Return and arglist of the RooFit objects."""
    arglist = RooArgList()
    for arg in args:
        if arg.InheritsFrom(RooAbsArg.Class()): arglist.add(arg)
        else: TypeError('%s should inherit from RooAbsArg' % arg.GetName())
    return arglist

def main ():
    N_bkg1 = 9000
    N_signal = 1000
    N_bkg1_obs = 10000
    N_signal_obs = 2000
    N_data = N_bkg1_obs + N_signal_obs
    mu1, mu2, sigma1, sigma2 = 100, 140, 15, 5
    x1 = mu1 + sigma1 * np.random.randn( N_bkg1 )
    x2 = mu2 + sigma2 * np.random.randn( N_signal )
    x1_obs = mu1 + sigma1 * np.random.randn( N_bkg1_obs )
    x2_obs = mu2 + sigma2 * np.random.randn( N_signal_obs )
    
    h1 = Hist( 100, 40, 200, title = 'Background' )
    h2 = h1.Clone( title = 'Signal' )
    h3 = h1.Clone( title = 'Data' )
    h3.markersize = 1.2
    
    # fill the histograms with our distributions
    map( h1.Fill, x1 )
    map( h2.Fill, x2 )
    map( h3.Fill, x1_obs )
    map( h3.Fill, x2_obs )
    
    histograms_1 = {'signal': h2,
                  'bkg1': h1,
                  'data': h3}
    
    histograms_2 = {'signal': h2,
                  'bkg1': h1,
                  'data': h3}
    
    # roofit_histograms contains RooDataHist
    # model = RooAddPdf
    model1, roofit_histograms_1,fit_variable_1 = get_roofit_model( histograms_1, fit_boundaries = ( 40, 200 ), name = 'm1' )
    model2, roofit_histograms_2, fit_variable_2 = get_roofit_model( histograms_2, fit_boundaries = ( 40, 200 ), name = 'm2' )
    sample = RooCategory( 'sample', 'sample' )
    sample.defineType( 'm1', 1 )
    sample.defineType( 'm2', 2 )
    combined_data = deepcopy( roofit_histograms_1['data'] )
    combined_data.add( roofit_histograms_2['data'] )
    # RooDataHist(const char* name, const char* title, const RooArgList& vars, RooCategory& indexCat, map<std::string,TH1*> histMap, Double_t initWgt = 1.0)
    sim_pdf = RooSimultaneous( "simPdf", "simultaneous pdf", sample )
    sim_pdf.addPdf( model1, 'm1' )
    sim_pdf.addPdf( model2, 'm2' )
    variables = RooArgList()
    variables.add(fit_variable_1)
    variables.add(fit_variable_2)
#     combined_data = RooDataHist('combined_data', 'combined_data',
#                                 variables, RooFit.Index(sample),
#                                 RooFit.Import('m1', roofit_histograms_1['data']),
#                                 RooFit.Import('m2', roofit_histograms_2['data']))
    fitResult = sim_pdf.fitTo( combined_data,
#                    RooFit.Minimizer( "Minuit2", "Migrad" ),
#                    RooFit.NumCPU( 1 ),
#                    RooFit.Extended(),
                    RooFit.Save(), 
                   )

def get_roofit_model( histograms, fit_boundaries, name = 'model' ):
    data_label = 'data'
    samples = sorted( histograms.keys() )
    samples.remove( data_label )
    roofit_histograms = {}
    roofit_pdfs = {}
    roofit_variables = {}
    variables = RooArgList()
    variable_set = RooArgSet()

    fit_variable = RooRealVar( name , name, fit_boundaries[0], fit_boundaries[1] )
    variables.add( fit_variable )
    variable_set.add( fit_variable )
    
    roofit_histograms[data_label] = RooDataHist( data_label,
                                                     data_label,
                                                     variables,
                                                     histograms[data_label] )
    
    pdf_arglist = RooArgList()
    variable_arglist = RooArgList()
    N_total = histograms[data_label].Integral() * 2
    N_min = 0
    for sample in samples:
        roofit_histogram = RooDataHist( sample, sample, variables, histograms[sample] )
        roofit_histograms[sample] = roofit_histogram
        roofit_pdf = RooHistPdf ( 'pdf' + sample, 'pdf' + sample, variable_set, roofit_histogram, 0 )
        roofit_pdfs[sample] = roofit_pdf
        roofit_variable = RooRealVar( sample, "number of " + sample + " events", histograms[sample].Integral(), N_min, N_total, "event" )
        roofit_variables[sample] = roofit_variable
        pdf_arglist.add( roofit_pdf )
        variable_arglist.add( roofit_variable )
        
    model = RooAddPdf( name, name, pdf_arglist, variable_arglist )
    return model, roofit_histograms, fit_variable

    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.0
        N_max = self.normalisation[self.data_label] * 2.0
        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

def getPdfInRegions(w,sample,region):
    if isinstance(sample,list):
        sampleArgList = RooArgList()
        sample_str="group"
        for s in sample:
            componentTmp = Util.GetComponent(w,s,region,True)
            sample_str=sample_str+"_"+s
            sampleArgList.add(componentTmp)
            pass
        pdfInRegion = RooAddition(sample_str,sample_str,sampleArgList)
    else:
        pdfInRegion  = Util.GetComponent(w,sample,region)
        pass
    return pdfInRegion

def getPdfInRegionsWithRangeName(w,sample,region,rangeName):
    """
    Should be moved to $HF/src/Utils.h -- FIXME
    """
    if isinstance(sample,list):
        sampleArgList = RooArgList()
        sample_str="group"
        for s in sample:
            componentTmp = Util.GetComponent(w,s,region,True,rangeName)
            sample_str=sample_str+"_"+s
            sampleArgList.add(componentTmp)
            pass
        pdfInRegion = RooAddition(sample_str,sample_str,sampleArgList)
    else:
        pdfInRegion  = Util.GetComponent(w,sample,region,False,rangeName)
        pass
    return pdfInRegion

    def addHistData(self, ws, ds_name, item_title, var_set_name, var_set_type, ds_object, sets = None):
        legend = '[exostConfig::addHistData]:'
        
        arg_list = RooArgList()
        if (var_set_type == 'set'):
            #arg_list.add(ws.set(var_set_name))
            arg_list.add(sets[var_set_name])
        elif (var_set_type == 'var'):
            arg_list.add(ws.var(var_set_name))
        else:
            print legend, 'error: unknown var_set_type, cannot create dataset', ds_name
            return -1

        #create the hist dataset
        ds = RooDataHist(ds_name, item_title, arg_list, ds_object, 1)
        
        # import the datahist. Note workaround 'import' being a reserved word
        getattr(ws, 'import')(ds)

    def Hist2Pdf(self, hist, pdfName, ws, order = 0):
        if ws.pdf(pdfName):
            return ws.pdf(pdfName)

        try:
            obs = [ self.pars.varNames[x] for x in self.pars.var ]
        except AttributeError:
            obs = self.pars.var

        varList = RooArgList()
        for v in obs:
            varList.add(ws.var(v))
        newHist = RooDataHist(pdfName + '_hist', pdfName + '_hist',
                              varList, hist)
        
        thePdf = RooHistPdf(pdfName, pdfName, RooArgSet(varList), 
                            newHist, order)
        getattr(ws, 'import')(thePdf)

        return ws.pdf(pdfName)

def getPdfInRegions(w,sample,region):
    """
    Return the PDF in a region for a sample
    Should be moved to $HF/src/Utils.h -- FIXME

    @param sample The sample to find
    @param region The region to use
    """
    if isinstance(sample,list):
        sampleArgList = RooArgList()
        sample_str="group"
        for s in sample:
            componentTmp = Util.GetComponent(w,s,region,True)
            sample_str=sample_str+"_"+s
            sampleArgList.add(componentTmp)
            pass
        pdfInRegion = RooAddition(sample_str,sample_str,sampleArgList)
    else:
        pdfInRegion  = Util.GetComponent(w,sample,region)
        pass
    return pdfInRegion

def performFitInLeptonAbsEta(data_histogram, signal_histogram, bkg1_histogram, bkg2_histogram):
    N_Data = data_histogram.Integral()
    N_signal = signal_histogram.Integral()
    N_bkg1 = bkg1_histogram.Integral()
    N_bkg2 = bkg2_histogram.Integral()
    leptonAbsEta = RooRealVar("leptonAbsEta", "leptonAbsEta", 0., 3.)
    variables = RooArgList()
    variables.add(leptonAbsEta)
    variable_set = RooArgSet()
    variable_set.add(leptonAbsEta)
    
    lowerBound = 0
    upperBound = N_Data*2
    
    data_RooDataHist = RooDataHist("data", "dataset with leptonAbsEta", variables, data_histogram)
    signal_RooDataHist =  RooDataHist("rh_signal", "signal", variables, signal_histogram);
    bkg1_RooDataHist =  RooDataHist("rh_bkg1", "bkg1", variables, bkg1_histogram);
    bkg2_RooDataHist =  RooDataHist("rh_bkg2", "bkg2", variables, bkg2_histogram);

    signal_RooHistPdf = RooHistPdf("pdf_signal", "Signal PDF", variable_set, signal_RooDataHist, 0)
    signal_RooHistPdf = RooHistPdf("pdf_signal", "Signal PDF", variable_set, signal_RooDataHist, 0)

def plotFitModel(model, frame, wksp, myconfigfile, log) :
    if debug :
        model.Print('t')
        frame.Print('v')

    dataset                             = w.data(dataSetToPlot)

    # plot model itself
    fr = model.plotOn(frame,
                      RooFit.LineColor(kBlue+3),RooFit.Name("FullPdf"))

    var = []
    tacc_list = RooArgList()
    numKnots = myconfigfile["Acceptance"]["knots"].__len__()
    for i in range(0,numKnots+1):
        varName = "var%d"%(int(i+1))
        var.append(wksp.var(varName))
        print "[INFO] Load %s with value %0.3lf"%(var[i].GetName(),var[i].getValV())
        tacc_list.add(var[i])

    varAdd = RooAddition(wksp.obj("var%d"%(numKnots+2)))
    #varAdd = RooRealVar(wksp.obj("var%d"%(numKnots+2)))
    print "[INFO] Load %s with value %0.3lf"%(varAdd.GetName(),varAdd.getValV())
    tacc_list.add(varAdd)

    #len = var.__len__()
    #tacc_list.add(var[len])
    spl = RooCubicSplineFun("splinePdf", "splinePdf", time, "splineBinning", tacc_list)
    if log:
        rel = 200
    else:
        rel = 1000
    rel = 40                                                                                                                                                                           
     
    fr = spl.plotOn(frame, RooFit.LineColor(kRed),  RooFit.Normalization(rel, RooAbsReal.Relative),RooFit.Name("sPline"))
    fr = model.plotOn(frame,
                      RooFit.LineColor(kBlue+3), RooFit.Name("FullPdf"))

def PlotSignalShapes(Selection):
    f__ = TFile.Open( "datacards/22June/2dPlots.root")
    signal_fname_1 = ("signals/22June/out_{sample:s}_syst.root", "cms_hgg_13TeV" )
    signal_fname_2 = ("signals/22June/out_ctcv_{sample:s}_syst.root" , "ctcv" )
    samples = {"thw":signal_fname_2, "thq":signal_fname_2, "tth":signal_fname_1 , "vh":signal_fname_1 }
    purity_h_name = "{sample:s}/"+Selection+"/h{sample:s}_"+Selection+"_purity_CtCv"
    purities = RooArgList()
    signalshapes = RooArgList()

    ctOverCvs = []

    mVar = None
    ctovercv_vals = None
    for sample in samples :
        purity = CtCvCpInfo("purity_" + sample)
        ctovercv_vals = sorted(purity.AllCtOverCVs.keys())
        purity.FillFrom2DHisto( f__.Get( purity_h_name.format( sample=sample ) ) )
        purity.GetCtOverCv()
        purities.add( purity.CtOverCvDataHistFunc )
        objsToKeep.append( purity )

        sFile = TFile.Open( samples[sample][0].format( sample=sample ) )
        ws = sFile.Get( samples[sample][1] )
        pdf = ws.pdf("RV{sample:s}_mh125".format( sample=sample) )
        objsToKeep.append(sFile)
        objsToKeep.append(ws)
        objsToKeep.append(pdf)
        signalshapes.add( pdf )

        ctOverCvs.append( ws.var( "CtOverCv" ) )
        mVar = ws.var("CMS_hgg_mass")
        
    ret = RooAddPdf("signal" , "signal" , signalshapes , purities )
    objsToKeep.append( ret )
    plot = mVar.frame()
    options = ""
    for ctovercv in ctovercv_vals :
        for var in ctOverCvs:
            var.setVal( ctovercv )
        name = "name%g" % ctovercv
        ret.plotOn( plot , RooFit.DrawOption(options) , RooFit.Name(name) )
        
        c = TCanvas()
        plot.Draw()
        c.SaveAs("a.gif+")

        if not "same" in options :
            options += " same"

    return c

temp_tt.Rebin(20)
temp_wj.Rebin(20)
temp_zj.Rebin(20)
temp_stop.Rebin(20)
temp_qcd.Rebin(20)
temp_qcd.Scale(800/temp_qcd.Integral())
temp_signal = temp_tt.Clone('signal')
temp_signal.Add(temp_stop)
#temp_tt.Scale(1/temp_tt.Integral())
#temp_wj.Scale(1/temp_wj.Integral())
#temp_zj.Scale(1/temp_zj.Integral())
#temp_stop.Scale(1/temp_stop.Integral())
#temp_qcd.Scale(1/temp_qcd.Integral())
leptonAbsEta = RooRealVar("leptonAbsEta", "leptonAbsEta", 0., 3.)
vars = RooArgList()
vars.add(leptonAbsEta)
vars_set = RooArgSet()
vars_set.add(leptonAbsEta)
n_event_obs = h_m3_data.GetEntries();
lowerBound = -10 * sqrt(n_event_obs); 
upperBound = n_event_obs + 10 * sqrt(n_event_obs);
n_init = n_event_obs / 2.;
data = RooDataHist("data", "dataset with leptonAbsEta", vars, h_m3_data)
rh_tt = RooDataHist("rh_tt", "tt", vars, temp_tt);
rh_wj = RooDataHist("rh_wj", "wj", vars, temp_wj);
rh_zj = RooDataHist("rh_zj", "zj", vars, temp_zj);
rh_qcd = RooDataHist("rh_qcd", "qcd", vars, temp_qcd);
rh_stop = RooDataHist("rh_stop", "singletop", vars, temp_stop);
rh_signal = RooDataHist("rh_signal", "signal", vars, temp_signal);

REBIN = 20
n_signal = n_ttbar + n_stop
temp_VPlusJets.Rebin(REBIN)
temp_tt.Rebin(REBIN)
temp_wj.Rebin(REBIN)
temp_zj.Rebin(REBIN)
temp_stop.Rebin(REBIN)
temp_qcd.Rebin(REBIN)
temp_qcd2.Rebin(REBIN)
temp_qcd.Scale(n_qcd / temp_qcd.Integral())
temp_qcd2.Scale(n_qcd/temp_qcd2.Integral())
temp_signal = temp_tt.Clone('signal')
temp_signal.Add(temp_stop)

leptonAbsEta = RooRealVar("leptonAbsEta", "leptonAbsEta", 0., 3.)
variables = RooArgList()
variables.add(leptonAbsEta)
vars_set = RooArgSet()
vars_set.add(leptonAbsEta)
n_event_obs = h_data.Integral();
lowerBound = -10 * sqrt(n_event_obs);
lowerBound = 0
upperBound = n_event_obs + 10 * sqrt(n_event_obs);
upperBound = n_event_obs
n_init = n_event_obs / 2.;
data = RooDataHist("data", "dataset with leptonAbsEta", variables, h_data)
rh_tt = RooDataHist("rh_tt", "tt", variables, temp_tt);
rh_wj = RooDataHist("rh_wj", "wj", variables, temp_wj);
rh_zj = RooDataHist("rh_zj", "zj", variables, temp_zj);
rh_VJ = RooDataHist("rh_VJ", "VJ", variables, temp_VPlusJets);
rh_qcd = RooDataHist("rh_qcd", "qcd", variables, temp_qcd);

def findOnePe(hist, ws, name='x', Npe = 1):
    fitPed(hist, ws, name)
    x = ws.var(name)

    ped = ws.pdf('ped')
    pedWidth = ws.var('pedWidth')

    pdfs = RooArgList(ped)
    pdfList = []

    fped = RooRealVar('fped', 'f_{ped}', 0.8, 0., 1.)
    fractions = RooArgList(fped)
    fList = []
    peList = []

    peMean = RooRealVar('peMean', 'mean_{pe}', 6., 0., 20.)
    peWidth = RooRealVar('peWidth', 'width_{pe}', pedWidth.getVal(), 0., 10.)

    for i in range(0, Npe):
        pem = RooFormulaVar('pem{0}'.format(i+1), '@0+{0}*@1'.format(i+1),
                            RooArgList(ws.var('pedMean'), peMean))
        peList.append(pem)
        npepdf = RooGaussian('pe{0}pdf'.format(i+1), 'pe{0}pdf'.format(i+1),
                             x, pem, pedWidth)
        pdfs.add(npepdf)
        pdfList.append(npepdf)

        fnpe = RooRealVar('fpe{0}'.format(i+1), 'fpe{0}'.format(i+1),
                          0.5, -0.1, 1.0)
        fractions.add(fnpe)
        fList.append(fnpe)

    #bgMean = RooRealVar("bgMean", "bgMean", 6.0, x.getMin(), x.getMax())
    bgScale = RooRealVar("bgScale", "bgScale", 0.5, -1.0, Npe + 1.0)
    bgMean = RooFormulaVar("bgMean", "@[email protected]*@2",
                           RooArgList(peMean, ws.var('pedMean'), bgScale))
    bgWidthL = RooRealVar("bgWidthL", "bgWidthL", pedWidth.getVal()*2,
                          0., 25.)
    bgWidthR = RooRealVar("bgWidthR", "bgWidthR", pedWidth.getVal()*7,
                          0., 25.)

    bgGauss = RooBifurGauss("bgGauss", "bgGauss", x, bgMean,
                            bgWidthR, bgWidthR)

    if (Npe > 1):
        pdfs.add(bgGauss)
    else:
        fractions.remove(fractions.at(fractions.getSize()-1))

##     pem = RooFormulaVar('pem', '@[email protected]', RooArgList(peMean, ws.var('pedMean')))
##     firstPe = RooGaussian('firstPe', 'firstPe', x, pem, peWidth)

##     pdfs.Print("v")
##     fractions.Print("v")
    pedPlusOne = RooAddPdf('pedPlusOne', 'pedPlusOne', pdfs, fractions, True)

    ## pedWidth = ped.GetParameter(2)
    ## pedMean = ped.GetParameter(1)
    ## pedA = ped.GetParameter(0)
    
    secondMax = hist.GetMaximumBin() + 1
    goingDown = True
    maxVal = hist.GetBinContent(secondMax)
    foundMax = False
    while (not foundMax) and (secondMax < hist.GetNbinsX()):
        tmpVal = hist.GetBinContent(secondMax+1)
        if (tmpVal < maxVal):
            if not goingDown:
                foundMax = True
            else:
                goingDown = True
                maxVal = tmpVal
                secondMax += 1
        elif (tmpVal > maxVal):
            goingDown = False
            maxVal = tmpVal
            secondMax += 1
        else:
            maxVal = tmpVal
            secondMax += 1

    secondMaxx = hist.GetBinCenter(secondMax)
    print 'found 2nd maximum in bin',secondMax,'value',secondMaxx

##     peMean.setVal(secondMaxx)
##     bgMean.setVal(secondMaxx*0.6)
    x.setRange('pedPlus_fit', x.getMin(), ws.var('pedMean').getVal()+pedWidth.getVal()*6.*(Npe+0))

    pedPlusOne.fitTo(ws.data('ds'), RooFit.Minos(False),
                     RooFit.Range('pedPlus_fit'),
                     RooFit.PrintLevel(1))

    getattr(ws, 'import')(pedPlusOne)