Switch to sliding_window_kwargs

Author: dcherian

xarray/core/dask_array_compat.py

@@ -5,7 +5,7 @@ def sliding_window_view(
     x, window_shape, axis=None, *, automatic_rechunk=True, **kwargs
 ):
     # Backcompat for handling `automatic_rechunk`, delete when dask>=2024.11.0
-    # subok, writeable are unsupported by dask
+    # Note that subok, writeable are unsupported by dask, so we ignore those in kwargs
     from dask.array.lib.stride_tricks import sliding_window_view
 
     if module_available("dask", "2024.11.0"):

xarray/core/duck_array_ops.py

@@ -30,7 +30,6 @@
     transpose,
     unravel_index,
 )
-from numpy.ma import masked_invalid  # noqa
 from packaging.version import Version
 from pandas.api.types import is_extension_array_dtype
 
@@ -107,7 +106,7 @@ def f(*args, **kwargs):
         else:
             wrapped = getattr(eager_module, name)
             for kwarg in dask_only_kwargs:
-                kwargs.pop(kwarg)
+                kwargs.pop(kwarg, None)
         return wrapped(*args, **kwargs)
 
     return f
@@ -125,6 +124,12 @@ def fail_on_dask_array_input(values, msg=None, func_name=None):
 # Requires special-casing because pandas won't automatically dispatch to dask.isnull via NEP-18
 pandas_isnull = _dask_or_eager_func("isnull", eager_module=pd, dask_module="dask.array")
 
+# TODO replace with simply np.ma.masked_invalid once numpy/numpy#16022 is fixed
+# TODO: replacing breaks iris + dask tests
+masked_invalid = _dask_or_eager_func(
+    "masked_invalid", eager_module=np.ma, dask_module="dask.array.ma"
+)
+
 # sliding_window_view will not dispatch arbitrary kwargs (automatic_rechunk),
 # so we need to hand-code this.
 sliding_window_view = _dask_or_eager_func(

xarray/core/rolling.py

@@ -187,7 +187,7 @@ def method(self, keep_attrs=None, **kwargs):
                 rolling_agg_func=rolling_agg_func,
                 keep_attrs=keep_attrs,
                 fillna=fillna,
-                automatic_rechunk=automatic_rechunk,
+                sliding_window_kwargs=dict(automatic_rechunk=automatic_rechunk),
                 **kwargs,
             )
 
@@ -329,7 +329,7 @@ def construct(
         stride: int | Mapping[Any, int] = 1,
         fill_value: Any = dtypes.NA,
         keep_attrs: bool | None = None,
-        automatic_rechunk: bool = True,
+        sliding_window_kwargs: Mapping[Any, Any] | None = None,
         **window_dim_kwargs: Hashable,
     ) -> DataArray:
         """
@@ -348,10 +348,9 @@ def construct(
             If True, the attributes (``attrs``) will be copied from the original
             object to the new one. If False, the new object will be returned
             without attributes. If None uses the global default.
-        automatic_rechunk: bool, default True
-            Whether dask should automatically rechunk the output to avoid
-            exploding chunk sizes. Importantly, each chunk will be a view of the data
-            so large chunk sizes are only safe if *no* copies are made later.
+        sliding_window_kwargs : Mapping
+            Keyword arguments that should be passed to the underlying array type's
+            ``sliding_window_view`` function.
         **window_dim_kwargs : Hashable, optional
             The keyword arguments form of ``window_dim`` {dim: new_name, ...}.
 
@@ -365,6 +364,15 @@ def construct(
         numpy.lib.stride_tricks.sliding_window_view
         dask.array.lib.stride_tricks.sliding_window_view
 
+        Notes
+        -----
+        With dask arrays, it's possible to pass the ``automatic_rechunk`` kwarg as
+        ``sliding_window_kwargs={"automatic_rechunk": True}``. This controls
+        whether dask should automatically rechunk the output to avoid
+        exploding chunk sizes. Automatically rechunking is the default behaviour.
+        Importantly, each chunk will be a view of the data so large chunk sizes are
+        only safe if *no* copies are made later.
+
         Examples
         --------
         >>> da = xr.DataArray(np.arange(8).reshape(2, 4), dims=("a", "b"))
@@ -399,13 +407,15 @@ def construct(
 
         """
 
+        if sliding_window_kwargs is None:
+            sliding_window_kwargs = {}
         return self._construct(
             self.obj,
             window_dim=window_dim,
             stride=stride,
             fill_value=fill_value,
             keep_attrs=keep_attrs,
-            automatic_rechunk=automatic_rechunk,
+            sliding_window_kwargs=sliding_window_kwargs,
             **window_dim_kwargs,
         )
 
@@ -417,11 +427,14 @@ def _construct(
         stride: int | Mapping[Any, int] = 1,
         fill_value: Any = dtypes.NA,
         keep_attrs: bool | None = None,
-        automatic_rechunk: bool = True,
+        sliding_window_kwargs: Mapping[Any, Any] | None = None,
         **window_dim_kwargs: Hashable,
     ) -> DataArray:
         from xarray.core.dataarray import DataArray
 
+        if sliding_window_kwargs is None:
+            sliding_window_kwargs = {}
+
         keep_attrs = self._get_keep_attrs(keep_attrs)
 
         if window_dim is None:
@@ -442,7 +455,7 @@ def _construct(
             window_dims,
             center=self.center,
             fill_value=fill_value,
-            automatic_rechunk=automatic_rechunk,
+            **sliding_window_kwargs,
         )
 
         attrs = obj.attrs if keep_attrs else {}
@@ -463,7 +476,7 @@ def reduce(
         func: Callable,
         keep_attrs: bool | None = None,
         *,
-        automatic_rechunk: bool = True,
+        sliding_window_kwargs: Mapping[Any, Any] | None = None,
         **kwargs: Any,
     ) -> DataArray:
         """Reduce each window by applying `func`.
@@ -480,10 +493,9 @@ def reduce(
             If True, the attributes (``attrs``) will be copied from the original
             object to the new one. If False, the new object will be returned
             without attributes. If None uses the global default.
-        automatic_rechunk: bool, default True
-            Whether dask should automatically rechunk the output of ``construct`` to avoid
-            exploding chunk sizes. Importantly, each chunk will be a view of the data
-            so large chunk sizes are only safe if *no* copies are made in ``func``.
+        sliding_window_kwargs
+            Keyword arguments that should be passed to the underlying array type's
+            ``sliding_window_view`` function.
         **kwargs : dict
             Additional keyword arguments passed on to `func`.
 
@@ -492,6 +504,15 @@ def reduce(
         reduced : DataArray
             Array with summarized data.
 
+        Notes
+        -----
+        With dask arrays, it's possible to pass the ``automatic_rechunk`` kwarg as
+        ``sliding_window_kwargs={"automatic_rechunk": True}``. This controls
+        whether dask should automatically rechunk the output to avoid
+        exploding chunk sizes. Automatically rechunking is the default behaviour.
+        Importantly, each chunk will be a view of the data so large chunk sizes are
+        only safe if *no* copies are made later.
+
         Examples
         --------
         >>> da = xr.DataArray(np.arange(8).reshape(2, 4), dims=("a", "b"))
@@ -541,7 +562,7 @@ def reduce(
             window_dim=rolling_dim,
             keep_attrs=keep_attrs,
             fill_value=fillna,
-            automatic_rechunk=automatic_rechunk,
+            sliding_window_kwargs=sliding_window_kwargs,
         )
 
         dim = list(rolling_dim.values())
@@ -873,7 +894,7 @@ def construct(
         stride: int | Mapping[Any, int] = 1,
         fill_value: Any = dtypes.NA,
         keep_attrs: bool | None = None,
-        automatic_rechunk: bool = True,
+        sliding_window_kwargs: Mapping[Any, Any] | None = None,
         **window_dim_kwargs: Hashable,
     ) -> Dataset:
         """
