add FFT and tests (keon#847)

ekorre1001 · web-flow · commit 51f93c600bd2 · 2023-03-08T01:33:02.000-08:00
diff --git a/algorithms/maths/fft.py b/algorithms/maths/fft.py
@@ -0,0 +1,32 @@
+"""
+Implementation of the Cooley-Tukey, which is the most common FFT algorithm.
+
+Input: an array of complex values which has a size of N, where N is an integer power of 2
+Output: an array of complex values which is the discrete fourier transform of the input
+
+Example 1
+Input: [2.0+2j, 1.0+3j, 3.0+1j, 2.0+2j]
+Output: [8+8j, 2j, 2-2j, -2+0j]
+
+
+Pseudocode: https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm
+"""
+from cmath import exp, pi
+
+def fft(x):
+    """ Recursive implementation of the Cooley-Tukey"""
+    N = len(x)
+    if N == 1:
+        return x
+
+    # get the elements at even/odd indices
+    even = fft(x[0::2])
+    odd = fft(x[1::2])
+
+    y = [0 for i in range(N)]
+    for k in range(N//2):
+        q = exp(-2j*pi*k/N)*odd[k]
+        y[k] = even[k] + q
+        y[k + N//2] = even[k] - q
+
+    return y
diff --git a/tests/test_maths.py b/tests/test_maths.py
@@ -26,6 +26,7 @@
     diffie_hellman_key_exchange, krishnamurthy_number,
     num_perfect_squares,
     chinese_remainder_theorem,
+    fft
 )
 
 import unittest
@@ -556,5 +557,41 @@ def test_empty_lists(self):
             chinese_remainder_theorem.solve_chinese_remainder(num, rem)
 
 
+class TestFFT(unittest.TestCase):
+    """[summary]
+    Test for the file fft.py
+
+    Arguments:
+        unittest {[type]} -- [description]
+    """
+    def test_real_numbers(self):
+        x = [1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]
+        y = [4.000, 2.613, 0.000, 1.082, 0.000, 1.082, 0.000, 2.613]
+        # abs(complex) returns the magnitude
+        result = [float("%.3f" % abs(f)) for f in fft.fft(x)]
+        self.assertEqual(result, y)
+    
+    def test_all_zero(self):
+        x = [0.0, 0.0, 0.0, 0.0]
+        y = [0.0, 0.0, 0.0, 0.0]
+        result = [float("%.1f" % abs(f)) for f in fft.fft(x)]
+        self.assertEqual(result, y)
+    
+    def test_all_ones(self):
+        x = [1.0, 1.0, 1.0, 1.0]
+        y = [4.0, 0.0, 0.0, 0.0]
+        result = [float("%.1f" % abs(f)) for f in fft.fft(x)]
+        self.assertEqual(result, y)
+
+    def test_complex_numbers(self):
+        x = [2.0+2j, 1.0+3j, 3.0+1j, 2.0+2j]
+        real = [8.0, 0.0, 2.0, -2.0]
+        imag = [8.0, 2.0, -2.0, 0.0]
+        realResult = [float("%.1f" % f.real) for f in fft.fft(x)]
+        imagResult = [float("%.1f" % f.imag) for f in fft.fft(x)]
+        self.assertEqual(real, realResult)
+        self.assertEqual(imag, imagResult)
+        
+
 if __name__ == "__main__":
     unittest.main()
