Created
February 16, 2018 20:56
-
-
Save sbmueller/e46be793b9eedcbbd3ac345ca8939fae to your computer and use it in GitHub Desktop.
Comparison of different FFT implementations in Python
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#/usr/bin/env python3 | |
import numpy as np | |
from scipy import fftpack | |
import pyfftw | |
import unittest | |
import time | |
class Timer: | |
def __enter__(self): | |
self.start = time.clock() | |
return self | |
def __exit__(self, *args): | |
self.end = time.clock() | |
self.interval = self.end - self.start | |
print(self.interval) | |
class FFTComparer(unittest.TestCase): | |
""" | |
Compare different FFT implementations for Python. | |
A random 3D cube of complex data is used and the FFTs are performed | |
over the 2nd axis to avoid adjacent data access in memory. | |
""" | |
def setUp(self): | |
self.data = np.random.rand(213, 517, 17) + 1j*np.random.rand(213, 517, 17) | |
def test_ffts(self): | |
in_data = pyfftw.empty_aligned((213, 517, 17), dtype='complex128') | |
out_data = pyfftw.empty_aligned((213, 517, 17), dtype='complex128') | |
api = pyfftw.FFTW(in_data, out_data, flags=('FFTW_EXHAUSTIVE',), axes=(1,)) | |
print("numpy: ") | |
with Timer(): | |
X_np = np.fft.fft(self.data, axis=1) | |
print("fftpack: ") | |
with Timer(): | |
X_pck = fftpack.fft(self.data, axis=1) | |
print("pyfftw: ") | |
with Timer(): | |
in_data[:] = self.data | |
X_fftw = api() | |
self.assertEqual(X_np.all(), X_pck.all()) | |
self.assertEqual(X_np.all(), X_fftw.all()) | |
if __name__ == '__main__': | |
unittest.main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment