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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import numpy as np # noqa: F401
import pytest # noqa: F401
import awkward as ak
to_list = ak.operations.to_list
def test():
array = ak.highlevel.Array(
[[[0.0, 1.1, 2.2], []], [[3.3, 4.4]], [], [[5.5], [], [6.6, 7.7, 8.8, 9.9]]]
)
assert to_list(ak.operations.local_index(array, axis=0)) == [
0,
1,
2,
3,
]
assert to_list(ak.operations.local_index(array, axis=1)) == [
[0, 1],
[0],
[],
[0, 1, 2],
]
assert to_list(ak.operations.local_index(array, axis=2)) == [
[[0, 1, 2], []],
[[0, 1]],
[],
[[0], [], [0, 1, 2, 3]],
]
assert to_list(ak.operations.local_index(array, axis=-1)) == [
[[0, 1, 2], []],
[[0, 1]],
[],
[[0], [], [0, 1, 2, 3]],
]
assert to_list(ak.operations.local_index(array, axis=-2)) == [
[0, 1],
[0],
[],
[0, 1, 2],
]
assert to_list(ak.operations.local_index(array, axis=-3)) == [
0,
1,
2,
3,
]
assert to_list(
ak.operations.zip(
[
ak.operations.local_index(array, axis=0),
ak.operations.local_index(array, axis=1),
ak.operations.local_index(array, axis=2),
]
)
) == [
[[(0, 0, 0), (0, 0, 1), (0, 0, 2)], []],
[[(1, 0, 0), (1, 0, 1)]],
[],
[[(3, 0, 0)], [], [(3, 2, 0), (3, 2, 1), (3, 2, 2), (3, 2, 3)]],
]
|
