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,
)
