Python ROOT.RooArgSet类代码示例

def applyDecayTimeErrPdf(config, name, ws, time, timeerr, qt, qf, mistagobs,
        timepdf, timeerrpdf, mistagpdf):
    """
    apply the per-event time error pdf

    config      -- configuration dictionary
    name        -- prefix to be used for new RooFit objects
    ws          -- workspace to import new objects into
    time        -- time observable
    timeerr     -- decay time error observable (or None for average decay time
                   error)
    qt          -- tagging decision
    qf          -- final state charge
    mistagobs   -- mistag observable (or None for average mistag)
    timepdf     -- decay time pdf
    timeerrpdf  -- decay time error pdf (or None for av. decay time error)
    mistagpdf   -- mistag pdf (or None for average mistag)

    returns the (possibly modified) decay time error pdf (with decay time
    error multiplied on, if applicable).
    """
    # no per-event time error is easy...
    if None == timeerrpdf: return timepdf
    from ROOT import RooFit, RooArgSet, RooProdPdf
    noncondset = RooArgSet(time, qf, qt)
    if None != mistagpdf:
        noncondset.add(mistagobs)
    timepdf = WS(ws, RooProdPdf('%s_TimeTimeerrPdf' % name,
        '%s (time,timeerr) pdf' % name, RooArgSet(timeerrpdf),
        RooFit.Conditional(RooArgSet(timepdf), noncondset)))
    return timepdf

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

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

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

 def cnvrt(i) :
     if not hasattr(i,'__iter__') or any( isinstance(i, t) for t in __doNotConvert ): return i
     _i = RooArgSet()
     for j in i : 
         from ROOT import RooAbsArg
         if not isinstance(j,RooAbsArg) : return i
         _i.add( j )
     return _i

    def get_dataset_from_tree(
        self,
        path_to_tree,
        tree_variables,
        weight="1==1",
        weight_var_name=0,
        dataset_name="my_dataset",
        basket=True,
        category=None,
    ):
        """
        Creates RooDataSet from a plain root tree given:
        - variables name list
        - weight expression. It works in the same way as TTree cut.

        Returns:
        --------
        - RooDataSet
        - also fills the basket with datasets (basket inhereted from RootHelperBase class)

        TODO
        ----
        - add implementation for category setting(check in prepare_toy_datasets_for_sync)
            - check if adding toy dataset to each channel workspace individually behaves well
              after combineCards.py.
        """

        # make RooRealVars from tree_variables
        my_arg_set = RooArgSet()
        my_rrv = dict()
        for var_name in tree_variables:
            # TODO implement check that branch exist
            my_rrv[var_name] = RooRealVar(var_name, var_name, -999999999, 999999999)
            my_arg_set.add(my_rrv[var_name])
        if self.DEBUG:
            self.log.debug("RooArgSet is now:")
            my_arg_set.Print()

        # get the tree from path_to_tree
        my_tree = self.get_TTree(path_to_tree, cut=weight)
        self.log.debug("Selected tree contains {0} events".format(my_tree.GetEntries()))
        # create RooDataSet and reduce tree if needed
        # self.dataset_from_tree =  RooDataSet(dataset_name, dataset_name, my_tree, my_arg_set, weight).reduce(my_arg_set)
        self.dataset_from_tree = RooDataSet(dataset_name, dataset_name, my_tree, my_arg_set)
        # self.dataset_from_tree =  RooDataSet(dataset_name, dataset_name, my_tree, my_arg_set, "", weight_var_name)
        # data[j]=new RooDataSet(Form("data%d",j),Form("data%d",j),outTree,RooArgSet(rCMS_zz4l_widthKD,rCMS_zz4l_widthMass,rweightFit),"","_weight_");
        self.log.debug("RooDataSet contains {0} events".format(self.dataset_from_tree.sumEntries()))
        # .reduce(ROOT.RooArgSet(self.D0))
        self.current_arg_set = my_arg_set

        # add dataset to basket
        if basket:
            self.add_to_basket(self.dataset_from_tree, new_name=dataset_name, new_title=dataset_name)

        return self.dataset_from_tree

