Namespace-aware xarray.ufuncs

xarray/__init__.py

@@ -1,6 +1,6 @@
 from importlib.metadata import version as _version
 
-from xarray import groupers, testing, tutorial
+from xarray import groupers, testing, tutorial, ufuncs
 from xarray.backends.api import (
     load_dataarray,
     load_dataset,
@@ -69,6 +69,7 @@
     "groupers",
     "testing",
     "tutorial",
+    "ufuncs",
     # Top-level functions
     "align",
     "apply_ufunc",

xarray/tests/test_dask.py

@@ -11,6 +11,7 @@
 import pytest
 
 import xarray as xr
+import xarray.ufuncs as xu
 from xarray import DataArray, Dataset, Variable
 from xarray.core import duck_array_ops
 from xarray.core.duck_array_ops import lazy_array_equiv
@@ -274,6 +275,17 @@ def test_bivariate_ufunc(self):
         self.assertLazyAndAllClose(np.maximum(u, 0), np.maximum(v, 0))
         self.assertLazyAndAllClose(np.maximum(u, 0), np.maximum(0, v))
 
+    def test_univariate_xufunc(self):
+        u = self.eager_var
+        v = self.lazy_var
+        self.assertLazyAndAllClose(np.sin(u), xu.sin(v))
+
+    def test_bivariate_xufunc(self):
+        u = self.eager_var
+        v = self.lazy_var
+        self.assertLazyAndAllClose(np.maximum(u, 0), xu.maximum(v, 0))
+        self.assertLazyAndAllClose(np.maximum(u, 0), xu.maximum(0, v))
+
     def test_compute(self):
         u = self.eager_var
         v = self.lazy_var

xarray/tests/test_sparse.py

@@ -9,6 +9,7 @@
 import pytest
 
 import xarray as xr
+import xarray.ufuncs as xu
 from xarray import DataArray, Variable
 from xarray.namedarray.pycompat import array_type
 from xarray.tests import assert_equal, assert_identical, requires_dask
@@ -294,6 +295,13 @@ def test_bivariate_ufunc(self):
         assert_sparse_equal(np.maximum(self.data, 0), np.maximum(self.var, 0).data)
         assert_sparse_equal(np.maximum(self.data, 0), np.maximum(0, self.var).data)
 
+    def test_univariate_xufunc(self):
+        assert_sparse_equal(xu.sin(self.var).data, np.sin(self.data))
+
+    def test_bivariate_xufunc(self):
+        assert_sparse_equal(xu.multiply(self.var, 0).data, np.multiply(self.data, 0))
+        assert_sparse_equal(xu.multiply(0, self.var).data, np.multiply(0, self.data))
+
     def test_repr(self):
         expected = dedent(
             """\

xarray/tests/test_ufuncs.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
+import pickle
+from unittest.mock import patch
+
 import numpy as np
 import pytest
 
 import xarray as xr
-from xarray.tests import assert_allclose, assert_array_equal, mock
+import xarray.ufuncs as xu
+from xarray.tests import assert_allclose, assert_array_equal, mock, requires_dask
 from xarray.tests import assert_identical as assert_identical_
 
 
@@ -155,3 +159,109 @@ def test_gufuncs():
     fake_gufunc = mock.Mock(signature="(n)->()", autospec=np.sin)
     with pytest.raises(NotImplementedError, match=r"generalized ufuncs"):
         xarray_obj.__array_ufunc__(fake_gufunc, "__call__", xarray_obj)
+
+
+class DuckArray(np.ndarray):
+    # Minimal subclassed duck array with its own self-contained namespace,
+    # which implements a few ufuncs
+    def __new__(cls, array):
+        obj = np.asarray(array).view(cls)
+        return obj
+
+    def __array_namespace__(self):
+        return DuckArray
+
+    @staticmethod
+    def sin(x):
+        return np.sin(x)
+
+    @staticmethod
+    def add(x, y):
+        return x + y
+
+
+class DuckArray2(DuckArray):
+    def __array_namespace__(self):
+        return DuckArray2
+
+
+class TestXarrayUfuncs:
+    @pytest.fixture(autouse=True)
+    def setUp(self):
+        self.x = xr.DataArray([1, 2, 3])
+        self.xd = xr.DataArray(DuckArray([1, 2, 3]))
+        self.xd2 = xr.DataArray(DuckArray2([1, 2, 3]))
+        self.xt = xr.DataArray(np.datetime64("2021-01-01", "ns"))
+
+    @pytest.mark.filterwarnings("ignore::RuntimeWarning")
+    @pytest.mark.parametrize("name", xu.__all__)
+    def test_ufuncs(self, name, request):
+        xu_func = getattr(xu, name)
+        if isinstance(xu_func, xu._UnavailableUfunc):
+            pytest.xfail(f"Ufunc {name} is not available in numpy {np.__version__}.")
+
+        np_func = getattr(np, name)
+
+        if name == "isnat":
+            args = (self.xt,)
+        elif hasattr(np_func, "nin") and np_func.nin == 2:
+            args = (self.x, self.x)
+        else:
+            args = (self.x,)
+
+        expected = np_func(*args)
+        actual = xu_func(*args)
+
+        if name in ["angle", "iscomplex"]:
+            np.testing.assert_equal(expected, actual.values)
+        else:
+            assert_identical(actual, expected)
+
+    def test_ufunc_pickle(self):
+        a = 1.0
+        cos_pickled = pickle.loads(pickle.dumps(xu.cos))
+        assert_identical(cos_pickled(a), xu.cos(a))
+
+    def test_ufunc_scalar(self):
+        actual = xu.sin(1)
+        assert isinstance(actual, float)
+
+    def test_ufunc_duck_array_dataarray(self):
+        actual = xu.sin(self.xd)
+        assert isinstance(actual.data, DuckArray)
+
+    def test_ufunc_duck_array_variable(self):
+        actual = xu.sin(self.xd.variable)
+        assert isinstance(actual.data, DuckArray)
+
+    def test_ufunc_duck_array_dataset(self):
+        ds = xr.Dataset({"a": self.xd})
+        actual = xu.sin(ds)
+        assert isinstance(actual.a.data, DuckArray)
+
+    @requires_dask
+    def test_ufunc_duck_dask(self):
+        import dask.array as da
+
+        x = xr.DataArray(da.from_array(DuckArray(np.array([1, 2, 3]))))
+        actual = xu.sin(x)
+        assert isinstance(actual.data._meta, DuckArray)
+
+    @requires_dask
+    @pytest.mark.xfail(reason="dask ufuncs currently dispatch to numpy")
+    def test_ufunc_duck_dask_no_array_ufunc(self):
+        import dask.array as da
+
+        # dask ufuncs currently only preserve duck arrays that implement __array_ufunc__
+        with patch.object(DuckArray, "__array_ufunc__", new=None, create=True):
+            x = xr.DataArray(da.from_array(DuckArray(np.array([1, 2, 3]))))
+            actual = xu.sin(x)
+            assert isinstance(actual.data._meta, DuckArray)
+
+    def test_ufunc_mixed_arrays_compatible(self):
+        actual = xu.add(self.xd, self.x)
+        assert isinstance(actual.data, DuckArray)
+
+    def test_ufunc_mixed_arrays_incompatible(self):
+        with pytest.raises(ValueError, match=r"Mixed array types"):
+            xu.add(self.xd, self.xd2)