在 xarray 上具有滚动窗口的分箱变量答案

【问题标题】：Binning variable with rolling window on xarray在 xarray 上具有滚动窗口的分箱变量
【发布时间】：2021-03-11 07:41:08
【问题描述】：

我有一个带有温度数据的 xarray.Dataset，并希望使用 7 天的滚动窗口计算数组中每个元素的分箱温度。

我有这种形式的数据：

import xarray as xr

ds = xr.Dataset(
    {'t2m': (['time', 'lat', 'lon'], t2m)},
    coords={
        'lon': lon,
        'lat': lat,
        'time': time,
    }
)

然后我使用rolling() 方法并在每个窗口数组上应用一个函数：

r = ds.t2m.\
chunk({'time': 10}).\
rolling(time=7)

window_results = []
for label, arr_window in tqdm(r):
    max_temp = arr_window.max(dim=...).values
    min_temp = arr_window.min(dim=...).values
    if not np.isnan(max_temp):
        bins = np.arange(min_temp, max_temp, 2)
        
        buckets = np.digitize(arr_window.isel(time=-1),
                              bins=bins)
        buckets_arr = xr.DataArray(buckets,
                                   dims={
                                       'lat': arr_window.lat.values,
                                       'lon': arr_window.lon.values
                                   })
        buckets_arr = buckets_arr.assign_coords({'time': label})

        window_results.append(buckets_arr)

最后，我得到了每个时间步长的列表，其中包含分箱数组的窗口计算：

ds_concat = xr.concat(window_results, dim='time')
ds_concat

>> <xarray.DataArray (time: 18, lat: 10, lon: 10)>
array([[[1, 2, 2, ..., 2, 2, 3],
        [1, 3, 3, ..., 1, 1, 2],
        [2, 3, 2, ..., 1, 2, 3],
        ...,
        [2, 2, 2, ..., 2, 2, 2],
        [2, 2, 2, ..., 1, 2, 2],
        [2, 2, 3, ..., 2, 3, 2]],
...

此代码正在产生我正在寻找的结果，但我相信必须有更好的替代方案来使用apply_ufunc 或dask 应用相同的过程。我也在使用dask.distributed.Client，所以我正在寻找一种方法来优化我的代码以快速运行。

感谢任何帮助

【问题讨论】：

标签： python python-3.x python-xarray

【解决方案1】：

我终于明白了！希望这可以帮助遇到同样问题的人。

dask.distributed 最酷的功能之一是dask.delayed。我可以重写上面的循环并使用惰性函数：

import dask
import xarray as xr

@dask.delayed
def create_bucket_window(arr, label):
    
    max_temp = arr.max(dim=...).values
    min_temp = arr.min(dim=...).values
    
    if not np.isnan(max_temp):
        bins = np.arange(min_temp, max_temp, 2)
        buckets = np.digitize(arr.isel(time=-1),
                              bins=bins)
        buckets_arr = xr.DataArray(buckets,
                                   dims={
                                       'lat': arr.lat.values,
                                       'lon': arr.lon.values
                                   })
        buckets_arr = buckets_arr.assign_coords({'time': label})

        return buckets_arr

然后：

window_results = []
for label, arr_window in tqdm(r):
    bucket_array = create_bucket_window(arr=arr_window,
                                        label=label)
    window_results.append(bucket_array)

一旦我这样做了，dask 将延迟生成这个数组，并且只会在需要时评估它们：

dask.compute(*window_results)

你会得到一个结果集合！

【讨论】：