test_basic.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549
  1. import queue
  2. import threading
  3. import multiprocessing
  4. import numpy as np
  5. import pytest
  6. from numpy.random import random
  7. from numpy.testing import assert_array_almost_equal, assert_allclose
  8. from pytest import raises as assert_raises
  9. import scipy.fft as fft
  10. from scipy._lib._array_api import (
  11. is_numpy, xp_size, xp_assert_close, xp_assert_equal, make_xp_test_case,
  12. make_xp_pytest_param
  13. )
  14. lazy_xp_modules = [fft]
  15. skip_xp_backends = pytest.mark.skip_xp_backends
  16. # Expected input dtypes. Note that `scipy.fft` is more flexible for numpy,
  17. # but for C2C transforms like `fft.fft`, the array API standard only mandates
  18. # that complex dtypes should work, float32/float64 aren't guaranteed to.
  19. def get_expected_input_dtype(func, xp):
  20. # use __name__ so that `lazy_xp_function` doesn't break things
  21. if func.__name__ in ["fft", "fftn", "fft2", "ifft", "ifftn", "ifft2", "hfft",
  22. "hfftn", "hfft2", "irfft", "irfftn", "irfft2"]:
  23. dtype = xp.complex128
  24. elif func.__name__ in ["rfft", "rfftn", "rfft2", "ihfft", "ihfftn", "ihfft2"]:
  25. dtype = xp.float64
  26. else:
  27. raise ValueError(f'Unknown FFT function: {func}')
  28. return dtype
  29. def fft1(x):
  30. L = len(x)
  31. phase = -2j*np.pi*(np.arange(L)/float(L))
  32. phase = np.arange(L).reshape(-1, 1) * phase
  33. return np.sum(x*np.exp(phase), axis=1)
  34. class TestFFT:
  35. @make_xp_test_case(fft.ifft, fft.fft, fft.rfft, fft.irfft)
  36. def test_identity(self, xp):
  37. maxlen = 512
  38. x = xp.asarray(random(maxlen) + 1j*random(maxlen))
  39. xr = xp.asarray(random(maxlen))
  40. # Check some powers of 2 and some primes
  41. for i in [1, 2, 16, 128, 512, 53, 149, 281, 397]:
  42. xp_assert_close(fft.ifft(fft.fft(x[0:i])), x[0:i])
  43. xp_assert_close(fft.irfft(fft.rfft(xr[0:i]), i), xr[0:i])
  44. @skip_xp_backends(np_only=True, reason='significant overhead for some backends')
  45. def test_identity_extensive(self, xp):
  46. maxlen = 512
  47. x = xp.asarray(random(maxlen) + 1j*random(maxlen))
  48. xr = xp.asarray(random(maxlen))
  49. for i in range(1, maxlen):
  50. xp_assert_close(fft.ifft(fft.fft(x[0:i])), x[0:i])
  51. xp_assert_close(fft.irfft(fft.rfft(xr[0:i]), i), xr[0:i])
  52. @make_xp_test_case(fft.fft)
  53. def test_fft(self, xp):
  54. x = random(30) + 1j*random(30)
  55. expect = xp.asarray(fft1(x))
  56. x = xp.asarray(x)
  57. xp_assert_close(fft.fft(x), expect)
  58. xp_assert_close(fft.fft(x, norm="backward"), expect)
  59. xp_assert_close(fft.fft(x, norm="ortho"),
  60. expect / xp.sqrt(xp.asarray(30, dtype=xp.float64)),)
  61. xp_assert_close(fft.fft(x, norm="forward"), expect / 30)
  62. @skip_xp_backends(np_only=True, reason='some backends allow `n=0`')
  63. def test_fft_n(self, xp):
  64. x = xp.asarray([1, 2, 3], dtype=xp.complex128)
  65. assert_raises(ValueError, fft.fft, x, 0)
  66. @make_xp_test_case(fft.fft, fft.ifft)
  67. def test_ifft(self, xp):
  68. x = xp.asarray(random(30) + 1j*random(30))
  69. xp_assert_close(fft.ifft(fft.fft(x)), x)
  70. for norm in ["backward", "ortho", "forward"]:
  71. xp_assert_close(fft.ifft(fft.fft(x, norm=norm), norm=norm), x)
  72. @make_xp_test_case(fft.fft, fft.fft2)
  73. def test_fft2(self, xp):
  74. x = xp.asarray(random((30, 20)) + 1j*random((30, 20)))
  75. expect = fft.fft(fft.fft(x, axis=1), axis=0)
  76. xp_assert_close(fft.fft2(x), expect)
  77. xp_assert_close(fft.fft2(x, norm="backward"), expect)
  78. xp_assert_close(fft.fft2(x, norm="ortho"),
  79. expect / xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)))
  80. xp_assert_close(fft.fft2(x, norm="forward"), expect / (30 * 20))
  81. @make_xp_test_case(fft.ifft, fft.ifft2)
  82. def test_ifft2(self, xp):
  83. x = xp.asarray(random((30, 20)) + 1j*random((30, 20)))
  84. expect = fft.ifft(fft.ifft(x, axis=1), axis=0)
  85. xp_assert_close(fft.ifft2(x), expect)
  86. xp_assert_close(fft.ifft2(x, norm="backward"), expect)
  87. xp_assert_close(fft.ifft2(x, norm="ortho"),
  88. expect * xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)))
  89. xp_assert_close(fft.ifft2(x, norm="forward"), expect * (30 * 20))
  90. @make_xp_test_case(fft.fft, fft.fftn)
  91. def test_fftn(self, xp):
  92. x = xp.asarray(random((30, 20, 10)) + 1j*random((30, 20, 10)))
  93. expect = fft.fft(fft.fft(fft.fft(x, axis=2), axis=1), axis=0)
  94. xp_assert_close(fft.fftn(x), expect)
  95. xp_assert_close(fft.fftn(x, norm="backward"), expect)
  96. xp_assert_close(fft.fftn(x, norm="ortho"),
  97. expect / xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64)))
  98. xp_assert_close(fft.fftn(x, norm="forward"), expect / (30 * 20 * 10))
  99. @make_xp_test_case(fft.ifft, fft.ifftn)
  100. def test_ifftn(self, xp):
  101. x = xp.asarray(random((30, 20, 10)) + 1j*random((30, 20, 10)))
  102. expect = fft.ifft(fft.ifft(fft.ifft(x, axis=2), axis=1), axis=0)
  103. xp_assert_close(fft.ifftn(x), expect, rtol=1e-7)
  104. xp_assert_close(fft.ifftn(x, norm="backward"), expect, rtol=1e-7)
  105. xp_assert_close(
  106. fft.ifftn(x, norm="ortho"),
  107. fft.ifftn(x) * xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64))
  108. )
  109. xp_assert_close(fft.ifftn(x, norm="forward"),
  110. expect * (30 * 20 * 10),
  111. rtol=1e-7)
  112. @make_xp_test_case(fft.fft, fft.rfft)
  113. def test_rfft(self, xp):
  114. x = xp.asarray(random(29), dtype=xp.float64)
  115. for n in [xp_size(x), 2*xp_size(x)]:
  116. for norm in [None, "backward", "ortho", "forward"]:
  117. xp_assert_close(fft.rfft(x, n=n, norm=norm),
  118. fft.fft(xp.asarray(x, dtype=xp.complex128),
  119. n=n, norm=norm)[:(n//2 + 1)])
  120. xp_assert_close(
  121. fft.rfft(x, n=n, norm="ortho"),
  122. fft.rfft(x, n=n) / xp.sqrt(xp.asarray(n, dtype=xp.float64))
  123. )
  124. @make_xp_test_case(fft.irfft, fft.rfft)
  125. def test_irfft(self, xp):
  126. x = xp.asarray(random(30))
  127. xp_assert_close(fft.irfft(fft.rfft(x)), x)
  128. for norm in ["backward", "ortho", "forward"]:
  129. xp_assert_close(fft.irfft(fft.rfft(x, norm=norm), norm=norm), x)
  130. @make_xp_test_case(fft.rfft2)
  131. def test_rfft2(self, xp):
  132. x = xp.asarray(random((30, 20)), dtype=xp.float64)
  133. expect = fft.fft2(xp.asarray(x, dtype=xp.complex128))[:, :11]
  134. xp_assert_close(fft.rfft2(x), expect)
  135. xp_assert_close(fft.rfft2(x, norm="backward"), expect)
  136. xp_assert_close(fft.rfft2(x, norm="ortho"),
  137. expect / xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)))
  138. xp_assert_close(fft.rfft2(x, norm="forward"), expect / (30 * 20))
  139. @make_xp_test_case(fft.rfft2, fft.irfft2)
  140. def test_irfft2(self, xp):
  141. x = xp.asarray(random((30, 20)))
  142. xp_assert_close(fft.irfft2(fft.rfft2(x)), x)
  143. for norm in ["backward", "ortho", "forward"]:
  144. xp_assert_close(fft.irfft2(fft.rfft2(x, norm=norm), norm=norm), x)
  145. @make_xp_test_case(fft.fftn, fft.rfftn)
  146. def test_rfftn(self, xp):
  147. x = xp.asarray(random((30, 20, 10)), dtype=xp.float64)
  148. expect = fft.fftn(xp.asarray(x, dtype=xp.complex128))[:, :, :6]
  149. xp_assert_close(fft.rfftn(x), expect)
  150. xp_assert_close(fft.rfftn(x, norm="backward"), expect)
  151. xp_assert_close(fft.rfftn(x, norm="ortho"),
  152. expect / xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64)))
  153. xp_assert_close(fft.rfftn(x, norm="forward"), expect / (30 * 20 * 10))
  154. @make_xp_test_case(fft.irfftn, fft.rfftn)
  155. def test_irfftn(self, xp):
  156. x = xp.asarray(random((30, 20, 10)))
  157. xp_assert_close(fft.irfftn(fft.rfftn(x)), x)
  158. for norm in ["backward", "ortho", "forward"]:
  159. xp_assert_close(fft.irfftn(fft.rfftn(x, norm=norm), norm=norm), x)
  160. @make_xp_test_case(fft.hfft, fft.fft)
  161. def test_hfft(self, xp):
  162. x = random(14) + 1j*random(14)
  163. x_herm = np.concatenate((random(1), x, random(1)))
  164. x = np.concatenate((x_herm, x[::-1].conj()))
  165. x = xp.asarray(x)
  166. x_herm = xp.asarray(x_herm)
  167. expect = xp.real(fft.fft(x))
  168. xp_assert_close(fft.hfft(x_herm), expect)
  169. xp_assert_close(fft.hfft(x_herm, norm="backward"), expect)
  170. xp_assert_close(fft.hfft(x_herm, norm="ortho"),
  171. expect / xp.sqrt(xp.asarray(30, dtype=xp.float64)))
  172. xp_assert_close(fft.hfft(x_herm, norm="forward"), expect / 30)
  173. @make_xp_test_case(fft.hfft, fft.ihfft)
  174. def test_ihfft(self, xp):
  175. x = random(14) + 1j*random(14)
  176. x_herm = np.concatenate((random(1), x, random(1)))
  177. x = np.concatenate((x_herm, x[::-1].conj()))
  178. x = xp.asarray(x)
  179. x_herm = xp.asarray(x_herm)
  180. xp_assert_close(fft.ihfft(fft.hfft(x_herm)), x_herm)
  181. for norm in ["backward", "ortho", "forward"]:
  182. xp_assert_close(fft.ihfft(fft.hfft(x_herm, norm=norm), norm=norm), x_herm)
  183. @make_xp_test_case(fft.hfft2, fft.ihfft2)
  184. def test_hfft2(self, xp):
  185. x = xp.asarray(random((30, 20)))
  186. xp_assert_close(fft.hfft2(fft.ihfft2(x)), x)
  187. for norm in ["backward", "ortho", "forward"]:
  188. xp_assert_close(fft.hfft2(fft.ihfft2(x, norm=norm), norm=norm), x)
  189. @make_xp_test_case(fft.ifft2)
  190. def test_ihfft2(self, xp):
  191. x = xp.asarray(random((30, 20)), dtype=xp.float64)
  192. expect = fft.ifft2(xp.asarray(x, dtype=xp.complex128))[:, :11]
  193. xp_assert_close(fft.ihfft2(x), expect)
  194. xp_assert_close(fft.ihfft2(x, norm="backward"), expect)
  195. xp_assert_close(
  196. fft.ihfft2(x, norm="ortho"),
  197. expect * xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64))
  198. )
  199. xp_assert_close(fft.ihfft2(x, norm="forward"), expect * (30 * 20))
  200. @make_xp_test_case(fft.hfftn, fft.ihfftn)
  201. def test_hfftn(self, xp):
  202. x = xp.asarray(random((30, 20, 10)))
  203. xp_assert_close(fft.hfftn(fft.ihfftn(x)), x)
  204. for norm in ["backward", "ortho", "forward"]:
  205. xp_assert_close(fft.hfftn(fft.ihfftn(x, norm=norm), norm=norm), x)
  206. @make_xp_test_case(fft.ifftn, fft.ihfftn)
  207. def test_ihfftn(self, xp):
  208. x = xp.asarray(random((30, 20, 10)), dtype=xp.float64)
  209. expect = fft.ifftn(xp.asarray(x, dtype=xp.complex128))[:, :, :6]
  210. xp_assert_close(expect, fft.ihfftn(x))
  211. xp_assert_close(expect, fft.ihfftn(x, norm="backward"))
  212. xp_assert_close(
  213. fft.ihfftn(x, norm="ortho"),
  214. expect * xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64))
  215. )
  216. xp_assert_close(fft.ihfftn(x, norm="forward"), expect * (30 * 20 * 10))
  217. def _check_axes(self, op, xp):
  218. dtype = get_expected_input_dtype(op, xp)
  219. x = xp.asarray(random((30, 20, 10)), dtype=dtype)
  220. axes = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)]
  221. for a in axes:
  222. op_tr = op(xp.permute_dims(x, axes=a))
  223. tr_op = xp.permute_dims(op(x, axes=a), axes=a)
  224. xp_assert_close(op_tr, tr_op)
  225. @pytest.mark.parametrize("op", [make_xp_pytest_param(fft.fftn),
  226. make_xp_pytest_param(fft.ifftn),
  227. make_xp_pytest_param(fft.rfftn),
  228. make_xp_pytest_param(fft.irfftn)])
  229. def test_axes_standard(self, op, xp):
  230. self._check_axes(op, xp)
  231. @pytest.mark.parametrize("op", [make_xp_pytest_param(fft.hfftn),
  232. make_xp_pytest_param(fft.ihfftn)])
  233. def test_axes_non_standard(self, op, xp):
  234. self._check_axes(op, xp)
  235. @pytest.mark.parametrize("op", [make_xp_pytest_param(fft.fftn),
  236. make_xp_pytest_param(fft.ifftn),
  237. make_xp_pytest_param(fft.rfftn),
  238. make_xp_pytest_param(fft.irfftn)])
  239. def test_axes_subset_with_shape_standard(self, op, xp):
  240. dtype = get_expected_input_dtype(op, xp)
  241. x = xp.asarray(random((16, 8, 4)), dtype=dtype)
  242. axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]
  243. for a in axes:
  244. # different shape on the first two axes
  245. shape = tuple([2*x.shape[ax] if ax in a[:2] else x.shape[ax]
  246. for ax in range(x.ndim)])
  247. # transform only the first two axes
  248. op_tr = op(xp.permute_dims(x, axes=a),
  249. s=shape[:2], axes=(0, 1))
  250. tr_op = xp.permute_dims(op(x, s=shape[:2], axes=a[:2]),
  251. axes=a)
  252. xp_assert_close(op_tr, tr_op)
  253. @pytest.mark.parametrize("op", [make_xp_pytest_param(fft.fft2),
  254. make_xp_pytest_param(fft.ifft2),
  255. make_xp_pytest_param(fft.rfft2),
  256. make_xp_pytest_param(fft.irfft2),
  257. make_xp_pytest_param(fft.hfft2),
  258. make_xp_pytest_param(fft.ihfft2),
  259. make_xp_pytest_param(fft.hfftn),
  260. make_xp_pytest_param(fft.ihfftn)])
  261. def test_axes_subset_with_shape_non_standard(self, op, xp):
  262. dtype = get_expected_input_dtype(op, xp)
  263. x = xp.asarray(random((16, 8, 4)), dtype=dtype)
  264. axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]
  265. for a in axes:
  266. # different shape on the first two axes
  267. shape = tuple([2*x.shape[ax] if ax in a[:2] else x.shape[ax]
  268. for ax in range(x.ndim)])
  269. # transform only the first two axes
  270. op_tr = op(xp.permute_dims(x, axes=a), s=shape[:2], axes=(0, 1))
  271. tr_op = xp.permute_dims(op(x, s=shape[:2], axes=a[:2]), axes=a)
  272. xp_assert_close(op_tr, tr_op)
  273. @make_xp_test_case(fft.rfft, fft.irfft, fft.ihfft, fft.hfft, fft.fft, fft.ifft)
  274. def test_all_1d_norm_preserving(self, xp):
  275. # verify that round-trip transforms are norm-preserving
  276. x = xp.asarray(random(30), dtype=xp.float64)
  277. x_norm = xp.linalg.vector_norm(x)
  278. n = xp_size(x) * 2
  279. func_pairs = [(fft.rfft, fft.irfft),
  280. # hfft: order so the first function takes x.size samples
  281. # (necessary for comparison to x_norm above)
  282. (fft.ihfft, fft.hfft),
  283. # functions that expect complex dtypes at the end
  284. (fft.fft, fft.ifft),
  285. ]
  286. for forw, back in func_pairs:
  287. if forw == fft.fft:
  288. x = xp.asarray(x, dtype=xp.complex128)
  289. x_norm = xp.linalg.vector_norm(x)
  290. for n in [xp_size(x), 2*xp_size(x)]:
  291. for norm in ['backward', 'ortho', 'forward']:
  292. tmp = forw(x, n=n, norm=norm)
  293. tmp = back(tmp, n=n, norm=norm)
  294. xp_assert_close(xp.linalg.vector_norm(tmp), x_norm)
  295. @pytest.mark.parametrize("dtype", [np.float16, np.longdouble])
  296. def test_dtypes_nonstandard(self, dtype):
  297. x = random(30).astype(dtype)
  298. out_dtypes = {np.float16: np.complex64, np.longdouble: np.clongdouble}
  299. x_complex = x.astype(out_dtypes[dtype])
  300. res_fft = fft.ifft(fft.fft(x))
  301. res_rfft = fft.irfft(fft.rfft(x))
  302. res_hfft = fft.hfft(fft.ihfft(x), x.shape[0])
  303. # Check both numerical results and exact dtype matches
  304. assert_array_almost_equal(res_fft, x_complex)
  305. assert_array_almost_equal(res_rfft, x)
  306. assert_array_almost_equal(res_hfft, x)
  307. assert res_fft.dtype == x_complex.dtype
  308. assert res_rfft.dtype == np.result_type(np.float32, x.dtype)
  309. assert res_hfft.dtype == np.result_type(np.float32, x.dtype)
  310. @make_xp_test_case(fft.irfft, fft.rfft)
  311. @pytest.mark.parametrize("dtype", ["float32", "float64"])
  312. def test_dtypes_real(self, dtype, xp):
  313. x = xp.asarray(random(30), dtype=getattr(xp, dtype))
  314. res_rfft = fft.irfft(fft.rfft(x))
  315. res_hfft = fft.hfft(fft.ihfft(x), x.shape[0])
  316. # Check both numerical results and exact dtype matches
  317. xp_assert_close(res_rfft, x)
  318. xp_assert_close(res_hfft, x)
  319. @make_xp_test_case(fft.fft, fft.ifft)
  320. @pytest.mark.parametrize("dtype", ["complex64", "complex128"])
  321. def test_dtypes_complex(self, dtype, xp):
  322. rng = np.random.default_rng(1234)
  323. x = xp.asarray(rng.random(30), dtype=getattr(xp, dtype))
  324. res_fft = fft.ifft(fft.fft(x))
  325. # Check both numerical results and exact dtype matches
  326. xp_assert_close(res_fft, x)
  327. @pytest.mark.parametrize("op", [fft.fft, fft.ifft,
  328. fft.fft2, fft.ifft2,
  329. fft.fftn, fft.ifftn,
  330. fft.rfft, fft.irfft,
  331. fft.rfft2, fft.irfft2,
  332. fft.rfftn, fft.irfftn,
  333. fft.hfft, fft.ihfft,
  334. fft.hfft2, fft.ihfft2,
  335. fft.hfftn, fft.ihfftn,])
  336. def test_array_like(self, op):
  337. x = [[[1.0, 1.0], [1.0, 1.0]],
  338. [[1.0, 1.0], [1.0, 1.0]],
  339. [[1.0, 1.0], [1.0, 1.0]]]
  340. xp_assert_close(op(x), op(np.asarray(x)))
  341. @pytest.mark.parametrize(
  342. "dtype",
  343. [np.float32, np.float64, np.longdouble,
  344. np.complex64, np.complex128, np.clongdouble])
  345. @pytest.mark.parametrize("order", ["F", 'non-contiguous'])
  346. @pytest.mark.parametrize(
  347. "fft",
  348. [fft.fft, fft.fft2, fft.fftn, fft.ifft, fft.ifft2, fft.ifftn])
  349. def test_fft_with_order(dtype, order, fft):
  350. # Check that FFT/IFFT produces identical results for C, Fortran and
  351. # non contiguous arrays
  352. rng = np.random.RandomState(42)
  353. X = rng.rand(8, 7, 13).astype(dtype, copy=False)
  354. if order == 'F':
  355. Y = np.asfortranarray(X)
  356. else:
  357. # Make a non contiguous array
  358. Y = X[::-1]
  359. X = np.ascontiguousarray(X[::-1])
  360. if fft.__name__.endswith('fft'):
  361. for axis in range(3):
  362. X_res = fft(X, axis=axis)
  363. Y_res = fft(Y, axis=axis)
  364. assert_array_almost_equal(X_res, Y_res)
  365. elif fft.__name__.endswith(('fft2', 'fftn')):
  366. axes = [(0, 1), (1, 2), (0, 2)]
  367. if fft.__name__.endswith('fftn'):
  368. axes.extend([(0,), (1,), (2,), None])
  369. for ax in axes:
  370. X_res = fft(X, axes=ax)
  371. Y_res = fft(Y, axes=ax)
  372. assert_array_almost_equal(X_res, Y_res)
  373. else:
  374. raise ValueError
  375. @skip_xp_backends(cpu_only=True)
  376. class TestFFTThreadSafe:
  377. threads = 16
  378. input_shape = (800, 200)
  379. def _test_mtsame(self, func, *args, xp=None):
  380. def worker(args, q):
  381. q.put(func(*args))
  382. q = queue.Queue()
  383. expected = func(*args)
  384. # Spin off a bunch of threads to call the same function simultaneously
  385. t = [threading.Thread(target=worker, args=(args, q))
  386. for i in range(self.threads)]
  387. [x.start() for x in t]
  388. [x.join() for x in t]
  389. # Make sure all threads returned the correct value
  390. for i in range(self.threads):
  391. xp_assert_equal(
  392. q.get(timeout=5), expected,
  393. err_msg='Function returned wrong value in multithreaded context'
  394. )
  395. @make_xp_test_case(fft.fft)
  396. def test_fft(self, xp):
  397. a = xp.ones(self.input_shape, dtype=xp.complex128)
  398. self._test_mtsame(fft.fft, a, xp=xp)
  399. @make_xp_test_case(fft.ifft)
  400. def test_ifft(self, xp):
  401. a = xp.full(self.input_shape, 1+0j)
  402. self._test_mtsame(fft.ifft, a, xp=xp)
  403. @make_xp_test_case(fft.rfft)
  404. def test_rfft(self, xp):
  405. a = xp.ones(self.input_shape)
  406. self._test_mtsame(fft.rfft, a, xp=xp)
  407. @make_xp_test_case(fft.irfft)
  408. def test_irfft(self, xp):
  409. a = xp.full(self.input_shape, 1+0j)
  410. self._test_mtsame(fft.irfft, a, xp=xp)
  411. @make_xp_test_case(fft.hfft)
  412. def test_hfft(self, xp):
  413. a = xp.ones(self.input_shape, dtype=xp.complex64)
  414. self._test_mtsame(fft.hfft, a, xp=xp)
  415. @make_xp_test_case(fft.ihfft)
  416. def test_ihfft(self, xp):
  417. a = xp.ones(self.input_shape)
  418. self._test_mtsame(fft.ihfft, a, xp=xp)
  419. @pytest.mark.parametrize("func", [fft.fft, fft.ifft, fft.rfft, fft.irfft])
  420. def test_multiprocess(func):
  421. # Test that fft still works after fork (gh-10422)
  422. with multiprocessing.Pool(2) as p:
  423. res = p.map(func, [np.ones(100) for _ in range(4)])
  424. expect = func(np.ones(100))
  425. for x in res:
  426. assert_allclose(x, expect)
  427. @make_xp_test_case(fft.irfftn)
  428. class TestIRFFTN:
  429. def test_not_last_axis_success(self, xp):
  430. ar, ai = np.random.random((2, 16, 8, 32))
  431. a = ar + 1j*ai
  432. a = xp.asarray(a)
  433. axes = (-2,)
  434. # Should not raise error
  435. fft.irfftn(a, axes=axes)
  436. @pytest.mark.parametrize("func", [make_xp_pytest_param(fft.fft),
  437. make_xp_pytest_param(fft.ifft),
  438. make_xp_pytest_param(fft.rfft),
  439. make_xp_pytest_param(fft.irfft),
  440. make_xp_pytest_param(fft.fftn),
  441. make_xp_pytest_param(fft.ifftn),
  442. make_xp_pytest_param(fft.rfftn),
  443. make_xp_pytest_param(fft.irfftn),
  444. make_xp_pytest_param(fft.hfft),
  445. make_xp_pytest_param(fft.ihfft)])
  446. def test_non_standard_params(func, xp):
  447. # use __name__ so that `lazy_xp_function` doesn't break things
  448. if func.__name__ in ["rfft", "rfftn", "ihfft"]:
  449. dtype = xp.float64
  450. else:
  451. dtype = xp.complex128
  452. x = xp.asarray([1, 2, 3], dtype=dtype)
  453. # func(x) should not raise an exception
  454. func(x)
  455. if is_numpy(xp):
  456. func(x, workers=2)
  457. else:
  458. assert_raises(ValueError, func, x, workers=2)
  459. # `plan` param is not tested since SciPy does not use it currently
  460. # but should be tested if it comes into use
  461. @pytest.mark.parametrize("dtype", ['float32', 'float64'])
  462. @pytest.mark.parametrize("func", [make_xp_pytest_param(fft.fft),
  463. make_xp_pytest_param(fft.ifft),
  464. make_xp_pytest_param(fft.irfft),
  465. make_xp_pytest_param(fft.fftn),
  466. make_xp_pytest_param(fft.ifftn),
  467. make_xp_pytest_param(fft.irfftn),
  468. make_xp_pytest_param(fft.hfft)])
  469. def test_real_input(func, dtype, xp):
  470. x = xp.asarray([1, 2, 3], dtype=getattr(xp, dtype))
  471. # func(x) should not raise an exception
  472. func(x)