Beispiel #1
0
def createTrees(hist_cfg, out_dir, verbose=False, friend_func=None):
    '''Writes out TTrees from histogram configuration for each contribution. 
    Takes list of variables attached to histogram config (hist_cfg.vars) to 
    create branches.
    '''
    plot = DataMCPlot(hist_cfg.name) # Used to cache TTrees
    vcfgs = hist_cfg.vars
    for cfg in hist_cfg.cfgs:
        
        out_file = TFile('/'.join([out_dir, hist_cfg.name + '_' + cfg.name + '.root']), 'RECREATE')
        out_tree = TTree('tree', '')

        # Create branches for all variables
        branches = [array('f', [0.]) for i in xrange(len(vcfgs))]
        for branch_name, branch in zip([v.name for v in vcfgs], branches):
            out_tree.Branch(branch_name, branch, branch_name+'/F')

        # Create branch with full weight including lumi x cross section
        full_weight = array('f', [0.])
        out_tree.Branch('full_weight', full_weight, 'full_weight/F')
        branches.append(full_weight)

        total_scale = hist_cfg.total_scale if hist_cfg.total_scale else 1.
        fillIntoTree(out_tree, branches, cfg, hist_cfg, vcfgs, total_scale, plot, verbose, friend_func)

        out_file.cd()
        out_tree.Write()
        out_file.Write()
        out_file.Close()
    return plot
Beispiel #2
0
def staff():

    staff = ROOT.staff_t()

    # The input file cern.dat is a copy of the CERN staff data base
    # from 1988

    f = TFile('staff.root', 'RECREATE')
    tree = TTree('T', 'staff data from ascii file')
    tree.Branch('staff', staff,
                'Category/I:Flag:Age:Service:Children:Grade:Step:Hrweek:Cost')
    tree.Branch('Divisions', AddressOf(staff, 'Division'), 'Division/C')
    tree.Branch('Nation', AddressOf(staff, 'Nation'), 'Nation/C')

    # note that the branches Division and Nation cannot be on the first branch
    fname = os.path.expandvars('$ROOTSYS/tutorials/tree/cernstaff.dat')
    for line in open(fname).readlines():
        t = filter(lambda x: x, re.split('\s+', line))
        staff.Category = int(t[0])  # assign as integers
        staff.Flag = int(t[1])
        staff.Age = int(t[2])
        staff.Service = int(t[3])
        staff.Children = int(t[4])
        staff.Grade = int(t[5])
        staff.Step = int(t[6])
        staff.Hrweek = int(t[7])
        staff.Cost = int(t[8])
        staff.Division = t[9]  # assign as strings
        staff.Nation = t[10]

        tree.Fill()

    tree.Print()
    tree.Write()
Beispiel #3
0
def writeTree():
    print ">>> Writing a tree."

    # make new root file with new tree
    file = TFile("tree.root", 'recreate')
    tree = TTree("tree_name", "tree title")

    # create 1 dimensional float arrays as fill variables, in this way the float array serves
    # as a pointer which can be passed to the branch
    n = array('d', [0])
    u = array('d', [0])

    # using numpy instead of array class (note python's float datatype corresponds to C++ doubles):
    #n = numpy.zeros(1, dtype=float)
    #u = numpy.zeros(1, dtype=float)

    # create the branches and assign the fill-variables to them as doubles (D)
    tree.Branch("normal", n, 'normal/D')
    tree.Branch("uniform", u, 'uniform/D')

    # create some random numbers, assign them into the fill variables and call Fill()
    for i in xrange(10000):
        n[0] = gRandom.Gaus()
        u[0] = gRandom.Uniform()
        tree.Fill()

    # write the tree into the output file and close the file
    file.Write()
    file.Close()
    def test05WriteSomeDataObjectBranched(self):
        """Test writing of a complex object across different branches"""

        f = TFile(self.fname, 'RECREATE')
        t = TTree(self.tname, self.ttitle)

        d = SomeDataStruct()

        # note: for p2.2, which has incomplete support of property types,
        # we need to keep a reference alive to the result of the property
        # call, or it will be deleted too soon; for later pythons, it is
        # safe to use d.Floats directly in the Branch() call
        fl = d.Floats
        t.Branch('floats', fl)
        t.Branch('nlabel', AddressOf(d, 'NLabel'), 'NLabel/I')
        t.Branch('label', AddressOf(d, 'Label'), 'Label/C')

        for i in range(self.N):
            for j in range(self.M):
                d.Floats.push_back(i * self.M + j)

            d.NLabel = i
            d.Label = '%d' % i

            t.Fill()

        f.Write()
        f.Close()
