FFT prime problem

fftprimeproblem.py

# https://github.com/fnielsen/everything
from everything import * 

import pyfftw.interfaces.scipy_fftpack


def is_prime(n):
    """https://stackoverflow.com/questions/18833759/"""
    if n % 2 == 0 and n > 2: 
        return False
    return all([bool(n % i) for i in range(3, int(math.sqrt(n)) + 1, 2)])


def profile(func):
    """https://stackoverflow.com/questions/5375624."""
    def wrap(*args, **kwargs):
        started_at = time.time()
        result = func(*args, **kwargs)
        print("%12s : %f" % (func.__name__, time.time() - started_at))
        return result
    return wrap


@profile
def czt(x, m=None, w=None, a=None):
    """https://math.stackexchange.com/questions/77118/"""
    # Translated from GNU Octave's czt.m
    n = len(x)
    if m is None: m = n
    if w is None: w = exp(-2j * pi / m)
    if a is None: a = 1
    chirp = w ** (arange(1 - n, max(m, n)) ** 2 / 2.0)
    N2 = int(2 ** ceil(log2(m + n - 1)))  # next power of 2
    xp = append(x * a ** -arange(n) * chirp[n - 1 : n + n - 1], zeros(N2 - n))
    ichirpp = append(1 / chirp[: m + n - 1], zeros(N2 - (m + n - 1)))
    r = ifft(fft(xp) * fft(ichirpp))
    return r[n - 1 : m + n - 1] * chirp[n - 1 : m + n - 1]


@profile
def padded_fft(x):
    axis = 0
    n_original = x.shape[axis]
    n_power_of_2 = 2 ** int(math.ceil(math.log(n_original, 2)))
    n_pad = n_power_of_2 - n_original
    z = np.zeros( (n_pad,) + x.shape[1:] )
    padded = np.concatenate((x, z), axis=axis)
    return scipy.fftpack.fft(padded, axis=axis)

@profile
def numpy_fft(x):
    return np.fft.fft(x)

@profile
def scipy_fft(x):
    return scipy.fftpack.fft(x)

@profile
def fftw_fft(x):
    return pyfftw.interfaces.scipy_fftpack.fft(x)



# Show some primes
print([n for n in range(20000, 22000) if is_prime(n)])

for n in [20000, 20011, 21803, 21804, 21997, 32768]:
    print("%d prime=%s" % (n, str(is_prime(n))))
    x = sin(linspace(0, 100, n) ** 2)
    y1 = padded_fft(x)
    y2 = numpy_fft(x)
    y3 = scipy_fft(x)
    y4 = czt(x)
    y5 = fftw_fft(x)
    print('-' * 60)



octave = """
x = sin(linspace(0, 100, 20000) .^ 2); 
tic; y = fft(x); toc
x = sin(linspace(0, 100, 20011) .^ 2); 
tic; y = fft(x); toc
x = sin(linspace(0, 100, 21803) .^ 2); 
tic; y = fft(x); toc
x = sin(linspace(0, 100, 21804) .^ 2); 
tic; y = fft(x); toc
x = sin(linspace(0, 100, 21997) .^ 2); 
tic; y = fft(x); toc
x = sin(linspace(0, 100, 32768) .^ 2); 
tic; y = fft(x); toc

"""