/
_helper.py
51 lines (38 loc) · 1.44 KB
/
_helper.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
from typing import Dict, List, Tuple
import math
_next_fast_len_cache: Dict[Tuple[int, List[int]], int] = {}
def _next_fast_len_impl(n, primes):
if len(primes) == 0:
return math.inf
result = _next_fast_len_cache.get((n, primes), None)
if result is None:
if n == 1:
result = 1
else:
p = primes[0]
result = min(
_next_fast_len_impl((n + p - 1) // p, primes) * p,
_next_fast_len_impl(n, primes[1:]))
_next_fast_len_cache[(n, primes)] = result
return result
def next_fast_len(target, real=False):
"""Find the next fast size to ``fft``.
Args:
target (int): The size of input array.
real (bool): ``True`` if the FFT involves real input or output.
This parameter is of no use, and only for compatibility to
SciPy's interface.
Returns:
int: The smallest fast length greater than or equal to the input value.
.. seealso:: :func:`scipy.fft.next_fast_len`
.. note::
It may return a different value to :func:`scipy.fft.next_fast_len`
as pocketfft's prime factors are different from cuFFT's factors.
For details, see the `cuFFT documentation`_.
.. _cuFFT documentation:
https://docs.nvidia.com/cuda/cufft/index.html#accuracy-and-performance
"""
if target == 0:
return 0
primes = (2, 3, 5, 7)
return _next_fast_len_impl(target, primes)