def test_lazily_indexed_array(self):
original = np.random.rand(10, 20, 30)
x = indexing.NumpyIndexingAdapter(original)
v = Variable(['i', 'j', 'k'], original)
lazy = indexing.LazilyOuterIndexedArray(x)
v_lazy = Variable(['i', 'j', 'k'], lazy)
I = ReturnItem() # noqa: E741 # allow ambiguous name
# test orthogonally applied indexers
indexers = [I[:], 0, -2, I[:3], [0, 1, 2, 3], [0], np.arange(10) < 5]
for i in indexers:
for j in indexers:
for k in indexers:
if isinstance(j, np.ndarray) and j.dtype.kind == 'b':
j = np.arange(20) < 5
if isinstance(k, np.ndarray) and k.dtype.kind == 'b':
k = np.arange(30) < 5
expected = np.asarray(v[i, j, k])
for actual in [
v_lazy[i, j, k], v_lazy[:, j, k][i],
v_lazy[:, :, k][:, j][i]
]:
assert expected.shape == actual.shape
assert_array_equal(expected, actual)
assert isinstance(actual._data,
indexing.LazilyOuterIndexedArray)
# make sure actual.key is appropriate type
if all(
isinstance(k, (
int,
slice,
)) for k in v_lazy._data.key.tuple):
assert isinstance(v_lazy._data.key,
indexing.BasicIndexer)
else:
assert isinstance(v_lazy._data.key,
indexing.OuterIndexer)
# test sequentially applied indexers
indexers = [(3, 2), (I[:], 0), (I[:2], -1), (I[:4], [0]), ([4, 5], 0),
([0, 1, 2], [0, 1]), ([0, 3, 5], I[:2])]
for i, j in indexers:
expected = v[i][j]
actual = v_lazy[i][j]
assert expected.shape == actual.shape
assert_array_equal(expected, actual)
# test transpose
if actual.ndim > 1:
order = np.random.choice(actual.ndim, actual.ndim)
order = np.array(actual.dims)
transposed = actual.transpose(*order)
assert_array_equal(expected.transpose(*order), transposed)
assert isinstance(actual._data,
(indexing.LazilyVectorizedIndexedArray,
indexing.LazilyOuterIndexedArray))
assert isinstance(actual._data, indexing.LazilyOuterIndexedArray)
assert isinstance(actual._data.array,
indexing.NumpyIndexingAdapter)
def test_implicit_indexing_adapter():
array = np.arange(10, dtype=np.int64)
implicit = indexing.ImplicitToExplicitIndexingAdapter(
indexing.NumpyIndexingAdapter(array), indexing.BasicIndexer
)
np.testing.assert_array_equal(array, np.asarray(implicit))
np.testing.assert_array_equal(array, implicit[:])
def test_decompose_indexers(shape, indexer_mode, indexing_support):
data = np.random.randn(*shape)
indexer = get_indexers(shape, indexer_mode)
backend_ind, np_ind = indexing.decompose_indexer(indexer, shape, indexing_support)
expected = indexing.NumpyIndexingAdapter(data)[indexer]
array = indexing.NumpyIndexingAdapter(data)[backend_ind]
if len(np_ind.tuple) > 0:
array = indexing.NumpyIndexingAdapter(array)[np_ind]
np.testing.assert_array_equal(expected, array)
if not all(isinstance(k, indexing.integer_types) for k in np_ind.tuple):
combined_ind = indexing._combine_indexers(backend_ind, shape, np_ind)
array = indexing.NumpyIndexingAdapter(data)[combined_ind]
np.testing.assert_array_equal(expected, array)
def test_lazily_indexed_array(self):
original = np.random.rand(10, 20, 30)
x = indexing.NumpyIndexingAdapter(original)
v = Variable(['i', 'j', 'k'], original)
lazy = indexing.LazilyIndexedArray(x)
v_lazy = Variable(['i', 'j', 'k'], lazy)
I = ReturnItem()
# test orthogonally applied indexers
indexers = [I[:], 0, -2, I[:3], [0, 1, 2, 3], [0], np.arange(10) < 5]
for i in indexers:
for j in indexers:
for k in indexers:
if isinstance(j, np.ndarray) and j.dtype.kind == 'b':
j = np.arange(20) < 5
if isinstance(k, np.ndarray) and k.dtype.kind == 'b':
k = np.arange(30) < 5
expected = np.asarray(v[i, j, k])
for actual in [v_lazy[i, j, k],
v_lazy[:, j, k][i],
v_lazy[:, :, k][:, j][i]]:
self.assertEqual(expected.shape, actual.shape)
self.assertArrayEqual(expected, actual)
assert isinstance(actual._data,
indexing.LazilyIndexedArray)
# test sequentially applied indexers
indexers = [(3, 2), (I[:], 0), (I[:2], -1), (I[:4], [0]), ([4, 5], 0),
([0, 1, 2], [0, 1]), ([0, 3, 5], I[:2])]
for i, j in indexers:
expected = np.asarray(v[i][j])
actual = v_lazy[i][j]
self.assertEqual(expected.shape, actual.shape)
self.assertArrayEqual(expected, actual)
assert isinstance(actual._data, indexing.LazilyIndexedArray)
assert isinstance(actual._data.array,
indexing.NumpyIndexingAdapter)
def setup(self):
self.data = indexing.NumpyIndexingAdapter(np.random.randn(10, 12, 13))
self.indexers = [
np.array([[0, 3, 2]]),
np.array([[0, 3, 3], [4, 6, 7]]),
slice(2, -2, 2),
slice(2, -2, 3),
slice(None),
]
def __getitem__(self, item):
key, np_inds = indexing.decompose_indexer(
item, self.shape, indexing.IndexingSupport.OUTER_1VECTOR)
array = self.backend_array[key.tuple]
if len(np_inds.tuple) > 0:
array = indexing.NumpyIndexingAdapter(array)[np_inds]
return array
def __getitem__(self, item):
"""Wrap getitem around the data."""
item, np_inds = indexing.decompose_indexer(
item, self.shape, indexing.IndexingSupport.BASIC)
with self.datastore:
array = self.get_array()[item.tuple]
if len(np_inds.tuple) > 0:
array = indexing.NumpyIndexingAdapter(array)[np_inds]
return array
def test_vectorized_lazily_indexed_array(self) -> None:
original = np.random.rand(10, 20, 30)
x = indexing.NumpyIndexingAdapter(original)
v_eager = Variable(["i", "j", "k"], x)
lazy = indexing.LazilyIndexedArray(x)
v_lazy = Variable(["i", "j", "k"], lazy)
arr = ReturnItem()
def check_indexing(v_eager, v_lazy, indexers):
for indexer in indexers:
actual = v_lazy[indexer]
expected = v_eager[indexer]
assert expected.shape == actual.shape
assert isinstance(
actual._data,
(
indexing.LazilyVectorizedIndexedArray,
indexing.LazilyIndexedArray,
),
)
assert_array_equal(expected, actual)
v_eager = expected
v_lazy = actual
# test orthogonal indexing
indexers = [(arr[:], 0, 1), (Variable("i", [0, 1]), )]
check_indexing(v_eager, v_lazy, indexers)
# vectorized indexing
indexers = [
(Variable("i", [0, 1]), Variable("i", [0, 1]), slice(None)),
(slice(1, 3, 2), 0),
]
check_indexing(v_eager, v_lazy, indexers)
indexers = [
(slice(None, None, 2), 0, slice(None, 10)),
(Variable("i", [3, 2, 4, 3]), Variable("i", [3, 2, 1, 0])),
(Variable(["i", "j"], [[0, 1], [1, 2]]), ),
]
check_indexing(v_eager, v_lazy, indexers)
indexers = [
(Variable("i", [3, 2, 4, 3]), Variable("i", [3, 2, 1, 0])),
(Variable(["i", "j"], [[0, 1], [1, 2]]), ),
]
check_indexing(v_eager, v_lazy, indexers)
def test_vectorized_lazily_indexed_array(self):
original = np.random.rand(10, 20, 30)
x = indexing.NumpyIndexingAdapter(original)
v_eager = Variable(['i', 'j', 'k'], x)
lazy = indexing.LazilyOuterIndexedArray(x)
v_lazy = Variable(['i', 'j', 'k'], lazy)
I = ReturnItem() # noqa: E741 # allow ambiguous name
def check_indexing(v_eager, v_lazy, indexers):
for indexer in indexers:
actual = v_lazy[indexer]
expected = v_eager[indexer]
assert expected.shape == actual.shape
assert isinstance(actual._data,
(indexing.LazilyVectorizedIndexedArray,
indexing.LazilyOuterIndexedArray))
assert_array_equal(expected, actual)
v_eager = expected
v_lazy = actual
# test orthogonal indexing
indexers = [(I[:], 0, 1), (Variable('i', [0, 1]), )]
check_indexing(v_eager, v_lazy, indexers)
# vectorized indexing
indexers = [
(Variable('i', [0, 1]), Variable('i', [0, 1]), slice(None)),
(slice(1, 3, 2), 0)]
check_indexing(v_eager, v_lazy, indexers)
indexers = [
(slice(None, None, 2), 0, slice(None, 10)),
(Variable('i', [3, 2, 4, 3]), Variable('i', [3, 2, 1, 0])),
(Variable(['i', 'j'], [[0, 1], [1, 2]]), )]
check_indexing(v_eager, v_lazy, indexers)
indexers = [
(Variable('i', [3, 2, 4, 3]), Variable('i', [3, 2, 1, 0])),
(Variable(['i', 'j'], [[0, 1], [1, 2]]), )]
check_indexing(v_eager, v_lazy, indexers)
def test_lazily_indexed_array(self):
x = indexing.NumpyIndexingAdapter(np.random.rand(10, 20, 30))
lazy = indexing.LazilyIndexedArray(x)
I = ReturnItem()
# test orthogonally applied indexers
indexers = [I[:], 0, -2, I[:3], [0, 1, 2, 3], np.arange(10) < 5]
for i in indexers:
for j in indexers:
for k in indexers:
expected = np.asarray(x[i, j, k])
for actual in [lazy[i, j, k],
lazy[:, j, k][i],
lazy[:, :, k][:, j][i]]:
self.assertEqual(expected.shape, actual.shape)
self.assertArrayEqual(expected, actual)
# test sequentially applied indexers
indexers = [(3, 2), (I[:], 0), (I[:2], -1), (I[:4], [0]), ([4, 5], 0),
([0, 1, 2], [0, 1]), ([0, 3, 5], I[:2])]
for i, j in indexers:
expected = np.asarray(x[i][j])
actual = lazy[i][j]
self.assertEqual(expected.shape, actual.shape)
self.assertArrayEqual(expected, actual)