Beispiel #5
0
def WriteTreeForEachScan():
    filenamedf=pd.read_csv(dfdir+'filename.csv')
    normvalue=np.loadtxt(dfdir+'normvalueinde.txt')

    wavech = array('f',720*[0.])
    intch = array('f',12*[0.])
    fileid = array('i',[0])

    for scanid in range(1,NbofscanID-1):
        relatedpixels=longimap[:,scanid-1:scanid+2].reshape(12)
        relatednorm=normvalue[:,scanid-1:scanid+2].reshape(12)
        f = TFile(rootdir+"lwaveform"+str(scanid)+".root", 'recreate' )
        tree = TTree( 'wavetree', 'wavetree' )
        tree.Branch('wavech',wavech,'wavech[720]/F')
        tree.Branch('intch',intch,'intch[12]/F')
        tree.Branch('fileid',fileid,'fileid/I')
        for fiid in range(0,33):
            print(scanid,fiid)
            fileid[0]=fiid
            wavelist,intlist=ReadSingleRawText(filenamedf['scanstart'+str(scanID[scanid])].iloc[fiid],relatedpixels,relatednorm)
            for itr in range(len(intlist)):
                for k in range(12):
                    intch[k]=intlist[itr][k]
                for k in range(720):
                    wavech[k]=wavelist[itr][k]
                tree.Fill()
                
        f.Write()
        f.Close()
def WriteTreeForEachScan3by4():
    Nbofevents = 10000
    filenamedf = pd.read_csv(dfdir + 'filename.csv')
    normvalue = np.loadtxt(dfdir + 'normvalueinde.txt')
    filerange = [k for k in range(0, 33)]
    Nbofwavepoint = 27
    intch = array('f', 12 * [0.])
    fileid = array('i', [0])

    for scanid in range(1, NbofscanID - 1):
        f = TFile(outdatadir + "lwaveform" + str(scanid) + "cutLED3by4.root",
                  'recreate')
        tree = TTree('wavetree', 'wavetree')
        # tree.Branch('wavech',wavech,'wavech[243]/F')
        tree.Branch('intch', intch, 'intch[12]/F')
        tree.Branch('fileid', fileid, 'fileid/I')
        relatedpixels = longimap[:, scanid - 1:scanid + 2].reshape(12)
        relatednorm = normvalue[:, scanid - 1:scanid + 2].reshape(12)
        for fiid in filerange:
            print(scanid, fiid)
            fileid[0] = fiid
            wavelist, intlist = ReadSingleRawText(
                filenamedf['scanstart' + str(scanID[scanid])].iloc[fiid],
                relatedpixels, relatednorm, Nbofwavepoint, Nbofevents)
            for itr in range(len(intlist)):
                for k in range(12):
                    intch[k] = intlist[itr][k]
                # for k in range(243):
                #     wavech[k]=wavelist[itr][k]
                tree.Fill()
        f.Write()
        f.Close()
Beispiel #7
0
def prep_AdSimple(host_file):
    '''
    Return a tuple of (adSimple, buffer) containing the TTree object and the
    buffer used to fill its TBranches.

    '''
    buf = TreeBuffer()
    buf.triggerType = unsigned_int_value()
    buf.site = int_value()
    buf.energy = float_value()
    buf.x = float_value()
    buf.y = float_value()
    buf.z = float_value()

    host_file.cd()
    event_subdir = host_file.Get('Event')
    if not bool(event_subdir):
        event_subdir = host_file.mkdir('Event')
    event_subdir.cd()
    rec_subdir = event_subdir.Get('Rec')
    if not bool(rec_subdir):
        rec_subdir = event_subdir.mkdir('Rec')
    rec_subdir.cd()
    adSimple = TTree('AdSimple', 'Tree at /Event/Rec/AdSimple holding '
            'Rec_AdSimple')
    adSimple.Branch('context.mSite', buf.site, 'context.mSite/I')
    adSimple.Branch('triggerType', buf.triggerType, 'triggerType/i')
    adSimple.Branch('energy', buf.energy, 'energy/F')
    adSimple.Branch('x', buf.x, 'x/F')
    adSimple.Branch('y', buf.y, 'y/F')
    adSimple.Branch('z', buf.z, 'z/F')
    return adSimple, buf
Beispiel #8
0
def WriteTreeForEachScan3by4():
    Nbofevents = 10000
    filenamedf = pd.read_csv(dfdatadir + 'energyfilename.csv')
    normvalue = np.loadtxt(dfdatadir + 'normvalueinde.txt')
    filerange = [k for k in range(Nbofenergypos)]
    intch = array('f', 12 * [0.])
    fileid = array('i', [0])
    relatedpixels = longienergymap[1:5, 3:6].reshape(12)
    relatednorm = normvalue[1:5, 3:6].reshape(12)
    for scanid in range(Nbofenergy):
        f = TFile(rootdatadir + "lwaveform" + str(scanid) + "energy3by4.root",
                  'recreate')
        tree = TTree('wavetree', 'wavetree')
        # tree.Branch('wavech',wavech,'wavech[243]/F')
        tree.Branch('intch', intch, 'intch[12]/F')
        tree.Branch('fileid', fileid, 'fileid/I')
        for fiid in filerange:
            print(scanid, fiid)
            fileid[0] = fiid
            intarray = ReadSingleRawTxTReturnIntegral(
                filenamedf['energy' + str(scanid)].iloc[fiid],
                relatedpixels,
                relatednorm,
                Nbofevents,
                thisledch=8,
                thisbeamch=7)
            for itr in range(intarray.shape[0]):
                for k in range(12):
                    intch[k] = intarray[itr, k]
                tree.Fill()
        f.Write()
        f.Close()
