StarMirror/algorithm-in-python
1
0
Fork 0
mirror of https://github.com/heqin-zhu/algorithm.git synced 2024-03-22 13:30:46 +08:00
Code Issues Projects Releases Wiki Activity

Update fft: iteration version

Browse Source
This commit is contained in:
mbinary 2019-06-11 12:48:40 +08:00
parent 26b50a1c81
commit df3f138550
1 changed files with 56 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
math/fft.py
View File

@ -1,45 +1,91 @@
import numpy as np
def _fft_n2(a, invert):
'''O(n^2)'''
N = len(a)
w = np.arange(N)
i = 2j if invert else -2j
m = w.reshape((N, 1)) * w
W = np.exp(m * i * np.pi / N)
return np.concatenate(np.dot(W, a.reshape((N,
1)))) # important, cannot use *
def _fft(a, invert=False):
'''fft, len(a) is power of two'''
'''recursion version'''
N = len(a)
if N == 1:
return [a[0]]
else:
elif N & (N - 1) == 0: # O(nlogn), 2^k
even = _fft(a[::2], invert)
odd = _fft(a[1::2], invert)
i = 2j if invert else -2j
factor = np.exp(i * np.pi * np.arange(N // 2) / N)
prod = factor * odd
return np.concatenate([even + prod, even - prod])
else:
return _fft_n2(a, invert)
def _fft2(a, invert=False):
''' iteration version'''
def rev(x):
ret = 0
for i in range(r):
ret <<= 1
if x & 1:
ret += 1
x >>= 1
return ret
N = len(a)
if N & (N - 1) == 0: # O(nlogn), 2^k
r = int(np.log(N))
c = np.array(a,dtype='complex')
i = 2j if invert else -2j
w = np.exp(i * np.pi / N)
for h in range(r - 1, -1, -1):
p = 2**h
z = w**(N / p / 2)
for k in range(N):
if k % p == k % (2 * p):
c[k], c[k + p] = c[k] + c[k + p], c[k] * z**(k % p)
return np.asarray([c[rev(i)] for i in range(N)])
else: # O(n^2)
return _fft_n2(a, invert)
def fft(a):
'''fourier[a]'''
n = len(a)
if n == 0 or n&(n-1)!=0:
raise Exception(f"[Error]: {n} is not power of 2")
if n == 0:
raise Exception("[Error]: Invalid length: 0")
return _fft(a)
def ifft(a):
'''invert fourier[a]'''
n = len(a)
if n == 0 or n&(n-1)!=0:
raise Exception(f"[Error]: {n} is not power of 2")
if n == 0:
raise Exception("[Error]: Invalid length: 0")
return _fft(a, True) / n
def fft2(arr):
return np.apply_along_axis(fft, 0, np.apply_along_axis(fft, 1, np.asarray(arr)))
return np.apply_along_axis(fft, 0,
np.apply_along_axis(fft, 1, np.asarray(arr)))
def ifft2(arr):
return np.apply_along_axis(ifft, 0, np.apply_along_axis(ifft, 1, np.asarray(arr)))
return np.apply_along_axis(ifft, 0,
np.apply_along_axis(ifft, 1, np.asarray(arr)))
def test(n=128):
print('\nsequence length:', n)
print('fft')
li = np.random.random(n)
print(np.allclose(fft(li), np.fft.fft(li)))
@ -58,4 +104,5 @@ def test(n=128):
if __name__ == '__main__':
test(128)
for i in range(1, 4):
test(i * 16)