npybench.py

import timeit

import numpy as np
from numpy.fft import fftshift, ifftshift, fft2

x, y = np.linspace(-1, 1, 128), np.linspace(-1, 1, 128)
xx, yy = np.meshgrid(x, y)
rho, phi = np.sqrt(xx**2 + yy**2), np.arctan2(yy, xx)
phase_arr = rho ** 2 * np.sin(phi)
complex_arr = np.exp(1j * 2 * np.pi * phase_arr)
padded = np.pad(complex_arr, ((64, 64), (64, 64)), mode='constant')
fft_result = fftshift(fft2(ifftshift(padded)))


def bench_trig(rho, phi):
    print('trigonometry')
    print(timeit.repeat('rho ** 2 * np.sin(phi)', number=1000, globals=globals()))


def bench_realtocomplex(phase_err):
    print('real to complex')
    print(timeit.repeat('np.exp(1j * 2 * np.pi * phase_arr)', number=1000, globals=globals()))


def bench_fft(padded):
    print('fft')
    print(timeit.repeat('fftshift(fft2(ifftshift(padded)))', number=1000, globals=globals()))


bench_trig(rho, phi)
bench_realtocomplex(phase_arr)
bench_fft(padded)