Beispiel #9
0
def WriteTreeForEachScan3by6():
    Nbofevents = 1000
    filenamedf = pd.read_csv(dfdatadir + 'filename.csv')
    normvalue = np.loadtxt(dfdatadir + 'normvalueinde.txt')
    intch = array('f', 18 * [0.])
    fileid = array('i', [0])
    for scanid in range(Nbofscan):
        relatedpixels = longimap[scanid:scanid + 3, 1:7].reshape(18)
        relatednorm = normvalue[scanid:scanid + 3, 1:7].reshape(18)

        f = TFile(rootdatadir + "lwaveform" + str(scanid) + "energy3by6.root",
                  'recreate')
        tree = TTree('wavetree', 'wavetree')
        # tree.Branch('wavech',wavech,'wavech[243]/F')
        tree.Branch('intch', intch, 'intch[18]/F')
        tree.Branch('fileid', fileid, 'fileid/I')
        for fiid in range(filenamedf.shape[0]):
            print(scanid, fiid)
            fileid[0] = fiid
            intarray = ReadSingleRawTxTReturnIntegral(
                filenamedf['scanstart' +
                           str(longimap[scanid + 1, 2])].iloc[fiid],
                relatedpixels,
                relatednorm,
                Nbofevents,
                thisledch=ledch,
                thisbeamch=beamch)
            for itr in range(intarray.shape[0]):
                for k in range(18):
                    intch[k] = intarray[itr, k]
                tree.Fill()
        f.Write()
        f.Close()
Beispiel #10
0
def test_minimizer(fin_name='gen_toy_10bin_highstat', fout_name='result'):

    fout = TFile('deltas/' + fout_name + '.root', 'RECREATE')
    tout = TTree('tree', 'tree')
    bin = array('i', [0])
    n = array('f', [0])
    delta_n_u = array('f', [0.])
    delta_n_d = array('f', [0.])
    delta_m_u = array('f', [0.])
    delta_m_d = array('f', [0.])
    var_delta_m = array('f', [0.])

    tout.Branch('bin', bin, 'bin/I')
    tout.Branch('n', n, 'n/F')
    tout.Branch('delta_n_u', delta_n_u, 'delta_n_u/F')
    tout.Branch('delta_n_d', delta_n_d, 'delta_n_d/F')
    tout.Branch('delta_m_u', delta_m_u, 'delta_m_u/F')
    tout.Branch('delta_m_d', delta_m_d, 'delta_m_d/F')
    tout.Branch('var_delta_m', var_delta_m, 'var_delta_m/F')

    f = TFile.Open('deltas/' + fin_name + ".root")
    tree = f.Get("tree")
    tree.SetBranchStatus("*", True)

    for iev in range(tree.GetEntries()):
        tree.GetEntry(iev)
        ev = tree
        n_toy = ev.n_toy
        n_bins = ev.n_bin
        n_deltas = ev.n_deltas
        deltas = ev.deltas
        bin_n = ev.bin_n
        bin_u = ev.bin_u
        bin_d = ev.bin_d

        print deltas
        fix_params = []
        #fix_params.extend([1, 2+3*n_bins])
        #fix_params.extend(range(n_bins+2, 2*n_bins+2))
        minimize_deltas = Minimze_deltas(deltas=deltas,
                                         nbins=n_bins,
                                         step=0.01,
                                         fix_params=fix_params)
        res = minimize_deltas.run()
        for b in range(n_bins):
            bin[0] = b + 1
            n[0] = bin_n[b]
            delta_n_u[0] = res[0][b]
            delta_n_d[0] = res[1][b]
            delta_m_u[0] = res[2][b]
            delta_m_d[0] = res[3][b]
            var_delta_m[0] = res[4][b]
            #print "Bin %s: [%s, %s] ==> %s +/- %s" % ( b+1, delta_n_u[0], delta_n_d[0], abs(delta_m_u[0]), math.sqrt(var_delta_m[0]))
            tout.Fill()

    f.Close()
    fout.cd()
    fout.Write()
    fout.Close()
Beispiel #11
0
def test_unsupported_branch_in_branches():
    tree = TTree('test', 'test')
    vect = TLorentzVector()
    double = np.array([0], dtype=float)
    tree.Branch('vector', vect)
    tree.Branch('double', double, 'double/D')
    rnp.tree2array(tree)
    assert_raises(TypeError, rnp.tree2array, tree, branches=['vector'])
def init_tree_from_table(hit_table,
                         meta_table,
                         chunk_size=1,
                         hit_tree_entry=None,
                         meta_tree_entry=None):
    ''' Initializes a ROOT tree from a HDF5 table.
    Takes the HDF5 table column names and types and creates corresponding 
    branches. If a chunk size is specified the branches will have the length of 
    the chunk size and an additional parameter is returned to change the chunk 
    size at a later stage. If a tree_entry is defined (a ROOT c-struct) the new 
    tree has the branches set to the corresponding tree entry address.
    
    Parameters
    ----------
    numpy_type_descriptor: np.dtype
    '''
    if (chunk_size > 1 and hit_tree_entry is not None
            and meta_tree_entry is not None):
        raise NotImplementedError()

    tree = TTree('Table', 'Converted HDF5 table')
    n_entries = None
    if chunk_size > 1:
        n_entries = ctypes.c_int(chunk_size) if chunk_size > 1 else 1
        # needs to be added, otherwise one cannot access chunk_size_tree:
        tree.Branch('n_entries', ctypes.addressof(n_entries), 'n_entries/I')

    # adding branches for hit table:
    for hit_column_name in hit_table.dtype.names:
        print hit_column_name, hit_table.dtype[
            hit_column_name], get_root_type_descriptor(
                hit_table.dtype[hit_column_name])
        tree.Branch(
            hit_column_name, 'NULL' if hit_tree_entry is None else AddressOf(
                hit_tree_entry, hit_column_name),
            hit_column_name + '[n_entries]/' +
            get_root_type_descriptor(hit_table.dtype[hit_column_name])
            if chunk_size > 1 else hit_column_name + '/' +
            get_root_type_descriptor(hit_table.dtype[hit_column_name]))

    # adding branches for meta table (timestamps):
    for meta_column_name in meta_table.dtype.names:

        print meta_column_name, meta_table.dtype[
            meta_column_name], get_root_type_descriptor(
                meta_table.dtype[meta_column_name])

        if meta_column_name == 'timestamp_start' or meta_column_name == 'timestamp_stop':
            tree.Branch(
                meta_column_name,
                'NULL' if meta_tree_entry is None else AddressOf(
                    meta_tree_entry, meta_column_name),
                meta_column_name + '[n_entries]/' +
                get_root_type_descriptor(meta_table.dtype[meta_column_name])
                if chunk_size > 1 else meta_column_name + '/' +
                get_root_type_descriptor(meta_table.dtype[meta_column_name]))

    return tree, n_entries
Beispiel #13
0
def cli():
    
    helpString = "This script takes an input Sim file from Cosima and spits out a Root\
    file which can be used in Alex MGeant analysis scripts.  Takes as\
    input the Simulation File name and the Root File name."

    import argparse

    parser = argparse.ArgumentParser(description=helpString)
    parser.add_argument("SimFile", help="Input Simulation File (from Cosima)")
    parser.add_argument("RootFile", help="Output Root File")

    args = parser.parse_args()
    
    from EventViewer import parse
    from ROOT import TFile, TTree, AddressOf

    sim = parse(args.SimFile)
    print "There are {} events in this file.".format(len(sim.events))

    #for event in sim.events[:10]:
    #    print event.hits

    #Make the tree
    f = TFile(args.RootFile, 'RECREATE')
    t = TTree('h10','Simulations')

    MGS = MGSimulation()

    t.Branch('Runevt', AddressOf(MGS.MGStruct, 'Runevt'), 'Runevt/F')
    t.Branch('Xcor', AddressOf(MGS.MGStruct, 'Xcor'), 'Xcor/F')
    t.Branch('Ycor', AddressOf(MGS.MGStruct, 'Ycor'), 'Ycor/F')
    t.Branch('Zcor', AddressOf(MGS.MGStruct, 'Zcor'), 'Zcor/F')
    t.Branch('Estep', AddressOf(MGS.MGStruct, 'Estep'), 'Estep/F')
    #t.Branch('Det', AddressOf(MGS.MGStruct, 'Det'), 'Det/F')
    
    for event in sim.events:

        hits = zip(event.hits.detector,
                   event.hits.x,
                   event.hits.y,
                   event.hits.z,
                   event.hits.energy)
    
        for hit in hits:
            MGS.MGStruct.Runevt = float(event.id_trigger)
            #MGS.MGStruct.Det = hit[0]
            MGS.MGStruct.Xcor = hit[1]
            MGS.MGStruct.Ycor = hit[2]
            MGS.MGStruct.Zcor = hit[3]
            MGS.MGStruct.Estep = hit[4]
            t.Fill()

    f.Write()
    f.Close()
Beispiel #14
0
def make_prim_tree():

    #input
    #inp = TFile.Open("../../lmon.root")
    inp = TFile.Open("../../data/ew/ew1bx1.root")
    tree = inp.Get("DetectorTree")

    #number of events, negative for all
    nev = -1

    #output
    out = TFile("prim_bx1.root", "recreate")
    gROOT.ProcessLine("struct EntryF {Float_t v;};")
    x = rt.EntryF()
    y = rt.EntryF()
    z = rt.EntryF()
    en = rt.EntryF()
    otree = TTree("prim_tree", "prim_tree")
    otree.Branch("x", addressof(x, "v"), "x/F")
    otree.Branch("y", addressof(y, "v"), "y/F")
    otree.Branch("z", addressof(z, "v"), "z/F")
    otree.Branch("en", addressof(en, "v"), "en/F")

    hits = ew_hits("ew", tree)

    if nev < 0: nev = tree.GetEntries()
    for ievt in range(nev):
        tree.GetEntry(ievt)

        #print("Next event")

        for ihit in range(hits.get_n()):

            hit = hits.get_hit(ihit)

            #hit by primary photon
            if hit.prim == 0 or hit.pdg != 22: continue

            hit.global_to_zpos(-18644)  # mm

            x.v = hit.x
            y.v = hit.y
            z.v = hit.z
            z.en = hit.en

            otree.Fill()

            #print(hit.x, hit.y, hit.z) # , hit.en
            #print(hit.pdg, hit.prim, hit.conv)

            #only first hit by primary particle
            break

    otree.Write()
    out.Close()
