Namespace-aware xarray.ufuncs (#9776)

* initial namespace-aware implementation

* use np subclass, test duck dask arrays

* remove dask special casing and numpy fallback

* add isnat

* hard code the supported ufuncs

* handle np versions, separate unary/binary path

* explicit unary/binary creators

* add to api docs

* add whats new

* move numpy version check to tests

* fix docs for aliased np funcs

* fix whats new

---------

Co-authored-by: Stephan Hoyer <shoyer@google.com>

Author: slevang

doc/api.rst

@@ -894,6 +894,120 @@ Methods copied from :py:class:`numpy.ndarray` objects, here applying to the data
 ..    DataTree.sortby
 ..    DataTree.broadcast_like
 
+Universal functions
+===================
+
+These functions are equivalent to their NumPy versions, but for xarray
+objects backed by non-NumPy array types (e.g. ``cupy``, ``sparse``, or ``jax``),
+they will ensure that the computation is dispatched to the appropriate
+backend. You can find them in the ``xarray.ufuncs`` module:
+
+.. autosummary::
+   :toctree: generated/
+
+   ufuncs.abs
+   ufuncs.absolute
+   ufuncs.acos
+   ufuncs.acosh
+   ufuncs.arccos
+   ufuncs.arccosh
+   ufuncs.arcsin
+   ufuncs.arcsinh
+   ufuncs.arctan
+   ufuncs.arctanh
+   ufuncs.asin
+   ufuncs.asinh
+   ufuncs.atan
+   ufuncs.atanh
+   ufuncs.bitwise_count
+   ufuncs.bitwise_invert
+   ufuncs.bitwise_not
+   ufuncs.cbrt
+   ufuncs.ceil
+   ufuncs.conj
+   ufuncs.conjugate
+   ufuncs.cos
+   ufuncs.cosh
+   ufuncs.deg2rad
+   ufuncs.degrees
+   ufuncs.exp
+   ufuncs.exp2
+   ufuncs.expm1
+   ufuncs.fabs
+   ufuncs.floor
+   ufuncs.invert
+   ufuncs.isfinite
+   ufuncs.isinf
+   ufuncs.isnan
+   ufuncs.isnat
+   ufuncs.log
+   ufuncs.log10
+   ufuncs.log1p
+   ufuncs.log2
+   ufuncs.logical_not
+   ufuncs.negative
+   ufuncs.positive
+   ufuncs.rad2deg
+   ufuncs.radians
+   ufuncs.reciprocal
+   ufuncs.rint
+   ufuncs.sign
+   ufuncs.signbit
+   ufuncs.sin
+   ufuncs.sinh
+   ufuncs.spacing
+   ufuncs.sqrt
+   ufuncs.square
+   ufuncs.tan
+   ufuncs.tanh
+   ufuncs.trunc
+   ufuncs.add
+   ufuncs.arctan2
+   ufuncs.atan2
+   ufuncs.bitwise_and
+   ufuncs.bitwise_left_shift
+   ufuncs.bitwise_or
+   ufuncs.bitwise_right_shift
+   ufuncs.bitwise_xor
+   ufuncs.copysign
+   ufuncs.divide
+   ufuncs.equal
+   ufuncs.float_power
+   ufuncs.floor_divide
+   ufuncs.fmax
+   ufuncs.fmin
+   ufuncs.fmod
+   ufuncs.gcd
+   ufuncs.greater
+   ufuncs.greater_equal
+   ufuncs.heaviside
+   ufuncs.hypot
+   ufuncs.lcm
+   ufuncs.ldexp
+   ufuncs.left_shift
+   ufuncs.less
+   ufuncs.less_equal
+   ufuncs.logaddexp
+   ufuncs.logaddexp2
+   ufuncs.logical_and
+   ufuncs.logical_or
+   ufuncs.logical_xor
+   ufuncs.maximum
+   ufuncs.minimum
+   ufuncs.mod
+   ufuncs.multiply
+   ufuncs.nextafter
+   ufuncs.not_equal
+   ufuncs.pow
+   ufuncs.power
+   ufuncs.remainder
+   ufuncs.right_shift
+   ufuncs.subtract
+   ufuncs.true_divide
+   ufuncs.angle
+   ufuncs.isreal
+   ufuncs.iscomplex
+
 IO / Conversion
 ===============
 

doc/whats-new.rst

@@ -41,6 +41,10 @@ New Features
 - Optimize :py:meth:`DataArray.polyfit` and :py:meth:`Dataset.polyfit` with dask, when used with
   arrays with more than two dimensions.
   (:issue:`5629`). By `Deepak Cherian <https://github.com/dcherian>`_.
+- Re-implement the :py:mod:`ufuncs` module, which now dynamically dispatches to the
+  underlying array's backend. Provides better support for certain wrapped array types
+  like ``jax.numpy.ndarray``. (:issue:`7848`, :pull:`9776`).
+  By `Sam Levang <https://github.com/slevang>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

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,108 @@ 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)
+        np_func = getattr(np, name, None)
+        if np_func is None and np.lib.NumpyVersion(np.__version__) < "2.0.0":
+            pytest.skip(f"Ufunc {name} is not available in numpy {np.__version__}.")
+
+        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)