def createTrees(hist_cfg, out_dir, verbose=False, friend_func=None):
'''Writes out TTrees from histogram configuration for each contribution.
Takes list of variables attached to histogram config (hist_cfg.vars) to
create branches.
'''
plot = DataMCPlot(hist_cfg.name) # Used to cache TTrees
vcfgs = hist_cfg.vars
for cfg in hist_cfg.cfgs:
out_file = TFile('/'.join([out_dir, hist_cfg.name + '_' + cfg.name + '.root']), 'RECREATE')
out_tree = TTree('tree', '')
# Create branches for all variables
branches = [array('f', [0.]) for i in xrange(len(vcfgs))]
for branch_name, branch in zip([v.name for v in vcfgs], branches):
out_tree.Branch(branch_name, branch, branch_name+'/F')
# Create branch with full weight including lumi x cross section
full_weight = array('f', [0.])
out_tree.Branch('full_weight', full_weight, 'full_weight/F')
branches.append(full_weight)
total_scale = hist_cfg.total_scale if hist_cfg.total_scale else 1.
fillIntoTree(out_tree, branches, cfg, hist_cfg, vcfgs, total_scale, plot, verbose, friend_func)
out_file.cd()
out_tree.Write()
out_file.Write()
out_file.Close()
return plot
def staff():
staff = ROOT.staff_t()
# The input file cern.dat is a copy of the CERN staff data base
# from 1988
f = TFile('staff.root', 'RECREATE')
tree = TTree('T', 'staff data from ascii file')
tree.Branch('staff', staff,
'Category/I:Flag:Age:Service:Children:Grade:Step:Hrweek:Cost')
tree.Branch('Divisions', AddressOf(staff, 'Division'), 'Division/C')
tree.Branch('Nation', AddressOf(staff, 'Nation'), 'Nation/C')
# note that the branches Division and Nation cannot be on the first branch
fname = os.path.expandvars('$ROOTSYS/tutorials/tree/cernstaff.dat')
for line in open(fname).readlines():
t = filter(lambda x: x, re.split('\s+', line))
staff.Category = int(t[0]) # assign as integers
staff.Flag = int(t[1])
staff.Age = int(t[2])
staff.Service = int(t[3])
staff.Children = int(t[4])
staff.Grade = int(t[5])
staff.Step = int(t[6])
staff.Hrweek = int(t[7])
staff.Cost = int(t[8])
staff.Division = t[9] # assign as strings
staff.Nation = t[10]
tree.Fill()
tree.Print()
tree.Write()
def writeTree():
print ">>> Writing a tree."
# make new root file with new tree
file = TFile("tree.root", 'recreate')
tree = TTree("tree_name", "tree title")
# create 1 dimensional float arrays as fill variables, in this way the float array serves
# as a pointer which can be passed to the branch
n = array('d', [0])
u = array('d', [0])
# using numpy instead of array class (note python's float datatype corresponds to C++ doubles):
#n = numpy.zeros(1, dtype=float)
#u = numpy.zeros(1, dtype=float)
# create the branches and assign the fill-variables to them as doubles (D)
tree.Branch("normal", n, 'normal/D')
tree.Branch("uniform", u, 'uniform/D')
# create some random numbers, assign them into the fill variables and call Fill()
for i in xrange(10000):
n[0] = gRandom.Gaus()
u[0] = gRandom.Uniform()
tree.Fill()
# write the tree into the output file and close the file
file.Write()
file.Close()
def test05WriteSomeDataObjectBranched(self):
"""Test writing of a complex object across different branches"""
f = TFile(self.fname, 'RECREATE')
t = TTree(self.tname, self.ttitle)
d = SomeDataStruct()
# note: for p2.2, which has incomplete support of property types,
# we need to keep a reference alive to the result of the property
# call, or it will be deleted too soon; for later pythons, it is
# safe to use d.Floats directly in the Branch() call
fl = d.Floats
t.Branch('floats', fl)
t.Branch('nlabel', AddressOf(d, 'NLabel'), 'NLabel/I')
t.Branch('label', AddressOf(d, 'Label'), 'Label/C')
for i in range(self.N):
for j in range(self.M):
d.Floats.push_back(i * self.M + j)
d.NLabel = i
d.Label = '%d' % i
t.Fill()
f.Write()
f.Close()
def WriteTreeForEachScan():
filenamedf=pd.read_csv(dfdir+'filename.csv')
normvalue=np.loadtxt(dfdir+'normvalueinde.txt')
wavech = array('f',720*[0.])
intch = array('f',12*[0.])
fileid = array('i',[0])
for scanid in range(1,NbofscanID-1):
relatedpixels=longimap[:,scanid-1:scanid+2].reshape(12)
relatednorm=normvalue[:,scanid-1:scanid+2].reshape(12)
f = TFile(rootdir+"lwaveform"+str(scanid)+".root", 'recreate' )
tree = TTree( 'wavetree', 'wavetree' )
tree.Branch('wavech',wavech,'wavech[720]/F')
tree.Branch('intch',intch,'intch[12]/F')
tree.Branch('fileid',fileid,'fileid/I')
for fiid in range(0,33):
print(scanid,fiid)
fileid[0]=fiid
wavelist,intlist=ReadSingleRawText(filenamedf['scanstart'+str(scanID[scanid])].iloc[fiid],relatedpixels,relatednorm)
for itr in range(len(intlist)):
for k in range(12):
intch[k]=intlist[itr][k]
for k in range(720):
wavech[k]=wavelist[itr][k]
tree.Fill()
f.Write()
f.Close()
def WriteTreeForEachScan3by4():
Nbofevents = 10000
filenamedf = pd.read_csv(dfdir + 'filename.csv')
normvalue = np.loadtxt(dfdir + 'normvalueinde.txt')
filerange = [k for k in range(0, 33)]
Nbofwavepoint = 27
intch = array('f', 12 * [0.])
fileid = array('i', [0])
for scanid in range(1, NbofscanID - 1):
f = TFile(outdatadir + "lwaveform" + str(scanid) + "cutLED3by4.root",
'recreate')
tree = TTree('wavetree', 'wavetree')
# tree.Branch('wavech',wavech,'wavech[243]/F')
tree.Branch('intch', intch, 'intch[12]/F')
tree.Branch('fileid', fileid, 'fileid/I')
relatedpixels = longimap[:, scanid - 1:scanid + 2].reshape(12)
relatednorm = normvalue[:, scanid - 1:scanid + 2].reshape(12)
for fiid in filerange:
print(scanid, fiid)
fileid[0] = fiid
wavelist, intlist = ReadSingleRawText(
filenamedf['scanstart' + str(scanID[scanid])].iloc[fiid],
relatedpixels, relatednorm, Nbofwavepoint, Nbofevents)
for itr in range(len(intlist)):
for k in range(12):
intch[k] = intlist[itr][k]
# for k in range(243):
# wavech[k]=wavelist[itr][k]
tree.Fill()
f.Write()
f.Close()
def prep_AdSimple(host_file):
'''
Return a tuple of (adSimple, buffer) containing the TTree object and the
buffer used to fill its TBranches.
'''
buf = TreeBuffer()
buf.triggerType = unsigned_int_value()
buf.site = int_value()
buf.energy = float_value()
buf.x = float_value()
buf.y = float_value()
buf.z = float_value()
host_file.cd()
event_subdir = host_file.Get('Event')
if not bool(event_subdir):
event_subdir = host_file.mkdir('Event')
event_subdir.cd()
rec_subdir = event_subdir.Get('Rec')
if not bool(rec_subdir):
rec_subdir = event_subdir.mkdir('Rec')
rec_subdir.cd()
adSimple = TTree('AdSimple', 'Tree at /Event/Rec/AdSimple holding '
'Rec_AdSimple')
adSimple.Branch('context.mSite', buf.site, 'context.mSite/I')
adSimple.Branch('triggerType', buf.triggerType, 'triggerType/i')
adSimple.Branch('energy', buf.energy, 'energy/F')
adSimple.Branch('x', buf.x, 'x/F')
adSimple.Branch('y', buf.y, 'y/F')
adSimple.Branch('z', buf.z, 'z/F')
return adSimple, buf
def WriteTreeForEachScan3by4():
Nbofevents = 10000
filenamedf = pd.read_csv(dfdatadir + 'energyfilename.csv')
normvalue = np.loadtxt(dfdatadir + 'normvalueinde.txt')
filerange = [k for k in range(Nbofenergypos)]
intch = array('f', 12 * [0.])
fileid = array('i', [0])
relatedpixels = longienergymap[1:5, 3:6].reshape(12)
relatednorm = normvalue[1:5, 3:6].reshape(12)
for scanid in range(Nbofenergy):
f = TFile(rootdatadir + "lwaveform" + str(scanid) + "energy3by4.root",
'recreate')
tree = TTree('wavetree', 'wavetree')
# tree.Branch('wavech',wavech,'wavech[243]/F')
tree.Branch('intch', intch, 'intch[12]/F')
tree.Branch('fileid', fileid, 'fileid/I')
for fiid in filerange:
print(scanid, fiid)
fileid[0] = fiid
intarray = ReadSingleRawTxTReturnIntegral(
filenamedf['energy' + str(scanid)].iloc[fiid],
relatedpixels,
relatednorm,
Nbofevents,
thisledch=8,
thisbeamch=7)
for itr in range(intarray.shape[0]):
for k in range(12):
intch[k] = intarray[itr, k]
tree.Fill()
f.Write()
f.Close()
def WriteTreeForEachScan3by6():
Nbofevents = 1000
filenamedf = pd.read_csv(dfdatadir + 'filename.csv')
normvalue = np.loadtxt(dfdatadir + 'normvalueinde.txt')
intch = array('f', 18 * [0.])
fileid = array('i', [0])
for scanid in range(Nbofscan):
relatedpixels = longimap[scanid:scanid + 3, 1:7].reshape(18)
relatednorm = normvalue[scanid:scanid + 3, 1:7].reshape(18)
f = TFile(rootdatadir + "lwaveform" + str(scanid) + "energy3by6.root",
'recreate')
tree = TTree('wavetree', 'wavetree')
# tree.Branch('wavech',wavech,'wavech[243]/F')
tree.Branch('intch', intch, 'intch[18]/F')
tree.Branch('fileid', fileid, 'fileid/I')
for fiid in range(filenamedf.shape[0]):
print(scanid, fiid)
fileid[0] = fiid
intarray = ReadSingleRawTxTReturnIntegral(
filenamedf['scanstart' +
str(longimap[scanid + 1, 2])].iloc[fiid],
relatedpixels,
relatednorm,
Nbofevents,
thisledch=ledch,
thisbeamch=beamch)
for itr in range(intarray.shape[0]):
for k in range(18):
intch[k] = intarray[itr, k]
tree.Fill()
f.Write()
f.Close()
def test_minimizer(fin_name='gen_toy_10bin_highstat', fout_name='result'):
fout = TFile('deltas/' + fout_name + '.root', 'RECREATE')
tout = TTree('tree', 'tree')
bin = array('i', [0])
n = array('f', [0])
delta_n_u = array('f', [0.])
delta_n_d = array('f', [0.])
delta_m_u = array('f', [0.])
delta_m_d = array('f', [0.])
var_delta_m = array('f', [0.])
tout.Branch('bin', bin, 'bin/I')
tout.Branch('n', n, 'n/F')
tout.Branch('delta_n_u', delta_n_u, 'delta_n_u/F')
tout.Branch('delta_n_d', delta_n_d, 'delta_n_d/F')
tout.Branch('delta_m_u', delta_m_u, 'delta_m_u/F')
tout.Branch('delta_m_d', delta_m_d, 'delta_m_d/F')
tout.Branch('var_delta_m', var_delta_m, 'var_delta_m/F')
f = TFile.Open('deltas/' + fin_name + ".root")
tree = f.Get("tree")
tree.SetBranchStatus("*", True)
for iev in range(tree.GetEntries()):
tree.GetEntry(iev)
ev = tree
n_toy = ev.n_toy
n_bins = ev.n_bin
n_deltas = ev.n_deltas
deltas = ev.deltas
bin_n = ev.bin_n
bin_u = ev.bin_u
bin_d = ev.bin_d
print deltas
fix_params = []
#fix_params.extend([1, 2+3*n_bins])
#fix_params.extend(range(n_bins+2, 2*n_bins+2))
minimize_deltas = Minimze_deltas(deltas=deltas,
nbins=n_bins,
step=0.01,
fix_params=fix_params)
res = minimize_deltas.run()
for b in range(n_bins):
bin[0] = b + 1
n[0] = bin_n[b]
delta_n_u[0] = res[0][b]
delta_n_d[0] = res[1][b]
delta_m_u[0] = res[2][b]
delta_m_d[0] = res[3][b]
var_delta_m[0] = res[4][b]
#print "Bin %s: [%s, %s] ==> %s +/- %s" % ( b+1, delta_n_u[0], delta_n_d[0], abs(delta_m_u[0]), math.sqrt(var_delta_m[0]))
tout.Fill()
f.Close()
fout.cd()
fout.Write()
fout.Close()
def test_unsupported_branch_in_branches():
tree = TTree('test', 'test')
vect = TLorentzVector()
double = np.array([0], dtype=float)
tree.Branch('vector', vect)
tree.Branch('double', double, 'double/D')
rnp.tree2array(tree)
assert_raises(TypeError, rnp.tree2array, tree, branches=['vector'])
def init_tree_from_table(hit_table,
meta_table,
chunk_size=1,
hit_tree_entry=None,
meta_tree_entry=None):
''' Initializes a ROOT tree from a HDF5 table.
Takes the HDF5 table column names and types and creates corresponding
branches. If a chunk size is specified the branches will have the length of
the chunk size and an additional parameter is returned to change the chunk
size at a later stage. If a tree_entry is defined (a ROOT c-struct) the new
tree has the branches set to the corresponding tree entry address.
Parameters
----------
numpy_type_descriptor: np.dtype
'''
if (chunk_size > 1 and hit_tree_entry is not None
and meta_tree_entry is not None):
raise NotImplementedError()
tree = TTree('Table', 'Converted HDF5 table')
n_entries = None
if chunk_size > 1:
n_entries = ctypes.c_int(chunk_size) if chunk_size > 1 else 1
# needs to be added, otherwise one cannot access chunk_size_tree:
tree.Branch('n_entries', ctypes.addressof(n_entries), 'n_entries/I')
# adding branches for hit table:
for hit_column_name in hit_table.dtype.names:
print hit_column_name, hit_table.dtype[
hit_column_name], get_root_type_descriptor(
hit_table.dtype[hit_column_name])
tree.Branch(
hit_column_name, 'NULL' if hit_tree_entry is None else AddressOf(
hit_tree_entry, hit_column_name),
hit_column_name + '[n_entries]/' +
get_root_type_descriptor(hit_table.dtype[hit_column_name])
if chunk_size > 1 else hit_column_name + '/' +
get_root_type_descriptor(hit_table.dtype[hit_column_name]))
# adding branches for meta table (timestamps):
for meta_column_name in meta_table.dtype.names:
print meta_column_name, meta_table.dtype[
meta_column_name], get_root_type_descriptor(
meta_table.dtype[meta_column_name])
if meta_column_name == 'timestamp_start' or meta_column_name == 'timestamp_stop':
tree.Branch(
meta_column_name,
'NULL' if meta_tree_entry is None else AddressOf(
meta_tree_entry, meta_column_name),
meta_column_name + '[n_entries]/' +
get_root_type_descriptor(meta_table.dtype[meta_column_name])
if chunk_size > 1 else meta_column_name + '/' +
get_root_type_descriptor(meta_table.dtype[meta_column_name]))
return tree, n_entries
def cli():
helpString = "This script takes an input Sim file from Cosima and spits out a Root\
file which can be used in Alex MGeant analysis scripts. Takes as\
input the Simulation File name and the Root File name."
import argparse
parser = argparse.ArgumentParser(description=helpString)
parser.add_argument("SimFile", help="Input Simulation File (from Cosima)")
parser.add_argument("RootFile", help="Output Root File")
args = parser.parse_args()
from EventViewer import parse
from ROOT import TFile, TTree, AddressOf
sim = parse(args.SimFile)
print "There are {} events in this file.".format(len(sim.events))
#for event in sim.events[:10]:
# print event.hits
#Make the tree
f = TFile(args.RootFile, 'RECREATE')
t = TTree('h10','Simulations')
MGS = MGSimulation()
t.Branch('Runevt', AddressOf(MGS.MGStruct, 'Runevt'), 'Runevt/F')
t.Branch('Xcor', AddressOf(MGS.MGStruct, 'Xcor'), 'Xcor/F')
t.Branch('Ycor', AddressOf(MGS.MGStruct, 'Ycor'), 'Ycor/F')
t.Branch('Zcor', AddressOf(MGS.MGStruct, 'Zcor'), 'Zcor/F')
t.Branch('Estep', AddressOf(MGS.MGStruct, 'Estep'), 'Estep/F')
#t.Branch('Det', AddressOf(MGS.MGStruct, 'Det'), 'Det/F')
for event in sim.events:
hits = zip(event.hits.detector,
event.hits.x,
event.hits.y,
event.hits.z,
event.hits.energy)
for hit in hits:
MGS.MGStruct.Runevt = float(event.id_trigger)
#MGS.MGStruct.Det = hit[0]
MGS.MGStruct.Xcor = hit[1]
MGS.MGStruct.Ycor = hit[2]
MGS.MGStruct.Zcor = hit[3]
MGS.MGStruct.Estep = hit[4]
t.Fill()
f.Write()
f.Close()
def make_prim_tree():
#input
#inp = TFile.Open("../../lmon.root")
inp = TFile.Open("../../data/ew/ew1bx1.root")
tree = inp.Get("DetectorTree")
#number of events, negative for all
nev = -1
#output
out = TFile("prim_bx1.root", "recreate")
gROOT.ProcessLine("struct EntryF {Float_t v;};")
x = rt.EntryF()
y = rt.EntryF()
z = rt.EntryF()
en = rt.EntryF()
otree = TTree("prim_tree", "prim_tree")
otree.Branch("x", addressof(x, "v"), "x/F")
otree.Branch("y", addressof(y, "v"), "y/F")
otree.Branch("z", addressof(z, "v"), "z/F")
otree.Branch("en", addressof(en, "v"), "en/F")
hits = ew_hits("ew", tree)
if nev < 0: nev = tree.GetEntries()
for ievt in range(nev):
tree.GetEntry(ievt)
#print("Next event")
for ihit in range(hits.get_n()):
hit = hits.get_hit(ihit)
#hit by primary photon
if hit.prim == 0 or hit.pdg != 22: continue
hit.global_to_zpos(-18644) # mm
x.v = hit.x
y.v = hit.y
z.v = hit.z
z.en = hit.en
otree.Fill()
#print(hit.x, hit.y, hit.z) # , hit.en
#print(hit.pdg, hit.prim, hit.conv)
#only first hit by primary particle
break
otree.Write()
out.Close()
def main():
parser = argparse.ArgumentParser(
description=
'Program that takes as an argument the json -->.txt<-- file and makes a root file out of it, which can be used with the main framework.'
)
parser.add_argument('--input',
'-i',
required=True,
type=str,
help='Name of the json converted to a .txt file')
parser.add_argument('--output',
'-o',
type=str,
help='Name of the target .root file')
args = parser.parse_args()
infilename = os.path.abspath(args.input)
if not infilename.endswith('.txt'):
raise ValueError('Infilename does not end with .txt')
outfilename = os.path.abspath(
args.output) if args.output is not None else '.root'.join(
infilename.rsplit('.txt', 1))
with open(infilename, 'r') as f:
dict_in_json = safe_load(f)
outfile = TFile(outfilename, 'RECREATE')
outtree = TTree("LumiTree", "LumiTree")
t_run = array('i', [0])
t_lb_low = array('i', [0])
t_lb_high = array('i', [0])
outtree.Branch("run", t_run, 'run number/i')
outtree.Branch("lb_low", t_lb_low,
'lumi block low number (including this one)/i')
outtree.Branch("lb_high", t_lb_high,
'lumi block high number (including this one)/i')
runs = dict_in_json.keys()
for run in dict_in_json:
parts = dict_in_json[run]
# print parts
for lumiblocks in parts:
t_run[0] = int(run)
t_lb_low[0] = int(lumiblocks[0])
t_lb_high[0] = int(lumiblocks[1])
outtree.Fill()
# for lb in range(lumiblocks[0], lumiblocks[1]+1):
# t_run[0] = int(run)
# t_lb[0] = int(lb)
# outtree.Fill()
# print 'filled: (%i, %i)' % (int(run), int(lb))
outfile.Write()
outfile.Close()
print green('--> Successfully created lumifile \'%s\'' % (outfilename))
def WriteTreeForEachScan():
Nbofevents = 30000
filenamedf = pd.read_csv(dfdatadir + 'filename.csv')
normvalue = np.loadtxt(dfdatadir + 'normvalueinde.txt')
intch = array('f', 9 * [0.])
fileid = array('i', [0])
for scanid in range(Nbofscan):
f = TFile(rootdatadir + "lwaveform" + str(scanid) + "energy.root",
'recreate')
tree1 = TTree('wavetreerow1', 'wavetree')
# tree1.Branch('wavech',wavech,'wavech[243]/F')
tree1.Branch('intch', intch, 'intch[9]/F')
tree1.Branch('fileid', fileid, 'fileid/I')
tree2 = TTree('wavetreerow2', 'wavetree')
# tree2.Branch('wavech',wavech,'wavech[243]/F')
tree2.Branch('intch', intch, 'intch[9]/F')
tree2.Branch('fileid', fileid, 'fileid/I')
tree3 = TTree('wavetreerow3', 'wavetree')
# tree3.Branch('wavech',wavech,'wavech[243]/F')
tree3.Branch('intch', intch, 'intch[9]/F')
tree3.Branch('fileid', fileid, 'fileid/I')
for fiid in range(filenamedf.shape[0]):
if fiid < 15:
relatedpixels = longimap[scanid:scanid + 3, 1:4].reshape(9)
relatednorm = normvalue[scanid:scanid + 3, 1:4].reshape(9)
elif fiid < 33:
relatedpixels = longimap[scanid:scanid + 3, 2:5].reshape(9)
relatednorm = normvalue[scanid:scanid + 3, 2:5].reshape(9)
else:
relatedpixels = longimap[scanid:scanid + 3, 3:6].reshape(9)
relatednorm = normvalue[scanid:scanid + 3, 3:6].reshape(9)
print(scanid, fiid)
fileid[0] = fiid
intarray = ReadSingleRawTxTReturnIntegral(
filenamedf['scanstart' +
str(longimap[scanid + 1, 2])].iloc[fiid],
relatedpixels,
relatednorm,
Nbofevents,
thisledch=ledch,
thisbeamch=beamch)
for itr in range(intarray.shape[0]):
for k in range(9):
intch[k] = intarray[itr, k]
if fiid < 15:
tree1.Fill()
elif fiid < 33:
tree2.Fill()
else:
tree3.Fill()
f.Write()
f.Close()
def plot_xsec_contour(sample):
output_file_name = "lhe.root"
output_file = TFile(output_file_name, "RECREATE")
output_tree = TTree("Physics", "Physics")
gROOT.ProcessLine(
"struct MyStruct{ Float_t xsec; Float_t mdm; Float_t mhs; Float_t mzp; Float_t whs; Float_t wzp; };"
)
from ROOT import MyStruct
s = MyStruct()
output_tree.Branch('xsec', AddressOf(s, 'xsec'), 'xsec/F')
output_tree.Branch('mdm', AddressOf(s, 'mdm'), 'mdm/F')
output_tree.Branch('mhs', AddressOf(s, 'mhs'), 'mhs/F')
output_tree.Branch('mzp', AddressOf(s, 'mzp'), 'mzp/F')
output_tree.Branch('whs', AddressOf(s, 'whs'), 'whs/F')
output_tree.Branch('wzp', AddressOf(s, 'wzp'), 'wzp/F')
for i, su in enumerate(sample):
s.xsec = 0.
s.mdm = 0.
s.mhs = 0.
s.mzp = 0.
s.whs = 0.
s.wzp = 0.
with open('../samples/' + su + '.lhe', "rt") as fin:
for line in fin:
if "Integrated weight (pb)" in line:
#print line.split(' ')[-1]
s.xsec = float(line.split(' ')[-1])
if "# mdm" in line:
#print line.split(' ')[-3]
s.mdm = float(line.split(' ')[-3])
if "# mhs" in line:
#print line.split(' ')[-3]
s.mhs = float(line.split(' ')[-3])
if "# mzp" in line:
#print line.split(' ')[-3]
s.mzp = float(line.split(' ')[-3])
#hs
if "DECAY 54" in line:
#print line.split(' ')[-1]
s.whs = float(line.split(' ')[-1])
#zp
if "DECAY 55" in line:
#print line.split(' ')[-1]
s.wzp = float(line.split(' ')[-1])
output_tree.Fill()
output_tree.Write()
output_file.Close()
#plot contour
gROOT.Macro('Contour.C("lhe.root","BBbarDM_hs50")')
def create_singles_TTree(host_file):
from ROOT import TTree
host_file.cd()
title = 'Single events by Sam Kohn (git: {})'.format(
translate.git_describe())
out = TTree('singles', title)
buf = TreeBuffer()
initialize_basic_TTree(out, buf)
buf.fQuad = float_value()
buf.fMax = float_value()
buf.fPSD_t1 = float_value()
buf.fPSD_t2 = float_value()
buf.f2inch_maxQ = float_value()
buf.x = float_value()
buf.y = float_value()
buf.z = float_value()
buf.x_AdTime = float_value()
buf.y_AdTime = float_value()
buf.z_AdTime = float_value()
buf.fID = float_value()
buf.fPSD = float_value()
buf.num_nearby_events = unsigned_int_value()
buf.nearby_dt = int_value(10)
buf.nearby_energy = float_value(10)
def branch(name, typecode):
out.Branch(name, getattr(buf, name), '{}/{}'.format(name,
typecode))
return
branch('fQuad', 'F')
branch('fMax', 'F')
branch('fPSD_t1', 'F')
branch('fPSD_t2', 'F')
branch('f2inch_maxQ', 'F')
branch('x', 'F')
branch('y', 'F')
branch('z', 'F')
branch('x_AdTime', 'F')
branch('y_AdTime', 'F')
branch('z_AdTime', 'F')
branch('fID', 'F')
branch('fPSD', 'F')
branch('num_nearby_events', 'I')
out.Branch('nearby_dt', buf.nearby_dt, 'nearby_dt[num_nearby_events]/I')
out.Branch('nearby_energy', buf.nearby_energy,
'nearby_energy[num_nearby_events]/F')
return out, buf
def convert(path_in, path_out, mode):
for i in xrange(len(path_in)):
try:
f_in = TFile(path_in[i])
t_in = f_in.Get('save')
entries = t_in.GetEntries()
logging.info(mode[i] + ' entries :' + str(entries))
except:
logging.error(path_in[i] + ' is invalid!')
sys.exit()
print '--> Begin to process file: ' + path_in[i]
f_out = TFile(path_out[i], 'RECREATE')
t_out = TTree('save', 'save')
m_rm_D = array('d', [999.])
t_out.Branch('rm_D', m_rm_D, 'm_rm_D/D')
nentries = t_in.GetEntries()
for ientry in xrange(nentries):
t_in.GetEntry(ientry)
m_rm_D[0] = t_in.m_rm_D
t_out.Fill()
m_rm_D[0] = t_in.m_rm_Dmiss
t_out.Fill()
f_out.cd()
t_out.Write()
f_out.Close()
print '--> End of processing file: ' + path_out[i]
def WriteTree(runSet):
print 'Writing tree'
# Make a tree
froot = TFile('mu.root', 'RECREATE')
tree = TTree("mu", "Flat ntuple for mu")
tree.SetDirectory(froot)
leaves = "run/F:lb/F:mu/F"
leafValues = array("f", [0.0, 0.0, 0.0])
newBranch = tree.Branch("mu_vars", leafValues, leaves)
for r, lbs in runSet.iteritems():
for i in range(lbs.FirstLB(), lbs.LastLB()):
leafValues[2] = lbs.GetMeanPileUp(i)
leafValues[0] = float(r)
leafValues[1] = i
print 'mu: %s run: %s lb: %s' % (leafValues[2], leafValues[0],
leafValues[1])
tree.Fill()
tree.Write()
#froot.Write()
froot.Close()
def _to_root(self, data, filename):
import ROOT
from ROOT import TFile, TTree, gROOT, AddressOf
columns = [
self._root_column(data, i, c['type'], c['name'])
for i, c in enumerate(data['cols'])
]
header = 'struct data_t { ' + ';'.join(columns) + '; };'
gROOT.ProcessLine(header)
row = ROOT.data_t()
f = TFile(filename, 'RECREATE')
tree = TTree('data', 'data from RHAPI')
tree.Branch('data', row)
for r in data['data']:
for i, c in enumerate(data['cols']):
v = r[i]
if v is None:
if c['type'] == 'NUMBER': v = -1
else: v = ''
try:
setattr(row, c['name'], v)
except Exception as e:
print(c['name'], '=', v)
print(c, v)
print(e)
tree.Fill()
tree.Print()
tree.Write()
def main(args):
n_evs = args.n_evs
n_files = args.n_files
base_dir = args.base_dir
file_name_tmpl = "file{}.root"
tree_name = "DummyTree"
print("Produce {} TTrees with {} events each".format(n_files, n_evs))
for i in range(n_files):
f = TFile(base_dir + "/" + file_name_tmpl.format(i), 'recreate')
t = TTree(tree_name, tree_name)
x = array('d', [0.])
y = array('d', [0.])
t.Branch('x', x, 'x/D')
t.Branch('y', y, 'y/D')
for k in range(n_evs):
x[0] = ROOT.gRandom.Gaus(5, 2)
# uniform [0,10]
y[0] = ROOT.gRandom.Rndm() * 10
t.Fill()
f.Write()
f.Close()
def createDecorationTree(self, dataset, treeName=None):
from ROOT import TTree
# logging.info("Decorating %s, tree:%s with %s..."% \
# (self.dataset.name,treeName,self.decorationList[i].branch))
# get primary tree
# recursive call on sub-datasets
if isinstance(dataset,PhysicsProcess):
for d in dataset.datasets:
self.createDecorationTree(d, treeName)
return
if dataset.name.endswith("2015"):
print "Skipping DS %s to avoid duplicated MC tree in decor file for 2015/16 splitting"%dataset.name
return
dsTree = dataset._open(treeName)
friendTree = TTree("DecorationFriend_"+dataset.name, "DecorationFriend_"+dataset.name)
values={}
branches={}
for i in range(0,len(self.decorationList)):
values[i] = array('f',[0.])
# print "register values",values[i],hex(id(values[i]))
branches[i] = friendTree.Branch(self.decorationList[i].branch,values[i],"%s/F"%(self.decorationList[i].branch))
# print "register branches",branches[i],hex(id(branches[i]))
# decorate
for i in dsTree:
for j in range(0,len(self.decorationList)):
values[j][0] = self.calc(dsTree,j)
friendTree.Fill()
for j in range(0,len(self.decorationList)):
val=branches[j].GetLeaf(self.decorationList[j].branch).GetValue()
# print "getting branch val:", j, val
self.writeDecorationFile(friendTree,dataset,treeName)
def addBranchToDirectory():
print ">>> Add branch to directory."
file = TFile("tree.root", 'update')
# make directory if it does not exist
dir = file.GetDirectory("dir1")
if not dir:
print ">>> created dir1"
dir = file.mkdir("dir1")
# make this directory the current one: everything you write will be saved here
dir.cd()
# make new tree with a branch
tree = TTree("tree2", "tree2")
n = array('d', [0])
#n = numpy.zeros(1, dtype=float)
tree.Branch("normal5", n, 'normal5/D')
# fill tree
for i in xrange(10000):
n[0] = gRandom.Gaus(3, 2)
tree.Fill()
file.Write("", TFile.kOverwrite)
file.Close()
def main(args):
n_branches = args.n_branches
n_evs = args.n_evs
n_files = args.n_files
base_dir = args.base_dir
file_name_tmpl = "file{}_{}branches_{}evs.root"
tree_name = "DummyTree"
print("Produce {} TTrees with {} branches, {} events each".format(
n_files, n_branches, n_evs))
for i in range(n_files):
f = TFile(base_dir + "/" + file_name_tmpl.format(i, n_branches, n_evs),
'recreate')
t = TTree(tree_name, tree_name)
arrays = []
for n in range(n_branches):
branch_name = "branch_{}".format(n)
br = array('d', [0.])
arrays.append(br)
t.Branch(branch_name, br, '{}/D'.format(branch_name))
for k in range(n_evs):
for br in arrays:
br[0] = ROOT.gRandom.Gaus(5, 2)
t.Fill()
f.Write()
f.Close()
def do_toymc(self, sigtree, hrec, hout0):
print('We are doing toy mc...')
if self.ntoys < 0:
print('WARNING: ntoys < 0, set it to 1')
self.ntoys = 1
nevents = int(hrec.Integral())
list_hout = []
toymcfile = TFile('toymcfile.root', 'recreate')
for itoy in range(self.ntoys):
print('doing toymc ', itoy, '...')
datatree = TTree(self.data_treename, 'toymc_' + str(itoy))
mrec = array('f', [0])
datatree.Branch('mrec', mrec, 'mrec/F')
for i in range(nevents):
mrec[0] = hrec.GetRandom()
datatree.Fill()
datatree.Write()
x, hout, binning = self.unfold_core(sigtree,
datatree,
self.nsplit,
dotoymc=1)
list_hout.append(hout)
if self.debug or 0:
self.showhist(hout0)
for hout in list_hout:
print('toymc results', hout)
self.showhist(hout)
self.Vx = self.get_newVx_from_toymc(hout0, list_hout)
toymcfile.Close()
def create_simple_tree(nevents):
"""Create a simple TTree with a couple of branches."""
t = TTree('aTTree','A TTree')
t.Branch('run', AddressOf(pytree, 'run'),'run/I')
t.Branch('evt', AddressOf(pytree, 'evt'),'evt/I')
t.Branch('x', AddressOf(pytree, 'x'),'x/F')
t.Branch('y', AddressOf(pytree, 'y'),'y/F')
t.Branch('z', AddressOf(pytree, 'z'),'z/F')
for i in range(nevents):
pytree.run = 1 if (i<500) else 2
pytree.evt = i
pytree.x = gauss(0., 10.)
pytree.y = gauss(0, 10.)
pytree.z = gauss(5., 50.)
t.Fill()
return t
def setUp(self):
super(TQHistogramObservableTest, self).setUp()
f = TFile(os.path.join(self.tempdir, 'test.root'), 'recreate')
t = TTree('testTree', 'testTree')
d = array('f', [0.])
t.Branch('testVal', d, 'testVal/F')
for i in range(1000):
d[0] = random.gauss(0., 1.)
t.Fill()
f.Write()
f.Close()
tf = TFile.Open(os.path.join(self.tempdir, "testHistogram.root"),
"RECREATE")
histMap = TH1F('mapping', '', 2, -20., 20.)
histMap.SetBinContent(1, 0.5)
histMap.SetBinContent(2, 2)
#tf.cd()
histMap.Write()
tf.Close()
class TreeProducerCommon(object):
"""Class to create a custom output file & tree; as well as create and contain branches."""
def __init__(self, filename, module, **kwargs):
print 'TreeProducerCommon is called for', name
ncuts = kwargs.get('ncuts',25)
self.filename = filename
self.module = module
self.outfile = TFile(filename,'RECREATE')
self.tree = TTree('tree','tree')
self.cutflow = TH1D('cutflow','cutflow',ncuts,0,ncuts)
def addBranch(self, name, dtype='f', default=None, arrname=None):
"""Add branch with a given name, and create an array of the same name as address."""
if hasattr(self,name):
raise IOError("Branch of name '%s' already exists!"%(name))
if not arrname:
arrname = name
if isinstance(dtype,str):
if dtype.lower()=='f': # 'f' is only a 'float32', and 'F' is a 'complex64', which do not work for filling float branches
dtype = float # float is a 'float64' ('f8')
elif dtype.lower()=='i': # 'i' is only a 'int32'
dtype = int # int is a 'int64' ('i8')
setattr(self,arrname,np.zeros(1,dtype=dtype))
self.tree.Branch(name, getattr(self,arrname), '%s/%s'%(name,root_dtype[dtype]))
if default!=None:
getattr(self,name)[0] = default
def endJob(self):
"""Write and close files after the job ends."""
self.outfile.Write()
self.outfile.Close()
def ConvertParquet2Root(parquetFile, outputName, treename='tree'):
'''
Converts a pandas dataframe to a TTree and saves it in a root file.
Warning: this function is very slow.
Arguments
----------
- parquetFile: the absolute path of the file that you wanto to convert
- outputName: the name of the root file
'''
oFile = TFile(outputName, 'recreate')
tree = TTree(treename, treename)
df = pd.read_parquet(parquetFile, engine='pyarrow')
print(f'\n\033[94mConversion of {parquetFile}\033[0m')
data = []
for iCol, colname in enumerate(list(df)):
data.append(np.ones((1), dtype="float32"))
tree.Branch(colname, data[iCol], colname + "/F")
with alive_bar(len(df)) as bar:
for i in range(len(df)):
for iCol, colname in enumerate(list(df)):
data[iCol][0] = df.values[i, iCol]
tree.Fill()
bar()
oFile.Write()
