diff --git a/math/fft.py b/math/fft.py new file mode 100644 index 0000000..5d6b9c6 --- /dev/null +++ b/math/fft.py @@ -0,0 +1,61 @@ +import numpy as np + + +def _fft(a, invert=False): + '''fft, len(a) is power of two''' + N = len(a) + if N == 1: + return [a[0]] + else: + 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]) + + +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") + 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") + return _fft(a, True) / n + + +def fft2(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))) + + +def test(n=128): + print('fft') + li = np.random.random(n) + print(np.allclose(fft(li), np.fft.fft(li))) + + print('ifft') + li = np.random.random(n) + print(np.allclose(ifft(li), np.fft.ifft(li))) + + print('fft2') + li = np.random.random(n * n).reshape((n, n)) + print(np.allclose(fft2(li), np.fft.fft2(li))) + + print('ifft2') + li = np.random.random(n * n).reshape((n, n)) + print(np.allclose(ifft2(li), np.fft.ifft2(li))) + + +if __name__ == '__main__': + test(128)