FFT timing comparison

fft_comparison.py

import numpy as np
from numpy.fft import fft as nfft, ifft as nifft
from accelerate.mkl.fftpack import fft as mfft, ifft as mifft
from pyfftw.interfaces.numpy_fft import fft as wfft, ifft as wifft
import matplotlib.pyplot as plt
import time

start = 3
end = 5
diff = 4*(end-start)+1
n_iter = 20
channels = 8

lengths = np.around(np.logspace(start, end, diff)).astype(int)
ffts = [nfft, nifft, mfft, mifft, wfft, wifft]

results = np.zeros((len(ffts), lengths.size))

for ii, l in enumerate(lengths):
    x = np.random.randn(channels, l)
    for jj, fft in enumerate(ffts):
        out = fft(x)
        t0 = time.time()
        for kk in range(n_iter):
            out = fft(x)
        t1 = time.time()
        results[jj, ii] = (t1-t0)/float(n_iter)
        
names = ['mkl vs. numpy, fft', 'mkl vs. numpy, ifft', 'mkl vs. pyfftw, fft', 'mkl vs. pyfftw, ifft']

plt.figure()
plt.loglog(lengths, results[0]/results[2], 'o', label=names[0])
plt.loglog(lengths, results[1]/results[3], 'o', label=names[1])

plt.loglog(lengths, results[4]/results[2], 'o', label=names[2])
plt.loglog(lengths, results[5]/results[3], 'o', label=names[3])

plt.loglog([.5*(10**start), 2*(10**end)], [1, 1], '-', c='black')
plt.legend(loc='best')
plt.ylabel(' time/mkl (above 1 is good)')
plt.xlabel('input shape = ('+str(channels)+', x)')
plt.xlim([.75*(10**start), 1.5*(10**end)])
plt.show()