Beispiel #15
0
def main():
    parser = argparse.ArgumentParser(
        description=
        'Program that takes as an argument the json -->.txt<-- file and makes a root file out of it, which can be used with the main framework.'
    )
    parser.add_argument('--input',
                        '-i',
                        required=True,
                        type=str,
                        help='Name of the json converted to a .txt file')
    parser.add_argument('--output',
                        '-o',
                        type=str,
                        help='Name of the target .root file')

    args = parser.parse_args()
    infilename = os.path.abspath(args.input)
    if not infilename.endswith('.txt'):
        raise ValueError('Infilename does not end with .txt')
    outfilename = os.path.abspath(
        args.output) if args.output is not None else '.root'.join(
            infilename.rsplit('.txt', 1))

    with open(infilename, 'r') as f:
        dict_in_json = safe_load(f)

    outfile = TFile(outfilename, 'RECREATE')
    outtree = TTree("LumiTree", "LumiTree")
    t_run = array('i', [0])
    t_lb_low = array('i', [0])
    t_lb_high = array('i', [0])

    outtree.Branch("run", t_run, 'run number/i')
    outtree.Branch("lb_low", t_lb_low,
                   'lumi block low number (including this one)/i')
    outtree.Branch("lb_high", t_lb_high,
                   'lumi block high number (including this one)/i')
    runs = dict_in_json.keys()
    for run in dict_in_json:
        parts = dict_in_json[run]
        # print parts
        for lumiblocks in parts:
            t_run[0] = int(run)
            t_lb_low[0] = int(lumiblocks[0])
            t_lb_high[0] = int(lumiblocks[1])
            outtree.Fill()
            # for lb in range(lumiblocks[0], lumiblocks[1]+1):
            #     t_run[0] = int(run)
            #     t_lb[0]  = int(lb)
            #     outtree.Fill()
            # print 'filled: (%i, %i)' % (int(run), int(lb))

    outfile.Write()
    outfile.Close()
    print green('--> Successfully created lumifile \'%s\'' % (outfilename))
Beispiel #16
0
def WriteTreeForEachScan():
    Nbofevents = 30000
    filenamedf = pd.read_csv(dfdatadir + 'filename.csv')
    normvalue = np.loadtxt(dfdatadir + 'normvalueinde.txt')
    intch = array('f', 9 * [0.])
    fileid = array('i', [0])

    for scanid in range(Nbofscan):
        f = TFile(rootdatadir + "lwaveform" + str(scanid) + "energy.root",
                  'recreate')
        tree1 = TTree('wavetreerow1', 'wavetree')
        # tree1.Branch('wavech',wavech,'wavech[243]/F')
        tree1.Branch('intch', intch, 'intch[9]/F')
        tree1.Branch('fileid', fileid, 'fileid/I')
        tree2 = TTree('wavetreerow2', 'wavetree')
        # tree2.Branch('wavech',wavech,'wavech[243]/F')
        tree2.Branch('intch', intch, 'intch[9]/F')
        tree2.Branch('fileid', fileid, 'fileid/I')
        tree3 = TTree('wavetreerow3', 'wavetree')
        # tree3.Branch('wavech',wavech,'wavech[243]/F')
        tree3.Branch('intch', intch, 'intch[9]/F')
        tree3.Branch('fileid', fileid, 'fileid/I')

        for fiid in range(filenamedf.shape[0]):
            if fiid < 15:
                relatedpixels = longimap[scanid:scanid + 3, 1:4].reshape(9)
                relatednorm = normvalue[scanid:scanid + 3, 1:4].reshape(9)
            elif fiid < 33:
                relatedpixels = longimap[scanid:scanid + 3, 2:5].reshape(9)
                relatednorm = normvalue[scanid:scanid + 3, 2:5].reshape(9)
            else:
                relatedpixels = longimap[scanid:scanid + 3, 3:6].reshape(9)
                relatednorm = normvalue[scanid:scanid + 3, 3:6].reshape(9)

            print(scanid, fiid)
            fileid[0] = fiid
            intarray = ReadSingleRawTxTReturnIntegral(
                filenamedf['scanstart' +
                           str(longimap[scanid + 1, 2])].iloc[fiid],
                relatedpixels,
                relatednorm,
                Nbofevents,
                thisledch=ledch,
                thisbeamch=beamch)
            for itr in range(intarray.shape[0]):
                for k in range(9):
                    intch[k] = intarray[itr, k]
                if fiid < 15:
                    tree1.Fill()
                elif fiid < 33:
                    tree2.Fill()
                else:
                    tree3.Fill()
        f.Write()
        f.Close()
