Python RooRealVar.setVal方法代码示例

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
def getTrueOmegasPerCat(config, mistagobs, mistag, mistagpdf):
    """
    estimate the per-category mistag (omega_i) as starting value for the fit

    calculate true per-category omegas based on the mistagpdf and the
    calibration that goes into the generation pdf

    config      -- config dictionary
    mistagobs   -- mistag observable (eta)
    mistag      -- calibrated mistag (omega(eta))
    mistagpdf   -- mistag pdf (P(eta))

    returns a pair of average mistag omega, list of per-category mistag
    (omega) averages

    relevant config dictionary keys:
    'MistagCategoryBinBounds':
        list of mistag category bin bounds (one more entry than mistag
        categories)
    """
    from ROOT import RooRealVar, RooCustomizer, RooProduct, RooArgList, RooArgSet
    eta1 = RooRealVar('eta1', 'eta1', mistagobs.getMin(),
            mistagobs.getMin(), mistagobs.getMax())
    eta2 = RooRealVar('eta2', 'eta2', mistagobs.getMax(),
            mistagobs.getMin(), mistagobs.getMax())
    prod = RooProduct('prod', 'prod', RooArgList(mistag, mistagpdf))
    oldmistagobs = mistagobs
    mistagobs = mistagobs.clone(mistagobs.GetName() + '_catclone')
    ROOT.SetOwnership(mistagobs, True)
    mistagobs.setRange(eta1, eta2)
    c = RooCustomizer(prod, 'cust')
    c.replaceArg(oldmistagobs, mistagobs)
    prod = c.build()
    ROOT.SetOwnership(prod, True)
    c = RooCustomizer(mistagpdf, 'cust2')
    c.replaceArg(oldmistagobs, mistagobs)
    pdf = c.build()
    ROOT.SetOwnership(pdf, True)
    if pdf.InheritsFrom('RooHistPdf'): pdf.forceNumInt()
    evnumer = prod.createIntegral(RooArgSet(mistagobs))
    evdenom = pdf.createIntegral(RooArgSet(mistagobs))
    totevdenom = evdenom.getVal()
    avomega = evnumer.getVal() / totevdenom
    omegas = [ ]
    for i in xrange(0, len(config['MistagCategoryBinBounds']) - 1):
        eta1.setVal(config['MistagCategoryBinBounds'][i])
        eta2.setVal(config['MistagCategoryBinBounds'][i + 1])
        omegas.append(evnumer.getVal() / evdenom.getVal())
    print 'INFO: Mistag calibration %s:' % mistag.GetName()
    print 'INFO:                Average omega (PDF): %g' % avomega
    print 'INFO:  Per category average omegas (PDF): %s' % str(omegas)
    return avomega, omegas

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
def roofit_poisson_unbinned(data):
    """returns lambda, error of lambda"""

    x = RooRealVar('x', 'x', 0, max(data) * 10)
    lam = RooRealVar('lambda', 'lambda', 0.1, 0.000001, max(data))

    model = RooPoisson('model', 'model', x, lam)

    dataset = RooDataSet('data', 'data', RooArgSet(x))
    for val in data:
        x.setVal(val)
        dataset.add(RooArgSet(x))

    model.fitTo(dataset, RooFit.Save(), RooFit.PrintLevel(-1))
    return lam.getVal(), lam.getError()

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
    def __init__( self, roovar, **kwargs ):
        '''
        To build the class, the variable to be blinded must be provided. In < scale >
        one modifies the range of values for the blinded variable. Also the names for
        the constant variable and the generated formula can be modified.
        '''
        
        name = roovar.GetName()
        print 'Defining a new blinded variable for <', name, '>'

        scale  = kwargs.get( 'scale', 1000 )
        eqname = kwargs.get( 'bvarname', name + '_BlindExpr' )

        ''' The seed is different for each call, which guarantees the blinding procedure '''
        print 'Generating new set of random bounds'
        rndm = TRandom3( 0 )
        vmin = roovar.getMin()
        boundlo = rndm.Uniform( vmin, vmin*scale )
        vmax = roovar.getMax()
        boundhi = rndm.Uniform( vmax, vmax*scale )

        clonename = name + '_BLIND'
        blindVar = RooRealVar( clonename, clonename, roovar.getVal(), boundlo, boundhi )
        print 'Created clone variable named <', clonename, '>'
        
        alpha = ( vmax - vmin )/( boundhi - boundlo )
        beta  = ( vmax*boundlo - vmin*boundhi )/( vmax - vmin )

        blindVar.setVal( alpha*( roovar.getVal() - beta ) )

        alpha = RooConstVar( 'alpha', 'alpha', alpha )
        beta  = RooConstVar( 'beta', 'beta', beta )
        
        formula = 'alpha*( %s - beta)' % clonename
        varlist = RooArgList( blindVar, alpha, beta )
        blindEq = RooFormulaVar( eqname, eqname, formula, varlist )
        print 'The blinding formula is:', formula
        
        self.Alpha    = alpha
        self.Beta     = beta
        self.BlindEq  = blindEq
        self.BlindVar = blindVar

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
def test_plot():
    c = TCanvas('c', 'canvas', 300, 300)

    x = RooRealVar("x","x",-10,10)
    mean = RooRealVar("mean","mean of gaussian",1,-10,10) 
    sigma = RooRealVar("sigma","width of gaussian",1,0.1,10)

    gauss = RooGaussian("gauss","gaussian PDF",x,mean,sigma) 

    xframe = x.frame(RooFit.Title("Gaussian p.d.f."))
    #gauss.plotOn(xframe)
    sigma.setVal(3)
    
    as_x = RooArgSet(x)
    data = gauss.generate(as_x,10000)
    data.plotOn(xframe, RooFit.MarkerSize(0.6), RooFit.MarkerStyle(20)) 
    xframe.Draw()

    pdffile = 'test.pdf'
    c.SaveAs(pdffile)

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
def dump_simple_simul():
    # idea: take two data sets, both low counts, but same lambda
    # and fit said lambda simultanously with both sets

    r_x = RooRealVar('x', 'x', 0, bins)

    r_lam = RooRealVar('lam', 'lam', 0.5, 0.0, 1.5)
    r_lam1 = RooRealVar('lam1', 'lam1', 0.5, 0.0, 1.5)
    r_lam2 = RooRealVar('lam2', 'lam2', 0.5, 0.0, 1.5)

    model1 = RooPoisson('pois1', 'pois1', r_x, r_lam1)
    model2 = RooPoisson('pois2', 'pois2', r_x, r_lam2)

    r_index = RooCategory('index', 'index')
    r_index.defineType('1')
    r_index.defineType('2')

    simul_model = RooSimultaneous('model', 'model', r_index)
    simul_model.addPdf(model1, '1')
    simul_model.addPdf(model2, '2')

    data1 = RooDataSet('', '', RooArgSet(r_x))
    for val in unbinned_from_binned(xdata, ydata):
        r_x.setVal(val)
        data1.add(RooArgSet(r_x))
    r_index.setLabel('1')
    data1.addColumn(r_index)

    data2 = RooDataSet('', '', RooArgSet(r_x))
    for val in unbinned_from_binned(xdata, ydata2):
        r_x.setVal(val)
        data2.add(RooArgSet(r_x))
    r_index.setLabel('2')
    data2.addColumn(r_index)

    data1.append(data2)

    _result = simul_model.fitTo(data1)
    print(r_lam.getVal(), '+-', r_lam.getError(), lam, np.abs(r_lam.getVal() - lam) / lam)
    print(r_lam1.getVal(), '+-', r_lam1.getError(), lam, np.abs(r_lam1.getVal() - lam) / lam)
    print(r_lam2.getVal(), '+-', r_lam2.getError(), lam, np.abs(r_lam2.getVal() - lam) / lam)

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
def dump_simple():
    # try poisson
    roo_lam = RooRealVar("lambda", ".", lam, 0.0, 1.5)
    roo_x = RooRealVar('x', 'x range', 0, bins)

    model = RooPoisson("poisson", "Poisson Model", roo_x, roo_lam)

    #data = model.generate(RooArgSet(roo_x), n_events)

    data = RooDataSet('data', 'Data', RooArgSet(roo_x))
    for val in unbinned_from_binned(xdata, ydata):
        roo_x.setVal(val)
        data.add(RooArgSet(roo_x))

    #// --- Perform extended ML fit of composite PDF to toy data ---
    fitresult = model.fitTo(data, RooFit.Save(), RooFit.PrintLevel(-1)) ;

    #// --- Plot toy data and composite PDF overlaid ---
    mesframe = roo_x.frame(bins) ;
    data.plotOn(mesframe) ;
    model.plotOn(mesframe) ;
    mesframe.Draw()

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
		ez1_idx_trk0 = aa.ez1_idx_trk0
		ez1_idx_trk1 = aa.ez1_idx_trk1
		ez1_idx_gam  = aa.ez1_idx_gam
                
		for j in range(len(ez1_m)):
			mu1_idx  = int(ez1_idx_trk0[j])
			mu2_idx  = int(ez1_idx_trk1[j])
			gam_idx  = int(ez1_idx_gam[j])

			mu1  = get_trk_vars(aa, mu1_idx)
			mu2  = get_trk_vars(aa, mu2_idx)
			gam  = get_part_vars(aa, 'gam', gam_idx)
			#print ez1_m[j]
			hist2.Fill(ez1_m[j])
			if ( ez1_m[j] > m2.getMin() and ez1_m[j] < m2.getMax() ):
				m2.setVal(ez1_m[j])
                                gpt.setVal(gam.pt)
				geta.setVal(gam.eta)
				gp.setVal(gam.p)
                                gamma_pt.Fill(gam.pt)
                                gamma_eta.Fill(gam.eta)
				ds2.add(argset2)
	  

