Nima Aghdaii

I’ve worked on ads at Meta, Snap, and Nextdoor. I also consult companies on building ads. If you are building your ad stack, feel free to reach out — let’s chat!

Discrete Fourier Transform

13 Feb 2024 » algorithm, fourier, math, learning, linear-algebra

Motivation

Let’s take a look at how we multiply two polynomials of degree N: f(x)=4x42x36x2 +4x + 3g(x)=x4+11x39x2+1x + 6

We can use the distributive property to multiply two polynomials and then sum up coefficients for identical terms. f(x)g(x)=(4x42x36x2 +4x + 3)(x4+11x39x2+1x + 6)=4x8+46x752x656x5+121x49x367x2+21x+18


However, this leads to O(N2) operations to compute the coefficients. In this section, we are going to look at how we can accomplish polynomial multiplication in O(NlogN) using Fast Fourier Transform (FFT) (an algorithm to compute DFT). We will also look at how we can solve this problem in ~10 lines of code!


Polynomial Representations

A polynomial of degree n1 can be represented via:

  • Coefficient Representation: n coefficients [a0,a1,a2,,an1]
  • Value Representation: n sample points [(x1,f(x1)),(x2,f(x2)),,(xn1,f(xn1))]

Note that any n unique points uniquely identify one and only one polynomial of degree n1 (Interpolation Theorem).


Polynomial Multiplication

We observe that computing polynomial multiplication using the coefficients leads to O(N2) operations. Now, let’s try the value representation.


These are the polynomials f(x) and g(x) that we want to multiply:


s1


Step1: Sample n points on each polynomial (at the same x coordinates).


s2


Step2: Compute the pairwise product of samples.


table s3


Step3: Find the unique polynomial of degree n1 (n coefficients) that goes through these n red points (interpolation).


fin


This new red curve is our polynomial f(x)g(x) and performing the multiplication using the value representation allowed us to multiply the two polynomials in O(N) operations.


fft


Now we need two magical operations to take us from the coefficient representation to the value representation and vice versa in less than O(N2). This magical algorithm is called Fast Fourier Transform or FFT.


Fast Fourier Transform


Evaluation (DFT)

We know the result of multiplying two degree n polynomials is going to be of degree 2n, so we need to take 2n+1 samples. However, each sample (for example evaluating f(1)) requires O(n) operations and we need 2n+1 samples, resulting in an O(n2) time complexity. On the other hand, we get to choose the set of x’s for which we want to sample f(x).


Observation

If we want to sample the function f(x)=x2, we can select specific samples by choosing positive values for x and their corresponding negative counterparts (given f(x)=f(x)), thus saving half of the computation. Also, note that the same property holds for any even function.

x2


Arbitrary Polynomials

Let’s say we want to evaluate the function f(x)=4x42x36x2+4x+3, let’s start by isolating the terms with even and odd powers, and then factor out x from the odd powers (to make them all even) and call them Fe and Fo respectively: F(x)=4x42x36x2+4x+3F(x)=(4x46x2+3)+(2x3+4x)F(x)=(4x46x2+3)+x(2x2+4)F(x)=Fe(x2)+xFo(x2)

Property 1: If we compute F(x) in this fashion, computing F(-x) becomes trivial. F(x)=Fe(x2)+xFo(x2)F(x)=Fe(x2)xFo(x2)

Property 2: Note that Fe(x2)=4x46x2+3 is actually a polynomial of degree 2! substituting x2 with v helps us see it better: Fe(v)=4v26v+3.

Recap: We start with a function F(x) and pick our samples as positive and negative pairs. We only need to compute the positive samples, cutting down our problem from n to two smaller subproblems of size n/2. Note that the sample size is reduced in half, as well as the degree of the polynomial. If we can make this work, we have an O(nlogn) recursive solution!

Challenge: The problem is as we go to recurse to the next step, when we decide to split our samples into positive and negative pairs [±x1,±x2,...,±xn/4], our new samples [x12,x22,...,xn/22] cannot become negative, unless if we allow the samples to be complex numbers.


Complex Samples

We observed that if we pick a sample pair x1=1 and x=1, in the next step of recursion we need x2=1 and x2=1. In order to achieve x2=1, we also need to pick a sample pair of i and i. So, our four samples will be [1,1,i,i]. Note that these are the 4 solutions to the equation x4=1, which are called the 4th roots of unity.

Nth Roots of Unity: In general, the solutions to the equation xn=1 are [ω0,ω1,,ωn1] where ω=e2πin. Below is the visualization of the roots for x8=1 on the complex plane.

unity

Basically, we are going to evaluate our polynomial at values [ω0,ω1,,ωn1] where the ± pairs are (wi,wi+n/2) because wi=wi+n/2.

pair


Implementation

We can implement this idea recursively in 10 lines of code. This implementation allows for evaluating n points from a degree n polynomial in O(nlog(n)). Keep in mind that if you need more than n points sampled, just imagine you have a higher degree polynomial with the extra coefficients of 0.

import math, numpy as np

