from numpy import * def f(x): return 1./(1.+cos(x)**2) def g(x): return x**-12 - x**-6 # Task 4.2.1 def dft_forw(f): N = len(f) fhat = zeros(N, dtype=complex) fac = -2.0j*pi/N k = arange(N) for n in range(N): fhat[k] += f[n]*exp(fac*k*n) return fhat def dft_back(fhat): N = len(fhat) f = zeros(N, dtype=complex) fac = 2.0j*pi/N n = arange(N) for k in range(N): f[n] += fhat[k]*exp(fac*k*n) f /= N return f # Task 4.2.2 def rdft_forw(f): N = len(f) Nk = N/2+1 fhat = zeros(Nk, dtype=complex) fac = -2.0j*pi/N k = arange(Nk) for n in range(N): fhat[k] += f[n]*exp(fac*k*n) return fhat def rdft_back(fhat, N=None): Nk = len(fhat) if N is None: N = 2*(Nk-1) f = zeros(N, dtype=complex) fac = 2.0j*pi/N n = arange(N) f[n] = fhat[0] for k in range(1,Nk): f[n] += 2.0*(fhat[k]*exp(fac*k*n)).real f /= N return f # Task 4.3 def fft_forw(f): N = len(f) if N == 1: return f else: fhat = zeros(N, dtype=complex) fac = -2.0j*pi/N g = fft_forw(f[0::2]) u = fft_forw(f[1::2]) k = arange(N/2-1) fhat[k] = g[k] + u[k]*exp(fac*k) fhat[k+N/2] = g[k] - u[k]*exp(fac*k) return fhat def fft_back(f): N = len(f) if N == 1: return f else: fhat = zeros(N, dtype=complex) fac = 2.0j*pi/N g = fft_forw(f[0::2]) u = fft_forw(f[1::2]) k = arange(N/2-1) fhat[k] = g[k] + u[k]*exp(fac*k) fhat[k+N/2] = g[k] - u[k]*exp(fac*k) return fhat def test_numpy(N): fft.ifft(fft.fft(f(linspace(-5,5,N)))) def test_dft(N): dft_back(dft_forw(f(linspace(-5,5,N)))) def test_fft(N): fft_back(fft_forw(f(linspace(-5,5,N))))