def open_root_file(self):
print("going to open a root file!")
self.file_handle = uproot.recreate(self.filename)
self.file_handle["EVENT_NTUPLE"] = uproot.newtree(
{
"pulse_height":
uproot.newbranch(numpy.dtype(">i8"), size="hit_count"),
"chan":
uproot.newbranch(numpy.dtype(">i8"), size="hit_count"),
"timestamp":
uproot.newbranch(numpy.dtype(">i8"))
},
compression=None)
def Export2TTree(sOutputName, fits):
with uproot.recreate(sOutputName) as f:
f["Recon"] = uproot.newtree({
"X":
uproot.newbranch(float, title="X"),
"Y":
uproot.newbranch(float, title="Y"),
"Z":
uproot.newbranch(float, title="Z"),
"T":
uproot.newbranch(float, title="T"),
"Theta":
uproot.newbranch(float, title="Theta"),
"Phi":
uproot.newbranch(float, title="Phi"),
"MCID":
uproot.newbranch(int, title="MCID")
})
f["Recon"].extend({
"X": fits[:, 0],
"Y": fits[:, 1],
"Z": fits[:, 2],
"T": fits[:, 3],
"Theta": fits[:, 4],
"Phi": fits[:, 5],
"MCID": fits[:, 6]
})
def write_tree(branches, filename, treename):
"""
Write a TTree to a new ROOT file from a collection of arrays.
Parameters
----------
branches : dict
Dictionary of `branchname: branch_data` pairs.
filename : str
Pathname of new ROOT file.
treename : str
Name of new TTree.
"""
branch_definition_dictionary = dict()
branch_content_dictionary = dict()
for name, content in branches.items():
if isinstance(content, np.ndarray) and len(content.shape) == 1:
branch_definition_dictionary[name] = uproot.newbranch(
content.dtype)
elif isinstance(content.content, np.ndarray):
if content.content.dtype == np.dtype('int64'):
raise NotImplementedError(
'Jagged arrays of 64-bit integers are not yet'
' supported due to a known bug in the tree-writing'
' code'
' (https://github.com/scikit-hep/uproot/issues/506)'
'.')
size_name = name + '_n'
branch_definition_dictionary[name] = uproot.newbranch(
content.content.dtype, size=size_name)
branch_content_dictionary[size_name] = content.count()
else:
raise NotImplementedError('Branch type (' + str(type(content)) +
') not supported')
branch_content_dictionary[name] = content
with uproot.recreate(filename) as file:
file[treename] = uproot.newtree(branch_definition_dictionary)
file[treename].extend(branch_content_dictionary)
#h["deepntuplizer/tree"] =
with open("../config/test-data-params.json") as fh:
params = json.load(fh)
with open("../config/test-compare-params.json") as fh:
compare = json.load(fh)
num_lines = 3
newtree = {}
extender = {}
for feat in params["features"]:
newtree[feat] = uproot.newbranch(np.int32)
extender[feat] = np.tile([0], num_lines)
for label in params["labels"]:
newtree[label] = uproot.newbranch(np.int32)
extender[label] = np.tile([1], num_lines)
for spec in params["spectators"]:
newtree[spec] = uproot.newbranch(np.int32)
extender[spec] = np.tile([1], num_lines)
for jet in compare["jet_features"]:
newtree[jet] = uproot.newbranch(np.dtype(">i8"))
extender[jet] = awkward.array.jagged.JaggedArray.fromiter(
[[250.51, 250.51], [250.51, 250.51, 250.51], [250.51, 250.51, 250.51]])
#np.tile([250.51], num_lines)
import uproot
b1 = uproot.newbranch("i4", compression=uproot.ZLIB(5))
b2 = uproot.newbranch("i8", compression=uproot.LZMA(4))
b3 = uproot.newbranch("f4")
branchdict = {"branch1": b1, "branch2": b2, "branch3": b3}
tree = uproot.newtree(branchdict, compression=uproot.LZ4(4))
with uproot.recreate("example.root", compression=uproot.LZMA(5)) as f:
f["t"] = tree
f["t"].extend({
"branch1": [1] * 1000,
"branch2": [2] * 1000,
"branch3": [3] * 1000
})
def main():
# start run time clock
tstart = time.time()
# get options from command line
parser = OptionParser()
parser.add_option('-d', '--dataset', help='dataset', dest='dataset')
parser.add_option('-N', '--nFiles', help='nFiles', dest='nFiles', type=int, default=-1)
parser.add_option('-M', '--startFile', help='startFile', dest='startFile', type=int, default=0)
parser.add_option( '--condor', help='running on condor', dest='condor', default=False, action='store_true')
parser.add_option( '--dask', help='run w/ dask', dest='dask', default=False, action='store_true')
parser.add_option( '--port', help='port for dask status dashboard (localhost:port)', dest='port', type=int, default=8787)
parser.add_option( '--mincores', help='dask waits for min # cores', dest='mincores', type=int, default=4)
parser.add_option( '--quiet', help='suppress status printouts', dest='quiet', default=False, action='store_true')
parser.add_option('-w', '--workers', help='Number of workers to use for multi-worker executors (e.g. futures or condor)', dest='workers', type=int, default=8)
parser.add_option('-s', '--chunksize', help='Chunk size', dest='chunksize', type=int, default=10000)
parser.add_option('-m', '--maxchunks', help='Max number of chunks (for testing)', dest='maxchunks', type=int, default=None)
options, args = parser.parse_args()
# set output root file
sample = options.dataset
# getting dictionary of files from a sample collection e.g. "2016_QCD, 2016_WJets, 2016_TTJets, 2016_ZJets"
fileset = s.getFileset(sample, True, options.startFile, options.nFiles)
outfile = "MyAnalysis_%s_%d" % (sample, options.startFile) if options.condor or options.dask else "test"
# get processor args
exe_args = {'workers': options.workers, 'flatten': False}
if options.dask:
exe_args = use_dask(options.condor,options.workers,options.port)
if options.quiet: exe_args['status'] = False
client = exe_args['client']
while len(client.ncores()) < options.mincores:
print('Waiting for more cores to spin up, currently there are {0} available...'.format(len(client.ncores())))
print('Dask client info ->', client)
time.sleep(10)
sf = s.sfGetter(sample)
print("scaleFactor = {}".format(sf))
# run processor
output = processor.run_uproot_job(
fileset,
treename='TreeMaker2/PreSelection',
processor_instance=MainProcessor(sample,sf),
executor=processor.dask_executor if options.dask else processor.futures_executor,
executor_args=exe_args,
chunksize=options.chunksize,
maxchunks=options.maxchunks,
)
# export the histograms to root files
## the loop makes sure we are only saving the histograms that are filled
values_dict = {}
branchdict = {}
for v in output.keys():
if len(output[v].value) > 0:
branchdict[v] = uproot.newbranch("f4")
values_dict[v] = output[v].value
tree = uproot.newtree(branchdict)
if values_dict != {}:
print("saving root files...")
with uproot.recreate("{}.root".format(outfile)) as f:
f["tree"] = tree
f["tree"].extend(values_dict)
# print run time in seconds
dt = time.time() - tstart
print("run time: %.2f [sec]" % (dt))
def clone_tree(tree,
new_filename,
new_treename=None,
branches=None,
selection=None,
new_branches=None):
"""
Copy a TTree to a new ROOT file.
Parameters
----------
tree : TTree
TTree to copy from.
new_filename : str
Pathname of new ROOT file.
new_treename : str, optional
Name of new TTree. If `None`, the new tree receives the same name as the old tree.
branches : list or tuple of strings, optional
List of branchnames to copy. If `None`, all branches are copied.
selection : array_like, optional
. Boolean mask or int array of entry indices to copy. If `None`, all entries are copied.
new_branches : dict, optional
Dictionary of `branchname: branch_data` pairs to insert into the new tree.
"""
if new_treename is None:
new_treename = tree.name.decode('utf-8')
if branches is None:
branchnames = list(map(lambda b: b.decode('utf-8'), tree.keys()))
else:
branchnames = branches
branch_definition_dictionary = dict()
branch_content_dictionary = dict()
for name in branchnames:
branch = tree[name]
if isinstance(branch.interpretation.type, np.dtype):
branch_definition_dictionary[name] = uproot.newbranch(
branch.interpretation.type)
elif isinstance(branch.interpretation.content.type, np.dtype):
if branch.interpretation.content.type == np.dtype('int64'):
raise NotImplementedError(
'Jagged arrays of 64-bit integers are not yet supported'
' due to a known bug in the tree-writing code'
' (https://github.com/scikit-hep/uproot/issues/506).')
size_name = name + '_n'
while size_name in branchnames:
size_name += '0'
branch_definition_dictionary[name] = uproot.newbranch(
branch.interpretation.content.type, size=size_name)
if selection is not None:
branch_content_dictionary[size_name] = branch.array().count(
)[selection]
else:
branch_content_dictionary[size_name] = branch.array().count()
else:
raise NotImplementedError('Branch type (' +
str(branch.interpretation) +
') not supported')
if selection is not None:
content = branch.array()[selection]
else:
content = branch.array()
branch_content_dictionary[name] = content
if new_branches is not None:
for name, content in new_branches.items():
if isinstance(content, np.ndarray) and len(content.shape) == 1:
branch_definition_dictionary[name] = uproot.newbranch(
content.dtype)
elif isinstance(content.content, np.ndarray):
if content.content.dtype == np.dtype('int64'):
raise NotImplementedError(
'Jagged arrays of 64-bit integers are not yet'
' supported due to a known bug in the tree-writing'
' code'
' (https://github.com/scikit-hep/uproot/issues/506)'
'.')
size_name = name + '_n'
while size_name in branchnames + list(new_branches.keys()):
size_name += '0'
branch_definition_dictionary[name] = uproot.newbranch(
content.content.dtype, size=size_name)
branch_content_dictionary[size_name] = content.count()
else:
raise NotImplementedError('Branch type (' +
str(type(content)) +
') not supported')
branch_content_dictionary[name] = content
with uproot.recreate(new_filename) as new_file:
new_file[new_treename] = uproot.newtree(branch_definition_dictionary)
new_file[new_treename].extend(branch_content_dictionary)
def save_to_ROOT(self, data, filename='muons.root'):
''' Use uproot to save a generated array to a ROOT file that is compalible with MuonBackGenerator.cxx from FairShip'''
shape = np.shape(data)[0]
data[:,
3] += 2084.5 # Shift target to 50m. In accordance with primGen.SetTarget(ship_geo.target.z0+50*u.m,0.) in run_simScript.py
# The start of target in the GAN training data is -7084.5.
dtype = '>f4'
Event_ID = uproot.newbranch(dtype)
ID = uproot.newbranch(dtype)
Parent_ID = uproot.newbranch(dtype)
Pythia_ID = uproot.newbranch(dtype)
ECut = uproot.newbranch(dtype)
W = uproot.newbranch(dtype)
X = uproot.newbranch(dtype)
Y = uproot.newbranch(dtype)
Z = uproot.newbranch(dtype)
PX = uproot.newbranch(dtype)
PY = uproot.newbranch(dtype)
PZ = uproot.newbranch(dtype)
Release_Time = uproot.newbranch(dtype)
Mother_ID = uproot.newbranch(dtype)
Process_ID = uproot.newbranch(dtype)
branchdict = {
"event_id": Event_ID,
"id": ID,
"parentid": Parent_ID,
"pythiaid": Pythia_ID,
"ecut": ECut,
"w": W,
"x": X,
"y": Y,
"z": Z,
"px": PX,
"py": PY,
"pz": PZ,
"release_time": Release_Time,
"mother_id": Mother_ID,
"process_id": Process_ID
}
tree = uproot.newtree(branchdict, title="pythia8-Geant4")
with uproot.recreate("example.root") as f:
f["pythia8-Geant4"] = tree
f["pythia8-Geant4"].extend({
"event_id":
np.ones(shape).astype(np.float64),
"id":
np.array(np.ones(shape) * 13).astype(np.float64),
"parentid":
np.zeros(shape).astype(np.float64),
"pythiaid":
data[:, 0].astype(np.float64),
"ecut":
np.array(np.ones(shape) * 0.00001).astype(np.float64),
"w":
np.ones(shape).astype(np.float64),
"x":
np.array(data[:, 1] * 0.01).astype(np.float64),
"y":
np.array(data[:, 2] * 0.01).astype(np.float64),
"z":
np.array(data[:, 3] * 0.01).astype(np.float64),
"px":
data[:, 4].astype(np.float64),
"py":
data[:, 5].astype(np.float64),
"pz":
data[:, 6].astype(np.float64),
"release_time":
np.zeros(shape).astype(np.float64),
"mother_id":
np.array(np.ones(shape) * 99).astype(np.float64),
"process_id":
np.array(np.ones(shape) * 99).astype(np.float64)
})
# Not clear if all the datatype formatting is needed. Can be fiddly with ROOT datatypes. This works so I left it.
# Add buffer event at the end. This will not be read into simulation.
f["pythia8-Geant4"].extend({
"event_id": [0],
"id": [0],
"parentid": [0],
"pythiaid": [0],
"ecut": [0],
"w": [0],
"x": [0],
"y": [0],
"z": [0],
"px": [0],
"py": [0],
"pz": [0],
"release_time": [0],
"mother_id": [0],
"process_id": [0]
})
print(' ')
print(' ')
print('Saved', shape, 'muons to', filename, '.')
print('run_simScript.py must be run with the option: -n', shape,
'(or lower)')
print(' ')
print(' ')