implement Gradient

doc/api.rst

@@ -25,6 +25,7 @@ Top-level functions
    zeros_like
    ones_like
    dot
+   differentiate
 
 Dataset
 =======
@@ -150,6 +151,7 @@ Computation
    Dataset.resample
    Dataset.diff
    Dataset.quantile
+   Dataset.differentiate
 
 **Aggregation**:
 :py:attr:`~Dataset.all`
@@ -317,6 +319,7 @@ Computation
    DataArray.diff
    DataArray.dot
    DataArray.quantile
+   DataArray.differentiate
 
 **Aggregation**:
 :py:attr:`~DataArray.all`

doc/whats-new.rst

@@ -36,6 +36,10 @@ Documentation
 Enhancements
 ~~~~~~~~~~~~
 
+- :py:func:`~xarray.differentiate`, :py:meth:`~xarray.DataArray.differentiate`,
+  and :py:meth:`~xarray.Dataset.differentiate` are newly added.
+  (:issue:`1332`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - Default colormap for sequential and divergent data can now be set via
   :py:func:`~xarray.set_options()`
   (:issue:`2394`)

xarray/__init__.py

@@ -10,7 +10,7 @@
 from .core.alignment import align, broadcast, broadcast_arrays
 from .core.common import full_like, zeros_like, ones_like
 from .core.combine import concat, auto_combine
-from .core.computation import apply_ufunc, where, dot
+from .core.computation import apply_ufunc, differentiate, dot, where
 from .core.extensions import (register_dataarray_accessor,
                               register_dataset_accessor)
 from .core.variable import as_variable, Variable, IndexVariable, Coordinate

xarray/core/computation.py

@@ -1137,3 +1137,81 @@ def where(cond, x, y):
                        join='exact',
                        dataset_join='exact',
                        dask='allowed')
+
+
+def _differentiate_variable(variable, coord, edge_order):
+    """ Compute the gradient along 1 dimension for Variable.
+    variable, coord: Variable
+    """
+    from .variable import Variable
+
+    result_array = duck_array_ops.gradient(
+        variable.data, coord.data, edge_order=edge_order,
+        axis=variable.get_axis_num(coord.dims[0]))
+    return Variable(variable.dims, result_array)
+
+
+def differentiate(dataarray, coord, edge_order=1):
+    """ Compute the gradient of the dataarray with the second order accurate
+    central differences.
+
+    Parameters
+    ----------
+    dataarray: xr.DataArray
+        Target array
+    coord: str
+        The coordinate to be used to differentiate the array.
+    edge_order: 1 or 2. Default 1
+        N-th order accurate differences at the boundaries.
+
+    Returns
+    -------
+    differentiate: DataArray or a sequence of DataArrays
+
+    See also
+    --------
+    numpy.gradient: corresponding numpy function
+
+    Examples
+    --------
+
+    >>> da = xr.DataArray(np.arange(12).reshape(4, 3), dims=['x', 'y'],
+    ...                   coords={'x': [0, 0.1, 1.1, 1.2]})
+    >>> da
+    <xarray.DataArray (x: 4, y: 3)>
+    array([[ 0,  1,  2],
+           [ 3,  4,  5],
+           [ 6,  7,  8],
+           [ 9, 10, 11]])
+    Coordinates:
+      * x        (x) float64 0.0 0.1 1.1 1.2
+    Dimensions without coordinates: y
+    >>>
+    >>> xr.differentiate(da, 'x')
+    <xarray.DataArray (x: 4, y: 3)>
+    array([[30.      , 30.      , 30.      ],
+           [27.545455, 27.545455, 27.545455],
+           [27.545455, 27.545455, 27.545455],
+           [30.      , 30.      , 30.      ]])
+    Coordinates:
+      * x        (x) float64 0.0 0.1 1.1 1.2
+    Dimensions without coordinates: y
+    """
+    from .dataarray import DataArray
+
+    if not isinstance(dataarray, DataArray):
+        raise TypeError(
+            'Only DataArray is supported. Given {}.'.format(type(dataarray)))
+
+    if coord not in dataarray.coords and coord not in dataarray.dims:
+        raise ValueError('Coordiante {} does not exist.'.format(coord))
+
+    coord_var = dataarray[coord].variable
+    if coord_var.ndim != 1:
+        raise ValueError(
+            'Only 1d-coordinate is supported. {} is {} '
+            'dimensional.'.format(coord, dataarray[coord].ndim))
+
+    return DataArray(
+        _differentiate_variable(dataarray.variable, coord_var, edge_order),
+        coords=dataarray.coords)

