Comparison of different FFT implementations in Python

fft_comparer.py

#/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()