Fast Fourier Transform Library in Python

myFFT.py

#################################################################################
#
# This is FFT library implemented in python.
# The Algorithm used is FFT, decimation in time, radix -2 algorithm
# can compute FFT of 1-d and 2-d lists, 1-d for 1-d signals , and 2-d for images
#
# by : Umanga Bista @ [email protected]
# 
#################################################################################

import cmath
import numpy as np
from math import log, ceil

def omega(p, q):
   ''' The omega term in DFT and IDFT formulas'''
   return cmath.exp((2.0 * cmath.pi * 1j * q) / p)

def pad(lst):
   '''padding the list to next nearest power of 2 as FFT implemented is radix 2'''
   k = 0
   while 2**k < len(lst):
      k += 1
   return np.concatenate((lst, ([0] * (2 ** k - len(lst)))))

def fft(x):
   ''' FFT of 1-d signals
   usage : X = fft(x)
   where input x = list containing sequences of a discrete time signals
   and output X = dft of x '''

   n = len(x)
   if n == 1:
      return x
   Feven, Fodd = fft(x[0::2]), fft(x[1::2])
   combined = [0] * n
   for m in xrange(n/2):
     combined[m] = Feven[m] + omega(n, -m) * Fodd[m]
     combined[m + n/2] = Feven[m] - omega(n, -m) * Fodd[m]
   return combined

def ifft(X):
   ''' IFFT of 1-d signals
   usage x = ifft(X) 
   unpadding must be done implicitly'''

   x = fft([x.conjugate() for x in X])
   return [x.conjugate()/len(X) for x in x]

def pad2(x):
   m, n = np.shape(x)
   M, N = 2 ** int(ceil(log(m, 2))), 2 ** int(ceil(log(n, 2)))
   F = np.zeros((M,N), dtype = x.dtype)
   F[0:m, 0:n] = x
   return F, m, n

def fft2(f):
   '''FFT of 2-d signals/images with padding
   usage X, m, n = fft2(x), where m and n are dimensions of original signal'''

   f, m, n = pad2(f)
   return np.transpose(fft(np.transpose(fft(f)))), m, n

def ifft2(F, m, n):
   ''' IFFT of 2-d signals
   usage x = ifft2(X, m, n) with unpaded, 
   where m and n are odimensions of original signal before padding'''

   f, M, N = fft2(np.conj(F))
   f = np.matrix(np.real(np.conj(f)))/(M*N)
   return f[0:m, 0:n]

def fftshift(F):
   ''' this shifts the centre of FFT of images/2-d signals'''
   M, N = F.shape
   R1, R2 = F[0: M/2, 0: N/2], F[M/2: M, 0: N/2]
   R3, R4 = F[0: M/2, N/2: N], F[M/2: M, N/2: N]
   sF = np.zeros(F.shape,dtype = F.dtype)
   sF[M/2: M, N/2: N], sF[0: M/2, 0: N/2] = R1, R4
   sF[M/2: M, 0: N/2], sF[0: M/2, N/2: N]= R3, R2
   return sF

myFFT_test.py

from myFFT import *
import numpy as np
print 'Testing for 1-d signals'
# Generating sin curve in range [-2p, 2pi] with 128 sample points
f = np.sin(np.linspace(-2*np.pi,2*np.pi,128))
# let us add some noise with mean 0.5 and sigma 0.75
f = f + 0.75 * np.random.rand(128) + 0.5
F = fft(f)
import pylab as plt
fig = plt.figure()
fig.add_subplot(311)
plt.plot(f)
plt.title('Original Signal')
fig.add_subplot(312)
plt.plot(np.log(np.abs(F[:64]) + 1))
plt.title('magnitude plot')
fig.add_subplot(313)
plt.plot(np.angle(F[:64]))
plt.title('Phase plot')
plt.show()

print '\ntesting for 2-d signals/images'
x = np.matrix([[1,2,1],[2,1,2],[0,1,1]])
X, m, n = fft2(x)
print '\nDFT is :'
print X
print '\nOriginal signal is :'
print ifft2(X, m, n)

Wouldn't it be more efficient to use numpy for the 1-d case?