xarray/core/dask_array_compat.py

@@ -1,6 +1,9 @@
 from __future__ import absolute_import, division, print_function
 
+from distutils.version import LooseVersion
+
 import numpy as np
+from dask import __version__ as dask_version
 import dask.array as da
 
 try:
@@ -30,3 +33,124 @@ def isin(element, test_elements, assume_unique=False, invert=False):
         if invert:
             result = ~result
         return result
+
+
+if LooseVersion(dask_version) > LooseVersion('1.19.2'):
+    gradient = da.gradient
+
+else:  # pragma: no cover
+    # Copied from dask v0.19.2
+    # Used under the terms of Dask's license, see licenses/DASK_LICENSE.
+    import math
+    from numbers import Integral, Real
+
+    AxisError = np.AxisError
+
+    def validate_axis(axis, ndim):
+        """ Validate an input to axis= keywords """
+        if isinstance(axis, (tuple, list)):
+            return tuple(validate_axis(ax, ndim) for ax in axis)
+        if not isinstance(axis, Integral):
+            raise TypeError("Axis value must be an integer, got %s" % axis)
+        if axis < -ndim or axis >= ndim:
+            raise AxisError("Axis %d is out of bounds for array of dimension "
+                            "%d" % (axis, ndim))
+        if axis < 0:
+            axis += ndim
+        return axis
+
+    def _gradient_kernel(x, block_id, coord, axis, array_locs, grad_kwargs):
+        """
+        x: nd-array
+            array of one block
+        coord: 1d-array or scalar
+            coordinate along which the gradient is computed.
+        axis: int
+            axis along which the gradient is computed
+        array_locs:
+            actual location along axis. None if coordinate is scalar
+        grad_kwargs:
+            keyword to be passed to np.gradient
+        """
+        block_loc = block_id[axis]
+        if array_locs is not None:
+            coord = coord[array_locs[0][block_loc]:array_locs[1][block_loc]]
+        grad = np.gradient(x, coord, axis=axis, **grad_kwargs)
+        return grad
+
+    def gradient(f, *varargs, **kwargs):
+        f = da.asarray(f)
+
+        kwargs["edge_order"] = math.ceil(kwargs.get("edge_order", 1))
+        if kwargs["edge_order"] > 2:
+            raise ValueError("edge_order must be less than or equal to 2.")
+
+        drop_result_list = False
+        axis = kwargs.pop("axis", None)
+        if axis is None:
+            axis = tuple(range(f.ndim))
+        elif isinstance(axis, Integral):
+            drop_result_list = True
+            axis = (axis,)
+
+        axis = validate_axis(axis, f.ndim)
+
+        if len(axis) != len(set(axis)):
+            raise ValueError("duplicate axes not allowed")
+
+        axis = tuple(ax % f.ndim for ax in axis)
+
+        if varargs == ():
+            varargs = (1,)
+        if len(varargs) == 1:
+            varargs = len(axis) * varargs
+        if len(varargs) != len(axis):
+            raise TypeError(
+                "Spacing must either be a single scalar, or a scalar / "
+                "1d-array per axis"
+            )
+
+        if issubclass(f.dtype.type, (np.bool8, Integral)):
+            f = f.astype(float)
+        elif issubclass(f.dtype.type, Real) and f.dtype.itemsize < 4:
+            f = f.astype(float)
+
+        results = []
+        for i, ax in enumerate(axis):
+            for c in f.chunks[ax]:
+                if np.min(c) < kwargs["edge_order"] + 1:
+                    raise ValueError(
+                        'Chunk size must be larger than edge_order + 1. '
+                        'Minimum chunk for aixs {} is {}. Rechunk to '
+                        'proceed.'.format(np.min(c), ax))
+
+            if np.isscalar(varargs[i]):
+                array_locs = None
+            else:
+                if isinstance(varargs[i], da.Array):
+                    raise NotImplementedError(
+                        'dask array coordinated is not supported.')
+                # coordinate position for each block taking overlap into
+                # account
+                chunk = np.array(f.chunks[ax])
+                array_loc_stop = np.cumsum(chunk) + 1
+                array_loc_start = array_loc_stop - chunk - 2
+                array_loc_stop[-1] -= 1
+                array_loc_start[0] = 0
+                array_locs = (array_loc_start, array_loc_stop)
+
+            results.append(f.map_overlap(
+                _gradient_kernel,
+                dtype=f.dtype,
+                depth={j: 1 if j == ax else 0 for j in range(f.ndim)},
+                boundary="none",
+                coord=varargs[i],
+                axis=ax,
+                array_locs=array_locs,
+                grad_kwargs=kwargs,
+            ))
+
+        if drop_result_list:
+            results = results[0]
+
+        return results

