/
helper.py
149 lines (119 loc) · 4.76 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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
from __future__ import division, print_function, absolute_import
from numpy import arange
from numpy.fft.helper import fftshift, ifftshift, fftfreq
from bisect import bisect_left
__all__ = ['fftshift', 'ifftshift', 'fftfreq', 'rfftfreq', 'next_fast_len']
def rfftfreq(n, d=1.0):
"""DFT sample frequencies (for usage with rfft, irfft).
The returned float array contains the frequency bins in
cycles/unit (with zero at the start) given a window length `n` and a
sample spacing `d`::
f = [0,1,1,2,2,...,n/2-1,n/2-1,n/2]/(d*n) if n is even
f = [0,1,1,2,2,...,n/2-1,n/2-1,n/2,n/2]/(d*n) if n is odd
Parameters
----------
n : int
Window length.
d : scalar, optional
Sample spacing. Default is 1.
Returns
-------
out : ndarray
The array of length `n`, containing the sample frequencies.
Examples
--------
>>> from scipy import fftpack
>>> sig = np.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)
>>> sig_fft = fftpack.rfft(sig)
>>> n = sig_fft.size
>>> timestep = 0.1
>>> freq = fftpack.rfftfreq(n, d=timestep)
>>> freq
array([ 0. , 1.25, 1.25, 2.5 , 2.5 , 3.75, 3.75, 5. ])
"""
if not isinstance(n, int) or n < 0:
raise ValueError("n = %s is not valid. "
"n must be a nonnegative integer." % n)
return (arange(1, n + 1, dtype=int) // 2) / float(n * d)
def next_fast_len(target):
"""
Find the next fast size of input data to `fft`, for zero-padding, etc.
SciPy's FFTPACK has efficient functions for radix {2, 3, 4, 5}, so this
returns the next composite of the prime factors 2, 3, and 5 which is
greater than or equal to `target`. (These are also known as 5-smooth
numbers, regular numbers, or Hamming numbers.)
Parameters
----------
target : int
Length to start searching from. Must be a positive integer.
Returns
-------
out : int
The first 5-smooth number greater than or equal to `target`.
Notes
-----
.. versionadded:: 0.18.0
Examples
--------
On a particular machine, an FFT of prime length takes 133 ms:
>>> from scipy import fftpack
>>> min_len = 10007 # prime length is worst case for speed
>>> a = np.random.randn(min_len)
>>> b = fftpack.fft(a)
Zero-padding to the next 5-smooth length reduces computation time to
211 us, a speedup of 630 times:
>>> fftpack.helper.next_fast_len(min_len)
10125
>>> b = fftpack.fft(a, 10125)
Rounding up to the next power of 2 is not optimal, taking 367 us to
compute, 1.7 times as long as the 5-smooth size:
>>> b = fftpack.fft(a, 16384)
"""
hams = (8, 9, 10, 12, 15, 16, 18, 20, 24, 25, 27, 30, 32, 36, 40, 45, 48,
50, 54, 60, 64, 72, 75, 80, 81, 90, 96, 100, 108, 120, 125, 128,
135, 144, 150, 160, 162, 180, 192, 200, 216, 225, 240, 243, 250,
256, 270, 288, 300, 320, 324, 360, 375, 384, 400, 405, 432, 450,
480, 486, 500, 512, 540, 576, 600, 625, 640, 648, 675, 720, 729,
750, 768, 800, 810, 864, 900, 960, 972, 1000, 1024, 1080, 1125,
1152, 1200, 1215, 1250, 1280, 1296, 1350, 1440, 1458, 1500, 1536,
1600, 1620, 1728, 1800, 1875, 1920, 1944, 2000, 2025, 2048, 2160,
2187, 2250, 2304, 2400, 2430, 2500, 2560, 2592, 2700, 2880, 2916,
3000, 3072, 3125, 3200, 3240, 3375, 3456, 3600, 3645, 3750, 3840,
3888, 4000, 4050, 4096, 4320, 4374, 4500, 4608, 4800, 4860, 5000,
5120, 5184, 5400, 5625, 5760, 5832, 6000, 6075, 6144, 6250, 6400,
6480, 6561, 6750, 6912, 7200, 7290, 7500, 7680, 7776, 8000, 8100,
8192, 8640, 8748, 9000, 9216, 9375, 9600, 9720, 10000)
if target <= 6:
return target
# Quickly check if it's already a power of 2
if not (target & (target-1)):
return target
# Get result quickly for small sizes, since FFT itself is similarly fast.
if target <= hams[-1]:
return hams[bisect_left(hams, target)]
match = float('inf') # Anything found will be smaller
p5 = 1
while p5 < target:
p35 = p5
while p35 < target:
# Ceiling integer division, avoiding conversion to float
# (quotient = ceil(target / p35))
quotient = -(-target // p35)
# Quickly find next power of 2 >= quotient
p2 = 2**((quotient - 1).bit_length())
N = p2 * p35
if N == target:
return N
elif N < match:
match = N
p35 *= 3
if p35 == target:
return p35
if p35 < match:
match = p35
p5 *= 5
if p5 == target:
return p5
if p5 < match:
match = p5
return match