随机 mask/set nan x% 的数据点很大 xarray.DataArray

Question

我有一个巨大的（约 20 亿个数据点）xarray.DataArray。我想随机删除（屏蔽或替换为 np.nan）给定百分比的数据，其中为 deletion/masking 选择每个数据点的概率在所有坐标中都是相同的。我可以将数组转换为 numpy.array，但我最好将其保存在 dask 块中以提高速度。

我的数据是这样的：

>> data
<xarray.DataArray 'stack-820860ba63bd07adc355885d96354267' (variable: 8, time: 228, latitude: 721, longitude: 1440)>
dask.array<stack, shape=(8, 228, 721, 1440), dtype=float64, chunksize=(1, 6, 721, 1440)>
Coordinates:
* latitude   (latitude) float32 90.0 89.75 89.5 89.25 89.0 88.75 88.5 ...
* variable   (variable) <U5 u'fal' u'swvl1' u'swvl3' u'e' u'swvl2' u'es' 
* longitude  (longitude) float32 0.0 0.25 0.5 0.75 1.0 1.25 1.5 1.75 2.0 
* time       (time) datetime64[ns] 2000-01-01 2000-02-01 2000-03-01 ...

我定义了

frac_missing = 0.2
k = int(frac_missing*data.size)

这是我已经尝试过的：

this solution works with np.ndindex but the np.ndindex object is converted to a list which is very slow. I tried circumventing the conversion and simply iterate over the np.ndindex object as described here and 但是遍历整个迭代器对于大约 20 亿个数据点来说很慢。
np.random.choice(data.stack(newdim=('latitude','variable','longitude','time')),k,replace=False) returns 所需的数据点子集，但未将它们设置为 nan

预期输出将是 xarray.DataArray，给定百分比的数据点设置为 np.nan 或被屏蔽，最好是相同的形状和相同的 dask 块。

Answer 1

user545424 的建议是一个很好的开始。为了不运行进入内存问题，您可以将它放在一个小的 user-defined 函数中并使用 apply_ufunc.

方法将其映射到 DataArray 上

import xarray as xr
import numpy as np

testdata = xr.DataArray(np.empty((100,1000,1000)), dims=['x','y','z'])

def set_random_fraction_to_nan(data):
    data[np.random.rand(*data.shape) < .8]=np.nan
    return data

# Set 80% of data randomly to nan
testdata = xr.apply_ufunc(set_random_fraction_to_nan, testdata, input_core_dims=[['x','y','z']],output_core_dims=[['x','y','z']], dask='parallelized')

有关包装自定义函数以使用 xarray 的更多说明，请参阅 here.

随机 mask/set nan x% 的数据点很大 xarray.DataArray

randomly mask/set nan x% of data points in huge xarray.DataArray

python

numpy

dask

python-xarray