FFT visualization in Python Mode

noahbuddy

I created this to get more familiar with FFT. Like the Fortran example at the DSP Guide, Python supports complex numbers directly.

Things to note: The forward and inverse FFT are very similar.

Pay close attention to how the sample sets ('signal' and 'wave' arrays) are displayed versus how they were created.

Included comments are my interpretation of the algorithm.

# Length of data to run through FFT
N = 256 # must be a power of 2 for FFT
M = int(log(N) / log(2)) + 1 # grab number of bit levels in N, plus one for range()

"""
Fast Fourier Transform
  Modifies data to represent the Complex Fourier Transform of the input
  Note: The inverse transform will also overwrite data passed to it

Reference: http://www.dspguide.com/ch12/3.htm
"""
def FFT( data ):
  # Sort by bit reversed index
  nd2 = N/2
  j = nd2
  k = 0
  for i in xrange( 1, N-2 ):
    if i < j:
      data[j], data[i] = data[i], data[j] # Pythonic swap
    k = nd2
    while not (k>j):
      j = j - k
      k = k / 2
    j = j + k


  # Bulk of the algorithm from here:
  # For each bit level...
  for L in xrange( 1, M ):

    # Calculate which frequencies to work on this round,
    le = 1<<L
    le2 = le / 2
    # Phase step size at this bit level, AKA: the frequency
    # A complex number
    s = cos(PI/le2) - sin(PI/le2)*1j

    # Init our complex multiplier
    u = 1+0j

    for j in xrange( 1, le2+1 ):
      jm1 = j - 1

      for i in xrange( jm1, N, le ):
        ip = i + le2 # where in data? i and ip

        # Complex multiplication
        # This is what creates constructive or destructive interference
        #   (if sample data is similar to selected frequency, they combine)
        #   (if sample data is _not_ similar to selected frequency, they cancel)
        t = data[ip] * u

        # Positive and Negative frequency bins
        # The FFT is symmetric
        data[ip] = data[i] - t
        data[i] = data[i] + t

      # With each step, rotate multiplier by the frequency step
      # Multiplying complex numbers is easier if you convert them to polar representation first
      #     In polar coordinates add lengths and angles
      #     Convert back to rectangular (complex)
      u = u * s


#
# Inverse FFT
#
def IFFT( data ):
  for i in xrange( len(data) ): # Mirror imaginary values of data
    data[i] = data[i].conjugate()

  FFT( data ); # FFT of mirrored data

  for i in xrange( len(data) ): # Mirror again and scale
    data[i] = data[i].conjugate() / N


# Init samples with complex numbers, length of N
signal = [ (0+0j) ]*N
wave = [ (0+0j) ]*N

# Build a signal in frequency domain
F = 7
signal[F] = 8j
signal[N-F] = -8j

# Create a sine wave in time domain
for i in xrange( N ):
    wave[i] = 0.5*sin( 4*(i*PI)/N )

# Arbitrary drawing size
scl = 1

def setup():
    size(512,512,P3D)
    noFill()
    FFT(signal)
    FFT(wave)
    noLoop()

def mouseDragged():
    redraw()

def draw():
    background(0)
    translate( (width/2), (height/2) )
    rotateY( (TWO_PI*mouseX) / width )
    rotateX( (TWO_PI*mouseY) / height )

    # Draw FFT of frequency domain signal in green
    x = 0
    lastx = 0
    lasty = 0
    lastz = 0
    stroke(0,255,0)
    for n in signal:
        x += 1
        y = n.real * scl
        z = n.imag * scl
        line(lastx-(N/2),lasty,lastz, x-(N/2),y,z)
        lastx = x
        lasty = y
        lastz = z

    # Draw FFT of time domain wave in blue
    x = 0
    lastx = 0
    lasty = 0
    lastz = 0
    stroke(0,0,255)
    for n in wave:
        x += 1
        y = n.real * scl
        z = n.imag * scl
        line(lastx-(N/2),lasty,lastz, x-(N/2),y,z)
        lastx = x
        lasty = y
        lastz = z