"""
Functions to pad and unpad a N-dimensional regular grid
"""
import numpy as np
from xarray.core.utils import either_dict_or_kwargs
from .utils import get_spacing
[docs]def pad(
da,
pad_width=None,
mode="constant",
stat_length=None,
constant_values=0,
end_values=None,
reflect_type=None,
**pad_width_kwargs,
):
"""
Pad array with evenly spaced coordinates
Wraps the :meth:`xarray.DataArray.pad` method but also pads the evenly
spaced coordinates by extrapolation using the same coordinate spacing.
The ``pad_width`` used for each coordinate is stored as one of its
attributes.
Parameters
----------
da : :class:`xarray.DataArray`
Array to be padded. The coordinates along which the array will be
padded must be evenly spaced.
pad_width : mapping of hashable to tuple of int
Mapping with the form of {dim: (pad_before, pad_after)}
describing the number of values padded along each dimension.
{dim: pad} is a shortcut for pad_before = pad_after = pad
mode : str, default: "constant"
One of the following string values (taken from numpy docs).
- constant: 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 : int, tuple or mapping of hashable to tuple, default: None
Used in 'maximum', 'mean', 'median', and 'minimum'. Number of
values at edge of each axis used to calculate the statistic value.
{dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)} unique
statistic lengths along each dimension.
((before, after),) yields same before and after statistic lengths
for each dimension.
(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 : scalar, tuple or mapping of hashable to tuple, default: 0
Used in 'constant'. The values to set the padded values for each
axis.
``{dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)}`` unique
pad constants along each dimension.
``((before, after),)`` yields same before and after constants for each
dimension.
``(constant,)`` or ``constant`` is a shortcut for ``before = after = constant`` for
all dimensions.
Default is 0.
end_values : scalar, tuple or mapping of hashable to tuple, default: 0
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.
``{dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)}`` unique
end values along each dimension.
``((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_width_kwargs
The keyword arguments form of ``pad_width``.
One of ``pad_width`` or ``pad_width_kwargs`` must be provided.
Returns
-------
da_padded : :class:`xarray.DataArray`
Examples
--------
>>> import xarray as xr
>>> da = xr.DataArray(
... [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
... coords={"x": [0, 1, 2], "y": [-5, -4, -3]},
... dims=("y", "x"),
... )
>>> da_padded = pad(da, x=2, y=1)
>>> da_padded
<xarray.DataArray (y: 5, x: 7)>
array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 2, 3, 0, 0],
[0, 0, 4, 5, 6, 0, 0],
[0, 0, 7, 8, 9, 0, 0],
[0, 0, 0, 0, 0, 0, 0]])
Coordinates:
* x (x) int64 -2 -1 0 1 2 3 4
* y (y) int64 -6 -5 -4 -3 -2
>>> da_padded.x
<xarray.DataArray 'x' (x: 7)>
array([-2, -1, 0, 1, 2, 3, 4])
Coordinates:
* x (x) int64 -2 -1 0 1 2 3 4
Attributes:
pad_width: 2
>>> da_padded.y
<xarray.DataArray 'y' (y: 5)>
array([-6, -5, -4, -3, -2])
Coordinates:
* y (y) int64 -6 -5 -4 -3 -2
Attributes:
pad_width: 1
Asymmetric padding
>>> da_padded = pad(da, x=(1, 4))
>>> da_padded
<xarray.DataArray (y: 3, x: 8)>
array([[0, 1, 2, 3, 0, 0, 0, 0],
[0, 4, 5, 6, 0, 0, 0, 0],
[0, 7, 8, 9, 0, 0, 0, 0]])
Coordinates:
* x (x) int64 -1 0 1 2 3 4 5 6
* y (y) int64 -5 -4 -3
>>> da_padded.x
<xarray.DataArray 'x' (x: 8)>
array([-1, 0, 1, 2, 3, 4, 5, 6])
Coordinates:
* x (x) int64 -1 0 1 2 3 4 5 6
Attributes:
pad_width: (1, 4)
"""
# Redefine pad_width if pad_width_kwargs were passed
pad_width = either_dict_or_kwargs(pad_width, pad_width_kwargs, "pad")
# Pad the array using the xarray.DataArray.pad method
padded_da = da.pad(
pad_width,
mode,
stat_length,
constant_values,
end_values,
reflect_type,
)
# Pad the coordinates selected in pad_width
padded_coords = _pad_coordinates(da.coords, pad_width)
# Assign the padded coordinates to the padded array
padded_da = padded_da.assign_coords(padded_coords)
# Edit the attributes of the padded array
for dim in pad_width.keys():
# Add attrs of the original coords to the padded array
padded_da.coords[dim].attrs.update(da.coords[dim].attrs)
# Add the pad_width used for this coordinate
padded_da.coords[dim].attrs.update({"pad_width": pad_width[dim]})
# Return padded array
return padded_da
def _pad_coordinates(coords, pad_width):
"""
Pad coordinates arrays according to the passed pad_width
Parameters
----------
coords : dict-like object
Dictionary with coordinates as :class:`xarray.DataArray`.
Only the coordinates specified through ``pad_width`` will be padded.
Every coordinate that will be padded should be evenly spaced.
pad_width : dict-like object
Dictionary with the same keys as ``coords``. The coordinates specified
through ``pad_width`` are returned as padded.
Returns
-------
padded_coords : dict-like object
Dictionary with 1d-arrays corresponding to the padded coordinates.
Examples
--------
>>> import numpy as np
>>> import xarray as xr
>>> x = np.linspace(-4, -1, 4)
>>> y = np.linspace(-1, 4, 6)
>>> coords = {
... "x": xr.DataArray(x, coords={"x": x}, dims=("x",)),
... "y": xr.DataArray(y, coords={"y": y}, dims=("y",)),
... }
>>> pad_width = {"x": 2}
>>> padded_coords = _pad_coordinates(coords, pad_width)
>>> padded_coords["x"]
array([-6., -5., -4., -3., -2., -1., 0., 1.])
>>> padded_coords["y"]
<xarray.DataArray (y: 6)>
array([-1., 0., 1., 2., 3., 4.])
Coordinates:
* y (y) float64 -1.0 0.0 1.0 2.0 3.0 4.0
"""
# Generate a dictionary with the original coordinates
padded_coords = {dim: coords[dim] for dim in coords}
# Start padding the coordinates that appear in pad_width
for dim in pad_width:
# Get the spacing of the selected coordinate
# (raises an error if not evenly spaced)
spacing = get_spacing(padded_coords[dim])
# Pad the coordinates using numpy.pad with the _pad_coordinates callback
padded_coords[dim] = np.pad(
padded_coords[dim],
pad_width=pad_width[dim],
mode=_pad_coordinates_callback,
spacing=spacing, # spacing is passed as a kwarg to the callback
)
return padded_coords
def _pad_coordinates_callback(vector, iaxis_pad_width, iaxis, kwargs):
"""
Callback for padding coordinates
This function is not intended to be called, but to be passed as the
``mode`` method to the :func:`numpy.pad`
Parameters
----------
vector : 1d-array
A rank 1 array already padded with zeros. Padded values are
``vector[:iaxis_pad_width[0]]`` and ``vector[-iaxis_pad_width[1]:]``.
iaxis_pad_width : tuple
A 2-tuple of ints, ``iaxis_pad_width[0]`` represents the number of
values padded at the beginning of vector where ``iaxis_pad_width[1]``
represents the number of values padded at the end of vector.
iaxis : int
The axis currently being calculated. This parameter is not used, but
the function will check if it's equal to zero. It exists for
compatibility with the ``padding_func`` callback that :func:`numpy.pad`
needs.
kwargs : dict
Any keyword arguments the function requires. The kwargs are ignored in
this function, they exist for compatibility with the ``padding_func``
callback that :func:`numpy.pad` needs.
Returns
-------
vector : 1d-array
Padded vector.
"""
assert iaxis == 0
spacing = kwargs["spacing"]
n_start, n_end = iaxis_pad_width[:]
vmin, vmax = vector[n_start], vector[-(n_end + 1)]
vector[:n_start] = (
vmin - n_start * spacing + np.linspace(0, spacing * (n_start - 1), n_start)
)
vector[-n_end:] = vmax + spacing + np.linspace(0, spacing * (n_end - 1), n_end)
return vector
[docs]def unpad(da, pad_width=None, **pad_width_kwargs):
"""
Unpad an array and its coordinates
Undo the padding process of the :func:`xrft.pad` function by slicing the
passed :class:`xarray.DataArray` and its coordinates.
Parameters
----------
da : :class:`xarray.DataArray`
Padded array. The coordinates along which the array will be
padded must be evenly spaced.
Returns
-------
da_unpaded : :class:`xarray.DataArray`
Unpadded array.
pad_width : mapping of hashable to tuple of int (optional)
Mapping with the form of {dim: (pad_before, pad_after)}
describing the number of values padded along each dimension.
{dim: pad} is a shortcut for pad_before = pad_after = pad.
If ``None``, then the *pad_width* for each coordinate is read from
their ``pad_width`` attribute.
**pad_width_kwargs (optional)
The keyword arguments form of ``pad_width``.
Pass ``pad_width`` or ``pad_width_kwargs``.
See also
--------
:func:`xrft.pad`
Examples
--------
>>> import xarray as xr
>>> da = xr.DataArray(
... [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
... coords={"x": [0, 1, 2], "y": [-5, -4, -3]},
... dims=("y", "x"),
... )
>>> da_padded = pad(da, x=2, y=1)
>>> da_padded
<xarray.DataArray (y: 5, x: 7)>
array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 2, 3, 0, 0],
[0, 0, 4, 5, 6, 0, 0],
[0, 0, 7, 8, 9, 0, 0],
[0, 0, 0, 0, 0, 0, 0]])
Coordinates:
* x (x) int64 -2 -1 0 1 2 3 4
* y (y) int64 -6 -5 -4 -3 -2
>>> unpad(da_padded)
<xarray.DataArray (y: 3, x: 3)>
array([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
Coordinates:
* x (x) int64 0 1 2
* y (y) int64 -5 -4 -3
Custom ``pad_width``
>>> unpad(da_padded, x=1, y=1)
<xarray.DataArray (y: 3, x: 5)>
array([[0, 1, 2, 3, 0],
[0, 4, 5, 6, 0],
[0, 7, 8, 9, 0]])
Coordinates:
* x (x) int64 -1 0 1 2 3
* y (y) int64 -5 -4 -3
"""
# Generate the pad_width dictionary
if pad_width is None and not pad_width_kwargs:
# Read the pad_width from each coordinate if pad_width is None and
# no pad_width_kwargs has been passed
pad_width = {
dim: coord.attrs["pad_width"]
for dim, coord in da.coords.items()
if "pad_width" in coord.attrs
}
# Raise error if there's no pad_width attribute in the coordinates
if not pad_width:
raise ValueError(
"The passed array doesn't seem to be a padded one: the 'pad_width' "
+ "attribute was missing on every one of its coordinates. "
)
else:
# Redefine pad_width if pad_width_kwargs were passed
pad_width = either_dict_or_kwargs(pad_width, pad_width_kwargs, "pad")
# Transform every pad_width into a tuple with indices
slices = {}
for dim in pad_width:
slices[dim] = _pad_width_to_slice(pad_width[dim], da.coords[dim].size)
# Slice the padded array
unpadded_da = da.isel(indexers=slices)
# Remove the pad_width attribute from coords since it's no longer necessary
for dim in pad_width:
if "pad_width" in unpadded_da.coords[dim].attrs:
unpadded_da.coords[dim].attrs.pop("pad_width")
return unpadded_da
def _pad_width_to_slice(pad_width, size):
"""
Create a slice for removing the padded elements of a coordinate array
Parameters
----------
pad_width : int or tuple
Integer or tuples with the width of the padding at each side of the
coordinate array. An integer means an equal padding at each side equal
to its value.
size : int
Number of elements of the coordinates array.
Returns
-------
coord_slice : slice
A slice object for removing the padded elements of the coordinate
array.
"""
if type(pad_width) == int:
pad_width = (pad_width, pad_width)
return slice(pad_width[0], size - pad_width[1])