Add DataArray.pad, Dataset.pad, Variable.pad

xarray/core/dask_array_compat.py

@@ -89,3 +89,63 @@ def meta_from_array(x, ndim=None, dtype=None):
             meta = meta.astype(dtype)
 
         return meta
+
+
+# TODO figure out how Dask versioning works
+# if LooseVersion(dask_version) >= LooseVersion("1.7.0"):
+try:
+    pad = da.pad
+except AttributeError:
+    def pad(array, pad_width, mode="constant", **kwargs):
+        """
+        Return a new dask.DataArray wit padding. This functions implements a
+        constant padding for versions of Dask that do not implement this yet.
+
+        Parameters
+        ----------
+        array: Array to pad
+
+        pad_width: List of the form [(before, after)]
+            Number of values padded to the edges of axis.
+        """
+        if mode != "constant":
+            raise NotImplementedError(
+                "Pad is not yet implemented for your current version of Dask. "
+                "Please update your version of Dask or use the "
+                "mode=`constant`, that is added by xarray."
+            )
+
+        try:
+            fill_value = kwargs["constant_values"]
+            dtype = array.dtype
+        except KeyError:
+            dtype, fill_value = dtypes.maybe_promote(array.dtype)
+
+        for axis, pad in enumerate(pad_width):
+            before_shape = list(array.shape)
+            before_shape[axis] = pad[0]
+            before_chunks = list(array.chunks)
+            before_chunks[axis] = (pad[0],)
+            after_shape = list(array.shape)
+            after_shape[axis] = pad[1]
+            after_chunks = list(array.chunks)
+            after_chunks[axis] = (pad[1],)
+
+            arrays = []
+            if pad[0] > 0:
+                arrays.append(
+                    da.full(
+                        before_shape, fill_value, dtype=dtype, chunks=before_chunks
+                    )
+                )
+            arrays.append(array)
+            if pad[1] > 0:
+                arrays.append(
+                    da.full(
+                        after_shape, fill_value, dtype=dtype, chunks=after_chunks
+                    )
+                )
+            if len(arrays) > 1:
+                array = da.concatenate(arrays, axis=axis)
+
+        return array

xarray/core/dataarray.py