def plot_xsec_contour(sample):
    output_file_name = "lhe.root"
    output_file = TFile(output_file_name, "RECREATE")
    output_tree = TTree("Physics", "Physics")
    gROOT.ProcessLine(
        "struct MyStruct{ Float_t xsec; Float_t mdm; Float_t mhs; Float_t mzp; Float_t whs; Float_t wzp; };"
    )
    from ROOT import MyStruct
    s = MyStruct()
    output_tree.Branch('xsec', AddressOf(s, 'xsec'), 'xsec/F')
    output_tree.Branch('mdm', AddressOf(s, 'mdm'), 'mdm/F')
    output_tree.Branch('mhs', AddressOf(s, 'mhs'), 'mhs/F')
    output_tree.Branch('mzp', AddressOf(s, 'mzp'), 'mzp/F')
    output_tree.Branch('whs', AddressOf(s, 'whs'), 'whs/F')
    output_tree.Branch('wzp', AddressOf(s, 'wzp'), 'wzp/F')

    for i, su in enumerate(sample):
        s.xsec = 0.
        s.mdm = 0.
        s.mhs = 0.
        s.mzp = 0.
        s.whs = 0.
        s.wzp = 0.
        with open('../samples/' + su + '.lhe', "rt") as fin:
            for line in fin:
                if "Integrated weight (pb)" in line:
                    #print line.split(' ')[-1]
                    s.xsec = float(line.split(' ')[-1])
                if "# mdm" in line:
                    #print line.split(' ')[-3]
                    s.mdm = float(line.split(' ')[-3])
                if "# mhs" in line:
                    #print line.split(' ')[-3]
                    s.mhs = float(line.split(' ')[-3])
                if "# mzp" in line:
                    #print line.split(' ')[-3]
                    s.mzp = float(line.split(' ')[-3])
                #hs
                if "DECAY  54" in line:
                    #print line.split(' ')[-1]
                    s.whs = float(line.split(' ')[-1])
                #zp
                if "DECAY  55" in line:
                    #print line.split(' ')[-1]
                    s.wzp = float(line.split(' ')[-1])
            output_tree.Fill()
    output_tree.Write()
    output_file.Close()

    #plot contour
    gROOT.Macro('Contour.C("lhe.root","BBbarDM_hs50")')
Beispiel #18
0
def create_singles_TTree(host_file):
    from ROOT import TTree
    host_file.cd()
    title = 'Single events by Sam Kohn (git: {})'.format(
            translate.git_describe())
    out = TTree('singles', title)
    buf = TreeBuffer()
    initialize_basic_TTree(out, buf)
    buf.fQuad = float_value()
    buf.fMax = float_value()
    buf.fPSD_t1 = float_value()
    buf.fPSD_t2 = float_value()
    buf.f2inch_maxQ = float_value()
    buf.x = float_value()
    buf.y = float_value()
    buf.z = float_value()
    buf.x_AdTime = float_value()
    buf.y_AdTime = float_value()
    buf.z_AdTime = float_value()
    buf.fID = float_value()
    buf.fPSD = float_value()
    buf.num_nearby_events = unsigned_int_value()
    buf.nearby_dt = int_value(10)
    buf.nearby_energy = float_value(10)

    def branch(name, typecode):
        out.Branch(name, getattr(buf, name), '{}/{}'.format(name,
            typecode))
        return

    branch('fQuad', 'F')
    branch('fMax', 'F')
    branch('fPSD_t1', 'F')
    branch('fPSD_t2', 'F')
    branch('f2inch_maxQ', 'F')
    branch('x', 'F')
    branch('y', 'F')
    branch('z', 'F')
    branch('x_AdTime', 'F')
    branch('y_AdTime', 'F')
    branch('z_AdTime', 'F')
    branch('fID', 'F')
    branch('fPSD', 'F')
    branch('num_nearby_events', 'I')
    out.Branch('nearby_dt', buf.nearby_dt, 'nearby_dt[num_nearby_events]/I')
    out.Branch('nearby_energy', buf.nearby_energy,
            'nearby_energy[num_nearby_events]/F')
    return out, buf
Beispiel #19
0
def convert(path_in, path_out, mode):
    for i in xrange(len(path_in)):
        try:
            f_in = TFile(path_in[i])
            t_in = f_in.Get('save')
            entries = t_in.GetEntries()
            logging.info(mode[i] + ' entries :' + str(entries))
        except:
            logging.error(path_in[i] + ' is invalid!')
            sys.exit()
        
        print '--> Begin to process file: ' + path_in[i]
        f_out = TFile(path_out[i], 'RECREATE')
        t_out = TTree('save', 'save')

        m_rm_D = array('d', [999.])
        t_out.Branch('rm_D', m_rm_D, 'm_rm_D/D')

        nentries = t_in.GetEntries()
        for ientry in xrange(nentries):
            t_in.GetEntry(ientry)
            m_rm_D[0] = t_in.m_rm_D
            t_out.Fill()
            m_rm_D[0] = t_in.m_rm_Dmiss
            t_out.Fill()

        f_out.cd()
        t_out.Write()
        f_out.Close()
        print '--> End of processing file: ' + path_out[i]
Beispiel #20
0
def WriteTree(runSet):

    print 'Writing tree'

    # Make a tree
    froot = TFile('mu.root', 'RECREATE')
    tree = TTree("mu", "Flat ntuple for mu")
    tree.SetDirectory(froot)

    leaves = "run/F:lb/F:mu/F"

    leafValues = array("f", [0.0, 0.0, 0.0])

    newBranch = tree.Branch("mu_vars", leafValues, leaves)

    for r, lbs in runSet.iteritems():
        for i in range(lbs.FirstLB(), lbs.LastLB()):
            leafValues[2] = lbs.GetMeanPileUp(i)
            leafValues[0] = float(r)
            leafValues[1] = i

            print 'mu: %s run: %s lb: %s' % (leafValues[2], leafValues[0],
                                             leafValues[1])

            tree.Fill()

    tree.Write()
    #froot.Write()
    froot.Close()