@@ -889,10 +910,9 @@ def construct(
             size of stride for the rolling window.
         fill_value : Any, default: dtypes.NA
             Filling value to match the dimension size.
-        automatic_rechunk: bool, default True
-            Whether dask should automatically rechunk the output to avoid
-            exploding chunk sizes. Importantly, each chunk will be a view of the data
-            so large chunk sizes are only safe if *no* copies are made later.
+        sliding_window_kwargs
+            Keyword arguments that should be passed to the underlying array type's
+            ``sliding_window_view`` function.
         **window_dim_kwargs : {dim: new_name, ...}, optional
             The keyword arguments form of ``window_dim``.
 
@@ -904,6 +924,15 @@ def construct(
         --------
         numpy.lib.stride_tricks.sliding_window_view
         dask.array.lib.stride_tricks.sliding_window_view
+
+        Notes
+        -----
+        With dask arrays, it's possible to pass the ``automatic_rechunk`` kwarg as
+        ``sliding_window_kwargs={"automatic_rechunk": True}``. This controls
+        whether dask should automatically rechunk the output to avoid
+        exploding chunk sizes. Automatically rechunking is the default behaviour.
+        Importantly, each chunk will be a view of the data so large chunk sizes are
+        only safe if *no* copies are made later.
         """
 
         from xarray.core.dataset import Dataset
@@ -935,6 +964,7 @@ def construct(
                     fill_value=fill_value,
                     stride=st,
                     keep_attrs=keep_attrs,
+                    sliding_window_kwargs=sliding_window_kwargs,
                 )
             else:
                 dataset[key] = da.copy()

xarray/core/variable.py

@@ -2019,7 +2019,7 @@ def rolling_window(
         *,
         center=False,
         fill_value=dtypes.NA,
-        automatic_rechunk: bool = True,
+        **kwargs,
     ):
         """
         Make a rolling_window along dim and add a new_dim to the last place.
@@ -2040,10 +2040,9 @@ def rolling_window(
             of the axis.
         fill_value
             value to be filled.
-        automatic_rechunk: bool, default True
-            Whether dask should automatically rechunk the output to avoid
-            exploding chunk sizes. Importantly, each chunk will be a view of the data
-            so large chunk sizes are only safe if *no* copies are made later.
+        **kwargs
+            Keyword arguments that should be passed to the underlying array type's
+            ``sliding_window_view`` function.
 
         Returns
         -------
@@ -2052,6 +2051,11 @@ def rolling_window(
         The return dim: self.dims + (window_dim, )
         The return shape: self.shape + (window, )
 
+        See Also
+        --------
+        numpy.lib.stride_tricks.sliding_window_view
+        dask.array.lib.stride_tricks.sliding_window_view
+
         Examples
         --------
         >>> v = Variable(("a", "b"), np.arange(8).reshape((2, 4)))
@@ -2132,10 +2136,7 @@ def rolling_window(
         return Variable(
             new_dims,
             duck_array_ops.sliding_window_view(
-                padded.data,
-                window_shape=window,
-                axis=axis,
-                automatic_rechunk=automatic_rechunk,
+                padded.data, window_shape=window, axis=axis, **kwargs
             ),
         )
 

xarray/tests/test_rolling.py

@@ -611,10 +611,10 @@ def test_rolling_construct_automatic_rechunk(self):
 
         # Dataset now has chunks of size (400, 400, 100 100) or 11.92 GiB
         rechunked = da.rolling(time=100, center=True).construct(
-            "window", automatic_rechunk=True
+            "window", sliding_window_kwargs=dict(automatic_rechunk=True)
         )
         not_rechunked = da.rolling(time=100, center=True).construct(
-            "window", automatic_rechunk=False
+            "window", sliding_window_kwargs=dict(automatic_rechunk=False)
         )
         assert rechunked.chunks != not_rechunked.chunks