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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134
|
# 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_empty_array_slice():
# inspired by PR021::test_getitem
a = ak.operations.from_json("[[], [[], []], [[], [], []]]")
assert to_list(a[2, 1, np.array([], dtype=np.int64)]) == []
# FIXME: assert [[]] == []
# assert to_list(a[2, np.array([1], dtype=np.int64), np.array([], dtype=np.int64)]) == []
a = ak.operations.from_iter([[], [[], []], [[], [], []]], highlevel=False)
assert to_list(a[2, 1, np.array([], dtype=np.int64)]) == []
assert (
a.to_typetracer()[2, 1, np.array([], dtype=np.int64)].form
== a[2, 1, np.array([], dtype=np.int64)].form
)
assert to_list(a[2, np.array([1], dtype=int), np.array([], dtype=int)]) == []
assert (
a.to_typetracer()[
2, np.array([1], dtype=np.int64), np.array([], dtype=np.int64)
].form
== a[2, np.array([1], dtype=np.int64), np.array([], dtype=np.int64)].form
)
# inspired by PR015::test_deep_numpy
content = ak.contents.NumpyArray(
np.array([[0.0, 1.1], [2.2, 3.3], [4.4, 5.5], [6.6, 7.7], [8.8, 9.9]])
)
listarray = ak.contents.ListArray(
ak.index.Index64(np.array([0, 3, 3])),
ak.index.Index64(np.array([3, 3, 5])),
content,
)
assert to_list(listarray[[2, 0, 0, -1], [1, -1, 0, 0], [0, 1, 0, 1]]) == [
8.8,
5.5,
0.0,
7.7,
]
assert (
listarray.to_typetracer()[[2, 0, 0, -1], [1, -1, 0, 0], [0, 1, 0, 1]].form
== listarray[[2, 0, 0, -1], [1, -1, 0, 0], [0, 1, 0, 1]].form
)
assert to_list(listarray[2, 1, np.array([], dtype=np.int64)]) == []
assert (
listarray.to_typetracer()[2, 1, np.array([], dtype=np.int64)].form
== listarray[2, 1, np.array([], dtype=np.int64)].form
)
assert to_list(listarray[2, 1, []]) == []
assert listarray.to_typetracer()[2, 1, []].form == listarray[2, 1, []].form
assert to_list(listarray[2, [1], []]) == []
assert listarray.to_typetracer()[2, [1], []].form == listarray[2, [1], []].form
assert to_list(listarray[2, [], []]) == []
assert listarray.to_typetracer()[2, [], []].form == listarray[2, [], []].form
def test_nonflat_slice():
array = np.arange(2 * 3 * 5).reshape(2, 3, 5)
content = ak.contents.NumpyArray(array.reshape(-1))
inneroffsets = ak.index.Index64(np.array([0, 5, 10, 15, 20, 25, 30]))
outeroffsets = ak.index.Index64(np.array([0, 3, 6]))
listoffsetarray = ak.contents.ListOffsetArray(
outeroffsets, ak.contents.ListOffsetArray(inneroffsets, content)
)
assert to_list(
array[[1, 0, 1, 1, 1, 0], [2, 0, 1, 1, 2, 0], [2, 4, 2, 4, 0, 1]]
) == [27, 4, 22, 24, 25, 1]
assert to_list(
array[
[[1, 0], [1, 1], [1, 0]], [[2, 0], [1, 1], [2, 0]], [[2, 4], [2, 4], [0, 1]]
]
) == [[27, 4], [22, 24], [25, 1]]
one = listoffsetarray[[1, 0, 1, 1, 1, 0], [2, 0, 1, 1, 2, 0], [2, 4, 2, 4, 0, 1]]
assert to_list(one) == [27, 4, 22, 24, 25, 1]
assert (
listoffsetarray.to_typetracer()[
[1, 0, 1, 1, 1, 0], [2, 0, 1, 1, 2, 0], [2, 4, 2, 4, 0, 1]
].form
== one.form
)
def test_nonflat_slice_2():
array = np.arange(2 * 3 * 5).reshape(2, 3, 5)
content = ak.contents.NumpyArray(array.reshape(-1))
inneroffsets = ak.index.Index64(np.array([0, 5, 10, 15, 20, 25, 30]))
outeroffsets = ak.index.Index64(np.array([0, 3, 6]))
listoffsetarray = ak.contents.ListOffsetArray(
outeroffsets, ak.contents.ListOffsetArray(inneroffsets, content)
)
two = listoffsetarray[
np.asarray([[1, 0], [1, 1], [1, 0]]),
np.asarray([[2, 0], [1, 1], [2, 0]]),
np.asarray([[2, 4], [2, 4], [0, 1]]),
]
assert to_list(two) == [[27, 4], [22, 24], [25, 1]]
def test_newaxis():
array = np.arange(2 * 3 * 5).reshape(2, 3, 5)
content = ak.contents.NumpyArray(array.reshape(-1))
inneroffsets = ak.index.Index64(np.array([0, 5, 10, 15, 20, 25, 30]))
outeroffsets = ak.index.Index64(np.array([0, 3, 6]))
listoffsetarray = ak.contents.ListOffsetArray(
outeroffsets, ak.contents.ListOffsetArray(inneroffsets, content)
)
assert to_list(array[:, np.newaxis]) == [
[[[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]]],
]
assert to_list(listoffsetarray[:, np.newaxis]) == [
[[[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]]],
]
assert (
listoffsetarray.to_typetracer()[:, np.newaxis].form
== listoffsetarray[:, np.newaxis].form
)
|