Exemplo n.º 1
0
def tree_to_recarray(trees, branches=None, include_weight=False, weight_name="weight", weight_dtype="f4"):
    """
    Convert a tree or a list of trees into a numpy.recarray
    with fields corresponding to the tree branches
    """
    if isinstance(trees, (list, tuple)):
        return np.concatenate(
            [
                _add_weight_field(tree2rec(tree, branches), tree, include_weight, weight_name, weight_dtype)
                for tree in trees
            ]
        )
    return _add_weight_field(tree2rec(trees, branches), trees, include_weight, weight_name, weight_dtype)
Exemplo n.º 2
0
def loadntuples( datadir, ymlfile='', branches=[], vetolist=[], acceptlist=[], verbose=True, cuts='' ):
    msg.overridelevel = 1-int(verbose)
    allsamples = sample.loadcollection( ymlfile=ymlfile, tag=globalcfg.config( 'HWW', datadir ), vetolist=vetolist, acceptlist=acceptlist )

    rec={}
    for smpl in allsamples:
        if cuts == '':
            cuts = smpl.get_cuts()
        else:
            cuts = '({0})*({1})'.format( smpl.get_cuts(), cuts )
        if branches != []:
            rec[smpl.get_name()] = rootnp.tree2rec( smpl.get_tree(), branches=branches, selection = cuts )
        else:
            rec[smpl.get_name()] = rootnp.tree2rec( smpl.get_tree(), selection = cuts )
    return rec
Exemplo n.º 3
0
def test_selection():
    chain = TChain('tree')
    chain.Add(load('single1.root'))
    chain.Add(load('single2.root'))
    a = rnp.tree2rec(chain, selection="d_double > 100")
    assert_equal((a['d_double'] <= 100).any(), False)

    # selection with differing variables in branches and expression
    a = rnp.tree2array(chain,
        branches=['d_double'],
        selection="f_float < 100 && n_int%2 == 1")

    # selection with TMath
    a = rnp.tree2rec(chain,
        selection="TMath::Erf(d_double) < 0.5")
Exemplo n.º 4
0
    def loadEventRange( self, start, end ):
        """ load events between range from tree into pandas data frame """
        # load event range from tree
        s = time.time()
        print "Getting Data between Event %d and %d" % ( start, end )
        if start>end:
            # come on, man
            tmp = end
            end = start
            start = tmp
        if start<self.first_event:
            start = self.first_event

        # definie new event range
        self.event_range = [ start, end ]
        # we find position in tree
        start_entry = self.searchEntryHistory( start )
        stop_entry = self.scanForEvent( end+1 )
        print "start/stop entry in tree: ",start_entry,stop_entry
        # load data into pandas dataframe
        numpy_rec_array= tree2rec( self.ttree, selection="event>=%d && event<=%d"%(self.event_range[0], self.event_range[1]), start=start_entry, stop=stop_entry )
        self.wf_df = pd.DataFrame(numpy_rec_array)
        # append the sample sizes of adc strings as a column to dataframe
        self.determineWaveformLengths()
        print "Time to load: ",time.time()-s
Exemplo n.º 5
0
    def getNextEntry(self):
        """ concrete instantiation of abc method """
        # get next event
        if self.event is not None:
            nextevent = self.event+1
        else:
            nextevent = self.first_event
        if self.maxevent is not None and nextevent>=self.maxevent:
            print "No events for ",nextevent,"! maxevent=",self.maxevent
            return False
        ok = True
        inc = 1

        q = self.wf_df.query( "event>=%d"%(nextevent) )
        if len( q )==0:
            print "empty query. must find next event number if tree"
            events = self.entry_points.keys()
            events.sort()
            last_tree_entry = self.entry_points[ events[-1] ]
            numpy_rec_array= tree2rec( self.ttree, selection="event>%d"%(self.event), start=last_tree_entry, stop=last_tree_entry+10 )
            wf_df = pd.DataFrame(numpy_rec_array)
            nextevent =  wf_df["event"].min()
            self.entry_points[ nextevent ] = last_tree_entry+1
            print "loaded edge of new dataframe. nextevent now ",nextevent
            #raw_input()
            if len(wf_df)==0:
                return False # no more
        else:
            nextevent = q["event"].min()
            #print  q["event"]
            print "next event is ",nextevent
        
        return self.gotoEvent( nextevent )
Exemplo n.º 6
0
def main(infile, outfile):
    # open my data file
    f = root_open(infile)
    # build a dictionary of the file structure

    dic = {p: t for p, d, t in f if p != ''}
    # open a new file to put the data in after the cut is applied
    f_store = pd.HDFStore(outfile)
    # iterate over the directoies
    for d, tl in dic.iteritems():

        # iterate over list of trees in the directory
        for t in tl:
            # tree is original data tree
            tree = f[d][t]
            #print tree.branchnames

            df = pd.DataFrame(
                tree2rec(
                    tree,
                    branches=tree.branchnames))
            print "writing to: {}/{}".format(d, t), tree.GetEntries()
            f_store.append('{}/{}'.format(d, t), df)
            tree.IsA().Destructor(tree)
    f_store.close()
    f.close()
Exemplo n.º 7
0
def getTreeToRec(tree):
    """
        Convert and return a tree into a record numpy array
        Inputs: TTree object
        Return: rec array
    """
    from root_numpy import tree2rec
    return tree2rec(treename)
Exemplo n.º 8
0
def test_branch_status():
    # test that original branch status is preserved
    chain = TChain('tree')
    chain.Add(load('single1.root'))
    chain.Add(load('single2.root'))
    chain.SetBranchStatus('d_double', False)
    a = rnp.tree2rec(chain, selection="d_double > 100")
    assert_equal(chain.GetBranchStatus('d_double'), False)
Exemplo n.º 9
0
def test_branch_status():
    # test that original branch status is preserved
    chain = TChain('tree')
    chain.Add(load('single1.root'))
    chain.Add(load('single2.root'))
    chain.SetBranchStatus('d_double', False)
    a = rnp.tree2rec(chain, selection="d_double > 100")
    assert_equal(chain.GetBranchStatus('d_double'), False)
