python - 将分块的 dask xarray 乘以掩码

Question

我有一个大型(> 100 GB)xarray 数据集保存天气预报数据(维度时间，预测步骤，纬度，经度，随着时间的推移，纬度和经度维度的 dask 块)并且想要计算平均天气(对于每次点)在不规则形状的区域(由维度为纬度和经度的二进制掩码数组定义)上。
这样做的天真方法是:

average_weather = weather.where(mask).mean(dim=('latitude', 'longitude'))

但是，对于大多数(纬度，经度)块，该区域中的掩码值都为零，因此无需加载该块。就我通过简要查看 xarray 和 dask 源可以看出，在加载块之前没有检查块的所有掩码值是否为零的优化，因此 naive 命令将使用大量不必要的数据传输和 CPU。

我确实看到可以使用 drop=True在 where命令将计算限制到掩码的边界框，但有可能比这做得更好吗？

Answer 1

默认使用 where 时，掩码为 False 的值替换为 NaN .如果您使用 drop=True关键字，它们将被完全删除。请注意，这可能会改变数据的形状，或导致意外 NaN值有时。

例如average_weather = weather.where(mask, drop=True).mean(dim=('latitude', 'longitude'))