Python RooDataSet.add方法代码示例

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
def fillDataSet(data, x, N, dsName = 'ds'):
    cols = RooArgSet(x)
    ds = RooDataSet(dsName, dsName, cols)
    #ds.Print()
    print 'length data:', N
    for datum in range(0,N):
        if (data[datum] < x.getMax()) and (data[datum] > x.getMin()):
            x.setVal(data[datum])
            ds.add(cols)
    ds.Print()
    return ds

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet 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 RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
def roofit_poisson_unbinned(data):
    """returns lambda, error of lambda"""

    x = RooRealVar('x', 'x', 0, max(data) * 10)
    lam = RooRealVar('lambda', 'lambda', 0.1, 0.000001, max(data))

    model = RooPoisson('model', 'model', x, lam)

    dataset = RooDataSet('data', 'data', RooArgSet(x))
    for val in data:
        x.setVal(val)
        dataset.add(RooArgSet(x))

    model.fitTo(dataset, RooFit.Save(), RooFit.PrintLevel(-1))
    return lam.getVal(), lam.getError()

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [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 add [as 别名]
def fill_dataset(varargset, ftree, wt, wtvar, cut=""):
    """Return a dataset (slow, get_dataset is more efficient).

    Return a dataset from the ntuple `ftree', also apply `cut'.  Use
    `wt' as the weight expression in the tree.  `wtvar' is the
    corresponding RooRealVar weight.  Note, varargset should contain
    wtvar.

    The dataset is filled by iterating over the tree.  This is needed
    when you want to ensure different datasets have the same weight
    variable names, so that they can be combined later on.  This is
    needed even if they are combined as different categories.

    """

    from rplot.fixes import ROOT
    from ROOT import RooDataSet, RooFit, TTreeFormula
    from helpers import suppress_warnings

    suppress_warnings()
    from rplot.tselect import Tsplice

    splice = Tsplice(ftree)
    splice.make_splice("sel", cut)

    formulae = {}
    wtname = wtvar.GetName()
    for var in varargset:
        name = var.GetName()
        expr = wt if name == wtname else name
        formulae[name] = TTreeFormula(name, expr, ftree)

    dataset = RooDataSet("dataset", "Dataset", varargset, RooFit.WeightVar(wtvar))
    for i in xrange(ftree.GetEntries()):
        ftree.GetEntry(i)
        for var, expr in formulae.iteritems():
            realvar = varargset.find(var)
            realvar.setVal(expr.EvalInstance())
        dataset.add(varargset, varargset[wtname].getVal())
    return dataset

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
def dump_simple():
    # try poisson
    roo_lam = RooRealVar("lambda", ".", lam, 0.0, 1.5)
    roo_x = RooRealVar('x', 'x range', 0, bins)

    model = RooPoisson("poisson", "Poisson Model", roo_x, roo_lam)

    #data = model.generate(RooArgSet(roo_x), n_events)

    data = RooDataSet('data', 'Data', RooArgSet(roo_x))
    for val in unbinned_from_binned(xdata, ydata):
        roo_x.setVal(val)
        data.add(RooArgSet(roo_x))

    #// --- Perform extended ML fit of composite PDF to toy data ---
    fitresult = model.fitTo(data, RooFit.Save(), RooFit.PrintLevel(-1)) ;

    #// --- Plot toy data and composite PDF overlaid ---
    mesframe = roo_x.frame(bins) ;
    data.plotOn(mesframe) ;
    model.plotOn(mesframe) ;
    mesframe.Draw()

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
def main(options,args):

    inputfiles = [ 'default', 'default' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_2011A.root', 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_2011B.root' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chib_2011A_gtV21A_2082pb.root', 'file:/scratch/knuenz/Polarization/RootInput/Chib/chib_2011B_gtV21A_2613pb.root' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chicstep5_chibV1_A.root', 'file:/scratch/knuenz/Polarization/RootInput/Chib/chicstep5_chibV1_B.root' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_v13d_A_17Jan.root', 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_v13d_B_17Jan.root' ]
    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_chibV1_A_20Jan.root', 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_chibV1_B_20Jan.root' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_pi0Rej_false_A_27jan.root', 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_pi0Rej_false_B_27jan.root' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_pi0Rej_false_A_27jan.root' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_pi0Rej_false_A_30jan.root', 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_pi0Rej_false_B_30jan.root' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_pi0Rej_false_B_30jan.root' ]
#    inputfiles = [ 'file:/scratch/knuenz/Polarization/RootInput/Chib/chibstep5_chibV1_A_20Jan.root' ]

