_helper.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. import operator
  2. from numpy.fft.helper import fftshift, ifftshift, fftfreq
  3. import scipy.fft._pocketfft.helper as _helper
  4. import numpy as np
  5. __all__ = ['fftshift', 'ifftshift', 'fftfreq', 'rfftfreq', 'next_fast_len']
  6. def rfftfreq(n, d=1.0):
  7. """DFT sample frequencies (for usage with rfft, irfft).
  8. The returned float array contains the frequency bins in
  9. cycles/unit (with zero at the start) given a window length `n` and a
  10. sample spacing `d`::
  11. f = [0,1,1,2,2,...,n/2-1,n/2-1,n/2]/(d*n) if n is even
  12. f = [0,1,1,2,2,...,n/2-1,n/2-1,n/2,n/2]/(d*n) if n is odd
  13. Parameters
  14. ----------
  15. n : int
  16. Window length.
  17. d : scalar, optional
  18. Sample spacing. Default is 1.
  19. Returns
  20. -------
  21. out : ndarray
  22. The array of length `n`, containing the sample frequencies.
  23. Examples
  24. --------
  25. >>> import numpy as np
  26. >>> from scipy import fftpack
  27. >>> sig = np.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)
  28. >>> sig_fft = fftpack.rfft(sig)
  29. >>> n = sig_fft.size
  30. >>> timestep = 0.1
  31. >>> freq = fftpack.rfftfreq(n, d=timestep)
  32. >>> freq
  33. array([ 0. , 1.25, 1.25, 2.5 , 2.5 , 3.75, 3.75, 5. ])
  34. """
  35. n = operator.index(n)
  36. if n < 0:
  37. raise ValueError("n = %s is not valid. "
  38. "n must be a nonnegative integer." % n)
  39. return (np.arange(1, n + 1, dtype=int) // 2) / float(n * d)
  40. def next_fast_len(target):
  41. """
  42. Find the next fast size of input data to `fft`, for zero-padding, etc.
  43. SciPy's FFTPACK has efficient functions for radix {2, 3, 4, 5}, so this
  44. returns the next composite of the prime factors 2, 3, and 5 which is
  45. greater than or equal to `target`. (These are also known as 5-smooth
  46. numbers, regular numbers, or Hamming numbers.)
  47. Parameters
  48. ----------
  49. target : int
  50. Length to start searching from. Must be a positive integer.
  51. Returns
  52. -------
  53. out : int
  54. The first 5-smooth number greater than or equal to `target`.
  55. Notes
  56. -----
  57. .. versionadded:: 0.18.0
  58. Examples
  59. --------
  60. On a particular machine, an FFT of prime length takes 133 ms:
  61. >>> from scipy import fftpack
  62. >>> import numpy as np
  63. >>> rng = np.random.default_rng()
  64. >>> min_len = 10007 # prime length is worst case for speed
  65. >>> a = rng.standard_normal(min_len)
  66. >>> b = fftpack.fft(a)
  67. Zero-padding to the next 5-smooth length reduces computation time to
  68. 211 us, a speedup of 630 times:
  69. >>> fftpack.next_fast_len(min_len)
  70. 10125
  71. >>> b = fftpack.fft(a, 10125)
  72. Rounding up to the next power of 2 is not optimal, taking 367 us to
  73. compute, 1.7 times as long as the 5-smooth size:
  74. >>> b = fftpack.fft(a, 16384)
  75. """
  76. # Real transforms use regular sizes so this is backwards compatible
  77. return _helper.good_size(target, True)
  78. def _good_shape(x, shape, axes):
  79. """Ensure that shape argument is valid for scipy.fftpack
  80. scipy.fftpack does not support len(shape) < x.ndim when axes is not given.
  81. """
  82. if shape is not None and axes is None:
  83. shape = _helper._iterable_of_int(shape, 'shape')
  84. if len(shape) != np.ndim(x):
  85. raise ValueError("when given, axes and shape arguments"
  86. " have to be of the same length")
  87. return shape