oFile = TFile("/user2/sfarry/workspaces/energyflow/tuples/eta_data2016_turbo_embedded_mumugammaonly.root", "RECREATE")
ds.Write()
ds2.Write()
hist.Write()
hist2.Write()
gamma_pt.Write()
fakegamma_pt.Write()

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
# gauss 15
mean15 = RooRealVar("mean15","mean of gaussian",4,-10,10)
sigma15 = RooRealVar("sigma15","width of gaussian",1.3,0.1,10)

gauss15 = RooGaussian("gauss15","gaussian PDF",x,mean15,sigma15)

# convert PDF to extended PDF 

nGauss14 =  RooRealVar("nGauss14", "Gauss 14 ext normalization",  2000,  0., 6000)
gauss14_ext = RooExtendPdf("gauss14_ext",  "extended p.d.f", gauss14,  nGauss14)

nGauss15 =  RooRealVar("nGauss15", "Gauss 15 ext normalization",  1000,  0., 6000)
gauss15_ext = RooExtendPdf("gauss15_ext",  "extended p.d.f", gauss15,  nGauss15)

#nGauss15.setVal(nDataGauss_right_SB)
nGauss15.setVal( n_generate*(1-frac_combine4.getVal() ) )
nGauss15.setConstant(True)

# comine ext PDF

