def test_create_mask_outer_indexer() -> None:
indexer = indexing.OuterIndexer((np.array([0, -1, 2]), ))
expected = np.array([False, True, False])
actual = indexing.create_mask(indexer, (5, ))
np.testing.assert_array_equal(expected, actual)
indexer = indexing.OuterIndexer((1, slice(2), np.array([0, -1, 2])))
expected = np.array(2 * [[False, True, False]])
actual = indexing.create_mask(indexer, (5, 5, 5))
np.testing.assert_array_equal(expected, actual)
def test_create_mask_dask():
da = pytest.importorskip('dask.array')
indexer = indexing.OuterIndexer((
1,
slice(2),
np.array([0, -1, 2]),
))
expected = np.array(2 * [[False, True, False]])
actual = indexing.create_mask(indexer, (
5,
5,
5,
),
chunks_hint=((1, 1), (2, 1)))
assert actual.chunks == ((1, 1), (2, 1))
np.testing.assert_array_equal(expected, actual)
indexer = indexing.VectorizedIndexer(
(np.array([0, -1, 2]), slice(None), np.array([0, 1, -1])))
expected = np.array([[False, True, True]] * 2).T
actual = indexing.create_mask(indexer, (5, 2), chunks_hint=((3, ), (2, )))
assert isinstance(actual, da.Array)
np.testing.assert_array_equal(expected, actual)
with pytest.raises(ValueError):
indexing.create_mask(indexer, (5, 2), chunks_hint=())
def test_outer_indexer_consistency_with_broadcast_indexes_vectorized():
def nonzero(x):
if isinstance(x, np.ndarray) and x.dtype.kind == 'b':
x = x.nonzero()[0]
return x
original = np.random.rand(10, 20, 30)
v = Variable(['i', 'j', 'k'], original)
I = ReturnItem() # noqa: E741 # allow ambiguous name
# test orthogonally applied indexers
indexers = [
I[:], 0, -2, I[:3],
np.array([0, 1, 2, 3]),
np.array([0]),
np.arange(10) < 5
]
for i, j, k in itertools.product(indexers, repeat=3):
if isinstance(j, np.ndarray) and j.dtype.kind == 'b': # match size
j = np.arange(20) < 4
if isinstance(k, np.ndarray) and k.dtype.kind == 'b':
k = np.arange(30) < 8
_, expected, new_order = v._broadcast_indexes_vectorized((i, j, k))
expected_data = nputils.NumpyVIndexAdapter(v.data)[expected.tuple]
if new_order:
old_order = range(len(new_order))
expected_data = np.moveaxis(expected_data, old_order, new_order)
outer_index = indexing.OuterIndexer(
(nonzero(i), nonzero(j), nonzero(k)))
actual = indexing._outer_to_numpy_indexer(outer_index, v.shape)
actual_data = v.data[actual]
np.testing.assert_array_equal(actual_data, expected_data)
def get_indexers(shape, mode):
if mode == "vectorized":
indexed_shape = (3, 4)
indexer = tuple(
np.random.randint(0, s, size=indexed_shape) for s in shape)
return indexing.VectorizedIndexer(indexer)
elif mode == "outer":
indexer = tuple(np.random.randint(0, s, s + 2) for s in shape)
return indexing.OuterIndexer(indexer)
elif mode == "outer_scalar":
indexer = (np.random.randint(0, 3, 4), 0, slice(None, None, 2))
return indexing.OuterIndexer(indexer[:len(shape)])
elif mode == "outer_scalar2":
indexer = (np.random.randint(0, 3, 4), -2, slice(None, None, 2))
return indexing.OuterIndexer(indexer[:len(shape)])
elif mode == "outer1vec":
indexer = [slice(2, -3) for s in shape]
indexer[1] = np.random.randint(0, shape[1], shape[1] + 2)
return indexing.OuterIndexer(tuple(indexer))
elif mode == "basic": # basic indexer
indexer = [slice(2, -3) for s in shape]
indexer[0] = 3
return indexing.BasicIndexer(tuple(indexer))
elif mode == "basic1": # basic indexer
return indexing.BasicIndexer((3, ))
elif mode == "basic2": # basic indexer
indexer = [0, 2, 4]
return indexing.BasicIndexer(tuple(indexer[:len(shape)]))
elif mode == "basic3": # basic indexer
indexer = [slice(None) for s in shape]
indexer[0] = slice(-2, 2, -2)
indexer[1] = slice(1, -1, 2)
return indexing.BasicIndexer(tuple(indexer[:len(shape)]))