Exemplo n.º 10
0
def make_JointPlot(plot, region, data, backgrounds) :

    sample_to_plot = []
    if data.name == plot.sample : sample_to_plot.append(data)
    if not len(sample_to_plot) :
        for bk in backgrounds :
            if bk.name == plot.sample : sample_to_plot.append(bk)
    if len(sample_to_plot) == 0 or len(sample_to_plot) > 1 :
        msg('ERROR make_JointPlot received %d samples to plot for plot with name %s'%(len(sample_to_plot), plot.name))
        sys.exit()

    # turn this tree into an array :)
    sample_to_plot = sample_to_plot[0]
    selection_ = '(' + region.tcut + ') * eventweight * ' + str(sample_to_plot.scale_factor)
    tree_array = tree2rec(sample_to_plot.tree, branches=[plot.x_var, plot.y_var],
                            selection=selection_)
    tree_array.dtype.names = (plot.x_var, plot.y_var)
    x_arr = tree_array[plot.x_var]
    y_arr = tree_array[plot.y_var]

    sns.set(style="white")

    # stats?
    stat_func_ = None
    if plot.stat_func == "kendalltau" :
        from scipy.stats import kendalltau
        stat_func_ = kendalltau
    elif plot.stat_func == None :
        from scipy.stats import pearsonr
        stat_func_ = pearsonr

    j_plot_grid = None
    if plot.cmap == None or plot.cmap == "default" :
        j_plot_grid = sns.jointplot(x_arr, y_arr, kind = plot.kind, stat_func=stat_func_, color = plot.color, linewidth = plot.line_width, ylim=[plot.y_range_min,plot.y_range_max], xlim=[plot.x_range_min,plot.x_range_max])
        #j_plot_grid = sns.jointplot(x_arr, y_arr, kind = plot.kind, stat_func=stat_func_, color = plot.color, linewidth = plot.line_width, joint_kws={"n_levels":plot.n_levels, "shade":True}, ylim=[plot.y_range_min,plot.y_range_max], xlim=[plot.x_range_min,plot.x_range_max])

    elif plot.cmap == "cubehelix" :
        cmap_ = sns.cubehelix_palette(as_cmap=True, dark=0, light=1, reverse = True)
        j_plot_grid = sns.jointplot(x_arr, y_arr, kind = plot.kind, stat_func=stat_func_, linewidth = plot.line_width, joint_kws={"cmap":cmap_, "n_levels":plot.n_levels, "shade":True}, ylim=[plot.y_range_min, plot.y_range_max], xlim=[plot.x_range_min,plot.x_range_max])
    elif plot.cmap == "blues" :
        j_plot_grid = sns.jointplot(x_arr, y_arr, kind = plot.kind, stat_func=stat_func_, linewidth = 1.0, joint_kws={"cmap":"Blues", "n_levels":plot.n_levels, "shade":True, "shade_lowest":False}, ylim=[plot.y_range_min, plot.y_range_max], xlim=[plot.x_range_min,plot.x_range_max])
    else :
        msg("cmap attribute of joint plot not yet added")
        sys.exit()

    j_plot_grid.fig.suptitle(plot.title)
    j_plot_grid.fig.subplots_adjust(top=0.935)
    j_plot_grid.set_axis_labels(plot.x_label, plot.y_label)


    # save the plot to file
    outname = plot.name + ".eps"
    j_plot_grid.savefig(outname)
    out = indir + "/plots/" + outdir 
    utils.mv_file_to_dir(outname, out, True)
    fullname = out + "/" + outname
    msg("%s saved to : %s"%(outname, os.path.abspath(fullname)))
Exemplo n.º 11
0
 def loadEvent( self, event ):
     start = time.time()
     print "loading event %d waveforms ..."%(event),
     self.current_rec = tree2rec( self.rawtree, selection="event==%d"%(event) )
     if len(self.current_rec)>0:
         self.waveforms_cached = False
         print " ok."
     else:
         print "No waveforms found from event %d in the file" %(event)
     print "Loading time: ",time.time()-start," secs"
Exemplo n.º 12
0
def test_weights():
    f = TFile(load('test.root'))
    tree = f.Get('tree')
    tree.SetWeight(5.)
    rec = rnp.tree2rec(tree, include_weight=True, weight_name='treeweight')
    assert_array_equal(rec['treeweight'], np.ones(100) * 5)
    f = load(['single1.root', 'single2.root'])
    a = rnp.root2array(f, include_weight=True)
    assert_array_equal(a['weight'],
        np.concatenate((np.ones(100) * 2., np.ones(100) * 3.)))
Exemplo n.º 13
0
    def get_data(filename="~/AMSDeutons/test.root", treename="data"):
        """ Converts a ROOT tree to a pandas dataframe. Unwraps vectors."""
        tfile = ROOT.TFile(filename)

        tree = tfile.Get(treename)
        data = root_numpy.tree2rec(tree)
        data = pd.DataFrame(data)
        row = data.ix[0]

        for c, t in data.dtypes.iteritems():
            if c == "fStatus": continue
            if not t == np.object: continue
            for i in range(len(row[c])):
                newc = "{0}_{1}".format(c,i)
                data[newc] = data[c].str.get(i)
            del data[c]

        return data 
Exemplo n.º 14
0
    def __init__(self,
                 bdsimFileName="output.pickle",
                 mad8TwissFileName="ebds1.twiss",
                 mad8EnvelopeFileName="ebds1.envelope"):
        # load bdsim data
        if bdsimFileName.find("pickle") != -1:
            f = open(bdsimFileName)
            self.bdsimData = _pkl.load(f)
            self.bdsimOptics = self.bdsimData['optics']
            f.close()
        elif bdsimFileName.find(".root") != -1:
            import ROOT as _ROOT
            import root_numpy as _root_numpy
            f = _ROOT.TFile(bdsimFileName)
            t = f.Get("optics")
            self.bdsimOptics = _root_numpy.tree2rec(t)

        # load mad8 data
        r = _pymad8.Mad8.OutputReader()
        [self.mad8Comm, self.mad8Envel] = r.readFile(mad8EnvelopeFileName,
                                                     "envel")
        [self.mad8Comm, self.mad8Twiss] = r.readFile(mad8TwissFileName,
                                                     "twiss")
Exemplo n.º 15
0
import numpy as np
import pandas as pd
from root_numpy import rec2array, root2array, root2rec, tree2rec, tree2array

sys.path.append('include/')

from chargedHbranches import *

# grab signal trees, merge them and assign a name
filename_signal = 'data/341524_0.root'
file_signal = ROOT.TFile(filename_signal)
sigTree = file_signal.Get('NOMINAL')

#convert signal tree to numpy array
signal = tree2rec(sigTree,
                  branches,
                  'met_et>100000.&&event_clean&&n_vx>=1&&bsm_tj_dirty_jet==0',
                  start=0)

# grab background trees, merge them and assign a name
filename_bkg = 'data/410000_0.root'
file_bkg = ROOT.TFile(filename_bkg)
bkgTree = file_bkg.Get('NOMINAL')

#convert background tree to numpy array
bkg = tree2rec(bkgTree,
               branches,
               'met_et>100000.&&event_clean&&n_vx>=1&&bsm_tj_dirty_jet==0',
               start=0)
branchnames = bkg.dtype.names

