python - 沿着 dim1 追踪 max，并在 dim3 中改变索引

Question

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

我想找到沿第一个维度的最大值的索引，但受到以下约束:对于每行，选择 z 索引以使列峰值最接近 0.5。

这是一个示例图像，它是具有固定 z 的行列平面:

这些数组有时会很大——比如 21x11x200 float64，所以我想向量化这个计算。用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

我们只需要使用带有相关ufunc的axis参数，这将引导我们得到一个矢量化的解决方案，就像这样 -

# 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)