xr.DataArray.where 使用多个 dask 块时将有效点设置为 nan

Question

我正在尝试随机删除 xr.DataArray 的一小部分（如所述并借助中的答案），然后仅访问原始数据集 [=14] 中的值=] 已删除。

只要数据不存储在 dask 数组中或仅存储在一个 dask 数组中，这就可以正常工作。一旦我定义了小于数据总大小的块，原始值就会设置为 nan。

data = xr.DataArray(np.arange(5*5*5.).reshape(5,5,5), dims=('time','latitude','longitude'))
data.to_netcdf('/path/to/file.nc')
#data = xr.open_dataarray('/path/to/file.nc', chunks={'time':5}) # creates expected output
data = xr.open_dataarray('/path/to/file.nc', chunks={'time':2}) # creates observed output 

def set_fraction_randomly_to_nan(data, frac_missing):
    np.random.seed(0)
    data[np.random.rand(*data.shape) < frac_missing] = np.nan
    return data

data_lost = xr.apply_ufunc(set_fraction_randomly_to_nan, data.copy(deep=True), output_core_dims=[['latitude','longitude']], dask='parallelized', input_core_dims=[['latitude','longitude']], output_dtypes=[data.dtype], kwargs={'frac_missing': 0.5})

print(data[0,-4:,-4:].values)
# >>
# [[ 6.  7.  8.  9.]
# [11. 12. 13. 14.]
# [16. 17. 18. 19.]
# [21. 22. 23. 24.]]

print(data.where(np.isnan(data_lost),0)[0,-4:,-4:].values)

最后一行的预期输出：保留所有 np.isnan(data_lost) 为 True 的值并将 rest 设置为零

[[ 6.  0.  0.  9.]
[ 0.  0.  0. 14.]
[16.  0.  0.  0.]
[ 0. 22.  0. 24.]]

最后一行的观察输出：将所有 np.isnan(data_lost) 为 True 的值设置为 nan 并将 rest 设置为零

[[nan  0.  0. nan]
[ 0.  0.  0. nan]
[nan  0.  0.  0.]
[ 0. nan  0. nan]]

非常感谢任何有关如何获得预期结果同时仍能够将我的（原本大得多的）数据分成块的帮助。

Answer 1

实际上并没有 "deep copying" dask 数组的概念。 Dask 假定您应用于 dask 数组的所有内容都是纯函数（尽管这不是直接强制执行的），因此如果您将变异函数映射到 dask 数组的块上，您将依赖未定义的行为。

解决方法是在应用函数内进行复制，例如

def set_fraction_randomly_to_nan(data, frac_missing):
    np.random.seed(0)
    data = data.copy()
    data[np.random.rand(*data.shape) < frac_missing] = np.nan
    return data

xr.DataArray.where 使用多个 dask 块时将有效点设置为 nan

xr.DataArray.where sets valid points to nan when using several dask chunks

python

dask

python-xarray