diff --git a/xarray/testing/duckarrays.py b/xarray/testing/duckarrays.py new file mode 100644 index 00000000000..cb6450a1d71 --- /dev/null +++ b/xarray/testing/duckarrays.py @@ -0,0 +1,145 @@ +from abc import abstractmethod +from typing import TYPE_CHECKING + +import hypothesis.extra.numpy as npst +import hypothesis.strategies as st +import numpy as np +import numpy.testing as npt +import pytest +from hypothesis import given, note + +import xarray as xr +import xarray.testing.strategies as xrst +from xarray.core.types import T_DuckArray +from xarray.testing.assertions import assert_identical + +if TYPE_CHECKING: + from xarray.core.types import _DTypeLikeNested, _ShapeLike + + +__all__ = [ + "ConstructorTests", + "ReduceTests", +] + + +class ArrayConstructorChecksMixin: + """Mixin for checking results of Variable/DataArray constructors.""" + + def check(self, var, arr): + self.check_types(var, arr) + self.check_values(var, arr) + self.check_attributes(var, arr) + + def check_types(self, var, arr): + # test type of wrapped array + assert isinstance( + var.data, type(arr) + ), f"found {type(var.data)}, expected {type(arr)}" + + def check_attributes(self, var, arr): + # test ndarray attributes are exposed correctly + assert var.ndim == arr.ndim + assert var.shape == arr.shape + assert var.dtype == arr.dtype + assert var.size == arr.size + + def check_values(self, var, arr): + # test coercion to numpy + npt.assert_equal(var.to_numpy(), np.asarray(arr)) + + +class ConstructorTests(ArrayConstructorChecksMixin): + shapes = npst.array_shapes() + dtypes = xrst.supported_dtypes() + + @staticmethod + @abstractmethod + def array_strategy_fn( + *, shape: "_ShapeLike", dtype: "_DTypeLikeNested" + ) -> st.SearchStrategy[T_DuckArray]: + # TODO can we just make this an attribute? + ... + + @given(st.data()) + def test_construct_variable(self, data) -> None: + shape = data.draw(self.shapes) + dtype = data.draw(self.dtypes) + arr = data.draw(self.array_strategy_fn(shape=shape, dtype=dtype)) + + dim_names = data.draw( + xrst.dimension_names(min_dims=len(shape), max_dims=len(shape)) + ) + var = xr.Variable(data=arr, dims=dim_names) + + self.check(var, arr) + + +def is_real_floating(dtype): + return np.issubdtype(dtype, np.number) and np.issubdtype(dtype, np.floating) + + +class ReduceTests: + dtypes = xrst.supported_dtypes() + + @staticmethod + @abstractmethod + def array_strategy_fn( + *, shape: "_ShapeLike", dtype: "_DTypeLikeNested" + ) -> st.SearchStrategy[T_DuckArray]: + # TODO can we just make this an attribute? + ... + + def check_reduce(self, var, op, dim, *args, **kwargs): + actual = getattr(var, op)(dim=dim, *args, **kwargs) + + data = np.asarray(var.data) + expected = getattr(var.copy(data=data), op)(*args, **kwargs) + + # create expected result (using nanmean because arrays with Nans will be generated) + reduce_axes = tuple(var.get_axis_num(d) for d in dim) + data = np.asarray(var.data) + expected = getattr(var.copy(data=data), op)(*args, axis=reduce_axes, **kwargs) + + note(f"actual:\n{actual}") + note(f"expected:\n{expected}") + + assert_identical(actual, expected) + + @pytest.mark.parametrize( + "method, dtype_assumption", + ( + ("all", lambda x: True), # should work for any dtype + ("any", lambda x: True), # should work for any dtype + # "cumprod", # not in array API + # "cumsum", # not in array API + ("max", is_real_floating), # only in array API for real numeric dtypes + # "median", # not in array API + ("min", is_real_floating), # only in array API for real numeric dtypes + ("prod", is_real_floating), # only in array API for real numeric dtypes + # "std", # TypeError: std() got an unexpected keyword argument 'ddof' + ("sum", is_real_floating), # only in array API for real numeric dtypes + # "var", # TypeError: std() got an unexpected keyword argument 'ddof' + ), + ) + @given(st.data()) + def test_reduce_variable(self, method, dtype_assumption, data): + """ + Test that the reduction applied to an xarray Variable is always equal + to the same reduction applied to the underlying array. + """ + + narrowed_dtypes = self.dtypes.filter(dtype_assumption) + + var = data.draw( + xrst.variables( + array_strategy_fn=self.array_strategy_fn, + dims=xrst.dimension_names(min_dims=1), + dtype=narrowed_dtypes, + ) + ) + + # specify arbitrary reduction along at least one dimension + reduce_dims = data.draw(xrst.unique_subset_of(var.dims, min_size=1)) + + self.check_reduce(var, method, dim=reduce_dims) diff --git a/xarray/testing/utils.py b/xarray/testing/utils.py new file mode 100644 index 00000000000..56dd1e83067 --- /dev/null +++ b/xarray/testing/utils.py @@ -0,0 +1,10 @@ +import warnings +from contextlib import contextmanager + + +@contextmanager +def suppress_warning(category, message=""): + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=category, message=message) + + yield diff --git a/xarray/tests/duckarrays/__init__.py b/xarray/tests/duckarrays/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/xarray/tests/duckarrays/test_array_api.py b/xarray/tests/duckarrays/test_array_api.py new file mode 100644 index 00000000000..4c217cf9645 --- /dev/null +++ b/xarray/tests/duckarrays/test_array_api.py @@ -0,0 +1,45 @@ +from typing import TYPE_CHECKING + +import hypothesis.strategies as st +from hypothesis.extra.array_api import make_strategies_namespace + +from xarray.core.types import T_DuckArray +from xarray.testing import duckarrays +from xarray.testing.utils import suppress_warning +from xarray.tests import _importorskip + +if TYPE_CHECKING: + from xarray.core.types import _DTypeLikeNested, _ShapeLike + + +# ignore the warning that the array_api is experimental raised by numpy +with suppress_warning( + UserWarning, "The numpy.array_api submodule is still experimental. See NEP 47." +): + _importorskip("numpy", "1.26.0") + import numpy.array_api as nxp + + +nxps = make_strategies_namespace(nxp) + + +class TestConstructors(duckarrays.ConstructorTests): + dtypes = nxps.scalar_dtypes() + + @staticmethod + def array_strategy_fn( + shape: "_ShapeLike", + dtype: "_DTypeLikeNested", + ) -> st.SearchStrategy[T_DuckArray]: + return nxps.arrays(shape=shape, dtype=dtype) + + +class TestReductions(duckarrays.ReduceTests): + dtypes = nxps.scalar_dtypes() + + @staticmethod + def array_strategy_fn( + shape: "_ShapeLike", + dtype: "_DTypeLikeNested", + ) -> st.SearchStrategy[T_DuckArray]: + return nxps.arrays(shape=shape, dtype=dtype) diff --git a/xarray/tests/duckarrays/test_numpy.py b/xarray/tests/duckarrays/test_numpy.py new file mode 100644 index 00000000000..f18b70fbc87 --- /dev/null +++ b/xarray/tests/duckarrays/test_numpy.py @@ -0,0 +1,33 @@ +from typing import TYPE_CHECKING + +import hypothesis.extra.numpy as npst +import hypothesis.strategies as st + +from xarray.core.types import T_DuckArray +from xarray.testing import duckarrays +from xarray.testing.strategies import supported_dtypes + +if TYPE_CHECKING: + from xarray.core.types import _DTypeLikeNested, _ShapeLike + + +class TestConstructors(duckarrays.ConstructorTests): + dtypes = supported_dtypes() + + @staticmethod + def array_strategy_fn( + shape: "_ShapeLike", + dtype: "_DTypeLikeNested", + ) -> st.SearchStrategy[T_DuckArray]: + return npst.arrays(shape=shape, dtype=dtype) + + +class TestReductions(duckarrays.ReduceTests): + dtypes = supported_dtypes() + + @staticmethod + def array_strategy_fn( + shape: "_ShapeLike", + dtype: "_DTypeLikeNested", + ) -> st.SearchStrategy[T_DuckArray]: + return npst.arrays(shape=shape, dtype=dtype) diff --git a/xarray/tests/duckarrays/test_sparse.py b/xarray/tests/duckarrays/test_sparse.py new file mode 100644 index 00000000000..1a91c6fd72d --- /dev/null +++ b/xarray/tests/duckarrays/test_sparse.py @@ -0,0 +1,76 @@ +from typing import TYPE_CHECKING + +import hypothesis.extra.numpy as npst +import hypothesis.strategies as st +import numpy as np +import numpy.testing as npt +import pytest + +import xarray.testing.strategies as xrst +from xarray.testing import duckarrays +from xarray.tests import _importorskip + +if TYPE_CHECKING: + from xarray.core.types import _DTypeLikeNested, _ShapeLike + + +_importorskip("sparse") +import sparse + + +@pytest.fixture(autouse=True) +def disable_bottleneck(): + from xarray import set_options + + with set_options(use_bottleneck=False): + yield + + +# sparse does not support float16 +sparse_dtypes = xrst.supported_dtypes().filter( + lambda dtype: (not np.issubdtype(dtype, np.float16)) +) + + +@st.composite +def sparse_arrays_fn( + draw: st.DrawFn, + *, + shape: "_ShapeLike", + dtype: "_DTypeLikeNested", +) -> sparse.COO: + """When called generates an arbitrary sparse.COO array of the given shape and dtype.""" + np_arr = draw(npst.arrays(dtype, shape)) + + def to_sparse(arr: np.ndarray) -> sparse.COO: + if arr.ndim == 0: + return arr + + return sparse.COO.from_numpy(arr) + + return to_sparse(np_arr) + + +class TestConstructors(duckarrays.ConstructorTests): + dtypes = sparse_dtypes() + + @staticmethod + def array_strategy_fn( + shape: "_ShapeLike", + dtype: "_DTypeLikeNested", + ) -> st.SearchStrategy[sparse.COO]: + return sparse_arrays_fn + + def check_values(self, var, arr): + npt.assert_equal(var.to_numpy(), arr.todense()) + + +class TestReductions(duckarrays.ReduceTests): + dtypes = nxps.scalar_dtypes() + + @staticmethod + def array_strategy_fn( + shape: "_ShapeLike", + dtype: "_DTypeLikeNested", + ) -> st.SearchStrategy[T_DuckArray]: + return nxps.arrays(shape=shape, dtype=dtype) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index af86c18668f..7b984d486f7 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -92,11 +92,14 @@ def array_strip_units(array): def array_attach_units(data, unit): + if unit is None or (isinstance(unit, int) and unit == 1): + return data + if isinstance(data, Quantity): raise ValueError(f"cannot attach unit {unit} to quantity {data}") try: - quantity = data * unit + quantity = unit._REGISTRY.Quantity(data, unit) except np.core._exceptions.UFuncTypeError: if isinstance(unit, unit_registry.Unit): raise @@ -181,36 +184,40 @@ def attach_units(obj, units): return array_attach_units(obj, units) if isinstance(obj, xr.Dataset): - data_vars = { - name: attach_units(value, units) for name, value in obj.data_vars.items() + variables = { + name: attach_units(value, {None: units.get(name)}) + for name, value in obj.variables.items() } - coords = { - name: attach_units(value, units) for name, value in obj.coords.items() + name: var for name, var in variables.items() if name in obj._coord_names + } + data_vars = { + name: var for name, var in variables.items() if name not in obj._coord_names } - new_obj = xr.Dataset(data_vars=data_vars, coords=coords, attrs=obj.attrs) elif isinstance(obj, xr.DataArray): # try the array name, "data" and None, then fall back to dimensionless - data_units = units.get(obj.name, None) or units.get(None, None) or 1 - - data = array_attach_units(obj.data, data_units) + units = units.copy() + THIS_ARRAY = xr.core.dataarray._THIS_ARRAY + unset = object() + if obj.name in units: + name = obj.name + elif None in units: + name = None + else: + name = unset - coords = { - name: ( - (value.dims, array_attach_units(value.data, units.get(name) or 1)) - if name in units - else (value.dims, value.data) - ) - for name, value in obj.coords.items() - } - dims = obj.dims - attrs = obj.attrs + if name is not unset: + units[THIS_ARRAY] = units.pop(name) - new_obj = xr.DataArray( - name=obj.name, data=data, coords=coords, attrs=attrs, dims=dims - ) + ds = obj._to_temp_dataset() + attached = attach_units(ds, units) + new_obj = obj._from_temp_dataset(attached, name=obj.name) else: + if isinstance(obj, xr.IndexVariable): + # no units for index variables + return obj + data_units = units.get("data", None) or units.get(None, None) or 1 data = array_attach_units(obj.data, data_units)