Beispiel #1
0
def buildDataAndCategories(ws,options,args):
    #Get the input data
    inputData = TChain(options.treeName,'The input data')
    for arg in args:
        print 'Adding data from: ',arg
        inputData.Add(arg)

    foldname = ''
    phirange = [0,90]
    
    if not options.folded:
        foldname=''
        phirange = [-180,180]
    
    #variables necessary for j/psi mass,lifetime,polarization fit
    jPsiMass      = RooRealVar('JpsiMass','M [GeV]',2.7,3.5)
    jPsiRap       = RooRealVar('JpsiRap','#nu',-2.3,2.3)
    jPsiPt        = RooRealVar("JpsiPt","pT [GeV]",0,40);
    jPsicTau      = RooRealVar('Jpsict','l_{J/#psi} [mm]',-1,2.5)
    jPsicTauError = RooRealVar('JpsictErr','Error on l_{J/#psi} [mm]',0,2)
    jPsiVprob     = RooRealVar('JpsiVprob','',.01,1)
    jPsiHXcosth   = None
    jPsiHXphi     = None

    jPsicTau.setBins(10000,"cache")
        
    if options.fitFrame is not None:
        jPsiHXcosth   = RooRealVar('costh_'+options.fitFrame+foldname,'cos(#theta)_{'+options.fitFrame+'}',-1,1)
        jPsiHXphi     = RooRealVar('phi_'+options.fitFrame+foldname,'#phi_{'+options.fitFrame+'}',phirange[0],phirange[1])
    else:
        jPsiHXcosth   = RooRealVar('costh_CS'+foldname,'cos(#theta)_{CS}',-1,1)
        jPsiHXphi     = RooRealVar('phi_CS'+foldname,'#phi_{CS}',phirange[0],phirange[1])
    
    #vars needed for on the fly calc of polarization variables
    jPsimuPosPx = RooRealVar('muPosPx','+ Muon P_{x} [GeV]',0)
    jPsimuPosPy = RooRealVar('muPosPy','+ Muon P_{y} [GeV]',0)
    jPsimuPosPz = RooRealVar('muPosPz','+ Muon P_{z} [GeV]',0)
    jPsimuNegPx = RooRealVar('muNegPx','- Muon P_{x} [GeV]',0)
    jPsimuNegPy = RooRealVar('muNegPy','- Muon P_{y} [GeV]',0)
    jPsimuNegPz = RooRealVar('muNegPz','- Muon P_{z} [GeV]',0)

    #create RooArgSet for eventual dataset creation
    dataVars = RooArgSet(jPsiMass,jPsiRap,jPsiPt,
                         jPsicTau,jPsicTauError,
                         jPsimuPosPx,jPsimuPosPy,jPsimuPosPz)
    
    #add trigger requirement if specified
    if options.triggerName:
        trigger = RooRealVar(options.triggerName,'Passes Trigger',0.5,1.5)
        dataVars.add(trigger)

    dataVars.add(jPsiVprob)
    dataVars.add(jPsimuNegPx)
    dataVars.add(jPsimuNegPy)
    dataVars.add(jPsimuNegPz)
    dataVars.add(jPsiHXcosth)
    dataVars.add(jPsiHXphi)
    
    
    redVars = RooArgSet(jPsiMass,jPsiRap,jPsiPt,
                        jPsicTau,jPsicTauError)
    redVars.add(jPsiHXcosth)
    redVars.add(jPsiHXphi)
    fitVars = redVars.Clone()    

    ### HERE IS WHERE THE BIT FOR CALCULATING POLARIZATION VARS GOES

    ctauStates = RooCategory('ctauRegion','Cut Region in lifetime')
    ctauStates.defineType('prompt',0)
    ctauStates.defineType('nonPrompt',1)

    massStates = RooCategory('massRegion','Cut Region in mass')
    massStates.defineType('signal',1)
    massStates.defineType('separation',0)
    massStates.defineType('leftMassSideBand',-2)
    massStates.defineType('rightMassSideBand',-1)

    states = RooCategory('mlRegion','Cut Region in mass')
    states.defineType('nonPromptSignal',2)
    states.defineType('promptSignal',1)
    states.defineType('separation',0)
    states.defineType('leftMassSideBand',-2)
    states.defineType('rightMassSideBand',-1)


    #define corresponding ranges in roorealvars
    #mass is a little tricky since the sidebands change definitions in each rap bin
    #define the names here and change as we do the fits
    #jPsiMass.setRange('NormalizationRangeFormlfit_promptSignal',2.7,3.5)
    #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 RooCategory, RooRealVar
from B2DXFitters.taggingutils import tagEfficiencyWeight


