helper.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. """
  2. Discrete Fourier Transforms - helper.py
  3. """
  4. from numpy.compat import integer_types
  5. from numpy.core import integer, empty, arange, asarray, roll
  6. from numpy.core.overrides import array_function_dispatch, set_module
  7. # Created by Pearu Peterson, September 2002
  8. __all__ = ['fftshift', 'ifftshift', 'fftfreq', 'rfftfreq']
  9. integer_types = integer_types + (integer,)
  10. def _fftshift_dispatcher(x, axes=None):
  11. return (x,)
  12. @array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')
  13. def fftshift(x, axes=None):
  14. """
  15. Shift the zero-frequency component to the center of the spectrum.
  16. This function swaps half-spaces for all axes listed (defaults to all).
  17. Note that ``y[0]`` is the Nyquist component only if ``len(x)`` is even.
  18. Parameters
  19. ----------
  20. x : array_like
  21. Input array.
  22. axes : int or shape tuple, optional
  23. Axes over which to shift. Default is None, which shifts all axes.
  24. Returns
  25. -------
  26. y : ndarray
  27. The shifted array.
  28. See Also
  29. --------
  30. ifftshift : The inverse of `fftshift`.
  31. Examples
  32. --------
  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. >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
  79. >>> freqs
  80. array([[ 0., 1., 2.],
  81. [ 3., 4., -4.],
  82. [-3., -2., -1.]])
  83. >>> np.fft.ifftshift(np.fft.fftshift(freqs))
  84. array([[ 0., 1., 2.],
  85. [ 3., 4., -4.],
  86. [-3., -2., -1.]])
  87. """
  88. x = asarray(x)
  89. if axes is None:
  90. axes = tuple(range(x.ndim))
  91. shift = [-(dim // 2) for dim in x.shape]
  92. elif isinstance(axes, integer_types):
  93. shift = -(x.shape[axes] // 2)
  94. else:
  95. shift = [-(x.shape[ax] // 2) for ax in axes]
  96. return roll(x, shift, axes)
  97. @set_module('numpy.fft')
  98. def fftfreq(n, d=1.0):
  99. """
  100. Return the Discrete Fourier Transform sample frequencies.
  101. The returned float array `f` contains the frequency bin centers in cycles
  102. per unit of the sample spacing (with zero at the start). For instance, if
  103. the sample spacing is in seconds, then the frequency unit is cycles/second.
  104. Given a window length `n` and a sample spacing `d`::
  105. f = [0, 1, ..., n/2-1, -n/2, ..., -1] / (d*n) if n is even
  106. f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n) if n is odd
  107. Parameters
  108. ----------
  109. n : int
  110. Window length.
  111. d : scalar, optional
  112. Sample spacing (inverse of the sampling rate). Defaults to 1.
  113. Returns
  114. -------
  115. f : ndarray
  116. Array of length `n` containing the sample frequencies.
  117. Examples
  118. --------
  119. >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)
  120. >>> fourier = np.fft.fft(signal)
  121. >>> n = signal.size
  122. >>> timestep = 0.1
  123. >>> freq = np.fft.fftfreq(n, d=timestep)
  124. >>> freq
  125. array([ 0. , 1.25, 2.5 , ..., -3.75, -2.5 , -1.25])
  126. """
  127. if not isinstance(n, integer_types):
  128. raise ValueError("n should be an integer")
  129. val = 1.0 / (n * d)
  130. results = empty(n, int)
  131. N = (n-1)//2 + 1
  132. p1 = arange(0, N, dtype=int)
  133. results[:N] = p1
  134. p2 = arange(-(n//2), 0, dtype=int)
  135. results[N:] = p2
  136. return results * val
  137. @set_module('numpy.fft')
  138. def rfftfreq(n, d=1.0):
  139. """
  140. Return the Discrete Fourier Transform sample frequencies
  141. (for usage with rfft, irfft).
  142. The returned float array `f` contains the frequency bin centers in cycles
  143. per unit of the sample spacing (with zero at the start). For instance, if
  144. the sample spacing is in seconds, then the frequency unit is cycles/second.
  145. Given a window length `n` and a sample spacing `d`::
  146. f = [0, 1, ..., n/2-1, n/2] / (d*n) if n is even
  147. f = [0, 1, ..., (n-1)/2-1, (n-1)/2] / (d*n) if n is odd
  148. Unlike `fftfreq` (but like `scipy.fftpack.rfftfreq`)
  149. the Nyquist frequency component is considered to be positive.
  150. Parameters
  151. ----------
  152. n : int
  153. Window length.
  154. d : scalar, optional
  155. Sample spacing (inverse of the sampling rate). Defaults to 1.
  156. Returns
  157. -------
  158. f : ndarray
  159. Array of length ``n//2 + 1`` containing the sample frequencies.
  160. Examples
  161. --------
  162. >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5, -3, 4], dtype=float)
  163. >>> fourier = np.fft.rfft(signal)
  164. >>> n = signal.size
  165. >>> sample_rate = 100
  166. >>> freq = np.fft.fftfreq(n, d=1./sample_rate)
  167. >>> freq
  168. array([ 0., 10., 20., ..., -30., -20., -10.])
  169. >>> freq = np.fft.rfftfreq(n, d=1./sample_rate)
  170. >>> freq
  171. array([ 0., 10., 20., 30., 40., 50.])
  172. """
  173. if not isinstance(n, integer_types):
  174. raise ValueError("n should be an integer")
  175. val = 1.0/(n*d)
  176. N = n//2 + 1
  177. results = arange(0, N, dtype=int)
  178. return results * val