pydata · shoyer · Mar 12, 2018 · Mar 5, 2018 · Mar 5, 2018 · Mar 5, 2018
diff --git a/asv_bench/benchmarks/rolling.py b/asv_bench/benchmarks/rolling.py
@@ -1,9 +1,8 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
+from __future__ import absolute_import, division, print_function
 
 import numpy as np
 import pandas as pd
+
 import xarray as xr
 
 from . import parameterized, randn, requires_dask

diff --git a/doc/api.rst b/doc/api.rst
@@ -357,6 +357,13 @@ Reshaping and reorganizing
 Universal functions
 ===================
 
+.. warning::
+
+   With recent versions of numpy, dask and xarray, NumPy ufuncs are now
+   supported directly on all xarray and dask objects. This obliviates the need
+   for the ``xarray.ufuncs`` module, which should not be used for new code
+   unless compatibility with versions of NumPy prior to v1.13 is required.
+
 This functions are copied from NumPy, but extended to work on NumPy arrays,
 dask arrays and all xarray objects. You can find them in the ``xarray.ufuncs``
 module:

diff --git a/doc/computation.rst b/doc/computation.rst
@@ -341,21 +341,15 @@ Datasets support most of the same methods found on data arrays:
     ds.mean(dim='x')
     abs(ds)
 
-Unfortunately, we currently do not support NumPy ufuncs for datasets [1]_.
-:py:meth:`~xarray.Dataset.apply` works around this
-limitation, by applying the given function to each variable in the dataset:
+Datasets also support NumPy ufuncs (requires NumPy v1.13 or newer), or
+alternatively you can use :py:meth:`~xarray.Dataset.apply` to apply a function
+to each variable in a dataset:
 
 .. ipython:: python
 
+    np.sin(ds)
     ds.apply(np.sin)
 
-You can also use the wrapped functions in the ``xarray.ufuncs`` module:
-
-.. ipython:: python
-
-    import xarray.ufuncs as xu
-    xu.sin(ds)
-
 Datasets also use looping over variables for *broadcasting* in binary
 arithmetic. You can do arithmetic between any ``DataArray`` and a dataset:
 
@@ -373,10 +367,6 @@ Arithmetic between two datasets matches data variables of the same name:
 Similarly to index based alignment, the result has the intersection of all
 matching data variables.
 
-.. [1] This was previously due to a limitation of NumPy, but with NumPy 1.13
-       we should be able to support this by leveraging ``__array_ufunc__``
-       (:issue:`1617`).
-
 .. _comput.wrapping-custom:
 
 Wrapping custom computation