signal = rec2array(signal)
Exemplo n.º 16
0
def root2hdf5(rfile, hfile, rpath="", entries=-1, userfunc=None, selection=None, show_progress=False):
    """
    Convert all trees in a ROOT file into tables in an HDF5 file.

    Parameters
    ----------

    rfile : string or asrootpy'd ROOT File
        A ROOT File handle or string path to an existing ROOT file.

    hfile : string or PyTables HDF5 File
        A PyTables HDF5 File handle or string path to an existing HDF5 file.

    rpath : string, optional (default='')
        Top level path to begin traversal through the ROOT file. By default
        convert everything in and below the root directory.

    entries : int, optional (default=-1)
        The number of entries to read at once while converting a ROOT TTree
        into an HDF5 table. By default read the entire TTree into memory (this
        may not be desired if your TTrees are large).

    userfunc : callable, optional (default=None)
        A function that will be called on every tree and that must return a
        tree or list of trees that will be converted instead of the original
        tree.

    selection : string, optional (default=None)
        A ROOT selection expression to be applied on all trees before
        conversion.

    show_progress : bool, optional (default=False)
        If True, then display and update a progress bar on stdout as each tree
        is converted.

    """
    show_progress = show_progress and check_tty(sys.stdout)
    if show_progress:
        widgets = [Percentage(), " ", Bar(), " ", ETA()]

    own_rootfile = False
    if isinstance(rfile, basestring):
        rfile = root_open(rfile)
        own_rootfile = True

    own_h5file = False
    if isinstance(hfile, basestring):
        hfile = tables_open(filename=hfile, mode="w", title="Data")
        own_h5file = True

    for dirpath, dirnames, treenames in rfile.walk(rpath, class_pattern="TTree"):

        # skip root
        if not dirpath and not treenames:
            continue

        # skip directories w/o trees or subdirs
        if not dirnames and not treenames:
            continue

        where_group = "/" + os.path.dirname(dirpath)
        current_dir = os.path.basename(dirpath)

        if not current_dir:
            group = hfile.root
        else:
            group = hfile.createGroup(where_group, current_dir, "")

        ntrees = len(treenames)
        log.info("Will convert {0:d} tree{1} in this directory".format(ntrees, "s" if ntrees != 1 else ""))

        for treename in treenames:

            input_tree = rfile.Get(os.path.join(dirpath, treename))

            if userfunc is not None:
                tmp_file = TemporaryFile()
                # call user-defined function on tree and get output trees
                log.info("Calling user function on tree '{0}'".format(input_tree.GetName()))
                trees = userfunc(input_tree)

                if not isinstance(trees, list):
                    trees = [trees]

            else:
                trees = [input_tree]
                tmp_file = None

            for tree in trees:

                log.info("Converting tree '{0}' with {1:d} entries ...".format(tree.GetName(), tree.GetEntries()))

                if tree.GetName() in group:
                    log.warning("skipping tree '{0}' that already exists " "in the output file".format(tree.GetName()))
                    continue

                total_entries = tree.GetEntries()
                pbar = None
                if show_progress and total_entries > 0:
                    pbar = ProgressBar(widgets=widgets, maxval=total_entries)

                if entries <= 0:
                    # read the entire tree
                    if pbar is not None:
                        pbar.start()
                    recarray = tree2rec(tree, selection=selection)
                    recarray = _drop_object_col(recarray)
                    if TABLES_NEW_API:
                        table = hfile.create_table(group, tree.GetName(), recarray, tree.GetTitle())
                    else:
                        table = hfile.createTable(group, tree.GetName(), recarray, tree.GetTitle())
                    # flush data in the table
                    table.flush()
                    # flush all pending data
                    hfile.flush()
                else:
                    # read the tree in chunks
                    start = 0
                    while start < total_entries or start == 0:
                        if start > 0:
                            with warnings.catch_warnings():
                                warnings.simplefilter("ignore", RootNumpyUnconvertibleWarning)
                                warnings.simplefilter("ignore", tables.NaturalNameWarning)
                                recarray = tree2rec(tree, selection=selection, start=start, stop=start + entries)
                            recarray = _drop_object_col(recarray, warn=False)
                            table.append(recarray)
                        else:
                            recarray = tree2rec(tree, selection=selection, start=start, stop=start + entries)
                            recarray = _drop_object_col(recarray)
                            if pbar is not None:
                                # start after any output from root_numpy
                                pbar.start()
                            if TABLES_NEW_API:
                                table = hfile.create_table(group, tree.GetName(), recarray, tree.GetTitle())
                            else:
                                table = hfile.createTable(group, tree.GetName(), recarray, tree.GetTitle())
                        start += entries
                        if start <= total_entries and pbar is not None:
                            pbar.update(start)
                        # flush data in the table
                        table.flush()
                        # flush all pending data
                        hfile.flush()

                if pbar is not None:
                    pbar.finish()

            input_tree.Delete()

            if userfunc is not None:
                for tree in trees:
                    tree.Delete()
                tmp_file.Close()

    if own_h5file:
        hfile.close()
    if own_rootfile:
        rfile.Close()
Exemplo n.º 17
0
def tree2hdf5(tree, hfile, group=None,
              entries=-1, selection=None,
              show_progress=False):
    """
    Convert a TTree into a HDF5 table.

    Parameters
    ----------

    tree : ROOT.TTree
        A ROOT TTree.

    hfile : string or PyTables HDF5 File
        A PyTables HDF5 File handle or string path to an existing HDF5 file.

    group : string or PyTables Group instance, optional (default=None)
        Write the table at this location in the HDF5 file.

    entries : int, optional (default=-1)
        The number of entries to read at once while converting a ROOT TTree
        into an HDF5 table. By default read the entire TTree into memory (this
        may not be desired if your TTrees are large).

    selection : string, optional (default=None)
        A ROOT selection expression to be applied on the TTree before
        conversion.

    show_progress : bool, optional (default=False)
        If True, then display and update a progress bar on stdout as the TTree
        is converted.

    """
    show_progress = show_progress and check_tty(sys.stdout)
    if show_progress:
        widgets = [Percentage(), ' ', Bar(), ' ', ETA()]

    own_h5file = False
    if isinstance(hfile, basestring):
        hfile = tables_open(filename=hfile, mode="w", title="Data")
        own_h5file = True

    log.info("Converting tree '{0}' with {1:d} entries ...".format(
        tree.GetName(),
        tree.GetEntries()))

    if not group:
        group = hfile.root
    elif isinstance(group, basestring):
        group_where = '/' + os.path.dirname(group)
        group_name = os.path.basename(group)
        if TABLES_NEW_API:
            group = hfile.create_group(group_where, group_name,
                                       createparents=True)
        else:
            group = hfile.createGroup(group_where, group_name)

    if tree.GetName() in group:
        log.warning(
            "Tree '{0}' already exists "
            "in the output file".format(tree.GetName()))
        return

    total_entries = tree.GetEntries()
    pbar = None
    if show_progress and total_entries > 0:
        pbar = ProgressBar(widgets=widgets, maxval=total_entries)

    if entries <= 0:
        # read the entire tree
        if pbar is not None:
            pbar.start()
        recarray = tree2rec(tree, selection=selection)
        recarray = _drop_object_col(recarray)
        if TABLES_NEW_API:
            table = hfile.create_table(
                group, tree.GetName(),
                recarray, tree.GetTitle())
        else:
            table = hfile.createTable(
                group, tree.GetName(),
                recarray, tree.GetTitle())
        # flush data in the table
        table.flush()
        # flush all pending data
        hfile.flush()
    else:
        # read the tree in chunks
        start = 0
        while start < total_entries or start == 0:
            if start > 0:
                with warnings.catch_warnings():
                    warnings.simplefilter(
                        "ignore",
                        RootNumpyUnconvertibleWarning)
                    warnings.simplefilter(
                        "ignore",
                        tables.NaturalNameWarning)
                    recarray = tree2rec(
                        tree,
                        selection=selection,
                        start=start,
                        stop=start + entries)
                recarray = _drop_object_col(recarray, warn=False)
                table.append(recarray)
            else:
                recarray = tree2rec(
                    tree,
                    selection=selection,
                    start=start,
                    stop=start + entries)
                recarray = _drop_object_col(recarray)
                if pbar is not None:
                    # start after any output from root_numpy
                    pbar.start()
                if TABLES_NEW_API:
                    table = hfile.create_table(
                        group, tree.GetName(),
                        recarray, tree.GetTitle())
                else:
                    table = hfile.createTable(
                        group, tree.GetName(),
                        recarray, tree.GetTitle())
            start += entries
            if start <= total_entries and pbar is not None:
                pbar.update(start)
            # flush data in the table
            table.flush()
            # flush all pending data
            hfile.flush()

    if pbar is not None:
        pbar.finish()

    if own_h5file:
        hfile.close()
Exemplo n.º 18
0

# define the model
class Event(TreeModel):

    x = FloatCol()
    y = FloatCol()
    z = FloatCol()
    i = IntCol()

tree = Tree("test", model=Event)

# fill the tree
for i in xrange(100000):
    tree.x = gauss(.5, 1.)
    tree.y = gauss(.3, 2.)
    tree.z = gauss(13., 42.)
    tree.i = i
    tree.fill()
tree.write()

# convert tree into a numpy record array
from root_numpy import tree2rec
array = tree2rec(tree)
print array
print array.x
print array.i
print tree.to_array()

f.close()
Exemplo n.º 19
0
f = root_open("test.root", "recreate")

# define the model
#class Event(TreeModel):
#    x = FloatCol()
#    y = FloatCol()
#    z = FloatCol()
#    i = IntCol()
#
#tree = Tree("test", model=Event)
#
## fill the tree
#for i in xrange(100):
#    tree.x = gauss(.5, 1.)
#    tree.y = gauss(.3, 2.)
#    tree.z = gauss(13., 42.)
#    tree.i = i
#    tree.fill()
#tree.write()

# convert tree into a numpy record array
from root_numpy import tree2rec
array = tree2rec(tree)
print array
print array.x
print array.i
print tree.to_array()

