_helper.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235
  1. """
  2. Discrete Fourier Transforms - _helper.py
  3. """
  4. from numpy._core import integer, empty, arange, asarray, roll
  5. from numpy._core.overrides import array_function_dispatch, set_module
  6. # Created by Pearu Peterson, September 2002
  7. __all__ = ['fftshift', 'ifftshift', 'fftfreq', 'rfftfreq']
  8. integer_types = (int, integer)
  9. def _fftshift_dispatcher(x, axes=None):
  10. return (x,)
  11. @array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')
  12. def fftshift(x, axes=None):
  13. """
  14. Shift the zero-frequency component to the center of the spectrum.
  15. This function swaps half-spaces for all axes listed (defaults to all).
  16. Note that ``y[0]`` is the Nyquist component only if ``len(x)`` is even.
  17. Parameters
  18. ----------
  19. x : array_like
  20. Input array.
  21. axes : int or shape tuple, optional
  22. Axes over which to shift. Default is None, which shifts all axes.
  23. Returns
  24. -------
  25. y : ndarray
  26. The shifted array.
  27. See Also
  28. --------
  29. ifftshift : The inverse of `fftshift`.
  30. Examples
  31. --------
  32. >>> import numpy as np
  33. >>> freqs = np.fft.fftfreq(10, 0.1)
  34. >>> freqs
  35. array([ 0., 1., 2., ..., -3., -2., -1.])
  36. >>> np.fft.fftshift(freqs)
  37. array([-5., -4., -3., -2., -1., 0., 1., 2., 3., 4.])
  38. Shift the zero-frequency component only along the second axis:
  39. >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
  40. >>> freqs
  41. array([[ 0., 1., 2.],
  42. [ 3., 4., -4.],
  43. [-3., -2., -1.]])
  44. >>> np.fft.fftshift(freqs, axes=(1,))
  45. array([[ 2., 0., 1.],
  46. [-4., 3., 4.],
  47. [-1., -3., -2.]])
  48. """
  49. x = asarray(x)
  50. if axes is None:
  51. axes = tuple(range(x.ndim))
  52. shift = [dim // 2 for dim in x.shape]
  53. elif isinstance(axes, integer_types):
  54. shift = x.shape[axes] // 2
  55. else:
  56. shift = [x.shape[ax] // 2 for ax in axes]
  57. return roll(x, shift, axes)
  58. @array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')
  59. def ifftshift(x, axes=None):
  60. """
  61. The inverse of `fftshift`. Although identical for even-length `x`, the
  62. functions differ by one sample for odd-length `x`.
  63. Parameters
  64. ----------
  65. x : array_like
  66. Input array.
  67. axes : int or shape tuple, optional
  68. Axes over which to calculate. Defaults to None, which shifts all axes.
  69. Returns
  70. -------
  71. y : ndarray
  72. The shifted array.
  73. See Also
  74. --------
  75. fftshift : Shift zero-frequency component to the center of the spectrum.
  76. Examples
  77. --------
  78. >>> import numpy as np
  79. >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
  80. >>> freqs
  81. array([[ 0., 1., 2.],
  82. [ 3., 4., -4.],
  83. [-3., -2., -1.]])
  84. >>> np.fft.ifftshift(np.fft.fftshift(freqs))
  85. array([[ 0., 1., 2.],
  86. [ 3., 4., -4.],
  87. [-3., -2., -1.]])
  88. """
  89. x = asarray(x)
  90. if axes is None:
  91. axes = tuple(range(x.ndim))
  92. shift = [-(dim // 2) for dim in x.shape]
  93. elif isinstance(axes, integer_types):
  94. shift = -(x.shape[axes] // 2)
  95. else:
  96. shift = [-(x.shape[ax] // 2) for ax in axes]
  97. return roll(x, shift, axes)
  98. @set_module('numpy.fft')
  99. def fftfreq(n, d=1.0, device=None):
  100. """
  101. Return the Discrete Fourier Transform sample frequencies.
  102. The returned float array `f` contains the frequency bin centers in cycles
  103. per unit of the sample spacing (with zero at the start). For instance, if
  104. the sample spacing is in seconds, then the frequency unit is cycles/second.
  105. Given a window length `n` and a sample spacing `d`::
  106. f = [0, 1, ..., n/2-1, -n/2, ..., -1] / (d*n) if n is even
  107. f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n) if n is odd
  108. Parameters
  109. ----------
  110. n : int
  111. Window length.
  112. d : scalar, optional
  113. Sample spacing (inverse of the sampling rate). Defaults to 1.
  114. device : str, optional
  115. The device on which to place the created array. Default: ``None``.
  116. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  117. .. versionadded:: 2.0.0
  118. Returns
  119. -------
  120. f : ndarray
  121. Array of length `n` containing the sample frequencies.
  122. Examples
  123. --------
  124. >>> import numpy as np
  125. >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)
  126. >>> fourier = np.fft.fft(signal)
  127. >>> n = signal.size
  128. >>> timestep = 0.1
  129. >>> freq = np.fft.fftfreq(n, d=timestep)
  130. >>> freq
  131. array([ 0. , 1.25, 2.5 , ..., -3.75, -2.5 , -1.25])
  132. """
  133. if not isinstance(n, integer_types):
  134. raise ValueError("n should be an integer")
  135. val = 1.0 / (n * d)
  136. results = empty(n, int, device=device)
  137. N = (n-1)//2 + 1
  138. p1 = arange(0, N, dtype=int, device=device)
  139. results[:N] = p1
  140. p2 = arange(-(n//2), 0, dtype=int, device=device)
  141. results[N:] = p2
  142. return results * val
  143. @set_module('numpy.fft')
  144. def rfftfreq(n, d=1.0, device=None):
  145. """
  146. Return the Discrete Fourier Transform sample frequencies
  147. (for usage with rfft, irfft).
  148. The returned float array `f` contains the frequency bin centers in cycles
  149. per unit of the sample spacing (with zero at the start). For instance, if
  150. the sample spacing is in seconds, then the frequency unit is cycles/second.
  151. Given a window length `n` and a sample spacing `d`::
  152. f = [0, 1, ..., n/2-1, n/2] / (d*n) if n is even
  153. f = [0, 1, ..., (n-1)/2-1, (n-1)/2] / (d*n) if n is odd
  154. Unlike `fftfreq` (but like `scipy.fftpack.rfftfreq`)
  155. the Nyquist frequency component is considered to be positive.
  156. Parameters
  157. ----------
  158. n : int
  159. Window length.
  160. d : scalar, optional
  161. Sample spacing (inverse of the sampling rate). Defaults to 1.
  162. device : str, optional
  163. The device on which to place the created array. Default: ``None``.
  164. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  165. .. versionadded:: 2.0.0
  166. Returns
  167. -------
  168. f : ndarray
  169. Array of length ``n//2 + 1`` containing the sample frequencies.
  170. Examples
  171. --------
  172. >>> import numpy as np
  173. >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5, -3, 4], dtype=float)
  174. >>> fourier = np.fft.rfft(signal)
  175. >>> n = signal.size
  176. >>> sample_rate = 100
  177. >>> freq = np.fft.fftfreq(n, d=1./sample_rate)
  178. >>> freq
  179. array([ 0., 10., 20., ..., -30., -20., -10.])
  180. >>> freq = np.fft.rfftfreq(n, d=1./sample_rate)
  181. >>> freq
  182. array([ 0., 10., 20., 30., 40., 50.])
  183. """
  184. if not isinstance(n, integer_types):
  185. raise ValueError("n should be an integer")
  186. val = 1.0/(n*d)
  187. N = n//2 + 1
  188. results = arange(0, N, dtype=int, device=device)
  189. return results * val