#    gROOT.ProcessLine('.L /afs/hephy.at/scratch/k/knuenz/CMSSW_4_2_4_patch2/src/CommonTools/Statistics/src/ChiSquaredProbability.cc+')
#    gROOT.ProcessLine('.L /afs/hephy.at/scratch/k/knuenz/CMSSW_4_2_4_patch2/src/CommonTools/Statistics/src/IncompleteGammaComplement.cc+')
#    from ROOT import ChiSquaredProbability
#    from ROOT import IncompleteGammaComplement
    Y1Smass0=9.46030
    Y2Smass0=10.02326
    Y3Smass0=10.3552
    Ymass_a=0.058
    Ymass_b=0.047
    Ymass_c=0.22
            
    events= Events(inputfiles)
    
    
    # candidate Chi
    chicCand_h  = Handle ("vector<reco::CompositeCandidate>")
    chicCand_l = ( "chib","chicCompCand","chibs5" )
    
    #candidate J/Psi
    jpsiCand_h  = Handle ("vector<reco::CompositeCandidate>")
    jpsiCand_l = ( "chib","jpsiCompCand" ,"chibs5")
    
    jpsiCandPAT_h  = Handle ("vector<pat::CompositeCandidate>")
    jpsiCandPAT_l = ( "onia2MuMuPatTrkTrk","" )
    
    
    #candidate gamma
    gammaCand_h  = Handle ("vector<reco::CompositeCandidate>")
    gammaCand_l = ( "chib","gammaCompCand" ,"chibs5")
    
    convCand_h = Handle ("vector<reco::Conversion>")
    convCand_l = ( "allConversions","" )
    
    
    # Create histograms, etc.
    gROOT.SetBatch()        # don't pop up canvases
    gROOT.SetStyle('Plain') # white background
    
    #fill a RooDataSet
    invm1S     = RooRealVar("invm1S",   "invm1S",9,50)
    invm2S     = RooRealVar("invm2S",   "invm2S",9,50)
    invm3S     = RooRealVar("invm3S",   "invm3S",9,50)
    jpsipt   = RooRealVar("jpsipt", "jpsipt",0,100)
    jpsimass   = RooRealVar("jpsimass", "jpsimass",8,12)
    gammapt   = RooRealVar("gammapt", "gammapt",0,100)
    deltaRChiJpsi  = RooRealVar("deltaRChiJpsi","deltaRChiJpsi",0,1)
    deltaRJpsig    = RooRealVar("deltaRJpsig","deltaRJpsig",0,1)
    jpsieta  = RooRealVar("jpsieta","jpsieta",-5,5)
    ctpv     = RooRealVar("ctpv",   "ctpv",-5,5)
    ctpverr  = RooRealVar("ctpverr","ctpverr",0,10)
    ctbs     = RooRealVar("ctbs",   "ctbs",-5,5)
    ctbserr  = RooRealVar("ctbserr","ctbserr",0,10)
    ctpvsig     = RooRealVar("ctpvsig",   "ctpvsig",0,1000)
    ctbssig     = RooRealVar("ctbssig",   "ctbssig",0,1000)
    weight   = RooRealVar("weight", "weight",0,1)
    Rconv    = RooRealVar("Rconv","Rconv",0,100)
    jpsiVprob    = RooRealVar("jpsiVprob","jpsiVprob",0,1.5)
    Y1Smass_nSigma    = RooRealVar("Y1Smass_nSigma","Y1Smass_nSigma",-100,100)
    Y2Smass_nSigma    = RooRealVar("Y2Smass_nSigma","Y2Smass_nSigma",-100,100)
    Y3Smass_nSigma    = RooRealVar("Y3Smass_nSigma","Y3Smass_nSigma",-100,100)
    jpsipx   = RooRealVar("jpsipx", "jpsipx",-100,100)
    jpsipy   = RooRealVar("jpsipy", "jpsipy",-100,100)
    jpsipz   = RooRealVar("jpsipz", "jpsipz",-100,100)
    gammapx   = RooRealVar("gammapx", "gammapx",-100,100)
    gammapy   = RooRealVar("gammapy", "gammapy",-100,100)
    gammapz   = RooRealVar("gammapz", "gammapz",-100,100)
    vertexChi2ProbGamma   = RooRealVar("vertexChi2ProbGamma", "vertexChi2ProbGamma",0,1)
    Q   = RooRealVar("Q", "Q",0,100)
    
    argSet = RooArgSet()
    argSet.add(invm1S)
    argSet.add(invm2S)
    argSet.add(invm3S)
    argSet.add(jpsipt)
    argSet.add(jpsieta)
    argSet.add(jpsimass)
    argSet.add(gammapt)
    argSet.add(deltaRJpsig)
    argSet.add(deltaRChiJpsi)
    argSet.add(ctpv)
    argSet.add(ctpverr)
    argSet.add(ctpvsig)
    argSet.add(ctbs)
    argSet.add(ctbserr)
    argSet.add(ctbssig)
    argSet.add(Rconv)
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
class DilutionToy(Toy):
    def __init__(self):
        Toy.__init__(self)
        self._gen_params = []

    def run(self, **kwargs):
        from ROOT import RooArgSet

        __check_req_kw__("Observables", kwargs)
        __check_req_kw__("Pdf", kwargs)
        __check_req_kw__("Sigmat", kwargs)
        __check_req_kw__("Time", kwargs)
        __check_req_kw__("SigmaGen", kwargs)
        sigma_gen = kwargs.pop("SigmaGen")
        observables = kwargs.pop("Observables")
        obs_set = RooArgSet(*observables)

        pdf = kwargs.pop("Pdf")
        sigmat = kwargs.pop("Sigmat")
        time = kwargs.pop("Time")

        gen_obs_set = RooArgSet(*observables)

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

        from P2VV.RooFitWrappers import RealVar

        da = RealVar("da", Observable=True, MinMax=(0.01, 1.1))
        dft = RealVar("dft", Observable=True, MinMax=(0.01, 1.1))
        result_params.add(da._target_())
        result_params.add(dft._target_())

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

        # Some extra numbers of interest
        from ROOT import RooRealVar

        seed = RooRealVar("seed", "random seed", 0.0)
        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

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

        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)

            data = pdf.generate(spec)
            if self.transform():
                old_data = data
                data = self.transform()(old_data)
                if not data:
                    transform.set_params(data_params)
                    self._data.add(data_params)
                    continue

            from P2VV import Dilution

            d_ft = Dilution.dilution_ft(data, time, t_range=2, quiet=True)
            d_a = Dilution.signal_dilution_dg(data, sigmat, *sigma_gen)
            da.setVal(d_a[0])
            da.setError(d_a[1] if d_a[1] != None else 0.0)
            dft.setVal(d_ft[0])
            dft.setError(d_ft[1] if d_ft[1] != None else 0.0)

            if transform:
                transform.set_params(data_params)

            self._data.add(result_params)

        return self.data()

    def gen_params(self):
        return self._gen_params

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
def makeRooDataSet(type,infile_name,outfile_name,tree_name,nevents):
    """ Make RooDataSets from TTrees"""




    inputfile = TFile.Open(infile_name,"READ")
    print "Importing tree"
    tree = TTree()
    inputfile.GetObject(tree_name, tree)  #get the tree from the data file

   

    #define variables for the RooDataSet
    m_mumu        = RooRealVar("m_mumu",   "m_mumu",   0.0, 4.0)
    y_mumu        = RooRealVar("y_mumu",   "y_mumu",   0.0, 2.0 )
    pt_mumu       = RooRealVar("pt_mumu",  "pt_mumu",  0.0, 260.0)
    eta_gamma     = RooRealVar("eta_gamma","eta_gamma",-3.5, 3.5)
    pt_gamma      = RooRealVar("pt_gamma", "pt_gamma", 0.0, 100.0)
    m_gamma       = RooRealVar("m_gamma",  "m_gamma",  -0.1,0.1)
  
    m_chi_rf1S    = RooRealVar("m_chi_rf1S", "m_chi_rf1S", 0.0, 7.0)
    m_chi_rf2S    = RooRealVar("m_chi_rf2S", "m_chi_rf2S", -1.0, 1.0)
    Qvalue        = RooRealVar("Qvalue","Q", -15., 15.) 

    ctpv          = RooRealVar("ctpv","ctpv", -1.0, 3.5)
    ctpv_error    = RooRealVar("ctpv_err","ctpv_err", -1.0, 1.0)
    pi0_abs_mass  = RooRealVar("pi0_abs_mass","pi0_abs_mass", 0.0, 2.2)
    psi1S_nsigma  = RooRealVar("psi1S_nsigma","psi1S_nsigma",0.0,1.0)
    psi2S_nsigma  = RooRealVar("psi2S_nsigma","psi2S_nsigma",0.0,1.0)
    psi3S_nsigma  = RooRealVar("psi3S_nsigma","psi3S_nsigma",0.0,1.0)
    rho_conv      = RooRealVar("rho_conv", "rho_conv", 0.0, 70.0)
    dz            = RooRealVar("dz","dz", -1.0, 1.0)

    probFit1S     = RooRealVar('probFit1S','probFit1S',0,1)
    probFit2S     = RooRealVar('probFit2S','probFit2S',0,1)
    probFit3S     = RooRealVar('probFit3S','probFit3S',0,1)

    dataArgSet = RooArgSet(m_mumu,
                           y_mumu,
                           pt_mumu,
                           eta_gamma,
                           pt_gamma,
                           m_gamma, 
                           m_chi_rf1S)
    
    dataArgSet.add( m_chi_rf2S )
    dataArgSet.add( Qvalue )
    dataArgSet.add( ctpv )
    dataArgSet.add( ctpv_error )
    dataArgSet.add( pi0_abs_mass )
    dataArgSet.add( psi1S_nsigma )
    dataArgSet.add( psi2S_nsigma )
    dataArgSet.add( psi3S_nsigma )
    dataArgSet.add( rho_conv )
    dataArgSet.add( dz )
    dataArgSet.add( probFit1S)
    dataArgSet.add( probFit2S)
    dataArgSet.add( probFit3S)

 

    print "Creating DataSet"
    dataSet = RooDataSet("chicds","Chic RooDataSet", dataArgSet)

    entries = tree.GetEntries()
    print entries

    if nevents is not 0:
        entries = nevents

    for ientry in range(0,entries):
        tree.GetEntry(ientry)

        # unfort ntuples are slightly different for chic and chib

        if applyscale:
            
            if usekinfit :    
                spatial = tree.rf3S_photon_p4.Vect()
                #spatial = tree.photon_p4.Vect()
                spatial *= (1/escale)
                corr_photon_p4=TLorentzVector()
                #corr_photon_p4.SetVectM(spatial,tree.rf1S_photon_p4.M())
                corr_photon_p4.SetVectM(spatial,0)
                corr_chi_p4 = tree.rf3S_dimuon_p4 + corr_photon_p4

            else:
                        
                spatial = tree.photon_p4.Vect()
                spatial *= (1/escale)
                corr_photon_p4=TLorentzVector()
                corr_photon_p4.SetVectM(spatial,tree.photon_p4.M())
                corr_chi_p4 = tree.dimuon_p4 + corr_photon_p4
                
            
        else :
            corr_chi_p4 = tree.chi_p4
        
        if type == 'chic':
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
def process(file_name, tree_name, max_events, n_bins, x_min, x_max, fit, background, peak_x_min, peak_x_max, draw_legend, plot_q_square, plot_momentum_resolution, mc_tree_name, verbose):
    """
        A function that forms the main logic of the script

        Args:
        file_name (str): the name of the file to process
        tree_name (str): the name of the tree to process
        n_bins (int): the number of bins to be used in the histogram
        x_min (float): the left bound of the histogram
        x_max (float): the right bound of the histogram
        fit (bool): the flag that determines whether the data will be fitted
        background (bool): the flag that determines whether signal or background b_mass_data is processed
        peak_x_min (float): the left bound of the peak
        peak_x_max (float): the right bound of the peak
        draw_legend (bool): the flag that determines whether the histogram legend will be drawn
        plot_q_square (bool): the flag that determines whether the q^2 distribution will be plotted
        plot_momentum_resolution (bool): the flag that determines whether the tau and neutrino momentum resolution distributions will be plotted
        max_events (int): the maximum number of events that will be processed
        verbose (bool): the flag that switches inreased verbosity
    """

    start_time = time.time()
    last_timestamp = time.time()

    # Opening the file and getting the branch
    input_file = TFile(file_name, 'read')
    event_tree = input_file.Get(tree_name)

    # Event counters
    processed_events = 0 # Number of processed events
    reconstructable_events = 0 # Events with valid tau+ and tau- decay vertex

    # Variables for RooFit
    b_mass = RooRealVar('mB', 'm_{B}', x_min, x_max)
    b_mass_data = RooDataSet('mB', 'm_{B} data', RooArgSet(b_mass)) # Storage for reconstructed B mass values

    if plot_q_square:
        q_square = RooRealVar('q2', 'q^{2}', 12.5, 22.5)
        q_square_data = RooDataSet('q2_data', 'q^{2} data', RooArgSet(q_square)) # q^2 values container

    if plot_momentum_resolution:
        error_p_tauplus_x = RooRealVar('error_p_tauplus_x', '#epsilon_{p_{#tau^{+}x}}', -2., 2.)
        error_p_tauplus_x_data = RooDataSet('error_p_tauplus_x_data', '#epsilon_{p_{#tau^{+}x}} data', RooArgSet(error_p_tauplus_x))
        error_p_tauplus_y = RooRealVar('error_p_tauplus_y', '#epsilon_{p_{#tau^{+}y}}', -2., 2.)
        error_p_tauplus_y_data = RooDataSet('error_p_tauplus_y_data', '#epsilon_{p_{#tau^{+}y}} data', RooArgSet(error_p_tauplus_y))
        error_p_tauplus_z = RooRealVar('error_p_tauplus_z', '#epsilon_{p_{#tau^{+}z}}', -2., 2.)
        error_p_tauplus_z_data = RooDataSet('error_p_tauplus_z_data', '#epsilon_{p_{#tau^{+}z}} data', RooArgSet(error_p_tauplus_z))
        error_p_tauminus_x = RooRealVar('error_p_tauminus_x', '#epsilon_{p_{#tau^{-}x}}', -2., 2.)
        error_p_tauminus_x_data = RooDataSet('error_p_tauminus_x_data', '#epsilon_{p_{#tau^{-}x}} data', RooArgSet(error_p_tauminus_x))
        error_p_tauminus_y = RooRealVar('error_p_tauminus_y', '#epsilon_{p_{#tau^{-}y}}', -2., 2.)
        error_p_tauminus_y_data = RooDataSet('error_p_tauminus_y_data', '#epsilon_{p_{#tau^{-}y}} data', RooArgSet(error_p_tauminus_y))
        error_p_tauminus_z = RooRealVar('error_p_tauminus_z', '#epsilon_{p_{#tau^{-}z}}', -2., 2.)
        error_p_tauminus_z_data = RooDataSet('error_p_tauminus_z_data', '#epsilon_{p_{#tau^{-}z}} data', RooArgSet(error_p_tauminus_z))
        error_p_nu_tauplus_x = RooRealVar('error_p_nu_tauplus_x', '#epsilon_{p_{#nu#tau^{+}x}}', -5., 5.)
        error_p_nu_tauplus_x_data = RooDataSet('error_p_nu_tauplus_x_data', '#epsilon_{p_{#nu#tau^{+}x}} data', RooArgSet(error_p_nu_tauplus_x))
        error_p_nu_tauplus_y = RooRealVar('error_p_nu_tauplus_y', '#epsilon_{p_{#nu#tau^{+}y}}', -5., 5.)
        error_p_nu_tauplus_y_data = RooDataSet('error_p_nu_tauplus_y_data', '#epsilon_{p_{#nu#tau^{+}y}} data', RooArgSet(error_p_nu_tauplus_y))
        error_p_nu_tauplus_z = RooRealVar('error_p_nu_tauplus_z', '#epsilon_{p_{#nu#tau^{+}z}}', -5., 5.)
        error_p_nu_tauplus_z_data = RooDataSet('error_p_nu_tauplus_z_data', '#epsilon_{p_{#nu#tau^{+}z}} data', RooArgSet(error_p_nu_tauplus_z))
        error_p_nu_tauminus_x = RooRealVar('error_p_nu_tauminus_x', '#epsilon_{p_{#nu#tau^{-}x}}', -5., 5.)
        error_p_nu_tauminus_x_data = RooDataSet('error_p_nu_tauminus_x_data', '#epsilon_{p_{#nu#tau^{-}x}} data', RooArgSet(error_p_nu_tauminus_x))
        error_p_nu_tauminus_y = RooRealVar('error_p_nu_tauminus_y', '#epsilon_{p_{#nu#tau^{-}y}}', -5., 5.)
        error_p_nu_tauminus_y_data = RooDataSet('error_p_nu_tauminus_y_data', '#epsilon_{p_{#nu#tau^{-}y}} data', RooArgSet(error_p_nu_tauminus_y))
        error_p_nu_tauminus_z = RooRealVar('error_p_nu_tauminus_z', '#epsilon_{p_{#nu#tau^{-}z}}', -5., 5.)
        error_p_nu_tauminus_z_data = RooDataSet('error_p_nu_tauminus_z_data', '#epsilon_{p_{#nu#tau^{-}z}} data', RooArgSet(error_p_nu_tauminus_z))

    # Loop through the events
    # for counter, (event, mc_event) in enumerate(zip(event_tree, mc_event_tree)): # this finest construction doesn't work for some reason
    for counter in xrange(event_tree.GetEntries()): # so we have to use an old one
        if counter < max_events:
            event_tree.GetEntry(counter)

            if plot_momentum_resolution:
                mc_event_tree.GetEntry(counter)

            processed_events += 1
            if (counter + 1) % 100 == 0: # print status message every 100 events
                print('Processing event {} ({:.1f} events / s)'.format(counter + 1, 100. / (time.time() - last_timestamp)))
                last_timestamp = time.time()

            try:
                rec_ev = reconstruct(event_tree, verbose)
                reconstructable_events += 1

                b_mass.setVal(rec_ev.m_b)
                b_mass_data.add(RooArgSet(b_mass))

                if plot_q_square:
                    q_square.setVal(rec_ev.q_square())
                    q_square_data.add(RooArgSet(q_square))

                if plot_momentum_resolution:
                    error_p_tauplus_x.setVal((rec_ev.p_tauplus.px - mc_event_tree.tauplus_px) / mc_event_tree.tauplus_px)
                    error_p_tauplus_x_data.add(RooArgSet(error_p_tauplus_x))
                    error_p_tauplus_y.setVal((rec_ev.p_tauplus.py - mc_event_tree.tauplus_py) / mc_event_tree.tauplus_py)
                    error_p_tauplus_y_data.add(RooArgSet(error_p_tauplus_y))
                    error_p_tauplus_z.setVal((rec_ev.p_tauplus.pz - mc_event_tree.tauplus_pz) / mc_event_tree.tauplus_pz)
                    error_p_tauplus_z_data.add(RooArgSet(error_p_tauplus_z))

                    error_p_tauminus_x.setVal((rec_ev.p_tauminus.px - mc_event_tree.tauminus_px) / mc_event_tree.tauminus_px)
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
class FitToy(Toy):
    def __init__(self):
        Toy.__init__(self)

    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
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
for i in range(n_toys % n_p):
    n_t[i] += 1

for n in n_t:
    c = Calculator(pdf, sigmat_cat, t, st, n)
    calculators.append(c)
    c.start()

args = RooArgSet(da, dft)
while len(calculators):
    for i, calculator in enumerate(calculators):
        msg = calculator.queue().get()
        if msg == 'done':
            calculator.join()
            calculators.pop(i)
            continue
        else:
            d_a, d_ft = msg
        da.setVal(d_a[0])
        dft.setVal(d_ft[0])
        dft.setError(d_ft[1])
        test_data.add(args)
    if test_data.numEntries() % 100 == 0:
        print 'completed ', test_data.numEntries()
        
diff = FormulaVar(Name = 'diff', Formula = '@0 - @1', Arguments = (dft, da), data = test_data)
from ROOT import TFile
f = TFile("dilution.root", "recreate")
f.WriteTObject(test_data, "data")
f.Close()

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
def process(file_name, max_events, n_bins, x_min, x_max, fit, peak_x_min, peak_x_max, draw_legend, verbose):
    """
        A function that forms the main logic of the script

        Args:
        file_name (str): the name of the file to process
        max_events (int): the maximum number of events that will be processed
        n_bins (int): the number of bins to be used in the histogram
        x_min (float): the left bound of the histogram
        x_max (float): the right bound of the histogram
        fit (bool): the flag that determines whether the data will be fitted
        peak_x_min (float): the left bound of the peak
        peak_x_max (float): the right bound of the peak
        draw_legend (bool): the flag that determines whether the histogram legend will be drawn
        verbose (bool): the flag that switches increased verbosity
    """

    start_time = time.time()
    last_timestamp = time.time()

    # Opening the file and getting the branch
    input_file = TFile(file_name, 'read')
    input_tree = input_file.Get('Events')

    # Event counters
    processed_events = 0 # Number of processed events
    reconstructable_events = 0 # Events with valid tau+ and tau- decay vertex

    # Variables for RooFit
    b_mass = RooRealVar('mB', 'm_{B}', x_min, x_max)
    b_mass_data = RooDataSet('mB', 'm_{B} data', RooArgSet(b_mass)) # reconstructed B mass values container
    q_square = RooRealVar('q2', 'q^{2}', 12.5, 17.5)
    q_square_data = RooDataSet('q2_data', 'q^{2} data', RooArgSet(q_square)) # q^2 values container

    # Loop through the events
    for counter, event in enumerate(input_tree):
        if counter < max_events:
            processed_events += 1
            if (counter + 1) % 100 == 0: # print status message every 100 events
                print('Processing event {:.1f} ({:.1f} events / s)'.format(counter + 1, 100. / (time.time() - last_timestamp)))
                last_timestamp = time.time()

            try:
                rec_ev = reconstruct(event, verbose)
                reconstructable_events += 1

                b_mass.setVal(rec_ev.m_b)
                b_mass_data.add(RooArgSet(b_mass))

                q_square.setVal(rec_ev.q_square())
                q_square_data.add(RooArgSet(q_square))
            except UnreconstructableEventError:
                pass

    end_time = time.time()

    # printing some useful statistics
    print('{} events have been processed'.format(processed_events))
    print('Elapsed time: {:.1f} s ({:.1f} events / s)'.format(end_time - start_time, float(processed_events) / (end_time - start_time)))
    print('Reconstruction efficiency: {} / {} = {:.3f}'.format(reconstructable_events, processed_events, float(reconstructable_events) / processed_events))

    if fit:
        # signal model
        # ILD-like
        signal_model = SignalModel(name = 'signal_model',
                                   title = 'Signal Model',
                                   x = b_mass,
                                   mean = RooRealVar('mean_signal', '#mu_{signal}', 5.279, peak_x_min, peak_x_max),
                                   width = RooRealVar('width_signal', '#sigma_{signal}', 0.03, 0.01, 0.1),
                                   width_wide = RooRealVar('width_wide_gauss', '#sigma_{wide}', 0.165),
                                   alpha = RooRealVar('alpha_signal', '#alpha_{signal}', -0.206),
                                   n = RooRealVar('n_signal', 'n_{signal}', 2.056),
                                   narrow_gauss_fraction = RooRealVar('narrow_gauss_fraction', 'Fraction of narrow Gaussian in signal', 0.127),
                                   cb_fraction = RooRealVar('signal_cb_fraction', 'Fraction of CB in signal', 0.5)
                                   )
        # progressive
        # signal_model = SignalModel(name = 'signal_model',
        #                            title = 'Signal Model',
        #                            x = b_mass,
        #                            mean = RooRealVar('mean_signal', '#mu_{signal}', 5.279, peak_x_min, peak_x_max),
        #                            width = RooRealVar('width_signal', '#sigma_{signal}', 0.03, 0.01, 0.1),
        #                            width_wide = RooRealVar('width_wide_gauss', '#sigma_{wide}', 0.151),
        #                            alpha = RooRealVar('alpha_signal', '#alpha_{signal}', -0.133),
        #                            n = RooRealVar('n_signal', 'n_{signal}', 2.891),
        #                            narrow_gauss_fraction = RooRealVar('narrow_gauss_fraction', 'Fraction of narrow Gaussian in signal', 0.301),
        #                            cb_fraction = RooRealVar('signal_cb_fraction', 'Fraction of CB in signal', 0.330)
        #                            )

        # Bs -> Ds Ds K* (with Ds -> tau nu) background model
        # ILD-like
        bs_ds2taunu_model = BackgroundModel(name = 'bs_ds2taunu_model',
                                            title = 'Bs (with Ds -> #tau #nu) background model',
                                            x = b_mass,
                                            mean = RooRealVar('mean_bs_ds2taunu', '#mu_{Bs (with Ds -> #tau #nu)}', 4.970),
                                            width_gauss = RooRealVar('width_bs_ds2taunu_gauss', '#sigma_{Bs (with Ds -> #tau #nu) Gauss}', 0.196),
                                            width_cb = RooRealVar('width_bs_ds2taunu_cb', '#sigma_{Bs (with Ds -> #tau #nu) CB}', 0.078),
                                            alpha = RooRealVar('alpha_bs_ds2taunu', '#alpha_{Bs (with Ds -> #tau #nu)}', -0.746),
                                            n = RooRealVar('n_bs_ds2taunu', 'n_{Bs (with Ds -> #tau #nu)}', 1.983),
                                            gauss_fraction = RooRealVar('bs_ds2taunu_gauss_fraction', 'Fraction of Gaussian in Bs (with Ds -> #tau #nu) background', 0.243)
                                            )
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
def gen_data_and_fit(ws, iterations,cat, mass,channel,turnon,truth):
    
    if channel in channels:        

        #fit gaus(x)bern to sigm(x)pow and sigm(x)exp
        #fit sigm(x)bern to erf(x)pow and erf(x)exp
        
        n_sample = int(ws.data('bkgdata_%s_%i'%(channel,cat)).sumEntries())
        if turnon == 'erf' and truth == 'exp':
            ws.factory('pull_ROI_erf_on_sigmexp_%s_cat%i[0]'%(channel,cat))
            ws.defineSet('biasVars_%s_cat%i'%(channel,cat),
                         'pull_ROI_erf_on_sigmexp_%s_cat%i,'%(channel,cat))
        if turnon == 'erf' and truth == 'pow':
            ws.factory('pull_ROI_erf_on_sigmpow_%s_cat%i[0]'%(channel,cat))
            ws.defineSet('biasVars_%s_cat%i'%(channel,cat),
                         'pull_ROI_erf_on_sigmpow_%s_cat%i,'%(channel,cat)
                         )
            
        if turnon == 'sigm' and truth == 'exp':        
            ws.factory('pull_ROI_sigm_on_erfexp_%s_cat%i[0]'%(channel,cat))
            ws.defineSet('biasVars_%s_cat%i'%(channel,cat),                         
                         'pull_ROI_sigm_on_erfexp_%s_cat%i,'%(channel,cat)
                         )
        if turnon == 'sigm' and truth == 'pow':
            ws.factory('pull_ROI_sigm_on_erfpow_%s_cat%i[0]'%(channel,cat))
            ws.defineSet('biasVars_%s_cat%i'%(channel,cat), 
                         'pull_ROI_sigm_on_erfpow_%s_cat%i'%(channel,cat)
                         )
        biasData = RooDataSet('biasData_%s_cat%i'%(channel,cat),
                              'bias data',
                              ws.set('biasVars_%s_cat%i'%(channel,cat)))
        
        for i in xrange(iterations):

            ### SIGM            
            #build ERF toy data to fit with SIGM
            if turnon == 'sigm' and truth == 'exp':
                truth_exp_erf = ws.pdf('MzgTruthModel_exp_erf_%s_cat%i'%(channel,cat))
                toy_data_exp_erf = truth_exp_erf.generate(                    
                    RooArgSet(ws.var('Mzg')),
                    n_sample,
                    RooFit.Extended(),
                    RooFit.Name('toy_erfexp_%s_cat%i_%i'%(channel,cat,i))
                    )
                true_ROI_yield_erfexp = ws.var('norm_truth_exp_%s_cat%i'%(channel,cat)).getVal()

                #fit logistics(x)bern to erfexp
                ws.var('norm_altrsb_cat%i'%cat).setMin(true_ROI_yield_erfexp*0.25)
                ws.var('norm_altrsb_cat%i'%cat).setMax(true_ROI_yield_erfexp*1.75)
                ws.var('norm_altrsb_cat%i'%cat).setVal(true_ROI_yield_erfexp)            
                
                minos_var = RooArgSet(ws.var('norm_altrsb_cat%i'%cat))
                sigm_nll = ws.pdf('RSBFitModelAlt_cat%i'%cat).createNLL(
                    toy_data_exp_erf
                    )
                sigm_min = RooMinimizer(sigm_nll)
                sigm_min.setPrintLevel(1)
                sigm_min.minimize('Minuit2','scan')
                sigm_min.minimize('Minuit2','simplex')                
                migrad_out = sigm_min.migrad()                
                migrad_out = sigm_min.migrad()
                hesse_out  = sigm_min.hesse()
                
                fit_sigm_norm = ws.var('norm_altrsb_cat%i'%cat).getVal()
                fit_sigm_err  = ws.var('norm_altrsb_cat%i'%cat).getError()
                fit_sigm_pull = (fit_sigm_norm - true_ROI_yield_erfexp)/fit_sigm_err

                #if fit failed run minos
                if migrad_out != 0 or hesse_out != 0:
                    migrad_out = sigm_min.migrad()                    
                    if migrad_out != 0:
                        sigm_min.minos(minos_var)
                        fit_sigm_err  = max(abs(ws.var('norm_altrsb_cat%i'%cat).getErrorHi()),
                                            abs(ws.var('norm_altrsb_cat%i'%cat).getErrorLo()))
                    else:
                        fit_sigm_err  = ws.var('norm_altrsb_cat%i'%cat).getError()
                    fit_sigm_norm = ws.var('norm_altrsb_cat%i'%cat).getVal()
                    fit_sigm_pull = (fit_sigm_norm - true_ROI_yield_erfexp)/fit_sigm_err
                
                print i, fit_sigm_norm, fit_sigm_err, true_ROI_yield_erfexp, fit_sigm_pull
                
                ws.var('pull_ROI_sigm_on_erfexp_%s_cat%i'%(channel,cat)).setVal(
                    fit_sigm_pull
                    )                

                biasData.add(ws.set('biasVars_%s_cat%i'%(channel,cat)))

                var = ws.var('pull_ROI_sigm_on_erfexp_%s_cat%i'%(channel,cat))

                print 'cumulative median bias: %.3f'%calculate_median(var,biasData)
                
                var = None                

                del sigm_min
                del sigm_nll
                del truth_exp_erf

            if turnon =='sigm' and truth =='pow':
                truth_pow_erf = ws.pdf('MzgTruthModel_pow_erf_%s_cat%i'%(channel,cat))                
                toy_data_pow_erf = truth_pow_erf.generate(                    
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooDataSet [as 别名]
# 或者: from ROOT.RooDataSet import add [as 别名]
			if k.pt < 250: continue
			counters[0] = counters[0] + 1
			if mu1.pt < 550 or mu2.pt < 550: continue
			counters[1] = counters[1] + 1
			if mu1.ipchi2 < 16 or mu2.ipchi2 < 16 or k.ipchi2 < 16: continue
			counters[2] = counters[2] + 1
		        if psi_dls[jpsi_idx] < 5: continue
			counters[3] = counters[3] + 1
			if abs(jpsimass - jpsi.m) > 50: continue
			counters[4] = counters[4] + 1
                        if k.pnn_k < 0.9: continue
                        counters[5] = counters[5] + 1
			hist.Fill(bp0_m[j])
			if ( bp0_m[j] > m.getMin() and bp0_m[j] < m.getMax() ):
				m.setVal(bp0_m[j])
				ds.add(argset)

	aa.SetBranchStatus("bp1*",1)
	aa.SetBranchStatus("piz*",1)
	aa.SetBranchStatus("bp0*",0)

	for i in range(aa.GetEntries()):
		s = aa.GetEntry(i)
		if s == -1:
			print "Error reading event ",i,", skipping"
			continue
		if i%10000 == 0:
			print i," entry of ",aa.GetEntries()
			print counters2
		bp1_m = aa.bp1_dtf_m
		bp1_idx_psi  = aa.bp1_idx_psi