def recalculate_flightDistance(tree, key):
vertex_mask = (tree["TagVarCSV_vertexCategory"].array()
== 1) | (tree["TagVarCSV_vertexCategory"].array() == 2)
branch = tree[key].array().copy()
zeros = ak.broadcast_arrays(0., branch)[0]
ret_arr = ak.where(vertex_mask, zeros, branch)
return ak.to_awkward0(ret_arr)
def track_var_to_flat(arr, idx_low, idx_high):
# arr = np.ones( (len(arr), 200))
arr = awkward.to_awkward0(
awkward.Array([[[a if a != -9999. else -999. for a in ev[ka:kb + 1]]
for ka, kb in zip(kidx_a, kidx_b)]
for ev, kidx_a, kidx_b in zip(arr, idx_low, idx_high)
])).flatten(axis=0)
# arr = awkward.to_awkward0(awkward.Array( [[ev[ka : kb] if len( ev[ka : kb] ) > 0 else [0.] for ka, kb in zip(kidx_a, kidx_b) ] for ev, kidx_a, kidx_b in zip(arr, idx_low, idx_high)])).flatten(axis=0)
# for i, a in enumerate(arr):
# if len(a)== 0:
# from IPython import embed;embed()
# arr[i] = np.array([0.])
return arr
def track_var_to_flat(arr, idx_low, idx_high):
return awkward.to_awkward0(
awkward.Array([[ev[ka:kb] for ka, kb in zip(kidx_a, kidx_b)]
for ev, kidx_a, kidx_b in zip(arr, idx_low, idx_high)
])).flatten(axis=0)
def test_toawkward0():
array = ak.from_iter([1.1, 2.2, 3.3, 4.4], highlevel=False)
assert isinstance(ak.to_awkward0(array), np.ndarray)
assert ak.to_awkward0(array).tolist() == [1.1, 2.2, 3.3, 4.4]
array = ak.from_numpy(np.arange(2 * 3 * 5).reshape(2, 3, 5),
highlevel=False).toRegularArray()
assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray)
assert ak.to_awkward0(array).tolist() == [
[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14]],
[[15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29]],
]
array = ak.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False)
assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray)
assert ak.to_awkward0(array).tolist() == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
array = ak.layout.ListArray64(
ak.layout.Index64(np.array([4, 999, 1], dtype=np.int64)),
ak.layout.Index64(np.array([7, 999, 3], dtype=np.int64)),
ak.layout.NumpyArray(
np.array([3.14, 4.4, 5.5, 123, 1.1, 2.2, 3.3, 321])),
)
assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray)
assert ak.to_awkward0(array).tolist() == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
array = ak.from_iter(
[
{
"x": 0,
"y": []
},
{
"x": 1.1,
"y": [1]
},
{
"x": 2.2,
"y": [2, 2]
},
{
"x": 3.3,
"y": [3, 3, 3]
},
],
highlevel=False,
)
assert isinstance(ak.to_awkward0(array[2]), dict)
assert ak.to_awkward0(array[2])["x"] == 2.2
assert isinstance(ak.to_awkward0(array[2])["y"], np.ndarray)
assert ak.to_awkward0(array[2])["y"].tolist() == [2, 2]
assert isinstance(ak.to_awkward0(array), awkward0.Table)
assert ak.to_awkward0(array).tolist() == [
{
"x": 0,
"y": []
},
{
"x": 1.1,
"y": [1]
},
{
"x": 2.2,
"y": [2, 2]
},
{
"x": 3.3,
"y": [3, 3, 3]
},
]
array = ak.from_iter([(0, []), (1.1, [1]), (2.2, [2, 2]),
(3.3, [3, 3, 3])],
highlevel=False)
assert isinstance(ak.to_awkward0(array), awkward0.Table)
assert ak.to_awkward0(array).tolist() == [
(0, []),
(1.1, [1]),
(2.2, [2, 2]),
(3.3, [3, 3, 3]),
]
assert isinstance(ak.to_awkward0(array[2]), tuple)
assert ak.to_awkward0(array[2])[0] == 2.2
assert ak.to_awkward0(array[2])[1].tolist() == [2, 2]
array = ak.from_iter([0.0, [], 1.1, [1], 2.2, [2, 2], 3.3, [3, 3, 3]],
highlevel=False)
assert isinstance(ak.to_awkward0(array), awkward0.UnionArray)
assert ak.to_awkward0(array).tolist() == [
0.0,
[],
1.1,
[1],
2.2,
[2, 2],
3.3,
[3, 3, 3],
]
array = ak.from_iter([1.1, 2.2, None, None, 3.3, None, 4.4],
highlevel=False)
assert isinstance(ak.to_awkward0(array), awkward0.IndexedMaskedArray)
assert ak.to_awkward0(array).tolist() == [
1.1, 2.2, None, None, 3.3, None, 4.4
]
content = ak.layout.NumpyArray(
np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]))
index = ak.layout.Index64(np.array([3, 2, 2, 5, 0], dtype=np.int64))
array = ak.layout.IndexedArray64(index, content)
assert isinstance(ak.to_awkward0(array), awkward0.IndexedArray)
assert ak.to_awkward0(array).tolist() == [3.3, 2.2, 2.2, 5.5, 0.0]
# for key in keys:
# arr = ak.to_awkward0( ak.Array( [np.zeros( (20))[:np.random.randint(4, 20)] for i in range(100)] ) )
# dtype = ur.newbranch( np.dtype("f8"), size="{}-counts".format(key) )
# counts = arr.counts
# branch_dict["{}-counts".format(key)] = counts
# branch_reg[key] = dtype
# branch_dict[key] = arr
# out_file["ttree"] = ur.newtree(branch_reg)
# out_file["ttree"].extend(branch_dict)
arr = ak.to_awkward0( ak.Array([np.zeros( (20), dtype=np.dtype("f8"))[:np.random.randint(4, 20)] for i in range(100)] ) )
# arr = ak.to_awkward0( ak.Array( [np.zeros( (20)) for i in range(100)] ) )
# arr = ak.to_awkward0( [np.zeros( (20), np.dtype("f4")) for i in range(100)] )
# arr = ak.to_awkward0( [np.zeros( (20), np.dtype("f4")) for i in range(100)] )
# from IPython import embed;embed()
out_file["ttree"] = ur.newtree( {"key-a": ur.newbranch(np.dtype("f8"), size="n")})
out_file["ttree"].extend({"key-a": arr, "n": arr.counts})
# out_file["ttree"] = ur.newtree( {"key-a": np.float32})
# out_file["ttree"].extend( {"key-a": np.zeros( (200, 20), dtype=np.float32)} )
print("Trying to load the file with root_numpy")
nparr = rn.root2array( [f_path], treename="ttree", stop=None, branches=keys)