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)