diff --git a/doc/api.rst b/doc/api.rst index 2eb953b8..8005fb76 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -28,3 +28,13 @@ You also may wish to use xrft's detrend function on its own. .. automodule:: xrft.detrend :members: + +padding +======= + +Pad and unpad arrays and its coordinates so they can be used for computing +FFTs. + +.. automodule:: xrft.padding + :members: + diff --git a/xrft/__init__.py b/xrft/__init__.py index f53aa986..a3e4c06e 100644 --- a/xrft/__init__.py +++ b/xrft/__init__.py @@ -5,3 +5,4 @@ from .xrft import * # noqa from .detrend import detrend +from .padding import pad, unpad diff --git a/xrft/padding.py b/xrft/padding.py new file mode 100644 index 00000000..c4b4937b --- /dev/null +++ b/xrft/padding.py @@ -0,0 +1,405 @@ +""" +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 + + +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 + + 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 + + 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 + + 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 + + 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 + + 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"] + + 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] = np.arange(vmin - spacing * n_start, vmin, spacing) + vector[-n_end:] = np.arange(vmax + spacing, vmax + spacing * (n_end + 1), spacing) + return vector + + +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 + + 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) + + 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) + + 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]) diff --git a/xrft/tests/test_padding.py b/xrft/tests/test_padding.py new file mode 100644 index 00000000..3a279fae --- /dev/null +++ b/xrft/tests/test_padding.py @@ -0,0 +1,230 @@ +""" +Unit tests for padding functions +""" +import pytest +import numpy as np +import xarray as xr +import xarray.testing as xrt +import numpy.testing as npt + +from ..padding import pad, _pad_coordinates, unpad, _pad_width_to_slice +from ..xrft import fft, ifft + + +@pytest.fixture +def sample_da_2d(): + """ + Defines a 2D sample xarray.DataArray + """ + x = np.linspace(0, 10, 11) + y = np.linspace(-4, 4, 17) + z = np.arange(11 * 17, dtype=float).reshape(17, 11) + # Create one xr.DataArray for each coordiante and add spacing and + # direct_lag attributes to them + dx, dy = x[1] - x[0], y[1] - y[0] + x = xr.DataArray( + x, coords={"x": x}, dims=("x",), attrs=dict(direct_lag=3.0, spacing=dx) + ) + y = xr.DataArray( + y, coords={"y": y}, dims=("y",), attrs=dict(direct_lag=-2.1, spacing=dy) + ) + return xr.DataArray(z, coords={"x": x, "y": y}, dims=("y", "x")) + + +def test_pad_coordinates(sample_da_2d): + """ + Test pad_coordinates function + """ + coords = sample_da_2d.coords + # Pad a single coordinate + padded_coords = _pad_coordinates(coords, {"x": 3}) + npt.assert_allclose(padded_coords["x"], np.linspace(-3, 13, 17)) + npt.assert_allclose(padded_coords["y"], coords["y"]) + # Pad two coordinates + padded_coords = _pad_coordinates(coords, {"x": 2, "y": 3}) + npt.assert_allclose(padded_coords["x"], np.linspace(-2, 12, 15)) + npt.assert_allclose(padded_coords["y"], np.linspace(-5.5, 5.5, 23)) + # Pad a single coordinate asymmetrically + padded_coords = _pad_coordinates(coords, {"x": (3, 2)}) + npt.assert_allclose(padded_coords["x"], np.linspace(-3, 12, 16)) + npt.assert_allclose(padded_coords["y"], coords["y"]) + # Pad two coordinates assymetrically + padded_coords = _pad_coordinates(coords, {"x": (2, 1), "y": (3, 4)}) + npt.assert_allclose(padded_coords["x"], np.linspace(-2, 11, 14)) + npt.assert_allclose(padded_coords["y"], np.linspace(-5.5, 6, 24)) + + +def test_pad_coordinates_invalid(sample_da_2d): + """ + Test if pad_coordinates raises error after unevenly spaced coords + """ + x = sample_da_2d.coords["x"].values + x[3] += 0.1 + sample_da_2d = sample_da_2d.assign_coords({"x": x}) + with pytest.raises(ValueError): + _pad_coordinates(sample_da_2d.coords, pad_width={"x": 2}) + + +def test_pad_with_kwargs(sample_da_2d): + """ + Test pad function by passing pad_width as kwargs + """ + padded_da = pad(sample_da_2d, x=2, y=1) + assert padded_da.shape == (19, 15) + npt.assert_allclose(padded_da.values[:1, :], 0) + npt.assert_allclose(padded_da.values[-1:, :], 0) + npt.assert_allclose(padded_da.values[:, :2], 0) + npt.assert_allclose(padded_da.values[:, -2:], 0) + npt.assert_allclose(padded_da.values[1:-1, 2:-2], sample_da_2d) + npt.assert_allclose(padded_da.x, np.linspace(-2, 12, 15)) + npt.assert_allclose(padded_da.y, np.linspace(-4.5, 4.5, 19)) + + +def test_pad_with_pad_width(sample_da_2d): + """ + Test pad function by passing pad_width as argument + """ + pad_width = {"x": (2, 3), "y": (1, 3)} + padded_da = pad(sample_da_2d, pad_width) + assert padded_da.shape == (21, 16) + npt.assert_allclose(padded_da.values[:1, :], 0) + npt.assert_allclose(padded_da.values[-3:, :], 0) + npt.assert_allclose(padded_da.values[:, :2], 0) + npt.assert_allclose(padded_da.values[:, -3:], 0) + npt.assert_allclose(padded_da.values[1:-3, 2:-3], sample_da_2d) + npt.assert_allclose(padded_da.x, np.linspace(-2, 13, 16)) + npt.assert_allclose(padded_da.y, np.linspace(-4.5, 5.5, 21)) + + +@pytest.mark.parametrize( + "pad_width", + ( + {"x": 2, "y": 3}, + {"x": 2}, + {"y": 3}, + {"x": (2, 3), "y": 3}, + {"x": (2, 3), "y": (1, 3)}, + {"x": (2, 3)}, + {"y": (1, 3)}, + ), +) +def test_coordinates_attrs_after_pad(sample_da_2d, pad_width): + """ + Test if the attributes of the coordinates are preserved after padding + and if the pad_width has been added + """ + padded_da = pad(sample_da_2d, pad_width) + # Check if the attrs in sample_da_2d is a subset of the attrs in padded_da + assert sample_da_2d.x.attrs.items() <= padded_da.x.attrs.items() + assert sample_da_2d.y.attrs.items() <= padded_da.y.attrs.items() + # Check if pad_width has been added to the attrs of each coordinate + for coord, width in pad_width.items(): + assert padded_da.coords[coord].attrs["pad_width"] == width + + +@pytest.mark.parametrize( + "pad_width", + ( + {"x": 2, "y": 3}, + {"x": 2}, + {"y": 3}, + {"x": (2, 3), "y": 3}, + {"x": (2, 3), "y": (1, 3)}, + {"x": (2, 3)}, + {"y": (1, 3)}, + ), +) +def test_pad_unpad_round_trip(sample_da_2d, pad_width): + """ + Test if applying pad and then unpad returns the original array + """ + unpadded = unpad(pad(sample_da_2d, pad_width)) + xrt.assert_allclose(sample_da_2d, unpadded) + + +def test_unpad_invalid_array(sample_da_2d): + """ + Test if error is raised when a not padded array is passed to unpad + """ + with pytest.raises(ValueError): + unpad(sample_da_2d) + + +@pytest.mark.parametrize( + "pad_width, size, expected_slice", + ( + [(1, 1), 4, slice(1, 3)], + [(1, 2), 5, slice(1, 3)], + [(2, 3), 10, slice(2, 7)], + [2, 10, slice(2, 8)], + ), +) +def test_pad_width_to_slice(pad_width, size, expected_slice): + """ + Test if _pad_width_to_slice work as expected + """ + assert _pad_width_to_slice(pad_width, size) == expected_slice + + +@pytest.mark.parametrize( + "kwargs", + ({"x": 1, "y": 1}, {"pad_width": {"x": 1, "y": 1}}), + ids=["pad_width_as_kwargs", "pad_width_as_argument"], +) +def test_unpad_custom_path_width(sample_da_2d, kwargs): + """ + Test the behaviour of unpad when passing a custom pad_width + """ + unpadded = unpad(sample_da_2d, **kwargs) + unpadded.shape == (15, 9) + npt.assert_allclose(unpadded.x, np.linspace(1, 9, 9)) + npt.assert_allclose(unpadded.y, np.linspace(-3.5, 3.5, 15)) + + +@pytest.mark.parametrize( + "pad_width_arg", + (None, "argument", "kwargs"), + ids=["pad_width_none", "pad_width_as_arg", "pad_width_as_kwargs"], +) +def test_unpad_pop_pad_width_attributes(sample_da_2d, pad_width_arg): + """ + Check if the unpadded array has no pad_width attributes + """ + pad_width = {"x": 2, "y": 1} + padded = pad(sample_da_2d, pad_width) + if pad_width_arg is None: + unpadded = unpad(padded) + elif pad_width_arg == "argument": + unpadded = unpad(padded, pad_width=pad_width) + else: + unpadded = unpad(padded, **pad_width) + # Check if unpadded doesn't have the pad_width attribtues + for dim in unpadded.coords: + assert "pad_width" not in unpadded.coords[dim].attrs + + +@pytest.mark.parametrize( + "pad_width", + ( + {"x": 4, "y": 3}, + {"x": 4}, + {"y": 3}, + {"x": (4, 3), "y": 3}, + {"x": (4, 3), "y": (5, 3)}, + {"x": (4, 3)}, + {"y": (5, 3)}, + ), +) +def test_unpad_ifft_fft_pad_round_trip(sample_da_2d, pad_width): + """ + Test if the round trip with padding and unpadding works + + This test passes a custom ``pad_width`` to the ``unpad`` function because + the ``fft`` doesn't support keeping the ``pad_width`` attribute on the + coordinates (at least for now). + """ + da_padded = pad(sample_da_2d, pad_width, constant_values=0) + da_fft = fft(da_padded, true_phase=True) + da_ifft = ifft(da_fft, true_phase=True) + da_unpadded = unpad(da_ifft, pad_width=pad_width) + xrt.assert_allclose(sample_da_2d, da_unpadded) diff --git a/xrft/tests/test_utils.py b/xrft/tests/test_utils.py new file mode 100644 index 00000000..57678767 --- /dev/null +++ b/xrft/tests/test_utils.py @@ -0,0 +1,55 @@ +""" +Test utility functions of xrft +""" +import pytest +import numpy as np +import pandas as pd +import xarray as xr +import numpy.testing as npt + +from ..utils import get_spacing + + +@pytest.fixture +def sample_da_2d(): + """ + Defines a 2D sample xarray.DataArray + """ + x = np.linspace(0, 10, 11) + y = np.linspace(-4, 4, 17) + z = np.arange(11 * 17, dtype=float).reshape(17, 11) + return xr.DataArray(z, coords={"x": x, "y": y}, dims=("y", "x")) + + +@pytest.fixture +def sample_da_time(): + """ + Defines a 1D sample xarray.DataArray with datetime coordinates + """ + time = pd.date_range( + "2021-01-01 00:00", "2021-01-01 23:00", periods=24 + ).to_pydatetime() + values = np.arange(24, dtype=float) + return xr.DataArray(values, coords={"time": time}, dims=("time",)) + + +def test_get_spacing(sample_da_2d, sample_da_time): + """ + Check if get_spacing function works as expected + """ + npt.assert_allclose(get_spacing(sample_da_2d.x), 1) + npt.assert_allclose(get_spacing(sample_da_2d.y), 0.5) + npt.assert_allclose(get_spacing(sample_da_time.time), 60 * 60) + + +def test_get_spacing_unvenly_spaced(sample_da_2d): + """ + Check if error is raised after unvenly spaced coordinates + """ + # Make the x coordinate unevenly spaced + x = sample_da_2d.x.values + x[0] += 0.2 + da = sample_da_2d.assign_coords({"x": x}) + # Check if error is raised + with pytest.raises(ValueError): + get_spacing(da.x) diff --git a/xrft/utils.py b/xrft/utils.py new file mode 100644 index 00000000..c53b3103 --- /dev/null +++ b/xrft/utils.py @@ -0,0 +1,19 @@ +""" +Utility functions for xrft +""" +import numpy as np + +from .xrft import _diff_coord + + +def get_spacing(coord): + """ + Return the spacing of evenly spaced coordinates array + """ + diff = _diff_coord(coord) + if not np.allclose(diff, diff[0]): + raise ValueError( + f"Found unevenly spaced coordinates '{coord.name}'. " + "These coordinates should be evenly spaced." + ) + return diff[0]