Beispiel #21
0
    def _to_root(self, data, filename):

        import ROOT
        from ROOT import TFile, TTree, gROOT, AddressOf

        columns = [
            self._root_column(data, i, c['type'], c['name'])
            for i, c in enumerate(data['cols'])
        ]
        header = 'struct data_t { ' + ';'.join(columns) + '; };'

        gROOT.ProcessLine(header)
        row = ROOT.data_t()
        f = TFile(filename, 'RECREATE')
        tree = TTree('data', 'data from RHAPI')
        tree.Branch('data', row)

        for r in data['data']:
            for i, c in enumerate(data['cols']):
                v = r[i]
                if v is None:
                    if c['type'] == 'NUMBER': v = -1
                    else: v = ''
                try:
                    setattr(row, c['name'], v)
                except Exception as e:
                    print(c['name'], '=', v)
                    print(c, v)
                    print(e)
            tree.Fill()

        tree.Print()
        tree.Write()
def main(args):
    n_evs = args.n_evs
    n_files = args.n_files
    base_dir = args.base_dir

    file_name_tmpl = "file{}.root"
    tree_name = "DummyTree"

    print("Produce {} TTrees with {} events each".format(n_files, n_evs))

    for i in range(n_files):
        f = TFile(base_dir + "/" + file_name_tmpl.format(i), 'recreate')
        t = TTree(tree_name, tree_name)

        x = array('d', [0.])
        y = array('d', [0.])
        t.Branch('x', x, 'x/D')
        t.Branch('y', y, 'y/D')

        for k in range(n_evs):
            x[0] = ROOT.gRandom.Gaus(5, 2)
            # uniform [0,10]
            y[0] = ROOT.gRandom.Rndm() * 10
            t.Fill()

        f.Write()
        f.Close()
Beispiel #23
0
    def createDecorationTree(self, dataset, treeName=None):
      from ROOT import TTree
      #  logging.info("Decorating %s, tree:%s with %s..."% \
      #        (self.dataset.name,treeName,self.decorationList[i].branch))
      # get primary tree
      # recursive call on sub-datasets
      if isinstance(dataset,PhysicsProcess):
          for d in dataset.datasets: 
              self.createDecorationTree(d, treeName)
          return
      if dataset.name.endswith("2015"):
        print "Skipping DS %s to avoid duplicated MC tree in decor file for 2015/16 splitting"%dataset.name
        return
      dsTree = dataset._open(treeName)
      friendTree = TTree("DecorationFriend_"+dataset.name, "DecorationFriend_"+dataset.name)
      
      values={}
      branches={}
      for i in range(0,len(self.decorationList)):
        values[i] = array('f',[0.])
       # print "register values",values[i],hex(id(values[i]))
        branches[i] = friendTree.Branch(self.decorationList[i].branch,values[i],"%s/F"%(self.decorationList[i].branch))
       # print "register branches",branches[i],hex(id(branches[i]))

      # decorate
      for i in dsTree:
          for j in range(0,len(self.decorationList)):
            values[j][0] = self.calc(dsTree,j)
          friendTree.Fill()
          for j in range(0,len(self.decorationList)):
            val=branches[j].GetLeaf(self.decorationList[j].branch).GetValue()
        #    print "getting branch val:", j, val
      
      self.writeDecorationFile(friendTree,dataset,treeName)
Beispiel #24
0
def addBranchToDirectory():
    print ">>> Add branch to directory."

    file = TFile("tree.root", 'update')

    # make directory if it does not exist
    dir = file.GetDirectory("dir1")
    if not dir:
        print ">>>   created dir1"
        dir = file.mkdir("dir1")

    # make this directory the current one: everything you write will be saved here
    dir.cd()

    # make new tree with a branch
    tree = TTree("tree2", "tree2")
    n = array('d', [0])
    #n = numpy.zeros(1, dtype=float)
    tree.Branch("normal5", n, 'normal5/D')

    # fill tree
    for i in xrange(10000):
        n[0] = gRandom.Gaus(3, 2)
        tree.Fill()

    file.Write("", TFile.kOverwrite)
    file.Close()
Beispiel #25
0
def main(args):
    n_branches = args.n_branches
    n_evs = args.n_evs
    n_files = args.n_files
    base_dir = args.base_dir

    file_name_tmpl = "file{}_{}branches_{}evs.root"
    tree_name = "DummyTree"

    print("Produce {} TTrees with {} branches, {} events each".format(
        n_files, n_branches, n_evs))

    for i in range(n_files):
        f = TFile(base_dir + "/" + file_name_tmpl.format(i, n_branches, n_evs),
                  'recreate')
        t = TTree(tree_name, tree_name)

        arrays = []
        for n in range(n_branches):
            branch_name = "branch_{}".format(n)
            br = array('d', [0.])
            arrays.append(br)
            t.Branch(branch_name, br, '{}/D'.format(branch_name))

        for k in range(n_evs):
            for br in arrays:
                br[0] = ROOT.gRandom.Gaus(5, 2)
            t.Fill()

        f.Write()
        f.Close()
