numpy：在3D阵列中沿一个轴查找最大值和对应位置

Question

我正在寻找一种使用Numpy沿着3D数组的第一轴查找最大值和相应索引的有效方法。

上下文：我正在处理存储在3D数组z.shape() = (t, x, y)并使用以下方法沿t轴进行了傅立叶变换

spectrum = np.fft.fft(z, axis=0)
spectrum.shape(f, x, y)

在此， f是频率空间，具有与t相同的大小。

我需要构造一个数组peak_vals.shape() = (x, y) ，该数组为每个x坐标和y坐标保存转换后的数据的最大值。 此外，我需要存储与每个峰值相对应的频率： peak_freqs.shape() = (x, y)

我已经有一个函数来计算与傅立叶变换相对应的频率：

def freq_arr(t):
    """Determine the frequency space that a DFT of vector t would have"""
    samp_freq = 1/(t[1]-t[0])
    return np.arange(len(t)) * samp_freq / len(t)

为了找到沿f轴的峰，我想出了：

peak_vals = np.amax(spectrum, axis=0)

如何找到与这些峰值相对应的频率？ 我想构造一个函数

def peaks(freqs, spectrum):
    peak_vals = np.amax(spectrum, axis=0)
    # do something
    return peak_vals, peak_freqs

我要避免for -loops，因为我将使用大型数组。

Answer 1

def peaks(freqs, spectrum):
    peak_vals = np.amax(spectrum, axis=0)

    # get the indices with the max value
    peak_idx = np.argmax(spectrum, axis=0)

    # now use the indices to get the corresponding frequencies
    # peak_freqs is now 2d even though freqs is 1d
    peak_freqs = freqs[peak_idx]

    return peak_vals, peak_freqs