Python RooDataSet.addColumn方法代码示例

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import addColumn [as 别名]
def get_dataset(varargset, ftree, cut="", wt="", scale=1):
    """Return a dataset.

    Return a dataset from the ntuple `ftree'. Apply a selection cut
    using the `cutVar' variable and the selection `cut'.

    """

    from rplot.fixes import ROOT
    from rplot.tselect import Tsplice

    splice = Tsplice(ftree)
    splice.make_splice("sel", cut)
    from ROOT import RooDataSet, RooFit, RooFormulaVar, RooArgList

    tmpdst = RooDataSet("tmpdataset", "", varargset, RooFit.Import(ftree))
    if wt:
        wtvar = RooFormulaVar("wt", "{}*@0".format(scale), RooArgList(varargset[wt]))
        wtvar = tmpdst.addColumn(wtvar)
        varargset.remove(varargset[wt])
        varargset.add(wtvar)
        dst = RooDataSet("dataset", "Dataset", varargset, RooFit.Import(tmpdst), RooFit.WeightVar(wtvar))
        varargset.remove(wtvar)
        dst = dst.reduce(varargset)
    return dst

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

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

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

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

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

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

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

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

    data1.append(data2)

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

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import addColumn [as 别名]
def prepare_truth_models(ws,cat,mass,channel,turnon,truth):    
    if channel in study_inputs:
        bkg_data = RooDataSet("bkgdata_%s"%(channel),
                              "M_{ll#gamma} with Errors",
                              ws.set("observables_weight"),
                              "weight")
        
        for k,file in enumerate(study_inputs[channel]):
            f_in = TFile.Open(file,"READ")
            gDirectory.cd(pwd)
            n_events = f_in.Get("eventCount").GetBinContent(1)
            tree = f_in.Get("selected_zg").CloneTree()
            tree.SetName("tree_%s_%i"%(channel,k))
            f_in.Close()            
            
            d_in = RooDataSet("temp",
                              "M_{ll#gamma} with Errors",
                              tree,
                              ws.set("observables"))
            norm = RooConstVar("norm","norm",n_events)
            weight = RooFormulaVar("weight",
                                   "weight",
                                   "1.3*@0*@1/@2", #1.3 gives 19.6/fb
                                   RooArgList(ws.var("puWeight"),
                                              ws.var("procWeight"),
                                              norm)
                                   )
            d_in.addColumn(weight)
            d_in_weight = RooDataSet("temp_weight",
                                     "M_{ll#gamma} with Errors",
                                     d_in,
                                     ws.set("observables_weight"),
                                     '','weight')
            bkg_data.append(d_in_weight)
            
        # split off data for each category, create the
        # toy dataset truth models
        data = bkg_data.reduce("Mz + Mzg > 185 && r94cat == %i"%cat)
        getattr(ws,'import')(data,
                             RooFit.Rename('bkgdata_%s_%i'%(channel,
                                                            cat)))
        nevts = data.sumEntries('Mzg > %f && Mzg < %f'%(mass-1.5,mass+1.5))
        #for sigm turn on we want erf truth
        if turnon == 'sigm' and truth == 'exp': 
            #make RooDecay 'truth' model with erf turn on
            ws.factory(
                'RooGaussModel::MzgResoShape_exp_erf_%s_cat%i(Mzg,'\
                'bias_exp_erf_%s_cat%i[120,90,150],sigma_exp_erf_%s_cat%i[1,0.01,10])'%(
                channel,cat,
                channel,cat,
                channel,cat)
                )
            ws.factory(
                'RooDecay::MzgTruthModelBase_exp_erf_%s_cat%i(Mzg,'\
                'tau_erf_%s_cat%i[25,0,50],MzgResoShape_exp_erf_%s_cat%i,'
                'RooDecay::SingleSided)'%(channel,cat,
                                          channel,cat,
                                          channel,cat)
                )            
            ws.factory(
                'RooExtendPdf::MzgTruthModel_exp_erf_%s_cat%i('\
                'MzgTruthModelBase_exp_erf_%s_cat%i,'\
                'norm_truth_exp_%s_cat%i[%f,%f,%f],"ROI")'%(channel,cat,
                                                            channel,cat,
                                                            channel,cat,
                                                            nevts,
                                                            0.25*nevts,1.75*nevts)
                )
            ws.pdf('MzgTruthModel_exp_erf_%s_cat%i'%(channel,cat)).fitTo(
                ws.data('bkgdata_%s_%i'%(channel,cat)),
                RooFit.Minimizer('Minuit2','scan'),
                RooFit.SumW2Error(False)
                )
            ws.pdf('MzgTruthModel_exp_erf_%s_cat%i'%(channel,cat)).fitTo(
                ws.data('bkgdata_%s_%i'%(channel,cat)),
                RooFit.Minimizer('Minuit','simplex'),
                RooFit.SumW2Error(False)
                )
            ws.pdf('MzgTruthModel_exp_erf_%s_cat%i'%(channel,cat)).fitTo(
                ws.data('bkgdata_%s_%i'%(channel,cat)),
                RooFit.SumW2Error(True)
                )
        if turnon == 'sigm' and truth == 'pow':
            #make power-law truth model with erf turn on
            ws.factory('EXPR::MzgTruthModelShape_pow_erf_%s_cat%i('\
                       '"1e-20 + (@0 > @1)*((@0)^([email protected]))",'\
                       '{Mzg,step_pow_erf_%s_cat%i[105,100,130],'\
                       'pow_%s_cat%i[2,0,10]})'\
                       %(channel,cat,
                         channel,cat,
                         channel,cat))
            ws.factory(
                'RooGaussModel::MzgResoShape_pow_erf_%s_cat%i(Mzg,'\
                'bias_pow_erf_%s_cat%i[0],sigma_pow_erf_%s_cat%i[1,0.01,10])'%(
                channel,cat,
                channel,cat,
                channel,cat)
                )
            ws.factory('FCONV::MzgTruthModelBase_pow_erf_%s_cat%i(Mzg,'\
                       'MzgTruthModelShape_pow_erf_%s_cat%i,'\
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import addColumn [as 别名]
        coefs = acc_fun.coefficients()[acc_fun.GetName()]
        for i, c in enumerate(coefs):
            c.setVal(hist.GetBinContent(i + 1))
        
    ## RooMsgService.instance().addStream(RooFit.DEBUG, RooFit.Topic(RooFit.Integration))
    result = sig_t.fitTo(data, **fitOpts)

    tau_result.setVal(tau.getVal())
    tau_result.setError(tau.getError())
    result_data.add(RooArgSet(tau_result))

    tau.setVal(tau_true.getVal())
    tau.setError(0.1)

tau_pull = RooPullVar("tau_pull", "tau_pull", tau_result._target_(), tau_true._target_())
tau_pull = result_data.addColumn(tau_pull)
tau_pull = w.put(tau_pull)

mean = RealVar('mean', Value = 0, MinMax = (-4, 4))
sigma = RealVar('sigma', Value = 1, MinMax = (0.1, 5))
from ROOT import RooGaussian
gauss = Pdf(Name = 'gauss', Type = RooGaussian, Parameters = (tau_pull, mean, sigma))
result = gauss.fitTo(result_data, **fitOpts)

from ROOT import TPostScript
ps = TPostScript("test.eps", 113)
ps.FontEmbed()
from ROOT import TCanvas
canvas = TCanvas('canvas', 'canvas', 600, 400)
frame = tau_pull.frame(Range = (-5, 5))
frame.GetXaxis().SetTitle('#tau pull')

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import addColumn [as 别名]
lab1_PT = RooRealVar("lab1_PT", "lab1_PT", 0., 100000., "MeV/c")
#lab2_P = RooRealVar("lab2_P", "lab2_P", 5000., 500000., "MeV/c")
allVars = [lab2_MM, lab2_IP_OWNPV, lab0_LifetimeFit_ctau0, lab0_LifetimeFit_ctauErr0, lab0_MM, lab1_P, lab1_PT]
varSet = RooArgSet(*allVars)
if is_MC:
    varSet.add(lab0_TRUETAU)
varSetMC = RooArgSet(*allVars)
varSetMC.add(lab0_TRUETAU)

# The actual RooDataSet containing the data (this may take a while)
print("Getting the data...")
dataset = RooDataSet("data", "data", t, varSet)
print("Done!")

# Add unit-corrected LifetimeFit variables (mm -> fs)
lab0_LifetimeFit_TAU = dataset.addColumn(RooFormulaVar("lab0_LifetimeFit_TAU", "lab0_LifetimeFit_TAU", "lab0_LifetimeFit_ctau0/{}".format(speedOfLight_mm_fs), RooArgList(lab0_LifetimeFit_ctau0)))
lab0_LifetimeFit_TAU.setUnit("fs")
lab0_LifetimeFit_TAU.setMin(-500.)
lab0_LifetimeFit_TAU.setMax(50000.)
lab0_LifetimeFit_TAUERR = dataset.addColumn(RooFormulaVar("lab0_LifetimeFit_TAUERR", "lab0_LifetimeFit_TAUERR", "lab0_LifetimeFit_ctauErr0/{}".format(speedOfLight_mm_fs), RooArgList(lab0_LifetimeFit_ctauErr0)))
lab0_LifetimeFit_TAUERR.setUnit("fs")
lab0_LifetimeFit_TAUERR.setMin(  0.)
lab0_LifetimeFit_TAUERR.setMax(150.)

# Add difference between measured and true lifetime as a variable
#taudiff = dataset.addColumn(RooFormulaVar("lab0_LifetimeFit_TAU_DIFF", "lab0_LifetimeFit_TAU_DIFF",\
        #"lab0_LifetimeFit_ctau0/{}-lab0_TRUETAU/{}".format(speedOfLight_mm_fs, ns_to_fs), RooArgList(lab0_LifetimeFit_ctau0, lab0_TRUETAU)))
#taudiff.setUnit("fs")
#taudiff.setMin(-500.)
#taudiff.setMax( 1500.)

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import addColumn [as 别名]
    # read data set with events in two asymmetry categories
    print 'plotAsymmetry: reading dataset with events in two asymmetry categories'
    dataFile = TFile.Open(dataSetFileOut)
    dataSetAsym = dataFile.Get('asymData')
    dataFile.Close()

dataSetAsym.Print()

# create weighted data set
from ROOT import RooProduct, RooArgSet, RooArgList
obsSet = RooArgSet( dataSetAsym.get() )
prodList = RooArgList( obsSet.find('sigWeight') )
if applyDilWeights : prodList.add( obsSet.find('dilution') )
if applyAngWeights : prodList.add( obsSet.find( 'angWeight_%s' % applyAngWeights ) )
weightVar = RooProduct( 'weightVar', 'weightVar', prodList )
weightVar = dataSetAsym.addColumn(weightVar)
obsSet.add(weightVar)

dataSetAsymW = RooDataSet( 'asymDataW', 'asymDataW', obsSet, Import = dataSetAsym, WeightVar = ( 'weightVar', True ) )
del dataSetAsym
ws.put(dataSetAsymW)
del dataSetAsymW
dataSetAsymW = ws['asymDataW']
obsSet = RooArgSet( dataSetAsymW.get() )
dataSetAsymW.Print()

# get sums of weights
sumW = dict( plus = 0., minus = 0. )
for evSet in dataSetAsymW :
    if evSet.getCatIndex('asymCat') == 1 :
        sumW['plus'] += dataSetAsymW.weight()