f.close()
Exemplo n.º 20
0
from plots.common.cuts import Cuts, Cut
from root_numpy import tree2rec

f = ROOT.TFile("exclusive/step3.root")

tree = f.Get("trees/Events")

interesting = [
    1798880,
    1808687,
    1838545
]

#interesting = [1837809]
for ev in tree:
    if ev.event_id in interesting:
        print "id =",ev.event_id, " pt =",ev.mu_pt, " iso =",ev.mu_iso, " eta =",ev.mu_eta

print "all events"
arr = tree2rec(tree, branches=["event_id"])
for i in interesting:
    if i in arr["event_id"]:
        print "%d is in tree" % i


print "HLT selection"
arr = tree2rec(tree, branches=["event_id"], selection=str(Cuts.hlt_isomu))
for i in interesting:
    if i in arr["event_id"]:
        print "%d passes HLT" % i
Exemplo n.º 21
0
def test_tree2rec():
    chain = TChain('tree')
    chain.Add(load('single1.root'))
    check_single(rnp.tree2rec(chain))
Exemplo n.º 22
0
# interactive matplotlib
plt.ion()

plt.rcParams.update({'font.size': 14})

fig, axarr = plt.subplots(2,2,figsize=(13,13))
scale2 = axarr[1,1].twinx()

f = ROOT.TFile(fin)

t   = f.Get( 'out_tree' )


arr = tree2rec(t,branches=['_michel_clustered_charge',
                           '_michel_Z','_michel_X','_X',
                           '_Z','_q_v','_s_v','_chi_v','_t_dqds_v',
                           '_t_q_v','_mean_chi','_lowest_chi',
                           '_rms_chi','_boundary','_chi_at_boundary',
                           '_event','_run','_subrun','_clus_idx','_forward'])

ismc = False;
if f.GetListOfKeys().Contains('_mc_tree'):
    ismc = True

all_hits = False
if (f.GetListOfKeys().Contains('_hit_tree')):
    all_hits = True
    #hit tree
    hit_tree    = f.Get('_hit_tree')


    # get hit information
Exemplo n.º 23
0
producers = []
efficiencies = [[], [], []]
purities = [[], [], []]

for line in txtfile:

    fname = line.split()[0]
    f = ROOT.TFile(fname)
    t = f.Get("fMMQualityTree")

    cluster_producer = fname.split("_")[3]
    producers.append(cluster_producer)
    print 'examining producer %s' % cluster_producer

    t_py = tree2rec(t)

    eff = [[], [], []]
    pur = [[], [], []]

    high_eff = [[], [], []]
    high_pur = [[], [], []]

    for n in xrange(len(t_py)):

        for pl in xrange(3):

            # efficiency
            for i in xrange(len(t_py['BestEff_%i' % pl][n])):
                eff_val = t_py['BestEff_%i' % pl][n][i]
                eff[pl].append(eff_val)
Exemplo n.º 24
0
if (len(sys.argv) < 2):
    print
    print "Incorrect Use"
    print "Please provide the path to the file containing the tree which stores info on the reco'ed Michel"
    print
    sys.exit(0)


f = ROOT.TFile(sys.argv[-1])
t   = f.Get( 'out_tree' )


arr = tree2rec(t,branches=['_michel_clustered_charge',
                           '_michel_n_hits',
                           'michel_start_X','michel_start_Z',
                           '_michel_Z','_michel_X','_X',
                           '_Z','_q_v','_s_v','_chi_v','_t_dqds_v',
                           '_t_q_v','_mean_chi','_lowest_chi',
                           '_rms_chi','_boundary','_chi_at_boundary',
                           '_event','_run','_subrun','_clus_idx','_forward'])


icarus_f = open('/home/david/Desktop/icarus_data.csv','r')

icarus_energies = []
icarus_numbers  = []
icarus_errors   = []
icarus_area = 0

e_prev = 0

for line in icarus_f:
Exemplo n.º 25
0
                         lambda x:y1,\
                         lambda x:x*((1-y1)/(x3-x2)) + (1-x3*((1-y1)/(x3-x2))),\
                         lambda x:1\
                         ])

BNB_events  = 1127678 + 795004
EXT_events  = 36206 + 9551
NUMI_events = 276205 + 91004
# POT: [4000,4500] : BNB 1.41E18, NUMI 2.46E18
# POT: [4500,4700] : BNB 3.10E18, NUMI 2.37E18
pot_BNB  = 1.41+3.10
pot_NUMI = 2.46+2.37

f = ROOT.TFile("/Users/tiamiceli/Data/TimingNeutrinos/daq_hist_small.root")
t = f.Get("opflash_tree_filtered")
df = pd.DataFrame(tree2rec(t,selection='pe_total > 50 & dt<23 & dt>0'))
del t

df_BNB  = df.query('trig_word == 2048')
df_NUMI = df.query('trig_word == 4096')
df_EXT  = df.query('trig_word == 512')

dt_BNB  = df_BNB['dt'].ravel()
dt_NUMI = df_NUMI['dt'].ravel()
dt_EXT = df_EXT['dt'].ravel()

# 20 is the usec that we are saving flashed for

# rate in counts / usec
EXT_rate = float(len(dt_EXT)) / ( 23. * float(EXT_events) )
# uncertainty on the EXT rate:
Exemplo n.º 26
0
def convert(rfile, hfile, rpath='', entries=-1, userfunc=None, selection=None):

    isatty = check_tty(sys.stdout)
    if isatty:
        widgets = [Percentage(), ' ', Bar(), ' ', ETA()]

    own_h5file = False
    if isinstance(hfile, basestring):
        hfile = tables.openFile(filename=hfile, mode="w", title="Data")
        own_h5file = True
    own_rootfile = False
    if isinstance(rfile, basestring):
        rfile = root_open(rfile)
        own_rootfile = True

    for dirpath, dirnames, treenames in utils.walk(rfile,
                                                   rpath,
                                                   class_pattern='TTree'):

        # skip root
        if not dirpath and not treenames:
            continue

        # skip directories w/o trees or subdirs
        if not dirnames and not treenames:
            continue

        where_group = '/' + os.path.dirname(dirpath)
        current_dir = os.path.basename(dirpath)

        if not current_dir:
            group = hfile.root
        else:
            group = hfile.createGroup(where_group, current_dir, "")

        ntrees = len(treenames)
        if ntrees > 1:
            log.info("Will convert %i trees in this directory" % ntrees)
        else:
            log.info("Will convert 1 tree in this directory")

        for treename in treenames:

            input_tree = rfile.Get(os.path.join(dirpath, treename))

            if userfunc is not None:
                tmp_file = TemporaryFile()
                # call user-defined function on tree and get output trees
                log.info("Calling user function on tree '%s'" %
                         input_tree.GetName())
                trees = userfunc(input_tree)

                if not isinstance(trees, list):
                    trees = [trees]

            else:
                trees = [input_tree]
                tmp_file = None

            for tree in trees:

                log.info("Converting tree '%s' with %i entries ..." %
                         (tree.GetName(), tree.GetEntries()))

                total_entries = tree.GetEntries()
                pbar = None
                if isatty and total_entries > 0:
                    pbar = ProgressBar(widgets=widgets, maxval=total_entries)

                if entries <= 0:
                    # read the entire tree
                    if pbar is not None:
                        pbar.start()
                    recarray = tree2rec(tree, selection=selection)
                    recarray = _drop_object_col(recarray)
                    table = hfile.createTable(group, tree.GetName(), recarray,
                                              tree.GetTitle())
                    # flush data in the table
                    table.flush()
                    # flush all pending data
                    hfile.flush()
                else:
                    # read the tree in chunks
                    offset = 0
                    while offset < total_entries or offset == 0:
                        if offset > 0:
                            with warnings.catch_warnings():
                                warnings.simplefilter(
                                    "ignore", RootNumpyUnconvertibleWarning)
                                recarray = tree2rec(tree,
                                                    entries=entries,
                                                    offset=offset,
                                                    selection=selection)
                            recarray = _drop_object_col(recarray, warn=False)
                            table.append(recarray)
                        else:
                            recarray = tree2rec(tree,
                                                entries=entries,
                                                offset=offset,
                                                selection=selection)
                            recarray = _drop_object_col(recarray)
                            if pbar is not None:
                                # start after any output from root_numpy
                                pbar.start()
                            table = hfile.createTable(group,
                                                      tree.GetName(), recarray,
                                                      tree.GetTitle())
                        offset += entries
                        if offset <= total_entries and pbar is not None:
                            pbar.update(offset)
                        # flush data in the table
                        table.flush()
                        # flush all pending data
                        hfile.flush()

                if pbar is not None:
                    pbar.finish()

            input_tree.Delete()

            if userfunc is not None:
                for tree in trees:
                    tree.Delete()
                tmp_file.Close()

    if own_h5file:
        hfile.close()
    if own_rootfile:
        rfile.Close()
