Friday, July 4, 2014

Fast Multiplication of Polynomials with Fast Fourier Transform and Complex Nth Roots of Unity



Introduction

The Python source code for this post is here. Two polynomials of degree n can be represented as arrays of n+1 coefficients. For example, 9 - 10x + 7x^2 + 6x^3 can be represented as [9, -10, 7, 6] and -5 + 4x - 2x^3 can be represented as [-5, 4, 0, -2]. The brute-force method to multiply two polynomials is to multiply each coefficient in the first array by each coefficient in the second array and then add them up, which runs in O(n^2).

def mult_poly(A, B):
    lna, lnb = len(A), len(B)
    n = lna + lnb - 1
    rslt = [0 for i in range(n)]
    for i in range(0, lna):
        for j in range(0, lnb):
            rslt[i+j] += A[i]*B[j]
    return rslt


Multiplying two polynomials by hand gives us -45 + 86x - 75x^2 - 20x^3 + 44x^2 - 14x^5 - 12x^6. Let us evaluate it in the Python interpreter.

>>> poly01 = [9, -10, 7, 6]

>>> poly02 = [-5, 4, 0, -2]
>>> mult_poly(poly01, poly02)
[-45, 86, -75, -20, 44, -14, -12]


A faster way to implement two polynomials is possible if we use the Fast Fourier Transform (FFT) to evaluate the two polynomials at the 2n complex 2n-th roots of unity with the FFT, do the pointwise multiplication of the resulting polynomials, and then use the Inverse Fast Fourier Transform (IFFT) to reconstruct (interpolate) the coefficients of the product polynomial. This algorithm is in O(nlogn).

Complex N-th Roots of Unity

To implement the O(nlogn) algorithm outlined above, we need a function to compute complex nth roots of unity. Here is an implementation of the formula for the n-th complex root of unity from Chapter 32 in "Introduction to Algorithms" by Cormen et al.  book. The error argument is added to control the complexity of the real and imaginary coefficients. When the real or imaginary parts of a complex number is less than the error, it is set to 0.

def omega_with_error(k, n, error=0):
    if n <= 0:
        return None
    else:
        x = cmath.exp((2 * cmath.pi * 1j * k)/n)
        xr = x.real
        xi = x.imag
        if abs(xr) < error: xr = 0
        if abs(xi) < error: xi = 0
        return complex(xr, xi)

 

Another commonly used version of omega has -1 in the exponent.

def omega2_with_error(k, n, error=0):
    if n <= 0:
        return None
    else:
        x = cmath.exp(-(2 * cmath.pi * 1j * k)/n)
        xr = x.real
        xi = x.imag
        if abs(xr) < error: xr = 0
        if abs(xi) < error: xi = 0
        return complex(xr, xi)

Fast Polynomial Multiplication

It is assumed that the two coefficient arrays, A and B, that represent the polynomials have a length equal to an integral power of 2. If they are not, they need to be padded with zeros at the right. This can be done as follows.

def pad_with_zeros(a):
    ln = len(a)
    padded_a = copy.copy(a)
    if ln > 2 and is_pow_of_2(ln):
        return padded_a
    while True:
        padded_a.append(0)
        if is_pow_of_2(len(padded_a)):
            return padded_a


Since the polynomials are evaluated with 2n complexs 2n-th roots of unity, the coefficients must be padded until they are of length 2n, which can be done with pad_with_zeros_until_required_length.

def pad_with_zeros_until_required_length(a, n):
    padded_a = copy.copy(a)
    while True:
        padded_a.append(0)
        if dft.is_pow_of_2(len(padded_a)):
            return padded_a


With these tools in place, we can implement a function convolve_poly(A, B) to multiply (convolve) two polynomials.

