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
|
from __future__ import annotations
import numpy as np
from numpy.testing import assert_array_equal
from xarray.core.nputils import NumpyVIndexAdapter, _is_contiguous
def test_is_contiguous() -> None:
assert _is_contiguous([1])
assert _is_contiguous([1, 2, 3])
assert not _is_contiguous([1, 3])
def test_vindex() -> None:
x = np.arange(3 * 4 * 5).reshape((3, 4, 5))
vindex = NumpyVIndexAdapter(x)
# getitem
assert_array_equal(vindex[0], x[0])
assert_array_equal(vindex[[1, 2], [1, 2]], x[[1, 2], [1, 2]])
assert vindex[[0, 1], [0, 1], :].shape == (2, 5)
assert vindex[[0, 1], :, [0, 1]].shape == (2, 4)
assert vindex[:, [0, 1], [0, 1]].shape == (2, 3)
# setitem
vindex[:] = 0
assert_array_equal(x, np.zeros_like(x))
# assignment should not raise
vindex[[0, 1], [0, 1], :] = vindex[[0, 1], [0, 1], :]
vindex[[0, 1], :, [0, 1]] = vindex[[0, 1], :, [0, 1]]
vindex[:, [0, 1], [0, 1]] = vindex[:, [0, 1], [0, 1]]
|
