Apply FFT transformation to a signal and find its top frequencies.

find_top_frequencies_using_fft.py

import heapq


def get_fft(signal, f_s=1.0):
    """Calculate FFT of the signal
    """
    fft = np.abs(np.fft.fft(signal))
    freq = np.fft.fftfreq(fft.size, 1.0 / f_s)
 
    half_n = np.ceil(fft.size / 2.0)
    fft_half = (2.0 / fft.size) * fft[:half_n]
    freq_half = freq[:half_n]

    return freq_half, fft_half


def get_top_freq(freq, amplitude, n=5):
    """Get top n frequencies by the amplitude.
    """
    local_max, = argrelextrema(amplitude, np.greater, order=15)
    ampl_max = amplitude[local_max]
    freq_max = freq[local_max]
    idx = heapq.nlargest(n, range(ampl_max.size), ampl_max.take)
    return np.array(zip(freq_max[idx], ampl_max[idx]))


# Calculate FFT transformation of the signal
# with frequency 1 sec.
f_s = 1.0
freq, fft = get_fft(signal, f_s)

# Visualize fft spectrum of the signal
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(freq, fft, 'b-', label='frequencies')
ax.legend(loc='best')

# Calculate top 10 frequencies of the signal
# based on the largest amplitudes in FFT spectrum.
top_freq = get_top_freq(freq, fft, n=10)
fq, ampl = top_freq.T

# Visualize top frequencies on the FFT spectrum.
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(freq, fft, 'b-', label='Frequencies')
ax.plot(fq, ampl, 'go', markersize = 5, label='Top freq')
ax.legend(loc='best')