fft.py

import math
import cmath

tau = 2 * cmath.pi

def wave(n, f):
    res = []
    j = 0
    while (j < n):
        res.append(math.sin(tau * f * j))
        j += 1
    return res

def isPowerOfTwo(n):
    if (n == 0):
        return False
    else:
        while (1 < n):
            if (n % 2 != 0):
                return False
            n = n >> 1
        return True

def cooley(x, n, off):
    if n <= 1: return
    m = n // 2
    theta = tau / n
    for s in range(0, m):
        omega = complex(cmath.cos(s * theta), -cmath.sin(s * theta))
        a = x[s + off + 0]
        b = x[s + off + m]
        x[s + off + 0] = a + b
        x[s + off + m] = (a - b) * omega
    cooley(x, m, off + 0)
    cooley(x, m, off + m)


def fft(x):
    n = len(x)
    assert isPowerOfTwo(n), "length of input must be a power of two"
    cooley(x, n, 0)

def test(x):
    print(x)
    fft(x)
    print(x)

seq = wave(8, 1)
test(seq)
# test([0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1])