102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
|
from functools import update_wrapper, lru_cache
|
||
|
|
||
|
from ._pocketfft import helper as _helper
|
||
|
|
||
|
|
||
|
def next_fast_len(target, real=False):
|
||
|
"""Find the next fast size of input data to ``fft``, for zero-padding, etc.
|
||
|
|
||
|
SciPy's FFT algorithms gain their speed by a recursive divide and conquer
|
||
|
strategy. This relies on efficient functions for small prime factors of the
|
||
|
input length. Thus, the transforms are fastest when using composites of the
|
||
|
prime factors handled by the fft implementation. If there are efficient
|
||
|
functions for all radices <= `n`, then the result will be a number `x`
|
||
|
>= ``target`` with only prime factors < `n`. (Also known as `n`-smooth
|
||
|
numbers)
|
||
|
|
||
|
Parameters
|
||
|
----------
|
||
|
target : int
|
||
|
Length to start searching from. Must be a positive integer.
|
||
|
real : bool, optional
|
||
|
True if the FFT involves real input or output (e.g., `rfft` or `hfft`
|
||
|
but not `fft`). Defaults to False.
|
||
|
|
||
|
Returns
|
||
|
-------
|
||
|
out : int
|
||
|
The smallest fast length greater than or equal to ``target``.
|
||
|
|
||
|
Notes
|
||
|
-----
|
||
|
The result of this function may change in future as performance
|
||
|
considerations change, for example, if new prime factors are added.
|
||
|
|
||
|
Calling `fft` or `ifft` with real input data performs an ``'R2C'``
|
||
|
transform internally.
|
||
|
|
||
|
Examples
|
||
|
--------
|
||
|
On a particular machine, an FFT of prime length takes 11.4 ms:
|
||
|
|
||
|
>>> from scipy import fft
|
||
|
>>> import numpy as np
|
||
|
>>> rng = np.random.default_rng()
|
||
|
>>> min_len = 93059 # prime length is worst case for speed
|
||
|
>>> a = rng.standard_normal(min_len)
|
||
|
>>> b = fft.fft(a)
|
||
|
|
||
|
Zero-padding to the next regular length reduces computation time to
|
||
|
1.6 ms, a speedup of 7.3 times:
|
||
|
|
||
|
>>> fft.next_fast_len(min_len, real=True)
|
||
|
93312
|
||
|
>>> b = fft.fft(a, 93312)
|
||
|
|
||
|
Rounding up to the next power of 2 is not optimal, taking 3.0 ms to
|
||
|
compute; 1.9 times longer than the size given by ``next_fast_len``:
|
||
|
|
||
|
>>> b = fft.fft(a, 131072)
|
||
|
|
||
|
"""
|
||
|
pass
|
||
|
|
||
|
|
||
|
# Directly wrap the c-function good_size but take the docstring etc., from the
|
||
|
# next_fast_len function above
|
||
|
next_fast_len = update_wrapper(lru_cache()(_helper.good_size), next_fast_len)
|
||
|
next_fast_len.__wrapped__ = _helper.good_size
|
||
|
|
||
|
|
||
|
def _init_nd_shape_and_axes(x, shape, axes):
|
||
|
"""Handle shape and axes arguments for N-D transforms.
|
||
|
|
||
|
Returns the shape and axes in a standard form, taking into account negative
|
||
|
values and checking for various potential errors.
|
||
|
|
||
|
Parameters
|
||
|
----------
|
||
|
x : array_like
|
||
|
The input array.
|
||
|
shape : int or array_like of ints or None
|
||
|
The shape of the result. If both `shape` and `axes` (see below) are
|
||
|
None, `shape` is ``x.shape``; if `shape` is None but `axes` is
|
||
|
not None, then `shape` is ``numpy.take(x.shape, axes, axis=0)``.
|
||
|
If `shape` is -1, the size of the corresponding dimension of `x` is
|
||
|
used.
|
||
|
axes : int or array_like of ints or None
|
||
|
Axes along which the calculation is computed.
|
||
|
The default is over all axes.
|
||
|
Negative indices are automatically converted to their positive
|
||
|
counterparts.
|
||
|
|
||
|
Returns
|
||
|
-------
|
||
|
shape : array
|
||
|
The shape of the result. It is a 1-D integer array.
|
||
|
axes : array
|
||
|
The shape of the result. It is a 1-D integer array.
|
||
|
|
||
|
"""
|
||
|
return _helper._init_nd_shape_and_axes(x, shape, axes)
|