myfft.py

Created by rr-inyo
Created on January 11, 2023
2.16 KB
Implementation of Fast Fourier Transform (FFT) based on decimation-in-frequency Cooley-Tukey algorithm as well as DFT
# myfft.py
# Implementation of FFT and DFT for NumWorks calculators
# (c) 2021 @RR_Inyo
# Released under the MIT lisence
# https://opensource.org/licenses/mit-license.php

import cmath

# FFT Type III, with cmath
def FFT(f):
  N_orig = len(f)
  m = int((cmath.log10(N_orig) / cmath.log10(2)).real)

  # FFT core calculation
  for s in range(0, m):  # Stage counter
    div = 2**s  # Number of divided data segments
    N = N_orig // div  # Number of data points in each segment
    n = N // 2  # Do no use '/', which gives float

    # Butterfly computations of each stage
    for j in range(div):  # Counter for divided time-domain data
      for i in range(n):  # Counter for each item in divided time-domain data
        # Prepare complex sinusoid, rotation operator
        W = cmath.exp(-2j * cmath.pi * i / N)

        # Butterfly computation
        f_tmp = f[N * j + i] - f[N * j + n + i]
        f[N * j + i] += f[N * j + n + i]
        f[N * j + n + i] = W * f_tmp

  # Bit-reverse reordering
  i = 0
  for j in range(1, N_orig - 1):
    k = N_orig >> 1
    while(True):
      i ^= k
      if k > i:
        k >>= 1
        continue
      else:
        break
    if j < i:
      f_tmp = f[j]
      f[j] = f[i]
      f[i] = f_tmp

# Test bench
def testFFT(points = 64):
  
  # Import modules used for test
  import time
  import matplotlib.pyplot as plt
  
  # Number of data points
  if points & (points - 1) != 0:
    print('Error: Not power of 2')
    exit()
  else:
    N = points

  k = range(N)  

  # Define function to calculate input digital signal
  def calc_f(i, N):
    f = 0.8 * cmath.exp(1 * 2j * cmath.pi * i / N) \
      + 0.6 * cmath.exp(7 * 2j * cmath.pi * i / N + 1j * cmath.pi / 2) \
      + 0.4 * cmath.exp(12 * 2j * cmath.pi * i / N + 1j * cmath.pi / 2) \
      + 0.3 * cmath.exp(17 * 2j * cmath.pi * i / N - 1j * cmath.pi / 6)
    return f
  
  # Prepare input digital signal
  f = [calc_f(i, N) for i in k]
  
  # Perform FFT and measure time
  t0 = time.monotonic()
  FFT(f)
  t1 = time.monotonic()
  print('Time for FFT: {:.2f} ms'.format((t1 - t0) * 1000))

  # Plot output spectrum
  plt.plot([val.real for val in f])
  plt.plot([val.imag for val in f])
  plt.show()