_helper.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. from functools import update_wrapper, lru_cache
  2. from ._pocketfft import helper as _helper
  3. def next_fast_len(target, real=False):
  4. """Find the next fast size of input data to ``fft``, for zero-padding, etc.
  5. SciPy's FFT algorithms gain their speed by a recursive divide and conquer
  6. strategy. This relies on efficient functions for small prime factors of the
  7. input length. Thus, the transforms are fastest when using composites of the
  8. prime factors handled by the fft implementation. If there are efficient
  9. functions for all radices <= `n`, then the result will be a number `x`
  10. >= ``target`` with only prime factors < `n`. (Also known as `n`-smooth
  11. numbers)
  12. Parameters
  13. ----------
  14. target : int
  15. Length to start searching from. Must be a positive integer.
  16. real : bool, optional
  17. True if the FFT involves real input or output (e.g., `rfft` or `hfft`
  18. but not `fft`). Defaults to False.
  19. Returns
  20. -------
  21. out : int
  22. The smallest fast length greater than or equal to ``target``.
  23. Notes
  24. -----
  25. The result of this function may change in future as performance
  26. considerations change, for example, if new prime factors are added.
  27. Calling `fft` or `ifft` with real input data performs an ``'R2C'``
  28. transform internally.
  29. Examples
  30. --------
  31. On a particular machine, an FFT of prime length takes 11.4 ms:
  32. >>> from scipy import fft
  33. >>> import numpy as np
  34. >>> rng = np.random.default_rng()
  35. >>> min_len = 93059 # prime length is worst case for speed
  36. >>> a = rng.standard_normal(min_len)
  37. >>> b = fft.fft(a)
  38. Zero-padding to the next regular length reduces computation time to
  39. 1.6 ms, a speedup of 7.3 times:
  40. >>> fft.next_fast_len(min_len, real=True)
  41. 93312
  42. >>> b = fft.fft(a, 93312)
  43. Rounding up to the next power of 2 is not optimal, taking 3.0 ms to
  44. compute; 1.9 times longer than the size given by ``next_fast_len``:
  45. >>> b = fft.fft(a, 131072)
  46. """
  47. pass
  48. # Directly wrap the c-function good_size but take the docstring etc., from the
  49. # next_fast_len function above
  50. next_fast_len = update_wrapper(lru_cache()(_helper.good_size), next_fast_len)
  51. next_fast_len.__wrapped__ = _helper.good_size
  52. def _init_nd_shape_and_axes(x, shape, axes):
  53. """Handle shape and axes arguments for N-D transforms.
  54. Returns the shape and axes in a standard form, taking into account negative
  55. values and checking for various potential errors.
  56. Parameters
  57. ----------
  58. x : array_like
  59. The input array.
  60. shape : int or array_like of ints or None
  61. The shape of the result. If both `shape` and `axes` (see below) are
  62. None, `shape` is ``x.shape``; if `shape` is None but `axes` is
  63. not None, then `shape` is ``numpy.take(x.shape, axes, axis=0)``.
  64. If `shape` is -1, the size of the corresponding dimension of `x` is
  65. used.
  66. axes : int or array_like of ints or None
  67. Axes along which the calculation is computed.
  68. The default is over all axes.
  69. Negative indices are automatically converted to their positive
  70. counterparts.
  71. Returns
  72. -------
  73. shape : array
  74. The shape of the result. It is a 1-D integer array.
  75. axes : array
  76. The shape of the result. It is a 1-D integer array.
  77. """
  78. return _helper._init_nd_shape_and_axes(x, shape, axes)