diff --git a/doc/whats-new.rst b/doc/whats-new.rst index cd01f0adaf1..1ef6f86f20a 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -23,6 +23,8 @@ v2024.03.0 (unreleased) New Features ~~~~~~~~~~~~ +- Allow control over padding in rolling. (:issue:`2007`, :pr:`5603`). + By `Kevin Squire `_. - Do not broadcast in arithmetic operations when global option ``arithmetic_broadcast=False`` (:issue:`6806`, :pull:`8784`). By `Etienne Schalk `_ and `Deepak Cherian `_. diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 7a0bdbc4d4c..bc07de8c908 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -6852,6 +6852,7 @@ def rolling( dim: Mapping[Any, int] | None = None, min_periods: int | None = None, center: bool | Mapping[Any, bool] = False, + pad: bool = True, **window_kwargs: int, ) -> DataArrayRolling: """ @@ -6919,7 +6920,9 @@ def rolling( from xarray.core.rolling import DataArrayRolling dim = either_dict_or_kwargs(dim, window_kwargs, "rolling") - return DataArrayRolling(self, dim, min_periods=min_periods, center=center) + return DataArrayRolling( + self, dim, min_periods=min_periods, center=center, pad=pad + ) def cumulative( self, diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 10bf1466156..54386e42206 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -10329,6 +10329,7 @@ def rolling( dim: Mapping[Any, int] | None = None, min_periods: int | None = None, center: bool | Mapping[Any, bool] = False, + pad: bool = True, **window_kwargs: int, ) -> DatasetRolling: """ @@ -10362,7 +10363,9 @@ def rolling( from xarray.core.rolling import DatasetRolling dim = either_dict_or_kwargs(dim, window_kwargs, "rolling") - return DatasetRolling(self, dim, min_periods=min_periods, center=center) + return DatasetRolling( + self, dim, min_periods=min_periods, center=center, pad=pad + ) def cumulative( self, diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 6cf49fc995b..7bc492d4296 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -4,7 +4,7 @@ import itertools import math import warnings -from collections.abc import Hashable, Iterator, Mapping +from collections.abc import Hashable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar import numpy as np @@ -53,6 +53,45 @@ """ +def get_pads( + dim: Sequence[Hashable], + window: Sequence[int], + center: Sequence[bool], + pad: Sequence[bool], +) -> dict[Hashable, tuple[int, int]]: + """The amount of padding to use at the each end of each dimension + + Parameters + ---------- + dim : sequence of str (or hashable) + dimension(s) for pads + window : sequence to int + Size of the window along a given dimension + center : sequence of bool + Whether or not to center the window on a particular dimension + pad : sequence of bool + Whether or not to pad a particular dimension + + Returns + ------- + pads: dict[Hashable, tuple[int, int]] + """ + pads = {} + for d, win, cent, p in zip(dim, window, center, pad): + if not p: + pads[d] = (0, 0) + continue + + if cent: + start = win // 2 # 10 -> 5, 9 -> 4 + end = (win - 1) // 2 # 10 -> 4, 9 -> 4 + pads[d] = (start, end) + else: + pads[d] = (win - 1, 0) + + return pads + + class Rolling(Generic[T_Xarray]): """A object that implements the moving window pattern. @@ -64,13 +103,14 @@ class Rolling(Generic[T_Xarray]): xarray.DataArray.rolling """ - __slots__ = ("obj", "window", "min_periods", "center", "dim") - _attributes = ("window", "min_periods", "center", "dim") + __slots__ = ("obj", "window", "min_periods", "center", "pad", "dim") + _attributes = ("window", "min_periods", "center", "pad", "dim") dim: list[Hashable] window: list[int] center: list[bool] obj: T_Xarray min_periods: int + pad: bool def __init__( self, @@ -78,6 +118,7 @@ def __init__( windows: Mapping[Any, int], min_periods: int | None = None, center: bool | Mapping[Any, bool] = False, + pad: bool = True, ) -> None: """ Moving window object. @@ -96,6 +137,9 @@ def __init__( center : bool or dict-like Hashable to bool, default: False Set the labels at the center of the window. If dict-like, set this property per rolling dimension. + pad : bool or mapping of hashable to bool, default: True + Pad the sides of the rolling window with ``NaN``. For different + padding, see ``DataArray.pad`` or ``Dataset.pad``. Returns ------- @@ -110,7 +154,8 @@ def __init__( self.window.append(w) self.center = self._mapping_to_list(center, default=False) - self.obj = obj + self.pad = self._mapping_to_list(pad, default=True) + self.obj: T_Xarray = obj missing_dims = tuple(dim for dim in self.dim if dim not in self.obj.dims) if missing_dims: @@ -132,8 +177,10 @@ def __repr__(self) -> str: """provide a nice str repr of our rolling object""" attrs = [ - "{k}->{v}{c}".format(k=k, v=w, c="(center)" if c else "") - for k, w, c in zip(self.dim, self.window, self.center) + "{k}->{v}{c}{p}".format( + k=k, v=w, c="(center)" if c else "", p="(no pad)" if not p else "" + ) + for k, w, c, p in zip(self.dim, self.window, self.center, self.pad) ] return "{klass} [{attrs}]".format( klass=self.__class__.__name__, attrs=",".join(attrs) @@ -246,6 +293,54 @@ def _get_keep_attrs(self, keep_attrs): return keep_attrs + def _get_output_dim_selector(self) -> dict[Hashable, Any]: + """Get output dimension selector, taking into account window size, window, centering, and padding. + + If any of the dimensions are not padded, the output size can be shorter than the input size + along that dimension. + """ + + def offsets_to_slice(start_offset: int, end_offset: int) -> slice: + # Turn start and end offsets into a slice object + slice_start = None if not start_offset else start_offset + slice_end = None if not end_offset else -end_offset + + return slice(slice_start, slice_end) + + # Dimensions which require offsets are those which are not padded, but the logic to determine + # the offset is very similar to determining padding sizes. + # So, we invert the `pad` flag(s), call `get_pads()`, and work from there. + + offset = [not p for p in self.pad] + offsets = get_pads(self.dim, self.window, self.center, offset) + + selector: dict[Hashable, slice] = { + dim: offsets_to_slice(start_offset, end_offset) + for dim, (start_offset, end_offset) in offsets.items() + } + + return selector + + def _get_output_coords(self) -> dict[str, Any]: + """Get output coordinates, taking into account window size, window, centering, and padding. + + If any of the dimensions are not padded, the output size can be shorter than the input size + along that dimension, so we need to shorten and properly label the corresponding coordinates. + + This useful for constructing a new DataArray or Dataset, where the data has already + been constructed to be the correct size along each dimension. + """ + # TODO: do we need to include dims without coordinates in output_coordinate_names? + # The code here does not include them. + selector = self._get_output_dim_selector() + + output_coords = { + k: self.obj.coords[k].isel(selector, missing_dims="ignore") + for k in self.obj.coords + } + + return output_coords + class DataArrayRolling(Rolling["DataArray"]): __slots__ = ("window_labels",) @@ -256,6 +351,7 @@ def __init__( windows: Mapping[Any, int], min_periods: int | None = None, center: bool | Mapping[Any, bool] = False, + pad: bool = True, ) -> None: """ Moving window object for DataArray. @@ -273,8 +369,11 @@ def __init__( Minimum number of observations in window required to have a value (otherwise result is NA). The default, None, is equivalent to setting min_periods equal to the size of the window. - center : bool, default: False + center : bool or mapping of hashable to bool, default: False Set the labels at the center of the window. + pad : bool or mapping of hashable to bool, default: True + Pad the sides of the rolling window with ``NaN``. For different + padding, see ``DataArray.pad``. Returns ------- @@ -287,7 +386,13 @@ def __init__( xarray.Dataset.rolling xarray.Dataset.groupby """ - super().__init__(obj, windows, min_periods=min_periods, center=center) + super().__init__( + obj, + windows, + min_periods=min_periods, + center=center, + pad=pad, + ) # TODO legacy attribute self.window_labels = self.obj[self.dim[0]] @@ -296,17 +401,48 @@ def __iter__(self) -> Iterator[tuple[DataArray, DataArray]]: if self.ndim > 1: raise ValueError("__iter__ is only supported for 1d-rolling") - dim0 = self.dim[0] - window0 = int(self.window[0]) - offset = (window0 + 1) // 2 if self.center[0] else 1 - stops = np.arange(offset, self.obj.sizes[dim0] + offset) - starts = stops - window0 - starts[: window0 - offset] = 0 + dim = self.dim[0] + center = self.center[0] + pad = self.pad[0] + window = self.window[0] + center_offset = window // 2 if center else 0 + + pads = get_pads(self.dim, self.window, self.center, self.pad) + start_pad, end_pad = pads[dim] + + # Select the proper subset of labels, based on whether or not to center and/or pad + first_label_idx = 0 if pad else center_offset if center else window - 1 + last_label_idx = ( + len(self.obj[dim]) + if pad or not center + else len(self.obj[dim]) - center_offset + ) - for label, start, stop in zip(self.window_labels, starts, stops): - window = self.obj.isel({dim0: slice(start, stop)}) + labels = ( + self.obj[dim][slice(first_label_idx, last_label_idx)] + if self.obj[dim].coords + else np.arange(last_label_idx - first_label_idx) + ) - counts = window.count(dim=[dim0]) + padded_obj = self.obj.pad(pads, mode="constant", constant_values=dtypes.NA) + + if pad and not center: + first_stop = 1 + last_stop = len(self.obj[dim]) + elif pad and center: + first_stop = end_pad + 1 + last_stop = len(self.obj[dim]) + end_pad + elif not pad: + first_stop = window + last_stop = len(self.obj[dim]) + + # These are indicies into the padded array, so we need to add start_pad + stops = np.arange(first_stop, last_stop + 1) + start_pad + starts = stops - window + + for label, start, stop in zip(labels, starts, stops): + window = padded_obj.isel({dim: slice(start, stop)}) + counts = window.count(dim=[dim]) window = window.where(counts >= self.min_periods) yield (label, window) @@ -412,15 +548,21 @@ def _construct( strides = self._mapping_to_list(stride, default=1) window = obj.variable.rolling_window( - self.dim, self.window, window_dims, self.center, fill_value=fill_value + self.dim, + self.window, + window_dims, + self.center, + self.pad, + fill_value=fill_value, ) attrs = obj.attrs if keep_attrs else {} + coords = self._get_output_coords() result = DataArray( window, dims=obj.dims + tuple(window_dims), - coords=obj.coords, + coords=coords, attrs=attrs, name=obj.name, ) @@ -522,6 +664,7 @@ def _counts(self, keep_attrs: bool | None) -> DataArray: .rolling( {d: w for d, w in zip(self.dim, self.window)}, center={d: self.center[i] for i, d in enumerate(self.dim)}, + pad={d: self.pad[i] for i, d in enumerate(self.dim)}, ) .construct(rolling_dim, fill_value=False, keep_attrs=keep_attrs) .sum(dim=dim, skipna=False, keep_attrs=keep_attrs) @@ -610,10 +753,11 @@ def _bottleneck_reduce(self, func, keep_attrs, **kwargs): values = values[valid] attrs = self.obj.attrs if keep_attrs else {} + output_selector = self._get_output_dim_selector() - return self.obj.__class__( + return type(self.obj)( values, self.obj.coords, attrs=attrs, name=self.obj.name - ) + ).isel(output_selector) def _array_reduce( self, @@ -700,6 +844,7 @@ def __init__( windows: Mapping[Any, int], min_periods: int | None = None, center: bool | Mapping[Any, bool] = False, + pad: bool = True, ) -> None: """ Moving window object for Dataset. @@ -719,6 +864,9 @@ def __init__( setting min_periods equal to the size of the window. center : bool or mapping of hashable to bool, default: False Set the labels at the center of the window. + pad : bool or mapping of hashable to bool, default: True + Pad the sides of the window with ``NaN``. For different + padding, see ``Dataset.pad``. Returns ------- @@ -731,21 +879,23 @@ def __init__( xarray.Dataset.groupby xarray.DataArray.groupby """ - super().__init__(obj, windows, min_periods, center) - + super().__init__(obj, windows, min_periods, center, pad) + if any(d not in self.obj.dims for d in self.dim): + raise KeyError(self.dim) # Keep each Rolling object as a dictionary self.rollings = {} for key, da in self.obj.data_vars.items(): # keeps rollings only for the dataset depending on self.dim - dims, center = [], {} + dims, center, pad = [], {}, {} for i, d in enumerate(self.dim): if d in da.dims: dims.append(d) center[d] = self.center[i] + pad[d] = self.pad[i] if dims: w = {d: windows[d] for d in dims} - self.rollings[key] = DataArrayRolling(da, w, min_periods, center) + self.rollings[key] = DataArrayRolling(da, w, min_periods, center, pad) def _dataset_implementation(self, func, keep_attrs, **kwargs): from xarray.core.dataset import Dataset @@ -763,7 +913,9 @@ def _dataset_implementation(self, func, keep_attrs, **kwargs): reduced[key].attrs = {} attrs = self.obj.attrs if keep_attrs else {} - return Dataset(reduced, coords=self.obj.coords, attrs=attrs) + coords = self._get_output_coords() + + return Dataset(reduced, coords=coords, attrs=attrs) def reduce( self, func: Callable, keep_attrs: bool | None = None, **kwargs: Any diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 5cb52cbd25c..1d109d304b3 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -57,6 +57,7 @@ Mapping, MutableMapping, MutableSet, + Sequence, ValuesView, ) from enum import Enum @@ -68,6 +69,7 @@ Generic, Literal, TypeVar, + Union, overload, ) @@ -1038,6 +1040,65 @@ def iterate_nested(nested_list): yield item +ItemOrSequence = Union[T, Sequence[T]] + + +def expand_args_to_dims( + dim: ItemOrSequence[Hashable], + arg_names: Sequence[str], + args: Sequence[ItemOrSequence[Any]], +) -> tuple[Sequence[Hashable], Sequence[Sequence[Any]]]: + """Expand dims and all elements in args to be arrays of the length of the number of dimensions + + Parameters + ---------- + dim : str or sequence of str + Dimension(s) + arg_names : sequence of str + Names of the arguments to expand + args : sequence of args, which may be individual items or lists of items + Arguments to expand + + Raises + ------ + ValueError: raised if dim is a scalar and any of the args are not scalars + ValueError: raised if any of the produced lists are of the wrong length + + Returns + ------- + list of dims, list of lists of arguments + """ + if is_scalar(dim): + dim_list: Sequence[Hashable] = [dim] + else: + assert isinstance(dim, Sequence) + dim_list = dim + + # dim is now a list + nroll = len(dim_list) + + def to_list(arg): + if is_scalar(arg): + return [arg] * nroll + elif isinstance(arg, list) and len(arg) == 1 and nroll > 1: + return [arg[0]] * nroll + return arg + + arr_args = [to_list(arg) for arg in args] + + if any(len(arg) != len(dim_list) for arg in arr_args): + names_args_len = ", ".join( + f"{name}={args!r} (len={len(args)})" + for name, args in zip(arg_names, arr_args) + if len(args) != len(dim_list) + ) + raise ValueError( + f"Expected all arguments to have len={len(dim_list)}. Received: {names_args_len}" + ) + + return dim_list, arr_args + + def contains_only_chunked_or_numpy(obj) -> bool: """Returns True if xarray object contains only numpy arrays or chunked arrays (i.e. pure dask or cubed). diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 2ac0c04d726..a5cf594c386 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -26,6 +26,7 @@ as_indexable, ) from xarray.core.options import OPTIONS, _get_keep_attrs +from xarray.core.rolling import get_pads from xarray.core.utils import ( OrderedSet, _default, @@ -34,14 +35,19 @@ drop_dims_from_indexers, either_dict_or_kwargs, ensure_us_time_resolution, + expand_args_to_dims, infix_dims, is_dict_like, is_duck_array, - is_duck_dask_array, maybe_coerce_to_str, ) from xarray.namedarray.core import NamedArray, _raise_if_any_duplicate_dimensions -from xarray.namedarray.pycompat import integer_types, is_0d_dask_array, to_duck_array +from xarray.namedarray.pycompat import ( + integer_types, + is_0d_dask_array, + is_duck_dask_array, + to_duck_array, +) NON_NUMPY_SUPPORTED_ARRAY_TYPES = ( indexing.ExplicitlyIndexed, @@ -1961,7 +1967,7 @@ def rank(self, dim, pct=False): return ranked def rolling_window( - self, dim, window, window_dim, center=False, fill_value=dtypes.NA + self, dim, window, window_dim, center=False, pad=True, fill_value=dtypes.NA ): """ Make a rolling_window along dim and add a new_dim to the last place. @@ -1973,14 +1979,18 @@ def rolling_window( For nd-rolling, should be list of dimensions. window : int Window size of the rolling - For nd-rolling, should be list of integers. + For nd-rolling, should be a list of integers. window_dim : str New name of the window dimension. - For nd-rolling, should be list of strings. + For nd-rolling, should be a list of strings. center : bool, default: False If True, pad fill_value for both ends. Otherwise, pad in the head of the axis. - fill_value + For nd-rolling, can be a list of bools + pad : bool, default: True + Pad the sides of the rolling_window with fill_value. + For nd-rolling, can be a list of bools + fill_value, default=NA value to be filled. Returns @@ -2016,6 +2026,13 @@ def rolling_window( [ 4., 5., 6.], [ 5., 6., 7.], [ 6., 7., nan]]]) + >>> v.rolling_window("b", 3, "window_dim", center=True, pad=False) + + array([[[0., 1., 2.], + [1., 2., 3.]], + + [[4., 5., 6.], + [5., 6., 7.]]]) """ if fill_value is dtypes.NA: # np.nan is passed dtype, fill_value = dtypes.maybe_promote(self.dtype) @@ -2024,46 +2041,19 @@ def rolling_window( dtype = self.dtype var = self - if utils.is_scalar(dim): - for name, arg in zip( - ["window", "window_dim", "center"], [window, window_dim, center] - ): - if not utils.is_scalar(arg): - raise ValueError( - f"Expected {name}={arg!r} to be a scalar like 'dim'." - ) - dim = (dim,) - - # dim is now a list - nroll = len(dim) - if utils.is_scalar(window): - window = [window] * nroll - if utils.is_scalar(window_dim): - window_dim = [window_dim] * nroll - if utils.is_scalar(center): - center = [center] * nroll - if ( - len(dim) != len(window) - or len(dim) != len(window_dim) - or len(dim) != len(center) - ): - raise ValueError( - "'dim', 'window', 'window_dim', and 'center' must be the same length. " - f"Received dim={dim!r}, window={window!r}, window_dim={window_dim!r}," - f" and center={center!r}." - ) + dim, (window, window_dim, center, pad) = expand_args_to_dims( + dim, + ["window", "window_dim", "center", "pad"], + [window, window_dim, center, pad], + ) - pads = {} - for d, win, cent in zip(dim, window, center): - if cent: - start = win // 2 # 10 -> 5, 9 -> 4 - end = win - 1 - start - pads[d] = (start, end) - else: - pads[d] = (win - 1, 0) + if all(not p for p in pad): + padded = var + else: + pads = get_pads(dim, window, center, pad) + padded = var.pad(pads, mode="constant", constant_values=fill_value) - padded = var.pad(pads, mode="constant", constant_values=fill_value) - axis = self.get_axis_num(dim) + axis = [self.get_axis_num(d) for d in dim] new_dims = self.dims + tuple(window_dim) return Variable( new_dims, diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 5007db9eeb2..e99f9ec3a22 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -340,6 +340,24 @@ def create_test_data( return obj +def get_expected_rolling_indices(length, window, center, pad, stride=1): + # Without padding, we lose (window-1) items from the index, either from the beginning + # (without centering) or from the beginning and end (with centering) + if pad: + start_index = 0 + end_index = length + elif center: + start_index = window // 2 # 10 -> 5, 9 -> 4 + end_index = length - (window - 1) // 2 # 10 -> 4, 9 -> 4 + else: + start_index = window - 1 + end_index = length + + expected_index = np.arange(start_index, end_index, stride) + + return expected_index + + _CFTIME_CALENDARS = [ "365_day", "360_day", diff --git a/xarray/tests/test_rolling.py b/xarray/tests/test_rolling.py index 79a5ba0a667..403a72f9028 100644 --- a/xarray/tests/test_rolling.py +++ b/xarray/tests/test_rolling.py @@ -1,10 +1,9 @@ from __future__ import annotations -from typing import Any - import numpy as np import pandas as pd import pytest +from numpy.testing import assert_array_equal import xarray as xr from xarray import DataArray, Dataset, set_options @@ -12,7 +11,9 @@ assert_allclose, assert_equal, assert_identical, + get_expected_rolling_indices, has_dask, + requires_bottleneck, requires_dask, requires_numbagg, ) @@ -51,9 +52,18 @@ def test_cumulative_vs_cum(d) -> None: class TestDataArrayRolling: @pytest.mark.parametrize("da", (1, 2), indirect=True) @pytest.mark.parametrize("center", [True, False]) - @pytest.mark.parametrize("size", [1, 2, 3, 7]) - def test_rolling_iter(self, da: DataArray, center: bool, size: int) -> None: - rolling_obj = da.rolling(time=size, center=center) + @pytest.mark.parametrize("pad", (True, False)) + @pytest.mark.parametrize("min_periods", (1, 6, None)) + @pytest.mark.parametrize("window", [6, 7]) + def test_rolling_iter( + self, + da: DataArray, + center: bool, + pad: bool, + min_periods: int | None, + window: int, + ) -> None: + rolling_obj = da.rolling(time=window, center=center) rolling_obj_mean = rolling_obj.mean() assert len(rolling_obj.window_labels) == len(da["time"]) @@ -65,7 +75,22 @@ def test_rolling_iter(self, da: DataArray, center: bool, size: int) -> None: actual = rolling_obj_mean.isel(time=i) expected = window_da.mean("time") - np.testing.assert_allclose(actual.values, expected.values) + # TODO add assert_allclose_with_nan, which compares nan position + # as well as the closeness of the values. + assert_array_equal(actual.isnull(), expected.isnull()) + if (~actual.isnull()).sum() > 0: + if actual.ndim == 0: + actual_values = actual.values + expected_values = expected.values + else: + actual_values = actual.values[actual.values.nonzero()] + expected_values = expected.values[expected.values.nonzero()] + + assert np.allclose(actual_values, expected_values) + + def test_rolling_nonunique_coords(self) -> None: + da = xr.DataArray([1, 2, 3], coords=dict(x=[1, 1, 1])).rolling(x=3).mean() + assert_array_equal(da.data, [np.nan, np.nan, 2.0]) @pytest.mark.parametrize("da", (1,), indirect=True) def test_rolling_repr(self, da) -> None: @@ -76,6 +101,24 @@ def test_rolling_repr(self, da) -> None: rolling_obj = da.rolling(time=7, x=3, center=True) assert repr(rolling_obj) == "DataArrayRolling [time->7(center),x->3(center)]" + rolling_obj = da.rolling(time=7, pad=False) + assert repr(rolling_obj) == "DataArrayRolling [time->7(no pad)]" + rolling_obj = da.rolling(time=7, center=True, pad=False) + assert repr(rolling_obj) == "DataArrayRolling [time->7(center)(no pad)]" + rolling_obj = da.rolling(time=7, x=3, center=True, pad=False) + assert ( + repr(rolling_obj) + == "DataArrayRolling [time->7(center)(no pad),x->3(center)(no pad)]" + ) + + rolling_obj = da.rolling( + time=7, + x=3, + center={"time": True, "x": False}, + pad={"time": True, "x": False}, + ) + assert repr(rolling_obj) == "DataArrayRolling [time->7(center),x->3(no pad)]" + @requires_dask def test_repeated_rolling_rechunks(self) -> None: # regression test for GH3277, GH2514 @@ -107,30 +150,42 @@ def test_rolling_properties(self, da) -> None: ): da.rolling(foo=2) + @requires_bottleneck @pytest.mark.parametrize( "name", ("sum", "mean", "std", "min", "max", "median", "argmin", "argmax") ) @pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) + @pytest.mark.parametrize("pad", (True, False)) def test_rolling_wrapped_bottleneck( - self, da, name, center, min_periods, compute_backend + self, da, name, min_periods, center, pad, compute_backend ) -> None: - bn = pytest.importorskip("bottleneck", minversion="1.1") + import bottleneck as bn + + window = 7 + # Test all bottleneck functions - rolling_obj = da.rolling(time=7, min_periods=min_periods) + rolling_obj = da.rolling( + time=window, min_periods=min_periods, center=center, pad=pad + ) + window = 7 func_name = f"move_{name}" actual = getattr(rolling_obj, name)() - window = 7 - expected = getattr(bn, func_name)( + expected_values = getattr(bn, func_name)( da.values, window=window, axis=1, min_count=min_periods ) + expected_indices = get_expected_rolling_indices( + da.sizes["time"], window, center, pad + ) # index 0 is at the rightmost edge of the window # need to reverse index here # see GH #8541 if func_name in ["move_argmin", "move_argmax"]: - expected = window - 1 - expected + expected_values = window - 1 - expected_values + expected = da.copy(data=expected_values).isel(time=expected_indices) + assert_equal(actual, expected) # Using assert_allclose because we get tiny (1e-17) differences in numbagg. np.testing.assert_allclose(actual.values, expected) @@ -147,19 +202,25 @@ def test_rolling_wrapped_bottleneck( @requires_dask @pytest.mark.parametrize("name", ("mean", "count")) @pytest.mark.parametrize("center", (True, False, None)) + @pytest.mark.parametrize("pad", (True, False)) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("window", (7, 8)) @pytest.mark.parametrize("backend", ["dask"], indirect=True) - def test_rolling_wrapped_dask(self, da, name, center, min_periods, window) -> None: + def test_rolling_wrapped_dask( + self, da, name, center, pad, min_periods, window + ) -> None: # dask version - rolling_obj = da.rolling(time=window, min_periods=min_periods, center=center) + assert da.chunks is not None + rolling_obj = da.rolling( + time=window, min_periods=min_periods, center=center, pad=pad + ) actual = getattr(rolling_obj, name)().load() if name != "count": with pytest.warns(DeprecationWarning, match="Reductions are applied"): getattr(rolling_obj, name)(dim="time") # numpy version rolling_obj = da.load().rolling( - time=window, min_periods=min_periods, center=center + time=window, min_periods=min_periods, center=center, pad=pad ) expected = getattr(rolling_obj, name)() @@ -169,28 +230,35 @@ def test_rolling_wrapped_dask(self, da, name, center, min_periods, window) -> No # with zero chunked array GH:2113 rolling_obj = da.chunk().rolling( - time=window, min_periods=min_periods, center=center + time=window, + min_periods=min_periods, + center=center, + pad=pad, ) actual = getattr(rolling_obj, name)().load() assert_allclose(actual, expected) + @requires_dask @pytest.mark.parametrize("center", (True, None)) - def test_rolling_wrapped_dask_nochunk(self, center) -> None: + @pytest.mark.parametrize("pad", (True, False)) + def test_rolling_wrapped_dask_nochunk(self, center, pad) -> None: # GH:2113 - pytest.importorskip("dask.array") da_day_clim = xr.DataArray( np.arange(1, 367), coords=[np.arange(1, 367)], dims="dayofyear" ) - expected = da_day_clim.rolling(dayofyear=31, center=center).mean() - actual = da_day_clim.chunk().rolling(dayofyear=31, center=center).mean() + expected = da_day_clim.rolling(dayofyear=31, center=center, pad=pad).mean() + actual = ( + da_day_clim.chunk().rolling(dayofyear=31, center=center, pad=pad).mean() + ) assert_allclose(actual, expected) @pytest.mark.parametrize("center", (True, False)) + @pytest.mark.parametrize("pad", (False,)) @pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) @pytest.mark.parametrize("window", (1, 2, 3, 4)) def test_rolling_pandas_compat( - self, center, window, min_periods, compute_backend + self, center, pad, window, min_periods, compute_backend ) -> None: s = pd.Series(np.arange(10)) da = DataArray.from_series(s) @@ -198,12 +266,18 @@ def test_rolling_pandas_compat( if min_periods is not None and window < min_periods: min_periods = window - s_rolling = s.rolling(window, center=center, min_periods=min_periods).mean() + expected_index = get_expected_rolling_indices(10, window, center, pad) + + s_rolling = ( + s.rolling(window, center=center, min_periods=min_periods) + .mean() + .iloc[expected_index] + ) da_rolling = da.rolling( - index=window, center=center, min_periods=min_periods + index=window, center=center, pad=pad, min_periods=min_periods ).mean() da_rolling_np = da.rolling( - index=window, center=center, min_periods=min_periods + index=window, center=center, pad=pad, min_periods=min_periods ).reduce(np.nanmean) np.testing.assert_allclose(s_rolling.values, da_rolling.values) @@ -214,8 +288,10 @@ def test_rolling_pandas_compat( @pytest.mark.parametrize("center", (True, False)) @pytest.mark.parametrize("window", (1, 2, 3, 4)) def test_rolling_construct(self, center: bool, window: int) -> None: - s = pd.Series(np.arange(10)) + length = 10 + s = pd.Series(np.arange(length)) da = DataArray.from_series(s) + da = da.assign_coords(time=("index", np.arange(1, length + 1))) s_rolling = s.rolling(window, center=center, min_periods=1).mean() da_rolling = da.rolling(index=window, center=center, min_periods=1) @@ -236,13 +312,33 @@ def test_rolling_construct(self, center: bool, window: int) -> None: assert da_rolling_mean.isnull().sum() == 0 assert (da_rolling_mean == 0.0).sum() >= 0 + # with no padding + da_rolling = da.rolling(index=window, center=center, min_periods=1, pad=False) + da_rolling_mean = da_rolling.construct("window", stride=2).mean("window") + + expected_index = get_expected_rolling_indices( + length, window, center, pad=False, stride=2 + ) + + assert da_rolling_mean.sizes["index"] == len(expected_index) + assert (da_rolling_mean.index.values == expected_index).all() + assert (da_rolling_mean.time.values == expected_index + 1).all() + + np.testing.assert_allclose( + s_rolling.values[expected_index], da_rolling_mean.values + ) + np.testing.assert_allclose( + s_rolling.index[expected_index], da_rolling_mean["index"] + ) + @pytest.mark.parametrize("da", (1, 2), indirect=True) @pytest.mark.parametrize("center", (True, False)) + @pytest.mark.parametrize("pad", (True, False)) @pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) @pytest.mark.parametrize("window", (1, 2, 3, 4)) @pytest.mark.parametrize("name", ("sum", "mean", "std", "max")) def test_rolling_reduce( - self, da, center, min_periods, window, name, compute_backend + self, da, center, min_periods, pad, window, name, compute_backend ) -> None: if min_periods is not None and window < min_periods: min_periods = window @@ -251,7 +347,9 @@ def test_rolling_reduce( # this causes all nan slices window = 2 - rolling_obj = da.rolling(time=window, center=center, min_periods=min_periods) + rolling_obj = da.rolling( + time=window, center=center, pad=pad, min_periods=min_periods + ) # add nan prefix to numpy methods to get similar # behavior as bottleneck actual = rolling_obj.reduce(getattr(np, "nan%s" % name)) @@ -260,11 +358,12 @@ def test_rolling_reduce( assert actual.sizes == expected.sizes @pytest.mark.parametrize("center", (True, False)) + @pytest.mark.parametrize("pad", (True, False)) @pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) @pytest.mark.parametrize("window", (1, 2, 3, 4)) @pytest.mark.parametrize("name", ("sum", "max")) def test_rolling_reduce_nonnumeric( - self, center, min_periods, window, name, compute_backend + self, center, pad, min_periods, window, name, compute_backend ) -> None: da = DataArray( [0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time" @@ -273,7 +372,9 @@ def test_rolling_reduce_nonnumeric( if min_periods is not None and window < min_periods: min_periods = window - rolling_obj = da.rolling(time=window, center=center, min_periods=min_periods) + rolling_obj = da.rolling( + time=window, center=center, pad=pad, min_periods=min_periods + ) # add nan prefix to numpy methods to get similar behavior as bottleneck actual = rolling_obj.reduce(getattr(np, "nan%s" % name)) @@ -281,56 +382,62 @@ def test_rolling_reduce_nonnumeric( assert_allclose(actual, expected) assert actual.sizes == expected.sizes - def test_rolling_count_correct(self, compute_backend) -> None: + @pytest.mark.parametrize( + "time, min_periods, pad, expected", + ( + [11, 1, True, DataArray([1, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8], dims="time")], + [11, 1, False, DataArray([8], dims="time")], + [ + 11, + None, + True, + DataArray( + [np.nan] * 11, + dims="time", + ), + ], + [11, None, False, DataArray([np.nan], dims="time")], + [ + 7, + 2, + True, + DataArray([np.nan, np.nan, 2, 3, 3, 4, 5, 5, 5, 5, 5], dims="time"), + ], + [7, 2, False, DataArray([5, 5, 5, 5, 5], dims="time")], + ), + ) + def test_rolling_count_correct( + self, time, min_periods, pad, expected, compute_backend + ) -> None: da = DataArray([0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time") + result = da.rolling(time=time, min_periods=min_periods, pad=pad).count() + assert_equal(result, expected) - kwargs: list[dict[str, Any]] = [ - {"time": 11, "min_periods": 1}, - {"time": 11, "min_periods": None}, - {"time": 7, "min_periods": 2}, - ] - expecteds = [ - DataArray([1, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8], dims="time"), - DataArray( - [ - np.nan, - np.nan, - np.nan, - np.nan, - np.nan, - np.nan, - np.nan, - np.nan, - np.nan, - np.nan, - np.nan, - ], - dims="time", - ), - DataArray([np.nan, np.nan, 2, 3, 3, 4, 5, 5, 5, 5, 5], dims="time"), - ] - - for kwarg, expected in zip(kwargs, expecteds): - result = da.rolling(**kwarg).count() - assert_equal(result, expected) - - result = da.to_dataset(name="var1").rolling(**kwarg).count()["var1"] - assert_equal(result, expected) + result = ( + da.to_dataset(name="var1") + .rolling(time=time, min_periods=min_periods, pad=pad) + .count()["var1"] + ) + assert_equal(result, expected) @pytest.mark.parametrize("da", (1,), indirect=True) - @pytest.mark.parametrize("center", (True, False)) + @pytest.mark.parametrize("center", (True, False, {"time": True, "x": False})) + @pytest.mark.parametrize("pad", (True, False, {"time": True, "x": False})) @pytest.mark.parametrize("min_periods", (None, 1)) @pytest.mark.parametrize("name", ("sum", "mean", "max")) def test_ndrolling_reduce( - self, da, center, min_periods, name, compute_backend + self, da, center, min_periods, pad, name, compute_backend ) -> None: - rolling_obj = da.rolling(time=3, x=2, center=center, min_periods=min_periods) + rolling_obj = da.rolling( + time=3, x=2, pad=pad, center=center, min_periods=min_periods + ) actual = getattr(rolling_obj, name)() expected = getattr( getattr( - da.rolling(time=3, center=center, min_periods=min_periods), name - )().rolling(x=2, center=center, min_periods=min_periods), + da.rolling(time=3, center=center, pad=pad, min_periods=min_periods), + name, + )().rolling(x=2, center=center, pad=pad, min_periods=min_periods), name, )() @@ -347,23 +454,33 @@ def test_ndrolling_reduce( min_periods = 1 assert_allclose(actual, expected.where(count >= min_periods)) - @pytest.mark.parametrize("center", (True, False, (True, False))) + @pytest.mark.parametrize("center", (True, False, {"x": True, "z": False})) + @pytest.mark.parametrize( + "pad", + ( + True, + False, + {"x": True, "z": False}, + ), + ) @pytest.mark.parametrize("fill_value", (np.nan, 0.0)) - def test_ndrolling_construct(self, center, fill_value) -> None: + def test_ndrolling_construct(self, center, pad, fill_value) -> None: da = DataArray( np.arange(5 * 6 * 7).reshape(5, 6, 7).astype(float), dims=["x", "y", "z"], coords={"x": ["a", "b", "c", "d", "e"], "y": np.arange(6)}, ) - actual = da.rolling(x=3, z=2, center=center).construct( + actual = da.rolling(x=3, z=2, center=center, pad=pad).construct( x="x1", z="z1", fill_value=fill_value ) - if not isinstance(center, tuple): - center = (center, center) + if not isinstance(center, dict): + center = {"x": center, "z": center} + if not isinstance(pad, dict): + pad = {"x": pad, "z": pad} expected = ( - da.rolling(x=3, center=center[0]) + da.rolling(x=3, center=center["x"], pad=pad["x"]) .construct(x="x1", fill_value=fill_value) - .rolling(z=2, center=center[1]) + .rolling(z=2, center=center["z"], pad=pad["z"]) .construct(z="z1", fill_value=fill_value) ) assert_allclose(actual, expected) @@ -596,37 +713,46 @@ def test_rolling_properties(self, ds) -> None: ): ds.rolling(foo=2) + @requires_bottleneck @pytest.mark.parametrize( "name", ("sum", "mean", "std", "var", "min", "max", "median") ) - @pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("key", ("z1", "z2")) + @pytest.mark.parametrize("center", (True, False, None)) + @pytest.mark.parametrize("pad", (False,)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_wrapped_bottleneck( - self, ds, name, center, min_periods, key, compute_backend + self, ds, name, pad, center, min_periods, key, compute_backend ) -> None: - bn = pytest.importorskip("bottleneck", minversion="1.1") + import bottleneck as bn # Test all bottleneck functions - rolling_obj = ds.rolling(time=7, min_periods=min_periods) + window = 7 + rolling_obj = ds.rolling( + time=window, min_periods=min_periods, center=center, pad=pad + ) func_name = f"move_{name}" actual = getattr(rolling_obj, name)() + expected_indices = get_expected_rolling_indices( + ds.sizes["time"], window, center, pad + ) if key == "z1": # z1 does not depend on 'Time' axis. Stored as it is. expected = ds[key] elif key == "z2": - expected = getattr(bn, func_name)( - ds[key].values, window=7, axis=0, min_count=min_periods + expected = ( + ds[key] + .copy( + data=getattr(bn, func_name)( + ds[key].values, window=7, axis=0, min_count=min_periods + ) + ) + .isel(time=expected_indices) ) else: raise ValueError - np.testing.assert_allclose(actual[key].values, expected) - - # Test center - rolling_obj = ds.rolling(time=7, center=center) - actual = getattr(rolling_obj, name)()["time"] - assert_allclose(actual, ds["time"]) + assert_equal(actual[key], expected) @pytest.mark.parametrize("center", (True, False)) @pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) diff --git a/xarray/tests/test_utils.py b/xarray/tests/test_utils.py index 50061c774a8..82841024182 100644 --- a/xarray/tests/test_utils.py +++ b/xarray/tests/test_utils.py @@ -7,7 +7,12 @@ import pytest from xarray.core import duck_array_ops, utils -from xarray.core.utils import either_dict_or_kwargs, infix_dims, iterate_nested +from xarray.core.utils import ( + either_dict_or_kwargs, + expand_args_to_dims, + infix_dims, + iterate_nested, +) from xarray.tests import assert_array_equal, requires_dask @@ -348,6 +353,27 @@ def test_iterate_nested(nested_list, expected): assert list(iterate_nested(nested_list)) == expected +def test_expand_args_to_dims() -> None: + dims, (arg1, arg2, arg3, arg4) = expand_args_to_dims( + ["a", "b"], + ["arg1", "arg2", "arg3", "arg4"], + [1, ["val2.1", "val2.2"], False, [True, False]], + ) + + assert dims == ["a", "b"] + assert arg1 == [1, 1] + assert arg2 == ["val2.1", "val2.2"] + assert arg3 == [False, False] + assert arg4 == [True, False] + + with pytest.raises(ValueError, match="Expected all arguments"): + expand_args_to_dims( + ["a", "b"], + ["arg1", "arg2", "arg3", "arg4"], + ["asdf", ["arg2.1", "arg2.2"], ["arg3.1"], ["arg4.1", "arg4.2", "arg4.3"]], + ) + + def test_find_stack_level(): assert utils.find_stack_level() == 1 assert utils.find_stack_level(test_mode=True) == 2