The FFT algorithm

In [1]:
def pad_poly(f, n):
    '''
    Given f as an array of coefficients, pad the polynomial to length n
    by adding zeros
    '''
    if len(f) >= n: return f
    
    return f + [f[0] - f[0] for _ in range(n - len(f))]

def is_power_of_two(n):
    if n == 1: return True
    if n % 2 == 1: return False
    return is_power_of_two(n/2)

def FFT(f, omega, n):
    '''
    Assuming that: 
     - f is an array of n coefficients
     - omega is an n-th principal root of unity
     - n is a power of 2
    '''
    
    if n == 1: return [f[0]]
    
    if not is_power_of_two(n): raise ValueError(f"{n} is not a power of two")
    
    f = pad_poly(f, n)
        
    
    # Precompute all powers of omega
    pow_omega = []
    omega_i = (omega / omega)
    for i in range(n):
        pow_omega.append(omega_i)
        omega_i = omega_i * omega
    
    f_even = f[::2]
    f_odd  = f[1::2]
    
    a_even = FFT(f_even, omega*omega, n/2)
    a_odd  = FFT(f_odd , omega*omega, n/2)
    
    output = []
    for i in range(n):
        output.append(a_even[i % (n/2)] + pow_omega[i] * a_odd[i % (n/2)])
    return output
In [2]:
R = CC
I**2
Out[2]:
-1
In [3]:
FFT([1,2,3,4], I, 4)
Out[3]:
[10, -2*I - 2, -2, 2*I - 2]

Polynomial multiplication

Let's first write down the function where we are provided a suitable root of unity.

In [4]:
def poly_mult_with_rou(f, g, omega, n):
    '''
    f and g are polynomials given as an array of n coefficients,
    and omega is a 2n-th root of unity.
    
    Assuming that n is a power of 2
    '''
    
    if n == 1: return [f[0]*g[0]]

    f_prime = pad_poly(f, 2*n)
    g_prime = pad_poly(g, 2*n)
        
    FFT_f = FFT(f_prime, omega, 2*n)
    FFT_g = FFT(g_prime, omega, 2*n)
    
    FFT_h_times_2n = [FFT_f[i] * FFT_g[i] for i in range(2*n)]
    
    h_times_2n = FFT(FFT_h_times_2n, omega^-1, 2*n)
    
    return [h_times_2n[i] / (2*n) for i in range(2*n) ]
In [5]:
f = [1,2] # 1 + 2x
g = [3,4] # 3 + 4x
poly_mult_with_rou(f, g, omega = I, n = 2)
Out[5]:
[3, 10, 8, 0]
In [6]:
R.<x> = PolynomialRing(QQ)
f = R(f)
g = R(g)
h = R([3,10,8,0])
h == f * g
Out[6]:
True

When the ring doesn't contain a ROU

We'll just look at an example where we wish to multiply polynomials $f$ and $g$ of degree < $n = 4$. This would need an $8$-th principal root of unity but the field $\mathbb{Q}$ does not have one.

In [7]:
R.<x,y> = PolynomialRing(QQ)

f = 1 + 2*x - x^2 + 4*x^3
g = 7 - 2*x + 0   + 3*x^3

Let's write the above polynomials as a bivariate polynomial ($k = m = 2$)

\begin{align*} f(x) &= 1 + 2x - x^2 + 4x^3 & g(x) &= 7 - 2x + 3x^3\\ \tilde{f}(x,y) & = (1 + 2x) + (-1 + 4x)y & \tilde{g}(x,y) & = (7 - 2x) + (3x)y\\ f(x) & = \tilde{f}(x,x^2) & g(x) & = \tilde{g}(x,x^2) \end{align*}
In [8]:
f_tilde = (1 + 2*x) + (-1 + 4*x)*y
g_tilde = (7 - 2*x) + (0  + 3*x)*y

f_tilde_coeff = [1 + 2*x , (-1 + 4*x)]
g_tilde_coeff = [7 - 2*x , (0  + 3*x)]

In order to multiply $\tilde{f}$ and $\tilde{g}$, which are polynomials in $y$ of degree < $2$, we need a $4$-th root of unity. So let's expand the ring to $$R' = \frac{\mathbb{Q}[x]}{(x^4 + 1)}.$$

In [9]:
R_prime = R.quotient(R.ideal(x^4 + 1))
R_prime
xbar, ybar = R_prime.gens()

In the above ring, $\overline{x}$ is actually an $8$-th root of unity but we only need a $4$-th root of unity so we will choose $\omega = \bar{x}^2$ when we need to multiply.

In [10]:
omega = xbar^2
print(f"Powers of omega where we evaluate: {[omega^i for i in range(4)]}")

Powers of omega where we evaluate: [1, xbar^2, -1, -xbar^2]

Let's now interpret the two polynomials $\tilde{f}$ and $\tilde{g}$ modulo $x^4 + 1$.

In [11]:
f_tilde_bar = R_prime(f_tilde)
print(f"f_tilde_bar = {f_tilde_bar}")
g_tilde_bar = R_prime(g_tilde)
print(f"g_tilde_bar = {g_tilde_bar}")


f_tilde_bar_coeff = [R_prime(a) for a in f_tilde_coeff]
g_tilde_bar_coeff = [R_prime(b) for b in g_tilde_coeff]

f_tilde_bar = 4*xbar*ybar + 2*xbar - ybar + 1
g_tilde_bar = 3*xbar*ybar - 2*xbar + 7
\begin{align*} \tilde{f} & = 4 \bar{x} \bar{y} + 2\bar{x} - \bar{y} + 1\\ \tilde{g} & = 3 \bar{x} \bar{y} - 2\bar{x} + 7 \end{align*}

Since we already have a 4th principal root of unity in our ring ($\omega = \overline{x}^2$) we can use FFT to compute the product.

In [12]:
# If you want to see what FFT does with omega = xbar^2, you can try out the line below:
#
# print(FFT(f_tilde_bar_coeff, omega = xbar^2, n = 4))

# # [
# #      6*xbar,
# #      4*xbar^3 - xbar^2 + 2*xbar + 1, 
# #     -2*xbar + 2,
# #     -4*xbar^3 + xbar^2 + 2*xbar + 1
# # ]
# # which is basically just
# # [f_tilde(xbar, 1), f_tilde(xbar, xbar^2), f_tilde(xbar, -1), f_tilde(xbar, -xbar^2)]
In [13]:
h_tilde_bar_coeff = poly_mult_with_rou(
                        f_tilde_bar_coeff, 
                        g_tilde_bar_coeff, 
                        omega = xbar^2, 
                        n = 2
                    )
h_tilde_bar_coeff
Out[13]:
[-4*xbar^2 + 12*xbar + 7, -2*xbar^2 + 33*xbar - 7, 12*xbar^2 - 3*xbar, 0]
In [14]:
h_tilde = (-4*x^2 + 12*x + 7) + (-2*x^2 + 33*x - 7)*y + (12*x^2 - 3*x)*y^2
h_tilde == f_tilde * g_tilde
Out[14]:
True
In [15]:
h = h_tilde.subs(y = x^2)
f*g == h
Out[15]:
True