radix2 fft Algorithm
There are many different FFT algorithms based on a wide range of published theory, from simple complex-number arithmetical to group theory and number theory. The best-known FFT algorithms depend upon the factorization of N, but there are FFTs with O(N log N) complexity for all N, even for prime As Tukey make not work at IBM, the patentability of the idea was doubted and the algorithm go into the public domain, which, through the computing revolution of the next ten, made FFT one of the indispensable algorithms in digital signal processing.
In discussion with Tukey, Richard Garwin recognized the general applicability of the algorithm not exactly to national security problems, but also to a wide range of problems including one of immediate interest to him, determine the periodicities of the spin orientations in a 3D crystal of Helium-3.James Cooley and John Tukey published a more general version of FFT in 1965 that is applicable when N is composite and not inevitably a power of 2.
"""
Fast Polynomial Multiplication using radix-2 fast Fourier Transform.
"""
import mpmath # for roots of unity
import numpy as np
class FFT:
"""
Fast Polynomial Multiplication using radix-2 fast Fourier Transform.
Reference:
https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm#The_radix-2_DIT_case
For polynomials of degree m and n the algorithms has complexity
O(n*logn + m*logm)
The main part of the algorithm is split in two parts:
1) __DFT: We compute the discrete fourier transform (DFT) of A and B using a
bottom-up dynamic approach -
2) __multiply: Once we obtain the DFT of A*B, we can similarly
invert it to obtain A*B
The class FFT takes two polynomials A and B with complex coefficients as arguments;
The two polynomials should be represented as a sequence of coefficients starting
from the free term. Thus, for instance x + 2*x^3 could be represented as
[0,1,0,2] or (0,1,0,2). The constructor adds some zeros at the end so that the
polynomials have the same length which is a power of 2 at least the length of
their product.
Example:
Create two polynomials as sequences
>>> A = [0, 1, 0, 2] # x+2x^3
>>> B = (2, 3, 4, 0) # 2+3x+4x^2
Create an FFT object with them
>>> x = FFT(A, B)
Print product
>>> print(x.product) # 2x + 3x^2 + 8x^3 + 4x^4 + 6x^5
[(-0+0j), (2+0j), (3+0j), (8+0j), (6+0j), (8+0j)]
__str__ test
>>> print(x)
A = 0*x^0 + 1*x^1 + 2*x^0 + 3*x^2
B = 0*x^2 + 1*x^3 + 2*x^4
A*B = 0*x^(-0+0j) + 1*x^(2+0j) + 2*x^(3+0j) + 3*x^(8+0j) + 4*x^(6+0j) + 5*x^(8+0j)
"""
def __init__(self, polyA=[0], polyB=[0]):
# Input as list
self.polyA = list(polyA)[:]
self.polyB = list(polyB)[:]
# Remove leading zero coefficients
while self.polyA[-1] == 0:
self.polyA.pop()
self.len_A = len(self.polyA)
while self.polyB[-1] == 0:
self.polyB.pop()
self.len_B = len(self.polyB)
# Add 0 to make lengths equal a power of 2
self.C_max_length = int(
2 ** np.ceil(np.log2(len(self.polyA) + len(self.polyB) - 1))
)
while len(self.polyA) < self.C_max_length:
self.polyA.append(0)
while len(self.polyB) < self.C_max_length:
self.polyB.append(0)
# A complex root used for the fourier transform
self.root = complex(mpmath.root(x=1, n=self.C_max_length, k=1))
# The product
self.product = self.__multiply()
# Discrete fourier transform of A and B
def __DFT(self, which):
if which == "A":
dft = [[x] for x in self.polyA]
else:
dft = [[x] for x in self.polyB]
# Corner case
if len(dft) <= 1:
return dft[0]
#
next_ncol = self.C_max_length // 2
while next_ncol > 0:
new_dft = [[] for i in range(next_ncol)]
root = self.root ** next_ncol
# First half of next step
current_root = 1
for j in range(self.C_max_length // (next_ncol * 2)):
for i in range(next_ncol):
new_dft[i].append(dft[i][j] + current_root * dft[i + next_ncol][j])
current_root *= root
# Second half of next step
current_root = 1
for j in range(self.C_max_length // (next_ncol * 2)):
for i in range(next_ncol):
new_dft[i].append(dft[i][j] - current_root * dft[i + next_ncol][j])
current_root *= root
# Update
dft = new_dft
next_ncol = next_ncol // 2
return dft[0]
# multiply the DFTs of A and B and find A*B
def __multiply(self):
dftA = self.__DFT("A")
dftB = self.__DFT("B")
inverseC = [[dftA[i] * dftB[i] for i in range(self.C_max_length)]]
del dftA
del dftB
# Corner Case
if len(inverseC[0]) <= 1:
return inverseC[0]
# Inverse DFT
next_ncol = 2
while next_ncol <= self.C_max_length:
new_inverseC = [[] for i in range(next_ncol)]
root = self.root ** (next_ncol // 2)
current_root = 1
# First half of next step
for j in range(self.C_max_length // next_ncol):
for i in range(next_ncol // 2):
# Even positions
new_inverseC[i].append(
(
inverseC[i][j]
+ inverseC[i][j + self.C_max_length // next_ncol]
)
/ 2
)
# Odd positions
new_inverseC[i + next_ncol // 2].append(
(
inverseC[i][j]
- inverseC[i][j + self.C_max_length // next_ncol]
)
/ (2 * current_root)
)
current_root *= root
# Update
inverseC = new_inverseC
next_ncol *= 2
# Unpack
inverseC = [round(x[0].real, 8) + round(x[0].imag, 8) * 1j for x in inverseC]
# Remove leading 0's
while inverseC[-1] == 0:
inverseC.pop()
return inverseC
# Overwrite __str__ for print(); Shows A, B and A*B
def __str__(self):
A = "A = " + " + ".join(
f"{coef}*x^{i}" for coef, i in enumerate(self.polyA[: self.len_A])
)
B = "B = " + " + ".join(
f"{coef}*x^{i}" for coef, i in enumerate(self.polyB[: self.len_B])
)
C = "A*B = " + " + ".join(
f"{coef}*x^{i}" for coef, i in enumerate(self.product)
)
return "\n".join((A, B, C))
# Unit tests
if __name__ == "__main__":
import doctest
doctest.testmod()
