Create fft.py

daedalus · Nov 3, 2023 · aa1b315 · aa1b315
1 parent 02087e3
commit aa1b315
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/fft.py b/fft.py
@@ -0,0 +1,17 @@
+# Author Dario Clavijo 2023
+# based on https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm
+from gmpy2 import *
+pi_i = (-2j) * complex(gmpy2.const_pi())
+
+def fft(X):
+  if (N:=len(X)) == 1: 
+    return X
+  else:
+    N2 = N >> 1
+    E = fft(X[::2])
+    O = fft(X[1::2])
+    for k in range(0, N2):
+      q = exp((pi_i * k) / N ) * O[k]
+      X[k] = E[k] + q
+      X[k + N2] = E[k] - q
+  return X