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)
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
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")
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
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 )
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()
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)
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)
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)
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)))
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"
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.)))
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
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")
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)
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()
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()
# 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()
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()
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
def test_tree2rec():
chain = TChain('tree')
chain.Add(load('single1.root'))
check_single(rnp.tree2rec(chain))
# 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
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)
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:
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:
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()
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
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()
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()
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
def get_event_ids(tree, cut):
arr = root_numpy.tree2rec(tree, branches=["event_id"], selection=cut)["event_id"]
arr.sort()
return arr
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()
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))
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)
#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)
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+') )
# 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,
def get_event_ids(tree, cut):
arr = root_numpy.tree2rec(tree, branches=["event_id"],
selection=cut)["event_id"]
arr.sort()
return arr
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()
def test_tree2rec():
chain = TChain('tree')
chain.Add(load('single1.root'))
check_single(rnp.tree2rec(chain))
#!/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()
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()
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)')
