108 lines
2.6 KiB
Python
Executable File
108 lines
2.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import numpy as np
|
|
|
|
# Example Functions
|
|
def func1(x):
|
|
return np.exp(20j * x) + np.exp(40j * x )
|
|
|
|
def func2(x):
|
|
return np.exp( 5j * x**2 )
|
|
|
|
####################
|
|
# Fourier Transforms
|
|
####################
|
|
def DFT(yk):
|
|
N = np.size(yk)
|
|
xk = 2*np.pi/N*np.arange(N)
|
|
|
|
beta = np.zeros(N, dtype=np.complex128)
|
|
for j in range(N):
|
|
beta[j] = np.dot(yk, np.exp(- 1j * j * xk ))
|
|
|
|
return beta
|
|
|
|
def FFT(yk):
|
|
N = np.size(yk)
|
|
|
|
if ( N%2 > 0 ):
|
|
raise ValueError("The data set does not consist of 2^M values:", N)
|
|
|
|
if N <= 2:
|
|
return DFT(yk)
|
|
|
|
else:
|
|
beta_even = FFT(yk[::2])
|
|
beta_odd = FFT(yk[1::2])
|
|
|
|
exp_terms = np.exp( -2j * np.pi * np.arange(0,N) / N )
|
|
|
|
beta_even_full = np.concatenate([beta_even, beta_even])
|
|
beta_odd_full = np.concatenate([beta_odd, beta_odd])
|
|
|
|
return beta_even_full + exp_terms * beta_odd_full
|
|
|
|
|
|
|
|
# Plotting functions
|
|
def plot_DFTs(N = 128, fft1 = DFT, fft2 = np.fft.fft, datafuncs=[func1,func2]):
|
|
from matplotlib import pyplot
|
|
|
|
j = np.arange(N)
|
|
xk = 2*np.pi/N * np.arange(N)
|
|
pyplot.subplot()
|
|
pyplot.title("Discrete Fourier Transform")
|
|
pyplot.xlabel("j")
|
|
pyplot.ylabel("|$\\beta_j$|")
|
|
|
|
# First Function
|
|
for i, dataf in enumerate(datafuncs):
|
|
pyplot.plot(j, np.abs(fft1(func1(xk))), label="{} {}".format(fft1.__name__, dataf.__name__), linewidth=3)
|
|
pyplot.plot(j, np.abs(fft2(func1(xk))), '--', label="{} {}".format(fft1.__name__, dataf.__name__))
|
|
|
|
|
|
pyplot.legend()
|
|
pyplot.show()
|
|
|
|
|
|
def time_FT_func(func, yk, *args, **kwargs):
|
|
import timeit
|
|
|
|
tOut = timeit.repeat(stmt=lambda: func(yk), *args, **kwargs)
|
|
|
|
return np.mean(tOut)
|
|
|
|
def plot_FT_func_timing(func = DFT, M_max = 4, repeat = 10, number = 5, show = True):
|
|
from matplotlib import pyplot
|
|
|
|
M = np.arange(1,M_max)
|
|
Ns = 2**M
|
|
|
|
times = np.zeros(M_max-1)
|
|
for i in range(np.size(times)):
|
|
xk = 2*np.pi/Ns[i] * np.arange(Ns[i])
|
|
yk = func1(xk)
|
|
|
|
times[i] = time_FT_func(func, yk, repeat=repeat,number=number)
|
|
|
|
pyplot.subplot()
|
|
pyplot.title("Evaluation Time for {}*{} DFTs for each N ".format(repeat, number))
|
|
pyplot.ylabel("Time (sec)")
|
|
pyplot.xlabel("$N^2$")
|
|
pyplot.ticklabel_format(scilimits=(0,0))
|
|
|
|
pyplot.plot(Ns**2, times, "o:", label=func.__name__)
|
|
pyplot.legend()
|
|
|
|
if show:
|
|
pyplot.show()
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
plot_DFTs()
|
|
plot_DFTs(fft1=DFT,fft2=FFT)
|
|
|
|
plot_FT_func_timing(M_max = 9, func=FFT, show = False) # 12 takes 10 seconds, 13 takes 50 seconds
|
|
plot_FT_func_timing(M_max = 9, func=DFT) # 12 takes 10 seconds, 13 takes 50 seconds
|