pdf_ext_combine5 = RooAddPdf("pdf_ext_combine5"," gauss14_ext + gauss15_ext ", RooArgList(gauss14_ext , gauss15_ext ), RooArgList( nGauss14 , nGauss15 ))

pdf_ext_combine5.fixAddCoefRange("whole_range")

# fit

#pdf_ext_combine5.fitTo(data4, RooFit.Range("whole_range") )
pdf_ext_combine5.fitTo(data4_SB, RooFit.Range("left_side_band_region,right_side_band_region") )

print ""
print "after fit"

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [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 RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
# make dataset from tree
wt = RooRealVar('wt', '', 1.0, 0.0, 1.0)
dst = RooDataSet('dst', '', RooArgSet(time, wt), RooFit.WeightVar(wt))
for tp in ntuples:
    print '{}: {}'.format(tp.GetName(), tp.GetEntries())
    tp.SetBranchStatus('*', 0)
    tp.SetBranchStatus('wt', 1)
    tp.SetBranchStatus('Bid', 1)
    tp.SetBranchStatus('hid', 1)
    tp.SetBranchStatus('time', 1)
    for i in range(tp.GetEntries()):
        tp.GetEntry(i)
        # choose decay equation (1 of 2 for Dspi)
        if (tp.Bid == 531 and tp.hid == 211):
            time.setVal(tp.time)
            dst.add(RooArgSet(time), tp.wt)
dst.Print()

# model.fitTo(dst)

# plot
tfr = time.frame()
dst.plotOn(tfr)
model.plotOn(tfr)
tfr.Draw()
mode_title = get_title_from_mode(mode)
if nokfactor:
    k_title = 'no'
else:
    k_title = 'with'

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

#.........这里部分代码省略.........
    dataArgSet.add( probFit3S)

 

    print "Creating DataSet"
    dataSet = RooDataSet("chicds","Chic RooDataSet", dataArgSet)

    entries = tree.GetEntries()
    print entries

    if nevents is not 0:
        entries = nevents

    for ientry in range(0,entries):
        tree.GetEntry(ientry)

        # unfort ntuples are slightly different for chic and chib

        if applyscale:
            
            if usekinfit :    
                spatial = tree.rf3S_photon_p4.Vect()
                #spatial = tree.photon_p4.Vect()
                spatial *= (1/escale)
                corr_photon_p4=TLorentzVector()
                #corr_photon_p4.SetVectM(spatial,tree.rf1S_photon_p4.M())
                corr_photon_p4.SetVectM(spatial,0)
                corr_chi_p4 = tree.rf3S_dimuon_p4 + corr_photon_p4

            else:
                        
                spatial = tree.photon_p4.Vect()
                spatial *= (1/escale)
                corr_photon_p4=TLorentzVector()
                corr_photon_p4.SetVectM(spatial,tree.photon_p4.M())
                corr_chi_p4 = tree.dimuon_p4 + corr_photon_p4
                
            
        else :
            corr_chi_p4 = tree.chi_p4
        
        if type == 'chic':
            
            m_mumu.setVal(tree.dimuon_p4.M())
            y_mumu.setVal(tree.dimuon_p4.Rapidity())        
            pt_mumu.setVal(tree.dimuon_p4.Pt())
            eta_gamma.setVal(tree.photon_p4.Eta())
            pt_gamma.setVal(tree.photon_p4.Pt())
            m_gamma.setVal(tree.photon_p4.M())
            m_chi_rf1S.setVal(tree.rf1S_chi_p4.M())
            m_chi_rf2S.setVal(tree.rf2S_chi_p4.M())
        
            if usekinfit : Qvalue.setVal(corr_chi_p4.M())
            else: Qvalue.setVal((corr_chi_p4).M() - tree.dimuon_p4.M())
                
            #Qvalue.setVal((tree.chi_p4).M()**2 - tree.dimuon_p4.M()**2)
            psi1S_nsigma.setVal(tree.psi1S_nsigma)
            psi2S_nsigma.setVal(tree.psi2S_nsigma)  
            psi3S_nsigma.setVal(0)
            
        elif type == 'chib':

            m_mumu.setVal(tree.dimuon_p4.M())
            y_mumu.setVal(tree.dimuon_p4.Rapidity())        
            pt_mumu.setVal(tree.dimuon_p4.Pt())
            eta_gamma.setVal(tree.photon_p4.Eta())
            pt_gamma.setVal(tree.photon_p4.Pt())
            m_chi_rf1S.setVal(tree.rf1S_chi_p4.M())
            m_chi_rf2S.setVal(tree.rf2S_chi_p4.M())

        
            if usekinfit : Qvalue.setVal(corr_chi_p4.M())
            #if usekinfit : Qvalue.setVal(tree.rf3S_chi_p4.M()) #uncorrected
            else: Qvalue.setVal(corr_chi_p4.M() - tree.dimuon_p4.M())
                
        
            psi1S_nsigma.setVal(tree.Y1S_nsigma)
            psi2S_nsigma.setVal(tree.Y2S_nsigma) 
            psi3S_nsigma.setVal(tree.Y3S_nsigma)
        
        ctpv.setVal(tree.ctpv)
        ctpv_error.setVal(tree.ctpv_error)
        pi0_abs_mass.setVal(tree.pi0_abs_mass)

        rho_conv.setVal(tree.conv_vertex)
        dz.setVal(tree.dz)
        
        probFit1S.setVal(tree.probFit1S)
        probFit2S.setVal(tree.probFit2S)
        probFit3S.setVal(tree.probFit3S)
        
        
        if selectchi1:
            if (  tree.chic_pdgId == 20443): dataSet.add(dataArgSet)
        else :
            dataSet.add(dataArgSet)
        

    outfile = TFile(outfile_name,'recreate')    
    dataSet.Write()

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
def fit_mbc():
    from tools import normalizedRooFitIntegral, RooFitIntegral
    from ROOT import RooRealVar, RooDataSet, RooArgList, RooArgSet, \
         RooGaussian, RooArgusBG, RooCBShape, RooAddPdf, RooPolynomial, \
         RooDataHist, RooFit, kTRUE, kFALSE

    signal_margins = [1.86, 1.87]
    sb_margins = [1.84, 1.85]

    # Right here we compute background yield
    mbc = RooRealVar('mbc', 'mbc', 1.83, 1.89, 'GeV')
    arg_cutoff = RooRealVar('arg_cutoff', 'Argus cutoff', 1.8869, 1.885, 1.888, 'GeV')
    arg_slope = RooRealVar('arg_slope', 'Argus slope', -13, -100, -1)
    mbc_d0 = RooRealVar('mbc_d0', 'D0 Mbc', 1.8647, 'GeV')
    mbc_dp = RooRealVar('mbc_dp', 'D+ Mbc', 1.8694, 'GeV')
    mbc_float = RooRealVar('mbc_float', 'Floating D mass', 1.869, 1.855, 1.875, 'GeV')
    sigma = RooRealVar('sigma', 'D width', 0.00145, 0.0001, 0.0025, 'GeV')
    sigma2 = RooRealVar('sigma2', 'CB width', 0.00145, 0.0001, 0.005,
                        'GeV')
    alpha = RooRealVar('alpha', 'CB shape cutoff', -1.515, -2., 2)
    n = RooRealVar('n', 'CB tail parameter', 6, 0, 20)
    gauss_d0 = RooGaussian('gauss_d0', 'D0 gaussian', mbc, mbc_d0, sigma2)
    gauss_dp = RooGaussian('gauss_dp', 'D+ gaussian', mbc, mbc_dp, sigma2)
    gauss_float = RooGaussian('gauss_float', 'Floating gaussian',
                                                    mbc, mbc_float, sigma2)
    cb_d0 = RooCBShape('cb_d0', 'D0 Crystal Barrel', mbc,
                   mbc_d0, sigma, alpha, n)
    cb_dp = RooCBShape('cb_dp', 'D+ Crystal Barrel', mbc,
                       mbc_dp, sigma, alpha, n)
    cb_float = RooCBShape('cb_float', 'Floating Crystal Barrel', mbc,
                          mbc_float, sigma, alpha, n)
    argus = RooArgusBG('argus', 'Argus BG', mbc, arg_cutoff, arg_slope)
    yld = RooRealVar('yield', 'D yield', 25700, 0, 100000)
    yld2 = RooRealVar('yield2', '2nd yield', 100, 0, 2000)
    bkg = RooRealVar('bkg', 'Background', 1300, 0, 40000)
    a = RooRealVar('a', 'Norm', 1)
    poly = RooPolynomial('poly', 'poly PDF', mbc, RooArgList(a), 0)
    sumpdf_d0 = RooAddPdf('sumpdf_d0', 'D0 sum pdf',
                          RooArgList(cb_d0, argus),
                          RooArgList(yld, bkg))
    sumpdf_dp = RooAddPdf('sumpdf_dp', 'Dp sum pdf',
                          RooArgList(cb_dp, argus),
                          RooArgList(yld, bkg))
    sumpdf_float = RooAddPdf('sumpdf_float', 'Generic D sum pdf',
                             RooArgList(cb_float, argus),
                             RooArgList(yld, bkg))

    width_modes = { '0': 0.00150, '1': 0.001831, '3': 0.001426, '200': 0.001387,
                    '202': 0.001407 }

    n_modes = { '0': 2.68, '1': 4.06, '3': 4.34, '200': 4.05, '202': 5.26 }
    alpha_modes = { '0': -1.6145, '1': -1.4562, '3': -1.5834, '200': -1.6538,
                    '202': -1.5598 }


    pdf = sumpdf_float


    sigma.setVal(width_modes['1']) # sigma.setConstant()


    n.setVal(n_modes['1']) #n.setConstant()

    alpha.setVal(alpha_modes['1']) #alpha.setConstant()
    #sigma.setConstant()
    #arg_cutoff.setVal(1.8865); #arg_cutoff.setConstant()

    c1.Divide(1,2)
    c1.cd(1)
    dset = RooDataHist('dsetmc', 'title', RooArgList(mbc), h_mbc['mc'])

    #pdf.fitTo(dset, 'eq')
    Extended = RooFit.Extended(kTRUE) # e
    Verbose = RooFit.Verbose(kFALSE) #q

    pdf.fitTo(dset, Extended, Verbose)

    # xframe = mbc.frame()
    # dset.plotOn(xframe)
    # pdf.plotOn(xframe)
    # pdf.paramOn(xframe,dset)
    # xframe.SetTitle('Fake type 1, MC')
    # xframe.Draw()
    c1.cd(2)
    dset = RooDataHist('dsetdata', 'title', RooArgList(mbc), h_mbc['data'])
    #pdf.fitTo(dset, 'eq')
    pdf.fitTo(dset, Extended, Verbose)
    # xframe = mbc.frame()
    # dset.plotOn(xframe)
    # pdf.plotOn(xframe)
    # pdf.paramOn(xframe,dset)
    # xframe.SetTitle('Fake type 1, data')
    # xframe.Draw()
    sb_scale = (normalizedRooFitIntegral(argus, mbc, signal_margins[0],
                                         signal_margins[1])/
                normalizedRooFitIntegral(argus, mbc, sb_margins[0], sb_margins[1]))

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
def mbc_dline_che(evtfile, mc, setMres, setGamma, setR, sp1, sp2, sp3, fa,
                  fb, setmd, setp, setxi, setN1, setN2, setNbkgd1, setNbkgd2,
                  title1, title2, epsfile, txtfile, ymin=0.5,
                  cuts=None, err_type='SYMM', test=False):
    
    from ROOT import (gROOT, RooRealVar, RooCategory, RooArgSet, RooDataSet,
                      RooFit, RooGaussian, RooArgList, RooAddPdf, RooSimultaneous,
                      RooArgusBG, RooFormulaVar, RooChebychev, RooAbsData,
                      RooDataHist, TCanvas, kRed, kBlue, kGreen, kMagenta,
                      TPaveText, RooDLineShape)
    set_root_style(stat=1, grid=0)

    mbc = RooRealVar('mbc', 'Beam constrained mass', 1.83, 1.89, 'GeV')
    ebeam = RooRealVar('ebeam','Ebeam', 1.8815, 1.892, 'GeV')

    dflav = RooCategory('dflav','D0 flavor')
    dflav.defineType('dflav',1)
    dflav.defineType('dbarflav',-1)

    if cuts != None:
        if 'kkmass' in cuts:
            kkmass = RooRealVar('kkmass', 'KK invariant mass', 0.97, 1.90, 'GeV')
            ras = RooArgSet(mbc, ebeam, kkmass, dflav)
            dataset = RooDataSet.read(evtfile, ras)
        else:
            raise NameError(cuts)

        sys.stdout.write('Using cuts: %s...' %cuts)
        dataset = dataset.reduce(cuts)
        sys.stdout.write(' selected %s events.\n' % dataset.numEntries())
    else:
        ras = RooArgSet(mbc, ebeam, dflav)
        dataset = RooDataSet.read(evtfile, ras)

    res = RooRealVar("datares", "datares", mc)
    mres = RooRealVar("mres","mres", setMres)
    gamma = RooRealVar('gamma', 'gamma', setGamma)

    r = RooRealVar('r', 'r', setR)
    sigmaE = RooRealVar("sigmaE","sigmaE", 0.0021)

    sigmap1 = RooRealVar("sigmap1","sigmap1", sp1, 0.002, 0.040)

    scalep2 = RooRealVar("scalep2","scalep2",2.00,1.500,5.500)
    scalep3 = RooRealVar("scalep3","scalep3",5.00,3.00,10.000)
   
    scalep2.setVal(sp2)
    scalep2.setConstant(1)
    scalep3.setVal(sp3)
    scalep3.setConstant(1)

    as12 = RooArgList(sigmap1,scalep2)
    sigmap2 = RooFormulaVar("sigmap2","sigma2","sigmap1*scalep2", as12)

    as123 = RooArgList(sigmap1,scalep2,scalep3)
    sigmap3 = RooFormulaVar("sigmap3","sigma3","sigmap1*scalep2*scalep3",
                            as123)
    
    md = RooRealVar("md","md", setmd,1.863,1.875)

    f2 = RooRealVar("f2","f2", fa)
    f3 = RooRealVar("f3","f3", fb)
    al23 = RooArgList(f2,f3)
    f1 = RooFormulaVar("f1","f1","1.0-f2-f3", al23)

    # Construct signal shape

    fcn1_1 = RooDLineShape("DLineshape1_1","DLineShape1_1",4,mbc,ebeam,
                           mres,gamma,r,sigmaE,sigmap1,md,res)
    fcn1_2 = RooDLineShape("DLineshape1_2","DLineShape1_2",4,mbc,ebeam,
                           mres,gamma,r,sigmaE,sigmap2,md,res)
    fcn1_3 = RooDLineShape("DLineshape1_3","DLineShape1_3",4,mbc,ebeam,
                           mres,gamma,r,sigmaE,sigmap3,md,res)

    fcn2_1 = RooDLineShape("DLineshape2_1","DLineShape2_1",4,mbc,ebeam,
                           mres,gamma,r,sigmaE,sigmap1,md,res)
    fcn2_2 = RooDLineShape("DLineshape2_2","DLineShape2_2",4,mbc,ebeam,
                           mres,gamma,r,sigmaE,sigmap2,md,res)
    fcn2_3 = RooDLineShape("DLineshape2_3","DLineShape2_3",4,mbc,ebeam,
                           mres,gamma,r,sigmaE,sigmap3,md,res)


    alf1_123 = RooArgList(fcn1_1,fcn1_2,fcn1_3)
    af12 = RooArgList(f1,f2) 
    sigpdf = RooAddPdf("signal1_3","signal1_3", alf1_123, af12)

    alf2_123 = RooArgList(fcn2_1,fcn2_2,fcn2_3)
    
    sigbarpdf = RooAddPdf("signal2_3","signal2_3", alf2_123, af12)

    con0 = RooRealVar('c0', 'constant', -1, 1)
    con1 = RooRealVar('c1', 'linear', -10, 10)
    con2 = RooRealVar('c2', 'quadratic', 1)

    bkgpdf = RooChebychev('bkgpdf', 'Background', mbc, RooArgList(con1, con2))
    
    bkgbarpdf = RooChebychev('bkgbarpdf', 'Background',
                             mbc, RooArgList(con1, con2))

    yld = RooRealVar('yld', 'D yield', 100, 0, 2000)
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
 (800, (4.256970608754619, 0.21145030944619475)),
 (900, (2.1511503582230613, 0.10749890705342333)),
 (1000, (1.2694182781976522, 0.06340890314406443)),
 (1100, (0.643681492794392, 0.03217545331472744)),
 (1500, (0.08287990689532962, 0.004117232695182858)),
 (1750, (0.026150402908862617, 0.0013037026701146663)),
 (2000, (0.008063601507094236, 0.00040306052931164015)),
 (2250, (0.003056637009041868, 0.00015253920943590022)),
 (4000, (1.5147902166034461e-06, 7.553510014249738e-08)),
 #(5000, (6.7752618437184105e-09, 3.3855147194005573e-10)),
 #(6000, (1.673646004156878e-10, 8.333326606880855e-12)),
 #(7000, (4.7172685387069244e-11, 2.356523443723669e-12))
]

