test_pocketfft.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589
  1. import queue
  2. import threading
  3. import pytest
  4. import numpy as np
  5. from numpy.random import random
  6. from numpy.testing import IS_WASM, assert_allclose, assert_array_equal, assert_raises
  7. def fft1(x):
  8. L = len(x)
  9. phase = -2j * np.pi * (np.arange(L) / L)
  10. phase = np.arange(L).reshape(-1, 1) * phase
  11. return np.sum(x * np.exp(phase), axis=1)
  12. class TestFFTShift:
  13. def test_fft_n(self):
  14. assert_raises(ValueError, np.fft.fft, [1, 2, 3], 0)
  15. class TestFFT1D:
  16. def test_identity(self):
  17. maxlen = 512
  18. x = random(maxlen) + 1j * random(maxlen)
  19. xr = random(maxlen)
  20. for i in range(1, maxlen):
  21. assert_allclose(np.fft.ifft(np.fft.fft(x[0:i])), x[0:i],
  22. atol=1e-12)
  23. assert_allclose(np.fft.irfft(np.fft.rfft(xr[0:i]), i),
  24. xr[0:i], atol=1e-12)
  25. @pytest.mark.parametrize("dtype", [np.single, np.double, np.longdouble])
  26. def test_identity_long_short(self, dtype):
  27. # Test with explicitly given number of points, both for n
  28. # smaller and for n larger than the input size.
  29. maxlen = 16
  30. atol = 5 * np.spacing(np.array(1., dtype=dtype))
  31. x = random(maxlen).astype(dtype) + 1j * random(maxlen).astype(dtype)
  32. xx = np.concatenate([x, np.zeros_like(x)])
  33. xr = random(maxlen).astype(dtype)
  34. xxr = np.concatenate([xr, np.zeros_like(xr)])
  35. for i in range(1, maxlen * 2):
  36. check_c = np.fft.ifft(np.fft.fft(x, n=i), n=i)
  37. assert check_c.real.dtype == dtype
  38. assert_allclose(check_c, xx[0:i], atol=atol, rtol=0)
  39. check_r = np.fft.irfft(np.fft.rfft(xr, n=i), n=i)
  40. assert check_r.dtype == dtype
  41. assert_allclose(check_r, xxr[0:i], atol=atol, rtol=0)
  42. @pytest.mark.parametrize("dtype", [np.single, np.double, np.longdouble])
  43. def test_identity_long_short_reversed(self, dtype):
  44. # Also test explicitly given number of points in reversed order.
  45. maxlen = 16
  46. atol = 6 * np.spacing(np.array(1., dtype=dtype))
  47. x = random(maxlen).astype(dtype) + 1j * random(maxlen).astype(dtype)
  48. xx = np.concatenate([x, np.zeros_like(x)])
  49. for i in range(1, maxlen * 2):
  50. check_via_c = np.fft.fft(np.fft.ifft(x, n=i), n=i)
  51. assert check_via_c.dtype == x.dtype
  52. assert_allclose(check_via_c, xx[0:i], atol=atol, rtol=0)
  53. # For irfft, we can neither recover the imaginary part of
  54. # the first element, nor the imaginary part of the last
  55. # element if npts is even. So, set to 0 for the comparison.
  56. y = x.copy()
  57. n = i // 2 + 1
  58. y.imag[0] = 0
  59. if i % 2 == 0:
  60. y.imag[n - 1:] = 0
  61. yy = np.concatenate([y, np.zeros_like(y)])
  62. check_via_r = np.fft.rfft(np.fft.irfft(x, n=i), n=i)
  63. assert check_via_r.dtype == x.dtype
  64. assert_allclose(check_via_r, yy[0:n], atol=atol, rtol=0)
  65. def test_fft(self):
  66. x = random(30) + 1j * random(30)
  67. assert_allclose(fft1(x), np.fft.fft(x), atol=1e-6)
  68. assert_allclose(fft1(x), np.fft.fft(x, norm="backward"), atol=1e-6)
  69. assert_allclose(fft1(x) / np.sqrt(30),
  70. np.fft.fft(x, norm="ortho"), atol=1e-6)
  71. assert_allclose(fft1(x) / 30.,
  72. np.fft.fft(x, norm="forward"), atol=1e-6)
  73. @pytest.mark.parametrize("axis", (0, 1))
  74. @pytest.mark.parametrize("dtype", (complex, float))
  75. @pytest.mark.parametrize("transpose", (True, False))
  76. def test_fft_out_argument(self, dtype, transpose, axis):
  77. def zeros_like(x):
  78. if transpose:
  79. return np.zeros_like(x.T).T
  80. else:
  81. return np.zeros_like(x)
  82. # tests below only test the out parameter
  83. if dtype is complex:
  84. y = random((10, 20)) + 1j * random((10, 20))
  85. fft, ifft = np.fft.fft, np.fft.ifft
  86. else:
  87. y = random((10, 20))
  88. fft, ifft = np.fft.rfft, np.fft.irfft
  89. expected = fft(y, axis=axis)
  90. out = zeros_like(expected)
  91. result = fft(y, out=out, axis=axis)
  92. assert result is out
  93. assert_array_equal(result, expected)
  94. expected2 = ifft(expected, axis=axis)
  95. out2 = out if dtype is complex else zeros_like(expected2)
  96. result2 = ifft(out, out=out2, axis=axis)
  97. assert result2 is out2
  98. assert_array_equal(result2, expected2)
  99. @pytest.mark.parametrize("axis", [0, 1])
  100. def test_fft_inplace_out(self, axis):
  101. # Test some weirder in-place combinations
  102. y = random((20, 20)) + 1j * random((20, 20))
  103. # Fully in-place.
  104. y1 = y.copy()
  105. expected1 = np.fft.fft(y1, axis=axis)
  106. result1 = np.fft.fft(y1, axis=axis, out=y1)
  107. assert result1 is y1
  108. assert_array_equal(result1, expected1)
  109. # In-place of part of the array; rest should be unchanged.
  110. y2 = y.copy()
  111. out2 = y2[:10] if axis == 0 else y2[:, :10]
  112. expected2 = np.fft.fft(y2, n=10, axis=axis)
  113. result2 = np.fft.fft(y2, n=10, axis=axis, out=out2)
  114. assert result2 is out2
  115. assert_array_equal(result2, expected2)
  116. if axis == 0:
  117. assert_array_equal(y2[10:], y[10:])
  118. else:
  119. assert_array_equal(y2[:, 10:], y[:, 10:])
  120. # In-place of another part of the array.
  121. y3 = y.copy()
  122. y3_sel = y3[5:] if axis == 0 else y3[:, 5:]
  123. out3 = y3[5:15] if axis == 0 else y3[:, 5:15]
  124. expected3 = np.fft.fft(y3_sel, n=10, axis=axis)
  125. result3 = np.fft.fft(y3_sel, n=10, axis=axis, out=out3)
  126. assert result3 is out3
  127. assert_array_equal(result3, expected3)
  128. if axis == 0:
  129. assert_array_equal(y3[:5], y[:5])
  130. assert_array_equal(y3[15:], y[15:])
  131. else:
  132. assert_array_equal(y3[:, :5], y[:, :5])
  133. assert_array_equal(y3[:, 15:], y[:, 15:])
  134. # In-place with n > nin; rest should be unchanged.
  135. y4 = y.copy()
  136. y4_sel = y4[:10] if axis == 0 else y4[:, :10]
  137. out4 = y4[:15] if axis == 0 else y4[:, :15]
  138. expected4 = np.fft.fft(y4_sel, n=15, axis=axis)
  139. result4 = np.fft.fft(y4_sel, n=15, axis=axis, out=out4)
  140. assert result4 is out4
  141. assert_array_equal(result4, expected4)
  142. if axis == 0:
  143. assert_array_equal(y4[15:], y[15:])
  144. else:
  145. assert_array_equal(y4[:, 15:], y[:, 15:])
  146. # Overwrite in a transpose.
  147. y5 = y.copy()
  148. out5 = y5.T
  149. result5 = np.fft.fft(y5, axis=axis, out=out5)
  150. assert result5 is out5
  151. assert_array_equal(result5, expected1)
  152. # Reverse strides.
  153. y6 = y.copy()
  154. out6 = y6[::-1] if axis == 0 else y6[:, ::-1]
  155. result6 = np.fft.fft(y6, axis=axis, out=out6)
  156. assert result6 is out6
  157. assert_array_equal(result6, expected1)
  158. def test_fft_bad_out(self):
  159. x = np.arange(30.)
  160. with pytest.raises(TypeError, match="must be of ArrayType"):
  161. np.fft.fft(x, out="")
  162. with pytest.raises(ValueError, match="has wrong shape"):
  163. np.fft.fft(x, out=np.zeros_like(x).reshape(5, -1))
  164. with pytest.raises(TypeError, match="Cannot cast"):
  165. np.fft.fft(x, out=np.zeros_like(x, dtype=float))
  166. @pytest.mark.parametrize('norm', (None, 'backward', 'ortho', 'forward'))
  167. def test_ifft(self, norm):
  168. x = random(30) + 1j * random(30)
  169. assert_allclose(
  170. x, np.fft.ifft(np.fft.fft(x, norm=norm), norm=norm),
  171. atol=1e-6)
  172. # Ensure we get the correct error message
  173. with pytest.raises(ValueError,
  174. match='Invalid number of FFT data points'):
  175. np.fft.ifft([], norm=norm)
  176. def test_fft2(self):
  177. x = random((30, 20)) + 1j * random((30, 20))
  178. assert_allclose(np.fft.fft(np.fft.fft(x, axis=1), axis=0),
  179. np.fft.fft2(x), atol=1e-6)
  180. assert_allclose(np.fft.fft2(x),
  181. np.fft.fft2(x, norm="backward"), atol=1e-6)
  182. assert_allclose(np.fft.fft2(x) / np.sqrt(30 * 20),
  183. np.fft.fft2(x, norm="ortho"), atol=1e-6)
  184. assert_allclose(np.fft.fft2(x) / (30. * 20.),
  185. np.fft.fft2(x, norm="forward"), atol=1e-6)
  186. def test_ifft2(self):
  187. x = random((30, 20)) + 1j * random((30, 20))
  188. assert_allclose(np.fft.ifft(np.fft.ifft(x, axis=1), axis=0),
  189. np.fft.ifft2(x), atol=1e-6)
  190. assert_allclose(np.fft.ifft2(x),
  191. np.fft.ifft2(x, norm="backward"), atol=1e-6)
  192. assert_allclose(np.fft.ifft2(x) * np.sqrt(30 * 20),
  193. np.fft.ifft2(x, norm="ortho"), atol=1e-6)
  194. assert_allclose(np.fft.ifft2(x) * (30. * 20.),
  195. np.fft.ifft2(x, norm="forward"), atol=1e-6)
  196. def test_fftn(self):
  197. x = random((30, 20, 10)) + 1j * random((30, 20, 10))
  198. assert_allclose(
  199. np.fft.fft(np.fft.fft(np.fft.fft(x, axis=2), axis=1), axis=0),
  200. np.fft.fftn(x), atol=1e-6)
  201. assert_allclose(np.fft.fftn(x),
  202. np.fft.fftn(x, norm="backward"), atol=1e-6)
  203. assert_allclose(np.fft.fftn(x) / np.sqrt(30 * 20 * 10),
  204. np.fft.fftn(x, norm="ortho"), atol=1e-6)
  205. assert_allclose(np.fft.fftn(x) / (30. * 20. * 10.),
  206. np.fft.fftn(x, norm="forward"), atol=1e-6)
  207. def test_ifftn(self):
  208. x = random((30, 20, 10)) + 1j * random((30, 20, 10))
  209. assert_allclose(
  210. np.fft.ifft(np.fft.ifft(np.fft.ifft(x, axis=2), axis=1), axis=0),
  211. np.fft.ifftn(x), atol=1e-6)
  212. assert_allclose(np.fft.ifftn(x),
  213. np.fft.ifftn(x, norm="backward"), atol=1e-6)
  214. assert_allclose(np.fft.ifftn(x) * np.sqrt(30 * 20 * 10),
  215. np.fft.ifftn(x, norm="ortho"), atol=1e-6)
  216. assert_allclose(np.fft.ifftn(x) * (30. * 20. * 10.),
  217. np.fft.ifftn(x, norm="forward"), atol=1e-6)
  218. def test_rfft(self):
  219. x = random(30)
  220. for n in [x.size, 2 * x.size]:
  221. for norm in [None, 'backward', 'ortho', 'forward']:
  222. assert_allclose(
  223. np.fft.fft(x, n=n, norm=norm)[:(n // 2 + 1)],
  224. np.fft.rfft(x, n=n, norm=norm), atol=1e-6)
  225. assert_allclose(
  226. np.fft.rfft(x, n=n),
  227. np.fft.rfft(x, n=n, norm="backward"), atol=1e-6)
  228. assert_allclose(
  229. np.fft.rfft(x, n=n) / np.sqrt(n),
  230. np.fft.rfft(x, n=n, norm="ortho"), atol=1e-6)
  231. assert_allclose(
  232. np.fft.rfft(x, n=n) / n,
  233. np.fft.rfft(x, n=n, norm="forward"), atol=1e-6)
  234. def test_rfft_even(self):
  235. x = np.arange(8)
  236. n = 4
  237. y = np.fft.rfft(x, n)
  238. assert_allclose(y, np.fft.fft(x[:n])[:n // 2 + 1], rtol=1e-14)
  239. def test_rfft_odd(self):
  240. x = np.array([1, 0, 2, 3, -3])
  241. y = np.fft.rfft(x)
  242. assert_allclose(y, np.fft.fft(x)[:3], rtol=1e-14)
  243. def test_irfft(self):
  244. x = random(30)
  245. assert_allclose(x, np.fft.irfft(np.fft.rfft(x)), atol=1e-6)
  246. assert_allclose(x, np.fft.irfft(np.fft.rfft(x, norm="backward"),
  247. norm="backward"), atol=1e-6)
  248. assert_allclose(x, np.fft.irfft(np.fft.rfft(x, norm="ortho"),
  249. norm="ortho"), atol=1e-6)
  250. assert_allclose(x, np.fft.irfft(np.fft.rfft(x, norm="forward"),
  251. norm="forward"), atol=1e-6)
  252. def test_rfft2(self):
  253. x = random((30, 20))
  254. assert_allclose(np.fft.fft2(x)[:, :11], np.fft.rfft2(x), atol=1e-6)
  255. assert_allclose(np.fft.rfft2(x),
  256. np.fft.rfft2(x, norm="backward"), atol=1e-6)
  257. assert_allclose(np.fft.rfft2(x) / np.sqrt(30 * 20),
  258. np.fft.rfft2(x, norm="ortho"), atol=1e-6)
  259. assert_allclose(np.fft.rfft2(x) / (30. * 20.),
  260. np.fft.rfft2(x, norm="forward"), atol=1e-6)
  261. def test_irfft2(self):
  262. x = random((30, 20))
  263. assert_allclose(x, np.fft.irfft2(np.fft.rfft2(x)), atol=1e-6)
  264. assert_allclose(x, np.fft.irfft2(np.fft.rfft2(x, norm="backward"),
  265. norm="backward"), atol=1e-6)
  266. assert_allclose(x, np.fft.irfft2(np.fft.rfft2(x, norm="ortho"),
  267. norm="ortho"), atol=1e-6)
  268. assert_allclose(x, np.fft.irfft2(np.fft.rfft2(x, norm="forward"),
  269. norm="forward"), atol=1e-6)
  270. def test_rfftn(self):
  271. x = random((30, 20, 10))
  272. assert_allclose(np.fft.fftn(x)[:, :, :6], np.fft.rfftn(x), atol=1e-6)
  273. assert_allclose(np.fft.rfftn(x),
  274. np.fft.rfftn(x, norm="backward"), atol=1e-6)
  275. assert_allclose(np.fft.rfftn(x) / np.sqrt(30 * 20 * 10),
  276. np.fft.rfftn(x, norm="ortho"), atol=1e-6)
  277. assert_allclose(np.fft.rfftn(x) / (30. * 20. * 10.),
  278. np.fft.rfftn(x, norm="forward"), atol=1e-6)
  279. # Regression test for gh-27159
  280. x = np.ones((2, 3))
  281. result = np.fft.rfftn(x, axes=(0, 0, 1), s=(10, 20, 40))
  282. assert result.shape == (10, 21)
  283. expected = np.fft.fft(np.fft.fft(np.fft.rfft(x, axis=1, n=40),
  284. axis=0, n=20), axis=0, n=10)
  285. assert expected.shape == (10, 21)
  286. assert_allclose(result, expected, atol=1e-6)
  287. def test_irfftn(self):
  288. x = random((30, 20, 10))
  289. assert_allclose(x, np.fft.irfftn(np.fft.rfftn(x)), atol=1e-6)
  290. assert_allclose(x, np.fft.irfftn(np.fft.rfftn(x, norm="backward"),
  291. norm="backward"), atol=1e-6)
  292. assert_allclose(x, np.fft.irfftn(np.fft.rfftn(x, norm="ortho"),
  293. norm="ortho"), atol=1e-6)
  294. assert_allclose(x, np.fft.irfftn(np.fft.rfftn(x, norm="forward"),
  295. norm="forward"), atol=1e-6)
  296. def test_hfft(self):
  297. x = random(14) + 1j * random(14)
  298. x_herm = np.concatenate((random(1), x, random(1)))
  299. x = np.concatenate((x_herm, x[::-1].conj()))
  300. assert_allclose(np.fft.fft(x), np.fft.hfft(x_herm), atol=1e-6)
  301. assert_allclose(np.fft.hfft(x_herm),
  302. np.fft.hfft(x_herm, norm="backward"), atol=1e-6)
  303. assert_allclose(np.fft.hfft(x_herm) / np.sqrt(30),
  304. np.fft.hfft(x_herm, norm="ortho"), atol=1e-6)
  305. assert_allclose(np.fft.hfft(x_herm) / 30.,
  306. np.fft.hfft(x_herm, norm="forward"), atol=1e-6)
  307. def test_ihfft(self):
  308. x = random(14) + 1j * random(14)
  309. x_herm = np.concatenate((random(1), x, random(1)))
  310. x = np.concatenate((x_herm, x[::-1].conj()))
  311. assert_allclose(x_herm, np.fft.ihfft(np.fft.hfft(x_herm)), atol=1e-6)
  312. assert_allclose(x_herm, np.fft.ihfft(np.fft.hfft(x_herm,
  313. norm="backward"), norm="backward"), atol=1e-6)
  314. assert_allclose(x_herm, np.fft.ihfft(np.fft.hfft(x_herm,
  315. norm="ortho"), norm="ortho"), atol=1e-6)
  316. assert_allclose(x_herm, np.fft.ihfft(np.fft.hfft(x_herm,
  317. norm="forward"), norm="forward"), atol=1e-6)
  318. @pytest.mark.parametrize("op", [np.fft.fftn, np.fft.ifftn,
  319. np.fft.rfftn, np.fft.irfftn])
  320. def test_axes(self, op):
  321. x = random((30, 20, 10))
  322. axes = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)]
  323. for a in axes:
  324. op_tr = op(np.transpose(x, a))
  325. tr_op = np.transpose(op(x, axes=a), a)
  326. assert_allclose(op_tr, tr_op, atol=1e-6)
  327. @pytest.mark.parametrize("op", [np.fft.fftn, np.fft.ifftn,
  328. np.fft.fft2, np.fft.ifft2])
  329. def test_s_negative_1(self, op):
  330. x = np.arange(100).reshape(10, 10)
  331. # should use the whole input array along the first axis
  332. assert op(x, s=(-1, 5), axes=(0, 1)).shape == (10, 5)
  333. @pytest.mark.parametrize("op", [np.fft.fftn, np.fft.ifftn,
  334. np.fft.rfftn, np.fft.irfftn])
  335. def test_s_axes_none(self, op):
  336. x = np.arange(100).reshape(10, 10)
  337. with pytest.warns(match='`axes` should not be `None` if `s`'):
  338. op(x, s=(-1, 5))
  339. @pytest.mark.parametrize("op", [np.fft.fft2, np.fft.ifft2])
  340. def test_s_axes_none_2D(self, op):
  341. x = np.arange(100).reshape(10, 10)
  342. with pytest.warns(match='`axes` should not be `None` if `s`'):
  343. op(x, s=(-1, 5), axes=None)
  344. @pytest.mark.parametrize("op", [np.fft.fftn, np.fft.ifftn,
  345. np.fft.rfftn, np.fft.irfftn,
  346. np.fft.fft2, np.fft.ifft2])
  347. def test_s_contains_none(self, op):
  348. x = random((30, 20, 10))
  349. with pytest.warns(match='array containing `None` values to `s`'):
  350. op(x, s=(10, None, 10), axes=(0, 1, 2))
  351. def test_all_1d_norm_preserving(self):
  352. # verify that round-trip transforms are norm-preserving
  353. x = random(30)
  354. x_norm = np.linalg.norm(x)
  355. n = x.size * 2
  356. func_pairs = [(np.fft.fft, np.fft.ifft),
  357. (np.fft.rfft, np.fft.irfft),
  358. # hfft: order so the first function takes x.size samples
  359. # (necessary for comparison to x_norm above)
  360. (np.fft.ihfft, np.fft.hfft),
  361. ]
  362. for forw, back in func_pairs:
  363. for n in [x.size, 2 * x.size]:
  364. for norm in [None, 'backward', 'ortho', 'forward']:
  365. tmp = forw(x, n=n, norm=norm)
  366. tmp = back(tmp, n=n, norm=norm)
  367. assert_allclose(x_norm,
  368. np.linalg.norm(tmp), atol=1e-6)
  369. @pytest.mark.parametrize("axes", [(0, 1), (0, 2), None])
  370. @pytest.mark.parametrize("dtype", (complex, float))
  371. @pytest.mark.parametrize("transpose", (True, False))
  372. def test_fftn_out_argument(self, dtype, transpose, axes):
  373. def zeros_like(x):
  374. if transpose:
  375. return np.zeros_like(x.T).T
  376. else:
  377. return np.zeros_like(x)
  378. # tests below only test the out parameter
  379. if dtype is complex:
  380. x = random((10, 5, 6)) + 1j * random((10, 5, 6))
  381. fft, ifft = np.fft.fftn, np.fft.ifftn
  382. else:
  383. x = random((10, 5, 6))
  384. fft, ifft = np.fft.rfftn, np.fft.irfftn
  385. expected = fft(x, axes=axes)
  386. out = zeros_like(expected)
  387. result = fft(x, out=out, axes=axes)
  388. assert result is out
  389. assert_array_equal(result, expected)
  390. expected2 = ifft(expected, axes=axes)
  391. out2 = out if dtype is complex else zeros_like(expected2)
  392. result2 = ifft(out, out=out2, axes=axes)
  393. assert result2 is out2
  394. assert_array_equal(result2, expected2)
  395. @pytest.mark.parametrize("fft", [np.fft.fftn, np.fft.ifftn, np.fft.rfftn])
  396. def test_fftn_out_and_s_interaction(self, fft):
  397. # With s, shape varies, so generally one cannot pass in out.
  398. if fft is np.fft.rfftn:
  399. x = random((10, 5, 6))
  400. else:
  401. x = random((10, 5, 6)) + 1j * random((10, 5, 6))
  402. with pytest.raises(ValueError, match="has wrong shape"):
  403. fft(x, out=np.zeros_like(x), s=(3, 3, 3), axes=(0, 1, 2))
  404. # Except on the first axis done (which is the last of axes).
  405. s = (10, 5, 5)
  406. expected = fft(x, s=s, axes=(0, 1, 2))
  407. out = np.zeros_like(expected)
  408. result = fft(x, s=s, axes=(0, 1, 2), out=out)
  409. assert result is out
  410. assert_array_equal(result, expected)
  411. @pytest.mark.parametrize("s", [(9, 5, 5), (3, 3, 3)])
  412. def test_irfftn_out_and_s_interaction(self, s):
  413. # Since for irfftn, the output is real and thus cannot be used for
  414. # intermediate steps, it should always work.
  415. x = random((9, 5, 6, 2)) + 1j * random((9, 5, 6, 2))
  416. expected = np.fft.irfftn(x, s=s, axes=(0, 1, 2))
  417. out = np.zeros_like(expected)
  418. result = np.fft.irfftn(x, s=s, axes=(0, 1, 2), out=out)
  419. assert result is out
  420. assert_array_equal(result, expected)
  421. @pytest.mark.parametrize(
  422. "dtype",
  423. [np.float32, np.float64, np.complex64, np.complex128])
  424. @pytest.mark.parametrize("order", ["F", 'non-contiguous'])
  425. @pytest.mark.parametrize(
  426. "fft",
  427. [np.fft.fft, np.fft.fft2, np.fft.fftn,
  428. np.fft.ifft, np.fft.ifft2, np.fft.ifftn])
  429. def test_fft_with_order(dtype, order, fft):
  430. # Check that FFT/IFFT produces identical results for C, Fortran and
  431. # non contiguous arrays
  432. rng = np.random.RandomState(42)
  433. X = rng.rand(8, 7, 13).astype(dtype, copy=False)
  434. # See discussion in pull/14178
  435. _tol = 8.0 * np.sqrt(np.log2(X.size)) * np.finfo(X.dtype).eps
  436. if order == 'F':
  437. Y = np.asfortranarray(X)
  438. else:
  439. # Make a non contiguous array
  440. Y = X[::-1]
  441. X = np.ascontiguousarray(X[::-1])
  442. if fft.__name__.endswith('fft'):
  443. for axis in range(3):
  444. X_res = fft(X, axis=axis)
  445. Y_res = fft(Y, axis=axis)
  446. assert_allclose(X_res, Y_res, atol=_tol, rtol=_tol)
  447. elif fft.__name__.endswith(('fft2', 'fftn')):
  448. axes = [(0, 1), (1, 2), (0, 2)]
  449. if fft.__name__.endswith('fftn'):
  450. axes.extend([(0,), (1,), (2,), None])
  451. for ax in axes:
  452. X_res = fft(X, axes=ax)
  453. Y_res = fft(Y, axes=ax)
  454. assert_allclose(X_res, Y_res, atol=_tol, rtol=_tol)
  455. else:
  456. raise ValueError
  457. @pytest.mark.parametrize("order", ["F", "C"])
  458. @pytest.mark.parametrize("n", [None, 7, 12])
  459. def test_fft_output_order(order, n):
  460. rng = np.random.RandomState(42)
  461. x = rng.rand(10)
  462. x = np.asarray(x, dtype=np.complex64, order=order)
  463. res = np.fft.fft(x, n=n)
  464. assert res.flags.c_contiguous == x.flags.c_contiguous
  465. assert res.flags.f_contiguous == x.flags.f_contiguous
  466. @pytest.mark.skipif(IS_WASM, reason="Cannot start thread")
  467. class TestFFTThreadSafe:
  468. threads = 16
  469. input_shape = (800, 200)
  470. def _test_mtsame(self, func, *args):
  471. def worker(args, q):
  472. q.put(func(*args))
  473. q = queue.Queue()
  474. expected = func(*args)
  475. # Spin off a bunch of threads to call the same function simultaneously
  476. t = [threading.Thread(target=worker, args=(args, q))
  477. for i in range(self.threads)]
  478. [x.start() for x in t]
  479. [x.join() for x in t]
  480. # Make sure all threads returned the correct value
  481. for i in range(self.threads):
  482. assert_array_equal(q.get(timeout=5), expected,
  483. 'Function returned wrong value in multithreaded context')
  484. def test_fft(self):
  485. a = np.ones(self.input_shape) * 1 + 0j
  486. self._test_mtsame(np.fft.fft, a)
  487. def test_ifft(self):
  488. a = np.ones(self.input_shape) * 1 + 0j
  489. self._test_mtsame(np.fft.ifft, a)
  490. def test_rfft(self):
  491. a = np.ones(self.input_shape)
  492. self._test_mtsame(np.fft.rfft, a)
  493. def test_irfft(self):
  494. a = np.ones(self.input_shape) * 1 + 0j
  495. self._test_mtsame(np.fft.irfft, a)
  496. def test_irfft_with_n_1_regression():
  497. # Regression test for gh-25661
  498. x = np.arange(10)
  499. np.fft.irfft(x, n=1)
  500. np.fft.hfft(x, n=1)
  501. np.fft.irfft(np.array([0], complex), n=10)
  502. def test_irfft_with_n_large_regression():
  503. # Regression test for gh-25679
  504. x = np.arange(5) * (1 + 1j)
  505. result = np.fft.hfft(x, n=10)
  506. expected = np.array([20., 9.91628173, -11.8819096, 7.1048486,
  507. -6.62459848, 4., -3.37540152, -0.16057669,
  508. 1.8819096, -20.86055364])
  509. assert_allclose(result, expected)
  510. @pytest.mark.parametrize("fft", [
  511. np.fft.fft, np.fft.ifft, np.fft.rfft, np.fft.irfft
  512. ])
  513. @pytest.mark.parametrize("data", [
  514. np.array([False, True, False]),
  515. np.arange(10, dtype=np.uint8),
  516. np.arange(5, dtype=np.int16),
  517. ])
  518. def test_fft_with_integer_or_bool_input(data, fft):
  519. # Regression test for gh-25819
  520. result = fft(data)
  521. float_data = data.astype(np.result_type(data, 1.))
  522. expected = fft(float_data)
  523. assert_array_equal(result, expected)