fft_comparison_tests.py

import numpy as np
import fftw3
import pyfftw
import multiprocessing
import matplotlib
import matplotlib.pyplot as pl
import time


def fft_comparison_tests(size=2048, dtype=np.complex128, byte_align=False):
    """ Compare speed and test the API of pyFFTW3 and PyFFTW
    which are, somewhat surprisingly, completely different independent modules"""

    

    test_array = np.ones( (size,size), dtype=dtype)
    test_array[size*3/8:size*5/8, size*3/8:size*5/8] = 1 # square aperture oversampling 2...

    for ncores in (1, multiprocessing.cpu_count()):

        pl.clf()
        for FFT_direction in ['forward']: #,'backward']:

            print "Array size: {1} x {1}\nDirection: {2}\ndtype: {3}\nncores: {4}".format(0, size, FFT_direction, dtype, ncores)
            print ""
            # Let's first deal with some planning to make sure wisdom is generated ahead of time.
            for i, fft_type in enumerate(['numpy','pyfftw3','pyfftw','pyfftw_noplan']):
                    # planning using PyFFTW3
                    p0 = time.time()
                    print "Now planning "+fft_type+" in the "+FFT_direction+" direction"
                    #if (test_array.shape, FFT_direction) not in _FFTW3_INIT.keys():
                    if fft_type=='numpy':
                        print "\tno planning required"
                    elif fft_type=='pyfftw3':
                        fftplan = fftw3.Plan(test_array.copy(), None, nthreads = ncores,direction=FFT_direction, flags=['measure'])
                    elif fft_type=='pyfftw_noplan':
                        print "\tno planning required"
                    else:
                        pyfftw.interfaces.cache.enable()
                        pyfftw.interfaces.cache.set_keepalive_time(30)

                        if byte_align: test_array = pyfftw.n_byte_align_empty( (size,size), 16, dtype=dtype)
                        test_array = pyfftw.interfaces.numpy_fft.fft2(test_array, overwrite_input=True, planner_effort='FFTW_MEASURE', threads=ncores)

                    p1 = time.time()
                    print "\tTime elapsed planning: {0:.4f} s".format(p1-p0)

            print ""

            # Now let's run some FFTs
            for i, fft_type in enumerate(['numpy','pyfftw3','pyfftw','pyfftw_noplan']):
                
                # display
                if fft_type == 'pyfftw' and byte_align:
                    test_array = pyfftw.n_byte_align_empty( (size,size), 16, dtype=dtype)
                    #output_array = pyfftw.n_byte_align_empty( (size,size), 16, dtype=dtype)
                    test_array[:,:] = 0

                else:
                    test_array = np.zeros( (size,size), dtype=np.complex128)
                test_array[size*3/8:size*5/8, size*3/8:size*5/8] = 1 # square aperture oversampling 2...
                pl.subplot(2,4, 1 + i)
                pl.imshow(np.abs(test_array), vmin=0, vmax=1)
                pl.title( "FFT type: {0:10s} input array".format(fft_type))
                pl.draw()


                # actual timed FFT section starts here:
                t0 = time.time()

                if fft_type=='numpy':
                    test_array = np.fft.fft2(test_array)
                elif fft_type=='pyfftw3':

                    fftplan = fftw3.Plan(test_array, None, nthreads = multiprocessing.cpu_count(),direction=FFT_direction, flags=['measure'])
                    fftplan.execute() # execute the plan
                elif fft_type=='pyfftw':
                    test_array = pyfftw.interfaces.numpy_fft.fft2(test_array, overwrite_input=True, planner_effort='FFTW_MEASURE', threads=ncores)
                elif fft_type=='pyfftw_noplan':
                    test_array = pyfftw.interfaces.numpy_fft.fftn(test_array, overwrite_input=True, planner_effort='FFTW_MEASURE', threads=ncores)
     
        
                t1 = time.time()

                if FFT_direction=='forward': test_array = np.fft.fftshift(test_array)

                # display
                t_elapsed = t1-t0
                summarytext = "FFT type: {0:10s}\tTime Elapsed: {3:.4f} s".format(fft_type, size, FFT_direction, t_elapsed)
                print summarytext

                pl.subplot(2,4, 1+i+4)



                psf = np.real( test_array * np.conjugate(test_array))
                norm=matplotlib.colors.LogNorm(vmin=psf.max()*1e-6, vmax=psf.max())

                cmap = matplotlib.cm.jet
                cmap.set_bad((0,0,0.5))

                imtoshow = np.abs(psf[size*3/8:size*5/8, size*3/8:size*5/8])+1e-7
                #imtoshow = psf[size*3/8:size*5/8, size*3/8:size*5/8]
                print "min(imtoshow):",np.min(imtoshow)
                if imtoshow.min() == imtoshow.max():
                    pl.imshow(imtoshow)
                else:
                    pl.imshow(imtoshow, norm=norm)
                pl.title(summarytext)
        pl.draw()
        pl.savefig('ncores%i_size%i_fft%s_comparison.png' % (ncores,size,FFT_direction),bbox_inches='tight')

                #stop()