def convolve_poly(A, B, wf=dft.omega2, error=0):
    ## 1. ensure that coefficient arrays A and B are of the same length
    lenA, lenB = len(A), len(B)
    if lenA != lenB:
        raise Exception('Coefficient arrays must be of the same length')
    ## 2. pad A and B until their lengths are an integral power of 2
    padded_A = dft.pad_with_zeros(A)
    padded_B = dft.pad_with_zeros(B)
    ## 3. pad A and B more so that their lengths are 2n, where n
    ## is the length of padded A

    padded_A = pad_with_zeros_until_required_length(A, 2*len(padded_A))
    padded_B = pad_with_zeros_until_required_length(B, 2*len(padded_A))
    ## 4. compute dft of padded A and B
    fft_A = dft.recursive_fft(padded_A, wf=wf, error=error)
    fft_B = dft.recursive_fft(padded_B, wf=wf, error=error)
    ## 5. compute pointwise multiplication of A and B
    C = []
    for a, b in zip(fft_A, fft_B):
        C.append(a*b)
    ## 6. compute inverse dft and return the first 2*lenA-2 elements.
    return dft.recursive_inverse_fft(C, wf=wf, error=error)[0:2*lenA-1]


If necessary, we can split the real and imaginary coefficients, which is done with convolve_poly_split(A, B).

def convolve_poly_split(A, B, wf=dft.omega2, error=0):
    lenA, lenB = len(A), len(B)
    if lenA != lenB:
        raise Exception('Coefficient arrays must be of the same length')
    ## 1. pad A and B with zeros until their lengths are an integral
    ## power of 2

    padded_A = dft.pad_with_zeros(A)
    padded_B = dft.pad_with_zeros(B)
    ## 2. pad A and B more so that their length is 2*n, where n is
    ## the length of padded_A.

    padded_A = pad_with_zeros_until_required_length(padded_A, 2*len(padded_A))
    padded_B = pad_with_zeros_until_required_length(padded_B, 2*len(padded_A))
    ## 3. compute dft of padded A and B
    fft_A = dft.recursive_fft(padded_A, wf=wf, error=error)
    fft_B = dft.recursive_fft(padded_B, wf=wf, error=error)
    ## 4. compute pointwise multiplication of A and B
    C = []
    for a, b in zip(fft_A, fft_B):
        C.append(a*b)
    ## 5. compute inverse dft of A*B and split the output into
    ## real and imaginary numbers.

    reals, imags = dft.recursive_inverse_fft_split(C, wf=wf, error=error)
    ## 6. the degree of A and B is lenA-1. the degree of A*B is
    ## 2*(lenA-1) = 2*lenA - 2.

    return reals[0:2*lenA-1], imags[0:2*lenA-1]
 



Example

Let us multiply two polynomials -10 + x - x^2 + 7x^3 and 3 - 6x + 8x^3. Doing it by hand gives us
-30 + 63x - 9x^2 - 53x^3 - 34x^4 - 8x^5 + 56x^6.

poly03 = [-10, 1, -1, 7]
poly04 = [3, -6, 0, 8]
print 'POLY03 x POLY04'
poly03xpoly04 = convolve_poly_split(poly03, poly04)
print poly03xpoly04
poly03xpoly04 = convolve_poly(poly03, poly04)
print poly03xpoly04
poly03xpoly04 = mult_poly(poly03, poly04)


Below is the output.

POLY03 x POLY04
([-29.999999999999993, 63.0, -9.0, -53.0, -34.00000000000001, -8.0, 56.0], [-7.105427357601002e-15, 5.452154178752075e-15, 1.887379141862766e-15, -9.004867857552575e-15, -1.4210854715202004e-14, 5.452154178752075e-15, 1.942890293094024e-14])
[(-29.999999999999993-7.105427357601002e-15j),(63+5.452154178752075e-15j), (-9+1.887379141862766e-15j),(-53-9.004867857552575e-15j),(-34.00000000000001-1.4210854715202004e-14j), (-8+5.452154178752075e-15j), (56+1.942890293094024e-14j)]