def expectedPlcLimit(obs_, poi_, model, ws, ntoys = 30, CL = 0.95):
    # obs : observable variable or RooArgSet of observables
    # poi : parameter of interest or RooArgSet of parameters
    # model : RooAbsPdf of model to consider including any constraints
    #         the parameters should have the values corresponding to the
    #         background-only hypothesis which will be used to  estimate the
    #         expected limit.
    # ntoys : number of toy datsets to generate to get expected limit
    # CL : confidence level for interval
    # returns a dictionary containing the expected limits and their 1 sigma
    # errors for the first/only parameter in poi_ and a list of the results
    # from the individual toys.

    from math import sqrt
    obs = RooArgSet(obs_)
    obs.setName('observables')
    mPars = model.getParameters(obs)
    genPars = mPars.snapshot()

    print "parameters for generating toy datasets"
    genPars.Print("v")

    limits = []
    sumUpper = 0.
    sumUpper2 = 0.
    sumLower = 0.
    sumLower2 = 0.
    nOK = 0
    for i in range(0,ntoys):
        print 'generate limit of toy %i of %i' % (i+1, ntoys)
        mPars.assignValueOnly(genPars)

        toyData = model.generate(obs, RooFit.Extended())
        toyData.SetName('data_obs_%i' % i)

        limits.append(plcLimit(obs_, poi_, model, ws, toyData, CL))

        if limits[-1]['limits'][poi_.GetName()]['ok']:
            nOK += 1
            sumUpper += limits[-1]['limits'][poi_.GetName()]['upper']
            sumUpper2 += limits[-1]['limits'][poi_.GetName()]['upper']**2
            sumLower += limits[-1]['limits'][poi_.GetName()]['lower']
            sumLower2 += limits[-1]['limits'][poi_.GetName()]['lower']**2

        toyData.IsA().Destructor(toyData)

    expLimits = {'upper' : sumUpper/nOK,
                 'upperErr' : sqrt(sumUpper2/(nOK-1)-sumUpper**2/nOK/(nOK-1)),
                 'lower' : sumLower/nOK,
                 'lowerErr' : sqrt(sumLower2/(nOK-1)-sumLower**2/nOK/(nOK-1)),
                 'ntoys': nOK
                 }
    return (expLimits, limits)

    def addAsciiData(self, ws, ds_name, item_title, var_set_name, var_set_type, ds_file_name, weight_var_name = None, debug = 0, sets = None):
        legend = '[exostConfig::addAsciiData]:'
        
        arg_set = RooArgSet()
        if (var_set_type == 'set'):
            #_var_set = ws.set(var_set_name)
            _var_set = sets[var_set_name]
            arg_set.add(_var_set)

            if _var_set.getSize() != 1:
                print legend, 'Error: too many or too few columns in the input ASCII file'
                print legend, 'Error: Smart program as I am, I can only handle ASCII files'
                print legend, 'Error: with exactly one column at the moment.'
                print legend, 'Error: Support for multiple columns will be implemented'
                print legend, 'Error: eventually, contact the developers.'
                print legend, 'Error: Better yet, switch to ROOT input files,'
                print legend, 'Error: all the cool kids are doing that!'

                return -1
            
        elif (var_set_type == 'var'):
            arg_set.add(ws.var(var_set_name))
        else:
            print legend, 'error: unknown var_set_type, cannot create dataset', ds_name
            return -1

        #create the dataset
        if weight_var_name == None: #no weight
            if (debug>0):
                print legend, 'no weight variable given'

            _arglist = RooArgList(arg_set)
            #ds = RooDataSet()ds_name, item_title)
            ds = RooDataSet.read(ds_file_name, _arglist)
            ds.SetName(ds_name)
            ds.SetTitle(item_title)
        else:
            if (debug>0):
                print legend, 'using variable', weight_var_name, 'as weight'
        
                print legend, 'Error: Smart program as I am, I cannot handle'
                print legend, 'Error: weights when loading from ASCII files yet.'
                print legend, 'Error: Support for weights from ASCII files will'
                print legend, 'Error: be implemented eventually, contact the developers.'
                print legend, 'Error: Better yet, switch to ROOT input files,'
                print legend, 'Error: all the cool kids are doing that!'

                return -1


        # import the datahist. Note workaround 'import' being a reserved word
        getattr(ws, 'import')(ds)

