从多维数组中提取特征索引

Question

我有一个数组，表示三维空间中云水浓度的值。在云水浓度高于某个阈值的地方，我说我有云（见下面的横截面）。大部分领域是干燥的，但在大部分领域都有层积云，底部在 400 米左右。

我要做的是提取云底和云顶的 (x,y,z) 坐标。然后我想在代表风速垂直分量的不同三维数组上使用这些坐标来获得云底的上升气流。

我现在正在做的工作有效但速度很慢。我觉得一定有办法利用 NumPy 来加快速度。

这就是我现在正在做的事情：

# 3d array representing cloud water at a particular timestep t
qc = QC(t)

# get the coordinates where there is cloud
cloud_coords = argwhere( qc > qc_thresh )

# Arrays to hold the z values of cloud base (cb) and cloud top (ct)
zcb = zeros((nx,ny))
zct = zeros((nx,ny))

# Since each coordinate (x,y) will in general have multiple z values
# for cloud I have to loop over all (x,y) and
# pull out max and min height for each point (x,y)
for x in range(nx):
    # Pull out all the coordinates with a given x value
    xslice = cloud_coords[ where(cloud_coords[:,0] == x) ]

    for y in range(ny):       
        # for the given x value select a particular y value
        column = xslice[ where(xslice[:,1] == y) ]

        try:
            zcb[x,y] = min( column[:,2] )
            zct[x,y] = max( column[:,2] )
        except:
            # Because there may not be any cloud at all
            # (a "hole") we fill the array with an average value
            zcb[x,y] = mean(zcb[zcb.nonzero()])
            zct[x,y] = mean(zct[zct.nonzero()])


# Because I intend to use these as indices I need them to be ints
zcb = array(zcb, dtype='int')
zct = array(zct, dtype='int')

输出是一个二维数组，包含云底（和顶部）的 z 坐标

然后我在另一个数组上使用这些索引来获取云底的风速等变量：

wind = W(t)
j,i = meshgrid(arange(ny),arange(nx))
wind_base = wind[i,j,zcb]

我在模拟中的许多时间步都这样做，最慢的部分是所有 (x,y) 坐标上的 python 循环。任何有关使用 NumPy 更快地提取这些值的帮助将不胜感激！

Answer 1

您怀疑 numpy 可以很好地解决您的问题是正确的。实际上，您正在做多种低效率的事情，例如，最后使用np.array() 显式创建一个新数组，以及一个使用int 的dtype 的新数组，这是python 3 中的一个复杂对象。

您可以在几行向量化的 numpy 行中完成大部分工作。这个想法是找到云出现的索引（沿高度轴）或云结束的位置就足够了。我们可以使用numpy.argmax 以矢量化方式做到这一点。这确实是有效解决方案的核心：

import numpy as np
import matplotlib.pyplot as plt

# generate dummy data
qc_thresh = 0.6
nx,ny,nz = 400,400,100
qc = np.zeros((nx,ny,nz))
# insert random cloud layer
qc[...,50:80] = np.random.rand(nx,ny,30)
# insert holes in clouds for completeness
qc[np.random.randint(nx,size=2*nx),np.random.randint(ny,size=2*nx),:] = 0

def compute_cloud_boundaries():
    cloud_arr = qc > qc_thresh

    # find boundaries by making use of np.argmax returning first maximum
    zcb = np.argmax(cloud_arr,axis=-1)
    zct = nz - 1 - np.argmax(cloud_arr[...,::-1],axis=-1)

    # logical (nx,ny)-shaped array where there's a cloud
    cloud_inds = (zcb | (zct!=nz-1)).astype(bool)
    # this is short for `(zcb==0) | (zct!=nz-1)`

    # fill the rest with the mean
    zcb[np.logical_not(cloud_inds)] = zcb[cloud_inds].mean()
    zct[np.logical_not(cloud_inds)] = zct[cloud_inds].mean()

    return zcb,zct

我根据您的方法检查了上述内容（以及相应的小示例），它给出了完全相同的结果。正如我所说，这个想法是cloud_arr = qc > qc_thresh 是一个逻辑数组，告诉我们哪里的湿度大到足以形成云。然后我们沿着最后一个（高度）轴查看这个（基本上是二进制！）矩阵的最大值。调用np.argmax 将告诉我们每个平面二维索引的第一个（最底部）高度值。为了得到云的顶部，我们需要反转我们的数组 并从另一边做同样的事情（注意转换回结果索引）。反转数组会创建一个视图而不是副本，因此这也很有效。最后，我们修正没有云的点；代替更好的约束，我们检查argmax 返回的最高索引对应于边缘点的位置。考虑到真实的天气数据，我们可以确定最底部和最顶部的测量值不对应于云，因此这应该是一个安全的标准。

这是展示的虚拟数据的横截面：

上述400x400x100案例的非代表性时序：

In [24]: %timeit compute_cloud_boundaries()
10 loops, best of 3: 29.1 ms per loop

In [25]: %timeit orig() # original loopy version from the question
1 loop, best of 3: 9.37 s per loop

这似乎比速度提高了 300 多倍。当然，您的实际用例将是对这种方法的适当测试，但应该没问题。

---

至于索引步骤，您可以通过为索引使用开放网格并利用数组广播来节省一些内存。不必分配额外的(nx,ny) 形数组也可能会加快这一步：

wind = W(t)
i,j = np.ogrid[:nx,:ny]
wind_base = wind[i,j,zcb]

如您所见，np.ogrid 创建了一个形状为(nx,1) 和(1,ny) 的开放网格，它们一起广播相当于meshgrid 调用。