mixState = RooCategory( 'mixstate', 'B/Bbar -> D pi mixing state' )
mixState.defineType( "unmixed" ,  1 )
mixState.defineType( "mixed"   , -1 )
mixState.defineType( "untagged",  0 )

SigTagEff =  0.25
sigTagEff = RooRealVar( "sigTagEff", "Signal tagging efficiency",
                        SigTagEff, 0., 1. )

sigTagWeight = tagEfficiencyWeight( mixState, sigTagEff, 'Bd2DPi' )

sigTagWeight.Print( 'v' )

print '\ntafEff =', sigTagEff.getVal()
mixState.setIndex(0)
print 'mixState = %2d  =>  tagWeight = %f' % \
      ( mixState.getIndex(), sigTagWeight.getVal() )
mixState.setIndex(1)
print 'mixState = %2d  =>  tagWeight = %f' % \
      ( mixState.getIndex(), sigTagWeight.getVal() )
mixState.setIndex(-1)
print 'mixState = %2d  =>  tagWeight = %f' % \
      ( mixState.getIndex(), sigTagWeight.getVal() )
Beispiel #3
0
def buildDataAndCategories(ws, options, args):
    #Get the input data
    inputData = TChain(options.treeName, 'The input data')
    for arg in args:
        print 'Adding data from: ', arg
        inputData.Add(arg)

    foldname = ''
    phirange = [0, 90]

    if not options.folded:
        foldname = ''
        phirange = [-180, 180]

    #variables necessary for j/psi mass,lifetime,polarization fit
    jPsiMass = RooRealVar('JpsiMass', 'M [GeV]', 2.7, 3.5)
    jPsiRap = RooRealVar('JpsiRap', '#nu', -2.3, 2.3)
    jPsiPt = RooRealVar("JpsiPt", "pT [GeV]", 0, 40)
    jPsicTau = RooRealVar('Jpsict', 'l_{J/#psi} [mm]', -1, 2.5)
    jPsicTauError = RooRealVar('JpsictErr', 'Error on l_{J/#psi} [mm]', 0, 2)
    jPsiVprob = RooRealVar('JpsiVprob', '', .01, 1)
    jPsiHXcosth = None
    jPsiHXphi = None

    jPsicTau.setBins(10000, "cache")

    if options.fitFrame is not None:
        jPsiHXcosth = RooRealVar('costh_' + options.fitFrame + foldname,
                                 'cos(#theta)_{' + options.fitFrame + '}', -1,
                                 1)
        jPsiHXphi = RooRealVar('phi_' + options.fitFrame + foldname,
                               '#phi_{' + options.fitFrame + '}', phirange[0],
                               phirange[1])
    else:
        jPsiHXcosth = RooRealVar('costh_CS' + foldname, 'cos(#theta)_{CS}', -1,
                                 1)
        jPsiHXphi = RooRealVar('phi_CS' + foldname, '#phi_{CS}', phirange[0],
                               phirange[1])

    #vars needed for on the fly calc of polarization variables
    jPsimuPosPx = RooRealVar('muPosPx', '+ Muon P_{x} [GeV]', 0)
    jPsimuPosPy = RooRealVar('muPosPy', '+ Muon P_{y} [GeV]', 0)
    jPsimuPosPz = RooRealVar('muPosPz', '+ Muon P_{z} [GeV]', 0)
    jPsimuNegPx = RooRealVar('muNegPx', '- Muon P_{x} [GeV]', 0)
    jPsimuNegPy = RooRealVar('muNegPy', '- Muon P_{y} [GeV]', 0)
    jPsimuNegPz = RooRealVar('muNegPz', '- Muon P_{z} [GeV]', 0)

    #create RooArgSet for eventual dataset creation
    dataVars = RooArgSet(jPsiMass, jPsiRap, jPsiPt, jPsicTau, jPsicTauError,
                         jPsimuPosPx, jPsimuPosPy, jPsimuPosPz)

    #add trigger requirement if specified
    if options.triggerName:
        trigger = RooRealVar(options.triggerName, 'Passes Trigger', 0.5, 1.5)
        dataVars.add(trigger)

    dataVars.add(jPsiVprob)
    dataVars.add(jPsimuNegPx)
    dataVars.add(jPsimuNegPy)
    dataVars.add(jPsimuNegPz)
    dataVars.add(jPsiHXcosth)
    dataVars.add(jPsiHXphi)

    redVars = RooArgSet(jPsiMass, jPsiRap, jPsiPt, jPsicTau, jPsicTauError)
    redVars.add(jPsiHXcosth)
    redVars.add(jPsiHXphi)
    fitVars = redVars.Clone()

    ### HERE IS WHERE THE BIT FOR CALCULATING POLARIZATION VARS GOES

    ctauStates = RooCategory('ctauRegion', 'Cut Region in lifetime')
    ctauStates.defineType('prompt', 0)
    ctauStates.defineType('nonPrompt', 1)

    massStates = RooCategory('massRegion', 'Cut Region in mass')
    massStates.defineType('signal', 1)
    massStates.defineType('separation', 0)
    massStates.defineType('leftMassSideBand', -2)
    massStates.defineType('rightMassSideBand', -1)

    states = RooCategory('mlRegion', 'Cut Region in mass')
    states.defineType('nonPromptSignal', 2)
    states.defineType('promptSignal', 1)
    states.defineType('separation', 0)
    states.defineType('leftMassSideBand', -2)
    states.defineType('rightMassSideBand', -1)

    #define corresponding ranges in roorealvars
    #mass is a little tricky since the sidebands change definitions in each rap bin
    #define the names here and change as we do the fits
    #jPsiMass.setRange('NormalizationRangeFormlfit_promptSignal',2.7,3.5)
    #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)