def plcLimit(obs_, poi_, model, ws, data, CL = 0.95, verbose = False):
    # obs : observable variable or RooArgSet of observables
    # poi : parameter of interest or RooArgSet of parameters
    # model : RooAbsPdf of model to consider including any constraints
    # data : RooAbsData of the data
    # CL : confidence level for interval
    # returns a dictionary with the upper and lower limits for the first/only
    # parameter in poi_ as well as the interval object and status flag
    
    obs = RooArgSet(obs_)
    obs.setName('observables')
    poi = RooArgSet(poi_)
    poi.setName('poi')
    poi.setAttribAll('Constant', False)
    nuis = model.getParameters(obs)
    nuis.remove(poi)
    nuis.remove(nuis.selectByAttrib('Constant', True))
    nuis.setName('nuisance')

    if verbose:
        print 'observables'
        obs.Print('v')
        print 'parameters of interest'
        poi.Print('v')
        print 'nuisance parameters'
        nuis.Print('v')
 
    mc = RooStats.ModelConfig('mc')
    mc.SetWorkspace(ws)
    mc.SetPdf(model)
    mc.SetObservables(obs)
    mc.SetParametersOfInterest(poi)
    mc.SetNuisanceParameters(nuis)

    plc = RooStats.ProfileLikelihoodCalculator(data, mc)
    plc.SetConfidenceLevel(CL)

    interval = plc.GetInterval()

    upperLimit = Double(999.)
    lowerLimit = Double(0.)
    Limits = {}

    paramIter = poi.createIterator()
    param = paramIter.Next()
    while param:
        ok = interval.FindLimits(param, lowerLimit, upperLimit)
        Limits[param.GetName()] = {'ok' : ok, 'upper' : float(upperLimit),
                                   'lower' : float(lowerLimit)}
        param = paramIter.Next()

    if verbose:
        print '%.0f%% CL limits' % (interval.ConfidenceLevel() * 100)
        print Limits

    Limits['interval'] = interval
    return Limits

    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 File2Dataset(self, fnames, dsName, ws, noCuts = False, 
                     weighted = False, CPweight = False, cutOverride = None,
                     interference = 0, additionalWgt = 1.0):
        if ws.data(dsName):
            return ws.data(dsName)

        cols = RooArgSet(ws.set('obsSet'))
        # print 'interference weight flag:',interference
        if (weighted):
            evtWgt = RooRealVar('evtWgt', 'evtWgt', 1.0)
            cols.add(evtWgt)
            ds = RooDataSet(dsName, dsName, cols, 'evtWgt')
            print 'including weights for eff',
            if CPweight:
                print 'and CP weight',
            if interference in [1,2,3]:
                print 'and interference',
            print
        else:
            ds = RooDataSet(dsName, dsName, cols)
            
        if not (type(fnames) == type([])):
            fnames = [fnames]

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

        for fname in fnames:
            for (row, effWgt, cpw, iwt) in \
                    self.TreeLoopFromFile(fname, noCuts,
                                          CPweight = CPweight,
                                          cutOverride = cutOverride,
                                          interference = interference):
                inRange = True
                for (i,v) in enumerate(obs):
                    inRange = (inRange and ws.var(v).inRange(row[i], ''))
                    cols.setRealValue(v, row[i])
                if CPweight:
                    effWgt *= cpw
                if interference in [1,2,3]:
                    effWgt *= iwt
                if inRange:
                    ds.add(cols, effWgt*additionalWgt)

        getattr(ws, 'import')(ds)

        return ws.data(dsName)

    def write_output(self):
        # Write the results to a file
        from ROOT import TFile

        output_file = TFile.Open(self.options().output, "recreate")
        output_file.WriteTObject(self._data, self._data.GetName())
        gp = self.gen_params()
        if gp:
            from ROOT import RooArgSet

            gpars = RooArgSet()
            for p in gp:
                gpars.add(p)
            gpars = gpars.snapshot(True)
            output_file.WriteTObject(gpars, "gen_params")
        output_file.Close()

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

        #create the dataset
        if weight_var_name == None: #no weight
            if (debug>0):
                print legend, 'no weight branch given'
            ds = RooDataSet(ds_name, item_title, ds_object, arg_set)
        else:
            if (debug>0):
                print legend, 'using weight branch', weight_var_name

            # old and incorrect(?) way of applying weights
            #
            #_var_formula = RooFormulaVar('f'+weight_var_name,
            #                             'f'+weight_var_name,
            #                             '@0',
            #                             RooArgList(ws.var(weight_var_name)))
            #ds = RooDataSet(ds_name, item_title, ds_object, arg_set, _var_formula, weight_var_name)
            
            arg_set.add(ws.var(weight_var_name))
            ds = RooDataSet(ds_name, item_title, ds_object, arg_set, "1>0", weight_var_name)
        
        # import the datahist. Note workaround 'import' being a reserved word
        getattr(ws, 'import')(ds)

