Source

pypy / site-packages / numpy / fft / test / test_fftpack_lite.py

#!/usr/bin/env python

try:
    import numpypy as np
except ImportError:
    import numpy as np

from random import random
from fftpack_lite import *
from test_helpers import fft_slow, assert_array_almost_equal

# Test FFTs directly through the fftpack wrapper

# real ffts
for n in [8, 11, 17]:
    print "Real", n,
    #a = np.random.random(n)
    a = np.array([random() for _ in xrange(n)])
    wrk = rffti(n)
    res = rfftf(a, wrk) 
    # slow version of fft
    tres = fft_slow(a)
    t = n//2 + 1
    assert_array_almost_equal(res, tres[:t])

    # inverse fft
    res = rfftb(res, wrk) / n
    assert_array_almost_equal(res, a)
    print "OK"

# complex ffts
for n in [8, 11, 17]:
    print "Complex", n, 
    #a = np.random.random(n) + 1j * np.random.random(n)
    a = np.array([random() for _ in xrange(n)]) + 1j * np.array([random() for _ in xrange(n)]) 
    wrk = cffti(n)
    res = cfftf(a, wrk)
    # slow version of fft
    tres = fft_slow(a)
    assert_array_almost_equal(res, tres)

    # inverse fft
    res = cfftb(res, wrk)  / n
    assert_array_almost_equal(res, a)
    print 'OK'