Exemplo n.º 27
0
EXT_events  = 62666
NUMI_events = 290771
# POT: [4000,4500] : BNB 1.41E18, NUMI 2.46E18
# POT: [4500,4700] : BNB 3.10E18, NUMI 2.37E18
pot_NUMI = 2.46+2.37

TMIN = 3
TMAX = 23

PEMIN = 50
PEMAX = 800

f_numi = ROOT.TFile("/Users/tiamiceli/Data/TimingNeutrinos/beam_timing_run2_numi.root")
t_numi = f_numi.Get("_tree")
df_NUMI = pd.DataFrame(tree2rec(t_numi,selection='_pe_total > %i & _pe_total < %i & dt<%i & dt>%i'%(PEMIN,PEMAX,TMAX,TMIN)))

f_ext = ROOT.TFile("/Users/tiamiceli/Data/TimingNeutrinos/beam_timing_run2_ext.root")
t_ext = f_ext.Get("_tree")
print f_ext
print t_ext

df_EXT = pd.DataFrame(tree2rec(t_ext,selection='_pe_total > %i & _pe_total < %i & dt<%i & dt>%i'%(PEMIN,PEMAX,TMAX,TMIN)))

dt_NUMI = df_NUMI['_dt'].ravel()
dt_EXT = df_EXT['_dt'].ravel()

# 20 is the usec that we are saving flashed for


Exemplo n.º 28
0
def tree2hdf5(tree,
              hfile,
              group=None,
              entries=-1,
              selection=None,
              show_progress=False):
    """
    Convert a TTree into a HDF5 table.

    Parameters
    ----------

    tree : ROOT.TTree
        A ROOT TTree.

    hfile : string or PyTables HDF5 File
        A PyTables HDF5 File handle or string path to an existing HDF5 file.

    group : string or PyTables Group instance, optional (default=None)
        Write the table at this location in the HDF5 file.

    entries : int, optional (default=-1)
        The number of entries to read at once while converting a ROOT TTree
        into an HDF5 table. By default read the entire TTree into memory (this
        may not be desired if your TTrees are large).

    selection : string, optional (default=None)
        A ROOT selection expression to be applied on the TTree before
        conversion.

    show_progress : bool, optional (default=False)
        If True, then display and update a progress bar on stdout as the TTree
        is converted.

    """
    show_progress = show_progress and check_tty(sys.stdout)
    if show_progress:
        widgets = [Percentage(), ' ', Bar(), ' ', ETA()]

    own_h5file = False
    if isinstance(hfile, string_types):
        hfile = tables_open(filename=hfile, mode="w", title="Data")
        own_h5file = True

    log.info("Converting tree '{0}' with {1:d} entries ...".format(
        tree.GetName(), tree.GetEntries()))

    if not group:
        group = hfile.root
    elif isinstance(group, string_types):
        group_where = '/' + os.path.dirname(group)
        group_name = os.path.basename(group)
        if TABLES_NEW_API:
            group = hfile.create_group(group_where,
                                       group_name,
                                       createparents=True)
        else:
            group = hfile.createGroup(group_where, group_name)

    if tree.GetName() in group:
        log.warning("Tree '{0}' already exists "
                    "in the output file".format(tree.GetName()))
        return

    total_entries = tree.GetEntries()
    pbar = None
    if show_progress and total_entries > 0:
        pbar = ProgressBar(widgets=widgets, maxval=total_entries)

    if entries <= 0:
        # read the entire tree
        if pbar is not None:
            pbar.start()
        recarray = tree2rec(tree, selection=selection)
        recarray = _drop_object_col(recarray)
        if TABLES_NEW_API:
            table = hfile.create_table(group, tree.GetName(), recarray,
                                       tree.GetTitle())
        else:
            table = hfile.createTable(group, tree.GetName(), recarray,
                                      tree.GetTitle())
        # flush data in the table
        table.flush()
        # flush all pending data
        hfile.flush()
    else:
        # read the tree in chunks
        start = 0
        while start < total_entries or start == 0:
            if start > 0:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore",
                                          RootNumpyUnconvertibleWarning)
                    warnings.simplefilter("ignore", tables.NaturalNameWarning)
                    recarray = tree2rec(tree,
                                        selection=selection,
                                        start=start,
                                        stop=start + entries)
                recarray = _drop_object_col(recarray, warn=False)
                table.append(recarray)
            else:
                recarray = tree2rec(tree,
                                    selection=selection,
                                    start=start,
                                    stop=start + entries)
                recarray = _drop_object_col(recarray)
                if pbar is not None:
                    # start after any output from root_numpy
                    pbar.start()
                if TABLES_NEW_API:
                    table = hfile.create_table(group, tree.GetName(), recarray,
                                               tree.GetTitle())
                else:
                    table = hfile.createTable(group, tree.GetName(), recarray,
                                              tree.GetTitle())
            start += entries
            if start <= total_entries and pbar is not None:
                pbar.update(start)
            # flush data in the table
            table.flush()
            # flush all pending data
            hfile.flush()

    if pbar is not None:
        pbar.finish()

    if own_h5file:
        hfile.close()
Exemplo n.º 29
0
def root2hdf5(rfile, hfile, rpath='',
              entries=-1, userfunc=None,
              selection=None,
              show_progress=False):

    show_progress = show_progress and check_tty(sys.stdout)
    if show_progress:
        widgets = [Percentage(), ' ', Bar(), ' ', ETA()]

    own_rootfile = False
    if isinstance(rfile, basestring):
        rfile = root_open(rfile)
        own_rootfile = True

    own_h5file = False
    if isinstance(hfile, basestring):
        hfile = tables.openFile(filename=hfile, mode="w", title="Data")
        own_h5file = True

    for dirpath, dirnames, treenames in rfile.walk(
            rpath, class_pattern='TTree'):

        # skip root
        if not dirpath and not treenames:
            continue

        # skip directories w/o trees or subdirs
        if not dirnames and not treenames:
            continue

        where_group = '/' + os.path.dirname(dirpath)
        current_dir = os.path.basename(dirpath)

        if not current_dir:
            group = hfile.root
        else:
            group = hfile.createGroup(where_group, current_dir, "")

        ntrees = len(treenames)
        log.info(
            "Will convert {0:d} tree{1} in this directory".format(
                ntrees, 's' if ntrees != 1 else ''))

        for treename in treenames:

            input_tree = rfile.Get(os.path.join(dirpath, treename))

            if userfunc is not None:
                tmp_file = TemporaryFile()
                # call user-defined function on tree and get output trees
                log.info("Calling user function on tree '{0}'".format(
                    input_tree.GetName()))
                trees = userfunc(input_tree)

                if not isinstance(trees, list):
                    trees = [trees]

            else:
                trees = [input_tree]
                tmp_file = None

            for tree in trees:

                log.info("Converting tree '{0}' with {1:d} entries ...".format(
                    tree.GetName(),
                    tree.GetEntries()))

                if tree.GetName() in group:
                    log.warning(
                        "skipping tree '{0}' that already exists "
                        "in the output file".format(tree.GetName()))
                    continue

                total_entries = tree.GetEntries()
                pbar = None
                if show_progress and total_entries > 0:
                    pbar = ProgressBar(widgets=widgets, maxval=total_entries)

                if entries <= 0:
                    # read the entire tree
                    if pbar is not None:
                        pbar.start()
                    recarray = tree2rec(tree, selection=selection)
                    recarray = _drop_object_col(recarray)
                    table = hfile.createTable(
                        group, tree.GetName(),
                        recarray, tree.GetTitle())
                    # flush data in the table
                    table.flush()
                    # flush all pending data
                    hfile.flush()
                else:
                    # read the tree in chunks
                    start = 0
                    while start < total_entries or start == 0:
                        if start > 0:
                            with warnings.catch_warnings():
                                warnings.simplefilter(
                                    "ignore",
                                    RootNumpyUnconvertibleWarning)
                                recarray = tree2rec(
                                    tree,
                                    selection=selection,
                                    start=start,
                                    stop=start + entries)
                            recarray = _drop_object_col(recarray, warn=False)
                            table.append(recarray)
                        else:
                            recarray = tree2rec(
                                tree,
                                selection=selection,
                                start=start,
                                stop=start + entries)
                            recarray = _drop_object_col(recarray)
                            if pbar is not None:
                                # start after any output from root_numpy
                                pbar.start()
                            table = hfile.createTable(
                                group, tree.GetName(),
                                recarray, tree.GetTitle())
                        start += entries
                        if start <= total_entries and pbar is not None:
                            pbar.update(start)
                        # flush data in the table
                        table.flush()
                        # flush all pending data
                        hfile.flush()

                if pbar is not None:
                    pbar.finish()

            input_tree.Delete()

            if userfunc is not None:
                for tree in trees:
                    tree.Delete()
                tmp_file.Close()

    if own_h5file:
        hfile.close()
    if own_rootfile:
        rfile.Close()
Exemplo n.º 30
0
def chainer(
        dtype='cf',
        izip=1,
        base='/tera2/data3/cdmsbatsProd/R133/dataReleases/Prodv5-3_June2013/merged/',
        data='all',
        productions=['all'],
        rqs=[],
        eventrqs=[],
        rrqs=[],
        eventrrqs=[],
        cuts=[],
        eventcuts=[],
        selections=[],
        cutrev='current',
        load_cut_to_ram=False):

    # time call
    t1 = time()

    # deal with data chains
    dchain = ROOT.TChain()  # initialize data chain
    dlist = []
    # initialize first chain with calibebent trees (because they are small)
    dpaths = expander(
        dtype=dtype,
        tree='rrqDir/calibevent',
        base=base, data=data,
        productions=productions)
    map(dchain.Add, dpaths)
    # then make a list of chains for the other types
    for i, v in {
            'rrqDir/calibzip{}': rrqs,
            'rqDir/zip{}': rqs,
            'rqDir/eventTree': eventrqs}.iteritems():
        if len(v) != 0:
            tmp = ROOT.TChain()
            dpaths = expander(
                dtype=dtype,
                tree=i.format(izip),
                base=base,
                productions=productions)
            map(tmp.Add, dpaths)
            dlist.append(tmp)
    # friend each other data tree with the original chain
    map(dchain.AddFriend, dlist)

    # deal with cuts
    clist = {}
    for i, v in {
            'cutDir/cutzip{}': cuts,
            'cutDir/cutevent': eventcuts}.iteritems():
        for c in v:
            cpaths = expander(
                data='cuts',
                dtype=dtype,
                tree=i.format(izip),
                base=base,
                productions=productions,
                cut=c + '/',
                cutrev='current/')
            tmp = ROOT.TChain()
            #print "cpaths ", cpaths
            map(tmp.Add, cpaths)
            clist[c] = tmp
    #print "adding cuts: ", clist
    [dchain.AddFriend(v, k) for k, v in clist.iteritems()]

    # build cut selection
    cut_string = None
    if len(selections) != 0:
        cut_string = reduce(
            lambda x, y: x & y,
            map(
                Cut,
                selections))
    # extract the desired variables from the file turn into a Data Frame
    rows = ['SeriesNumber', 'EventNumber']
    branches = rrqs+rqs+eventrqs+eventrrqs+rows
    if load_cut_to_ram:
        branches += cuts+eventcuts
    df = pd.pivot_table(
        pd.DataFrame(
            tree2rec(
                dchain,
                branches=list(set(branches)),
                selection=cut_string)),
        rows=rows)
    t2 = time()
    print "Load time: ", t2-t1, "s"
    return df
Exemplo n.º 31
0
def get_event_ids(tree, cut):
    arr = root_numpy.tree2rec(tree, branches=["event_id"], selection=cut)["event_id"]
    arr.sort()
    return arr
Exemplo n.º 32
0
def convert(rfile, hfile, rpath='', entries=-1, userfunc=None, selection=None):

    isatty = check_tty(sys.stdout)
    if isatty:
        widgets = [Percentage(), ' ', Bar(), ' ', ETA()]

    if isinstance(hfile, basestring):
        hfile = tables.openFile(filename=hfile, mode="w", title="Data")
    if isinstance(rfile, basestring):
        rfile = ropen(rfile)

    for dirpath, dirnames, treenames in utils.walk(
            rfile, rpath, class_pattern='TTree'):

        # skip root
        if not dirpath and not treenames:
            continue

        # skip directories w/o trees or subdirs
        if not dirnames and not treenames:
            continue

        where_group = '/' + os.path.dirname(dirpath)
        current_dir = os.path.basename(dirpath)

        if not current_dir:
            group = hfile.root
        else:
            group = hfile.createGroup(where_group, current_dir, "")

        ntrees = len(treenames)
        if ntrees > 1:
            log.info("Will convert %i trees in this directory" % ntrees)
        else:
            log.info("Will convert 1 tree in this directory")

        for treename in treenames:

            tree = rfile.Get(os.path.join(dirpath, treename))

            if userfunc is not None:
                tmp_file = TemporaryFile()
                # call user-defined function on tree and get output trees
                log.info("Calling user function on tree '%s'" % tree.GetName())
                trees = userfunc(tree)

                if not isinstance(trees, list):
                    trees = [trees]
            else:
                trees = [tree]
                tmp_file = None

            for tree in trees:

                log.info("Converting tree '%s' with %i entries ..." % (
                    tree.GetName(),
                    tree.GetEntries()))

                total_entries = tree.GetEntries()
                pbar = None
                if isatty and total_entries > 0:
                    pbar = ProgressBar(widgets=widgets, maxval=total_entries)

                if entries <= 0:
                    # read the entire tree
                    if pbar is not None:
                        pbar.start()
                    recarray = tree2rec(tree, selection=selection)
                    table = hfile.createTable(
                        group, tree.GetName(),
                        recarray, tree.GetTitle())
                    table.flush()
                else:
                    # read the tree in chunks
                    offset = 0
                    while offset < total_entries or offset == 0:
                        if offset > 0:
                            with warnings.catch_warnings():
                                warnings.simplefilter("ignore",
                                        RootNumpyUnconvertibleWarning)
                                recarray = tree2rec(tree,
                                        entries=entries,
                                        offset=offset,
                                        selection=selection)
                            table.append(recarray)
                        else:
                            recarray = tree2rec(tree,
                                    entries=entries,
                                    offset=offset,
                                    selection=selection)
                            if pbar is not None:
                                # start after any output from root_numpy
                                pbar.start()
                            table = hfile.createTable(
                                group, tree.GetName(),
                                recarray, tree.GetTitle())
                        offset += entries
                        if offset <= total_entries and pbar is not None:
                            pbar.update(offset)
                        table.flush()

                if pbar is not None:
                    pbar.finish()

            if tmp_file is not None:
                tmp_file.Close()