xarray/core/dataarray.py

@@ -2010,6 +2010,9 @@ def diff(self, dim, n=1, label='upper'):
         Coordinates:
         * x        (x) int64 3 4
 
+        See Also
+        --------
+        DataArray.differentiate
         """
         ds = self._to_temp_dataset().diff(n=n, dim=dim, label=label)
         return self._from_temp_dataset(ds)
@@ -2289,6 +2292,54 @@ def rank(self, dim, pct=False, keep_attrs=False):
         ds = self._to_temp_dataset().rank(dim, pct=pct, keep_attrs=keep_attrs)
         return self._from_temp_dataset(ds)
 
+    def differentiate(self, coord, edge_order=1):
+        """ Differentiate the array with the second order accurate central
+        differences.
+
+        Parameters
+        ----------
+        coord: str
+            The coordinate to be used to compute the gradient.
+        edge_order: 1 or 2. Default 1
+            N-th order accurate differences at the boundaries.
+
+        Returns
+        -------
+        differentiated: DataArray
+
+        See also
+        --------
+        numpy.gradient: corresponding numpy function
+        xr.differentiate: more numpy-like function for xarray object.
+
+        Examples
+        --------
+
+        >>> da = xr.DataArray(np.arange(12).reshape(4, 3), dims=['x', 'y'],
+        ...                   coords={'x': [0, 0.1, 1.1, 1.2]})
+        >>> da
+        <xarray.DataArray (x: 4, y: 3)>
+        array([[ 0,  1,  2],
+               [ 3,  4,  5],
+               [ 6,  7,  8],
+               [ 9, 10, 11]])
+        Coordinates:
+          * x        (x) float64 0.0 0.1 1.1 1.2
+        Dimensions without coordinates: y
+        >>>
+        >>> da.differentiate('x')
+        <xarray.DataArray (x: 4, y: 3)>
+        array([[30.      , 30.      , 30.      ],
+               [27.545455, 27.545455, 27.545455],
+               [27.545455, 27.545455, 27.545455],
+               [30.      , 30.      , 30.      ]])
+        Coordinates:
+          * x        (x) float64 0.0 0.1 1.1 1.2
+        Dimensions without coordinates: y
+        """
+        ds = self._to_temp_dataset().differentiate(coord, edge_order)
+        return self._from_temp_dataset(ds)
+
 
 # priority most be higher than Variable to properly work with binary ufuncs
 ops.inject_all_ops_and_reduce_methods(DataArray, priority=60)