Beispiel #1
0
def doFit(ws,options):
    rap_bins = range(1,len(jpsi.pTRange))
    pt_bins = None
    
    if options.testBin is not None:
        rap_bins = [int(options.testBin.split(',')[0])]
        pt_bins  = [int(options.testBin.split(',')[1])-1]
        
    for rap_bin in rap_bins:
        if options.testBin is None:
            pt_bins = range(len(jpsi.pTRange[rap_bin]))
        for pt_bin in pt_bins:

            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]

            jPsiMass = ws.var('JpsiMass')
            jPsicTau = ws.var('Jpsict')

            jPsiMass.setRange('mlfit_prompt',2.7,3.5)
            jPsiMass.setRange('mlfit_nonPrompt',2.7,3.5)            

            jPsiMass.setRange('NormalizationRangeFormlfit_prompt',2.7,3.5)
            jPsiMass.setRange('NormalizationRangeFormlfit_nonPrompt',2.7,3.5)

            jPsicTau.setRange('mlfit_signal',-1,2.5)
            jPsicTau.setRange('mlfit_leftMassSideBand',-1,2.5)
            jPsicTau.setRange('mlfit_rightMassSideBand',-1,2.5)

            jPsicTau.setRange('NormalizationRangeFormlfit_signal',-1,2.5)
            jPsicTau.setRange('NormalizationRangeFormlfit_leftMassSideBand',-1,2.5)
            jPsicTau.setRange('NormalizationRangeFormlfit_rightMassSideBand',-1,2.5)

            #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)
            
            #jPsicTau.setRange('mlfit_promptSignal',-1,.1)
            #jPsicTau.setRange('mlfit_nonPromptSignal',.1,2.5)
            #jPsicTau.setRange('mlfit_leftMassSideBand',-1,2.5)
            #jPsicTau.setRange('mlfit_rightMassSideBand',-1,2.5)

            #reset parameters
            ws.var('CBn').setVal(.5)
            ws.var('CBalpha').setVal(.5)
            ws.var('CBmass').setVal(3.1)
            ws.var('CBsigma').setVal(.02)
            ws.var('bkgLambda').setVal(0)
            
            ws.var('bkgTauSSDL').setVal(.5)
            #ws.var('bkgTauFDL').setVal(.5)
            ws.var('bkgTauDSDL').setVal(.5)
            ws.var('fBkgSSDL').setVal(.5)
            ws.var('fBkgLR').setVal(.5)

            ws.var('bkgTauSSDR').setVal(.5)
            #ws.var('bkgTauFDR').setVal(.5)
            ws.var('bkgTauDSDR').setVal(.5)
            ws.var('fBkgSSDR').setVal(.5)
            #ws.var('fBkgFDR').setVal(.25)
            
            #ws.var('nPrompt').setVal(5000)
            #ws.var('nNonPrompt').setVal(500)
            #ws.var('nBackground').setVal(100)
            #ws.var('nBackgroundL').setVal(50)
            #ws.var('nBackgroundR').setVal(50)
            
            ws.var('nonPromptTau').setVal(.5)
            ws.var('promptMean').setVal(0)
            ws.var('ctResolution').setVal(1)

                        
            LPdf = ws.pdf('LPdf')
            MPdf = ws.pdf('MPdf')            
            
                
            data = ws.data('data_rap'+str(rap_bin)+'_pt'+str(pt_bin+1))

            NLLs = RooArgSet()

            MassNLL = MPdf.createNLL(data,
                                     ROOT.RooFit.Range('mlfit'),
                                     ROOT.RooFit.SplitRange(True),                                    
                                     ROOT.RooFit.ConditionalObservables(RooArgSet(ws.var('JpsictErr'))),
                                     ROOT.RooFit.NumCPU(2))
                                     

            CTauNLL = LPdf.createNLL(data,
                                     ROOT.RooFit.Range('mlfit'),
                                     ROOT.RooFit.SplitRange(True),                                    
                                     ROOT.RooFit.ConditionalObservables(RooArgSet(ws.var('JpsictErr'))),
                                     ROOT.RooFit.NumCPU(2))
                
            NLLs.add(MassNLL)
            NLLs.add(CTauNLL)

            simNLL = RooAddition('add','add',NLLs)

            minuit = RooMinuit(simNLL)
            minuit.setStrategy(2)
            minuit.setPrintEvalErrors(-1)

            minuit.simplex()
            minuit.migrad()
            minuit.migrad()
            minuit.hesse()
            
            fitresult = minuit.save('polfitresult_rap'+str(rap_bin)+'_pt'+str(pt_bin+1))
            getattr(ws,'import')(fitresult)           

            ws.saveSnapshot('snapshot_rap'+str(rap_bin)+'_pt'+str(pt_bin+1),ws.allVars())
phoERes.setRange(phoERes.getMin() - 0.1 * phoEResRange,
                 phoERes.getMax() + 0.1 * phoEResRange)
