simple dft example with NodeBox

dft1.py

from cmath import exp, sin, cos, pi

N = 512
PADDING = 16
global old_x, old_y
old_x, old_y = 0, N

def circle(x, y, r):
    hr = r / 2
    oval(x - hr, y - hr, r, r)

def plot(x, y, is_line=False):
    global old_x, old_y
    if isinstance(x, complex):
        x = x.real
    if isinstance(y, complex):
        y = y.real

    if is_line:
        y = N - (y * N / 2.0)
        line(old_x, old_y, x, y)
        old_x, old_y = x, y
    else:
        circle(x, y * N / 2.0 + N / 2.0, 3)

def init_input():
    ary = []
    for i in range(N):
        ary.append(sin(32.77 * pi * i / N) * 0.5 + 
                   cos(67 * pi * i / N) * 0.5)
    return tuple(ary)

def draw(ary):
    stroke(0.8)
    translate(PADDING, PADDING)
    for i in range(len(ary)):
        plot(i, ary[i])
    translate(-PADDING, -PADDING)

def draw_power(ary):
    x_power = map(lambda x: abs(x) ** 2, ary)
    power_max = max(x_power)
    stroke(0.3, 0.4, 1.0)
    translate(PADDING, PADDING)

    for i, x in enumerate(ary):
        x_power = abs(x) ** 2
        plot(i, x_power / power_max * 2, is_line=True)

    translate(-PADDING, -PADDING)

def dft(ary):
    out = []
    for k in range(N):
        xk = 0
        for p, xp in enumerate(ary):
            xk += xp * exp(-2j * pi * p * k / N)
        out.append(xk)
    return out

# main
size(N + PADDING * 2, N + PADDING * 2)
print "processing..."
fill(0.2)

x_array = init_input()
out_array = dft(x_array)
draw_power(out_array)
draw(x_array)