def test_boxing3(): asnumpy = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) assert sys.getrefcount(asnumpy) == 2 aslayout = ak.layout.NumpyArray(asnumpy) aspart = ak.repartition(aslayout, 3, highlevel=False) asarray = ak.Array(aspart) aspart = asarray._layout assert ( sys.getrefcount(asnumpy), sys.getrefcount(aslayout), sys.getrefcount(aspart), ) == (3, 2, 3) @numba.njit def f3(x): return x, x for _ in range(10): out1, out2 = f3(asarray) assert isinstance(out1.layout, ak.partition.PartitionedArray) assert isinstance(out2.layout, ak.partition.PartitionedArray) assert ak.to_list(out1) == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] assert ak.to_list(out2) == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] assert ( sys.getrefcount(asnumpy), sys.getrefcount(aslayout), sys.getrefcount(aspart), ) == (3, 2, 3) del out1 del out2 del asarray del aspart del aslayout import gc gc.collect() assert sys.getrefcount(asnumpy) == 2
def test_11(one, two, three, tmp_path): filename = os.path.join(str(tmp_path), "test11.parquet") data = [ {"x": []}, {"x": [{"z": 1.1, "y": {"q": one}}]}, { "x": [ {"z": 1.1, "y": {"q": one}}, {"z": 2.2, "y": {"q": two}}, {"z": 3.3, "y": {"q": three}}, ] }, ] ak.to_parquet(ak.Array(data), filename) array = ak.from_parquet(filename, lazy=True, lazy_cache_key="tmp") assert set(array.caches[0].keys()) == set() array.layout.field("x").array assert len(set(array.caches[0].keys())) == 1 assert np.asarray(array.layout.field("x").array.offsets).tolist() == [0, 0, 1, 4] assert len(set(array.caches[0].keys())) == 1 array.layout.field("x").array.content.field("y").array assert len(set(array.caches[0].keys())) == 1 array.layout.field("x").array.content.field("y").array.field("q").array assert len(set(array.caches[0].keys())) == 2 array.layout.field("x").array.content.field("z").array assert len(set(array.caches[0].keys())) == 3 assert array.tolist() == data array = ak.from_parquet(filename, lazy=True, lazy_cache_key="tmp") assert set(array.caches[0].keys()) == set() array.layout.field("x").array assert len(set(array.caches[0].keys())) == 1 assert np.asarray(array.layout.field("x").array.offsets).tolist() == [0, 0, 1, 4] assert len(set(array.caches[0].keys())) == 1 array.layout.field("x").array.content.field("y").array assert len(set(array.caches[0].keys())) == 1 array.layout.field("x").array.content.field("z").array assert len(set(array.caches[0].keys())) == 2 array.layout.field("x").array.content.field("y").array.field("q").array assert len(set(array.caches[0].keys())) == 3 assert array.tolist() == data
def test_nested_virtualness(): counter = [0, 0] 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])) def materialize1(): counter[1] += 1 return content generator1 = ak.layout.ArrayGenerator(materialize1, form=content.form, length=len(content)) virtual1 = ak.layout.VirtualArray(generator1) offsets = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 10], dtype=np.int64)) listarray = ak.layout.ListOffsetArray64(offsets, virtual1) def materialize2(): counter[0] += 1 return listarray generator2 = ak.layout.ArrayGenerator(materialize2, form=listarray.form, length=len(listarray)) virtual2 = ak.layout.VirtualArray(generator2) array = ak.Array(virtual2) assert counter == [0, 0] @numba.njit def f1(x, i): return x[i] tmp1 = f1(array, 2) assert counter == [1, 0] tmp2 = f1(tmp1, 1) assert tmp2 == 4.4 assert counter == [1, 1]
def test_14(one, two, three, tmp_path): filename = os.path.join(str(tmp_path), "test6.parquet") data = [ {"x": [{"y": [], "z": 1.1}]}, {"x": []}, {"x": [{"y": [one, two, three], "z": 3.3}]}, ] ak.to_parquet(ak.Array(data), filename) array = ak.from_parquet(filename, lazy=True, lazy_cache_key="tmp") assert set(array.caches[0].keys()) == set() array.layout.field("x").array assert set(array.caches[0].keys()) == {"tmp:off:x.list.item.y.list.item:x[0]"} array.layout.field("x").array.content.field("z").array assert set(array.caches[0].keys()) == { "tmp:off:x.list.item.y.list.item:x[0]", "tmp:col:x.list.item.z[0]", } array.layout.field("x").array.content.field("y").array assert set(array.caches[0].keys()) == { "tmp:off:x.list.item.y.list.item:x[0]", "tmp:col:x.list.item.z[0]", "tmp:lst:x.list.item.y[0]", } assert array.tolist() == data array = ak.from_parquet(filename, lazy=True, lazy_cache_key="tmp") assert set(array.caches[0].keys()) == set() array.layout.field("x").array assert set(array.caches[0].keys()) == {"tmp:off:x.list.item.y.list.item:x[0]"} array.layout.field("x").array.content.field("y").array assert set(array.caches[0].keys()) == { "tmp:off:x.list.item.y.list.item:x[0]", "tmp:lst:x.list.item.y[0]", } array.layout.field("x").array.content.field("z").array assert set(array.caches[0].keys()) == { "tmp:off:x.list.item.y.list.item:x[0]", "tmp:lst:x.list.item.y[0]", "tmp:col:x.list.item.z[0]", } assert array.tolist() == data
def exclusive_subjets_up_to(self, data, nsub): try: px = data.px py = data.py pz = data.pz E = data.E except AttributeError: raise AttributeError("Lorentz vector not found") from None np_results = self._results.to_numpy_exclusive_subjets_up_to(px, py, pz, E, nsub) return ak.Array( ak.layout.RecordArray( [ ak.layout.NumpyArray(np_results[0]), ak.layout.NumpyArray(np_results[1]), ak.layout.NumpyArray(np_results[2]), ak.layout.NumpyArray(np_results[3]), ], ["px", "py", "pz", "E"], parameters={"__record__": "Momentum4D"}, ), behavior=self.data.behavior, )
def nestedindex(stack): """Concatenate a list of indices along a new axis Signature: index1,index2,...,!nestedindex Index arrays should all be same shape flat arrays Outputs a jagged array with same outer shape as index arrays """ indexers = stack[:] stack.clear() # return awkward.concatenate([idx[:, None] for idx in indexers], axis=1) n = len(indexers) out = numpy.empty(n * len(indexers[0]), dtype="int64") for i, idx in enumerate(indexers): out[i::n] = idx offsets = numpy.arange(0, len(out) + 1, n, dtype=numpy.int64) out = awkward.Array( awkward.layout.ListOffsetArray64( awkward.layout.Index64(offsets), awkward.layout.NumpyArray(out), ) ) stack.append(out)
def test_argcombinations(): array = ak.Array([[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6], [7.7], [8.8, 9.9, 10.0, 11.1, 12.2]]) assert ak.to_list(ak.argcombinations(array, 2, replacement=False)) == [ [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)], [], [(0, 1), (0, 2), (1, 2)], [], [ (0, 1), (0, 2), (0, 3), (0, 4), (1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4), ], ]
def _createSpectators(self, tree): starttime = time.time() recHitX = self._readAndSplit(tree,"RecHitHGC_x") recHitY = self._readAndSplit(tree,"RecHitHGC_y") recHitZ = self._readAndSplit(tree,"RecHitHGC_z") recHitSimClusIdx = self._readAndSplit(tree,"RecHitHGC_BestMergedSimClusterIdx") #Define spectators recHit_df_events = [pd.DataFrame({"recHitX":recHitX[i], "recHitY":recHitY[i], "recHitZ":recHitZ[i], "recHitSimClusIdx":recHitSimClusIdx[i] }) for i in range(len(recHitX))] for ievent in range(len(recHit_df_events)): df_event = recHit_df_events[ievent] unique_shower_idx = np.unique(df_event['recHitSimClusIdx']) df_event['spectator_distance'] = 0. # df_event['recHitSimClus_nHits'] = df_event.groupby('recHitSimClusIdx').recHitX.transform(len) #adding number of rec hits that are associated to this truth cluster for idx in unique_shower_idx: df_shower = df_event[df_event['recHitSimClusIdx']==idx] x_to_fit = df_shower[['recHitX','recHitY','recHitZ']].to_numpy() spectators_shower_dist = None try: spectators_shower_dist = find_pcas(x_to_fit,PCA_n=2,min_hits=10) except: pass if (spectators_shower_dist is not None) : spectators_idx = (df_shower.index.tolist()) df_event.loc[spectators_idx,'spectator_distance'] = spectators_shower_dist del df_shower del df_event print('spectators calculated after',time.time()-starttime,'s') recHitSpectatorFlag = ak1.Array([np.expand_dims(recHit_df_events[i]['spectator_distance'].to_numpy(),axis=1) for i in range(len(recHit_df_events))]) print('ended spectators after', time.time()-starttime,'s') return recHitSpectatorFlag
def test(): a = ak._v2.operations.structure.values_astype(ak.Array([1, None]), np.float32) assert ak._v2.operations.structure.fill_none( a, np.float32(0)).tolist() == [1, 0] assert (str(ak._v2.operations.structure.fill_none( a, np.float32(0)).type) == "2 * float32") assert ak._v2.operations.structure.fill_none(a, np.array( 0, np.float32)).tolist() == [1, 0] assert (str( ak._v2.operations.structure.fill_none(a, np.array( 0, np.float32)).type) == "2 * float32") assert ak._v2.operations.structure.fill_none(a, np.array( [0], np.float32)).tolist() == [1, [0]] assert (str( ak._v2.operations.structure.fill_none(a, np.array( [0], np.float32)).type) == "2 * union[float32, 1 * float32]") assert ak._v2.operations.structure.fill_none(a, np.array( [[0]], np.float32)).tolist() == [1, [[0]]] assert (str( ak._v2.operations.structure.fill_none(a, np.array( [[0]], np.float32)).type) == "2 * union[float32, 1 * 1 * float32]") assert ak._v2.operations.structure.fill_none(a, 0).tolist() == [1, 0] assert str(ak._v2.operations.structure.fill_none(a, 0).type) == "2 * float64" assert ak._v2.operations.structure.fill_none(a, [0]).tolist() == [1, [0]] assert (str(ak._v2.operations.structure.fill_none( a, [0]).type) == "2 * union[float32, 1 * int64]") assert ak._v2.operations.structure.fill_none(a, [[0]]).tolist() == [1, [[0]]] assert (str(ak._v2.operations.structure.fill_none( a, [[0]]).type) == "2 * union[float32, 1 * var * int64]")
def test_argsort_2d(): array = ak.Array([[1, 2, None, 3, 0, None], [1, 2, None, 3, 0, None]]) assert ak.argsort(array).tolist() == [[4, 0, 1, 3, 2, 5], [4, 0, 1, 3, 2, 5]] assert array[ak.argsort(array)].tolist() == [ [ 0, 1, 2, 3, None, None, ], [ 0, 1, 2, 3, None, None, ], ]
def test_toregular(): array = ak.Array([{ "x": np.arange(2 * 3 * 5).reshape(2, 3, 5).tolist(), "y": np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5, 7), }]) assert str(array.type) in ( '1 * {"x": var * var * var * int64, "y": var * var * var * var * int64}', '1 * {"y": var * var * var * var * int64, "x": var * var * var * int64}', ) assert str(ak.to_regular(array, axis=-1).type) in ( '1 * {"x": var * var * 5 * int64, "y": var * var * var * 7 * int64}', '1 * {"y": var * var * var * 7 * int64, "x": var * var * 5 * int64}', ) assert str(ak.to_regular(array, axis=-2).type) in ( '1 * {"x": var * 3 * var * int64, "y": var * var * 5 * var * int64}', '1 * {"y": var * var * 5 * var * int64, "x": var * 3 * var * int64}', ) assert str(ak.to_regular(array, axis=-3).type) in ( '1 * {"x": 2 * var * var * int64, "y": var * 3 * var * var * int64}', '1 * {"y": var * 3 * var * var * int64, "x": 2 * var * var * int64}', )
def test(): array = ak.Array([{ "x": "one" }, { "x": "two" }, { "x": "three" }], check_valid=True) assert ak.to_list(array) == [{"x": "one"}, {"x": "two"}, {"x": "three"}] assert ak.to_list(ak.from_iter(ak.to_list(array))) == [ { "x": "one" }, { "x": "two" }, { "x": "three" }, ] assert ak.to_list(array.layout) == [{ "x": "one" }, { "x": "two" }, { "x": "three" }] assert ak.to_list(ak.from_iter(ak.to_list(array.layout))) == [ { "x": "one" }, { "x": "two" }, { "x": "three" }, ]
def test_9(one, two, three, tmp_path): filename = os.path.join(str(tmp_path), "test9.parquet") data = [ {"x": []}, {"x": [{"y": {"q": one}}]}, {"x": [{"y": {"q": one}}, {"y": {"q": two}}, {"y": {"q": three}}]}, ] ak.to_parquet(ak.Array(data), filename) array = ak.from_parquet(filename, lazy=True, lazy_cache_key="tmp") assert set(array.caches[0].keys()) == set() array.layout.field("x").array assert set(array.caches[0].keys()) == {"tmp:off:x.list.item.y.q:x[0]"} assert np.asarray(array.layout.field("x").array.offsets).tolist() == [0, 0, 1, 4] assert set(array.caches[0].keys()) == {"tmp:off:x.list.item.y.q:x[0]"} array.layout.field("x").array.content.field("y").array assert set(array.caches[0].keys()) == {"tmp:off:x.list.item.y.q:x[0]"} array.layout.field("x").array.content.field("y").array.field("q").array assert set(array.caches[0].keys()) == { "tmp:off:x.list.item.y.q:x[0]", "tmp:col:x.list.item.y.q[0]", } assert array.tolist() == data
def get_child(self, data): try: px = data.px py = data.py pz = data.pz E = data.E except AttributeError: raise AttributeError("Lorentz vector not found") from None np_results = self._results.to_numpy_get_child(px, py, pz, E) return ak.Array( ak.layout.RecordArray( ( ak.layout.NumpyArray(np_results[0]), ak.layout.NumpyArray(np_results[1]), ak.layout.NumpyArray(np_results[2]), ak.layout.NumpyArray(np_results[3]), ), ("px", "py", "pz", "E"), parameters={"__record__": "Momentum4D"}, ), behavior=self.data.behavior, )
def test_union(): one = ak.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False) two = ak.from_iter([[6.6], [7.7, 8.8], [], [9.9, 10.0, 11.1, 12.2]], highlevel=False) tags = ak.layout.Index8(np.array([0, 0, 0, 1, 1, 1, 1], dtype=np.int8)) index = ak.layout.Index64(np.array([0, 1, 2, 0, 1, 2, 3], dtype=np.int64)) unionarray = ak.layout.UnionArray8_64(tags, index, [one, two]) assert ak.to_list(unionarray) == [ [1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8], [], [9.9, 10.0, 11.1, 12.2], ] assert ak.to_list( unionarray[ak.Array([[0, -1], [], [1, 1], [], [-1], [], [1, -2, -1]], check_valid=True)]) == [[1.1, 3.3], [], [5.5, 5.5], [], [8.8], [], [10.0, 11.1, 12.2]]
def test_string2(): content = ak.layout.NumpyArray( np.array([ord(x) for x in "heythere"], dtype=np.uint8), parameters={ "__array__": "char", "encoding": "utf-8" }, ) listoffsetarray = ak.layout.ListOffsetArray64( ak.layout.Index64(np.array([0, 3, 3, 8])), content, parameters={"__array__": "string"}, ) a = ak.Array(listoffsetarray, check_valid=True) assert isinstance(a, ak.Array) assert ak.to_list(a) == ["hey", "", "there"] assert str(a) == "['hey', '', 'there']" assert repr(a[0]) == "'hey'" assert repr(a[1]) == "''" assert repr(a[2]) == "'there'"
def test_constructed(): @numba.njit def f1(array): out = np.ones(6, np.float64) i = 0 for x in array: out[i] = x i += 1 return out content = ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5])) virtual1 = ak.layout.VirtualArray( ak.layout.ArrayGenerator(lambda: content, form=content.form, length=len(content))) virtual2 = ak.layout.VirtualArray( ak.layout.ArrayGenerator(lambda: virtual1, form=virtual1.form, length=len(virtual1))) array = ak.Array(virtual2) assert f1(array).tolist() == [0.0, 1.1, 2.2, 3.3, 4.4, 5.5]
def test_categorical_form(): form = ak.forms.Form.fromjson(""" { "class": "IndexedArray64", "index": "i64", "content": "int64", "parameters": { "__array__": "categorical" } } """) builder = ak.layout.TypedArrayBuilder(form) builder.int64(2019) builder.int64(2020) builder.int64(2021) builder.int64(2020) builder.int64(2019) builder.int64(2020) builder.int64(2020) builder.int64(2020) builder.int64(2020) builder.int64(2020) assert ak.to_list(builder.snapshot()) == [ 2019, 2020, 2021, 2020, 2019, 2020, 2020, 2020, 2020, 2020, ] assert str(ak.type(ak.Array( builder.snapshot()))) == "10 * categorical[type=int64]"
def test_other_axis_values(): ak_array = ak.Array([ [[2, 3, 5], [], [], [7, 11, 13]], [[17, 19, 23], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-1).tolist() == [ [2 * 3 * 5, 1, 1, 7 * 11 * 13], [17 * 19 * 23, 1, 1, 29 * 31 * 37], ] assert ak.prod(ak_array, axis=-2).tolist() == [ [2 * 7, 3 * 11, 5 * 13], [17 * 29, 19 * 31, 23 * 37], ] assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19, 5 * 23], [], [], [7 * 29, 11 * 31, 13 * 37], ]
def calculate_selection(self, syst_tag, events): """ """ electron_cut = lepton_selections.select_electrons( electrons=events.Electron, options=self.options["electrons"], clean={}, name="ele", tagger=self) electrons = awkward_utils.add_field(events=events, name="ele", data=events.Electron[electron_cut]) electrons = awkward.Array(electrons, with_name="Momentum4D") presel_cut = awkward.num(electrons) >= 1 self.register_cuts(names=["n_electrons >= 1"], results=[presel_cut]) return presel_cut, events
def test_jagged_mask(): array = ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]], check_valid=True) assert ak.to_list(array[[[True, True, True], [], [True, True], [True], [True, True, True]]]) == [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]] assert ak.to_list(array[[[False, True, True], [], [True, True], [True], [True, True, True]]]) == [[2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]] assert ak.to_list(array[[[True, False, True], [], [True, True], [True], [True, True, True]]]) == [[1.1, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]] assert ak.to_list(array[[[True, True, True], [], [False, True], [True], [True, True, True]]]) == [[1.1, 2.2, 3.3], [], [5.5], [6.6], [7.7, 8.8, 9.9]] assert ak.to_list(array[[[True, True, True], [], [False, False], [True], [True, True, True]]]) == [[1.1, 2.2, 3.3], [], [], [6.6], [7.7, 8.8, 9.9]]
def test(): example = ak.Array([[{"x": 1, "y": []}, {"x": 2, "y": [1, 1]}], []]) cache = {} virtualarray = ak.virtual(lambda: example, form=example.layout.form, length=len(example)) assert len(cache) == 0 tmp1 = virtualarray["x"] assert json.loads(str(tmp1.layout.form)) == { "class": "VirtualArray", "form": { "class": "ListOffsetArray64", "offsets": "i64", "content": "int64" }, "has_length": True, } assert len(cache) == 0 tmp2 = virtualarray["y"] assert json.loads(str(tmp2.layout.form)) == { "class": "VirtualArray", "form": { "class": "ListOffsetArray64", "offsets": "i64", "content": { "class": "ListOffsetArray64", "offsets": "i64", "content": "int64", }, }, "has_length": True, } assert len(cache) == 0 assert tmp1.tolist() == [[1, 2], []] assert tmp2.tolist() == [[[], [1, 1]], []]
def test_from_buffers(): array = ak.Array([ [{ "x": 0.0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": None }], [], [{ "x": 3.3, "y": [1, 2, 3] }, None, { "x": 4.4, "y": [1, 2, 3, 4] }], ]) assert str( array.type) == '3 * var * ?{"x": float64, "y": option[var * int64]}' pa_table = ak.to_arrow_table(array) awkward_array = ak.from_arrow(pa_table) form, length, container = ak.to_buffers(awkward_array) reconstituted = ak.from_buffers(form, length, container, lazy=True) assert reconstituted[2].tolist() == [ { "x": 3.3, "y": [1, 2, 3] }, None, { "x": 4.4, "y": [1, 2, 3, 4] }, ]
def test_string(): array = ak.Array(["one", "two", "three", "four", "five"], check_valid=True) def f1(x, i): return x[i] assert f1(array, 0) == "one" assert f1(array, 1) == "two" assert f1(array, 2) == "three" f1 = numba.njit(f1) assert f1(array, 0) == "one" assert f1(array, 1) == "two" assert f1(array, 2) == "three" if not ak._util.py27: def f2(x, i, j): return x[i] + x[j] assert f2(array, 1, 3) == "twofour" assert numba.njit(f2)(array, 1, 3) == "twofour"
def test_IndexedArray_getitem(): content = ak.from_iter( [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], highlevel=False ) index = ak.layout.Index64(np.array([3, 2, 2, 5, 0, 7], dtype=np.int64)) array = ak.Array(ak.layout.IndexedArray64(index, content), check_valid=True) @numba.njit def f1(x, i): return x[i] assert [f1(array, 0), f1(array, 1), f1(array, 2), f1(array, 3)] == [ 3.3, 2.2, 2.2, 5.5, ] @numba.njit def f2(x, i1, i2): return x[i1:i2] assert ak.to_list(f2(array, 1, 5)) == [2.2, 2.2, 5.5, 0]
def test_firsts_singletons(): array = ak.Array([None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None]) one = ak.singletons(array) assert ak.to_list(one) == [[], [1.1], [2.2], [], [3.3], [], [], [4.4], [5.5], []] two = ak.firsts(one) assert ak.to_list(two) == [ None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None ] array = ak.repartition(array, 3) assert isinstance(array.layout, ak.partition.PartitionedArray) one = ak.singletons(array) assert isinstance(one.layout, ak.partition.PartitionedArray) assert ak.to_list(one) == [[], [1.1], [2.2], [], [3.3], [], [], [4.4], [5.5], []] two = ak.firsts(one) assert isinstance(two.layout, ak.partition.PartitionedArray) assert ak.to_list(two) == [ None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None ]
def test_to_categorical_none(): array = ak.Array([ "one", "two", "three", None, "one", "two", "three", None, "one", "two", "three", None, ]) assert not ak.is_categorical(array) categorical = ak.to_categorical(array) assert ak.is_categorical(categorical) assert ak.to_list(array) == categorical.tolist() assert ak.to_list(categorical.layout.content) == ["one", "two", "three"] not_categorical = ak.from_categorical(categorical) assert not ak.is_categorical(not_categorical) assert ak.categories(categorical).tolist() == ["one", "two", "three"]
def test_5(one, two, three, tmp_path): filename = os.path.join(str(tmp_path), "test5.parquet") data = [{ "x": { "y": [] } }, { "x": { "y": [one] } }, { "x": { "y": [one, two, three] } }] ak.to_parquet(ak.Array(data), filename) array = ak.from_parquet(filename, lazy=True, lazy_cache_key="tmp") assert set(array.caches[0].keys()) == set() array.layout.field("x").array assert set(array.caches[0].keys()) == set() array.layout.field("x").array.field("y").array assert set(array.caches[0].keys()) == set(["tmp:lst:x.y[0]"]) assert array.tolist() == data
def test_IndexedOptionArray_getitem(): array = ak.Array([1.1, 2.2, None, 3.3, None, None, 4.4, 5.5], check_valid=True) @numba.njit def f1(x, i): return x[i] assert [f1(array, 0), f1(array, 1), f1(array, 2), f1(array, 3)] == [ 1.1, 2.2, None, 3.3, ] @numba.njit def f2(x, i1, i2): return x[i1:i2] assert ak.to_list(f2(array, 1, 5)) == [2.2, None, 3.3, None]
def exclusive_jets(self, n_jets, dcut): np_results = 0 if n_jets == 0: raise ValueError("Njets cannot be 0") from None if dcut == -1 and n_jets != -1: np_results = self._results.to_numpy_exclusive_njet(n_jets) if n_jets == -1 and dcut != -1: np_results = self._results.to_numpy_exclusive_dcut(dcut) if np_results == 0: raise ValueError("Either Dcut or Njets should be entered") from None return ak.Array( ak.layout.RecordArray( ( ak.layout.NumpyArray(np_results[0]), ak.layout.NumpyArray(np_results[1]), ak.layout.NumpyArray(np_results[2]), ak.layout.NumpyArray(np_results[3]), ), ("px", "py", "pz", "E"), parameters={"__record__": "Momentum4D"}, ), behavior=self.data.behavior, )