uft.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451
  1. r"""Function of unitary fourier transform (uft) and utilities
  2. This module implements the unitary fourier transform, also known as
  3. the ortho-normal transform. It is especially useful for convolution
  4. [1], as it respects the Parseval equality. The value of the null
  5. frequency is equal to
  6. .. math:: \frac{1}{\sqrt{n}} \sum_i x_i
  7. so the Fourier transform has the same energy as the original image
  8. (see ``image_quad_norm`` function). The transform is applied from the
  9. last axis for performance (assuming a C-order array input).
  10. References
  11. ----------
  12. .. [1] B. R. Hunt "A matrix theory proof of the discrete convolution
  13. theorem", IEEE Trans. on Audio and Electroacoustics,
  14. vol. au-19, no. 4, pp. 285-288, dec. 1971
  15. """
  16. import numpy as np
  17. import scipy.fft as fft
  18. from .._shared.utils import _supported_float_type
  19. def ufftn(inarray, dim=None):
  20. """N-dimensional unitary Fourier transform.
  21. Parameters
  22. ----------
  23. inarray : ndarray
  24. The array to transform.
  25. dim : int, optional
  26. The last axis along which to compute the transform. All
  27. axes by default.
  28. Returns
  29. -------
  30. outarray : ndarray (same shape than inarray)
  31. The unitary N-D Fourier transform of ``inarray``.
  32. Examples
  33. --------
  34. >>> input = np.ones((3, 3, 3))
  35. >>> output = ufftn(input)
  36. >>> np.allclose(np.sum(input) / np.sqrt(input.size), output[0, 0, 0])
  37. True
  38. >>> output.shape
  39. (3, 3, 3)
  40. """
  41. if dim is None:
  42. dim = inarray.ndim
  43. outarray = fft.fftn(inarray, axes=range(-dim, 0), norm='ortho')
  44. return outarray
  45. def uifftn(inarray, dim=None):
  46. """N-dimensional unitary inverse Fourier transform.
  47. Parameters
  48. ----------
  49. inarray : ndarray
  50. The array to transform.
  51. dim : int, optional
  52. The last axis along which to compute the transform. All
  53. axes by default.
  54. Returns
  55. -------
  56. outarray : ndarray
  57. The unitary inverse nD Fourier transform of ``inarray``. Has the same shape as
  58. ``inarray``.
  59. Examples
  60. --------
  61. >>> input = np.ones((3, 3, 3))
  62. >>> output = uifftn(input)
  63. >>> np.allclose(np.sum(input) / np.sqrt(input.size), output[0, 0, 0])
  64. True
  65. >>> output.shape
  66. (3, 3, 3)
  67. """
  68. if dim is None:
  69. dim = inarray.ndim
  70. outarray = fft.ifftn(inarray, axes=range(-dim, 0), norm='ortho')
  71. return outarray
  72. def urfftn(inarray, dim=None):
  73. """N-dimensional real unitary Fourier transform.
  74. This transform considers the Hermitian property of the transform on
  75. real-valued input.
  76. Parameters
  77. ----------
  78. inarray : ndarray, shape (M[, ...], P)
  79. The array to transform.
  80. dim : int, optional
  81. The last axis along which to compute the transform. All
  82. axes by default.
  83. Returns
  84. -------
  85. outarray : ndarray, shape (M[, ...], P / 2 + 1)
  86. The unitary N-D real Fourier transform of ``inarray``.
  87. Notes
  88. -----
  89. The ``urfft`` functions assume an input array of real
  90. values. Consequently, the output has a Hermitian property and
  91. redundant values are not computed or returned.
  92. Examples
  93. --------
  94. >>> input = np.ones((5, 5, 5))
  95. >>> output = urfftn(input)
  96. >>> np.allclose(np.sum(input) / np.sqrt(input.size), output[0, 0, 0])
  97. True
  98. >>> output.shape
  99. (5, 5, 3)
  100. """
  101. if dim is None:
  102. dim = inarray.ndim
  103. outarray = fft.rfftn(inarray, axes=range(-dim, 0), norm='ortho')
  104. return outarray
  105. def uirfftn(inarray, dim=None, shape=None):
  106. """N-dimensional inverse real unitary Fourier transform.
  107. This transform considers the Hermitian property of the transform
  108. from complex to real input.
  109. Parameters
  110. ----------
  111. inarray : ndarray
  112. The array to transform.
  113. dim : int, optional
  114. The last axis along which to compute the transform. All
  115. axes by default.
  116. shape : tuple of int, optional
  117. The shape of the output. The shape of ``rfft`` is ambiguous in
  118. case of odd-valued input shape. In this case, this parameter
  119. should be provided. See ``np.fft.irfftn``.
  120. Returns
  121. -------
  122. outarray : ndarray
  123. The unitary N-D inverse real Fourier transform of ``inarray``.
  124. Notes
  125. -----
  126. The ``uirfft`` function assumes that the output array is
  127. real-valued. Consequently, the input is assumed to have a Hermitian
  128. property and redundant values are implicit.
  129. Examples
  130. --------
  131. >>> input = np.ones((5, 5, 5))
  132. >>> output = uirfftn(urfftn(input), shape=input.shape)
  133. >>> np.allclose(input, output)
  134. True
  135. >>> output.shape
  136. (5, 5, 5)
  137. """
  138. if dim is None:
  139. dim = inarray.ndim
  140. outarray = fft.irfftn(inarray, shape, axes=range(-dim, 0), norm='ortho')
  141. return outarray
  142. def ufft2(inarray):
  143. """2-dimensional unitary Fourier transform.
  144. Compute the Fourier transform on the last 2 axes.
  145. Parameters
  146. ----------
  147. inarray : ndarray
  148. The array to transform.
  149. Returns
  150. -------
  151. outarray : ndarray (same shape as inarray)
  152. The unitary 2-D Fourier transform of ``inarray``.
  153. See Also
  154. --------
  155. uifft2, ufftn, urfftn
  156. Examples
  157. --------
  158. >>> input = np.ones((10, 128, 128))
  159. >>> output = ufft2(input)
  160. >>> np.allclose(np.sum(input[1, ...]) / np.sqrt(input[1, ...].size),
  161. ... output[1, 0, 0])
  162. True
  163. >>> output.shape
  164. (10, 128, 128)
  165. """
  166. return ufftn(inarray, 2)
  167. def uifft2(inarray):
  168. """2-dimensional inverse unitary Fourier transform.
  169. Compute the inverse Fourier transform on the last 2 axes.
  170. Parameters
  171. ----------
  172. inarray : ndarray
  173. The array to transform.
  174. Returns
  175. -------
  176. outarray : ndarray (same shape as inarray)
  177. The unitary 2-D inverse Fourier transform of ``inarray``.
  178. See Also
  179. --------
  180. uifft2, uifftn, uirfftn
  181. Examples
  182. --------
  183. >>> input = np.ones((10, 128, 128))
  184. >>> output = uifft2(input)
  185. >>> np.allclose(np.sum(input[1, ...]) / np.sqrt(input[1, ...].size),
  186. ... output[0, 0, 0])
  187. True
  188. >>> output.shape
  189. (10, 128, 128)
  190. """
  191. return uifftn(inarray, 2)
  192. def urfft2(inarray):
  193. """2-dimensional real unitary Fourier transform
  194. Compute the real Fourier transform on the last 2 axes. This
  195. transform considers the Hermitian property of the transform from
  196. complex to real-valued input.
  197. Parameters
  198. ----------
  199. inarray : ndarray, shape (M[, ...], P)
  200. The array to transform.
  201. Returns
  202. -------
  203. outarray : ndarray, shape (M[, ...], 2 * (P - 1))
  204. The unitary 2-D real Fourier transform of ``inarray``.
  205. See Also
  206. --------
  207. ufft2, ufftn, urfftn
  208. Examples
  209. --------
  210. >>> input = np.ones((10, 128, 128))
  211. >>> output = urfft2(input)
  212. >>> np.allclose(np.sum(input[1,...]) / np.sqrt(input[1,...].size),
  213. ... output[1, 0, 0])
  214. True
  215. >>> output.shape
  216. (10, 128, 65)
  217. """
  218. return urfftn(inarray, 2)
  219. def uirfft2(inarray, shape=None):
  220. """2-dimensional inverse real unitary Fourier transform.
  221. Compute the real inverse Fourier transform on the last 2 axes.
  222. This transform considers the Hermitian property of the transform
  223. from complex to real-valued input.
  224. Parameters
  225. ----------
  226. inarray : ndarray, shape (M[, ...], P)
  227. The array to transform.
  228. shape : tuple of int, optional
  229. The shape of the output. The shape of ``rfft`` is ambiguous in
  230. case of odd-valued input shape. In this case, this parameter
  231. should be provided. See ``np.fft.irfftn``.
  232. Returns
  233. -------
  234. outarray : ndarray, shape (M[, ...], 2 * (P - 1))
  235. The unitary 2-D inverse real Fourier transform of ``inarray``.
  236. See Also
  237. --------
  238. urfft2, uifftn, uirfftn
  239. Examples
  240. --------
  241. >>> input = np.ones((10, 128, 128))
  242. >>> output = uirfftn(urfftn(input), shape=input.shape)
  243. >>> np.allclose(input, output)
  244. True
  245. >>> output.shape
  246. (10, 128, 128)
  247. """
  248. return uirfftn(inarray, 2, shape=shape)
  249. def image_quad_norm(inarray):
  250. """Return the quadratic norm of images in Fourier space.
  251. This function detects whether the input image satisfies the
  252. Hermitian property.
  253. Parameters
  254. ----------
  255. inarray : ndarray
  256. Input image. The image data should reside in the final two
  257. axes.
  258. Returns
  259. -------
  260. norm : float
  261. The quadratic norm of ``inarray``.
  262. Examples
  263. --------
  264. >>> input = np.ones((5, 5))
  265. >>> image_quad_norm(ufft2(input)) == np.sum(np.abs(input)**2)
  266. True
  267. >>> image_quad_norm(ufft2(input)) == image_quad_norm(urfft2(input))
  268. True
  269. """
  270. # If there is a Hermitian symmetry
  271. if inarray.shape[-1] != inarray.shape[-2]:
  272. return 2 * np.sum(np.sum(np.abs(inarray) ** 2, axis=-1), axis=-1) - np.sum(
  273. np.abs(inarray[..., 0]) ** 2, axis=-1
  274. )
  275. else:
  276. return np.sum(np.sum(np.abs(inarray) ** 2, axis=-1), axis=-1)
  277. def ir2tf(imp_resp, shape, dim=None, is_real=True):
  278. """Compute the transfer function of an impulse response (IR).
  279. This function makes the necessary correct zero-padding, zero
  280. convention, correct fft2, etc... to compute the transfer function
  281. of IR. To use with unitary Fourier transform for the signal (ufftn
  282. or equivalent).
  283. Parameters
  284. ----------
  285. imp_resp : ndarray
  286. The impulse responses.
  287. shape : tuple of int
  288. A tuple of integer corresponding to the target shape of the
  289. transfer function.
  290. dim : int, optional
  291. The last axis along which to compute the transform. All
  292. axes by default.
  293. is_real : bool, optional
  294. If True (default), imp_resp is supposed real and the Hermitian property
  295. is used with rfftn Fourier transform.
  296. Returns
  297. -------
  298. y : complex ndarray
  299. The transfer function of shape ``shape``.
  300. See Also
  301. --------
  302. ufftn, uifftn, urfftn, uirfftn
  303. Examples
  304. --------
  305. >>> np.all(np.array([[4, 0], [0, 0]]) == ir2tf(np.ones((2, 2)), (2, 2)))
  306. True
  307. >>> ir2tf(np.ones((2, 2)), (512, 512)).shape == (512, 257)
  308. True
  309. >>> ir2tf(np.ones((2, 2)), (512, 512), is_real=False).shape == (512, 512)
  310. True
  311. Notes
  312. -----
  313. The input array can be composed of multiple-dimensional IR with
  314. an arbitrary number of IR. The individual IR must be accessed
  315. through the first axes. The last ``dim`` axes contain the space
  316. definition.
  317. """
  318. if not dim:
  319. dim = imp_resp.ndim
  320. # Zero padding and fill
  321. irpadded_dtype = _supported_float_type(imp_resp.dtype)
  322. irpadded = np.zeros(shape, dtype=irpadded_dtype)
  323. irpadded[tuple([slice(0, s) for s in imp_resp.shape])] = imp_resp
  324. # Roll for zero convention of the fft to avoid the phase
  325. # problem. Work with odd and even size.
  326. for axis, axis_size in enumerate(imp_resp.shape):
  327. if axis >= imp_resp.ndim - dim:
  328. irpadded = np.roll(irpadded, shift=-int(np.floor(axis_size / 2)), axis=axis)
  329. func = fft.rfftn if is_real else fft.fftn
  330. out = func(irpadded, axes=(range(-dim, 0)))
  331. # TODO: remove .astype call once SciPy >= 1.4 is required
  332. cplx_dtype = np.promote_types(irpadded_dtype, np.complex64)
  333. return out.astype(cplx_dtype, copy=False)
  334. def laplacian(ndim, shape, is_real=True):
  335. """Return the transfer function of the Laplacian.
  336. Laplacian is the second order difference, on row and column.
  337. Parameters
  338. ----------
  339. ndim : int
  340. The dimension of the Laplacian.
  341. shape : tuple
  342. The support on which to compute the transfer function.
  343. is_real : bool, optional
  344. If True (default), imp_resp is assumed to be real-valued and
  345. the Hermitian property is used with rfftn Fourier transform
  346. to return the transfer function.
  347. Returns
  348. -------
  349. tf : array_like, complex
  350. The transfer function.
  351. impr : array_like, real
  352. The Laplacian.
  353. Examples
  354. --------
  355. >>> tf, ir = laplacian(2, (32, 32))
  356. >>> np.all(ir == np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]]))
  357. True
  358. >>> np.all(tf == ir2tf(ir, (32, 32)))
  359. True
  360. """
  361. impr = np.zeros([3] * ndim)
  362. for dim in range(ndim):
  363. idx = tuple(
  364. [slice(1, 2)] * dim + [slice(None)] + [slice(1, 2)] * (ndim - dim - 1)
  365. )
  366. impr[idx] = np.array([-1.0, 0.0, -1.0]).reshape(
  367. [-1 if i == dim else 1 for i in range(ndim)]
  368. )
  369. impr[(slice(1, 2),) * ndim] = 2.0 * ndim
  370. return ir2tf(impr, shape, is_real=is_real), impr