test_pocketfft.py 24 KB

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