Python RooDataSet.reduce方法代码示例

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]
def get_dataset(varargset, ftree, cut="", wt="", scale=1):
    """Return a dataset.

    Return a dataset from the ntuple `ftree'. Apply a selection cut
    using the `cutVar' variable and the selection `cut'.

    """

    from rplot.fixes import ROOT
    from rplot.tselect import Tsplice

    splice = Tsplice(ftree)
    splice.make_splice("sel", cut)
    from ROOT import RooDataSet, RooFit, RooFormulaVar, RooArgList

    tmpdst = RooDataSet("tmpdataset", "", varargset, RooFit.Import(ftree))
    if wt:
        wtvar = RooFormulaVar("wt", "{}*@0".format(scale), RooArgList(varargset[wt]))
        wtvar = tmpdst.addColumn(wtvar)
        varargset.remove(varargset[wt])
        varargset.add(wtvar)
        dst = RooDataSet("dataset", "Dataset", varargset, RooFit.Import(tmpdst), RooFit.WeightVar(wtvar))
        varargset.remove(wtvar)
        dst = dst.reduce(varargset)
    return dst

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]
def main() :
    inputfile_name = "/data1/chibdata/collision/v2/2012_AllData_v2.root"

    chibTree_name = "rootuple/chibTree"
    print "Opening file"
    inputfile = TFile.Open(inputfile_name,"READ")
    print "Importing tree"
    tree = TTree()
    inputfile.GetObject(chibTree_name, tree)

    invm1S = RooRealVar("invm1S", "invm1S", 9.5, 11.5)
    invm2S = RooRealVar("invm2S", "invm2S", 9.5, 20.0)
    invm3S = RooRealVar("invm3S", "invm3S", 9.5, 20.0)
    dimuon_mass = RooRealVar("dimuon_mass","dimuon_mass", 8.0, 12.0)
    dimuon_rapidity = RooRealVar("dimuon_rapidity", "dimuon_rapidity", -5.0, 5.0)
    dimuon_pt = RooRealVar("dimuon_pt","dimuon_pt", 0.0, 100.0)
    photon_eta = RooRealVar("photon_eta","photon_eta", -5.0, 5.0)
    photon_pt = RooRealVar("photon_pt","photon_pt", 0.0, 100.0)
    ctpv = RooRealVar("ctpv","ctpv", -5.0, 5.0)
    ctpv_error = RooRealVar("ctpv_err","ctpv_err", -5.0, 5.0)
    pi0_abs_mass = RooRealVar("pi0_abs_mass","pi0_abs_mass", 0.0, 2.0)
    Y1S_nsigma = RooRealVar("Y1S_nsigma","Y1S_nsigma",0.0,30.0)
    Y2S_nsigma = RooRealVar("Y2S_nsigma","Y2S_nsigma",0.0,30.0)
    Y3S_nsigma = RooRealVar("Y3S_nsigma","Y3S_nsigma",0.0,35.0)
    conv_vertex = RooRealVar("conv_vertex", "conv_vertex", 0.0, 70.0)
    dz = RooRealVar("dz","dz", -50.0, 50.0)


    print "Assigning argset"
    dataArgSet = RooArgSet(invm1S, invm2S, invm3S, dimuon_mass, dimuon_rapidity, dimuon_pt, photon_eta, photon_pt, ctpv)
    dataArgSet.add( ctpv_error )
    dataArgSet.add( pi0_abs_mass )
    dataArgSet.add( Y1S_nsigma )
    dataArgSet.add( Y2S_nsigma )
    dataArgSet.add( Y3S_nsigma )
    dataArgSet.add( conv_vertex )
    dataArgSet.add( dz )

    print "Creating DataSet"
    dataSet = RooDataSet("chibds","Chib RooDataSet", tree, dataArgSet)


    # Selection
    cuts = "photon_pt > 0.0" +                   \
           "&&dimuon_pt > 9.5"                   \
           "&& abs(photon_eta) < 1.4" +          \
           "&& abs(dimuon_rapidity) < 1.25 " +   \
           "&& Y3S_nsigma < 2.5 " +              \
           "&& abs(dz) < 0.1"

    rds_cutted = dataSet.reduce( RooFit.Cut(cuts) )

    dofit(rds_cutted, "Chib_fit_2012_3S")

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]
def get_dataset(varargset, ftree, cut='', cutVars=[], *cmdArgs):
    """Return a dataset.

    Return a dataset from the ntuple `ftree'. Apply a selection cut
    using the `cutVar' variable and the selection `cut'.

    """

    varargsetclone = varargset.clone('varargsetclone')
    for cvar in cutVars:
        varargsetclone.add(cvar) # Add selVar to apply cut

    tmpdataset = RooDataSet('dataset', 'Dataset', varargsetclone,
                            RooFit.Import(ftree), RooFit.Cut(cut), *cmdArgs)
    dataset = tmpdataset.reduce(varargset)
    del tmpdataset
    return dataset

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]
		ApProduced = TH2F("ApProduced","ApProduced",nApMass,Medges,NepsBins,Epsedges)
		VProducedRho = TH2F("VProducedRho","VProducedRho",nApMass,Medges,NepsBins,Epsedges)
		detectableRho = TH2F("detectableRho","detectableRho",nApMass,Medges,NepsBins,Epsedges)
		VProducedPhi = TH2F("VProducedPhi","VProducedPhi",nApMass,Medges,NepsBins,Epsedges)
		detectablePhi = TH2F("detectablePhi","detectablePhi",nApMass,Medges,NepsBins,Epsedges)
		for k in range(0,nApMass):
			massAp = (massApmax - massApmin)/float(nApMass - 1) * k + massApmin
			massV = 0.6*massAp
			massPi = massAp/3.
			fPi = massPi/fpiArr[i]
			massVres = massresfile.Get("mres").GetFunction("pol1").Eval(massV)
			breakz = tailsfile.Get("breakz").GetFunction("pol3").Eval(massV)
			length = tailsfile.Get("length").GetFunction("pol3").Eval(massV)

			frame = uncVZ.frame()
			dataInRange = dataset.reduce(obs,"abs({0}-{1})<{2}/2*{3}".format("uncM",massV,masscut_nsigma,massVres))
			binnedData = dataInRange.binnedClone()
			binnedData.plotOn(frame)
			mean = binnedData.mean(uncVZ)
			sigma = binnedData.sigma(uncVZ)
			uncVZ.setRange("fitRange",mean-2*sigma,mean+2*sigma)
			gauss_params.setRealValue("mean",mean)
			gauss_params.setRealValue("sigma",sigma)
			gauss_pdf.fitTo(binnedData,RooFit.Range("fitRange"),RooFit.PrintLevel(-1))
			mean = gauss_params.getRealValue("mean")
			sigma = gauss_params.getRealValue("sigma")
			gaussexp_params.setRealValue("gauss_mean",mean)
			gaussexp_params.setRealValue("gauss_sigma",sigma)
			gaussexp_params.setRealValue("exp_breakpoint",breakz)
			gaussexp_params.setRealValue("exp_length",length)
			w.var("gauss_mean").setConstant(True)

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]
    #c.Print(remainder[0]+".pdf","Title:mass_{0}".format(mass))

    num_pairs = gDirectory.Get("mass").GetEntries()*scale_factor
    radfrac = radfracFile.Get("radfrac").GetFunction("pol3").Eval(mass)
    num_rad = radfrac*num_pairs
    ap_yield= 3*math.pi/(2*(1/137.0))*num_rad*(mass/deltaM)
    print "{0} pairs, {1} radfrac, {2} rad, {3} A'".format(num_pairs,radfrac,num_rad,ap_yield)


    breakz = tailsFile.Get("breakz").GetFunction("pol3").Eval(mass)
    length = tailsFile.Get("length").GetFunction("pol3").Eval(mass)

    c.cd(2)
    gPad.SetLogy()
    frame=uncVZ.frame()
    dataInRange = dataset.reduce(obs,"abs({0}-{1})<{2}/2*({3}+{4}*uncVZ)".format(massVar,mass,masscut_nsigma,mres_p0,mres_p1))
    binnedData = dataInRange.binnedClone()
    binnedData.plotOn(frame)
    mean = binnedData.mean(uncVZ)
    sigma = binnedData.sigma(uncVZ)
    uncVZ.setRange("fitRange",mean-2*sigma,mean+2*sigma)
    gauss_params.setRealValue("mean",mean)
    gauss_params.setRealValue("sigma",sigma)
    gauss_pdf.fitTo(binnedData,RooFit.Range("fitRange"),RooFit.PrintLevel(-1))
    mean= gauss_params.getRealValue("mean")
    sigma= gauss_params.getRealValue("sigma")
    gaussexp_params.setRealValue("gauss_mean",mean)
    gaussexp_params.setRealValue("gauss_sigma",sigma)
    gaussexp_params.setRealValue("exp_breakpoint",breakz)
    gaussexp_params.setRealValue("exp_length",length)
    w.var("gauss_mean").setConstant(True)

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]

