Numpy FFT

numpy-04_fft.py

# -*- coding: utf-8 -*-

from numpy import *
from numpy.random import *
from numpy.fft import *
from pylab import *

N = 2**9 # FFT works better with signals whose length is power of 2
T = 10   # Length of time interval
t = linspace(0,T,N)
dt = t[1] - t[0]

# Frequencies
f1 = 5
f2 = 10
f3 = 20
A1 = 1
A2 = 1
A3 = 2

# Gaussian Noise
sigma = 0.5

# Signal
x = A1*sin(2*pi*f1*t) + A2*sin(2*pi*f2*t) + A3*sin(2*pi*f3*t) + sigma*randn(N)


# Exact calculation of the DFT
xt = zeros(N,dtype='complex')
for k in range(N):
    xt[k] = 0
    for m in range(N):
        xt[k] += x[m]*exp(-2*1j*pi*m*k/N)

# FFT
xt2 = fft(x)  # = xt[k]


# Plot : xt[k] corresponds to the frequency f = k/T = k/N*(1/dt)
subplot(2,1,1)
stem(arange(N/2)/(N*dt),abs(xt[:N/2]),linefmt='-b',markerfmt='.')
subplot(2,1,2)
stem(arange(N/2)/(N*dt),abs(xt2[:N/2]),linefmt='-b',markerfmt='.')
show()