phoEResShape = w.factory('KeysPdf::phoEResShape(phoERes, phoEResData, NoMirror, 1.5)')
#phoERes.setRange(-10,10)
phoEResShape2 = w.factory('KeysNDPdf::phoEResShape2(phoERes, phoEResData, NoMirror, 2)')

w.Print()

## Find the mode of the photon eneregy resolution function.
## This is equal to the photon energy scale in MC.
phoEResShapeNegative = w.factory('''FormulaVar::phoEResShapeNegative::(
    "-phoEResShape", {phoEResShape}
    )''')

## Find the minimum of - phoEResShape = maximum of phoEResShape
minuit = RooMinuit(phoEResShapeNegative)
minuit.migrad()
phoEScaleMC = phoERes.getVal()

## Turn the photon resolution function into a mass smearing while introducing
## location and scale parameters.  Do this by turning it into a RooHistPdf.
## Start a hack to make sure that the photon resolution shape is 0
## outside of the training range.
phoEResHist = phoEResShape.createHistogram('phoERes', 10000)
bw = phoEResHist.GetBinWidth(1)
nbins = phoEResHist.GetNbinsX()
phoEResHistExtended = TH1F('phoEResHistExtended', 'phoEResHistExtended',
                           int(1.2 * nbins),
                           phoERes.getMin() - 0.1 * nbins * bw,
                           phoERes.getMax() + 0.1 * nbins * bw)
for b in range(1, phoEResHist.GetNbinsX() + 1):
Beispiel #3
0
def main(infiles=None):

    infile = infiles[0]

    var = "leadmupt"
    bounds = [25, 300]

    c1 = ROOT.TCanvas("NLL", "NLL", 1000, 1000)

    x = ROOT.RooRealVar(var, var, bounds[0], bounds[1])
    aset = ROOT.RooArgSet(x, "aset")

    frame = x.frame()

    f = ROOT.TFile.Open(infile)
    tree = f.Get(f.GetListOfKeys().At(0).GetName())

    tree.Print()
    nentries = tree.GetEntries()

    y = []

    dh2 = ROOT.TH1F()

    data = ROOT.RooDataSet("Data", "Data", aset)
    for n in range(nentries):
        tree.GetEntry(n)
        y.append(getattr(tree, var))
        if y[n] <= bounds[1] and y[n] >= bounds[0]:
            x.setVal(y[n])
            data.add(aset)
            dh2.Fill(y[n])

    data.plotOn(frame)

    dh = RooDataHist("dh", "dh", RooArgSet(x), data)

    nbins = dh2.GetNbinsX()
    nbinsy = dh2.GetNbinsX()
    print("nbins: ", nbins)
    print("nbinsy: ", nbinsy)
    for i in range(nbins):
        if dh2.GetBinContent(dh2.GetBin(i)) == 0:
            print("bin: ", i)
            #dh2.SetBinError(bin,0.01)

    ## CREATE GAUSSIAN MODEL
    mx = RooRealVar("mx", "mx", 10, 0, 350)
    sx = RooRealVar("sx", "sx", 3, 0, 10)
    gx = RooGaussian("gx", "gx", x, mx, sx)

    ## CREATE EXPONENTIAL MODEL
    lambda1 = RooRealVar("lambda1", "slope1", -100, 100)
    expo1 = RooExponential("expo1", "exponential PDF 1", x, lambda1)
    lambda2 = RooRealVar("lambda2", "slope2", -.03, -1000, 1000)
    expo2 = RooExponential("expo2", "exponential PDF 2", x, lambda2)

    l1 = RooRealVar("l1", "yield1", 100, 0, 10000)
    l2 = RooRealVar("l2", "yield2", 100, 0, 10000)

    #sum = RooAddPdf("sum","exp and gauss",RooArgList(expo1,gx),RooArgList(l1,l2))
    sum = RooAddPdf("sum", "2 exps", RooArgList(expo1, expo2),
                    RooArgList(l1, l2))

    ## Construct binned likelihood
    nll = RooNLLVar("nll", "nll", expo1, data, ROOT.RooFit.Extended(True))

    ## Start Minuit session on NLL
    m = RooMinuit(nll)
    m.migrad()
    m.hesse()
    r1 = m.save()

    #sum.plotOn(frame,ROOT.RooFit.LineColor(1))
    #sum.plotOn(frame,ROOT.RooFit.Components("expo1"),ROOT.RooFit.LineColor(2))
    #sum.plotOn(frame,ROOT.RooFit.Components("expo2"),ROOT.RooFit.LineColor(3))

    expo1.plotOn(frame)

    ## Construct Chi2
    chi2 = RooChi2Var("chi2", "chi2", expo2, dh)

    ## Start Minuit session on Chi2
    m2 = RooMinuit(chi2)
    m2.migrad()
    m2.hesse()
    r2 = m2.save()

    frame.Draw()

    c2 = ROOT.TCanvas("Chi2", "Chi2", 1000, 1000)
    frame2 = x.frame()

    data.plotOn(frame2)

    expo2.plotOn(frame2)
    #sum.plotOn(frame2,ROOT.RooFit.LineColor(4))
    #sum.plotOn(frame2,ROOT.RooFit.Components("expo1"),ROOT.RooFit.LineColor(5))
    #sum.plotOn(frame2,ROOT.RooFit.Components("expo2"),ROOT.RooFit.LineColor(6))

    ## Print results
    print("result of likelihood fit")
    r1.Print("v")
    print("result of chi2 fit")
    r2.Print("v")

    frame2.Draw()

    c1.Draw()
    c2.Draw()

    rep = ''
    while not rep in ['q', 'Q']:
        rep = input('enter "q" to quit: ')
        if 1 < len(rep):
            rep = rep[0]