Beispiel #4
0
class SWeightTransform(object):
    def __init__(self, pdf, component, fit_opts, correct_weights=True):
        self.__comp = component
        self.__pdf = pdf
        self.__fit_opts = fit_opts
        self.__result = None
        self.__correct_weights = correct_weights

        from ROOT import RooCategory

        self.__status = RooCategory("sweight_status", "sweight fit status")
        self.__status.defineType("success", 0)
        self.__status.defineType("one", 1)
        self.__status.defineType("two", 2)
        self.__status.defineType("three", 3)
        self.__status.defineType("other", 4)
        self.__parameters = None

    def __call__(self, data):
        pdf_pars = self.__pdf.getParameters(data.get())
        if not self.__parameters:
            self.__parameters = pdf_pars
            self.__parameters = [p for p in pdf_pars] + [self.__status]
        else:
            for p in self.__parameters:
                pdf_par = pdf_pars.find(p.GetName())
                if not pdf_par:
                    continue
                pdf_par.setVal(p.getVal())
                pdf_par.setError(p.getError())

        success = False
        status = 4
        for i in range(3):
            self.__result = self.__pdf.fitTo(data, **self.__fit_opts)
            status = self.__result.status()
            if status == 0:
                success = True
                break

        if status < 4:
            self.__status.setIndex(self.__result.status())
        else:
            self.__status.setIndex(4)
        if success:
            from P2VV.Utilities.SWeights import SData

            self.__sData = SData(Pdf=self.__pdf, Data=data, Name="MassSPlot")
            data = self.__sData.data(self.__comp)
            if self.__correct_weights:
                from P2VV.Utilities.DataHandling import correctWeights

                data = correctWeights(data, splitCatNames=None, ImportIntoWS=False)
            return data
        else:
            return None

    def gen_params(self, observables=None):
        from ROOT import RooArgSet

        if self.__parameters:
            return self.__parameters
        else:
            if observables and not isinstance(observables, RooArgSet):
                obs = RooArgSet()
                for o in observables:
                    obs.add(o._target_() if hasattr(o, "_target_") else o)
                observables = obs
            params = self.__pdf.getParameters(observables)
            self.__parameters = [p for p in params] + [self.__status]
            return self.__parameters

    def result_params(self):
        if not self.__result:
            return []
        else:
            return [p for p in self.__result.floatParsFinal()] + [self.__status]

    def set_params(self, data_params):
        from ROOT import RooCategory

        for trans_param in self.result_params():
            data_param = data_params.find(trans_param.GetName())
            if isinstance(trans_param, RooCategory):
                data_param.setIndex(trans_param.getIndex())
            else:
                data_param.setVal(trans_param.getVal())
                # This sets a symmetric error, but since we don't run Minos, that's ok
                data_param.setError(trans_param.getError())
