From 039a9bf6d33275da3f904f9149fcef740d9925c3 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Fri, 9 Jul 2021 11:47:58 -0700 Subject: [PATCH 01/25] Allow padding to be turned off for rolling windows * Adds a `pad: bool` parameter to DataArray.rolling / Dataset.rolling --- xarray/core/common.py | 13 +++- xarray/core/rolling.py | 107 +++++++++++++++++++++----- xarray/core/utils.py | 136 +++++++++++++++++++++++++++++++++ xarray/core/variable.py | 63 ++++++--------- xarray/tests/test_dataarray.py | 58 +++++++++++--- 5 files changed, 310 insertions(+), 67 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index 7b6e9198b43..028362c9c6b 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -821,6 +821,7 @@ def rolling( dim: Mapping[Hashable, int] = None, min_periods: int = None, center: Union[bool, Mapping[Hashable, bool]] = False, + pad: Union[bool, Mapping[Hashable, bool]] = True, keep_attrs: bool = None, **window_kwargs: int, ): @@ -838,6 +839,9 @@ def rolling( setting min_periods equal to the size of the window. center : bool or mapping, default: False Set the labels at the center of the window. + pad : bool or mapping, default: True + Pad the sides of the window with ``NaN``. For different + padding, see ``DataArray.pad`` or ``Dataset.pad``. **window_kwargs : optional The keyword arguments form of ``dim``. One of dim or window_kwargs must be provided. @@ -886,11 +890,18 @@ def rolling( -------- core.rolling.DataArrayRolling core.rolling.DatasetRolling + DataArray.pad + Dataset.pad """ dim = either_dict_or_kwargs(dim, window_kwargs, "rolling") return self._rolling_cls( - self, dim, min_periods=min_periods, center=center, keep_attrs=keep_attrs + self, + dim, + min_periods=min_periods, + center=center, + pad=pad, + keep_attrs=keep_attrs, ) def rolling_exp( diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index b87dcda24b0..e71e40252c4 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -48,10 +48,12 @@ class Rolling: xarray.DataArray.rolling """ - __slots__ = ("obj", "window", "min_periods", "center", "dim", "keep_attrs") - _attributes = ("window", "min_periods", "center", "dim", "keep_attrs") + __slots__ = ("obj", "window", "min_periods", "center", "pad", "dim", "keep_attrs") + _attributes = ("window", "min_periods", "center", "pad", "dim", "keep_attrs") - def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None): + def __init__( + self, obj, windows, min_periods=None, center=False, pad=True, keep_attrs=None + ): """ Moving window object. @@ -66,8 +68,11 @@ def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None 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`` or ``Dataset.pad``. Returns ------- @@ -81,6 +86,7 @@ def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None self.window.append(w) self.center = self._mapping_to_list(center, default=False) + self.pad = self._mapping_to_list(pad, default=True) self.obj = obj # attributes @@ -102,8 +108,10 @@ def __repr__(self): """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) @@ -204,7 +212,9 @@ def _get_keep_attrs(self, keep_attrs): class DataArrayRolling(Rolling): __slots__ = ("window_labels",) - def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None): + def __init__( + self, obj, windows, min_periods=None, center=False, pad=True, keep_attrs=None + ): """ Moving window object for DataArray. You should use DataArray.rolling() method to construct this object @@ -221,8 +231,11 @@ def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None 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 ------- @@ -236,7 +249,12 @@ def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None xarray.Dataset.groupby """ super().__init__( - obj, windows, min_periods=min_periods, center=center, keep_attrs=keep_attrs + obj, + windows, + min_periods=min_periods, + center=center, + pad=pad, + keep_attrs=keep_attrs, ) # TODO legacy attribute @@ -357,15 +375,21 @@ def _construct( stride = self._mapping_to_list(stride, default=1) window = obj.variable.rolling_window( - self.dim, self.window, window_dim, self.center, fill_value=fill_value + self.dim, + self.window, + window_dim, + self.center, + self.pad, + fill_value=fill_value, ) attrs = obj.attrs if keep_attrs else {} + coords = self._get_rolling_dim_coords(all=True) result = DataArray( window, dims=obj.dims + tuple(window_dim), - coords=obj.coords, + coords=coords, attrs=attrs, name=obj.name, ) @@ -512,8 +536,11 @@ def _bottleneck_reduce(self, func, keep_attrs, **kwargs): values = values[valid] attrs = self.obj.attrs if keep_attrs else {} + output_dim_coords = self._get_rolling_dim_coords() - return DataArray(values, self.obj.coords, attrs=attrs, name=self.obj.name) + return DataArray(values, self.obj.coords, attrs=attrs, name=self.obj.name).sel( + output_dim_coords + ) def _numpy_or_bottleneck_reduce( self, @@ -557,11 +584,51 @@ def _numpy_or_bottleneck_reduce( return self.reduce(array_agg_func, keep_attrs=keep_attrs, **kwargs) + def _get_rolling_dim_coords(self, all=False) -> Dict[str, Any]: + # If any of the dimensions are not padded, the output size can be shorter than the input size + # along that dimension, so we also need to shorten the corresponding coordinates. + + # 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. + + coords = self.obj.coords + dim = list(self.obj.coords.keys()) if all else self.dim + window = self.window + center = self.center + pad = self.pad + + if pad is False: + pad = [True] + elif utils.is_list_like(pad): + pad = [not p for p in pad] + + offset_pads = utils.get_pads(self.dim, window, center, pad) + + output_coords: Dict[str, Any] = {} + + for d in dim: + obj_coords = coords[d] + if d not in self.dim: + output_coords[d] = obj_coords + continue + + pad_start_offset, pad_end_offset = offset_pads[d] + + start_offset = None if not pad_start_offset else pad_start_offset + end_offset = None if not pad_end_offset else -pad_end_offset + + output_coords[d] = obj_coords[slice(start_offset, end_offset)] + + return output_coords + class DatasetRolling(Rolling): __slots__ = ("rollings",) - def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None): + def __init__( + self, obj, windows, min_periods=None, center=False, pad=True, keep_attrs=None + ): """ Moving window object for Dataset. You should use Dataset.rolling() method to construct this object @@ -580,6 +647,9 @@ def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None 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 ------- @@ -592,22 +662,23 @@ def __init__(self, obj, windows, min_periods=None, center=False, keep_attrs=None xarray.Dataset.groupby xarray.DataArray.groupby """ - super().__init__(obj, windows, min_periods, center, keep_attrs) + super().__init__(obj, windows, min_periods, center, pad, keep_attrs) 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 .dataset import Dataset @@ -625,7 +696,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_rolling_dim_coords(all_dims=True) + + return Dataset(reduced, coords=coords, attrs=attrs) def reduce(self, func, keep_attrs=None, **kwargs): """Reduce the items in this group by applying `func` along some diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 72e34932579..292bf5e7ff8 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -8,12 +8,14 @@ import re import sys import warnings +from collections import defaultdict from enum import Enum from typing import ( Any, Callable, Collection, Container, + DefaultDict, Dict, Hashable, Iterable, @@ -915,3 +917,137 @@ def iterate_nested(nested_list): yield from iterate_nested(item) else: yield item + + +ItemOrSequence = Union[T, Sequence[T]] + + +def expand_args_to_num_dims( + dim: ItemOrSequence[str], + arg_names: Sequence[str], + args: Sequence[ItemOrSequence[Any]], +) -> Tuple[Sequence[str], 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): + for name, arg in zip(arg_names, args): + if not is_scalar(arg): + raise ValueError(f"Expected {name}={arg!r} to be a scalar like 'dim'.") + assert isinstance(dim, str) or isinstance(dim, bytes) + dim = [dim] + + # dim is now a list + nroll = len(dim) + + def to_array(arg): + if is_scalar(arg): + return [arg] * nroll + return arg + + arr_args = [to_array(arg) for arg in args] + + if any(len(dim) != len(arg) for arg in arr_args): + names_vals = ", ".join( + f"{name}={val!r}" for name, val in zip(arg_names, arr_args) + ) + raise ValueError( + "Arguments must all be the same length. " f"Received {names_vals}." + ) + + return dim, arr_args + + +def get_pads( + dim: Sequence[str], + window: Sequence[int], + center: Sequence[bool], + pad: Sequence[bool], +) -> Dict[str, Tuple[int, int]]: + """Return a mapping from dim to the amount of padding to use at the each end of that dimension + + Parameters + ---------- + dim : sequence of str + 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 + ------- + Dict[str, Tuple[int, int]] + """ + pads = {} + for d, win, cent, p in zip(dim, window, center, pad): + if cent: + start = win // 2 # 10 -> 5, 9 -> 4 + end = (win - 1) // 2 # 10 -> 4, 9 -> 4 + pads[d] = (start, end) if p else (0, 0) + else: + pads[d] = (win - 1, 0) if p else (0, 0) + + return pads + + +def get_slice_offsets( + dim: Sequence[str], + window: Sequence[int], + center: Sequence[bool], + pad: Sequence[bool], +) -> DefaultDict[Hashable, Tuple[Optional[int], Optional[int]]]: + """Return a mapping from dims to the start and ends of the output along those dimensions + + The end of the input is indicated using a negative index. + + Parameters + ---------- + dim : sequence of str + 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 + ------- + Dict[str, Tuple[int, int]] + """ + if pad is False: + pad = [True] + elif is_list_like(pad): + pad = [not p for p in pad] + + pads = get_pads(dim, window, center, pad) + + offsets: DefaultDict[Hashable, Tuple[Optional[int], Optional[int]]] = defaultdict( + lambda: (None, None) + ) + for d, (start_offset, end_offset) in pads.items(): + _start_offset = None if start_offset == 0 else start_offset + _end_offset = None if end_offset == 0 else -end_offset + offsets[d] = (_start_offset, _end_offset) + + return offsets diff --git a/xarray/core/variable.py b/xarray/core/variable.py index ace09c6f482..aeb1536061d 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -42,6 +42,8 @@ drop_dims_from_indexers, either_dict_or_kwargs, ensure_us_time_resolution, + expand_args_to_num_dims, + get_pads, infix_dims, is_duck_array, maybe_coerce_to_str, @@ -2028,7 +2030,7 @@ def rank(self, dim, pct=False): return Variable(self.dims, 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. @@ -2040,14 +2042,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 @@ -2083,6 +2089,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) @@ -2091,43 +2104,13 @@ 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_num_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) + pads = get_pads(dim, window, center, pad) padded = var.pad(pads, mode="constant", constant_values=fill_value) axis = [self.get_axis_num(d) for d in dim] diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index b9f04085935..4cfccf7a85f 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6504,12 +6504,22 @@ def test_isin(da): @pytest.mark.parametrize("da", (1, 2), indirect=True) -def test_rolling_iter(da): - rolling_obj = da.rolling(time=7) +@pytest.mark.parametrize("center", (True, False, None)) +@pytest.mark.parametrize("pad", (True, False, None)) +def test_rolling_iter(da, center, pad): + rolling_obj = da.rolling(time=7, center=center, pad=pad) rolling_obj_mean = rolling_obj.mean() - assert len(rolling_obj.window_labels) == len(da["time"]) - assert_identical(rolling_obj.window_labels, da["time"]) + if pad: + expected_times = da["time"] + else: + if center: + expected_times = da["time"][slice(3, -3)] + else: + expected_times = da["time"][slice(6, None)] + + assert len(rolling_obj.window_labels) == len(expected_times) + assert_identical(rolling_obj.window_labels, expected_times) for i, (label, window_da) in enumerate(rolling_obj): assert label == da["time"].isel(time=i) @@ -6536,6 +6546,21 @@ def test_rolling_repr(da): 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(): @@ -6567,10 +6592,9 @@ def test_rolling_properties(da): @pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) -@pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) -def test_rolling_wrapped_bottleneck(da, name, center, min_periods): +def test_rolling_wrapped_bottleneck(da, name, min_periods): bn = pytest.importorskip("bottleneck", minversion="1.1") # Test all bottleneck functions @@ -6586,10 +6610,26 @@ def test_rolling_wrapped_bottleneck(da, name, center, min_periods): with pytest.warns(DeprecationWarning, match="Reductions are applied"): getattr(rolling_obj, name)(dim="time") - # Test center - rolling_obj = da.rolling(time=7, center=center) + +@pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) +@pytest.mark.parametrize("center", (True, False, None)) +@pytest.mark.parametrize("pad", (True, False, None)) +@pytest.mark.parametrize("backend", ["numpy"], indirect=True) +def test_rolling_wrapped_bottleneck_center_pad(da, name, center, pad): + pytest.importorskip("bottleneck", minversion="1.1") + + rolling_obj = da.rolling(time=7, center=center, pad=pad) actual = getattr(rolling_obj, name)()["time"] - assert_equal(actual, da["time"]) + + if pad: + expected = da["time"] + else: + if center: + expected = da["time"][slice(3, -3)] + else: + expected = da["time"][slice(6, None)] + + assert_equal(actual, expected) @requires_dask From 90228b09935a1551d66cf43173cc7372919524a2 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Wed, 14 Jul 2021 11:10:00 -0700 Subject: [PATCH 02/25] Fix iteration on rolling windows to handle no padding Also changes the object returned during iteration, so that it is the same (modulo transposition) as the view created during rolling_window.construct() along the rolling axis. --- xarray/core/rolling.py | 142 ++++++++++++++++++++++----------- xarray/tests/test_dataarray.py | 33 ++++---- 2 files changed, 111 insertions(+), 64 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index e71e40252c4..3124367f871 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -208,6 +208,54 @@ def _get_keep_attrs(self, keep_attrs): return keep_attrs + def _get_rolling_dim_coords(self, all_dims=False) -> Dict[str, Any]: + # If any of the dimensions are not padded, the output size can be shorter than the input size + # along that dimension, so we also need to shorten the corresponding coordinates. + + # 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. + + coords = self.obj.coords + dim = list(self.obj.coords.keys()) if all_dims else self.dim + window = self.window + center = self.center + pad = self.pad + + if pad is False: + pad = [True] + elif utils.is_list_like(pad): + pad = [not p for p in pad] + + offset_pads = utils.get_pads(self.dim, window, center, pad) + + output_coords: Dict[str, Any] = {} + + def get_offsets(d: str): + pad_start_offset, pad_end_offset = offset_pads[d] + + start_offset = None if not pad_start_offset else pad_start_offset + end_offset = None if not pad_end_offset else -pad_end_offset + + return start_offset, end_offset + + for d in dim: + obj_coords = coords[d] + if d in self.dim: + start_offset, end_offset = get_offsets(d) + obj_coords = obj_coords.isel({d: slice(start_offset, end_offset)}) + else: + for obj_d in obj_coords.dims: + if obj_d in self.dim: + start_offset, end_offset = get_offsets(obj_d) + obj_coords = obj_coords.isel( + {obj_d: slice(start_offset, end_offset)} + ) + + output_coords[d] = obj_coords + + return output_coords + class DataArrayRolling(Rolling): __slots__ = ("window_labels",) @@ -263,11 +311,47 @@ def __init__( def __iter__(self): if len(self.dim) > 1: raise ValueError("__iter__ is only supported for 1d-rolling") - stops = np.arange(1, len(self.window_labels) + 1) - starts = stops - int(self.window[0]) - starts[: int(self.window[0])] = 0 - for (label, start, stop) in zip(self.window_labels, starts, stops): - window = self.obj.isel(**{self.dim[0]: slice(start, stop)}) + 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 = utils.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 + ) + + 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) + ) + + 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({self.dim[0]: slice(start, stop)}) counts = window.count(dim=self.dim[0]) window = window.where(counts >= self.min_periods) @@ -384,7 +468,7 @@ def _construct( ) attrs = obj.attrs if keep_attrs else {} - coords = self._get_rolling_dim_coords(all=True) + coords = self._get_rolling_dim_coords(all_dims=True) result = DataArray( window, @@ -486,15 +570,17 @@ def _counts(self, keep_attrs): # array is faster to be reduced than object array. # The use of skipna==False is also faster since it does not need to # copy the strided array. + output_dim_coords = self._get_rolling_dim_coords() counts = ( self.obj.notnull(keep_attrs=keep_attrs) .rolling( center={d: self.center[i] for i, d in enumerate(self.dim)}, + pad={p: self.pad[i] for i, p in enumerate(self.pad)}, **{d: w for d, w in zip(self.dim, self.window)}, ) .construct(rolling_dim, fill_value=False, keep_attrs=keep_attrs) .sum(dim=list(rolling_dim.values()), skipna=False, keep_attrs=keep_attrs) - ) + ).sel(output_dim_coords) return counts def _bottleneck_reduce(self, func, keep_attrs, **kwargs): @@ -584,44 +670,6 @@ def _numpy_or_bottleneck_reduce( return self.reduce(array_agg_func, keep_attrs=keep_attrs, **kwargs) - def _get_rolling_dim_coords(self, all=False) -> Dict[str, Any]: - # If any of the dimensions are not padded, the output size can be shorter than the input size - # along that dimension, so we also need to shorten the corresponding coordinates. - - # 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. - - coords = self.obj.coords - dim = list(self.obj.coords.keys()) if all else self.dim - window = self.window - center = self.center - pad = self.pad - - if pad is False: - pad = [True] - elif utils.is_list_like(pad): - pad = [not p for p in pad] - - offset_pads = utils.get_pads(self.dim, window, center, pad) - - output_coords: Dict[str, Any] = {} - - for d in dim: - obj_coords = coords[d] - if d not in self.dim: - output_coords[d] = obj_coords - continue - - pad_start_offset, pad_end_offset = offset_pads[d] - - start_offset = None if not pad_start_offset else pad_start_offset - end_offset = None if not pad_end_offset else -pad_end_offset - - output_coords[d] = obj_coords[slice(start_offset, end_offset)] - - return output_coords - class DatasetRolling(Rolling): __slots__ = ("rollings",) @@ -820,7 +868,9 @@ def construct( attrs = self.obj.attrs if keep_attrs else {} - return Dataset(dataset, coords=self.obj.coords, attrs=attrs).isel( + coords = self._get_rolling_dim_coords(all_dims=True) + + return Dataset(dataset, coords=coords, attrs=attrs).isel( **{d: slice(None, None, s) for d, s in zip(self.dim, stride)} ) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 4cfccf7a85f..94004295c29 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6506,23 +6506,16 @@ def test_isin(da): @pytest.mark.parametrize("da", (1, 2), indirect=True) @pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("pad", (True, False, None)) -def test_rolling_iter(da, center, pad): - rolling_obj = da.rolling(time=7, center=center, pad=pad) +@pytest.mark.parametrize("min_periods", (1, 6, None)) +@pytest.mark.parametrize("window", (6, 7)) +def test_rolling_iter(da, center, pad, min_periods, window): + rolling_obj = da.rolling( + time=window, center=center, pad=pad, min_periods=min_periods + ) rolling_obj_mean = rolling_obj.mean() - if pad: - expected_times = da["time"] - else: - if center: - expected_times = da["time"][slice(3, -3)] - else: - expected_times = da["time"][slice(6, None)] - - assert len(rolling_obj.window_labels) == len(expected_times) - assert_identical(rolling_obj.window_labels, expected_times) - for i, (label, window_da) in enumerate(rolling_obj): - assert label == da["time"].isel(time=i) + assert label == rolling_obj_mean["time"].isel(time=i) actual = rolling_obj_mean.isel(time=i) expected = window_da.mean("time") @@ -6531,10 +6524,14 @@ def test_rolling_iter(da, center, pad): # as well as the closeness of the values. assert_array_equal(actual.isnull(), expected.isnull()) if (~actual.isnull()).sum() > 0: - np.allclose( - actual.values[actual.values.nonzero()], - expected.values[expected.values.nonzero()], - ) + 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) @pytest.mark.parametrize("da", (1,), indirect=True) From c43db785ff1a751e49992e174f44d8e71387255e Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Thu, 15 Jul 2021 16:38:54 -0700 Subject: [PATCH 03/25] Add tests for rolling(..., pad=False) --- xarray/tests/__init__.py | 18 +++++ xarray/tests/test_dataarray.py | 138 +++++++++++++++++++++++---------- xarray/tests/test_dataset.py | 43 ++++++++-- 3 files changed, 154 insertions(+), 45 deletions(-) diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 9029dc1c621..3e9912d2be9 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -231,6 +231,24 @@ def create_test_data(seed=None, add_attrs=True): return obj +def get_expected_rolling_indices(count, 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 = count + elif center: + start_index = window // 2 # 10 -> 5, 9 -> 4 + end_index = count - (window - 1) // 2 # 10 -> 4, 9 -> 4 + else: + start_index = window - 1 + end_index = count + + expected_index = np.arange(start_index, end_index, stride) + + return expected_index + + _CFTIME_CALENDARS = [ "365_day", "360_day", diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 94004295c29..f92c3fcdc74 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -33,6 +33,7 @@ assert_array_equal, assert_equal, assert_identical, + get_expected_rolling_indices, has_dask, raise_if_dask_computes, requires_bottleneck, @@ -6505,7 +6506,7 @@ def test_isin(da): @pytest.mark.parametrize("da", (1, 2), indirect=True) @pytest.mark.parametrize("center", (True, False, None)) -@pytest.mark.parametrize("pad", (True, False, None)) +@pytest.mark.parametrize("pad", (True, False)) @pytest.mark.parametrize("min_periods", (1, 6, None)) @pytest.mark.parametrize("window", (6, 7)) def test_rolling_iter(da, center, pad, min_periods, window): @@ -6610,21 +6611,18 @@ def test_rolling_wrapped_bottleneck(da, name, min_periods): @pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) @pytest.mark.parametrize("center", (True, False, None)) -@pytest.mark.parametrize("pad", (True, False, None)) +@pytest.mark.parametrize("pad", (True, False)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_wrapped_bottleneck_center_pad(da, name, center, pad): pytest.importorskip("bottleneck", minversion="1.1") + window = 7 + count = len(da["time"]) rolling_obj = da.rolling(time=7, center=center, pad=pad) actual = getattr(rolling_obj, name)()["time"] - if pad: - expected = da["time"] - else: - if center: - expected = da["time"][slice(3, -3)] - else: - expected = da["time"][slice(6, None)] + expected_index = get_expected_rolling_indices(count, window, center, pad) + expected = da["time"][expected_index] assert_equal(actual, expected) @@ -6632,18 +6630,23 @@ def test_rolling_wrapped_bottleneck_center_pad(da, name, center, pad): @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(da, name, center, min_periods, window): +def test_rolling_wrapped_dask(da, name, center, pad, min_periods, window): # dask version - rolling_obj = da.rolling(time=window, min_periods=min_periods, center=center) + 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) + rolling_obj = da.load().rolling( + time=window, min_periods=min_periods, center=center, pad=pad + ) expected = getattr(rolling_obj, name)() # using all-close because rolling over ghost cells introduces some @@ -6652,39 +6655,52 @@ def test_rolling_wrapped_dask(da, name, center, min_periods, window): # 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) @pytest.mark.parametrize("center", (True, None)) -def test_rolling_wrapped_dask_nochunk(center): +@pytest.mark.parametrize("pad", (True, False)) +def test_rolling_wrapped_dask_nochunk(center, pad): # 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(center, window, min_periods): +@pytest.mark.parametrize("window", (2, 3, 4)) +def test_rolling_pandas_compat(center, pad, window, min_periods): s = pd.Series(np.arange(10)) da = DataArray.from_series(s) 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() - da_rolling = da.rolling(index=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, 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) @@ -6694,10 +6710,12 @@ def test_rolling_pandas_compat(center, window, min_periods): @pytest.mark.parametrize("center", (True, False)) -@pytest.mark.parametrize("window", (1, 2, 3, 4)) +@pytest.mark.parametrize("window", (2, 3, 4)) def test_rolling_construct(center, window): - s = pd.Series(np.arange(10)) + count = 10 + s = pd.Series(np.arange(count)) da = DataArray.from_series(s) + da = da.assign_coords(time=("index", np.arange(1, count + 1))) s_rolling = s.rolling(window, center=center, min_periods=1).mean() da_rolling = da.rolling(index=window, center=center, min_periods=1) @@ -6718,13 +6736,31 @@ def test_rolling_construct(center, window): 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( + count, 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(da, center, min_periods, window, name): +def test_rolling_reduce(da, center, pad, min_periods, window, name): if min_periods is not None and window < min_periods: min_periods = window @@ -6732,7 +6768,9 @@ def test_rolling_reduce(da, center, min_periods, window, name): # 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)) @@ -6742,10 +6780,11 @@ def test_rolling_reduce(da, center, min_periods, window, name): @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(center, min_periods, window, name): +def test_rolling_reduce_nonnumeric(center, pad, min_periods, window, name): da = DataArray( [0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time" ).isnull() @@ -6753,7 +6792,9 @@ def test_rolling_reduce_nonnumeric(center, min_periods, window, name): 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)) @@ -6767,11 +6808,15 @@ def test_rolling_count_correct(): kwargs = [ {"time": 11, "min_periods": 1}, + {"time": 11, "min_periods": 1, "pad": False}, {"time": 11, "min_periods": None}, + {"time": 11, "min_periods": None, "pad": False}, {"time": 7, "min_periods": 2}, + {"time": 7, "min_periods": 2, "pad": False}, ] expecteds = [ DataArray([1, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8], dims="time"), + DataArray([8], dims="time"), DataArray( [ np.nan, @@ -6788,7 +6833,9 @@ def test_rolling_count_correct(): ], dims="time", ), + DataArray([np.nan], dims="time"), DataArray([np.nan, np.nan, 2, 3, 3, 4, 5, 5, 5, 5, 5], dims="time"), + DataArray([5, 5, 5, 5, 5], dims="time"), ] for kwarg, expected in zip(kwargs, expecteds): @@ -6800,17 +6847,20 @@ def test_rolling_count_correct(): @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(da, center, min_periods, name): - rolling_obj = da.rolling(time=3, x=2, center=center, min_periods=min_periods) +def test_ndrolling_reduce(da, center, pad, min_periods, name): + rolling_obj = da.rolling( + time=3, x=2, center=center, pad=pad, 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, )() @@ -6828,23 +6878,33 @@ def test_ndrolling_reduce(da, center, min_periods, name): 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(center, fill_value): +def test_ndrolling_construct(center, pad, fill_value): 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) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index b08ce9ea730..e1fe989e31c 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -40,6 +40,7 @@ assert_equal, assert_identical, create_test_data, + get_expected_rolling_indices, has_cftime, has_dask, requires_bottleneck, @@ -6146,11 +6147,10 @@ def test_rolling_properties(ds): @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("backend", ["numpy"], indirect=True) -def test_rolling_wrapped_bottleneck(ds, name, center, min_periods, key): +def test_rolling_wrapped_bottleneck(ds, name, min_periods, key): bn = pytest.importorskip("bottleneck", minversion="1.1") # Test all bottleneck functions @@ -6168,10 +6168,23 @@ def test_rolling_wrapped_bottleneck(ds, name, center, min_periods, key): raise ValueError assert_array_equal(actual[key].values, expected) - # Test center - rolling_obj = ds.rolling(time=7, center=center) + +@pytest.mark.parametrize("name", ("sum", "mean", "std", "var", "min", "max", "median")) +@pytest.mark.parametrize("center", (True, False, None)) +@pytest.mark.parametrize("pad", (False,)) +@pytest.mark.parametrize("backend", ["numpy"], indirect=True) +def test_rolling_wrapped_bottleneck_center_pad(ds, name, center, pad): + pytest.importorskip("bottleneck", minversion="1.1") + + window = 7 + count = len(ds["time"]) + rolling_obj = ds.rolling(time=window, center=center, pad=pad) actual = getattr(rolling_obj, name)()["time"] - assert_equal(actual, ds["time"]) + + expected_index = get_expected_rolling_indices(count, window, center, pad) + expected = ds["time"][expected_index] + + assert_equal(actual, expected) @requires_numbagg @@ -6249,8 +6262,9 @@ def test_rolling_pandas_compat(center, window, min_periods): @pytest.mark.parametrize("center", (True, False)) -@pytest.mark.parametrize("window", (1, 2, 3, 4)) +@pytest.mark.parametrize("window", (2, 3, 4)) def test_rolling_construct(center, window): + count = 20 df = pd.DataFrame( { "x": np.random.randn(20), @@ -6278,6 +6292,23 @@ def test_rolling_construct(center, window): assert (ds_rolling_mean.isnull().sum() == 0).to_array(dim="vars").all() assert (ds_rolling_mean["x"] == 0.0).sum() >= 0 + # with no padding + ds_rolling = ds.rolling(index=window, center=center, pad=False) + ds_rolling_mean = ds_rolling.construct("window", stride=2).mean("window") + + expected_index = get_expected_rolling_indices( + count, window, center, pad=False, stride=2 + ) + + assert ds_rolling_mean.sizes["index"] == len(expected_index) + assert (ds_rolling_mean.index.values == expected_index).all() + np.testing.assert_allclose( + df_rolling["x"][expected_index].values, ds_rolling_mean["x"].values + ) + np.testing.assert_allclose( + df_rolling.index[expected_index], ds_rolling_mean["index"] + ) + @pytest.mark.slow @pytest.mark.parametrize("ds", (1, 2), indirect=True) From 30a0b47de5deb8cf87af959774c677577e605078 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 13:49:21 -0700 Subject: [PATCH 04/25] Update xarray/core/rolling.py Co-authored-by: Deepak Cherian --- xarray/core/rolling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 3124367f871..b3fb4cfb624 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -217,7 +217,7 @@ def _get_rolling_dim_coords(self, all_dims=False) -> Dict[str, Any]: # So, we invert the `pad` flag(s), call `get_pads()`, and work from there. coords = self.obj.coords - dim = list(self.obj.coords.keys()) if all_dims else self.dim + dim = list(self.obj.coords) if all_dims else self.dim window = self.window center = self.center pad = self.pad From 0314d2fa42e77b83ea1c7ee4ae4c99e4b2fde001 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 13:49:48 -0700 Subject: [PATCH 05/25] Update xarray/core/rolling.py Co-authored-by: Deepak Cherian --- xarray/core/rolling.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index b3fb4cfb624..757a9113f5a 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -223,11 +223,11 @@ def _get_rolling_dim_coords(self, all_dims=False) -> Dict[str, Any]: pad = self.pad if pad is False: - pad = [True] + offset = [True] elif utils.is_list_like(pad): - pad = [not p for p in pad] + offset = [not p for p in pad] - offset_pads = utils.get_pads(self.dim, window, center, pad) + offset_pads = utils.get_pads(self.dim, window, center, offset) output_coords: Dict[str, Any] = {} From 18562b8fb429d32c4ca12925da183dedf35dd9b2 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 15:48:33 -0700 Subject: [PATCH 06/25] Rename _get_rolling_dim_coords -> _get_output_coords --- xarray/core/rolling.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 757a9113f5a..59e27b3d850 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -208,7 +208,7 @@ def _get_keep_attrs(self, keep_attrs): return keep_attrs - def _get_rolling_dim_coords(self, all_dims=False) -> Dict[str, Any]: + def _get_output_coords(self, all_dims=False) -> Dict[str, Any]: # If any of the dimensions are not padded, the output size can be shorter than the input size # along that dimension, so we also need to shorten the corresponding coordinates. @@ -468,7 +468,7 @@ def _construct( ) attrs = obj.attrs if keep_attrs else {} - coords = self._get_rolling_dim_coords(all_dims=True) + coords = self._get_output_coords(all_dims=True) result = DataArray( window, @@ -570,7 +570,7 @@ def _counts(self, keep_attrs): # array is faster to be reduced than object array. # The use of skipna==False is also faster since it does not need to # copy the strided array. - output_dim_coords = self._get_rolling_dim_coords() + output_dim_coords = self._get_output_coords() counts = ( self.obj.notnull(keep_attrs=keep_attrs) .rolling( @@ -622,7 +622,7 @@ def _bottleneck_reduce(self, func, keep_attrs, **kwargs): values = values[valid] attrs = self.obj.attrs if keep_attrs else {} - output_dim_coords = self._get_rolling_dim_coords() + output_dim_coords = self._get_output_coords() return DataArray(values, self.obj.coords, attrs=attrs, name=self.obj.name).sel( output_dim_coords @@ -744,7 +744,7 @@ def _dataset_implementation(self, func, keep_attrs, **kwargs): reduced[key].attrs = {} attrs = self.obj.attrs if keep_attrs else {} - coords = self._get_rolling_dim_coords(all_dims=True) + coords = self._get_output_coords(all_dims=True) return Dataset(reduced, coords=coords, attrs=attrs) @@ -868,7 +868,7 @@ def construct( attrs = self.obj.attrs if keep_attrs else {} - coords = self._get_rolling_dim_coords(all_dims=True) + coords = self._get_output_coords(all_dims=True) return Dataset(dataset, coords=coords, attrs=attrs).isel( **{d: slice(None, None, s) for d, s in zip(self.dim, stride)} From 90ae11d633ccce2fd0990713b62503f82e718065 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 15:55:54 -0700 Subject: [PATCH 07/25] Add comment better explaining _get_output_coords --- xarray/core/rolling.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 59e27b3d850..0fbe5bce3c6 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -209,12 +209,18 @@ def _get_keep_attrs(self, keep_attrs): return keep_attrs def _get_output_coords(self, all_dims=False) -> Dict[str, Any]: - # If any of the dimensions are not padded, the output size can be shorter than the input size - # along that dimension, so we also need to shorten the corresponding coordinates. + """Get output coordinates, taking into account window size, window, centering, and padding. - # 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. + 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. + + If `all_dims` is False, returns coordinates only for the dimension(s) used for the rolling + window. This is most useful if the coordinates will be used in a `da.sel()` call. + + If `all_dims` is True, returns all coordinates. This is most useful for constructing a new + DataArray or Dataset, where the data has already been constructed to be the correct size + along each dimension. + """ coords = self.obj.coords dim = list(self.obj.coords) if all_dims else self.dim @@ -222,6 +228,10 @@ def _get_output_coords(self, all_dims=False) -> Dict[str, Any]: center = self.center pad = self.pad + # 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. + if pad is False: offset = [True] elif utils.is_list_like(pad): From 33ccbf6f085229c22fa07b948ddca710a239099e Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 15:57:02 -0700 Subject: [PATCH 08/25] Simplify _get_output_coords --- xarray/core/rolling.py | 64 +++++++++++++++++------------------------- 1 file changed, 26 insertions(+), 38 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 0fbe5bce3c6..a79bdf33f72 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -208,61 +208,49 @@ def _get_keep_attrs(self, keep_attrs): return keep_attrs - def _get_output_coords(self, all_dims=False) -> Dict[str, Any]: + def _get_output_coords(self, all_coords=False) -> 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. - If `all_dims` is False, returns coordinates only for the dimension(s) used for the rolling + If `all_coords` is False, returns coordinates only for the dimension(s) used for the rolling window. This is most useful if the coordinates will be used in a `da.sel()` call. - If `all_dims` is True, returns all coordinates. This is most useful for constructing a new + If `all_coords` is True, returns all coordinates. This is most useful for constructing a new DataArray or Dataset, where the data has already been constructed to be the correct size along each dimension. """ - - coords = self.obj.coords - dim = list(self.obj.coords) if all_dims else self.dim + # TODO: do we need to include dims without coordinates in output_coordinate_names + # when all_coords is True? (The code here does not include them.) + output_coord_names = list(self.obj.coords) if all_coords else self.dim window = self.window center = self.center pad = self.pad + 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. - if pad is False: - offset = [True] - elif utils.is_list_like(pad): - offset = [not p for p in pad] - - offset_pads = utils.get_pads(self.dim, window, center, offset) - - output_coords: Dict[str, Any] = {} - - def get_offsets(d: str): - pad_start_offset, pad_end_offset = offset_pads[d] + offset = [not p for p in pad] + offsets = utils.get_pads(self.dim, window, center, offset) - start_offset = None if not pad_start_offset else pad_start_offset - end_offset = None if not pad_end_offset else -pad_end_offset - - return start_offset, end_offset - - for d in dim: - obj_coords = coords[d] - if d in self.dim: - start_offset, end_offset = get_offsets(d) - obj_coords = obj_coords.isel({d: slice(start_offset, end_offset)}) - else: - for obj_d in obj_coords.dims: - if obj_d in self.dim: - start_offset, end_offset = get_offsets(obj_d) - obj_coords = obj_coords.isel( - {obj_d: slice(start_offset, end_offset)} - ) + selector: Dict[str, slice] = { + dim: offsets_to_slice(start_offset, end_offset) + for dim, (start_offset, end_offset) in offsets.items() + } - output_coords[d] = obj_coords + output_coords = { + k: self.obj.coords[k].isel(selector, missing_dims="ignore") + for k in output_coord_names + } return output_coords @@ -478,7 +466,7 @@ def _construct( ) attrs = obj.attrs if keep_attrs else {} - coords = self._get_output_coords(all_dims=True) + coords = self._get_output_coords(all_coords=True) result = DataArray( window, @@ -754,7 +742,7 @@ def _dataset_implementation(self, func, keep_attrs, **kwargs): reduced[key].attrs = {} attrs = self.obj.attrs if keep_attrs else {} - coords = self._get_output_coords(all_dims=True) + coords = self._get_output_coords(all_coords=True) return Dataset(reduced, coords=coords, attrs=attrs) @@ -878,7 +866,7 @@ def construct( attrs = self.obj.attrs if keep_attrs else {} - coords = self._get_output_coords(all_dims=True) + coords = self._get_output_coords(all_coords=True) return Dataset(dataset, coords=coords, attrs=attrs).isel( **{d: slice(None, None, s) for d, s in zip(self.dim, stride)} From 5df7b12b539f9be2ae873f70eeaf6897c53b79d8 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 16:42:38 -0700 Subject: [PATCH 09/25] Remove utils.get_slice_offsets (unused) --- xarray/core/utils.py | 45 -------------------------------------------- 1 file changed, 45 deletions(-) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 292bf5e7ff8..809821af9c2 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -8,14 +8,12 @@ import re import sys import warnings -from collections import defaultdict from enum import Enum from typing import ( Any, Callable, Collection, Container, - DefaultDict, Dict, Hashable, Iterable, @@ -1008,46 +1006,3 @@ def get_pads( pads[d] = (win - 1, 0) if p else (0, 0) return pads - - -def get_slice_offsets( - dim: Sequence[str], - window: Sequence[int], - center: Sequence[bool], - pad: Sequence[bool], -) -> DefaultDict[Hashable, Tuple[Optional[int], Optional[int]]]: - """Return a mapping from dims to the start and ends of the output along those dimensions - - The end of the input is indicated using a negative index. - - Parameters - ---------- - dim : sequence of str - 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 - ------- - Dict[str, Tuple[int, int]] - """ - if pad is False: - pad = [True] - elif is_list_like(pad): - pad = [not p for p in pad] - - pads = get_pads(dim, window, center, pad) - - offsets: DefaultDict[Hashable, Tuple[Optional[int], Optional[int]]] = defaultdict( - lambda: (None, None) - ) - for d, (start_offset, end_offset) in pads.items(): - _start_offset = None if start_offset == 0 else start_offset - _end_offset = None if end_offset == 0 else -end_offset - offsets[d] = (_start_offset, _end_offset) - - return offsets From 174750af3b5bc9fc665d45c97c8d356b198afc46 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 21:40:00 -0700 Subject: [PATCH 10/25] Simplify get_pads --- xarray/core/utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 809821af9c2..6ebf08754c4 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -998,11 +998,15 @@ def get_pads( """ 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) if p else (0, 0) + pads[d] = (start, end) else: - pads[d] = (win - 1, 0) if p else (0, 0) + pads[d] = (win - 1, 0) return pads From 564833aa716d17c39ae55998e7eae5576e50912e Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 23:11:20 -0700 Subject: [PATCH 11/25] Fix rolling object construction in _count --- xarray/core/rolling.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index a79bdf33f72..787d93336bd 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -568,17 +568,16 @@ def _counts(self, keep_attrs): # array is faster to be reduced than object array. # The use of skipna==False is also faster since it does not need to # copy the strided array. - output_dim_coords = self._get_output_coords() counts = ( self.obj.notnull(keep_attrs=keep_attrs) .rolling( center={d: self.center[i] for i, d in enumerate(self.dim)}, - pad={p: self.pad[i] for i, p in enumerate(self.pad)}, + pad={d: self.pad[i] for i, d in enumerate(self.dim)}, **{d: w for d, w in zip(self.dim, self.window)}, ) .construct(rolling_dim, fill_value=False, keep_attrs=keep_attrs) .sum(dim=list(rolling_dim.values()), skipna=False, keep_attrs=keep_attrs) - ).sel(output_dim_coords) + ) return counts def _bottleneck_reduce(self, func, keep_attrs, **kwargs): From 68b112a6abca8f7f8650d69ee136f13a5fe3d2e7 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Mon, 19 Jul 2021 23:22:43 -0700 Subject: [PATCH 12/25] Skip padding in variable.rolling if not required --- xarray/core/variable.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/xarray/core/variable.py b/xarray/core/variable.py index aeb1536061d..6d2b8427291 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -46,6 +46,7 @@ get_pads, infix_dims, is_duck_array, + is_list_like, maybe_coerce_to_str, ) @@ -2110,9 +2111,12 @@ def rolling_window( [window, window_dim, center, pad], ) - pads = get_pads(dim, window, center, pad) + if not pad or is_list_like(pad) and 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(d) for d in dim] new_dims = self.dims + tuple(window_dim) return Variable( From ea332a6929bf6fee2ca01e7f06c2a2b6313fd633 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Tue, 20 Jul 2021 14:48:15 -0700 Subject: [PATCH 13/25] Fix rolling on DataArrays with nonunique dimension coordinates --- xarray/core/rolling.py | 52 ++++++++++++++++++---------------- xarray/tests/test_dataarray.py | 5 ++++ 2 files changed, 32 insertions(+), 25 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 787d93336bd..d9eee306ec5 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -208,25 +208,12 @@ def _get_keep_attrs(self, keep_attrs): return keep_attrs - def _get_output_coords(self, all_coords=False) -> Dict[str, Any]: - """Get output coordinates, taking into account window size, window, centering, and padding. + def _get_output_dim_selector(self) -> Dict[str, 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, so we need to shorten and properly label the corresponding coordinates. - - If `all_coords` is False, returns coordinates only for the dimension(s) used for the rolling - window. This is most useful if the coordinates will be used in a `da.sel()` call. - - If `all_coords` is True, returns all coordinates. This is most useful for constructing a new - DataArray or Dataset, where the data has already been constructed to be the correct size - along each dimension. + along that dimension. """ - # TODO: do we need to include dims without coordinates in output_coordinate_names - # when all_coords is True? (The code here does not include them.) - output_coord_names = list(self.obj.coords) if all_coords else self.dim - window = self.window - center = self.center - pad = self.pad def offsets_to_slice(start_offset: int, end_offset: int) -> slice: # Turn start and end offsets into a slice object @@ -239,17 +226,32 @@ def offsets_to_slice(start_offset: int, end_offset: int) -> slice: # 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 pad] - offsets = utils.get_pads(self.dim, window, center, offset) + offset = [not p for p in self.pad] + offsets = utils.get_pads(self.dim, self.window, self.center, offset) selector: Dict[str, 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 output_coord_names + for k in self.obj.coords } return output_coords @@ -466,7 +468,7 @@ def _construct( ) attrs = obj.attrs if keep_attrs else {} - coords = self._get_output_coords(all_coords=True) + coords = self._get_output_coords() result = DataArray( window, @@ -619,10 +621,10 @@ def _bottleneck_reduce(self, func, keep_attrs, **kwargs): values = values[valid] attrs = self.obj.attrs if keep_attrs else {} - output_dim_coords = self._get_output_coords() + output_selector = self._get_output_dim_selector() - return DataArray(values, self.obj.coords, attrs=attrs, name=self.obj.name).sel( - output_dim_coords + return DataArray(values, self.obj.coords, attrs=attrs, name=self.obj.name).isel( + output_selector ) def _numpy_or_bottleneck_reduce( @@ -741,7 +743,7 @@ def _dataset_implementation(self, func, keep_attrs, **kwargs): reduced[key].attrs = {} attrs = self.obj.attrs if keep_attrs else {} - coords = self._get_output_coords(all_coords=True) + coords = self._get_output_coords() return Dataset(reduced, coords=coords, attrs=attrs) @@ -865,7 +867,7 @@ def construct( attrs = self.obj.attrs if keep_attrs else {} - coords = self._get_output_coords(all_coords=True) + coords = self._get_output_coords() return Dataset(dataset, coords=coords, attrs=attrs).isel( **{d: slice(None, None, s) for d, s in zip(self.dim, stride)} diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index f92c3fcdc74..15d774f1cc5 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -7358,6 +7358,11 @@ def test_rolling_exp_keep_attrs(da, func): da.rolling_exp(time=10, keep_attrs=True) +def test_rolling_nonunique_coords(): + 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]) + + def test_no_dict(): d = DataArray() with pytest.raises(AttributeError): From ef0ac63337abe7964f82e04121322ebc0e69a7ab Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Tue, 20 Jul 2021 14:56:49 -0700 Subject: [PATCH 14/25] Rename expand_args_to_num_dims -> expand_args_to_dims --- xarray/core/utils.py | 2 +- xarray/core/variable.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 6ebf08754c4..47365585185 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -920,7 +920,7 @@ def iterate_nested(nested_list): ItemOrSequence = Union[T, Sequence[T]] -def expand_args_to_num_dims( +def expand_args_to_dims( dim: ItemOrSequence[str], arg_names: Sequence[str], args: Sequence[ItemOrSequence[Any]], diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 6d2b8427291..70f0f57f35b 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -42,7 +42,7 @@ drop_dims_from_indexers, either_dict_or_kwargs, ensure_us_time_resolution, - expand_args_to_num_dims, + expand_args_to_dims, get_pads, infix_dims, is_duck_array, @@ -2105,7 +2105,7 @@ def rolling_window( dtype = self.dtype var = self - dim, (window, window_dim, center, pad) = expand_args_to_num_dims( + dim, (window, window_dim, center, pad) = expand_args_to_dims( dim, ["window", "window_dim", "center", "pad"], [window, window_dim, center, pad], From cdadc60a51863c16931db71c83cf4823d7d3fcb9 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Tue, 20 Jul 2021 15:33:00 -0700 Subject: [PATCH 15/25] Update expand_args_to_dims * Generalize type of dim (str -> Hashable) * Rename to_array -> to_list * Make error message clearer (only print unexpected inputs) * Allow lists of length one as inputs * Add tests --- xarray/core/utils.py | 32 +++++++++++++++++--------------- xarray/tests/test_utils.py | 23 ++++++++++++++++++++++- 2 files changed, 39 insertions(+), 16 deletions(-) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 47365585185..b9111e38511 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -921,10 +921,10 @@ def iterate_nested(nested_list): def expand_args_to_dims( - dim: ItemOrSequence[str], + dim: ItemOrSequence[Hashable], arg_names: Sequence[str], args: Sequence[ItemOrSequence[Any]], -) -> Tuple[Sequence[str], Sequence[Sequence[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 @@ -946,31 +946,33 @@ def expand_args_to_dims( list of dims, list of lists of arguments """ if is_scalar(dim): - for name, arg in zip(arg_names, args): - if not is_scalar(arg): - raise ValueError(f"Expected {name}={arg!r} to be a scalar like 'dim'.") - assert isinstance(dim, str) or isinstance(dim, bytes) - dim = [dim] + dim_list: Sequence[Hashable] = [dim] + else: + assert isinstance(dim, list) + dim_list = dim # dim is now a list - nroll = len(dim) + nroll = len(dim_list) - def to_array(arg): + 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_array(arg) for arg in args] + arr_args = [to_list(arg) for arg in args] - if any(len(dim) != len(arg) for arg in arr_args): - names_vals = ", ".join( - f"{name}={val!r}" for name, val in zip(arg_names, arr_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( - "Arguments must all be the same length. " f"Received {names_vals}." + f"Expected all arguments to have len={len(dim_list)}. Received: {names_args_len}" ) - return dim, arr_args + return dim_list, arr_args def get_pads( diff --git a/xarray/tests/test_utils.py b/xarray/tests/test_utils.py index 9c78caea4d6..7a06520bc95 100644 --- a/xarray/tests/test_utils.py +++ b/xarray/tests/test_utils.py @@ -8,7 +8,7 @@ from xarray.coding.cftimeindex import CFTimeIndex from xarray.core import duck_array_ops, utils from xarray.core.indexes import PandasIndex -from xarray.core.utils import either_dict_or_kwargs, iterate_nested +from xarray.core.utils import either_dict_or_kwargs, expand_args_to_dims, iterate_nested from . import assert_array_equal, requires_cftime, requires_dask from .test_coding_times import _all_cftime_date_types @@ -333,3 +333,24 @@ def test_infix_dims_errors(supplied, all_): ) def test_iterate_nested(nested_list, expected): assert list(iterate_nested(nested_list)) == expected + + +def test_expand_args_to_dims(): + dims, (arg1, arg2, arg3, arg4) = expanded_args = 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"]], + ) From a2b235c215aac9993ba2d5f2faaf275b3d3e5850 Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Tue, 20 Jul 2021 15:41:37 -0700 Subject: [PATCH 16/25] Update get_pads() types, comments --- xarray/core/utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index b9111e38511..8986aa57046 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -976,16 +976,16 @@ def to_list(arg): def get_pads( - dim: Sequence[str], + dim: Sequence[Hashable], window: Sequence[int], center: Sequence[bool], pad: Sequence[bool], -) -> Dict[str, Tuple[int, int]]: - """Return a mapping from dim to the amount of padding to use at the each end of that dimension +) -> Dict[Hashable, Tuple[int, int]]: + """The amount of padding to use at the each end of each dimension Parameters ---------- - dim : sequence of str + dim : sequence of str (or hashable) dimension(s) for pads window : sequence to int Size of the window along a given dimension @@ -996,7 +996,7 @@ def get_pads( Returns ------- - Dict[str, Tuple[int, int]] + pads: Dict[Hashable, Tuple[int, int]] """ pads = {} for d, win, cent, p in zip(dim, window, center, pad): From b98ab3ab50f3ba0d5989ff8582a31fe4ad2a22ce Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Tue, 20 Jul 2021 15:43:48 -0700 Subject: [PATCH 17/25] Use "length" instead of "count" to refer to the length of a coordinate --- xarray/tests/__init__.py | 8 ++++---- xarray/tests/test_dataarray.py | 12 ++++++------ 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 3e9912d2be9..68e333c05d0 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -231,18 +231,18 @@ def create_test_data(seed=None, add_attrs=True): return obj -def get_expected_rolling_indices(count, window, center, pad, stride=1): +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 = count + end_index = length elif center: start_index = window // 2 # 10 -> 5, 9 -> 4 - end_index = count - (window - 1) // 2 # 10 -> 4, 9 -> 4 + end_index = length - (window - 1) // 2 # 10 -> 4, 9 -> 4 else: start_index = window - 1 - end_index = count + end_index = length expected_index = np.arange(start_index, end_index, stride) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 15d774f1cc5..55fe1ff081c 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6617,11 +6617,11 @@ def test_rolling_wrapped_bottleneck_center_pad(da, name, center, pad): pytest.importorskip("bottleneck", minversion="1.1") window = 7 - count = len(da["time"]) + length = len(da["time"]) rolling_obj = da.rolling(time=7, center=center, pad=pad) actual = getattr(rolling_obj, name)()["time"] - expected_index = get_expected_rolling_indices(count, window, center, pad) + expected_index = get_expected_rolling_indices(length, window, center, pad) expected = da["time"][expected_index] assert_equal(actual, expected) @@ -6712,10 +6712,10 @@ def test_rolling_pandas_compat(center, pad, window, min_periods): @pytest.mark.parametrize("center", (True, False)) @pytest.mark.parametrize("window", (2, 3, 4)) def test_rolling_construct(center, window): - count = 10 - s = pd.Series(np.arange(count)) + length = 10 + s = pd.Series(np.arange(length)) da = DataArray.from_series(s) - da = da.assign_coords(time=("index", np.arange(1, count + 1))) + 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) @@ -6741,7 +6741,7 @@ def test_rolling_construct(center, window): da_rolling_mean = da_rolling.construct("window", stride=2).mean("window") expected_index = get_expected_rolling_indices( - count, window, center, pad=False, stride=2 + length, window, center, pad=False, stride=2 ) assert da_rolling_mean.sizes["index"] == len(expected_index) From b36a049aa468bb376b5efb536d8b4be29bdb841f Mon Sep 17 00:00:00 2001 From: Kevin Squire Date: Tue, 20 Jul 2021 15:50:17 -0700 Subject: [PATCH 18/25] Update type of _get_output_dim_selector --- xarray/core/rolling.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index d9eee306ec5..98c25142929 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -1,7 +1,7 @@ import functools import itertools import warnings -from typing import Any, Callable, Dict +from typing import Any, Callable, Dict, Hashable import numpy as np @@ -208,7 +208,7 @@ def _get_keep_attrs(self, keep_attrs): return keep_attrs - def _get_output_dim_selector(self) -> Dict[str, Any]: + 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 @@ -229,7 +229,7 @@ def offsets_to_slice(start_offset: int, end_offset: int) -> slice: offset = [not p for p in self.pad] offsets = utils.get_pads(self.dim, self.window, self.center, offset) - selector: Dict[str, slice] = { + selector: Dict[Hashable, slice] = { dim: offsets_to_slice(start_offset, end_offset) for dim, (start_offset, end_offset) in offsets.items() } From 923a59866626936166c67de88a05722dce906a56 Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 11 Aug 2021 09:13:56 -0600 Subject: [PATCH 19/25] fixes. --- xarray/core/common.py | 4 ++-- xarray/core/rolling.py | 45 ++++++++++++++++++++++++++++++++++++++--- xarray/core/utils.py | 41 +------------------------------------ xarray/core/variable.py | 5 ++--- 4 files changed, 47 insertions(+), 48 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index cd46c50a983..03c429ff665 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -836,9 +836,9 @@ def rolling( 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 or mapping, 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, default: True + pad : bool or mapping of hashable to bool, default: True Pad the sides of the window with ``NaN``. For different padding, see ``DataArray.pad`` or ``Dataset.pad``. **window_kwargs : optional diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 500a5f87d44..9d598e7cce9 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -1,7 +1,7 @@ import functools import itertools import warnings -from typing import Any, Callable, Dict, Hashable +from typing import Any, Callable, Dict, Hashable, Sequence, Tuple import numpy as np @@ -37,6 +37,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: """A object that implements the moving window pattern. @@ -209,7 +248,7 @@ def offsets_to_slice(start_offset: int, end_offset: int) -> slice: # So, we invert the `pad` flag(s), call `get_pads()`, and work from there. offset = [not p for p in self.pad] - offsets = utils.get_pads(self.dim, self.window, self.center, offset) + offsets = get_pads(self.dim, self.window, self.center, offset) selector: Dict[Hashable, slice] = { dim: offsets_to_slice(start_offset, end_offset) @@ -296,7 +335,7 @@ def __iter__(self): window = self.window[0] center_offset = window // 2 if center else 0 - pads = utils.get_pads(self.dim, self.window, self.center, self.pad) + 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 diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 64676c36704..c66d0d2b355 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -970,7 +970,7 @@ def expand_args_to_dims( if is_scalar(dim): dim_list: Sequence[Hashable] = [dim] else: - assert isinstance(dim, list) + assert isinstance(dim, Sequence) dim_list = dim # dim is now a list @@ -996,42 +996,3 @@ def to_list(arg): ) return dim_list, arr_args - - -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 diff --git a/xarray/core/variable.py b/xarray/core/variable.py index d7a0f920c03..45099061d84 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -42,6 +42,7 @@ is_duck_dask_array, sparse_array_type, ) +from .rolling import get_pads from .utils import ( NdimSizeLenMixin, OrderedSet, @@ -51,10 +52,8 @@ either_dict_or_kwargs, ensure_us_time_resolution, expand_args_to_dims, - get_pads, infix_dims, is_duck_array, - is_list_like, maybe_coerce_to_str, ) @@ -2159,7 +2158,7 @@ def rolling_window( [window, window_dim, center, pad], ) - if not pad or is_list_like(pad) and all(not p for p in pad): + if all(not p for p in pad): padded = var else: pads = get_pads(dim, window, center, pad) From fea9baf0a49fe4dbeff67f5f220c8b549e8ef05a Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 11 Aug 2021 09:19:44 -0600 Subject: [PATCH 20/25] Add whats-new --- doc/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 3dad685aaf7..68393192480 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -23,6 +23,8 @@ v0.19.1 (unreleased) New Features ~~~~~~~~~~~~ +- Allow control over padding in rolling. (:issue:`2007`, :pr:`5603`). + By `Kevin Squire `_. Breaking changes ~~~~~~~~~~~~~~~~ From 2084c0867f0a25755dc906341845d0c301218239 Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 11 Aug 2021 09:57:52 -0600 Subject: [PATCH 21/25] minor test cleanup --- xarray/tests/test_dataarray.py | 8 +++++--- xarray/tests/test_dataset.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 2c44e552310..16da3c7aa17 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6584,11 +6584,12 @@ def test_rolling_properties(da): da.rolling(time=2, min_periods=0) +@requires_bottleneck @pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_wrapped_bottleneck(da, name, min_periods): - bn = pytest.importorskip("bottleneck", minversion="1.1") + import bottleneck as bn # Test all bottleneck functions rolling_obj = da.rolling(time=7, min_periods=min_periods) @@ -6604,12 +6605,13 @@ def test_rolling_wrapped_bottleneck(da, name, min_periods): getattr(rolling_obj, name)(dim="time") +@requires_bottleneck @pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) @pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("pad", (True, False)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_wrapped_bottleneck_center_pad(da, name, center, pad): - pytest.importorskip("bottleneck", minversion="1.1") + import bottleneck as bn window = 7 length = len(da["time"]) @@ -6659,11 +6661,11 @@ def test_rolling_wrapped_dask(da, name, center, pad, min_periods, window): assert_allclose(actual, expected) +@requires_dask @pytest.mark.parametrize("center", (True, None)) @pytest.mark.parametrize("pad", (True, False)) def test_rolling_wrapped_dask_nochunk(center, pad): # GH:2113 - pytest.importorskip("dask.array") da_day_clim = xr.DataArray( np.arange(1, 367), coords=[np.arange(1, 367)], dims="dayofyear" diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 583bcbd9a93..d7b1cab1bd7 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6123,12 +6123,13 @@ def test_rolling_properties(ds): ds.rolling(time2=2) +@requires_bottleneck @pytest.mark.parametrize("name", ("sum", "mean", "std", "var", "min", "max", "median")) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("key", ("z1", "z2")) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_wrapped_bottleneck(ds, name, min_periods, key): - 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) @@ -6146,12 +6147,13 @@ def test_rolling_wrapped_bottleneck(ds, name, min_periods, key): assert_array_equal(actual[key].values, expected) +@requires_bottleneck @pytest.mark.parametrize("name", ("sum", "mean", "std", "var", "min", "max", "median")) @pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("pad", (False,)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_wrapped_bottleneck_center_pad(ds, name, center, pad): - pytest.importorskip("bottleneck", minversion="1.1") + import bottleneck as bn window = 7 count = len(ds["time"]) @@ -6239,7 +6241,7 @@ def test_rolling_pandas_compat(center, window, min_periods): @pytest.mark.parametrize("center", (True, False)) -@pytest.mark.parametrize("window", (2, 3, 4)) +@pytest.mark.parametrize("window", (1, 2, 3, 4)) def test_rolling_construct(center, window): count = 20 df = pd.DataFrame( From c29d1fbb9a6efa418f31e6ed5004952b928f663d Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 11 Aug 2021 10:00:43 -0600 Subject: [PATCH 22/25] consolidate bottleneck dataarray test Failing for center=True+pad=True; and maybe all pad=False? --- xarray/tests/test_dataarray.py | 38 +++++++++++++--------------------- 1 file changed, 14 insertions(+), 24 deletions(-) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 16da3c7aa17..2308e3b64df 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6588,42 +6588,32 @@ def test_rolling_properties(da): @pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) -def test_rolling_wrapped_bottleneck(da, name, min_periods): +@pytest.mark.parametrize("center", (True, False, None)) +@pytest.mark.parametrize("pad", (True, False)) +def test_rolling_wrapped_bottleneck(da, name, min_periods, center, pad): 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 + ) func_name = f"move_{name}" actual = getattr(rolling_obj, name)() - expected = getattr(bn, func_name)( - da.values, window=7, axis=1, min_count=min_periods + expected_values = getattr(bn, func_name)( + da.values, window=window, axis=1, min_count=min_periods ) - assert_array_equal(actual.values, expected) + expected_indices = get_expected_rolling_indices( + da.sizes["time"], window, center, pad + ) + expected = da.copy(data=expected_values).isel(time=expected_indices) + assert_equal(actual, expected) with pytest.warns(DeprecationWarning, match="Reductions are applied"): getattr(rolling_obj, name)(dim="time") -@requires_bottleneck -@pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) -@pytest.mark.parametrize("center", (True, False, None)) -@pytest.mark.parametrize("pad", (True, False)) -@pytest.mark.parametrize("backend", ["numpy"], indirect=True) -def test_rolling_wrapped_bottleneck_center_pad(da, name, center, pad): - import bottleneck as bn - - window = 7 - length = len(da["time"]) - rolling_obj = da.rolling(time=7, center=center, pad=pad) - actual = getattr(rolling_obj, name)()["time"] - - expected_index = get_expected_rolling_indices(length, window, center, pad) - expected = da["time"][expected_index] - - assert_equal(actual, expected) - - @requires_dask @pytest.mark.parametrize("name", ("mean", "count")) @pytest.mark.parametrize("center", (True, False, None)) From b509ebf8badd94ba1cd88d03a882df1796024ac0 Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 11 Aug 2021 10:32:33 -0600 Subject: [PATCH 23/25] consolidate bottleneck dataset tests --- xarray/tests/test_dataset.py | 33 +++++++++++---------------------- 1 file changed, 11 insertions(+), 22 deletions(-) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index d7b1cab1bd7..d3232ba3ccb 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6127,15 +6127,23 @@ def test_rolling_properties(ds): @pytest.mark.parametrize("name", ("sum", "mean", "std", "var", "min", "max", "median")) @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(ds, name, min_periods, key): +def test_rolling_wrapped_bottleneck(ds, name, min_periods, key, center, pad): 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": @@ -6144,26 +6152,7 @@ def test_rolling_wrapped_bottleneck(ds, name, min_periods, key): ) else: raise ValueError - assert_array_equal(actual[key].values, expected) - - -@requires_bottleneck -@pytest.mark.parametrize("name", ("sum", "mean", "std", "var", "min", "max", "median")) -@pytest.mark.parametrize("center", (True, False, None)) -@pytest.mark.parametrize("pad", (False,)) -@pytest.mark.parametrize("backend", ["numpy"], indirect=True) -def test_rolling_wrapped_bottleneck_center_pad(ds, name, center, pad): - import bottleneck as bn - - window = 7 - count = len(ds["time"]) - rolling_obj = ds.rolling(time=window, center=center, pad=pad) - actual = getattr(rolling_obj, name)()["time"] - - expected_index = get_expected_rolling_indices(count, window, center, pad) - expected = ds["time"][expected_index] - - assert_equal(actual, expected) + assert_equal(actual[key], expected[key].isel(time=expected_indices)) @requires_numbagg From 5a2eadc4477f608d8ba2120b4a914a5387640de9 Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 11 Aug 2021 10:37:55 -0600 Subject: [PATCH 24/25] more test cleanup --- xarray/tests/test_dataarray.py | 72 ++++++++++++++++------------------ 1 file changed, 33 insertions(+), 39 deletions(-) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 2308e3b64df..5293f8c2063 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6790,47 +6790,41 @@ def test_rolling_reduce_nonnumeric(center, pad, min_periods, window, name): assert actual.dims == expected.dims -def test_rolling_count_correct(): +@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(time, min_periods, pad, expected): 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 = [ - {"time": 11, "min_periods": 1}, - {"time": 11, "min_periods": 1, "pad": False}, - {"time": 11, "min_periods": None}, - {"time": 11, "min_periods": None, "pad": False}, - {"time": 7, "min_periods": 2}, - {"time": 7, "min_periods": 2, "pad": False}, - ] - expecteds = [ - DataArray([1, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8], dims="time"), - DataArray([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], dims="time"), - DataArray([np.nan, np.nan, 2, 3, 3, 4, 5, 5, 5, 5, 5], dims="time"), - DataArray([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) From b963b29b622693af066ff62cf16cbacb442730a8 Mon Sep 17 00:00:00 2001 From: dcherian Date: Mon, 27 Mar 2023 21:52:58 -0600 Subject: [PATCH 25/25] merge tests --- xarray/tests/test_dataarray.py | 444 --------------------------------- xarray/tests/test_dataset.py | 319 ----------------------- xarray/tests/test_rolling.py | 291 ++++++++++++++------- 3 files changed, 200 insertions(+), 854 deletions(-) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index e4ca3a472a5..378d471ba6b 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -39,7 +39,6 @@ assert_equal, assert_identical, assert_no_warnings, - get_expected_rolling_indices, has_dask, raise_if_dask_computes, requires_bottleneck, @@ -6157,449 +6156,6 @@ def test_isin(da) -> None: assert_equal(result, expected) -@pytest.mark.parametrize("da", (1, 2), indirect=True) -@pytest.mark.parametrize("center", (True, False, None)) -@pytest.mark.parametrize("pad", (True, False)) -@pytest.mark.parametrize("min_periods", (1, 6, None)) -@pytest.mark.parametrize("window", (6, 7)) -def test_rolling_iter(da, center, pad, min_periods, window): - rolling_obj = da.rolling( - time=window, center=center, pad=pad, min_periods=min_periods - ) - rolling_obj_mean = rolling_obj.mean() - - for i, (label, window_da) in enumerate(rolling_obj): - assert label == rolling_obj_mean["time"].isel(time=i) - - actual = rolling_obj_mean.isel(time=i) - expected = window_da.mean("time") - - # 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) - - -@pytest.mark.parametrize("da", (1,), indirect=True) -def test_rolling_repr(da) -> None: - rolling_obj = da.rolling(time=7) - assert repr(rolling_obj) == "DataArrayRolling [time->7]" - rolling_obj = da.rolling(time=7, center=True) - assert repr(rolling_obj) == "DataArrayRolling [time->7(center)]" - 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() -> None: - # regression test for GH3277, GH2514 - dat = DataArray(np.random.rand(7653, 300), dims=("day", "item")) - dat_chunk = dat.chunk({"item": 20}) - dat_chunk.rolling(day=10).mean().rolling(day=250).std() - - -def test_rolling_doc(da) -> None: - rolling_obj = da.rolling(time=7) - - # argument substitution worked - assert "`mean`" in rolling_obj.mean.__doc__ - - -def test_rolling_properties(da): - rolling_obj = da.rolling(time=4) - - assert rolling_obj.obj.get_axis_num("time") == 1 - - # catching invalid args - with pytest.raises(ValueError, match="window must be > 0"): - da.rolling(time=-2) - - with pytest.raises(ValueError, match="min_periods must be greater than zero"): - da.rolling(time=2, min_periods=0) - - -@requires_bottleneck -@pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) -@pytest.mark.parametrize("min_periods", (1, None)) -@pytest.mark.parametrize("backend", ["numpy"], indirect=True) -@pytest.mark.parametrize("center", (True, False, None)) -@pytest.mark.parametrize("pad", (True, False)) -def test_rolling_wrapped_bottleneck(da, name, min_periods, center, pad) -> None: - import bottleneck as bn - - window = 7 - # Test all bottleneck functions - rolling_obj = da.rolling( - time=window, min_periods=min_periods, center=center, pad=pad - ) - - func_name = f"move_{name}" - actual = getattr(rolling_obj, 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 - ) - expected = da.copy(data=expected_values).isel(time=expected_indices) - assert_equal(actual, expected) - - with pytest.warns(DeprecationWarning, match="Reductions are applied"): - getattr(rolling_obj, name)(dim="time") - - -@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(da, name, center, pad, min_periods, window) -> None: - # dask version - 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, pad=pad - ) - expected = getattr(rolling_obj, name)() - - # using all-close because rolling over ghost cells introduces some - # precision errors - assert_allclose(actual, expected) - - # with zero chunked array GH:2113 - rolling_obj = da.chunk().rolling( - 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)) -@pytest.mark.parametrize("pad", (True, False)) -def test_rolling_wrapped_dask_nochunk(center, pad) -> None: - # GH:2113 - - 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, 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", (2, 3, 4)) -def test_rolling_pandas_compat(center, pad, window, min_periods) -> None: - s = pd.Series(np.arange(10)) - da = DataArray.from_series(s) - - if min_periods is not None and window < min_periods: - min_periods = window - - 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, pad=pad, min_periods=min_periods - ).mean() - da_rolling_np = da.rolling( - index=window, center=center, pad=pad, min_periods=min_periods - ).reduce(np.nanmean) - - np.testing.assert_allclose(s_rolling.values, da_rolling.values) - np.testing.assert_allclose(s_rolling.index, da_rolling["index"]) - np.testing.assert_allclose(s_rolling.values, da_rolling_np.values) - np.testing.assert_allclose(s_rolling.index, da_rolling_np["index"]) - - -@pytest.mark.parametrize("center", (True, False)) -@pytest.mark.parametrize("window", (2, 3, 4)) -def test_rolling_construct(center, window) -> None: - 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) - - da_rolling_mean = da_rolling.construct("window").mean("window") - np.testing.assert_allclose(s_rolling.values, da_rolling_mean.values) - np.testing.assert_allclose(s_rolling.index, da_rolling_mean["index"]) - - # with stride - da_rolling_mean = da_rolling.construct("window", stride=2).mean("window") - np.testing.assert_allclose(s_rolling.values[::2], da_rolling_mean.values) - np.testing.assert_allclose(s_rolling.index[::2], da_rolling_mean["index"]) - - # with fill_value - da_rolling_mean = da_rolling.construct("window", stride=2, fill_value=0.0).mean( - "window" - ) - 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(da, center, pad, min_periods, window, name) -> None: - if min_periods is not None and window < min_periods: - min_periods = window - - if da.isnull().sum() > 1 and window == 1: - # this causes all nan slices - window = 2 - - 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)) - expected = getattr(rolling_obj, name)() - assert_allclose(actual, expected) - assert actual.dims == expected.dims - - -@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(center, pad, min_periods, window, name) -> None: - da = DataArray( - [0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time" - ).isnull() - - if min_periods is not None and window < min_periods: - min_periods = window - - 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)) - expected = getattr(rolling_obj, name)() - assert_allclose(actual, expected) - assert actual.dims == expected.dims - - -@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(time, min_periods, pad, expected) -> 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) - - 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, {"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(da, center, pad, min_periods, name) -> None: - rolling_obj = da.rolling( - time=3, x=2, center=center, pad=pad, min_periods=min_periods - ) - - actual = getattr(rolling_obj, name)() - expected = getattr( - getattr( - 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, - )() - - assert_allclose(actual, expected) - assert actual.dims == expected.dims - - if name in ["mean"]: - # test our reimplementation of nanmean using np.nanmean - expected = getattr(rolling_obj.construct({"time": "tw", "x": "xw"}), name)( - ["tw", "xw"] - ) - count = rolling_obj.count() - if min_periods is None: - min_periods = 1 - assert_allclose(actual, expected.where(count >= min_periods)) - - -@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(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, pad=pad).construct( - x="x1", z="z1", fill_value=fill_value - ) - 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["x"], pad=pad["x"]) - .construct(x="x1", fill_value=fill_value) - .rolling(z=2, center=center["z"], pad=pad["z"]) - .construct(z="z1", fill_value=fill_value) - ) - assert_allclose(actual, expected) - - -@pytest.mark.parametrize( - "funcname, argument", - [ - ("reduce", (np.mean,)), - ("mean", ()), - ("construct", ("window_dim",)), - ("count", ()), - ], -) -def test_rolling_keep_attrs(funcname, argument) -> None: - attrs_da = {"da_attr": "test"} - - data = np.linspace(10, 15, 100) - coords = np.linspace(1, 10, 100) - - da = DataArray( - data, dims=("coord"), coords={"coord": coords}, attrs=attrs_da, name="name" - ) - - # attrs are now kept per default - func = getattr(da.rolling(dim={"coord": 5}), funcname) - result = func(*argument) - assert result.attrs == attrs_da - assert result.name == "name" - - # discard attrs - func = getattr(da.rolling(dim={"coord": 5}), funcname) - result = func(*argument, keep_attrs=False) - assert result.attrs == {} - assert result.name == "name" - - # test discard attrs using global option - func = getattr(da.rolling(dim={"coord": 5}), funcname) - with set_options(keep_attrs=False): - result = func(*argument) - assert result.attrs == {} - assert result.name == "name" - - # keyword takes precedence over global option - func = getattr(da.rolling(dim={"coord": 5}), funcname) - with set_options(keep_attrs=False): - result = func(*argument, keep_attrs=True) - assert result.attrs == attrs_da - assert result.name == "name" - - func = getattr(da.rolling(dim={"coord": 5}), funcname) - with set_options(keep_attrs=True): - result = func(*argument, keep_attrs=False) - assert result.attrs == {} - assert result.name == "name" - - def test_raise_no_warning_for_nan_in_binary_ops() -> None: with assert_no_warnings(): xr.DataArray([1, 2, np.NaN]) > 0 diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 40fcd18cd40..2e23d02a261 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -44,7 +44,6 @@ assert_identical, assert_no_warnings, create_test_data, - get_expected_rolling_indices, has_cftime, has_dask, raise_if_dask_computes, @@ -6435,324 +6434,6 @@ def test_dir_unicode(ds) -> None: assert "unicode" in result -@pytest.fixture(params=[1]) -def ds(request, backend): - if request.param == 1: - ds = Dataset( - dict( - z1=(["y", "x"], np.random.randn(2, 8)), - z2=(["time", "y"], np.random.randn(10, 2)), - ), - dict( - x=("x", np.linspace(0, 1.0, 8)), - time=("time", np.linspace(0, 1.0, 10)), - c=("y", ["a", "b"]), - y=range(2), - ), - ) - elif request.param == 2: - ds = Dataset( - dict( - z1=(["time", "y"], np.random.randn(10, 2)), - z2=(["time"], np.random.randn(10)), - z3=(["x", "time"], np.random.randn(8, 10)), - ), - dict( - x=("x", np.linspace(0, 1.0, 8)), - time=("time", np.linspace(0, 1.0, 10)), - c=("y", ["a", "b"]), - y=range(2), - ), - ) - elif request.param == 3: - ds = create_test_data() - else: - raise ValueError - - if backend == "dask": - return ds.chunk() - - return ds - - -@pytest.mark.parametrize( - "funcname, argument", - [ - ("reduce", (np.mean,)), - ("mean", ()), - ("construct", ("window_dim",)), - ("count", ()), - ], -) -def test_rolling_keep_attrs(funcname, argument) -> None: - global_attrs = {"units": "test", "long_name": "testing"} - da_attrs = {"da_attr": "test"} - da_not_rolled_attrs = {"da_not_rolled_attr": "test"} - - data = np.linspace(10, 15, 100) - coords = np.linspace(1, 10, 100) - - ds = Dataset( - data_vars={"da": ("coord", data), "da_not_rolled": ("no_coord", data)}, - coords={"coord": coords}, - attrs=global_attrs, - ) - ds.da.attrs = da_attrs - ds.da_not_rolled.attrs = da_not_rolled_attrs - - # attrs are now kept per default - func = getattr(ds.rolling(dim={"coord": 5}), funcname) - result = func(*argument) - assert result.attrs == global_attrs - assert result.da.attrs == da_attrs - assert result.da_not_rolled.attrs == da_not_rolled_attrs - assert result.da.name == "da" - assert result.da_not_rolled.name == "da_not_rolled" - - # discard attrs - func = getattr(ds.rolling(dim={"coord": 5}), funcname) - result = func(*argument, keep_attrs=False) - assert result.attrs == {} - assert result.da.attrs == {} - assert result.da_not_rolled.attrs == {} - assert result.da.name == "da" - assert result.da_not_rolled.name == "da_not_rolled" - - # test discard attrs using global option - func = getattr(ds.rolling(dim={"coord": 5}), funcname) - with set_options(keep_attrs=False): - result = func(*argument) - - assert result.attrs == {} - assert result.da.attrs == {} - assert result.da_not_rolled.attrs == {} - assert result.da.name == "da" - assert result.da_not_rolled.name == "da_not_rolled" - - # keyword takes precedence over global option - func = getattr(ds.rolling(dim={"coord": 5}), funcname) - with set_options(keep_attrs=False): - result = func(*argument, keep_attrs=True) - - assert result.attrs == global_attrs - assert result.da.attrs == da_attrs - assert result.da_not_rolled.attrs == da_not_rolled_attrs - assert result.da.name == "da" - assert result.da_not_rolled.name == "da_not_rolled" - - func = getattr(ds.rolling(dim={"coord": 5}), funcname) - with set_options(keep_attrs=True): - result = func(*argument, keep_attrs=False) - - assert result.attrs == {} - assert result.da.attrs == {} - assert result.da_not_rolled.attrs == {} - assert result.da.name == "da" - assert result.da_not_rolled.name == "da_not_rolled" - - -def test_rolling_properties(ds) -> None: - # catching invalid args - with pytest.raises(ValueError, match="window must be > 0"): - ds.rolling(time=-2) - with pytest.raises(ValueError, match="min_periods must be greater than zero"): - ds.rolling(time=2, min_periods=0) - with pytest.raises(KeyError, match="time2"): - ds.rolling(time2=2) - - -@requires_bottleneck -@pytest.mark.parametrize("name", ("sum", "mean", "std", "var", "min", "max", "median")) -@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(ds, name, min_periods, key, center, pad) -> None: - import bottleneck as bn - - # Test all bottleneck functions - 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 - ) - else: - raise ValueError - assert_equal(actual[key], expected[key].isel(time=expected_indices)) - - -@pytest.mark.parametrize("center", (True, False)) -@pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) -@pytest.mark.parametrize("window", (1, 2, 3, 4)) -def test_rolling_pandas_compat(center, window, min_periods) -> None: - df = pd.DataFrame( - { - "x": np.random.randn(20), - "y": np.random.randn(20), - "time": np.linspace(0, 1, 20), - } - ) - ds = Dataset.from_dataframe(df) - - if min_periods is not None and window < min_periods: - min_periods = window - - df_rolling = df.rolling(window, center=center, min_periods=min_periods).mean() - ds_rolling = ds.rolling(index=window, center=center, min_periods=min_periods).mean() - - np.testing.assert_allclose(df_rolling["x"].values, ds_rolling["x"].values) - np.testing.assert_allclose(df_rolling.index, ds_rolling["index"]) - - -@pytest.mark.parametrize("center", (True, False)) -@pytest.mark.parametrize("window", (1, 2, 3, 4)) -def test_rolling_construct(center, window) -> None: - count = 20 - df = pd.DataFrame( - { - "x": np.random.randn(20), - "y": np.random.randn(20), - "time": np.linspace(0, 1, 20), - } - ) - - ds = Dataset.from_dataframe(df) - df_rolling = df.rolling(window, center=center, min_periods=1).mean() - ds_rolling = ds.rolling(index=window, center=center) - - ds_rolling_mean = ds_rolling.construct("window").mean("window") - np.testing.assert_allclose(df_rolling["x"].values, ds_rolling_mean["x"].values) - np.testing.assert_allclose(df_rolling.index, ds_rolling_mean["index"]) - - # with stride - ds_rolling_mean = ds_rolling.construct("window", stride=2).mean("window") - np.testing.assert_allclose(df_rolling["x"][::2].values, ds_rolling_mean["x"].values) - np.testing.assert_allclose(df_rolling.index[::2], ds_rolling_mean["index"]) - # with fill_value - ds_rolling_mean = ds_rolling.construct("window", stride=2, fill_value=0.0).mean( - "window" - ) - assert (ds_rolling_mean.isnull().sum() == 0).to_array(dim="vars").all() - assert (ds_rolling_mean["x"] == 0.0).sum() >= 0 - - # with no padding - ds_rolling = ds.rolling(index=window, center=center, pad=False) - ds_rolling_mean = ds_rolling.construct("window", stride=2).mean("window") - - expected_index = get_expected_rolling_indices( - count, window, center, pad=False, stride=2 - ) - - assert ds_rolling_mean.sizes["index"] == len(expected_index) - assert (ds_rolling_mean.index.values == expected_index).all() - np.testing.assert_allclose( - df_rolling["x"][expected_index].values, ds_rolling_mean["x"].values - ) - np.testing.assert_allclose( - df_rolling.index[expected_index], ds_rolling_mean["index"] - ) - - -@pytest.mark.slow -@pytest.mark.parametrize("ds", (1, 2), indirect=True) -@pytest.mark.parametrize("center", (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", "var", "min", "max", "median")) -def test_rolling_reduce(ds, center, min_periods, window, name) -> None: - if min_periods is not None and window < min_periods: - min_periods = window - - if name == "std" and window == 1: - pytest.skip("std with window == 1 is unstable in bottleneck") - - rolling_obj = ds.rolling(time=window, center=center, 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)) - expected = getattr(rolling_obj, name)() - assert_allclose(actual, expected) - assert ds.dims == actual.dims - # make sure the order of data_var are not changed. - assert list(ds.data_vars.keys()) == list(actual.data_vars.keys()) - - # Make sure the dimension order is restored - for key, src_var in ds.data_vars.items(): - assert src_var.dims == actual[key].dims - - -@pytest.mark.parametrize("ds", (2,), indirect=True) -@pytest.mark.parametrize("center", (True, False)) -@pytest.mark.parametrize("min_periods", (None, 1)) -@pytest.mark.parametrize("name", ("sum", "max")) -@pytest.mark.parametrize("dask", (True, False)) -def test_ndrolling_reduce(ds, center, min_periods, name, dask) -> None: - if dask and has_dask: - ds = ds.chunk({"x": 4}) - - rolling_obj = ds.rolling(time=4, x=3, center=center, min_periods=min_periods) - - actual = getattr(rolling_obj, name)() - expected = getattr( - getattr( - ds.rolling(time=4, center=center, min_periods=min_periods), name - )().rolling(x=3, center=center, min_periods=min_periods), - name, - )() - assert_allclose(actual, expected) - assert actual.dims == expected.dims - - # Do it in the opposite order - expected = getattr( - getattr( - ds.rolling(x=3, center=center, min_periods=min_periods), name - )().rolling(time=4, center=center, min_periods=min_periods), - name, - )() - - assert_allclose(actual, expected) - assert actual.dims == expected.dims - - -@pytest.mark.parametrize("center", (True, False, (True, False))) -@pytest.mark.parametrize("fill_value", (np.nan, 0.0)) -@pytest.mark.parametrize("dask", (True, False)) -def test_ndrolling_construct(center, fill_value, dask) -> 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)}, - ) - ds = xr.Dataset({"da": da}) - if dask and has_dask: - ds = ds.chunk({"x": 4}) - - actual = ds.rolling(x=3, z=2, center=center).construct( - x="x1", z="z1", fill_value=fill_value - ) - if not isinstance(center, tuple): - center = (center, center) - expected = ( - ds.rolling(x=3, center=center[0]) - .construct(x="x1", fill_value=fill_value) - .rolling(z=2, center=center[1]) - .construct(z="z1", fill_value=fill_value) - ) - assert_allclose(actual, expected) - - def test_raise_no_warning_for_nan_in_binary_ops() -> None: with assert_no_warnings(): Dataset(data_vars={"x": ("y", [1, 2, np.NaN])}) > 0 diff --git a/xarray/tests/test_rolling.py b/xarray/tests/test_rolling.py index 15548c3523c..d78afa36011 100644 --- a/xarray/tests/test_rolling.py +++ b/xarray/tests/test_rolling.py @@ -1,7 +1,5 @@ from __future__ import annotations -from typing import Any - import numpy as np import pandas as pd import pytest @@ -14,7 +12,9 @@ assert_array_equal, assert_equal, assert_identical, + get_expected_rolling_indices, has_dask, + requires_bottleneck, requires_dask, requires_numbagg, ) @@ -28,9 +28,18 @@ 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"]) @@ -42,7 +51,18 @@ 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() @@ -57,6 +77,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 @@ -82,47 +120,59 @@ def test_rolling_properties(self, da) -> None: with pytest.raises(ValueError, match="min_periods must be greater than zero"): da.rolling(time=2, min_periods=0) + @requires_bottleneck @pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) @pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) - def test_rolling_wrapped_bottleneck(self, da, name, center, min_periods) -> None: - bn = pytest.importorskip("bottleneck", minversion="1.1") + @pytest.mark.parametrize("pad", (True, False)) + def test_rolling_wrapped_bottleneck( + self, da, name, min_periods, center, pad + ) -> None: + 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 + ) func_name = f"move_{name}" actual = getattr(rolling_obj, name)() - expected = getattr(bn, func_name)( - da.values, window=7, axis=1, min_count=min_periods + 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 ) - assert_array_equal(actual.values, expected) + expected = da.copy(data=expected_values).isel(time=expected_indices) + assert_equal(actual, expected) with pytest.warns(DeprecationWarning, match="Reductions are applied"): getattr(rolling_obj, name)(dim="time") - # Test center - rolling_obj = da.rolling(time=7, center=center) - actual = getattr(rolling_obj, name)()["time"] - assert_equal(actual, da["time"]) - @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)() @@ -132,39 +182,52 @@ 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) -> None: + @pytest.mark.parametrize("window", (2, 3, 4)) + def test_rolling_pandas_compat(self, center, pad, window, min_periods) -> None: s = pd.Series(np.arange(10)) da = DataArray.from_series(s) 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) @@ -173,10 +236,12 @@ def test_rolling_pandas_compat(self, center, window, min_periods) -> None: np.testing.assert_allclose(s_rolling.index, da_rolling_np["index"]) @pytest.mark.parametrize("center", (True, False)) - @pytest.mark.parametrize("window", (1, 2, 3, 4)) + @pytest.mark.parametrize("window", (2, 3, 4)) def test_rolling_construct(self, center, window) -> 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) @@ -197,12 +262,32 @@ def test_rolling_construct(self, center, window) -> 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) -> None: + def test_rolling_reduce(self, da, center, pad, min_periods, window, name) -> None: if min_periods is not None and window < min_periods: min_periods = window @@ -210,7 +295,9 @@ def test_rolling_reduce(self, da, center, min_periods, window, name) -> None: # 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)) @@ -219,10 +306,13 @@ def test_rolling_reduce(self, da, center, min_periods, window, name) -> None: assert actual.dims == expected.dims @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) -> None: + def test_rolling_reduce_nonnumeric( + self, center, pad, min_periods, window, name + ) -> None: da = DataArray( [0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time" ).isnull() @@ -230,7 +320,9 @@ def test_rolling_reduce_nonnumeric(self, center, min_periods, window, name) -> N 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)) @@ -238,54 +330,58 @@ def test_rolling_reduce_nonnumeric(self, center, min_periods, window, name) -> N assert_allclose(actual, expected) assert actual.dims == expected.dims - def test_rolling_count_correct(self) -> 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) -> 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) -> None: - rolling_obj = da.rolling(time=3, x=2, center=center, min_periods=min_periods) + def test_ndrolling_reduce(self, da, center, pad, min_periods, name) -> None: + rolling_obj = da.rolling( + time=3, x=2, center=center, pad=pad, 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, )() @@ -302,23 +398,33 @@ def test_ndrolling_reduce(self, da, center, min_periods, name) -> None: 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) @@ -562,23 +668,31 @@ def test_rolling_properties(self, ds) -> None: with pytest.raises(KeyError, match="time2"): ds.rolling(time2=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 + self, ds, name, min_periods, key, center, pad ) -> 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": @@ -587,12 +701,7 @@ def test_rolling_wrapped_bottleneck( ) else: raise ValueError - assert_array_equal(actual[key].values, expected) - - # Test center - rolling_obj = ds.rolling(time=7, center=center) - actual = getattr(rolling_obj, name)()["time"] - assert_equal(actual, ds["time"]) + assert_equal(actual[key], expected[key].isel(time=expected_indices)) @pytest.mark.parametrize("center", (True, False)) @pytest.mark.parametrize("min_periods", (None, 1, 2, 3))