Chris@87: from __future__ import division, absolute_import, print_function
Chris@87: 
Chris@87: import numpy as np
Chris@87: from numpy.testing import TestCase, run_module_suite, assert_array_almost_equal
Chris@87: from numpy.testing import assert_array_equal
Chris@87: import threading
Chris@87: import sys
Chris@87: if sys.version_info[0] >= 3:
Chris@87:     import queue
Chris@87: else:
Chris@87:     import Queue as queue
Chris@87: 
Chris@87: 
Chris@87: def fft1(x):
Chris@87:     L = len(x)
Chris@87:     phase = -2j*np.pi*(np.arange(L)/float(L))
Chris@87:     phase = np.arange(L).reshape(-1, 1) * phase
Chris@87:     return np.sum(x*np.exp(phase), axis=1)
Chris@87: 
Chris@87: 
Chris@87: class TestFFTShift(TestCase):
Chris@87: 
Chris@87:     def test_fft_n(self):
Chris@87:         self.assertRaises(ValueError, np.fft.fft, [1, 2, 3], 0)
Chris@87: 
Chris@87: 
Chris@87: class TestFFT1D(TestCase):
Chris@87: 
Chris@87:     def test_basic(self):
Chris@87:         rand = np.random.random
Chris@87:         x = rand(30) + 1j*rand(30)
Chris@87:         assert_array_almost_equal(fft1(x), np.fft.fft(x))
Chris@87: 
Chris@87: 
Chris@87: class TestFFTThreadSafe(TestCase):
Chris@87:     threads = 16
Chris@87:     input_shape = (800, 200)
Chris@87: 
Chris@87:     def _test_mtsame(self, func, *args):
Chris@87:         def worker(args, q):
Chris@87:             q.put(func(*args))
Chris@87: 
Chris@87:         q = queue.Queue()
Chris@87:         expected = func(*args)
Chris@87: 
Chris@87:         # Spin off a bunch of threads to call the same function simultaneously
Chris@87:         t = [threading.Thread(target=worker, args=(args, q))
Chris@87:              for i in range(self.threads)]
Chris@87:         [x.start() for x in t]
Chris@87: 
Chris@87:         [x.join() for x in t]
Chris@87:         # Make sure all threads returned the correct value
Chris@87:         for i in range(self.threads):
Chris@87:             assert_array_equal(q.get(timeout=5), expected,
Chris@87:                 'Function returned wrong value in multithreaded context')
Chris@87: 
Chris@87:     def test_fft(self):
Chris@87:         a = np.ones(self.input_shape) * 1+0j
Chris@87:         self._test_mtsame(np.fft.fft, a)
Chris@87: 
Chris@87:     def test_ifft(self):
Chris@87:         a = np.ones(self.input_shape) * 1+0j
Chris@87:         self._test_mtsame(np.fft.ifft, a)
Chris@87: 
Chris@87:     def test_rfft(self):
Chris@87:         a = np.ones(self.input_shape)
Chris@87:         self._test_mtsame(np.fft.rfft, a)
Chris@87: 
Chris@87:     def test_irfft(self):
Chris@87:         a = np.ones(self.input_shape) * 1+0j
Chris@87:         self._test_mtsame(np.fft.irfft, a)
Chris@87: 
Chris@87: 
Chris@87: if __name__ == "__main__":
Chris@87:     run_module_suite()