Support for __array_function__ implementers (sparse arrays) [WIP] (#3117)

* Support for __array_function__ implementers

* Pep8

* Consistent naming

* Check for NEP18 enabled and nep18 non-numpy arrays

* Replace .values with .data

* Add initial test for nep18

* Fix linting issues

* Add parameterized tests

* Internal clean-up of isnull() to avoid relying on pandas

This version should be much more compatible out of the box with duck typing.

* Add sparse to ci requirements

* Moar tests

* Two more patches for __array_function__ duck-arrays

* Don't use coords attribute from duck-arrays that aren't derived from DataWithCoords

* Improve checking for coords, and autopep8

* Skip tests if NEP-18 envvar is not set

* flake8

* Update xarray/core/dataarray.py

Co-Authored-By: Stephan Hoyer <[email protected]>

* Fix coords parsing

* More tests

* Add align tests

* Replace nep18 tests with more extensive tests on pydata/sparse

* Add xfails for missing np.result_type (fixed by pydata/sparse/pull/261)

* Fix xpasses

* Revert isnull/notnull

* Fix as_like_arrays by coercing dense arrays to COO if any sparse

* Make Variable.load a no-op for non-dask duck arrays

* Add additional method tests

* Fix utils.as_scalar to handle duck arrays with ndim>0

Author: nvictus

ci/requirements/py37.yml

@@ -21,6 +21,7 @@ dependencies:
   - pip
   - scipy
   - seaborn
+  - sparse
   - toolz
   - rasterio
   - boto3

xarray/core/dataarray.py

@@ -260,15 +260,18 @@ def __init__(
         else:
             # try to fill in arguments from data if they weren't supplied
             if coords is None:
-                coords = getattr(data, 'coords', None)
-                if isinstance(data, pd.Series):
+
+                if isinstance(data, DataArray):
+                    coords = data.coords
+                elif isinstance(data, pd.Series):
                     coords = [data.index]
                 elif isinstance(data, pd.DataFrame):
                     coords = [data.index, data.columns]
                 elif isinstance(data, (pd.Index, IndexVariable)):
                     coords = [data]
                 elif isinstance(data, pdcompat.Panel):
                     coords = [data.items, data.major_axis, data.minor_axis]
+
             if dims is None:
                 dims = getattr(data, 'dims', getattr(coords, 'dims', None))
             if name is None:

xarray/core/duck_array_ops.py

@@ -13,7 +13,7 @@
 
 from . import dask_array_ops, dtypes, npcompat, nputils
 from .nputils import nanfirst, nanlast
-from .pycompat import dask_array_type
+from .pycompat import dask_array_type, sparse_array_type
 
 try:
     import dask.array as dask_array
@@ -64,6 +64,7 @@ def fail_on_dask_array_input(values, msg=None, func_name=None):
 around = _dask_or_eager_func('around')
 isclose = _dask_or_eager_func('isclose')
 
+
 if hasattr(np, 'isnat') and (
         dask_array is None or hasattr(dask_array_type, '__array_ufunc__')):
     # np.isnat is available since NumPy 1.13, so __array_ufunc__ is always
@@ -153,7 +154,11 @@ def trapz(y, x, axis):
 
 
 def asarray(data):
-    return data if isinstance(data, dask_array_type) else np.asarray(data)
+    return (
+        data if (isinstance(data, dask_array_type)
+                 or hasattr(data, '__array_function__'))
+        else np.asarray(data)
+    )
 
 
 def as_shared_dtype(scalars_or_arrays):
@@ -170,6 +175,9 @@ def as_shared_dtype(scalars_or_arrays):
 def as_like_arrays(*data):
     if all(isinstance(d, dask_array_type) for d in data):
         return data
+    elif any(isinstance(d, sparse_array_type) for d in data):
+        from sparse import COO
+        return tuple(COO(d) for d in data)
     else:
         return tuple(np.asarray(d) for d in data)
 

xarray/core/formatting.py

@@ -357,7 +357,10 @@ def set_numpy_options(*args, **kwargs):
 
 
 def short_array_repr(array):
-    array = np.asarray(array)
+
+    if not hasattr(array, '__array_function__'):
+        array = np.asarray(array)
+
     # default to lower precision so a full (abbreviated) line can fit on
     # one line with the default display_width
     options = {
@@ -394,7 +397,7 @@ def short_data_repr(array):
     if isinstance(getattr(array, 'variable', array)._data, dask_array_type):
         return short_dask_repr(array)
     elif array._in_memory or array.size < 1e5:
-        return short_array_repr(array.values)
+        return short_array_repr(array.data)
     else:
         return u'[{} values with dtype={}]'.format(array.size, array.dtype)
 

xarray/core/indexing.py

@@ -657,6 +657,9 @@ def as_indexable(array):
         return PandasIndexAdapter(array)
     if isinstance(array, dask_array_type):
         return DaskIndexingAdapter(array)
+    if hasattr(array, '__array_function__'):
+        return NdArrayLikeIndexingAdapter(array)
+
     raise TypeError('Invalid array type: {}'.format(type(array)))
 
 
@@ -1189,6 +1192,16 @@ def __setitem__(self, key, value):
                 raise
 
 
+class NdArrayLikeIndexingAdapter(NumpyIndexingAdapter):
+    def __init__(self, array):
+        if not hasattr(array, '__array_function__'):
+            raise TypeError(
+                'NdArrayLikeIndexingAdapter must wrap an object that '
+                'implements the __array_function__ protocol'
+            )
+        self.array = array
+
+
 class DaskIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
     """Wrap a dask array to support explicit indexing."""
 

xarray/core/npcompat.py

@@ -357,3 +357,18 @@ def moveaxis(a, source, destination):
 # https://github.com/numpy/numpy/issues/7370
 # https://github.com/numpy/numpy-stubs/
 DTypeLike = Union[np.dtype, str]
+
+
+# from dask/array/utils.py
+def _is_nep18_active():
+    class A:
+        def __array_function__(self, *args, **kwargs):
+            return True
+
+    try:
+        return np.concatenate([A()])
+    except ValueError:
+        return False
+
+
+IS_NEP18_ACTIVE = _is_nep18_active()

xarray/core/pycompat.py

@@ -8,3 +8,10 @@
     dask_array_type = (dask.array.Array,)
 except ImportError:  # pragma: no cover
     dask_array_type = ()
+
+try:
+    # solely for isinstance checks
+    import sparse
+    sparse_array_type = (sparse.SparseArray,)
+except ImportError:  # pragma: no cover
+    sparse_array_type = ()

xarray/core/utils.py

@@ -243,7 +243,9 @@ def is_scalar(value: Any) -> bool:
     return (
         getattr(value, 'ndim', None) == 0 or
         isinstance(value, (str, bytes)) or not
-        isinstance(value, (Iterable, ) + dask_array_type))
+        (isinstance(value, (Iterable, ) + dask_array_type) or
+         hasattr(value, '__array_function__'))
+    )
 
 
 def is_valid_numpy_dtype(dtype: Any) -> bool:

xarray/core/variable.py

@@ -17,6 +17,7 @@
     as_indexable)
 from .options import _get_keep_attrs
 from .pycompat import dask_array_type, integer_types
+from .npcompat import IS_NEP18_ACTIVE
 from .utils import (
     OrderedSet, decode_numpy_dict_values, either_dict_or_kwargs,
     ensure_us_time_resolution)
@@ -179,6 +180,18 @@ def as_compatible_data(data, fastpath=False):
         else:
             data = np.asarray(data)
 
+    if not isinstance(data, np.ndarray):
+        if hasattr(data, '__array_function__'):
+            if IS_NEP18_ACTIVE:
+                return data
+            else:
+                raise TypeError(
+                    'Got an NumPy-like array type providing the '
+                    '__array_function__ protocol but NEP18 is not enabled. '
+                    'Check that numpy >= v1.16 and that the environment '
+                    'variable "NUMPY_EXPERIMENTAL_ARRAY_FUNCTION" is set to '
+                    '"1"')
+
     # validate whether the data is valid data types
     data = np.asarray(data)
 
@@ -288,7 +301,7 @@ def _in_memory(self):
 
     @property
     def data(self):
-        if isinstance(self._data, dask_array_type):
+        if hasattr(self._data, '__array_function__'):
             return self._data
         else:
             return self.values
@@ -320,7 +333,7 @@ def load(self, **kwargs):
         """
         if isinstance(self._data, dask_array_type):
             self._data = as_compatible_data(self._data.compute(**kwargs))
-        elif not isinstance(self._data, np.ndarray):
+        elif not hasattr(self._data, '__array_function__'):
             self._data = np.asarray(self._data)
         return self
 
@@ -705,8 +718,8 @@ def __setitem__(self, key, value):
 
         if new_order:
             value = duck_array_ops.asarray(value)
-            value = value[(len(dims) - value.ndim) * (np.newaxis,) +
-                          (Ellipsis,)]
+            value = value[(len(dims) - value.ndim) * (np.newaxis,)
+                          + (Ellipsis,)]
             value = duck_array_ops.moveaxis(
                 value, new_order, range(len(new_order)))
 
@@ -805,7 +818,8 @@ def copy(self, deep=True, data=None):
                 data = indexing.MemoryCachedArray(data.array)
 
             if deep:
-                if isinstance(data, dask_array_type):
+                if (hasattr(data, '__array_function__')
+                        or isinstance(data, dask_array_type)):
                     data = data.copy()
                 elif not isinstance(data, PandasIndexAdapter):
                     # pandas.Index is immutable
@@ -1494,9 +1508,10 @@ def equals(self, other, equiv=duck_array_ops.array_equiv):
         """
         other = getattr(other, 'variable', other)
         try:
-            return (self.dims == other.dims and
-                    (self._data is other._data or
-                     equiv(self.data, other.data)))
+            return (
+                self.dims == other.dims and
+                (self._data is other._data or equiv(self.data, other.data))
+            )
         except (TypeError, AttributeError):
             return False
 
@@ -1517,8 +1532,8 @@ def identical(self, other):
         """Like equals, but also checks attributes.
         """
         try:
-            return (utils.dict_equiv(self.attrs, other.attrs) and
-                    self.equals(other))
+            return (utils.dict_equiv(self.attrs, other.attrs)
+                    and self.equals(other))
         except (TypeError, AttributeError):
             return False
 
@@ -1959,8 +1974,8 @@ def equals(self, other, equiv=None):
         # otherwise use the native index equals, rather than looking at _data
         other = getattr(other, 'variable', other)
         try:
-            return (self.dims == other.dims and
-                    self._data_equals(other))
+            return (self.dims == other.dims
+                    and self._data_equals(other))
         except (TypeError, AttributeError):
             return False