In this notebook we are learning about the Discrete Fourier Transform (DFT). We will also learn how to compute the DFT of a $n$ dimentional vector in $\log n$ time complexity. The algorithm for computing the DFT of a vector $v$ is one of the most important algorithms in computer science. The objective of this notebook is to develop the algorithm and apply it to the problem of polynomial multiplication. There is supplementary document with additional mathematical details, here we are focused on the code.
We begin with the following imports and shorthands.
import numpy as np
from numpy.polynomial import Polynomial
exp = np.exp
pi = np.pi
An $n-1$ degree polynomial $\displaystyle A(x) = \sum_{k=1}^{n-1} a_k x^k$ can be described completely by an $n$ dimentional vector where the $k$'th entry (zero based) is given by the $k$'th coefficient, for example, $\vec A = (a_0, \ldots, a_n)$.
P = Polynomial([4,3,0,1])
P
Q = Polynomial([1,6,3,3])
Q
Let $\omega$ denote an $n$'th complex root of unity. We begin by defining the matrix $V_\omega$ where the entry in the $i$'th row and $j$'th column is $\omega^{(i-1)(j-1)}$. A key property of the matrix $V_\omega$ is the identity $V_\omega^{-1} = \frac{1}{n} V_{\omega^{-1}}$, it is easy to see that this inverse will always exist. Suppose that we have an $n-1$ degree polynomial $A$. Notice we have $V_\omega \vec A = (A(\omega^0), \ldots, A(\omega^n))$, we will call this vector an evaluation vector.
Since the matrix $V_\omega$ is invertable it has non zero determinant, therefore we have the not so obvious fact: there is one to one coorespondence between coefficient vectors and evaluation vectors. This is a powerful fact. Notice that when multiplying polynomials given in evaluation vector form, we can simply perform a term by term multiplication. This is easily done in $O(n)$, while standard polynomial multiplication is $O(n^2)$. The big question we answer is how to compute the evaluation vector quickly.
Note: the matrix $V_\omega$ is sometimes called a Vandermonde matrix.
n = 7 # Multiplying 2 degree 3 polynomials will result in a degree 6 polynomial. We should make sure that we can recover all 7 coefficients of the result.
omega = exp(2*pi*1j/n)
V = np.vander([omega**i for i in range(n)]).transpose()
Vinv = np.vander([omega**(-i) for i in range(n)]) / n
np.round(V @ Vinv) # Should be identity (it is).
# We pad our coef vectors with 0 so that the matrix operations are properly defined.
Pe = V @ np.pad(P.coef, (0, 3))
Qe = V @ np.pad(Q.coef, (0, 3))
result = np.round(Vinv @ np.multiply(Pe, Qe))
print(Polynomial(np.real(result)), P*Q) # It's almost magic.
poly([ 4. 27. 30. 22. 15. 3. 3.]) poly([ 4. 27. 30. 22. 15. 3. 3.])
The above code done naively is still $O(n^2)$. In order to reduce the complexity of the algorithm we will need to be more clever. Let us express the polynomial $A$ in two ways \begin{align*} A(x) &= q_1(x)(x^{n/2}-1)+r_1(x) \\ A(x) &= q_2(x)(x^{n/2}+1)+r_2(x). \end{align*} The above equations result from the division algorithm. The degree of $r_1, r_2$ will be less than $n/2$. Next we will leverage some of the special properties of roots of unity. We have $$A(\omega^{2k}) = q_1(\omega^{2k}) (\omega^{2kn/2} - 1) + r_1(\omega^{2k}) = r_1(\omega^{2k}),$$ $$A(\omega^{2k+1}) = q_2(\omega^{2k+1}) (\omega^{2kn/2}\omega^{n/2} + 1) + r_2(\omega^{2k+1}) = r_2(\omega^{2k+1}) = r_2(\omega \omega^{2k}).$$ Amazingly, the divisions by $x^2 - 1$ and $x^2 + 1$ can be done easily. We can evaluate the even and odd entries of $V_\omega \vec A$ seperately by recursively computing the DFT of vectors half the size! This algorithm is called the Fast Fourier Transform (FFT) and has a runtime complexity of $O(\log n)$.
# Only works on vectors with size that is an integer power of 2.
def fft(A, omega):
if len(A) == 1: return A
n = len(A)
# Each has degree n/2.
r1 = [A[i] + A[i+(n//2)] for i in range(n//2)]
r2 = [omega**i * (A[i] - A[i+(n//2)]) for i in range(n//2)]
# Compute the fft of r1, r2, rearrange the entries in proper order.
res = [fft(r1, omega**2), fft(r2, omega**2)]
return np.array([res[i % 2][i // 2] for i in range(n)])
def mult(f, g):
# We will need the vectors to have size 2^k for some k.
n, m = f.degree(), g.degree()
N = int(2 ** np.ceil(np.log2(n+m)))
f = np.pad(f.coef, (0, N - n - 1))
g = np.pad(g.coef, (0, N - m - 1))
# Point wise multiply transformed vectors.
fg = np.multiply(fft(f, exp(2*pi*1j/N)), fft(g, exp(2*pi*1j/N)))
# Reverse the transform on the result.
return np.round(fft(fg, exp(-2*pi*1j/N)) / N)
np.real(mult(P, Q)), (P * Q).coef # Yes!
(array([ 4., 27., 30., 22., 15., 3., 3., -0.]), array([ 4., 27., 30., 22., 15., 3., 3.]))
A common application of the FFT algorithm is in signal processing. Observe the example below where the FFT function is able to detect the constants inside of the sin functions.
import matplotlib.pyplot as plt
n = 256
t = np.arange(n)
sp = fft(np.sin(t)+np.sin(2*t)+np.sin(3*t), exp(-2*pi*1j/n))
freq = np.fft.fftfreq(n, d=1/(2*pi))
fig, ax = plt.subplots(2)
fig.set_figwidth(16)
ax[0].plot(t, np.sin(t)+np.sin(2*t)+np.sin(3*t))
ax[1].plot(freq, sp.real, freq, sp.imag)
plt.show()
A another application of the DFT is in algorithmic number theory. Sums of the form $$x_k = \sum_{j=0}^{n-1} a_j \exp\Big(\frac{j k 2 \pi i}{n}\Big), \ k = 0, \ldots, n-1$$ where $a_j$ are known constants, appear frequently due to the number theoretic significance of the roots of unity. The Fast Fourier Transform algorithm is used to significantly improve the runtime of computations.