if __name__ == '__main__':

    for size in [1024, 2048, 4096]:
        fft_comparison_tests(size=size)
        print "--------------------------------------------------------"

results.txt

Array size: 1024 x 1024
Direction: forward
dtype: <type 'numpy.complex128'>
ncores: 1

Now planning numpy in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s
Now planning pyfftw3 in the forward direction
	Time elapsed planning: 0.9599 s
Now planning pyfftw in the forward direction
	Time elapsed planning: 0.9168 s
Now planning pyfftw_noplan in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s

FFT type: numpy     	Time Elapsed: 0.0744 s
min(imtoshow): 1e-07
FFT type: pyfftw3   	Time Elapsed: 1.6931 s
min(imtoshow): 1e-07
FFT type: pyfftw    	Time Elapsed: 0.0472 s
min(imtoshow): 1e-07
FFT type: pyfftw_noplan	Time Elapsed: 0.0567 s
min(imtoshow): 1e-07
Array size: 1024 x 1024
Direction: forward
dtype: <type 'numpy.complex128'>
ncores: 16

Now planning numpy in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s
Now planning pyfftw3 in the forward direction
	Time elapsed planning: 0.0206 s
Now planning pyfftw in the forward direction
	Time elapsed planning: 1.1963 s
Now planning pyfftw_noplan in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s

FFT type: numpy     	Time Elapsed: 0.1294 s
min(imtoshow): 1e-07
FFT type: pyfftw3   	Time Elapsed: 0.0331 s
min(imtoshow): 1e-07
FFT type: pyfftw    	Time Elapsed: 0.0241 s
min(imtoshow): 1e-07
FFT type: pyfftw_noplan	Time Elapsed: 0.0409 s
min(imtoshow): 1e-07
--------------------------------------------------------
Array size: 2048 x 2048
Direction: forward
dtype: <type 'numpy.complex128'>
ncores: 1

Now planning numpy in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s
Now planning pyfftw3 in the forward direction
	Time elapsed planning: 4.2021 s
Now planning pyfftw in the forward direction
	Time elapsed planning: 3.6183 s
Now planning pyfftw_noplan in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s

FFT type: numpy     	Time Elapsed: 0.3888 s
min(imtoshow): 1e-07
FFT type: pyfftw3   	Time Elapsed: 6.1868 s
min(imtoshow): 1e-07
FFT type: pyfftw    	Time Elapsed: 0.2257 s
min(imtoshow): 1e-07
FFT type: pyfftw_noplan	Time Elapsed: 0.3075 s
min(imtoshow): 1e-07
Array size: 2048 x 2048
Direction: forward
dtype: <type 'numpy.complex128'>
ncores: 16

Now planning numpy in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s
Now planning pyfftw3 in the forward direction
	Time elapsed planning: 0.0743 s
Now planning pyfftw in the forward direction
	Time elapsed planning: 4.4689 s
Now planning pyfftw_noplan in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s

FFT type: numpy     	Time Elapsed: 0.4722 s
min(imtoshow): 1e-07
FFT type: pyfftw3   	Time Elapsed: 0.0718 s
min(imtoshow): 1e-07
FFT type: pyfftw    	Time Elapsed: 0.1765 s
min(imtoshow): 1e-07
FFT type: pyfftw_noplan	Time Elapsed: 0.2580 s
min(imtoshow): 1e-07
--------------------------------------------------------
Array size: 4096 x 4096
Direction: forward
dtype: <type 'numpy.complex128'>
ncores: 1

Now planning numpy in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s
Now planning pyfftw3 in the forward direction
	Time elapsed planning: 11.9428 s
Now planning pyfftw in the forward direction
	Time elapsed planning: 11.9802 s
Now planning pyfftw_noplan in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s

FFT type: numpy     	Time Elapsed: 2.2598 s
min(imtoshow): 1e-07
FFT type: pyfftw3   	Time Elapsed: 18.9230 s
min(imtoshow): 1e-07
FFT type: pyfftw    	Time Elapsed: 1.5436 s
min(imtoshow): 1e-07
FFT type: pyfftw_noplan	Time Elapsed: 1.5424 s
min(imtoshow): 1e-07
Array size: 4096 x 4096
Direction: forward
dtype: <type 'numpy.complex128'>
ncores: 16

Now planning numpy in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s
Now planning pyfftw3 in the forward direction
	Time elapsed planning: 0.2651 s
Now planning pyfftw in the forward direction
	Time elapsed planning: 16.0517 s
Now planning pyfftw_noplan in the forward direction
	no planning required
	Time elapsed planning: 0.0000 s

FFT type: numpy     	Time Elapsed: 2.0909 s
min(imtoshow): 1e-07
FFT type: pyfftw3   	Time Elapsed: 0.1479 s
min(imtoshow): 1e-07
FFT type: pyfftw    	Time Elapsed: 0.3140 s
min(imtoshow): 1e-07
FFT type: pyfftw_noplan	Time Elapsed: 0.9082 s
min(imtoshow): 1e-07
--------------------------------------------------------