FFT basics

fft_analysis.py

import numpy as np
import pylab

# make some random data
x = np.random.randn(15)

# get length of data
n = len(x)

# compute PSD (square of the abs of the fft values)
# Here I use the rfft, which simplifies this if you are taking
# the fft of real numbers only (we are)

# np.fft.fft works too, but you'll notice the output is redundant, 
# n/2 of it's values are exact copies of the other n/2. 
xhat = np.abs(np.fft.rfft(x))**2

# sample rate (in Hz, or recordings per second)
fps = 30.

# frequencies which correspond to what's in "xhat"
freqs = fps/n*np.arange(n/2+1.)

print freqs
print
print xhat

#==================================

# make some more random data
n = 100
fps = 100
x = np.random.randn(n)

# now ill add in two strong waveforms
# (2 Hz and 10 Hz), with amplitudes 4 and 6
t = np.arange(n)
y1 = 4*np.sin(3*t * 2.*np.pi/n)
y2 = 6*np.sin(10*t * 2.*np.pi/n)

data = x + y1 + y2

# get PSD, and freqs array
xhat= (np.abs(np.fft.rfft(data))**2)
freqs = fps/n*np.arange(n/2 + 1.)

pylab.subplot(211)
pylab.plot(data)

pylab.subplot(212)
# ignore the zero frequence
pylab.plot(freqs[1:],xhat[1:])
pylab.show()

# you should see peaks at 3 Hz, and 10 Hz in the bottom plot!

fft_filtering.py

import numpy as np
import pylab

# make some random data
x = np.random.randn(100)

# take the fft or rfft 
xhat = np.fft.rfft(x)

# simple low-pass filter (some low frequencies stay, everything
# else gets shrunk or zeroed out)
xhat[15:] = 0

# invert the fft
x_filt = np.fft.irfft(xhat)

# plot the result
pylab.plot(x, label="original")
pylab.plot(x_filt, linewidth=4, label="low-pass filtered")
pylab.legend()

pylab.show()