沿维度操作而不在 Xarray 上的 dask 数组中写回数据

如何解决沿维度操作而不在 Xarray 上的 dask 数组中写回数据

我有一个具有 3 个维度（“时间”、“x”和“y”）的数据集。我想沿时间维度应用此函数 foo：

def foo(arr):
    lo,hi = np.percentile(arr,(1,99))
    arr = np.clip(arr,lo,hi)
    arr = (arr - lo) / (hi - lo)
    return arr

基本上，我想用 foo 之类的函数在时间维度上缩放每个“图像”

在 numpy 中，我可以这样做：

for i in range(data.shape[0]):
   data[i] = foo(data[i])

但由于数据存储在 dask 数组中，我无法写回修改后的数据。我遇到了这个错误：

TypeError: this variable's data is stored in a dask array,which does not support item assignment. To assign to this variable,you must first load it into memory explicitly using the .load() method or accessing its .values attribute.

如何在 xarray/dask 中执行此操作？

解决方法

不需要在时间维度上循环，你可以用矢量化的方式来做到这一点：

da = xr.tutorial.open_dataset(
    "air_temperature",chunks={"lat": -1,"lon": -1,"time": 10}
)["air"]


def scale_image(da,quantiles):
    quantiles = da.quantile(quantiles,dim=("lat","lon"))
    lower = quantiles.isel(quantile=0,drop=True)
    upper = quantiles.isel(quantile=1,drop=True)
    clipped = xr.apply_ufunc(np.clip,da,lower,upper,dask="allowed")
    return (clipped - lower) / (upper - lower)


scaled = scale_image(da,quantiles=[0.01,0.99])

这样就没有必要将整个数组加载到内存中。

我意识到您也可以将 xarray 的 apply_ufunc 直接与 foo 函数一起使用，如果您向 axis 提供 np.percentile 参数并负责制作数组形状一致。看来，percentile 函数的 dask 版本没有为多维数组实现，但是您可以使用 parallelized 的 apply_ufunc 选项使其与 numpy 函数一起使用:

def foo(arr):
    lo,hi = np.percentile(arr,(1,99),axis=[1,2])
    arr = np.clip(arr,lo[:,None,None],hi[:,None])
    return (arr - lo[:,None]) / (hi[:,None] - lo[:,None])

scaled2 = xr.apply_ufunc(foo,dask="parallelized")