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

Question

我在每个网格框中输入了一些总水量的大气模型数据。我正试图根据这些输入数据计算云顶高度；所以对于每一列，我需要找到该输入数据大于阈值的最高实例。

对于nz x ny x nx，我的输入数据是100 x 900 x 900。我的数据通过带有100 x 50 x 50块的dask加载到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, 0, -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

然而，使用dask/numpy/xarray确实效率低下。不过，我一直在努力寻找替代者。我看到了各种建议，建议我将argmax与3D布尔索引结合使用，但我认为这与我想要的相反，而且xarray无论如何都不支持3D布尔索引。

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

Answer 1

怎么样：

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.,  1.],
[nan,  1., nan,  0., nan],
[ 1.,  1.,  2., nan,  1.],
[nan, nan,  2.,  1.,  2.],
[ 2.,  0., nan,  2.,  1.]])
Dimensions without coordinates: a, b

如果不是这样，你能发布一个可复制的例子吗？