@@ -3110,6 +3110,158 @@ def map_blocks(
 
         return map_blocks(func, self, args, kwargs)
 
+    def pad(
+        self,
+        pad_widths: Mapping[Hashable, Tuple[int, int]] = None,
+        mode: str = "constant",
+        stat_length: Union[
+            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
+        ] = None,
+        constant_values: Union[
+            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
+        ] = None,
+        end_values: Union[
+            int, Tuple[int, int], Mapping[Hashable, Tuple[int, int]]
+        ] = None,
+        reflect_type: str = None,
+        **pad_widths_kwargs: Any,
+    ) -> "DataArray":
+        """Pad this array along one or more dimensions.
+
+        When using one of the modes ("edge", "reflect", "symmetric", "wrap"),
+        coordinates will be padded with the same mode, otherwise coordinates
+        are padded using the "constant" mode with fill_value dtypes.NA.
+
+        Parameters
+        ----------
+        pad_widths : Mapping with the form of {dim: (pad_before, pad_after)}
+            Number of values padded along each dimension.
+        mode : str (taken from numpy docs)
+            One of the following string values or a user supplied function.
+            'constant' (default)
+                Pads with a constant value.
+            'edge'
+                Pads with the edge values of array.
+            'linear_ramp'
+                Pads with the linear ramp between end_value and the
+                array edge value.
+            'maximum'
+                Pads with the maximum value of all or part of the
+                vector along each axis.
+            'mean'
+                Pads with the mean value of all or part of the
+                vector along each axis.
+            'median'
+                Pads with the median value of all or part of the
+                vector along each axis.
+            'minimum'
+                Pads with the minimum value of all or part of the
+                vector along each axis.
+            'reflect'
+                Pads with the reflection of the vector mirrored on
+                the first and last values of the vector along each
+                axis.
+            'symmetric'
+                Pads with the reflection of the vector mirrored
+                along the edge of the array.
+            'wrap'
+                Pads with the wrap of the vector along the axis.
+                The first values are used to pad the end and the
+                end values are used to pad the beginning.
+        stat_length : sequence or int, optional
+            Used in 'maximum', 'mean', 'median', and 'minimum'.  Number of
+            values at edge of each axis used to calculate the statistic value.
+            ((before_1, after_1), ... (before_N, after_N)) unique statistic
+            lengths for each axis.
+            ((before, after),) yields same before and after statistic lengths
+            for each axis.
+            (stat_length,) or int is a shortcut for before = after = statistic
+            length for all axes.
+            Default is ``None``, to use the entire axis.
+        constant_values : sequence or scalar, optional
+            Used in 'constant'.  The values to set the padded values for each
+            axis.
+            ``((before_1, after_1), ... (before_N, after_N))`` unique pad constants
+            for each axis.
+            ``((before, after),)`` yields same before and after constants for each
+            axis.
+            ``(constant,)`` or ``constant`` is a shortcut for ``before = after = constant`` for
+            all axes.
+            Default is 0.
+        end_values : sequence or scalar, optional
+            Used in 'linear_ramp'.  The values used for the ending value of the
+            linear_ramp and that will form the edge of the padded array.
+            ``((before_1, after_1), ... (before_N, after_N))`` unique end values
+            for each axis.
+            ``((before, after),)`` yields same before and after end values for each
+            axis.
+            ``(constant,)`` or ``constant`` is a shortcut for ``before = after = constant`` for
+            all axes.
+            Default is 0.
+        reflect_type : {'even', 'odd'}, optional
+            Used in 'reflect', and 'symmetric'.  The 'even' style is the
+            default with an unaltered reflection around the edge value.  For
+            the 'odd' style, the extended part of the array is created by
+            subtracting the reflected values from two times the edge value.
+
+        **pad_widths_kwargs:
+            The keyword arguments form of ``pad_widths``.
+            One of pad_widths or pad_widths_kwarg must be provided.
+
+        Returns
+        -------
+        padded : DataArray
+            DataArray with the padded coordinates and data.
+
+        See also
+        --------
+        shift
+        roll
+
+        Examples
+        --------
+
+        >>> arr = xr.DataArray([5, 6, 7], coords=[("x", [0,1,2])])
+        >>> arr.pad(x=(1,2), pad_options={"constant_values" : 0})
+        <xarray.DataArray (x: 6)>
+        array([0, 5, 6, 7, 0, 0])
+        Coordinates:
+          * x        (x) float64 nan 0.0 1.0 2.0 nan nan
+        """
+        pad_widths = either_dict_or_kwargs(pad_widths, pad_widths_kwargs, "pad")
+
+        variable = self.variable.pad(
+            pad_widths=pad_widths,
+            mode=mode,
+            stat_length=stat_length,
+            constant_values=constant_values,
+            end_values=end_values,
+            reflect_type=reflect_type,
+        )
+
+        if mode in ("edge", "reflect", "symmetric", "wrap"):
+            coord_pad_mode = mode
+            coord_pad_options = {
+                "stat_length": stat_length,
+                "constant_values": constant_values,
+                "end_values": end_values,
+                "reflect_type": reflect_type,
+            }
+        else:
+            coord_pad_mode = "constant"
+            coord_pad_options = {}
+
+        coords = {}
+        for name, dim in self.coords.items():
+            if name in pad_widths:
+                coords[name] = dim.variable.pad(
+                    {name: pad_widths[name]}, mode=coord_pad_mode, **coord_pad_options
+                )
+            else:
+                coords[name] = as_variable(dim, name=name)
+
+        return self._replace(variable=variable, coords=coords)
+
     # this needs to be at the end, or mypy will confuse with `str`
     # https://mypy.readthedocs.io/en/latest/common_issues.html#dealing-with-conflicting-names
     str = property(StringAccessor)

xarray/core/duck_array_ops.py

@@ -17,8 +17,10 @@
 
 try:
     import dask.array as dask_array
+    from . import dask_array_compat
 except ImportError:
     dask_array = None  # type: ignore
+    dask_array_compat = None
 
 
 def _dask_or_eager_func(
@@ -454,3 +456,5 @@ def rolling_window(array, axis, window, center, fill_value):
         return dask_array_ops.rolling_window(array, axis, window, center, fill_value)
     else:  # np.ndarray
         return nputils.rolling_window(array, axis, window, center, fill_value)
+
+pad = _dask_or_eager_func("pad", dask_module=dask_array_compat)

xarray/core/rolling.py

@@ -330,7 +330,7 @@ def _bottleneck_reduce(self, func, **kwargs):
             else:
                 shift = (-self.window // 2) + 1
                 valid = (slice(None),) * axis + (slice(-shift, None),)
-            padded = padded.pad_with_fill_value({self.dim: (0, -shift)})
+            padded = padded.pad({self.dim: (0, -shift)}, mode="constant")
 
         if isinstance(padded.data, dask_array_type):
             raise AssertionError("should not be reachable")