147 lines
4.4 KiB
Python
147 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
import xarray as xr
|
|
from xarray.testing import assert_equal
|
|
|
|
np = pytest.importorskip("numpy", minversion="1.22")
|
|
xp = pytest.importorskip("array_api_strict")
|
|
|
|
from array_api_strict._array_object import Array # isort:skip # type: ignore[no-redef]
|
|
|
|
|
|
@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(xp.int64)
|
|
assert actual.dtype == xp.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, strict=True):
|
|
assert isinstance(a.data, Array)
|
|
assert_equal(a, e)
|
|
|
|
|
|
def test_broadcast_during_arithmetic(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 = np_arr * np_arr2
|
|
actual = xp_arr * xp_arr2
|
|
assert isinstance(actual.data, Array)
|
|
assert_equal(actual, expected)
|
|
|
|
expected = np_arr2 * np_arr
|
|
actual = xp_arr2 * xp_arr
|
|
assert isinstance(actual.data, Array)
|
|
assert_equal(actual, expected)
|
|
|
|
|
|
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
|
|
|
|
expected = np_arr.data.nbytes
|
|
assert np_arr.nbytes == expected
|
|
assert xp_arr.nbytes == expected
|
|
|
|
|
|
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_unstack(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
|
|
np_arr, xp_arr = arrays
|
|
expected = np_arr.stack(z=("x", "y")).unstack()
|
|
actual = xp_arr.stack(z=("x", "y")).unstack()
|
|
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)
|