#.........这里部分代码省略.........
    #jPsiMass.setRange('NormalizationRangeFormlfit_nonPromptSignal',2.7,3.5)
    #jPsiMass.setRange('NormalizationRangeFormlfit_leftMassSideBand',2.7,3.1)
    #jPsiMass.setRange('NormalizationRangeFormlfit_rightMassSideBand',3.1,3.5)

    #want the prompt fit only done in prompt region
    #non-prompt only in non-prompt region
    #background over entire cTau range
    #jPsicTau.setRange('NormalizationRangeFormlfit_promptSignal',-1,.1)
    #jPsicTau.setRange('NormalizationRangeFormlfit_nonPromptSignal',.1,2.5)
    #jPsicTau.setRange('NormalizationRangeFormlfit_leftMassSideBand',-1,2.5)
    #jPsicTau.setRange('NormalizationRangeFormlfit_rightMassSideBand',-1,2.5)

    #redVars.add(ctauStates)
    #redVars.add(massStates)
    #redVars.add(states)
    fitVars.add(ctauStates)
    fitVars.add(massStates)
    fitVars.add(states)
    
    fullData = RooDataSet('fullData','The Full Data From the Input ROOT Trees',
                          dataVars,
                          ROOT.RooFit.Import(inputData))    

    for rap_bin in range(1,len(jpsi.pTRange)):
        yMin  = jpsi.rapForPTRange[rap_bin-1][0]
        yMax  = jpsi.rapForPTRange[rap_bin-1][-1]
        for pt_bin in range(len(jpsi.pTRange[rap_bin])):

            ptMin = jpsi.pTRange[rap_bin][pt_bin][0]
            ptMax = jpsi.pTRange[rap_bin][pt_bin][-1]               

            sigMaxMass = jpsi.polMassJpsi[rap_bin] + jpsi.nSigMass*jpsi.sigmaMassJpsi[rap_bin]
            sigMinMass = jpsi.polMassJpsi[rap_bin] - jpsi.nSigMass*jpsi.sigmaMassJpsi[rap_bin]

            sbHighMass = jpsi.polMassJpsi[rap_bin] + jpsi.nSigBkgHigh*jpsi.sigmaMassJpsi[rap_bin]
            sbLowMass  = jpsi.polMassJpsi[rap_bin] - jpsi.nSigBkgLow*jpsi.sigmaMassJpsi[rap_bin]

            ctauNonPrompt = .1
            
            massFun = RooFormulaVar('massRegion','Function that returns the mass state.',
                                     '('+jPsiMass.GetName()+' < '+str(sigMaxMass)+' && '+jPsiMass.GetName()+' > '+str(sigMinMass)+
                                     ') - ('+jPsiMass.GetName()+' > '+str(sbHighMass)+')'+
                                     '-2*('+jPsiMass.GetName()+' < '+str(sbLowMass)+')',
                                     RooArgList(jPsiMass,jPsicTau))
            
            ctauFun = RooFormulaVar('ctauRegion','Function that returns the ctau state.',
                                     '('+jPsicTau.GetName()+' > '+str(ctauNonPrompt)+')',
                                     RooArgList(jPsiMass,jPsicTau))

            mlFun = RooFormulaVar('mlRegion','Function that returns the mass and lifetime state.',
                                  '('+jPsiMass.GetName()+' < '+str(sigMaxMass)+' && '+jPsiMass.GetName()+' > '+str(sigMinMass)+
                                  ') + ('+jPsiMass.GetName()+' < '+str(sigMaxMass)+' && '+jPsiMass.GetName()+' > '+
                                  str(sigMinMass)+' && '+jPsicTau.GetName()+' > '+str(ctauNonPrompt)+
                                  ') - ('+jPsiMass.GetName()+' > '+str(sbHighMass)+')'+
                                  '-2*('+jPsiMass.GetName()+' < '+str(sbLowMass)+')',
                                  RooArgList(jPsiMass,jPsicTau))
            

            cutStringPt = '('+jPsiPt.GetName()+' > '+str(ptMin)+' && '+jPsiPt.GetName()+' < '+str(ptMax)+')'
            cutStringY  = '( abs('+jPsiRap.GetName()+') > '+str(yMin)+' && abs('+jPsiRap.GetName()+') < '+str(yMax)+')'
            #cutStringM1 = '('+jPsiMass.GetName()+' < '+str(sigMinMass)+' && '+jPsiMass.GetName()+' > '+str(sbLowMass)+')'
            #cutStringM2 = '('+jPsiMass.GetName()+' < '+str(sbHighMass)+' && '+jPsiMass.GetName()+' > '+str(sigMaxMass)+')'
            #cutStringMT = '!('+cutStringM1+' || '+cutStringM2+')'
            cutString   = cutStringPt+' && '+cutStringY #+' && '+cutStringMT

            print cutString

            #get the reduced dataset we'll do the fit on
            binData = fullData.reduce(ROOT.RooFit.SelectVars(redVars),
                                      ROOT.RooFit.Cut(cutString),
                                      ROOT.RooFit.Name('data_rap'+str(rap_bin)+'_pt'+str(pt_bin+1)),
                                      ROOT.RooFit.Title('Data For Fitting'))

            binDataWithCategory = RooDataSet('data_rap'+str(rap_bin)+'_pt'+str(pt_bin+1),
                                             'Data For Fitting',
                                             fitVars)
            #categorize
            binData.addColumn(ctauStates)
            binData.addColumn(massStates)
            binData.addColumn(states)
            for ev in range(binData.numEntries()):
                args = binData.get(ev)
                                
                jPsiMass.setVal(args.find(jPsiMass.GetName()).getVal())
                jPsiRap.setVal(args.find(jPsiRap.GetName()).getVal())
                jPsiPt.setVal(args.find(jPsiPt.GetName()).getVal())
                jPsicTau.setVal(args.find(jPsicTau.GetName()).getVal())                    
                jPsicTauError.setVal(args.find(jPsicTauError.GetName()).getVal())
            
                jPsiHXcosth.setVal(args.find(jPsiHXcosth.GetName()).getVal())
                jPsiHXphi.setVal(args.find(jPsiHXphi.GetName()).getVal())

                massStates.setIndex(int(massFun.getVal()))
                ctauStates.setIndex(int(ctauFun.getVal()))
                states.setIndex(int(mlFun.getVal()))
                
                binDataWithCategory.add(fitVars)
            

            getattr(ws,'import')(binDataWithCategory)

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]
def prepare_truth_models(ws,cat,mass,channel,turnon,truth):    
    if channel in study_inputs:
        bkg_data = RooDataSet("bkgdata_%s"%(channel),
                              "M_{ll#gamma} with Errors",
                              ws.set("observables_weight"),
                              "weight")
        
        for k,file in enumerate(study_inputs[channel]):
            f_in = TFile.Open(file,"READ")
            gDirectory.cd(pwd)
            n_events = f_in.Get("eventCount").GetBinContent(1)
            tree = f_in.Get("selected_zg").CloneTree()
            tree.SetName("tree_%s_%i"%(channel,k))
            f_in.Close()            
            
            d_in = RooDataSet("temp",
                              "M_{ll#gamma} with Errors",
                              tree,
                              ws.set("observables"))
            norm = RooConstVar("norm","norm",n_events)
            weight = RooFormulaVar("weight",
                                   "weight",
                                   "1.3*@0*@1/@2", #1.3 gives 19.6/fb
                                   RooArgList(ws.var("puWeight"),
                                              ws.var("procWeight"),
                                              norm)
                                   )
            d_in.addColumn(weight)
            d_in_weight = RooDataSet("temp_weight",
                                     "M_{ll#gamma} with Errors",
                                     d_in,
                                     ws.set("observables_weight"),
                                     '','weight')
            bkg_data.append(d_in_weight)
            
        # split off data for each category, create the
        # toy dataset truth models
        data = bkg_data.reduce("Mz + Mzg > 185 && r94cat == %i"%cat)
        getattr(ws,'import')(data,
                             RooFit.Rename('bkgdata_%s_%i'%(channel,
                                                            cat)))
        nevts = data.sumEntries('Mzg > %f && Mzg < %f'%(mass-1.5,mass+1.5))
        #for sigm turn on we want erf truth
        if turnon == 'sigm' and truth == 'exp': 
            #make RooDecay 'truth' model with erf turn on
            ws.factory(
                'RooGaussModel::MzgResoShape_exp_erf_%s_cat%i(Mzg,'\
                'bias_exp_erf_%s_cat%i[120,90,150],sigma_exp_erf_%s_cat%i[1,0.01,10])'%(
                channel,cat,
                channel,cat,
                channel,cat)
                )
            ws.factory(
                'RooDecay::MzgTruthModelBase_exp_erf_%s_cat%i(Mzg,'\
                'tau_erf_%s_cat%i[25,0,50],MzgResoShape_exp_erf_%s_cat%i,'
                'RooDecay::SingleSided)'%(channel,cat,
                                          channel,cat,
                                          channel,cat)
                )            
            ws.factory(
                'RooExtendPdf::MzgTruthModel_exp_erf_%s_cat%i('\
                'MzgTruthModelBase_exp_erf_%s_cat%i,'\
                'norm_truth_exp_%s_cat%i[%f,%f,%f],"ROI")'%(channel,cat,
                                                            channel,cat,
                                                            channel,cat,
                                                            nevts,
                                                            0.25*nevts,1.75*nevts)
                )
            ws.pdf('MzgTruthModel_exp_erf_%s_cat%i'%(channel,cat)).fitTo(
                ws.data('bkgdata_%s_%i'%(channel,cat)),
                RooFit.Minimizer('Minuit2','scan'),
                RooFit.SumW2Error(False)
                )
            ws.pdf('MzgTruthModel_exp_erf_%s_cat%i'%(channel,cat)).fitTo(
                ws.data('bkgdata_%s_%i'%(channel,cat)),
                RooFit.Minimizer('Minuit','simplex'),
                RooFit.SumW2Error(False)
                )
            ws.pdf('MzgTruthModel_exp_erf_%s_cat%i'%(channel,cat)).fitTo(
                ws.data('bkgdata_%s_%i'%(channel,cat)),
                RooFit.SumW2Error(True)
                )
        if turnon == 'sigm' and truth == 'pow':
            #make power-law truth model with erf turn on
            ws.factory('EXPR::MzgTruthModelShape_pow_erf_%s_cat%i('\
                       '"1e-20 + (@0 > @1)*((@0)^([email protected]))",'\
                       '{Mzg,step_pow_erf_%s_cat%i[105,100,130],'\
                       'pow_%s_cat%i[2,0,10]})'\
                       %(channel,cat,
                         channel,cat,
                         channel,cat))
            ws.factory(
                'RooGaussModel::MzgResoShape_pow_erf_%s_cat%i(Mzg,'\
                'bias_pow_erf_%s_cat%i[0],sigma_pow_erf_%s_cat%i[1,0.01,10])'%(
                channel,cat,
                channel,cat,
                channel,cat)
                )
            ws.factory('FCONV::MzgTruthModelBase_pow_erf_%s_cat%i(Mzg,'\
                       'MzgTruthModelShape_pow_erf_%s_cat%i,'\
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]
obs=w.set("obs_1d")
uncVZ = w.var("uncVZ")
uncVZ.setBins(1000)
gauss_params = gauss_pdf.getParameters(obs)
gaussexp_params = gaussexp_pdf.getParameters(obs)