diff --git a/doc/gallery/control_plot_colorbar.py b/doc/gallery/control_plot_colorbar.py
@@ -7,9 +7,10 @@
 Use ``cbar_kwargs`` keyword to specify the number of ticks.
 The ``spacing`` kwarg can be used to draw proportional ticks.
 """
-import xarray as xr
 import matplotlib.pyplot as plt
 
+import xarray as xr
+
 # Load the data
 air_temp = xr.tutorial.load_dataset('air_temperature')
 air2d = air_temp.air.isel(time=500)

diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -38,6 +38,18 @@ Documentation
 Enhancements
 ~~~~~~~~~~~~
 
+- Implemented NumPy's ``__array_ufunc__`` protocol for all xarray objects
+  (:issue:`1617`). This enables using NumPy ufuncs directly on
+  ``xarray.Dataset`` objects with recent versions of NumPy (v1.13 and newer):
+
+  .. ipython:: python
+
+    ds = xr.Dataset({'a': 1})
+    np.sin(ds)
+
+  This obliviates the need for the ``xarray.ufuncs`` module, which will be
+  deprecated in the future when xarray drops support for older versions of
+  NumPy. By `Stephan Hoyer <https://github.com/shoyer>`_.
 - Improve :py:func:`~xarray.DataArray.rolling` logic.
   :py:func:`~xarray.DataArrayRolling` object now supports
   :py:func:`~xarray.DataArrayRolling.construct` method that returns a view

diff --git a/xarray/core/common.py b/xarray/core/common.py
@@ -1,12 +1,16 @@
 from __future__ import absolute_import, division, print_function
 
+import numbers
 import warnings
 
 import numpy as np
 import pandas as pd
 
 from . import dtypes, formatting, ops
-from .pycompat import OrderedDict, basestring, dask_array_type, suppress
+from .options import OPTIONS
+from .pycompat import (
+    OrderedDict, basestring, bytes_type, dask_array_type, suppress,
+    unicode_type)
 from .utils import Frozen, SortedKeysDict, not_implemented
 
 
@@ -235,7 +239,65 @@ def get_squeeze_dims(xarray_obj, dim, axis=None):
     return dim
 
 
-class BaseDataObject(AttrAccessMixin):
+class SupportsArithmetic(object):  # noqa: W1641
+    """Base class for Dataset, DataArray, Variable and GroupBy."""
+
+    # TODO: implement special methods for arithmetic here rather than injecting
+    # them in xarray/core/ops.py. Ideally, do so by inheriting from
+    # numpy.lib.mixins.NDArrayOperatorsMixin.
+
+    # TODO: allow extending this with some sort of registration system
+    _HANDLED_TYPES = (np.ndarray, np.generic, numbers.Number, bytes_type,
+                      unicode_type) + dask_array_type
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        from .computation import apply_ufunc
+
+        # See the docstring example for numpy.lib.mixins.NDArrayOperatorsMixin.
+        out = kwargs.get('out', ())
+        for x in inputs + out:
+            if not isinstance(x, self._HANDLED_TYPES + (SupportsArithmetic,)):
+                return NotImplemented
+
+        if ufunc.signature is not None:
+            raise NotImplementedError(
+                '{} not supported: xarray objects do not directly implement '
+                'generalized ufuncs. Instead, use xarray.apply_ufunc.'
+                .format(ufunc))
+
+        if method != '__call__':
+            # TODO: support other methods, e.g., reduce and accumulate.
+            raise NotImplementedError(
+                '{} method for ufunc {} is not implemented on xarray objects, '
+                'which currently only support the __call__ method.'
+                .format(method, ufunc))
+
+        if any(isinstance(o, SupportsArithmetic) for o in out):
+            # TODO: implement this with logic like _inplace_binary_op. This
+            # will be necessary to use NDArrayOperatorsMixin.
+            raise NotImplementedError(
+                'xarray objects are not yet supported in the `out` argument '
+                'for ufuncs.')
+
+        join = dataset_join = OPTIONS['arithmetic_join']
+
+        return apply_ufunc(ufunc, *inputs,
+                           input_core_dims=((),) * ufunc.nin,
+                           output_core_dims=((),) * ufunc.nout,
+                           join=join,
+                           dataset_join=dataset_join,
+                           dataset_fill_value=np.nan,
+                           kwargs=kwargs,
+                           dask='allowed')
+
+    # this has no runtime function - these are listed so IDEs know these
+    # methods are defined and don't warn on these operations
+    __lt__ = __le__ = __ge__ = __gt__ = __add__ = __sub__ = __mul__ = \
+        __truediv__ = __floordiv__ = __mod__ = __pow__ = __and__ = __xor__ = \
+        __or__ = __div__ = __eq__ = __ne__ = not_implemented
+
+
+class DataWithCoords(SupportsArithmetic, AttrAccessMixin):
     """Shared base class for Dataset and DataArray."""
 
     def squeeze(self, dim=None, drop=False, axis=None):
@@ -749,12 +811,6 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         self.close()
 
-    # this has no runtime function - these are listed so IDEs know these
-    # methods are defined and don't warn on these operations
-    __lt__ = __le__ = __ge__ = __gt__ = __add__ = __sub__ = __mul__ = \
-        __truediv__ = __floordiv__ = __mod__ = __pow__ = __and__ = __xor__ = \
-        __or__ = __div__ = __eq__ = __ne__ = not_implemented
-
 
 def full_like(other, fill_value, dtype=None):
     """Return a new object with the same shape and type as a given object.

diff --git a/xarray/core/dask_array_ops.py b/xarray/core/dask_array_ops.py
@@ -1,8 +1,7 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
+from __future__ import absolute_import, division, print_function
 
 import numpy as np
+
 from . import nputils
 
 try:

diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
@@ -10,7 +10,7 @@
 from ..plot.plot import _PlotMethods
 from .accessors import DatetimeAccessor
 from .alignment import align, reindex_like_indexers
-from .common import AbstractArray, BaseDataObject
+from .common import AbstractArray, DataWithCoords
 from .coordinates import (
     DataArrayCoordinates, Indexes, LevelCoordinatesSource,
     assert_coordinate_consistent, remap_label_indexers)
@@ -117,7 +117,7 @@ def __setitem__(self, key, value):
 _THIS_ARRAY = utils.ReprObject('<this-array>')
 
 
-class DataArray(AbstractArray, BaseDataObject):
+class DataArray(AbstractArray, DataWithCoords):
     """N-dimensional array with labeled coordinates and dimensions.
 
     DataArray provides a wrapper around numpy ndarrays that uses labeled

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -17,7 +17,7 @@
     rolling, utils)
 from .. import conventions
 from .alignment import align
-from .common import BaseDataObject, ImplementsDatasetReduce
+from .common import DataWithCoords, ImplementsDatasetReduce
 from .coordinates import (
     DatasetCoordinates, Indexes, LevelCoordinatesSource,
     assert_coordinate_consistent, remap_label_indexers)
@@ -298,7 +298,7 @@ def __getitem__(self, key):
         return self.dataset.sel(**key)
 
 
-class Dataset(Mapping, ImplementsDatasetReduce, BaseDataObject,
+class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords,
               formatting.ReprMixin):
     """A multi-dimensional, in memory, array database.
 
@@ -2362,7 +2362,7 @@ def dropna(self, dim, how='any', thresh=None, subset=None):
             array = self._variables[k]
             if dim in array.dims:
                 dims = [d for d in array.dims if d != dim]
-                count += array.count(dims)
+                count += np.asarray(array.count(dims))
                 size += np.prod([self.dims[d] for d in dims])
 
         if thresh is not None:

diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py
@@ -7,7 +7,8 @@
 
 from . import dtypes, duck_array_ops, nputils, ops
 from .combine import concat
-from .common import ImplementsArrayReduce, ImplementsDatasetReduce
+from .common import (
+    ImplementsArrayReduce, ImplementsDatasetReduce, SupportsArithmetic)
 from .pycompat import integer_types, range, zip
 from .utils import hashable, maybe_wrap_array, peek_at, safe_cast_to_index
 from .variable import IndexVariable, Variable, as_variable
@@ -151,7 +152,7 @@ def _unique_and_monotonic(group):
         return index.is_unique and index.is_monotonic
 
 
-class GroupBy(object):
+class GroupBy(SupportsArithmetic):
     """A object that implements the split-apply-combine pattern.
 
     Modeled after `pandas.GroupBy`. The `GroupBy` object can be iterated over

diff --git a/xarray/core/npcompat.py b/xarray/core/npcompat.py
@@ -1,8 +1,8 @@
 from __future__ import absolute_import, division, print_function
 
-import numpy as np
 from distutils.version import LooseVersion
 
+import numpy as np
 
 if LooseVersion(np.__version__) >= LooseVersion('1.12'):
     as_strided = np.lib.stride_tricks.as_strided

diff --git a/xarray/core/variable.py b/xarray/core/variable.py
@@ -216,8 +216,8 @@ def _as_array_or_item(data):
     return data
 
 
-class Variable(common.AbstractArray, utils.NdimSizeLenMixin):
-
+class Variable(common.AbstractArray, common.SupportsArithmetic,
+               utils.NdimSizeLenMixin):
     """A netcdf-like variable consisting of dimensions, data and attributes
     which describe a single Array. A single Variable object is not fully
     described outside the context of its parent Dataset (if you want such a

diff --git a/xarray/tests/test_nputils.py b/xarray/tests/test_nputils.py
@@ -1,8 +1,8 @@
 import numpy as np
5DF5from numpy.testing import assert_array_equal
 
-from xarray.core.nputils import (NumpyVIndexAdapter, _is_contiguous,
-                                 rolling_window)
+from xarray.core.nputils import (
+    NumpyVIndexAdapter, _is_contiguous, rolling_window)
 
 
 def test_is_contiguous():