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
68
69
70
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import numpy as np
import pytest # noqa: F401
import awkward as ak
to_list = ak.operations.to_list
def test():
one_content = ak.highlevel.Array(
[[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9, 10.0]]
).layout
one_starts = ak.index.Index64(np.array([0, 2, 3, 3], dtype=np.int64))
one_stops = ak.index.Index64(np.array([2, 3, 3, 5], dtype=np.int64))
one = ak.contents.ListArray(one_starts, one_stops, one_content)
assert to_list(one) == [
[[1.1, 2.2, 3.3], []],
[[4.4, 5.5]],
[],
[[6.6], [7.7, 8.8, 9.9, 10.0]],
]
two_content = ak.highlevel.Array(
[
[123],
[1.1, 2.2, 3.3],
[],
[234],
[4.4, 5.5],
[345],
[6.6],
[7.7, 8.8, 9.9, 10.0],
[456],
]
).layout
two_starts = ak.index.Index64(np.array([1, 4, 5, 6], dtype=np.int64))
two_stops = ak.index.Index64(np.array([3, 5, 5, 8], dtype=np.int64))
two = ak.contents.ListArray(two_starts, two_stops, two_content)
assert to_list(two) == [
[[1.1, 2.2, 3.3], []],
[[4.4, 5.5]],
[],
[[6.6], [7.7, 8.8, 9.9, 10.0]],
]
assert to_list(one[[[[0, 1, 2], []], [[0, 1]], [], [[0], [0, 1, 2, 3]]]]) == [
[[1.1, 2.2, 3.3], []],
[[4.4, 5.5]],
[],
[[6.6], [7.7, 8.8, 9.9, 10.0]],
]
assert (
one.to_typetracer()[[[[0, 1, 2], []], [[0, 1]], [], [[0], [0, 1, 2, 3]]]].form
== one[[[[0, 1, 2], []], [[0, 1]], [], [[0], [0, 1, 2, 3]]]].form
)
assert to_list(two[[[[0, 1, 2], []], [[0, 1]], [], [[0], [0, 1, 2, 3]]]]) == [
[[1.1, 2.2, 3.3], []],
[[4.4, 5.5]],
[],
[[6.6], [7.7, 8.8, 9.9, 10.0]],
]
assert (
two.to_typetracer()[[[[0, 1, 2], []], [[0, 1]], [], [[0], [0, 1, 2, 3]]]].form
== two[[[[0, 1, 2], []], [[0, 1]], [], [[0], [0, 1, 2, 3]]]].form
)
|