Beispiel #26
0
 def do_toymc(self, sigtree, hrec, hout0):
     print('We are doing toy mc...')
     if self.ntoys < 0:
         print('WARNING: ntoys < 0, set it to 1')
         self.ntoys = 1
     nevents = int(hrec.Integral())
     list_hout = []
     toymcfile = TFile('toymcfile.root', 'recreate')
     for itoy in range(self.ntoys):
         print('doing toymc ', itoy, '...')
         datatree = TTree(self.data_treename, 'toymc_' + str(itoy))
         mrec = array('f', [0])
         datatree.Branch('mrec', mrec, 'mrec/F')
         for i in range(nevents):
             mrec[0] = hrec.GetRandom()
             datatree.Fill()
         datatree.Write()
         x, hout, binning = self.unfold_core(sigtree,
                                             datatree,
                                             self.nsplit,
                                             dotoymc=1)
         list_hout.append(hout)
     if self.debug or 0:
         self.showhist(hout0)
         for hout in list_hout:
             print('toymc results', hout)
             self.showhist(hout)
     self.Vx = self.get_newVx_from_toymc(hout0, list_hout)
     toymcfile.Close()
Beispiel #27
0
def create_simple_tree(nevents):
    """Create a simple TTree with a couple of branches."""
    t = TTree('aTTree','A TTree')
    t.Branch('run', AddressOf(pytree, 'run'),'run/I')
    t.Branch('evt', AddressOf(pytree, 'evt'),'evt/I')
    t.Branch('x', AddressOf(pytree, 'x'),'x/F')
    t.Branch('y', AddressOf(pytree, 'y'),'y/F')
    t.Branch('z', AddressOf(pytree, 'z'),'z/F')
    for i in range(nevents):
        pytree.run = 1 if (i<500) else 2
        pytree.evt = i
        pytree.x = gauss(0., 10.)
        pytree.y = gauss(0, 10.)
        pytree.z = gauss(5., 50.)
        t.Fill()
    return t
Beispiel #28
0
    def setUp(self):
        super(TQHistogramObservableTest, self).setUp()

        f = TFile(os.path.join(self.tempdir, 'test.root'), 'recreate')
        t = TTree('testTree', 'testTree')

        d = array('f', [0.])
        t.Branch('testVal', d, 'testVal/F')

        for i in range(1000):

            d[0] = random.gauss(0., 1.)
            t.Fill()

        f.Write()
        f.Close()

        tf = TFile.Open(os.path.join(self.tempdir, "testHistogram.root"),
                        "RECREATE")
        histMap = TH1F('mapping', '', 2, -20., 20.)
        histMap.SetBinContent(1, 0.5)
        histMap.SetBinContent(2, 2)
        #tf.cd()
        histMap.Write()
        tf.Close()
Beispiel #29
0
class TreeProducerCommon(object):
    """Class to create a custom output file & tree; as well as create and contain branches."""
    
    def __init__(self, filename, module, **kwargs):
        print 'TreeProducerCommon is called for', name
        
        ncuts         = kwargs.get('ncuts',25)
        self.filename = filename
        self.module   = module
        self.outfile  = TFile(filename,'RECREATE')
        self.tree     = TTree('tree','tree')
        self.cutflow  = TH1D('cutflow','cutflow',ncuts,0,ncuts)
    
    def addBranch(self, name, dtype='f', default=None, arrname=None):
        """Add branch with a given name, and create an array of the same name as address."""
        if hasattr(self,name):
          raise IOError("Branch of name '%s' already exists!"%(name))
        if not arrname:
          arrname = name
        if isinstance(dtype,str):
          if dtype.lower()=='f':   # 'f' is only a 'float32', and 'F' is a 'complex64', which do not work for filling float branches
            dtype = float          # float is a 'float64' ('f8')
          elif dtype.lower()=='i': # 'i' is only a 'int32'
            dtype = int            # int is a 'int64' ('i8')
        setattr(self,arrname,np.zeros(1,dtype=dtype))
        self.tree.Branch(name, getattr(self,arrname), '%s/%s'%(name,root_dtype[dtype]))
        if default!=None:
          getattr(self,name)[0] = default
    
    def endJob(self):
        """Write and close files after the job ends."""
        self.outfile.Write()
        self.outfile.Close()
def ConvertParquet2Root(parquetFile, outputName, treename='tree'):
    '''
    Converts a pandas dataframe to a TTree and saves it in a root file.
    Warning: this function is very slow.
    Arguments
    ----------
    - parquetFile: the absolute path of the file that you wanto to convert
    - outputName: the name of the root file

    '''

    oFile = TFile(outputName, 'recreate')
    tree = TTree(treename, treename)

    df = pd.read_parquet(parquetFile, engine='pyarrow')

    print(f'\n\033[94mConversion of {parquetFile}\033[0m')

    data = []
    for iCol, colname in enumerate(list(df)):
        data.append(np.ones((1), dtype="float32"))
        tree.Branch(colname, data[iCol], colname + "/F")

    with alive_bar(len(df)) as bar:
        for i in range(len(df)):
            for iCol, colname in enumerate(list(df)):
                data[iCol][0] = df.values[i, iCol]
            tree.Fill()
            bar()

    oFile.Write()