def prepareTrees(year):
d_sig = {}
d_bkg = {}
for signal in signals_original:
# if(opt.AC==False):
fname = eos_path_sig + '%i_MELA' % year
# else:
# fname = eos_path_sig + 'AC%i' %year
if (year == 2017) & (signal == 'VBFH125'):
fname += '/' + signal + 'ext/' + signal + 'ext_reducedTree_MC_' + str(
year) + '.root'
else:
fname += '/' + signal + '/' + signal + '_reducedTree_MC_' + str(
year) + '.root'
d_sig[signal] = uproot.open(fname)[key]
for bkg in bkgs:
fname = eos_path_sig + '%i_MELA' % year
# if year == 2016:
# fname += '_CorrectBTag'
if (year == 2018) & (bkg == 'ZZTo4lext'):
fname += '/' + bkg + '1/' + bkg + '1_reducedTree_MC_' + str(
year) + '.root'
else:
fname += '/' + bkg + '/' + bkg + '_reducedTree_MC_' + str(
year) + '.root'
d_bkg[bkg] = uproot.open(fname)[key]
return d_sig, d_bkg
def test_issue55(self):
withoffsets = uproot3.open(
"tests/samples/small-dy-withoffsets.root")["tree"]
nooffsets = uproot3.open(
"tests/samples/small-dy-nooffsets.root")["tree"]
assert numpy.array_equal(withoffsets.array("nJet"),
nooffsets.array("nJet"))
assert numpy.array_equal(withoffsets.array("nMuon"),
nooffsets.array("nMuon"))
def equal(left, right):
if len(left) != len(right):
return False
for x, y in zip(left, right):
if not numpy.array_equal(x, y):
return False
return True
assert equal(withoffsets.array("Jet_jetId"),
nooffsets.array("Jet_jetId"))
assert equal(withoffsets.array("Jet_pt"), nooffsets.array("Jet_pt"))
assert equal(withoffsets.array("MET_pt"), nooffsets.array("MET_pt"))
assert equal(withoffsets.array("Muon_charge"),
nooffsets.array("Muon_charge"))
assert equal(withoffsets.array("Muon_pt"), nooffsets.array("Muon_pt"))
assert equal(withoffsets.array("event"), nooffsets.array("event"))
def test_merge_root_pet_incrementRun(self):
try:
import uproot3 as uproot
except:
print("uproot3 is mandatory to merge root file. Please, do:")
print("pip install uproot3")
logger.info('Test_MergeRoot test_merge_root_pet')
tmpdirpath = tempfile.mkdtemp()
print(tmpdirpath)
filenameRoot = wget.download(
"https://gitlab.in2p3.fr/opengate/gatetools_data/-/raw/master/pet.root?inline=false",
out=tmpdirpath,
bar=None)
gt.merge_root([filenameRoot, filenameRoot],
os.path.join(tmpdirpath, "output.root"), True)
input = uproot.open(filenameRoot)
output = uproot.open(os.path.join(tmpdirpath, "output.root"))
inputTree = input[input.keys()[0]]
outputTree = output[output.keys()[0]]
inputRunBranch = inputTree.array(inputTree.keys()[0])
outputRunBranch = outputTree.array(outputTree.keys()[0])
self.assertTrue(2 * max(inputRunBranch) + 1 == max(outputRunBranch))
self.assertTrue(2 * len(inputRunBranch) == len(outputRunBranch))
inputEventBranch = inputTree.array(inputTree.keys()[1])
outputEventBranch = outputTree.array(outputTree.keys()[1])
self.assertTrue(max(inputEventBranch) == max(outputEventBranch))
self.assertTrue(2 * len(inputEventBranch) == len(outputEventBranch))
def test_compression_branches(self):
branches = list(
uproot3.open("tests/samples/Zmumu-uncompressed.root")
["events"].keys())
assert list(
uproot3.open("tests/samples/Zmumu-zlib.root")
["events"].keys()) == branches
assert list(
uproot3.open("tests/samples/Zmumu-lzma.root")
["events"].keys()) == branches
assert list(
uproot3.open("tests/samples/Zmumu-lz4.root")
["events"].keys()) == branches
assert list(
uproot3.open("tests/samples/Zmumu-zstd.root")
["events"].keys()) == branches
branches = list(
uproot3.open("tests/samples/HZZ-uncompressed.root")
["events"].keys())
assert list(
uproot3.open("tests/samples/HZZ-zlib.root")
["events"].keys()) == branches
assert list(
uproot3.open("tests/samples/HZZ-lzma.root")
["events"].keys()) == branches
assert list(
uproot3.open("tests/samples/HZZ-lz4.root")
["events"].keys()) == branches
assert list(
uproot3.open("tests/samples/HZZ-zstd.root")
["events"].keys()) == branches
def compare(p: Path, orig: Path, key: str):
f_new = uproot.open(p)["kernel"]
f_old = uproot.open(orig)["kernel"]
d_new = f_new[key].array()
d_old = f_old[key].array()
return plot(d_new, d_old, 2000)
def create_dataframe(file, family):
# --- import trees and variables
T_pot = uproot.open(file)['wcpselection/T_pot']
df_pot = T_pot.pandas.df(pot_vars, flatten=False)
T_KINE = uproot.open(file)['wcpselection/T_KINEvars']
df_KINE = T_KINE.pandas.df(kine_vars, flatten=False)
T_BDT = uproot.open(file)['wcpselection/T_BDTvars']
df_BDT = T_BDT.pandas.df(bdt_vars, flatten=False)
T_PFeval = uproot.open(file)['wcpselection/T_PFeval']
df_PFeval = T_PFeval.pandas.df(pfeval_vars, flatten=False)
T_eval = uproot.open(file)['wcpselection/T_eval']
df_eval = T_eval.pandas.df(eval_vars, flatten=False)
# --- merge dataframes
df = pd.concat([df_KINE, df_PFeval, df_BDT, df_eval, df_pot], axis=1)
# --- calculate POT
POT = sum(df_pot.pot_tor875)
# --- fix weight variables, make sure they are valid numbers
df.loc[df['weight_cv'] <= 0, 'weight_cv'] = 1
df.loc[df['weight_cv'] > 30, 'weight_cv'] = 1
df.loc[df['weight_cv'] == np.nan, 'weight_cv'] = 1
df.loc[df['weight_cv'] == np.inf, 'weight_cv'] = 1
df.loc[df['weight_cv'].isna(), 'weight_cv'] = 1
df.loc[df['weight_spline'] <= 0, 'weight_spline'] = 1
df.loc[df['weight_spline'] > 30, 'weight_spline'] = 1
df.loc[df['weight_spline'] == np.nan, 'weight_spline'] = 1
df.loc[df['weight_spline'] == np.inf, 'weight_spline'] = 1
df.loc[df['weight_spline'].isna(), 'weight_spline'] = 1
# --- calculate weight
if (family == 'NUE'): W_ = 1
elif (family == 'MC'): W_ = 1
df.loc[:, 'weight_genie'] = df['weight_cv'] * df[
'weight_spline'] # calculate GENIE weight
df.loc[:, 'weight'] = [
W_
] * df.shape[0] * df['weight_genie'] # should I POT normalise it?
# --- create variable to track file of origin
if (family == 'NUE'): df.loc[:, 'original_file'] = 0
elif (family == 'MC'): df.loc[:, 'original_file'] = 1
# --- delete dataframes to save memory space
del df_pot
del df_KINE
del df_BDT
del df_PFeval
del df_eval
return df, POT
def test_vector_of_numbers(self):
branch = uproot3.open("tests/samples/small-evnt-tree-fullsplit.root")["tree"]["StlVecU32"]
a = branch.array()
for i in range(100):
assert a[i].tolist() == [i] * (i % 10)
branch = uproot3.open("tests/samples/small-evnt-tree-fullsplit.root")["tree"]["StlVecF64"]
a = branch.array()
for i in range(100):
assert a[i].tolist() == [i] * (i % 10)
def test_compression_content2(self):
array = uproot3.open("tests/samples/HZZ-uncompressed.root"
)["events"].array("Electron_Px").tolist()
assert uproot3.open("tests/samples/HZZ-zlib.root")["events"].array(
"Electron_Px").tolist() == array
assert uproot3.open("tests/samples/HZZ-lzma.root")["events"].array(
"Electron_Px").tolist() == array
assert uproot3.open("tests/samples/HZZ-lz4.root")["events"].array(
"Electron_Px").tolist() == array
assert uproot3.open("tests/samples/HZZ-zstd.root")["events"].array(
"Electron_Px").tolist() == array
def prepareTrees(year):
d_sig = {}
d_sig_failed = {}
for signal in signals_original:
# if opt.AC_HYP: fname = eos_path_sig + 'AC%i_MELA' %year
# else: fname = eos_path_sig + '%i_MELA' %year
fname = eos_path_sig + '%i_MELA' %year
if not opt.NNLOPS: fname += '/'+signal+'/'+signal+'_reducedTree_MC_'+str(year)+'.root'
else: fname += '/'+signal+'_NNLOPS/'+signal+'_NNLOPS_reducedTree_MC_'+str(year)+'.root'
d_sig[signal] = uproot.open(fname)[key]
d_sig_failed[signal] = uproot.open(fname)[key_failed]
return d_sig, d_sig_failed
def test_compression_content1(self):
for name, array in uproot3.open(
"tests/samples/Zmumu-uncompressed.root")["events"].arrays(
["Type", "Event", "E1", "px1", "Q1", "M"]).items():
array = array.tolist()
assert uproot3.open("tests/samples/Zmumu-zlib.root"
)["events"].array(name).tolist() == array
assert uproot3.open("tests/samples/Zmumu-lzma.root"
)["events"].array(name).tolist() == array
assert uproot3.open("tests/samples/Zmumu-lz4.root"
)["events"].array(name).tolist() == array
assert uproot3.open("tests/samples/Zmumu-zstd.root"
)["events"].array(name).tolist() == array
def test_issue38(self):
before_hadd = uproot3.open(
"tests/samples/issue38a.root")["ntupler/tree"]
after_hadd = uproot3.open(
"tests/samples/issue38b.root")["ntupler/tree"]
before = before_hadd.arrays()
after = after_hadd.arrays()
assert set(before.keys())
assert set(after.keys())
for key in before.keys():
assert before[key].tolist() * 3 == after[key].tolist()
def test_issue21(self):
t = uproot3.open("tests/samples/issue21.root")["nllscan"]
### Explicit recover removed
# assert t.array("mH").tolist() == []
# t.recover()
assert t["mH"].numbaskets == 1
assert t["mH"].basket_entrystart(0) == 0
assert t["mH"].basket_entrystop(0) == 61
assert t["mH"].basket_numentries(0) == 61
assert t.array("mH").tolist() == [
124.0, 124.09089660644531, 124.18180084228516, 124.27269744873047,
124.36360168457031, 124.45449829101562, 124.54550170898438,
124.63639831542969, 124.72730255126953, 124.81819915771484,
124.87000274658203, 124.87550354003906, 124.88089752197266,
124.88639831542969, 124.89179992675781, 124.89730072021484,
124.90270233154297, 124.908203125, 124.90910339355469,
124.9135971069336, 124.91909790039062, 124.92449951171875,
124.93000030517578, 124.98739624023438, 124.9906997680664,
124.99349975585938, 124.99590301513672, 124.9977035522461,
124.9990005493164, 124.99970245361328, 125.0, 125.00029754638672,
125.0009994506836, 125.0022964477539, 125.00409698486328,
125.00650024414062, 125.0093002319336, 125.01260375976562,
125.06999969482422, 125.07550048828125, 125.08090209960938,
125.0864028930664, 125.09089660644531, 125.091796875,
125.09729766845703, 125.10269927978516, 125.10820007324219,
125.11360168457031, 125.11910247802734, 125.12449645996094,
125.12999725341797, 125.18180084228516, 125.27269744873047,
125.36360168457031, 125.45449829101562, 125.54550170898438,
125.63639831542969, 125.72730255126953, 125.81819915771484,
125.90910339355469, 126.0
]
def get_df(Class,
file,
xsecwt,
selection,
treepath=None,
branches=['ele*'],
multfactor=1):
tree = uproot.open(file)[treepath]
if debug:
ddd = tree.pandas.df(branches=branches,
flatten=flatten,
entrystop=1000).query(selection)
else:
ddd = tree.pandas.df(branches=branches,
flatten=flatten).query(selection)
#ddd["Category"]=Category
ddd["Class"] = Class
if type(xsecwt) == type(('xsec', 2)):
ddd["xsecwt"] = ddd[xsecwt[0]] * xsecwt[1]
elif type(xsecwt) == type("hello"):
ddd["xsecwt"] = ddd[xsecwt]
elif type(xsecwt) == type(0.1):
ddd["xsecwt"] = xsecwt
elif type(xsecwt) == type(1):
ddd["xsecwt"] = xsecwt
else:
print(
"CAUTION: xsecwt should be a branch name or a number or a tuple... Assigning the weight as 1"
)
print(file)
return ddd
def inputfile(nanofile):
tested = False
forceaaa = False
pfn = os.popen("edmFileUtil -d %s" % (nanofile)).read()
pfn = re.sub("\n", "", pfn)
print((nanofile, " -> ", pfn))
if (os.getenv("GLIDECLIENT_Group", "") != "overflow"
and os.getenv("GLIDECLIENT_Group", "") != "overflow_conservative"
and not forceaaa):
if not tested:
print("Testing file open")
testfile = uproot.open(pfn)
if testfile:
print("Test OK")
nanofile = pfn
else:
if "root://cms-xrd-global.cern.ch/" not in nanofile:
nanofile = "root://cms-xrd-global.cern.ch/" + nanofile
forceaaa = True
else:
nanofile = pfn
else:
if "root://cms-xrd-global.cern.ch/" not in nanofile:
nanofile = "root://cms-xrd-global.cern.ch/" + nanofile
return nanofile
def __init__(self, example_file):
# Define jet constants
self.delta_r = .4
self.delphes = uproot.open(example_file)['Delphes']
self.min_eta = -3
self.max_eta = 3
self.min_pt = {'q': 30., 'h': 200., 't': 200., 'W': 200}
self.settings = {
't': {
'id': 0,
'pid': [6],
'cut_m': [105., 210.]
},
'V': {
'id': 1,
'pid': [23, 24],
'cut_m': [65., 105.]
},
'H': {
'id': 2,
'pid': [25],
'cut_m': [105., 140.]
}
}
def test_strings_array(self):
branch = uproot3.open("tests/samples/sample-6.10.05-uncompressed.root")["sample"]["str"]
expectation = [b"hey-0", b"hey-1", b"hey-2", b"hey-3", b"hey-4", b"hey-5", b"hey-6", b"hey-7", b"hey-8", b"hey-9", b"hey-10", b"hey-11", b"hey-12", b"hey-13", b"hey-14", b"hey-15", b"hey-16", b"hey-17", b"hey-18", b"hey-19", b"hey-20", b"hey-21", b"hey-22", b"hey-23", b"hey-24", b"hey-25", b"hey-26", b"hey-27", b"hey-28", b"hey-29"]
cache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(cache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, cache=cache).tolist() == expectation[entrystart:entrystop]
assert len(cache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, cache=cache).tolist() == expectation[entrystart:entrystop]
cache = {}
basketcache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(basketcache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, basketcache=basketcache).tolist() == expectation[entrystart:entrystop]
assert len(basketcache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, basketcache=basketcache).tolist() == expectation[entrystart:entrystop]
basketcache = {}
keycache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(keycache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, keycache=keycache).tolist() == expectation[entrystart:entrystop]
assert len(keycache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, keycache=keycache).tolist() == expectation[entrystart:entrystop]
keycache = {}
def test_flat_array(self):
branch = uproot3.open("tests/samples/sample-6.10.05-uncompressed.root")["sample"]["i8"]
expectation = [-15, -14, -13, -12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
cache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(cache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, cache=cache).tolist() == expectation[entrystart:entrystop]
assert len(cache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, cache=cache).tolist() == expectation[entrystart:entrystop]
cache = {}
basketcache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(basketcache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, basketcache=basketcache).tolist() == expectation[entrystart:entrystop]
assert len(basketcache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, basketcache=basketcache).tolist() == expectation[entrystart:entrystop]
basketcache = {}
keycache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(keycache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, keycache=keycache).tolist() == expectation[entrystart:entrystop]
assert len(keycache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, keycache=keycache).tolist() == expectation[entrystart:entrystop]
keycache = {}
def tree_to_pandas(file_name,
tree_name,
columns,
exclude_columns="",
**kwargs): # pylint: disable=dangerous-default-value
"""
Parameters
----------
file_name : str
Path to root file
tree_name : str
Name of TTree
columns: sequence of str
Names of branches or aliases to be read. Can also use wildcards like ["*"]
for all branches or ["*fluc*"] for all branches containing the string "fluc".
exclude_columns: str
Regular expression of columns to be ignored, e.g. ".*Id".
Returns
--------
pandas.DataFrame
Data frame with specified columns
"""
with uproot3.open(file_name) as file:
data = file[tree_name].pandas.df(columns, **kwargs)
if exclude_columns != "":
data = data.filter([
col for col in data.columns
if not re.compile(exclude_columns).match(col)
])
return data
def writeOutPrediction(self, predicted, features, truth, weights,
outfilename, inputfile):
# predicted will be a list
spectator_branches = ['jet_pt', 'jet_eta']
from root_numpy import array2root
if inputfile[-5:] == 'djctd':
print(
"storing normed pt and eta... run on root files if you want something else for now"
)
spectators = features[0][:, 0:2].transpose()
else:
import uproot3 as uproot
print(inputfile)
urfile = uproot.open(inputfile)["deepntuplizer/tree"]
spectator_arrays = urfile.arrays(spectator_branches)
print(spectator_arrays)
spectators = [
spectator_arrays[a.encode()] for a in spectator_branches
]
out = np.core.records.fromarrays(
np.vstack(
(predicted[0].transpose(), truth[0].transpose(), spectators)),
names=
'prob_isB, prob_isBB, prob_isC,prob_isUDSG,isB, isBB, isC,isUDSG,jet_pt, jet_eta'
)
array2root(out, outfilename, 'tree')
def get_geometry_df(fname):
"""
Get pixel geometry from inputs made in `geometry/`
"""
import uproot3
import numpy as np
f = uproot3.open(fname)
t = f["idToGeo"]
df = t.pandas.df(branches=["shape", "translation", "matrix"],
flatten=False)
df["translation_x"] = df["translation"].str[0]
df["translation_y"] = df["translation"].str[1]
df["translation_z"] = df["translation"].str[2]
df["translation_rho"] = np.hypot(df["translation_x"], df["translation_y"])
df = df[df["shape"].apply(lambda x: x[0]) == 2.]
df["endcap"] = df.eval("abs(translation_z)>25")
# layer 1-4 for barrel, 5,7,8 for z disks
layers = []
for rho, z in df[["translation_rho", "translation_z"]].values:
ilayer = -1
if abs(z) < 25:
if 0 < rho < 5: ilayer = 1
if 5 < rho < 9: ilayer = 2
if 9 < rho < 14: ilayer = 3
if 14 < rho < 25: ilayer = 4
else:
if 25 < abs(z) < 35: ilayer = 5
if 35 < abs(z) < 45: ilayer = 6
if 45 < abs(z) < 70: ilayer = 7
layers.append(ilayer)
df["layer"] = np.array(layers)
df = df.query("translation_rho<18") # 4 pixel layers
return df
def get_config(sdir):
config = {}
folders = [f for f in listdir(sdir) if not isfile(join(sdir, f))]
for folder_name in folders:
print('Processing %s folder' % folder_name)
current_path = join(sdir, folder_name)
config[current_path] = {}
config[current_path]['files'] = {}
total_events_in_file = 0
files = [f for f in listdir(current_path) if f[-5:] == '.root']
for file_name in files:
try:
file_path = join(current_path, file_name)
rfile = uproot.open(file_path)
events = len(rfile['Delphes']['Event'].array())
config[current_path]['files'][file_path] = events
total_events_in_file = total_events_in_file + events
except ValueError:
pass
except KeyError:
pass
except OSError:
pass
config[current_path]['events'] = total_events_in_file
return config
def test_export1d():
with pytest.warns(FutureWarning):
import uproot3
import os
from coffea.hist.export import export1d
counts, test_eta, test_pt = dummy_jagged_eta_pt()
h_regular_bins = hist.Hist("regular_joe", hist.Bin("x", "x", 20, 0, 200))
h_regular_bins.fill(x=test_pt)
hout = export1d(h_regular_bins)
filename = 'test_export1d.root'
with uproot3.create(filename) as fout:
fout['regular_joe'] = hout
fout.close()
with uproot3.open(filename) as fin:
hin = fin['regular_joe']
assert (np.all(hin.edges == hout.edges))
assert (np.all(hin.values == hout.values))
del hin
del fin
if os.path.exists(filename):
os.remove(filename)
def test_selection_jagged_index(config_jagged_index, filename):
selection = filters.build_selection("test_selection_jagged", config_jagged_index)
infile = uproot3.open(filename)["events"]
mask = selection(infile, is_mc=False)
# Compare to: events->Draw("", "Muon_Px[1] > 0.300")
assert len(mask) == len(infile)
assert np.count_nonzero(mask) == 144
def test_issue371(self):
t = uproot3.open("tests/samples/issue371.root")["Event"]
obj = t["DRIFT_0."].array()[0]
assert obj._samplerName == b'DRIFT_0'
assert obj._n == 1
assert obj._energy[0] == numpy.array([2.3371024],
dtype=numpy.float32)[0]
def test_issue243_new(self):
t = uproot3.open("tests/samples/issue243-new.root")["triggerList"]
first = t["triggerMap.first"].array()
second = t["triggerMap.second"].array()
for i in range(t.numentries):
x = dict(zip(first[i], second[i]))
assert all(y == 1.0 for y in x.values())
def test_selection_1(config_1, filename):
selection = filters.build_selection("test_selection_1", config_1)
infile = uproot3.open(filename)["events"]
mask = selection(infile, is_mc=False)
assert np.count_nonzero(mask) == 289
columns = selection.columns
values = selection.values
index = selection.index_values
assert len(values) == 1
assert len(index) == 1
assert index[0][0] == "0"
assert index[0][1] == 0
assert index[0][2] == "NMuon > 1"
assert values[0][columns[0].index("passed_incl")] == 289
assert values[0][columns[0].index("passed_only_cut")] == 289
assert values[0][columns[0].index("totals_incl")] == 4580
df = selection.to_dataframe()
assert len(df) == 1
assert all(df.index.get_level_values("depth") == [0])
assert all(df.index.get_level_values("cut") == ["NMuon > 1"])
row = ("0", 0, "NMuon > 1")
assert df.loc[row, ("passed_incl", "unweighted")] == 289
assert df.loc[row, ("passed_only_cut", "unweighted")] == 289
assert df.loc[row, ("totals_incl", "unweighted")] == 4580
def test_issue124_and_followup_issue419_with_pr420(self):
f = uproot3.open("tests/samples/issue124.root")
branch = f[b'KM3NET_TIMESLICE;1'][b'KM3NET_TIMESLICE']
assert branch.interpretation is None
assert 0 == branch.compressedbytes()
assert 0 == branch.uncompressedbytes()
assert 0 == branch.numbaskets
def test_regular_array(self):
branch = uproot3.open("tests/samples/sample-6.10.05-uncompressed.root")["sample"]["ai8"]
expectation = [[-14, -13, -12], [-13, -12, -11], [-12, -11, -10], [-11, -10, -9], [-10, -9, -8], [-9, -8, -7], [-8, -7, -6], [-7, -6, -5], [-6, -5, -4], [-5, -4, -3], [-4, -3, -2], [-3, -2, -1], [-2, -1, 0], [-1, 0, 1], [0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6], [5, 6, 7], [6, 7, 8], [7, 8, 9], [8, 9, 10], [9, 10, 11], [10, 11, 12], [11, 12, 13], [12, 13, 14], [13, 14, 15], [14, 15, 16], [15, 16, 17]]
cache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(cache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, cache=cache).tolist() == expectation[entrystart:entrystop]
assert len(cache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, cache=cache).tolist() == expectation[entrystart:entrystop]
cache = {}
basketcache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(basketcache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, basketcache=basketcache).tolist() == expectation[entrystart:entrystop]
assert len(basketcache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, basketcache=basketcache).tolist() == expectation[entrystart:entrystop]
basketcache = {}
keycache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(keycache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, keycache=keycache).tolist() == expectation[entrystart:entrystop]
assert len(keycache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, keycache=keycache).tolist() == expectation[entrystart:entrystop]
keycache = {}
def test_irregular_array(self):
branch = uproot3.open("tests/samples/sample-6.10.05-uncompressed.root")["sample"]["Ai8"]
expectation = [[], [-15], [-15, -13], [-15, -13, -11], [-15, -13, -11, -9], [], [-10], [-10, -8], [-10, -8, -6], [-10, -8, -6, -4], [], [-5], [-5, -3], [-5, -3, -1], [-5, -3, -1, 1], [], [0], [0, 2], [0, 2, 4], [0, 2, 4, 6], [], [5], [5, 7], [5, 7, 9], [5, 7, 9, 11], [], [10], [10, 12], [10, 12, 14], [10, 12, 14, 16]]
assert [len(x) for x in expectation] == [0, 1, 2, 3, 4] * 6
cache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(cache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, cache=cache).tolist() == expectation[entrystart:entrystop]
assert len(cache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, cache=cache).tolist() == expectation[entrystart:entrystop]
cache = {}
basketcache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(basketcache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, basketcache=basketcache).tolist() == expectation[entrystart:entrystop]
assert len(basketcache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, basketcache=basketcache).tolist() == expectation[entrystart:entrystop]
basketcache = {}
keycache = {}
for entrystart, entrystop in [(None, None), (1, None), (1, 2), (1, 10), (10, 11), (10, 20), (6, 12), (6, 13)]:
assert len(keycache) == 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, keycache=keycache).tolist() == expectation[entrystart:entrystop]
assert len(keycache) > 0
assert branch.array(entrystart=entrystart, entrystop=entrystop, keycache=keycache).tolist() == expectation[entrystart:entrystop]
keycache = {}
def convertFromSourceFile(self, filename, weighterobjects, istraining, treename="Events"):
fileTimeOut(filename, 10)#10 seconds for eos to recover
tree = uproot.open(filename)[treename]
rechitcoll = RecHitCollection(use_true_muon_momentum=self.include_tracks,
cp_plus_pu_mode=self.cp_plus_pu_mode,
tree=tree)
#in a similar manner, we can also add tracks from conversions etc here
if self.include_tracks:
trackcoll = TrackCollection(tree=tree)
rechitcoll.append(trackcoll)
# adds t_is_unique
rechitcoll.addUniqueIndices()
# converts to DeepJetCore.SimpleArray
farr = rechitcoll.getFinalFeaturesSA()
t = rechitcoll.getFinalTruthDictSA()
return [farr,
t['t_idx'], t['t_energy'], t['t_pos'], t['t_time'],
t['t_pid'], t['t_spectator'], t['t_fully_contained'],
t['t_rec_energy'], t['t_is_unique'] ],[], []