def FFT(c): # coefficients: [-1, 2, 3, 0] for 3x^2+2x-1
    n = len(c) # n must be a power of 2
    if n == 1:
        return c
    w = np.exp(2j * math.pi / n)
    fe, fo = FFT(c[::2]), FFT(c[1::2])
    f = [0] * n
    for i in range(n//2):
        f[i] = fe[i] + w**i * fo[i]
        f[i + n//2] = fe[i] - w**i * fo[i]
    return f


Interpolation

Having transformed the coefficients into points and performed pairwise multiplication, the next step involves determining the polynomial that passes through these points, a process known as “interpolation”. Although interpolation may initially appear to be a more challenging task, we will explore the close relationship between evaluation and interpolation, leading to a surprisingly simple solution.


DFT Matrix

Note that while sampling converts a polynomial (coefficients) to a set of points, interpolation acts as the inverse of that function converting points to the polynomial. Since sampling is a linear transformation, naturally the next step would be write the sampling/evaluation operation as a matrix and try to find its inverse to achieve the interpolation matrix.

Let’s take another look at sampling/evaluation for polynomial F(x)=c0+c1x+c2x2++cn1xn1 at points [x0,x1,,xn1], using the Vandermonde matrix. (F(x0)F(x1)F(x2)F(xn1))=(1x0x02x0n11x1x12x1n11x2x22x2n11xn1xn12xn1n1)(c0c1c2cn1)

And in the FFT algorithm, the i-th evaluation point is the corresponding root of unity (xi=ωi), leading to the matrix below known as the Discrete Fourier Transform Matrix (The DFT Matrix). DFT=(11111ωω2ωn11ω2ω4ω2(n1)1ωn1ω2(n1)ω(n1)(n1))


Inverse DFT Matrix

To be more precise, the DFT matrix also has a normalization factor of 1n as shown below. 1n(11111ωω2ωn11ω2ω4ω2(n1)1ωn1ω2(n1)ω(n1)(n1))

This normalization factor is often added to achieve a unitary matrix. Considering that the inverse of a unitary matrix is the conjugate transpose, and given our matrix is symmetric, we can achieve this by computing the conjugate of each term (negating the imaginary part). For the specific case of ω=e2πin the conjugate would be e2πin which is just ω1. So, the resulting matrix would be the same as the DFT matrix but with negative powers. 1n(11111ω1ω2ω(n1)1ω2ω4ω2(n1)1ω(n1)ω2(n1)ω(n1)(n1))

Yet, it’s crucial to recognize that our initial DFT transformation was essentially an evaluation, devoid of the 1n term. Therefore, to align with the problem at hand, we must provide another multiplier of 1n, resulting in the formulation: DFT1=1n(11111ω1ω2ω(n1)1ω2ω4ω2(n1)1ω(n1)ω2(n1)ω(n1)(n1))

Multiplying this matrix by the evaluated points will yield the coefficients of the desired polynomial. (c0c1c2cn1)=1n(11111ω1ω2ω(n1)1ω2ω4ω2(n1)1ω(n1)ω2(n1)ω(n1)(n1))(F(x0)F(x1)F(x2)F(xn1))

Now, comparing this with the original DFT matrix, the only differences are:

  • negative powers
  • multiply by 1n at the end


Full Implementation

With a one line change in the original implementation of our FFT function, we can have it compute inverse FFT as well. Just keep in mind that we have to divide the final result by n (which is why we introduced a separate function called IFFT).

import math, numpy as np

def FFT(c, inv=False):
    n = len(c) # coefficients: [-1, 2, 3, 0] for 3x^2+2x-1
    if n == 1: # n must be a power of 2
        return c
    w = np.exp(2j * math.pi / n * (-1 if inv else 1))
    fe, fo = FFT(c[::2], inv), FFT(c[1::2], inv)
    f = [0] * n
    for i in range(n//2):
        f[i] = fe[i] + w**i * fo[i]
        f[i + n//2] = fe[i] - w**i * fo[i]
    return f

def IFFT(c):
    f = FFT(c, inv=True)
    return [i / len(c) for i in f]


Testing

Let’s put our original polynomials to the test: f(x)=4x42x36x2 +4x + 3g(x)=x4+11x39x2+1x + 6f(x)g(x)=4x8+46x752x656x5+121x49x367x2+21x+18


Input

Note that the coefficients need to be sorted from the lowest power to the highest.

poly1 = [3, 4, -6, -2, 4]
poly2 = [6, -1, -9, 11, -1]


Code

Just keep in mind the output of IFFT is complex numbers, so we need to extract the real parts and round them to the nearest integer, since the types are float.

def make_power_of_two(arr):
    target_len = 2 ** (len(arr) - 1).bit_length()
    return arr + [0] * (target_len - len(arr))

def multiply_polynomials(poly1, poly2):
    # make them equal length and powers of 2
    target_len = 2 ** (len(poly1) + len(poly2) - 1).bit_length()
    poly1 += [0] * (target_len - len(poly1))
    poly2 += [0] * (target_len - len(poly2))

    p1, p2 = FFT(poly1), FFT(poly2)
    p_result = [p1[i] * p2[i] for i in range(target_len)]
    return IFFT(p_result)


Output

[18, 21, -67, -9, 121, -56, -52, 46, -4, 0, 0, 0, 0, 0, 0, 0]