helper.py 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326
  1. """
  2. Discrete Fourier Transforms - helper.py
  3. """
  4. from __future__ import division, absolute_import, print_function
  5. import collections
  6. try:
  7. import threading
  8. except ImportError:
  9. import dummy_threading as threading
  10. from numpy.compat import integer_types
  11. from numpy.core import integer, empty, arange, asarray, roll
  12. from numpy.core.overrides import array_function_dispatch, set_module
  13. # Created by Pearu Peterson, September 2002
  14. __all__ = ['fftshift', 'ifftshift', 'fftfreq', 'rfftfreq']
  15. integer_types = integer_types + (integer,)
  16. def _fftshift_dispatcher(x, axes=None):
  17. return (x,)
  18. @array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')
  19. def fftshift(x, axes=None):
  20. """
  21. Shift the zero-frequency component to the center of the spectrum.
  22. This function swaps half-spaces for all axes listed (defaults to all).
  23. Note that ``y[0]`` is the Nyquist component only if ``len(x)`` is even.
  24. Parameters
  25. ----------
  26. x : array_like
  27. Input array.
  28. axes : int or shape tuple, optional
  29. Axes over which to shift. Default is None, which shifts all axes.
  30. Returns
  31. -------
  32. y : ndarray
  33. The shifted array.
  34. See Also
  35. --------
  36. ifftshift : The inverse of `fftshift`.
  37. Examples
  38. --------
  39. >>> freqs = np.fft.fftfreq(10, 0.1)
  40. >>> freqs
  41. array([ 0., 1., 2., 3., 4., -5., -4., -3., -2., -1.])
  42. >>> np.fft.fftshift(freqs)
  43. array([-5., -4., -3., -2., -1., 0., 1., 2., 3., 4.])
  44. Shift the zero-frequency component only along the second axis:
  45. >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
  46. >>> freqs
  47. array([[ 0., 1., 2.],
  48. [ 3., 4., -4.],
  49. [-3., -2., -1.]])
  50. >>> np.fft.fftshift(freqs, axes=(1,))
  51. array([[ 2., 0., 1.],
  52. [-4., 3., 4.],
  53. [-1., -3., -2.]])
  54. """
  55. x = asarray(x)
  56. if axes is None:
  57. axes = tuple(range(x.ndim))
  58. shift = [dim // 2 for dim in x.shape]
  59. elif isinstance(axes, integer_types):
  60. shift = x.shape[axes] // 2
  61. else:
  62. shift = [x.shape[ax] // 2 for ax in axes]
  63. return roll(x, shift, axes)
  64. @array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')
  65. def ifftshift(x, axes=None):
  66. """
  67. The inverse of `fftshift`. Although identical for even-length `x`, the
  68. functions differ by one sample for odd-length `x`.
  69. Parameters
  70. ----------
  71. x : array_like
  72. Input array.
  73. axes : int or shape tuple, optional
  74. Axes over which to calculate. Defaults to None, which shifts all axes.
  75. Returns
  76. -------
  77. y : ndarray
  78. The shifted array.
  79. See Also
  80. --------
  81. fftshift : Shift zero-frequency component to the center of the spectrum.
  82. Examples
  83. --------
  84. >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
  85. >>> freqs
  86. array([[ 0., 1., 2.],
  87. [ 3., 4., -4.],
  88. [-3., -2., -1.]])
  89. >>> np.fft.ifftshift(np.fft.fftshift(freqs))
  90. array([[ 0., 1., 2.],
  91. [ 3., 4., -4.],
  92. [-3., -2., -1.]])
  93. """
  94. x = asarray(x)
  95. if axes is None:
  96. axes = tuple(range(x.ndim))
  97. shift = [-(dim // 2) for dim in x.shape]
  98. elif isinstance(axes, integer_types):
  99. shift = -(x.shape[axes] // 2)
  100. else:
  101. shift = [-(x.shape[ax] // 2) for ax in axes]
  102. return roll(x, shift, axes)
  103. @set_module('numpy.fft')
  104. def fftfreq(n, d=1.0):
  105. """
  106. Return the Discrete Fourier Transform sample frequencies.
  107. The returned float array `f` contains the frequency bin centers in cycles
  108. per unit of the sample spacing (with zero at the start). For instance, if
  109. the sample spacing is in seconds, then the frequency unit is cycles/second.
  110. Given a window length `n` and a sample spacing `d`::
  111. f = [0, 1, ..., n/2-1, -n/2, ..., -1] / (d*n) if n is even
  112. f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n) if n is odd
  113. Parameters
  114. ----------
  115. n : int
  116. Window length.
  117. d : scalar, optional
  118. Sample spacing (inverse of the sampling rate). Defaults to 1.
  119. Returns
  120. -------
  121. f : ndarray
  122. Array of length `n` containing the sample frequencies.
  123. Examples
  124. --------
  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, -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)
  137. N = (n-1)//2 + 1
  138. p1 = arange(0, N, dtype=int)
  139. results[:N] = p1
  140. p2 = arange(-(n//2), 0, dtype=int)
  141. results[N:] = p2
  142. return results * val
  143. @set_module('numpy.fft')
  144. def rfftfreq(n, d=1.0):
  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. Returns
  163. -------
  164. f : ndarray
  165. Array of length ``n//2 + 1`` containing the sample frequencies.
  166. Examples
  167. --------
  168. >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5, -3, 4], dtype=float)
  169. >>> fourier = np.fft.rfft(signal)
  170. >>> n = signal.size
  171. >>> sample_rate = 100
  172. >>> freq = np.fft.fftfreq(n, d=1./sample_rate)
  173. >>> freq
  174. array([ 0., 10., 20., 30., 40., -50., -40., -30., -20., -10.])
  175. >>> freq = np.fft.rfftfreq(n, d=1./sample_rate)
  176. >>> freq
  177. array([ 0., 10., 20., 30., 40., 50.])
  178. """
  179. if not isinstance(n, integer_types):
  180. raise ValueError("n should be an integer")
  181. val = 1.0/(n*d)
  182. N = n//2 + 1
  183. results = arange(0, N, dtype=int)
  184. return results * val
  185. class _FFTCache(object):
  186. """
  187. Cache for the FFT twiddle factors as an LRU (least recently used) cache.
  188. Parameters
  189. ----------
  190. max_size_in_mb : int
  191. Maximum memory usage of the cache before items are being evicted.
  192. max_item_count : int
  193. Maximum item count of the cache before items are being evicted.
  194. Notes
  195. -----
  196. Items will be evicted if either limit has been reached upon getting and
  197. setting. The maximum memory usages is not strictly the given
  198. ``max_size_in_mb`` but rather
  199. ``max(max_size_in_mb, 1.5 * size_of_largest_item)``. Thus the cache will
  200. never be completely cleared - at least one item will remain and a single
  201. large item can cause the cache to retain several smaller items even if the
  202. given maximum cache size has been exceeded.
  203. """
  204. def __init__(self, max_size_in_mb, max_item_count):
  205. self._max_size_in_bytes = max_size_in_mb * 1024 ** 2
  206. self._max_item_count = max_item_count
  207. self._dict = collections.OrderedDict()
  208. self._lock = threading.Lock()
  209. def put_twiddle_factors(self, n, factors):
  210. """
  211. Store twiddle factors for an FFT of length n in the cache.
  212. Putting multiple twiddle factors for a certain n will store it multiple
  213. times.
  214. Parameters
  215. ----------
  216. n : int
  217. Data length for the FFT.
  218. factors : ndarray
  219. The actual twiddle values.
  220. """
  221. with self._lock:
  222. # Pop + later add to move it to the end for LRU behavior.
  223. # Internally everything is stored in a dictionary whose values are
  224. # lists.
  225. try:
  226. value = self._dict.pop(n)
  227. except KeyError:
  228. value = []
  229. value.append(factors)
  230. self._dict[n] = value
  231. self._prune_cache()
  232. def pop_twiddle_factors(self, n):
  233. """
  234. Pop twiddle factors for an FFT of length n from the cache.
  235. Will return None if the requested twiddle factors are not available in
  236. the cache.
  237. Parameters
  238. ----------
  239. n : int
  240. Data length for the FFT.
  241. Returns
  242. -------
  243. out : ndarray or None
  244. The retrieved twiddle factors if available, else None.
  245. """
  246. with self._lock:
  247. if n not in self._dict or not self._dict[n]:
  248. return None
  249. # Pop + later add to move it to the end for LRU behavior.
  250. all_values = self._dict.pop(n)
  251. value = all_values.pop()
  252. # Only put pack if there are still some arrays left in the list.
  253. if all_values:
  254. self._dict[n] = all_values
  255. return value
  256. def _prune_cache(self):
  257. # Always keep at least one item.
  258. while len(self._dict) > 1 and (
  259. len(self._dict) > self._max_item_count or self._check_size()):
  260. self._dict.popitem(last=False)
  261. def _check_size(self):
  262. item_sizes = [sum(_j.nbytes for _j in _i)
  263. for _i in self._dict.values() if _i]
  264. if not item_sizes:
  265. return False
  266. max_size = max(self._max_size_in_bytes, 1.5 * max(item_sizes))
  267. return sum(item_sizes) > max_size