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循环，因为我将使用大型数组。

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

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

1 个答案: