Search code examples
pythonnumpysignal-processingfft

How to get the frequencies and corresponding amplitudes from the FFT of a signal?

I have the following 1.2Hz/9Hz/15Hz signal

import numpy as np

total_time = 5
sampling_frequency = 100

t = np.linspace(0, total_time, total_time * sampling_frequency, endpoint=False)
signal = np.sin(2 * np.pi * 1.2 * t) + 0.5*np.cos(2 * np.pi * 9 * t) + 0.75*np.sin(2 * np.pi * 15.0 * t)
plt.plot(t, signal)

enter image description here

and the discrete FFT

fft_spectrum = np.fft.rfft(signal)
fft_spectrum_abs = np.abs(fft_spectrum) * 2 / (total_time * sampling_frequency)

freq = np.fft.rfftfreq(signal.size, d=1./sampling_frequency)

plt.plot(freq, fft_spectrum_abs)
plt.xlabel("frequency, Hz")
plt.ylabel("amplitude, units")
plt.show()

enter image description here

How can I get the frequencies and amplitudes of the FFT?

I'm using

[
    (amp, freq) for amp, freq in sorted(zip(fft_spectrum_abs, freq), 
    key=lambda pair: pair[0]) if amp > 0.1
]
# [(0.4999999999999984, 9.0), (0.749999999999999, 15.0), (1.0, 1.2000000000000002)]

which is correct but seems a bit of a hack, is there any better way to approach the problem?

Solution

  • You can use scipy's scipy.signal.find_peaks. You want peaks with a height of at least 0.1, so you can specify that using the heights kwarg. When heights is specified, the properties dict that is returned will have a "peak_heights" key containing a numpy array of the peak heights (as the name suggests).

    from scipy.signal import find_peaks

peaks, properties = find_peaks(fft_spectrum_abs, height=0.1)
peak_freq = freq[peaks]
peak_amp = properties["peak_heights"]
print(peak_freq)  # [ 1.2  9.  15. ]
print(peak_amp)   # [1.   0.5  0.75]

    This solution might be more robust if you have messier/noisier data.