Exemplo n.º 33
0
    print "You pressed Ctrl+C. Exiting!"
    sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)

# interactive matplotlib
plt.ion()

plt.rcParams.update({'font.size': 14})

fig, axarr = plt.subplots(2,2,figsize=(13,13))

f = ROOT.TFile(fin)

t = f.Get('out_tree')

arr = tree2rec(t,branches=['_michel_clustered_charge','_michel_Z','_michel_X','_X','_Z','_q_v','_s_v',
                           '_boundary','_event','_run','_subrun','_clus_idx','_forward'])

all_hits = False
if (f.GetListOfKeys().Contains('_hit_tree')):
    all_hits = True
    #hit tree
    hit_tree = f.Get('_hit_tree')
    # get hit information
    hits = pd.DataFrame(tree2rec(hit_tree,branches=['_run','_subrun','_event','_q_v','_w_v','_t_v']))
    # fill a dictionary that goes from (run,subrun,event) -> data frame with hit info for that event
    hitDict = {}
    for n in xrange(len(hits)):
        run    = hits['_run'][n]
        subrun = hits['_subrun'][n]
        event  = hits['_event'][n]
        hitDict[(run,subrun,event)] = hits.query('_event == %i and _run == %i and _subrun == %i'%(event,run,subrun))
Exemplo n.º 34
0
    
    outfname = "fulldata." + str(n) + ".csv.gz" 
    print "Current target is " + outfname
    if os.path.exists(workdir + outfname):
        print " already exists skipping."
        continue
    os.system("touch " + workdir + outfname)
    
    data = []
    for filename in filenames[i:j]:
        print "Reading " + filename.split('/')[-1]
        sys.stdout.flush()
        f = ROOT.TFile(filename)
        if f.IsZombie(): continue
        tree = f.Get("selections")
        data.append(pd.DataFrame(root_numpy.tree2rec(tree)))
    data = pd.concat(data)

    print "Writing " + str(len(data)) + " rows into " + outfname + " ... ",
    sys.stdout.flush()
    with gzip.open(outfname, 'w') as csvOut: data.to_csv(csvOut)  
    print " done."
    
    print "Copying ...",
    sys.stdout.flush()
    os.system("cp " + outfname + " " + workdir)
    print " done."
    
    print "Cleanup ...",
    sys.stdout.flush()
    os.system("rm " + outfname)
Exemplo n.º 35
0
        #print "cut %.3f sig_eff=%.3f Nbg=%i bg_eff=%.3f"%(cut_value,
        #                                                  Nsig_pass/float(Nsig),
        #                                                  Nbg_pass,
        #                                                  Nbg_pass/float(Nbg))
        results.append((method, Nsig_pass/float(Nsig), Nbg_pass/float(Nbg)))

    for res in sorted(results, key=operator.itemgetter(2)):
        print res
        
    
if __name__ == "__main__":
    import sys
    tmva_fname = sys.argv[1]
    f = R.TFile(tmva_fname)
    
    test_tree = f.TestTree
    train_tree = f.TrainTree

    test = root_numpy.tree2rec(test_tree)
    train = root_numpy.tree2rec(train_tree)

    bdt_name = "BDT_ada_640_2_1"
    for bdt_name in (bdt_name, "BDT_grad_40_2_1", "BDT_random_320_2"):
    #draw_all(test, train)

    #rank_all(test, sig_eff=0.9)

        draw_overtraining(bdt_name, test, train)
        
        draw_roccurve(bdt_name, test)
Exemplo n.º 36
0
tree = 'mytree'

# getting 5 branches, but we're only going to make a histogram of the E
branches = ['jet_AntiKt10LCTopo_%s' % col for col in ['E', 'pt', 'm', 'eta', 'phi']]

# set up bins for the histogram (10 GeV bins)
bins = np.arange(0.,2010.,10.)
# initialize to zero, since we're going to accumulate when we loop over
hist_data = np.zeros(len(bins)-1).astype(float)

# open up the file and grab the tree
f = TFile.Open(args.filename)
t = f.Get('%s/%s' % (directory,tree))

startTime_wall      = time.time()
startTime_processor = time.clock()

for event_num in xrange(args.event_start, args.event_start+args.num_events, args.step_size):
  data = rnp.tree2rec(t, branches=branches, start=(event_num), stop=(event_num+args.step_size))
  # grab 0th element since there's only one event
  for energies in data['jet_AntiKt10LCTopo_E']/1000.:
    # accumulate, np.histogram's first argument is the histogram data
    hist_data = map(add, hist_data, np.histogram(energies, bins=bins)[0] )

endTime_wall      = time.time()
endTime_processor = time.clock()
print "Finished job %d in:\n\t Wall time: %0.2f s \n\t Clock Time: %0.2f s" % (args.process_num, (endTime_wall - startTime_wall), (endTime_processor - startTime_processor))

# dump results into file uniquely named by the process number of Condor job
pickle.dump({'hist': hist_data, 'bins': bins}, file('hist_jet_AntiKt10LCTopo_E_%d.pkl' % args.process_num, 'w+') )
Exemplo n.º 37
0
# set up bins for the histogram (10 GeV bins)
bins = np.arange(0., 2010., 10.)
# initialize to zero, since we're going to accumulate when we loop over
hist_data = np.zeros(len(bins) - 1).astype(float)

# open up the file and grab the tree
f = TFile.Open(args.filename)
t = f.Get('%s/%s' % (directory, tree))

startTime_wall = time.time()
startTime_processor = time.clock()

for event_num in xrange(args.event_start, args.event_start + args.num_events,
                        args.step_size):
    data = rnp.tree2rec(t,
                        branches=branches,
                        start=(event_num),
                        stop=(event_num + args.step_size))
    # grab 0th element since there's only one event
    for energies in data['jet_AntiKt10LCTopo_E'] / 1000.:
        # accumulate, np.histogram's first argument is the histogram data
        hist_data = map(add, hist_data, np.histogram(energies, bins=bins)[0])

endTime_wall = time.time()
endTime_processor = time.clock()
print "Finished job %d in:\n\t Wall time: %0.2f s \n\t Clock Time: %0.2f s" % (
    args.process_num, (endTime_wall - startTime_wall),
    (endTime_processor - startTime_processor))