for i in range(1,n_massbins):
#for i in range(0,totalH.GetXaxis().GetNbins()-binning+2,binning):
    print i
    centermass=yieldhist.GetXaxis().GetBinCenter(i)
    massrange=2.5*(massres_a*centermass + massres_b)
    lowedge = centermass-massrange/2.0
    highedge = centermass+massrange/2.0
    name="Radiative vertex Z, mass [{}, {}] GeV".format(lowedge,highedge)
    #dataset = RooDataSet("data","data",events,w.set("myVars"),"abs(uncM-{0})<{1}".format(centermass,massrange/2))
    dataInRange = dataset.reduce(obs,"abs(uncM-{0})<{1}".format(centermass,massrange/2))
    if dataInRange.sumEntries()<100:
        continue

    frame=uncVZ.frame()
    frame.SetTitle(name)

    binnedData = dataInRange.binnedClone()
    binnedData.plotOn(frame)
    mean = binnedData.mean(uncVZ)
    sigma = binnedData.sigma(uncVZ)
    print "before gaussian fit: mean={0}, sigma={1}".format(mean,sigma)
    uncVZ.setRange("fitRange",mean-3*sigma,mean+3*sigma)
    gauss_params.setRealValue("mean",mean)
    gauss_params.setRealValue("sigma",sigma)
    #gauss_params.printLatex()

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import reduce [as 别名]
pdfConfig['timeEffHistFiles'].getSettings( [ ( 'runPeriod', 'p2012' ) ] )['file'] = timeEffFile2012
pdfConfig['anglesEffType'] = 'basisSig6'
pdfConfig['angEffMomsFiles'] = dataPath + 'angEffNominalRew_5thOrder.par'
pdfConfig['signalData'] = dataSetAsymW
pdfConfig['readFromWS'] = True