Beispiel #4
0
def doFit(ws,options):
    rap_bins = range(1,len(jpsi.pTRange))
    pt_bins = None
    
    
    if options.testBin is not None:
        rap_bins = [int(options.testBin.split(',')[0])]
        pt_bins  = [int(options.testBin.split(',')[1])-1]
        
    for rap_bin in rap_bins:
        if options.testBin is None:
            pt_bins = range(len(jpsi.pTRange[rap_bin]))
        for pt_bin in pt_bins:

            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]

            jPsiMass = ws.var('JpsiMass')
            jPsicTau = ws.var('Jpsict')

            jPsiMass.setRange('mlfit_prompt',2.7,3.5)
            jPsiMass.setRange('mlfit_nonPrompt',2.7,3.5)            

            jPsiMass.setRange('NormalizationRangeFormlfit_prompt',2.7,3.5)
            jPsiMass.setRange('NormalizationRangeFormlfit_nonPrompt',2.7,3.5)

            jPsicTau.setRange('mlfit_signal',-1,2.5)
            jPsicTau.setRange('mlfit_leftMassSideBand',-1,2.5)
            jPsicTau.setRange('mlfit_rightMassSideBand',-1,2.5)

            jPsicTau.setRange('NormalizationRangeFormlfit_signal',-1,2.5)
            jPsicTau.setRange('NormalizationRangeFormlfit_leftMassSideBand',-1,2.5)
            jPsicTau.setRange('NormalizationRangeFormlfit_rightMassSideBand',-1,2.5)

            #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)
            
            #jPsicTau.setRange('mlfit_promptSignal',-1,.1)
            #jPsicTau.setRange('mlfit_nonPromptSignal',.1,2.5)
            #jPsicTau.setRange('mlfit_leftMassSideBand',-1,2.5)
            #jPsicTau.setRange('mlfit_rightMassSideBand',-1,2.5)

            #reset parameters
            ws.var('CBn').setVal(.5)
            ws.var('CBalpha').setVal(.5)
            ws.var('CBmass').setVal(3.1)
            ws.var('CBsigma').setVal(.02)
            ws.var('bkgLambda').setVal(0)
            
            ws.var('bkgTauSSDL').setVal(.5)
            #ws.var('bkgTauFDL').setVal(.5)
            ws.var('bkgTauDSDL').setVal(.5)
            ws.var('fBkgSSDL').setVal(.5)
            ws.var('fBkgLR').setVal(.5)

            ws.var('bkgTauSSDR').setVal(.5)
            #ws.var('bkgTauFDR').setVal(.5)
            ws.var('bkgTauDSDR').setVal(.5)
            ws.var('fBkgSSDR').setVal(.5)
            #ws.var('fBkgFDR').setVal(.25)
            
            #ws.var('nPrompt').setVal(5000)
            #ws.var('nNonPrompt').setVal(500)
            #ws.var('nBackground').setVal(100)
            #ws.var('nBackgroundL').setVal(50)
            #ws.var('nBackgroundR').setVal(50)
            
            ws.var('nonPromptTau').setVal(.5)
            ws.var('promptMean').setVal(0)
            ws.var('ctResolution').setVal(1)





            if options.fixBfrac:
                cutStringM1 = '('+jPsiMass.GetName()+' > '+str(sigMinMass)+' && '+jPsiMass.GetName()+' < '+str(sigMaxMass)+')'
                data2 = ws.data('data_rap'+str(rap_bin)+'_pt'+str(pt_bin+1))
                dataSizeBfrac1 = data2.numEntries()
                print dataSizeBfrac1
                bfracVars = RooArgSet(jPsiMass)
                bfracData = data2.reduce(ROOT.RooFit.SelectVars(bfracVars),
                                         ROOT.RooFit.Cut(cutStringM1),
                                         ROOT.RooFit.Name('data_for_normalizing_bfrac'),
                                         ROOT.RooFit.Title('data_for_normalizing_bfrac'))
                print rap_bin
                print pt_bin
                bfrac = jpsi.Bfrac[rap_bin-1][pt_bin]
                print bfrac
                dataSizeBfrac = bfracData.numEntries()
                factPrompt = 1-bfrac
                dataSizePrompt = dataSizeBfrac*factPrompt
                dataSizeNonPrompt = dataSizeBfrac*bfrac          
                
                print factPrompt
                print dataSizeBfrac
                print dataSizePrompt
                print dataSizeNonPrompt
                
                ws.var('nPromptSignal').setVal(dataSizePrompt)
                ws.var('nNonPromptSignal').setVal(dataSizeNonPrompt)
                ws.var('nPromptSignal').setConstant()
                ws.var('nNonPromptSignal').setConstant()  
        
        
        
        
            
            LPdf = ws.pdf('LPdf')
            MPdf = ws.pdf('MPdf')            
            
                
            data = ws.data('data_rap'+str(rap_bin)+'_pt'+str(pt_bin+1))

            NLLs = RooArgSet()

            MassNLL = MPdf.createNLL(data,
                                     ROOT.RooFit.Range('mlfit'),
                                     ROOT.RooFit.SplitRange(True),                                    
                                     ROOT.RooFit.ConditionalObservables(RooArgSet(ws.var('JpsictErr'))),
                                     ROOT.RooFit.NumCPU(2))
                                     

            CTauNLL = LPdf.createNLL(data,
                                     ROOT.RooFit.Range('mlfit'),
                                     ROOT.RooFit.SplitRange(True),                                    
                                     ROOT.RooFit.ConditionalObservables(RooArgSet(ws.var('JpsictErr'))),
                                     ROOT.RooFit.NumCPU(2))
                
            NLLs.add(MassNLL)
            NLLs.add(CTauNLL)

            simNLL = RooAddition('add','add',NLLs)

            minuit = RooMinuit(simNLL)
            minuit.setStrategy(2)
            minuit.setPrintEvalErrors(-1)

            if options.do_fit:
                minuit.simplex()
                minuit.migrad()
                minuit.migrad()
                minuit.hesse()
            
            fitresult = minuit.save('polfitresult_rap'+str(rap_bin)+'_pt'+str(pt_bin+1))
            getattr(ws,'import')(fitresult)           

            ws.saveSnapshot('snapshot_rap'+str(rap_bin)+'_pt'+str(pt_bin+1),ws.allVars())
            
            fitresult.Print()
phoEResShape = w.factory("KeysPdf::phoEResShape(phoERes, phoEResData, NoMirror, 1.5)")
# phoERes.setRange(-10,10)
phoEResShape2 = w.factory("KeysNDPdf::phoEResShape2(phoERes, phoEResData, NoMirror, 2)")

w.Print()

## Find the mode of the photon eneregy resolution function.
## This is equal to the photon energy scale in MC.
phoEResShapeNegative = w.factory(
    """FormulaVar::phoEResShapeNegative::(
    "-phoEResShape", {phoEResShape}
    )"""
)

## Find the minimum of - phoEResShape = maximum of phoEResShape
minuit = RooMinuit(phoEResShapeNegative)
minuit.migrad()
phoEScaleMC = phoERes.getVal()

## Turn the photon resolution function into a mass smearing while introducing
## location and scale parameters.  Do this by turning it into a RooHistPdf.
## Start a hack to make sure that the photon resolution shape is 0
## outside of the training range.
phoEResHist = phoEResShape.createHistogram("phoERes", 10000)
bw = phoEResHist.GetBinWidth(1)
nbins = phoEResHist.GetNbinsX()
phoEResHistExtended = TH1F(
    "phoEResHistExtended",
    "phoEResHistExtended",
    int(1.2 * nbins),
    phoERes.getMin() - 0.1 * nbins * bw,
Beispiel #6
0
def rooFit601():

    print ">>> setup pdf and likelihood..."
    x = RooRealVar("x", "x", -20, 20)
    mean = RooRealVar("mean", "mean of g1 and g2", 0)
    sigma1 = RooRealVar("sigma1", "width of g1", 3)
    sigma2 = RooRealVar("sigma2", "width of g2", 4, 3.0,
                        6.0)  # intentional strong correlations
    gauss1 = RooGaussian("gauss1", "gauss1", x, mean, sigma1)
    gauss2 = RooGaussian("gauss2", "gauss2", x, mean, sigma2)
    frac = RooRealVar("frac", "frac", 0.5, 0.0, 1.0)
    model = RooAddPdf("model", "model", RooArgList(gauss1, gauss2),
                      RooArgList(frac))

    print ">>> generate to data..."
    data = model.generate(RooArgSet(x), 1000)  # RooDataSet

    print ">>> construct unbinned likelihood of model wrt data..."
    nll = model.createNLL(data)  # RooAbsReal

    print ">>> interactive minimization and error analysis with MINUIT interface object..."
    minuit = RooMinuit(nll)

    print ">>> set avtive verbosity for logging of MINUIT parameter space stepping..."
    minuit.setVerbose(kTRUE)

    print ">>> call MIGRAD to minimize the likelihood..."
    minuit.migrad()

    print "\n>>> parameter values and error estimates that are back propagated from MINUIT:"
    model.getParameters(RooArgSet(x)).Print("s")

    print "\n>>> disable verbose logging..."
    minuit.setVerbose(kFALSE)

    print ">>> run HESSE to calculate errors from d2L/dp2..."
    minuit.hesse()

    print ">>> value of and error on sigma2 (back propagated from MINUIT):"
    sigma2.Print()

    print "\n>>> run MINOS on sigma2 parameter only..."
    minuit.minos(RooArgSet(sigma2))

    print "\n>>> value of and error on sigma2 (back propagated from MINUIT after running MINOS):"
    sigma2.Print()

    print "\n>>> saving results, contour plots..."
    # Save a snapshot of the fit result. This object contains the initial
    # fit parameters, the final fit parameters, the complete correlation
    # matrix, the EDM, the minimized FCN , the last MINUIT status code and
    # the number of times the RooFit function object has indicated evaluation
    # problems (e.g. zero probabilities during likelihood evaluation)
    result = minuit.save()  # RooFitResult

    # Make contour plot of mx vs sx at 1,2,3 sigma
    frame1 = minuit.contour(frac, sigma2, 1, 2, 3)  # RooPlot
    frame1.SetTitle("RooMinuit contour plot")

    # Print the fit result snapshot
    result.Print("v")

    print "\n>>> change value of \"mean\" parameter..."
    mean.setVal(0.3)

    # Rerun MIGRAD,HESSE
    print ">>> rerun MIGRAD, HESSE..."
    minuit.migrad()
    minuit.hesse()

    print ">>> value on and error of frac:"
    frac.Print()

    print "\n>>> fix value of \"sigma\" parameter (setConstant)..."
    sigma2.setConstant(kTRUE)

    print ">>> rerun MIGRAD, HESSE..."
    minuit.migrad()
    minuit.hesse()
    frac.Print()
Beispiel #7
0
def doFit(ws, options):
    rap_bins = range(1, len(jpsi.pTRange))
    pt_bins = None

    if options.testBin is not None:
        rap_bins = [int(options.testBin.split(',')[0])]
        pt_bins = [int(options.testBin.split(',')[1]) - 1]

    for rap_bin in rap_bins:
        if options.testBin is None:
            pt_bins = range(len(jpsi.pTRange[rap_bin]))
        for pt_bin in pt_bins:

            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]

            jPsiMass = ws.var('JpsiMass')
            jPsicTau = ws.var('Jpsict')

            jPsiMass.setRange('mlfit_prompt', 2.7, 3.5)
            jPsiMass.setRange('mlfit_nonPrompt', 2.7, 3.5)

            jPsiMass.setRange('NormalizationRangeFormlfit_prompt', 2.7, 3.5)
            jPsiMass.setRange('NormalizationRangeFormlfit_nonPrompt', 2.7, 3.5)

            jPsicTau.setRange('mlfit_signal', -1, 2.5)
            jPsicTau.setRange('mlfit_leftMassSideBand', -1, 2.5)
            jPsicTau.setRange('mlfit_rightMassSideBand', -1, 2.5)

            jPsicTau.setRange('NormalizationRangeFormlfit_signal', -1, 2.5)
            jPsicTau.setRange('NormalizationRangeFormlfit_leftMassSideBand',
                              -1, 2.5)
            jPsicTau.setRange('NormalizationRangeFormlfit_rightMassSideBand',
                              -1, 2.5)

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

            #jPsicTau.setRange('mlfit_promptSignal',-1,.1)
            #jPsicTau.setRange('mlfit_nonPromptSignal',.1,2.5)
            #jPsicTau.setRange('mlfit_leftMassSideBand',-1,2.5)
            #jPsicTau.setRange('mlfit_rightMassSideBand',-1,2.5)

            #reset parameters
            ws.var('CBn').setVal(.5)
            ws.var('CBalpha').setVal(.5)
            ws.var('CBmass').setVal(3.1)
            ws.var('CBsigma').setVal(.02)
            ws.var('bkgLambda').setVal(0)

            ws.var('bkgTauSSDL').setVal(.5)
            #ws.var('bkgTauFDL').setVal(.5)
            ws.var('bkgTauDSDL').setVal(.5)
            ws.var('fBkgSSDL').setVal(.5)
            ws.var('fBkgLR').setVal(.5)

            ws.var('bkgTauSSDR').setVal(.5)
            #ws.var('bkgTauFDR').setVal(.5)
            ws.var('bkgTauDSDR').setVal(.5)
            ws.var('fBkgSSDR').setVal(.5)
            #ws.var('fBkgFDR').setVal(.25)

            #ws.var('nPrompt').setVal(5000)
            #ws.var('nNonPrompt').setVal(500)
            #ws.var('nBackground').setVal(100)
            #ws.var('nBackgroundL').setVal(50)
            #ws.var('nBackgroundR').setVal(50)

            ws.var('nonPromptTau').setVal(.5)
            ws.var('promptMean').setVal(0)
            ws.var('ctResolution').setVal(1)

            LPdf = ws.pdf('LPdf')
            MPdf = ws.pdf('MPdf')

            data = ws.data('data_rap' + str(rap_bin) + '_pt' + str(pt_bin + 1))

            NLLs = RooArgSet()

            MassNLL = MPdf.createNLL(
                data, ROOT.RooFit.Range('mlfit'), ROOT.RooFit.SplitRange(True),
                ROOT.RooFit.ConditionalObservables(
                    RooArgSet(ws.var('JpsictErr'))), ROOT.RooFit.NumCPU(2))

            CTauNLL = LPdf.createNLL(
                data, ROOT.RooFit.Range('mlfit'), ROOT.RooFit.SplitRange(True),
                ROOT.RooFit.ConditionalObservables(
                    RooArgSet(ws.var('JpsictErr'))), ROOT.RooFit.NumCPU(2))

            NLLs.add(MassNLL)
            NLLs.add(CTauNLL)

            simNLL = RooAddition('add', 'add', NLLs)

            minuit = RooMinuit(simNLL)
            minuit.setStrategy(2)
            minuit.setPrintEvalErrors(-1)

            minuit.simplex()
            minuit.migrad()
            minuit.migrad()
            minuit.hesse()

            fitresult = minuit.save('polfitresult_rap' + str(rap_bin) + '_pt' +
                                    str(pt_bin + 1))
            getattr(ws, 'import')(fitresult)

            ws.saveSnapshot(
                'snapshot_rap' + str(rap_bin) + '_pt' + str(pt_bin + 1),
                ws.allVars())
