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
|
from __future__ import annotations
import warnings
import pytest
import xarray as xr
from xarray.testing import assert_equal
np = pytest.importorskip("numpy", minversion="1.22")
with warnings.catch_warnings():
warnings.simplefilter("ignore")
import numpy.array_api as xp # isort:skip
from numpy.array_api._array_object import Array # isort:skip
@pytest.fixture
def arrays() -> tuple[xr.DataArray, xr.DataArray]:
np_arr = xr.DataArray(
np.array([[1.0, 2.0, 3.0], [4.0, 5.0, np.nan]]),
dims=("x", "y"),
coords={"x": [10, 20]},
)
xp_arr = xr.DataArray(
xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, np.nan]]),
dims=("x", "y"),
coords={"x": [10, 20]},
)
assert isinstance(xp_arr.data, Array)
return np_arr, xp_arr
def test_arithmetic(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
np_arr, xp_arr = arrays
expected = np_arr + 7
actual = xp_arr + 7
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
def test_aggregation(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
np_arr, xp_arr = arrays
expected = np_arr.sum()
actual = xp_arr.sum()
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
def test_aggregation_skipna(arrays) -> None:
np_arr, xp_arr = arrays
expected = np_arr.sum(skipna=False)
actual = xp_arr.sum(skipna=False)
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
def test_astype(arrays) -> None:
np_arr, xp_arr = arrays
expected = np_arr.astype(np.int64)
actual = xp_arr.astype(np.int64)
assert actual.dtype == np.int64
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
def test_broadcast(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
np_arr, xp_arr = arrays
np_arr2 = xr.DataArray(np.array([1.0, 2.0]), dims="x")
xp_arr2 = xr.DataArray(xp.asarray([1.0, 2.0]), dims="x")
expected = xr.broadcast(np_arr, np_arr2)
actual = xr.broadcast(xp_arr, xp_arr2)
assert len(actual) == len(expected)
for a, e in zip(actual, expected):
assert isinstance(a.data, Array)
assert_equal(a, e)
def test_concat(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
np_arr, xp_arr = arrays
expected = xr.concat((np_arr, np_arr), dim="x")
actual = xr.concat((xp_arr, xp_arr), dim="x")
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
def test_indexing(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
np_arr, xp_arr = arrays
expected = np_arr[:, 0]
actual = xp_arr[:, 0]
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
def test_properties(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
np_arr, xp_arr = arrays
assert np_arr.nbytes == np_arr.data.nbytes
assert xp_arr.nbytes == np_arr.data.nbytes
def test_reorganizing_operation(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
np_arr, xp_arr = arrays
expected = np_arr.transpose()
actual = xp_arr.transpose()
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
def test_stack(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
np_arr, xp_arr = arrays
expected = np_arr.stack(z=("x", "y"))
actual = xp_arr.stack(z=("x", "y"))
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
def test_where() -> None:
np_arr = xr.DataArray(np.array([1, 0]), dims="x")
xp_arr = xr.DataArray(xp.asarray([1, 0]), dims="x")
expected = xr.where(np_arr, 1, 0)
actual = xr.where(xp_arr, 1, 0)
assert isinstance(actual.data, Array)
assert_equal(actual, expected)
|