Beispiel #5
0
    def run(self, **kwargs):
        from ROOT import RooArgSet

        __check_req_kw__("Observables", kwargs)
        __check_req_kw__("Pdf", kwargs)

        observables = kwargs.pop("Observables")
        obs_set = RooArgSet(*observables)

        pdf = kwargs.pop("Pdf")
        genPdf = kwargs.pop("GenPdf", pdf)

        gen_obs_set = RooArgSet()
        for o in list(observables) + list(genPdf.ConditionalObservables()):
            gen_obs_set.add(o._target_())
        gen_pdf_params = genPdf.getParameters(gen_obs_set).snapshot(True)

        genPdf = genPdf.clone(genPdf.GetName() + "_toy_clone")
        genPdf.recursiveRedirectServers(gen_pdf_params)

        fit_obs_set = RooArgSet()
        for o in list(observables) + list(pdf.ConditionalObservables()):
            fit_obs_set.add(o._target_())
        params = pdf.getParameters(fit_obs_set)

        pdf_params = RooArgSet()
        for p in params:
            if p.isConstant():
                continue
            pdf_params.add(p)
        ## for param in pdf_params:
        ##     if param.GetName() not in ['Gamma', 'dGamma']:
        ##         param.setConstant()
        self._gen_params = pdf_params.snapshot(True)

        # Make another ArgSet to put the fit results in
        result_params = RooArgSet(pdf_params, "result_params")

        transform = self.transform()
        if transform:
            trans_params = transform.gen_params(gen_obs_set)
            for p in trans_params:
                result_params.add(p)

        # Some extra numbers of interest
        from ROOT import RooRealVar

        NLL = RooRealVar("NLL", "-log(Likelihood)", 1.0)
        ngen = RooRealVar("ngen", "number of generated events", self.options().nevents)
        seed = RooRealVar("seed", "random seed", 0.0)
        from ROOT import RooCategory

        status = RooCategory("status", "fit status")
        status.defineType("success", 0)
        status.defineType("one", 1)
        status.defineType("two", 2)
        status.defineType("three", 3)
        status.defineType("other", 4)
        result_params.add(status)
        result_params.add(NLL)
        result_params.add(ngen)
        result_params.add(seed)

        # The dataset to store the results
        from ROOT import RooDataSet

        self._data = RooDataSet("result_data", "result_data", result_params)
        data_params = self._data.get()

        from ROOT import RooRandom
        import struct, os

        while self._data.numEntries() < self.options().ntoys:
            # Get a good random seed, set it and store it
            s = struct.unpack("I", os.urandom(4))[0]
            RooRandom.randomGenerator().SetSeed(s)
            seed.setVal(s)

            # Reset pdf parameters to initial values. Note: this does not reset the estimated errors...
            pdf_params.assignValueOnly(self.gen_params())
            args = dict(NumEvents=self.options().nevents)
            if "ProtoData" in kwargs:
                args["ProtoData"] = kwargs.pop("ProtoData")

            genPdf.getParameters(obs_set).assignValueOnly(gen_pdf_params)
            data = genPdf.generate(obs_set, **args)
            if transform:
                data = transform(data)
                if not data:
                    # Transform has failed
                    transform.set_params(data_params)
                    self._data.add(data_params)
                    continue

            if data.isWeighted() and "SumW2Error" not in self.fit_opts():
                self.fit_opts()["SumW2Error"] = False

            j = 0
            while j < 4:
                fit_result = pdf.fitTo(data, NumCPU=self.options().ncpu, **(self.fit_opts()))
                if fit_result.status() == 0:
                    fit_result.Print()
                    break
                j += 1
            if fit_result.status() != 0:
                print "Fit result status = %s" % fit_result.status()
            NLL.setVal(fit_result.minNll())
            if fit_result.status() < 4:
                status.setIndex(fit_result.status())
            else:
                status.setIndex(4)
            for result_param in result_params:
                data_param = data_params.find(result_param.GetName())
                if isinstance(result_param, RooCategory):
                    data_param.setIndex(result_param.getIndex())
                else:
                    data_param.setVal(result_param.getVal())
                    # This sets a symmetric error, but since we don't run Minos, that's ok
                    data_param.setError(result_param.getError())
            if transform:
                transform.set_params(data_params)

            self._data.add(data_params)

        return self.data()
    1: 0.,  # {1 : 2.37,
    2: 0.,  #  2 : 0.69,
    3: 0.,  #  3 : -0.57,
    4: 0.
}  #  4 : -1.96}
res_cor_match = {1: 0, 2: 0, 3: 0, 4: 0}

ds_res = {
    k: RooDataSet("ds_M%d" % (k + 1), "ds_M%d" % (k + 1), obs)
    for k in range(1, 5)
}
for entry in tree:
    if entry.p < 3000:
        continue
    for i in range(entry.n_hits):
        q.setIndex(entry.q)
        s = entry.res_s[i]
        res_x.setVal(entry.res_x[i] - entry.q * res_cor_true[s])
        res_y.setVal(entry.res_y[i])
        pull_x.setVal(entry.res_x[i] / (2 * entry.res_dx[i]) -
                      entry.q * pull_cor_true[s])
        pull_y.setVal(entry.res_y[i] / (2 * entry.res_dy[i]))
        match_pull_x.setVal(entry.match_res_x[i] / (2 * entry.res_dx[i]) -
                            entry.q * pull_cor_match[s])
        match_pull_y.setVal(entry.match_res_y[i] / (2 * entry.res_dy[i]))
        match_res_x.setVal(entry.match_res_x[i] - entry.q * res_cor_match[s])
        match_res_y.setVal(entry.match_res_y[i])
        ds_res[s].add(obs)

x_obs = match_pull_x
y_obs = match_pull_y