Python RooArgSet.add方法代码示例

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
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

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
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

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
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

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
 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

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
    def fit(self):
        fit_variable = RooRealVar("fit_variable", "fit_variable", self.fit_boundaries[0], self.fit_boundaries[1])
        fit_variable.setBins(self.histograms[self.data_label].nbins())
        variables = RooArgList()
        variables.add(fit_variable)
        variable_set = RooArgSet()
        variable_set.add(fit_variable)

        roofit_histograms = {}
        roofit_pdfs = {}
        roofit_variables = {}

        N_min = 0.0
        N_max = self.normalisation[self.data_label] * 2.0
        pdf_arglist = RooArgList()
        variable_arglist = RooArgList()

        roofit_histograms[self.data_label] = RooDataHist(
            self.data_label, self.data_label, variables, self.histograms[self.data_label]
        )
        for sample in self.samples:
            roofit_histogram = RooDataHist(sample, sample, variables, self.histograms[sample])
            roofit_histograms[sample] = roofit_histogram
            roofit_pdf = RooHistPdf("pdf" + sample, "pdf" + sample, variable_set, roofit_histogram)
            roofit_pdfs[sample] = roofit_pdf
            roofit_variable = RooRealVar(sample, sample + " events", self.normalisation[sample], N_min, N_max)
            roofit_variables[sample] = roofit_variable
            pdf_arglist.add(roofit_pdf)
            variable_arglist.add(roofit_variable)

        model = RooAddPdf("model", "sum of all known", pdf_arglist, variable_arglist)
        use_model = model
        if self.constraints:
            arg_set = RooArgSet(model)
            constraints = self.get_fit_normalisation_constraints(model, roofit_variables)
            for constraint in constraints:
                arg_set.add(constraint)
            model_with_constraints = RooProdPdf(
                "model_with_constraints", "model  with gaussian constraints", arg_set, RooLinkedList()
            )
            use_model = model_with_constraints

        if self.method == "TMinuit":
            # WARNING: number of cores changes the results!!!
            self.saved_result = use_model.fitTo(
                roofit_histograms[self.data_label],
                RooFit.Minimizer("Minuit2", "Migrad"),
                RooFit.NumCPU(1),
                RooFit.Extended(),
                RooFit.Save(),
            )

        results = {}
        for sample in self.samples:
            results[sample] = (roofit_variables[sample].getVal(), roofit_variables[sample].getError())
        self.results = results

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
    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

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
def makeRooDataset(inputfile_name):
    chibTree_name = 'rootuple/chibTree'
    inputfile = TFile.Open(inputfile_name,"READ")
    tree = TTree()
    inputfile.GetObject(chibTree_name, tree)
    
    invm1S = RooRealVar("invm1S", "invm1S", 9.7, 10.1)#9.5, 11.5
    invm2S = RooRealVar("invm2S", "invm2S", 9.5, 13.0)
    invm3S = RooRealVar("invm3S", "invm3S", 9.5, 13.0)
    chib_mass = RooRealVar("chib_mass", "chib_mass", 5., 15.0)
    chib_pt = RooRealVar("chib_pt", "chib_pt", 0., 300.)
    chib_eta = RooRealVar("chib_eta", "chib_eta", -5., 5.)
    chib_phi = RooRealVar("chib_phi", "chib_phi", -3.2, 3.2)
    dimuon_mass = RooRealVar("dimuon_mass","dimuon_mass", 8.0, 12.0)
    dimuon_rapidity = RooRealVar("dimuon_rapidity", "dimuon_rapidity", -5.0, 5.0)
    dimuon_pt = RooRealVar("dimuon_pt","dimuon_pt", 0.0, 100.0)
    photon_eta = RooRealVar("photon_eta","photon_eta", -5.0, 5.0)
    photon_pt = RooRealVar("photon_pt","photon_pt", 0.0, 100.0)
    ctpv = RooRealVar("ctpv","ctpv", -5.0, 5.0)
    ctpv_error = RooRealVar("ctpv_err","ctpv_err", -5.0, 5.0)
    pi0_abs_mass = RooRealVar("pi0_abs_mass","pi0_abs_mass", 0.0, 2.0)
    Y1S_nsigma = RooRealVar("Y1S_nsigma","Y1S_nsigma",0.0,30.0)
    Y2S_nsigma = RooRealVar("Y2S_nsigma","Y2S_nsigma",0.0,30.0)
    Y3S_nsigma = RooRealVar("Y3S_nsigma","Y3S_nsigma",0.0,35.0)
    conv_vertex = RooRealVar("conv_vertex", "conv_vertex", 0.0, 70.0)
    dz = RooRealVar("dz","dz", 0., 0.6)
    numPrimaryVertices = RooRealVar("numPrimaryVertices","numPrimaryVertices",0,60)
    rf1S_chib_mass = RooRealVar("rf1S_chib_mass","rf1S_chib_mass",9.5,60.)
    probFit1S = RooRealVar("probFit1S","probFit1S",0.,1.)
    rf1S_chib_pt = RooRealVar("rf1S_chib_pt", "rf1S_chib_pt", 0., 300.)
    rf1S_chib_eta = RooRealVar("rf1S_chib_eta", "rf1S_chib_eta", -5., 5.)
    rf1S_dimuon_mass = RooRealVar("rf1S_dimuon_mass","rf1S_dimuon_mass", 8.0, 12.0)
    rf1S_dimuon_rapidity = RooRealVar("rf1S_dimuon_rapidity", "rf1S_dimuon_rapidity", -5.0, 5.0)
    rf1S_dimuon_pt = RooRealVar("rf1S_dimuon_pt","rf1S_dimuon_pt", 0.0, 230.0)
    rf1S_photon_eta = RooRealVar("rf1S_photon_eta","rf1S_photon_eta", -5.0, 5.0)
    rf1S_photon_pt = RooRealVar("rf1S_photon_pt","rf1S_photon_pt", 0.0, 100.0)
    rf2S_chib_mass = RooRealVar("rf2S_chib_mass","rf2S_chib_mass",9.5,60.)
    probFit2S = RooRealVar("probFit2S","probFit2S",0.,1.)
    rf3S_chib_mass = RooRealVar("rf3S_chib_mass","rf3S_chib_mass",9.5,60.)
    probFit3S = RooRealVar("probFit3S","probFit3S",0.,1.)

    muonP_pt = RooRealVar("muonP_pt","muonP_pt",0.,100.)
    muonM_pt = RooRealVar("muonM_pt","muonM_pt",0.,100.)
    muonP_eta = RooRealVar("muonP_eta","muonP_eta",-5.,5.)
    muonM_eta = RooRealVar("muonM_eta","muonM_eta",-5.,5.)

    dataArgSet = RooArgSet(invm1S, invm2S, invm3S, dimuon_mass, dimuon_rapidity, dimuon_pt, photon_eta, photon_pt, ctpv)
    dataArgSet.add(RooArgSet(chib_mass, chib_pt, chib_eta, chib_phi))
    dataArgSet.add(RooArgSet(ctpv_error, pi0_abs_mass, Y1S_nsigma, Y2S_nsigma, Y3S_nsigma, conv_vertex, dz, numPrimaryVertices))
    dataArgSet.add(RooArgSet(rf1S_chib_mass, probFit1S, rf2S_chib_mass, probFit2S, rf3S_chib_mass, probFit3S))
    dataArgSet.add(RooArgSet(rf1S_chib_pt, rf1S_chib_eta, rf1S_dimuon_mass, rf1S_dimuon_rapidity, rf1S_dimuon_pt, rf1S_photon_eta, rf1S_photon_pt))
    dataArgSet.add(RooArgSet(muonP_pt, muonM_pt, muonP_eta, muonM_eta))
    dataSet = RooDataSet("chibds","Chib RooDataSet", tree, dataArgSet)
    return dataSet

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
    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)

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
    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

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
    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)

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
    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()

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
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

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
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)

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
    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)

# 需要导入模块: from ROOT import RooArgSet [as 别名]
# 或者: from ROOT.RooArgSet import add [as 别名]
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