如何在Dask / xarray的轴上有效地计算值的第一个实例？

问题描述

我在每个网格框中都有一些输入的大气总水模型数据。我正在尝试根据此输入数据计算云顶高度；因此，对于每一列，我需要找到此输入数据大于阈值的最高实例。

我的输入数据是100 x 900 x 900的{{1}}。我的数据通过大块nz x ny x nx的dask加载到xarray中。传统上，我会这样：

100 x 50 x 50

但是，对于dask / numpy / xarray来说，这确实效率低下。不过，我一直在努力寻找替代者。我看到了将cloud_top_height = numpy.zeros((900,900) for x in range(0,nx): for y in range(0,ny): cloud_top_found = false for z in range(nz,-1): if cloud_val > threshold: cloud_top_height[x,y] = z cloud_top_found = true if not cloud_top_found: cloud_top_height = np.nan与3D布尔索引一起使用的各种建议，但是我认为这与我想要的相反，而且xarray仍然不支持3D布尔索引。

使用xarray / dask计算值大于阈值的轴上最大索引的最佳方法是什么？

解决方法

怎么样：

In [2]: da = xr.DataArray(np.random.rand(5,5,5),dims=list('abc'),coords=dict(c=range(5)))

In [3]: (
    ...:     da
    ...:     .where(lambda x: x>0.8)
    ...:     .idxmax(dim='c')
    ...: )
Out[3]:
<xarray.DataArray 'c' (a: 5,b: 5)>
array([[ 4.,2.,1.,1.],[nan,nan,0.,nan],[ 1.,2.],[ 2.,1.]])
Dimensions without coordinates: a,b

如果不是那样，您可以发表可复制的示例吗？

dask python python-xarray