Python/maths/radix2_fft.py

181 lines
6.0 KiB
Python
Raw Normal View History

"""
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(
2019-10-05 05:14:13 +00:00
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
2019-10-05 05:14:13 +00:00
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
2019-10-05 05:14:13 +00:00
for j in range(self.C_max_length // (next_ncol * 2)):
for i in range(next_ncol):
2019-10-05 05:14:13 +00:00
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
2019-10-05 05:14:13 +00:00
for j in range(self.C_max_length // (next_ncol * 2)):
for i in range(next_ncol):
2019-10-05 05:14:13 +00:00
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")
2019-10-05 05:14:13 +00:00
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]
2019-10-05 05:14:13 +00:00
+ inverseC[i][j + self.C_max_length // next_ncol]
)
/ 2
)
# Odd positions
new_inverseC[i + next_ncol // 2].append(
(
inverseC[i][j]
2019-10-05 05:14:13 +00:00
- inverseC[i][j + self.C_max_length // next_ncol]
)
/ (2 * current_root)
)
current_root *= root
# Update
inverseC = new_inverseC
next_ncol *= 2
# Unpack
2019-10-05 05:14:13 +00:00
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(
2019-10-05 05:14:13 +00:00
f"{coef}*x^{i}" for coef, i in enumerate(self.polyA[: self.len_A])
)
B = "B = " + " + ".join(
2019-10-05 05:14:13 +00:00
f"{coef}*x^{i}" for coef, i in enumerate(self.polyB[: self.len_B])
)
C = "A*B = " + " + ".join(
2019-10-05 05:14:13 +00:00
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()