from P2VV.Parameterizations.FullPDFs import Bs2Jpsiphi_PdfBuilder as PdfBuilder
pdfBuild = PdfBuilder( **pdfConfig )
pdf = pdfBuild.pdf()
pdfConfig.readParametersFromFile( filePath = parFilePath )
pdfConfig.setParametersInPdf(pdf)

# create projection data sets
projSet = [ obs for obs in pdf.ConditionalObservables() ] + [ cat for cat in pdf.indexCat().inputCatList() ]
projSet = RooArgSet( obsSet.find( var.GetName() ) for var in projSet )
projDataSets = dict(  plus  = dataSetAsymW.reduce( Name = 'asymDataWPlus',  SelectVars = projSet, Cut = 'asymCat==+1' )
                    , minus = dataSetAsymW.reduce( Name = 'asymDataWMinus', SelectVars = projSet, Cut = 'asymCat==-1' )
                   )
print 'pdfAsymmetry: projection data sets:'
projDataSets['plus'].Print()
projDataSets['minus'].Print()

# create time ranges for time bins
rangeNames = dict( )
timePoints = dict( )
lastBin = False
if not timeBins : timeBins = range(8)
if not periods : periods = range( int( ws['time'].getMax() / oscPeriod - binOffset / float(numTimeBinsTot) ) + 1 )
for per in periods :
    for bin in timeBins :
        timeLo = ( float(per) + ( float(bin) + binOffset ) / float(numTimeBinsTot) ) * oscPeriod