From 26b50a1c814284af50aa05a7746e334e0ce4cf6c Mon Sep 17 00:00:00 2001
From: mbinary <zhuheqin1@gmail.com>
Date: Tue, 11 Jun 2019 10:51:46 +0800
Subject: [PATCH] :sparkles: Add fft

---
 math/fft.py | 61 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)
 create mode 100644 math/fft.py

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)