跟踪dim1的最大值，dim3中的索引不同

Question

我有一个三维数据的“立方体”，其中列中有一些峰值，或第一维。峰的索引可能根据检查的行而移动。第三个维度可能会做一些更复杂的事情，但现在可以被认为只是通过一些线性函数来缩放事物。

我想在第一个维度上找到最大值的索引，但要受到限制，每个行都选择z索引，使得列峰值最接近0.5。

这是一个样本图像，它是行中的平面，具有固定z的列：

这些数组有时会很大 - 例如21x11x200 float64s，所以我想对这个计算进行矢量化。用for循环编写，看起来像这样：

cols, rows, zs = data.shape
for i in range(rows):
    # for each field point, make an intermediate array that is 2D with focus,frequency dimensions
    arr = data[:,i,:]

    # compute the thru-focus max and find the peak closest to 0.5
    maxs = np.max(arr, axis=0)
    max_manip = np.abs(maxs-0.5)
    freq_idx = np.argmin(max_manip)

    # take the thru-focus slice that peaks closest to 0.5
    arr2 = data[:,i,freq_idx]
    focus_idx = np.argmax(arr2)
    print(focus_idx)

我的问题是我不知道如何将这些计算转换为向量操作。我要感谢任何帮助，谢谢！

Answer 1

我们只需要在那里使用axis param和相关的ufunc，这将引导我们使用矢量化解决方案，就像这样 -

# Get freq indices along all rows in one go
idx = np.abs(data.max(0)-0.5).argmin(1)

# Index into data with those and get the argmax indices
out = data[:,np.arange(data.shape[1]), idx].argmax(0)