Code: Select all
# algorithms.py
# 14th March 2015
# Examples of Discrete Fourier Transform code in Python
# http://forum.micropython.org/viewtopic.php?f=2&t=208&hilit=fft
#
import math, cmath, time
import gc
# Recursive algorithm: simple but slow and stack-hungry
# Enter with an array of complex numbers of length 2**N
def fft_recursive(x):
n = len(x)
if n <= 1:
return x
even = fft_recursive(x[0::2])
odd = fft_recursive(x[1::2])
return [even[m] + math.e**(-2j*math.pi*m/n)*odd[m] for m in range(n//2)] + \
[even[m] - math.e**(-2j*math.pi*m/n)*odd[m] for m in range(n//2)]
# Code used as the basis for the assembler routine. Cooley-Tukey algorithm with
# twiddle factors precomputed.
def fft(nums, roots, forward=True):
n = len(nums)
m = int(math.log(n)/math.log(2))
# n= 2**m
# Do the bit reversal
i2 = n >> 1
j = 0
for i in range(n-1):
if i < j:
nums[i], nums[j] = nums[j], nums[i]
k = i2
while (k <= j):
j -= k
k >>= 1
j += k
# Compute the FFT
l2 = 1
for l in range(m):
c = roots[l]
if forward:
c = c.real - c.imag * 1j
l1 = l2
l2 <<= 1
u = 0j+1
for j in range(l1):
for i in range(j, n, l2):
i1 = i+l1
t1 = u*nums[i1]
nums[i1] = nums[i] - t1
nums[i] += t1
u *= c
# Scaling for forward transform
if forward:
for i in range(n):
nums[i] /= n
return nums
def buildarrays(length):
bits = int(math.log(length)/math.log(2))
roots = []
c = -1+0j
roots.append(c) # Complex roots of unity
for x in range(bits):
cimag = math.sqrt((1.0 - c.real) / 2.0) # Imaginary part
creal = math.sqrt((1.0 + c.real) / 2.0) # Real part
c = creal + cimag*1j
roots.append(c)
re = [0+0j]*length
return re, roots
# Test/demo for above algorithm
def printlist(q):
for t in q:
tc = t + (0 + 0j)
print("[{:5.2f}{:5.2f}] ".format(tc.real, tc.imag), end="")
print()
def run(n=1024):
nums, roots = buildarrays(n)
print("Initial")
gc.collect()
try:
for x in range(len(nums)):
nums[x] = 0.1 + math.cos(2*math.pi*x/len(nums))
start = time.ticks_ms()
fft_recursive(nums)
print("Recursive fft", time.ticks_diff(time.ticks_ms(), start)/1000)
except Exception as e:
print("Recursive fft failed: ", e)
for x in range(len(nums)):
nums[x] = 0.1 + math.cos(2*math.pi*x/len(nums))
gc.collect()
start = time.ticks_ms()
fft(nums, roots, True)
print("Forward fft", time.ticks_diff(time.ticks_ms(), start)/1000)
gc.collect()
start = time.ticks_ms()
# printlist(nums)
fft(nums, roots, False)
print("Reverse fft", time.ticks_diff(time.ticks_ms(), start)/1000)
# start = time.ticks_ms()
# printlist(nums)
# print("List", time.ticks_diff(time.ticks_ms(), start)/1000)