def main(options,args):

    eta_bins = [[0.0,1.4442]] # try to split material dependence
#                [1.560,2.5]]

    pt_bins = [[12, 15],
               [15, 20]]

    #get input workspace
    input = TFile.Open(options.inputFile)    
    ws = input.Get(options.workspaceName).Clone()
    input.Close()

    #ws.Print("v")

    sieie = RooRealVar(options.showerShapeName,options.showerShapeName,0)
    pt = RooRealVar(options.ptName,options.ptName,0)
    eta = RooRealVar(options.etaName,options.etaName,0)    
    chIso = RooRealVar(options.isoName,options.isoName,0)

    vars = RooArgSet()
    vars.add(sieie)
    vars.add(pt)
    vars.add(eta)
    vars.add(chIso)    

    if options.MCFakeRate:
        #chIso = RooRealVar(options.isoName, 'phHIso', 100, 0, 20)
        fake = RooRealVar('FakePho','',-0.5,1.5)
        vars.add(fake)

    data = None

    if options.weight is not None:
        weight = RooRealVar(options.weight,'',.5)
        vars.add(weight)
        data = RooDataSet('data','All Input Data',
                          vars,
                          RooFit.ImportFromFile(options.dataFile,
                                                options.dataTreeName),
                          RooFit.WeightVar(options.weight)
                          )
    else:
        data = RooDataSet('data','All Input Data',
                          vars,
                          RooFit.ImportFromFile(options.dataFile,
                                                options.dataTreeName)
                          )

    data.Print("v")
    
    # put in the raw datasets, no cuts or manipulation
    getattr(ws,'import')(data,
                         RooFit.RenameVariable(options.showerShapeName,'Pho_SigmaIEtaIEta'),
                         RooFit.RenameVariable(options.ptName,'Pho_Pt'),
                         RooFit.RenameVariable(options.etaName,'Pho_Eta'))

    sieie = ws.var('Pho_SigmaIEtaIEta')
    sieie.setRange('barrelSel',0,0.011)
    sieie.setRange('endcapSel',0,0.033)

    output = TFile.Open(options.outputFile,'RECREATE')
    output.cd()

    #histograms
    fakeRate   = []
    nFakes     = []
    nTot       = []
    fakeRateMC = []
    #counters
    totBkg     = [] #central value in each bin
    totBkgUp   = [] #upper error
    totBkgLo   = [] #lower error
    ietabin = 0

    #setup histograms and counters
    for i in range(len(eta_bins)):
        fakeRate.append(TGraphAsymmErrors())
        fakeRate[i].SetName('fakeRate_etabin%d'%i)

        nFakes.append(TGraphAsymmErrors())
        nFakes[i].SetName('nFakes_etabin%d'%i)

        nTot.append(TGraphAsymmErrors())
        
        nTot[i].SetName('nTot_etabin%d'%i)

        if options.MCFakeRate:
            fakeRateMC.append(TGraphAsymmErrors())
            fakeRateMC[i].SetName('fakeRateMC_etabin%d'%i)

        totBkg.append(0)
        totBkgUp.append(0)
        totBkgLo.append(0)

    #loop through bins making all fits
    for etabin in eta_bins:
        #pt bin number for TGraphs
        iptbin = 0
        
        for ptbin in pt_bins:
            
            if 'abs(Pho_Eta) < 1.4442':
                phoselsig = ' abs(Pho_Eta) > %f  && abs(Pho_Eta) < %f && Pho_Pt > %f && Pho_Pt < %f'%(etabin[0],
                                                                                                     etabin[1],
                                                                                                     ptbin[0],
                                                                                                      ptbin[1])
       #         if options.MCFakeRate:
                phoselbkg = 'phoCHIso > 2 && phoCHIso < 6 && abs(Pho_Eta) > %f  && abs(Pho_Eta) < %f && Pho_Pt > %f && Pho_Pt < %f'%(etabin[0],
                                                                                                                                      etabin[1],
                                                                                                                                      ptbin[0],
                                                                                                                                      ptbin[1])
                
            phosel= phoselsig or phoselbkg
            #phosel = 'phoselsig || phoselbkg'
            #phosel = 'abs(Pho_Eta) > %f  && abs(Pho_Eta) < %f && Pho_Pt > %f && Pho_Pt < %f'%(etabin[0],
            #    etabin[1],
            #   ptbin[0],
            #  ptbin[1])
            postfix = '_Eta_%.4f_%.4f_Pt_%.2f_%.2f'%(etabin[0],etabin[1],ptbin[0],ptbin[1])


            binData = ws.data('data').reduce(RooFit.Cut(phosel),
                                             RooFit.Name('binData'+postfix),
                                             RooFit.Title('Data Events'))        
                        
            #save it all in the workspace
            getattr(ws,'import')(binData)            
            
            #create the PDFs we are going to fit!
            options.pdfModule.makePdf(ws,etabin,postfix,options.extraInput)
            output.cd()
            
            ws.factory('RooExtendPdf::fullPDFExt'+postfix+'('+
                       'fullPDF'+postfix+',nTot'+postfix+'[2000,1e-8,10000])')

            
            bNLL = ws.pdf('fullPDFExt'+postfix).createNLL(ws.data('binData'+postfix),                                                          
                                                          RooFit.Extended(True), RooFit.NumCPU(4))
                        
            bFit = RooMinuit(bNLL)
        

            #bFit.setVerbose(True)
            if ws.data('binData'+postfix).sumEntries() > 0:
                bFit.setPrintLevel(-1)
                bFit.migrad()
                bFit.hesse()
                bFit.setPrintLevel(0)
                bFit.minos(RooArgSet(ws.var('fBkg'+postfix)))
                
                bRes = bFit.save('binFitResult'+postfix)
                getattr(ws,'import')(bRes)
                
                #plot fits on top of data
                bFrame = sieie.frame()
                bFrame.SetName('binFrame'+postfix)
                
                binData.plotOn(bFrame)
                ws.pdf('fullPDFExt'+postfix).plotOn(bFrame, RooFit.LineColor(ROOT.kYellow))
                #ws.pdf('fullPDFExt'+postfix).plotOn(bFramea)
                ws.pdf('fullPDFExt'+postfix).plotOn(bFrame,
                                                    RooFit.Components('bkgTemplatePDF'+postfix),
                                                    RooFit.LineColor(ROOT.kBlue))        
                ws.pdf('fullPDFExt'+postfix).plotOn(bFrame,
                                                    RooFit.Components('sigTemplatePDF'+postfix),
                                                    RooFit.LineColor(ROOT.kRed))        

                ws.pdf('fullPDFExt'+postfix).plotOn(bFrame,
                                                    RooFit.Components('fullPDF'+postfix),
                                                    RooFit.LineColor(ROOT.kGreen))        
                bFrame.Write()

                bSigInt = None
                bBkgInt = None
                
                #make plot of fake rate in signal region as a function of pT
                #< 0.011 (EB), <0.030 (EE)
                if etabin[0] >= 1.560:
                    bSigInt = ws.pdf('sigTemplatePDF'+
                                     postfix).createIntegral(RooArgSet(sieie),
                                                             RooArgSet(sieie),
                                                             'endcapSel')
                    bBkgInt = ws.pdf('bkgTemplatePDF'+
                                     postfix).createIntegral(RooArgSet(sieie),
                                                             RooArgSet(sieie),
                                                             'endcapSel')
                else:
                    bSigInt = ws.pdf('sigTemplatePDF'+
                                     postfix).createIntegral(RooArgSet(sieie),
                                                             RooArgSet(sieie),
                                                             'barrelSel')
                    bBkgInt = ws.pdf('bkgTemplatePDF'+
                                     postfix).createIntegral(RooArgSet(sieie),
                                                             RooArgSet(sieie),
                                                             'barrelSel')
                
                nEventsB = ws.var('nTot'+postfix).getVal()                
                nSignalB = (1.0 - ws.var('fBkg'+postfix).getVal())*nEventsB*bSigInt.getVal()
                nBkgB    = ws.var('fBkg'+postfix).getVal()*nEventsB*bBkgInt.getVal()
                nBkgBUp  = ws.var('fBkg'+postfix).getErrorHi()*nEventsB*bBkgInt.getVal()
                nBkgBLo  = -ws.var('fBkg'+postfix).getErrorLo()*nEventsB*bBkgInt.getVal()
                print "Events Sig Bkg: %.3f + %.3f - %.3f"%(nEventsB, nSignalB, nBkgB)

                trash = (1.0 - ws.var('fBkg'+postfix).getVal())
                trash1 = bSigInt.getVal()
                trash2 = bBkgInt.getVal()
                trash3 = ws.var('fBkg'+postfix).getVal()
                print nEventsB
                print trash
                print trash1
                print trash2
                print trash3

                if nBkgBLo == 0.0: #catch cases when minos dies on lower limits
                    parb_err = ws.var('fBkg'+postfix).getError()*nEventsB*bBkgInt.getVal()
                    nBkgBLo = -max(nBkgB - parb_err,-nBkgB)
                    
                totBkg[ietabin]   += nBkgB
                totBkgUp[ietabin] = sqrt(nBkgBUp*nBkgBUp + totBkgUp[ietabin]*totBkgUp[ietabin])
                totBkgLo[ietabin] = sqrt(nBkgBLo*nBkgBLo + totBkgLo[ietabin]*totBkgLo[ietabin])
                
                print "Background Events: %.3f + %.3f - %.3f"%(nBkgB,nBkgBUp,nBkgBLo)
                fakeRate[ietabin].SetPoint(iptbin,(float(ptbin[1])+float(ptbin[0]))/2,nBkgB/(nSignalB+nBkgB))
                fakeRate[ietabin].SetPointError(iptbin,
                                                (float(ptbin[1])+float(ptbin[0]))/2 - float(ptbin[0]),
                                                float(ptbin[1]) - (float(ptbin[1])+float(ptbin[0]))/2,
                                                -ws.var('fBkg'+postfix).getErrorLo(),
                                                ws.var('fBkg'+postfix).getErrorHi())
                
                #nFakes[ietabin].SetPoint(iptbin,(float(ptbin[1])+float(ptbin[0]))/2,nBkgB/(float(ptbin[1]) - float(ptbin[0])))                
                nFakes[ietabin].SetPoint(iptbin,(float(ptbin[1])+float(ptbin[0]))/2,nBkgB)                
                nFakes[ietabin].SetPointError(iptbin,
                                              (float(ptbin[1])+float(ptbin[0]))/2 - float(ptbin[0]),
                                              float(ptbin[1]) - (float(ptbin[1])+float(ptbin[0]))/2,
                                              nBkgBLo/(float(ptbin[1]) - float(ptbin[0])),
                                              nBkgBUp/(float(ptbin[1]) - float(ptbin[0])))
                
                #nTot[ietabin].SetPoint(iptbin,(float(ptbin[1])+float(ptbin[0]))/2,nEventsB/(float(ptbin[1]) - float(ptbin[0])))
                nTot[ietabin].SetPoint(iptbin,(float(ptbin[1])+float(ptbin[0]))/2,nEventsB)

                nTot[ietabin].SetPointError(iptbin,
                                            (float(ptbin[1])+float(ptbin[0]))/2 - float(ptbin[0]),
                                            float(ptbin[1]) - (float(ptbin[1])+float(ptbin[0]))/2,
                                            ws.var('nTot'+postfix).getError()/(float(ptbin[1]) - float(ptbin[0])),
                                            ws.var('nTot'+postfix).getError()/(float(ptbin[1]) - float(ptbin[0])))
            
                if options.MCFakeRate:
                    fr = 0
                    if etabin[0] > 1.560:
                        fr = (ws.data('binData'+postfix).sumEntries('FakePho > 0.5 && Pho_SigmaIEtaIEta < 0.030')/
                              ws.data('binData'+postfix).sumEntries('FakePho > -1 && Pho_SigmaIEtaIEta < 0.030'))
                    else:
                        fr = (ws.data('binData'+postfix).sumEntries('FakePho > 0.5 && Pho_SigmaIEtaIEta < 0.011')/
                              ws.data('binData'+postfix).sumEntries('FakePho > -1 && Pho_SigmaIEtaIEta < 0.011'))
                    
                    fakeRateMC[ietabin].SetPoint(iptbin,(float(ptbin[1])+float(ptbin[0]))/2,fr)            
                    fakeRateMC[ietabin].SetPointError(iptbin,
                                                      (float(ptbin[1])+float(ptbin[0]))/2 - float(ptbin[0]),
                                                      float(ptbin[1]) - (float(ptbin[1])+float(ptbin[0]))/2,
                                                      0,
                                                      0)       
        
            iptbin += 1
        ietabin += 1
        
    for i in range(len(eta_bins)):
        print " %.4f < |Eta| < %.4f  Total Background = %.3f +/- (%.3f,%.3f)"%(eta_bins[i][0],eta_bins[i][1],
                                                                               totBkg[i],totBkgUp[i],totBkgLo[i])
        
        fakeRate[i].Write()
        nFakes[i].Write()
        nTot[i].Write()
        
        if options.MCFakeRate:
            fakeRateMC[i].Write()
    
    ws.Write()
    output.Close()