def make_roodataset(python_list_of_vars, file_name, 
                    tree_name = "BdToKstarMuMu_Data", 
                    path_to_save = "select_kstarmumu_roodataset.root", cuts=None):
    """
    Makes a RooDataSet in the current directory. 
    I only really want to do this once since python is slow.
    """

    from ROOT import RooArgSet, RooDataSet

    # open the files and get the trees
    from sam.pp.utils.root import get_tree_from_file
    (f,t) = get_tree_from_file(file_name, tree_name)

    # get the variables into the stupid roofit format
    ras = RooArgSet()
    for var in python_list_of_vars:
        ras.add(var)
    ras.Print()
   
    from ROOT import TFile
    fcut = TFile('/tmp/temp.root', 'RECREATE')
    if cuts:
        tr = t.CopyTree(cuts)
    else:
        tr = t

    # make a RooDataSet from both trees
    print("making the RooArgSet... yes this takes a long time because python"\
        + " is slower than C++ \nperhaps you could use the time to reflect"\
        + "on how much you hate coding C++ and how beautiful python is")
    ds = RooDataSet("ds", "data set", tr, ras)
    print("made the RooDataSet")
    fcut.Close()

    # save it
    #from ROOT import TFile
    fds = TFile(path_to_save, "RECREATE")
    fds.cd()
    ds.Write()
    print("saved the RooDataSet")

    return ds

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

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

def hessian(nll, parameters, N = 9):
    values = dict([(p.GetName(), (p.getVal(), p.getError())) for p in parameters])
    nll_min = nll.getVal()
    __parameters = RooArgSet()
    for p in parameters:
        __parameters.add(p)
    nll_params = nll.getObservables(__parameters)
    
    H = TMatrixTSym('double')(parameters.getSize())
    diagonals = {}
    for i, x in enumerate(parameters):
        p = nll_params.find(x.GetName())
        r = fit_parabola(nll, nll_min, p, N)
        diagonals[p.GetName()] = r
        H[i][i] = 2 * r[1]
    
    for i, x in enumerate(parameters):
        x = nll_params.find(x.GetName())
        for j in range(i + 1, parameters.getSize()):
            y = nll_params.find(parameters.at(j).GetName())
            H[i][j] = H[j][i] = fit_cross(nll, nll_min, x, y, diagonals, N)
    
    return H