多维Xarray滚动构造

Question

让xarray滚动构造多维滚动窗口的最好方法是什么？下面是一个numpy示例：

import numpy as np
from numpy.lib.stride_tricks import as_strided

data = np.array(np.arange(6).reshape(2, 3),dtype="float64")


win_size = (
    3  # Size of the window (e.g. 3*3)
)
win_size_half = int(np.floor(win_size / 2))

# pad with nan to get correct window for the edges
data = np.pad(
    data,
    (win_size_half, win_size_half),
    "constant",
    constant_values=(np.nan),
)

sub_shape = (win_size, win_size)
view_shape = tuple(np.subtract(data.shape, sub_shape) + 1) + sub_shape
data_view = as_strided(
    data, view_shape, data.strides * 2
)
data_view = data_view.reshape((-1,) + sub_shape)

#Expected results
>>> data_view
array([[[nan, nan, nan],
        [nan,  0.,  1.],
        [nan,  3.,  4.]],

       [[nan, nan, nan],
        [ 0.,  1.,  2.],
        [ 3.,  4.,  5.]],

       [[nan, nan, nan],
        [ 1.,  2., nan],
        [ 4.,  5., nan]],

       [[nan,  0.,  1.],
        [nan,  3.,  4.],
        [nan, nan, nan]],

       [[ 0.,  1.,  2.],
        [ 3.,  4.,  5.],
        [nan, nan, nan]],

       [[ 1.,  2., nan],
        [ 4.,  5., nan],
        [nan, nan, nan]]])

我想知道如何使用xarray来实现同样的目的。例如，使用xarray执行与上面相同的操作：

import xarray as xr

da =xr.DataArray(np.array(np.arange(6).reshape(2, 3),dtype="float64"),dims=("a","b"))

# And something like
rolling = da.rolling({"a":win_size,"b":win_size})

# producing same results as in numpy example
rolling.construct("window_dim")

据我所知，xr.rolling不允许多维。请让我知道是否有其他方法来做这种类型的操作。

谢谢

Answer 1

xr.rolling现在接受多个维度。您必须提供到rolling.construct的字典映射(或基于关键字)。

您的center=True示例以windows center为例，它不是xr.rolling的默认设置，因此您必须显式地提供numpy

下面的代码给出了和你的numpy代码一样的结果：

import xarray as xr
import numpy as np

da =xr.DataArray(np.array(np.arange(6).reshape(2, 3),dtype="float64"),dims=("a","b"))

rolling = da.rolling({"a":3,"b":3}, center=True)

# producing same results as in numpy example
da_roll = rolling.construct(a='ka',b='kb')
da_roll
Out[2]: 
<xarray.DataArray (a: 2, b: 3, ka: 3, kb: 3)>
array([[[[nan, nan, nan],
         [nan,  0.,  1.],
         [nan,  3.,  4.]],
        [[nan, nan, nan],
         [ 0.,  1.,  2.],
         [ 3.,  4.,  5.]],
        [[nan, nan, nan],
         [ 1.,  2., nan],
         [ 4.,  5., nan]]],
       [[[nan,  0.,  1.],
         [nan,  3.,  4.],
         [nan, nan, nan]],
        [[ 0.,  1.,  2.],
         [ 3.,  4.,  5.],
         [nan, nan, nan]],
        [[ 1.,  2., nan],
         [ 4.,  5., nan],
         [nan, nan, nan]]]])
Dimensions without coordinates: a, b, ka, kb