# dump results into file uniquely named by the process number of Condor job
pickle.dump({
    'hist': hist_data,
Exemplo n.º 38
0
def get_event_ids(tree, cut):
    arr = root_numpy.tree2rec(tree, branches=["event_id"],
                              selection=cut)["event_id"]
    arr.sort()
    return arr
Exemplo n.º 39
0
def root2hdf5(rfile,
              hfile,
              rpath='',
              entries=-1,
              userfunc=None,
              selection=None,
              show_progress=False):
    """
    Convert all trees in a ROOT file into tables in an HDF5 file.

    Parameters
    ----------

    rfile : string or asrootpy'd ROOT File
        A ROOT File handle or string path to an existing ROOT file.

    hfile : string or PyTables HDF5 File
        A PyTables HDF5 File handle or string path to an existing HDF5 file.

    rpath : string, optional (default='')
        Top level path to begin traversal through the ROOT file. By default
        convert everything in and below the root directory.

    entries : int, optional (default=-1)
        The number of entries to read at once while converting a ROOT TTree
        into an HDF5 table. By default read the entire TTree into memory (this
        may not be desired if your TTrees are large).

    userfunc : callable, optional (default=None)
        A function that will be called on every tree and that must return a
        tree or list of trees that will be converted instead of the original
        tree.

    selection : string, optional (default=None)
        A ROOT selection expression to be applied on all trees before
        conversion.

    show_progress : bool, optional (default=False)
        If True, then display and update a progress bar on stdout as each tree
        is converted.

    """
    show_progress = show_progress and check_tty(sys.stdout)
    if show_progress:
        widgets = [Percentage(), ' ', Bar(), ' ', ETA()]

    own_rootfile = False
    if isinstance(rfile, basestring):
        rfile = root_open(rfile)
        own_rootfile = True

    own_h5file = False
    if isinstance(hfile, basestring):
        hfile = tables_open(filename=hfile, mode="w", title="Data")
        own_h5file = True

    for dirpath, dirnames, treenames in rfile.walk(rpath,
                                                   class_pattern='TTree'):

        # skip root
        if not dirpath and not treenames:
            continue

        # skip directories w/o trees or subdirs
        if not dirnames and not treenames:
            continue

        where_group = '/' + os.path.dirname(dirpath)
        current_dir = os.path.basename(dirpath)

        if not current_dir:
            group = hfile.root
        else:
            group = hfile.createGroup(where_group, current_dir, "")

        ntrees = len(treenames)
        log.info("Will convert {0:d} tree{1} in this directory".format(
            ntrees, 's' if ntrees != 1 else ''))

        for treename in treenames:

            input_tree = rfile.Get(os.path.join(dirpath, treename))

            if userfunc is not None:
                tmp_file = TemporaryFile()
                # call user-defined function on tree and get output trees
                log.info("Calling user function on tree '{0}'".format(
                    input_tree.GetName()))
                trees = userfunc(input_tree)

                if not isinstance(trees, list):
                    trees = [trees]

            else:
                trees = [input_tree]
                tmp_file = None

            for tree in trees:

                log.info("Converting tree '{0}' with {1:d} entries ...".format(
                    tree.GetName(), tree.GetEntries()))

                if tree.GetName() in group:
                    log.warning("skipping tree '{0}' that already exists "
                                "in the output file".format(tree.GetName()))
                    continue

                total_entries = tree.GetEntries()
                pbar = None
                if show_progress and total_entries > 0:
                    pbar = ProgressBar(widgets=widgets, maxval=total_entries)

                if entries <= 0:
                    # read the entire tree
                    if pbar is not None:
                        pbar.start()
                    recarray = tree2rec(tree, selection=selection)
                    recarray = _drop_object_col(recarray)
                    if TABLES_NEW_API:
                        table = hfile.create_table(group, tree.GetName(),
                                                   recarray, tree.GetTitle())
                    else:
                        table = hfile.createTable(group, tree.GetName(),
                                                  recarray, tree.GetTitle())
                    # flush data in the table
                    table.flush()
                    # flush all pending data
                    hfile.flush()
                else:
                    # read the tree in chunks
                    start = 0
                    while start < total_entries or start == 0:
                        if start > 0:
                            with warnings.catch_warnings():
                                warnings.simplefilter(
                                    "ignore", RootNumpyUnconvertibleWarning)
                                warnings.simplefilter(
                                    "ignore", tables.NaturalNameWarning)
                                recarray = tree2rec(tree,
                                                    selection=selection,
                                                    start=start,
                                                    stop=start + entries)
                            recarray = _drop_object_col(recarray, warn=False)
                            table.append(recarray)
                        else:
                            recarray = tree2rec(tree,
                                                selection=selection,
                                                start=start,
                                                stop=start + entries)
                            recarray = _drop_object_col(recarray)
                            if pbar is not None:
                                # start after any output from root_numpy
                                pbar.start()
                            if TABLES_NEW_API:
                                table = hfile.create_table(
                                    group, tree.GetName(), recarray,
                                    tree.GetTitle())
                            else:
                                table = hfile.createTable(
                                    group, tree.GetName(), recarray,
                                    tree.GetTitle())
                        start += entries
                        if start <= total_entries and pbar is not None:
                            pbar.update(start)
                        # flush data in the table
                        table.flush()
                        # flush all pending data
                        hfile.flush()

                if pbar is not None:
                    pbar.finish()

            input_tree.Delete()

            if userfunc is not None:
                for tree in trees:
                    tree.Delete()
                tmp_file.Close()

    if own_h5file:
        hfile.close()
    if own_rootfile:
        rfile.Close()
Exemplo n.º 40
0
def test_tree2rec():
    chain = TChain('tree')
    chain.Add(load('single1.root'))
    check_single(rnp.tree2rec(chain))
Exemplo n.º 41
0
#!/usr/bin/python -i

from ROOT import *  #import it all who gives a hoot
import numpy as np
#from looks import   *
#from methods import *
from root_numpy import root2array, root2rec, tree2rec

#looks_minos()

tfile = TFile("outfile.root", "READ")
ttree = tfile.Get("data")

rec = tree2rec(ttree)

c1 = TCanvas()
c1.cd()

tg = TGraph()
tg.SetTitle("")

c = 0

for e in ttree:
    for o in xrange(len(e.y)):
        tg.SetPoint(c, e.x[o], e.y[o])
        c += 1
    tg.Draw("AL")
    tg.GetYaxis().SetRangeUser(0, 1.0)

    c1.Update()
Exemplo n.º 42
0
    outfname = "fulldata." + str(n) + ".csv.gz"
    print "Current target is " + outfname
    if os.path.exists(workdir + outfname):
        print " already exists skipping."
        continue
    os.system("touch " + workdir + outfname)

    data = []
    for filename in filenames[i:j]:
        print "Reading " + filename.split('/')[-1]
        sys.stdout.flush()
        f = ROOT.TFile(filename)
        if f.IsZombie(): continue
        tree = f.Get("selections")
        data.append(pd.DataFrame(root_numpy.tree2rec(tree)))
    data = pd.concat(data)

    print "Writing " + str(len(data)) + " rows into " + outfname + " ... ",
    sys.stdout.flush()
    with gzip.open(outfname, 'w') as csvOut:
        data.to_csv(csvOut)
    print " done."

    print "Copying ...",
    sys.stdout.flush()
    os.system("cp " + outfname + " " + workdir)
    print " done."

    print "Cleanup ...",
    sys.stdout.flush()
Exemplo n.º 43
0
                         lambda x:y1,\
                         lambda x:x*((1-y1)/(x3-x2)) + (1-x3*((1-y1)/(x3-x2))),\
                         lambda x:1\
                         ])

BNB_events  = 1127678 + 795004
EXT_events  = 36206 + 9551
NUMI_events = 276205 + 91004
# POT: [4000,4500] : BNB 1.41E18, NUMI 2.46E18
# POT: [4500,4700] : BNB 3.10E18, NUMI 2.37E18


#f = ROOT.TFile("beam_timing.root")
f = ROOT.TFile("/Users/tiamiceli/Data/TimingNeutrinos/beam_timing_run2_bnb.root")
t = f.Get("_tree")
df = pd.DataFrame(tree2rec(t))#,selection='_pe_total > 50 and _dt<23 and _dt>0'))

print df.shape

good_runs = [8317,8318,8319,8320,8321,8322,8324,8325,8329,8330,8331,8333,8336,8338,8346,8347,8348,8349,8350,8351,8352,8354,8355,8356,8357,8358,8383,8385,8386,8388,8393,8394]

good_string = '(_run == %i) '%good_runs[0]
for n in xrange(1,len(good_runs)):
    good_string += ' or (_run == %i)'%(good_runs[n])

print good_string

#df = df.query(good_string)

#df = df.query('(_run == 8317) or (_run == 8318)  or (_run == 8319) or (_run == 8320) or (_run == 8321) or (_run == 8322) or (_run == 8324) or (_run == 8325) or (_run == 8329) or (_run == 8330) or (_run == 8331) or (_run == 8333) or (_run == 8336) or (_run == 8338)')