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

【问题标题】：Operate along a dimension without writing back data in dask array on Xarray沿维度操作，而不在 Xarray 上的 dask 数组中写回数据
【发布时间】：2021-08-17 01:44:58
【问题描述】：

我有一个具有 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 中执行此操作？

【问题讨论】：

标签： python-3.x data-science dask python-xarray

【解决方案1】：

时间维度上不需要循环，可以用向量化的方式来做：

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])

这样就不需要将整个数组加载到内存中。

我意识到，如果您向np.percentile 提供axis 参数并注意使数组形状保持一致，您也可以直接将xarray 的apply_ufunc 与foo 函数一起使用。看来，percentile 函数的 dask 版本并未针对多维数组实现，但您可以使用 apply_ufunc 的 parallelized 选项使其与 numpy 函数一起使用：

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

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

【讨论】：