for massval, (xsval, xserrval) in data:
    mass.setVal(massval); xs.setVal(xsval); xs.setError(xserrval)
    ds.add(RooArgSet(mass, xs))

#c0pr = RooRealVar("c0pr", "c0pr", 1, 1e-5, 1e5)
c0 = RooRealVar("c0", "c0",  3.13e-06,  1e-8, 1e5)
c1 = RooRealVar("c1", "c1", -1.59e+01, -1e+2, 1e2)
c2 = RooRealVar("c2", "c2",  1.08e+00,  1e-5, 1e1)
c3 = RooRealVar("c3", "c3",  7.66e+01,  0, 1e2)
c4 = RooRealVar("c4", "c4", -4.97e-03, -1e-5, 1)
#c5 = RooRealVar("c5", "c5",  0.01); c5.setConstant()

partial_pdf = RooGenericPdf("xspdf", "xspdf", "(mass^(c1+c2*log(mass))*exp(c3+c4*mass))",
    RooArgList(mass, c1, c2, c3, c4))
    
pdf = R.RooExtendPdf("fullpdf", "fullpdf", partial_pdf, c0)

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setVal [as 别名]
def makeRooDataSet(type,infile_name,outfile_name,tree_name,nevents):
    """ Make RooDataSets from TTrees"""




    inputfile = TFile.Open(infile_name,"READ")
    print "Importing tree"
    tree = TTree()
    inputfile.GetObject(tree_name, tree)  #get the tree from the data file

   

    #define variables for the RooDataSet
    m_mumu        = RooRealVar("m_mumu",   "m_mumu",   0.0, 4.0)
    y_mumu        = RooRealVar("y_mumu",   "y_mumu",   0.0, 2.0 )
    pt_mumu       = RooRealVar("pt_mumu",  "pt_mumu",  0.0, 260.0)
    eta_gamma     = RooRealVar("eta_gamma","eta_gamma",-3.5, 3.5)
    pt_gamma      = RooRealVar("pt_gamma", "pt_gamma", 0.0, 100.0)
    m_gamma       = RooRealVar("m_gamma",  "m_gamma",  -0.1,0.1)
  
    m_chi_rf1S    = RooRealVar("m_chi_rf1S", "m_chi_rf1S", 0.0, 7.0)
    m_chi_rf2S    = RooRealVar("m_chi_rf2S", "m_chi_rf2S", -1.0, 1.0)
    #Qvalue        = RooRealVar("Qvalue","Q", -15., 15.)
    s             = RooRealVar("s","s", -10., 10.)
    
    ctpv          = RooRealVar("ctpv","ctpv", -1.0, 3.5)
    ctpv_error    = RooRealVar("ctpv_err","ctpv_err", -1.0, 1.0)
    pi0_abs_mass  = RooRealVar("pi0_abs_mass","pi0_abs_mass", 0.0, 2.2)
    psi1S_nsigma  = RooRealVar("psi1S_nsigma","psi1S_nsigma",0.0,1.0)
    psi2S_nsigma  = RooRealVar("psi2S_nsigma","psi2S_nsigma",0.0,1.0)
    psi3S_nsigma  = RooRealVar("psi3S_nsigma","psi3S_nsigma",0.0,1.0)
    rho_conv      = RooRealVar("rho_conv", "rho_conv", 0.0, 70.0)
    dz            = RooRealVar("dz","dz", -1.0, 1.0)
    probFit1S     = RooRealVar("probFit1S","probFit1S",0,1)
    probFit2S     = RooRealVar("probFit2S","probFit2S",0,1)

    dataArgSet = RooArgSet(m_mumu,
                           y_mumu,
                           pt_mumu,
                           eta_gamma,
                           pt_gamma,
                           m_gamma, 
                           m_chi_rf1S)
    
    dataArgSet.add( m_chi_rf2S )
    dataArgSet.add( s )
    dataArgSet.add( ctpv )
    dataArgSet.add( ctpv_error )
    dataArgSet.add( pi0_abs_mass )
    dataArgSet.add( psi1S_nsigma )
    dataArgSet.add( psi2S_nsigma )
    dataArgSet.add( rho_conv )
    dataArgSet.add( dz )
    dataArgSet.add( probFit1S )
    dataArgSet.add( probFit2S )


    print "Creating DataSet"
    dataSet = RooDataSet("chicds","Chic RooDataSet", dataArgSet)

    entries = tree.GetEntries()
    print entries

    if nevents is not 0:
        entries = nevents

    for ientry in range(0,entries):
        tree.GetEntry(ientry)

        # unfort ntuples are slightly different for chic and chib
        
        if type == 'chic':
            
            m_mumu.setVal(tree.dimuon_mass)
            y_mumu.setVal(tree.dimuon_rapidity)        
            pt_mumu.setVal(tree.dimuon_pt)
            eta_gamma.setVal(tree.photon_eta)
            pt_gamma.setVal(tree.photon_pt)
            #m_gamma.setVal(tree.photon_p4.M())
            m_chi_rf1S.setVal(tree.rf1S_chic_mass)
            #m_chi_rf1S.setVal(tree.rf2S_chi_p4.M())
            #Qvalue.setVal((tree.chi_p4).M() - tree.dimuon_p4.M())
            #Qvalue.setVal((tree.chi_p4).M()**2 - tree.dimuon_p4.M()**2)
            #Qvalue.setVal((tree.rf1S_chic_mass**2 -tree.dimuon_mass**2)
            #              / (3.5107**2 - 3.0969**2 ) -1)

            
#           this should be the correct one if the refitted variable was available
#           s.setVal((tree.rf1S_chic_mass**2 - tree.rf1S_dimuon_p4.M()**2)/ (3.5107**2 - 3.0969**2 ) -1)
             
            s.setVal((tree.rf1S_chic_mass**2 - 3.0969**2)/ (3.5107**2 - 3.0969**2 ) -1)


            
            psi1S_nsigma.setVal(tree.psi1S_nsigma)
            psi2S_nsigma.setVal(0)  
            psi3S_nsigma.setVal(0)
            
        elif type == 'chib':
#.........这里部分代码省略.........