Discrete Fast Fourier Transform À Partir de Zéro

Motivation: Most Underrated Algorithm | Veritasium

Derivation (of the radix-2 Cooley-Tukey algorithm)

We start by considering the definition of discrete fourier transform:

This reads as: the $k$th component of the frequency vector is the dot product of the sample vector and its corresponding angle vector, in case you are a bit rusty with notations.

If $e^{-i2\pi\frac{kn}{N}}$ looks scary to you, we can rewrite it with Euler's formula ( $e^{i\theta}=\cos{\theta}+i\sin{\theta}$):

Without complex exponentiation, it looks more approachable now.

We now formulate an implementation in C:

• We will need to perform the summation for each bin in the output, hence we use nested loops. The outer loop goes over each output bin to be calculated, and the inner loop performs the summation.
• The output is complex, we thus store the real and imaginary components with a two dimensional array.

#include <math.h>

void dft(const double *data, double freq[2][SAMPLE_SIZE], size_t samples) {
  double *real = freq[0];
  double *imag = freq[1];
  // for each bin
  for (size_t k = 0; k < samples; k++) {
    double re = 0.0;
    double im = 0.0;
    // summation
    for (size_t n = 0; n < samples; n++) {
      double angle = 2.0 * M_PI * k * n / samples;
      re += data[n] * cos(angle);
      im -= data[n] * sin(angle);
    }
    real[k] = re;
    imag[k] = im;
  }
}

It should produce the correct output. However, with the time complexity of $O(n^2)$, it will be the slowest thing after compiling hello_world.rs. We now address this problem by implementing a fast version of the DFT.

For the sake of tidyness, we will use Euler form here. Notice that if we uninterleave the terms, we can construct two sequences similar to the original sequence:

Note that:

What if we subsitute $k$ for $k+N/2$? Don't be scared by the algebra, all steps are elementary:

I think you can now spot that $X_{k}$ and $X_{k+N/2}$ share the same terms. For clarity, we introduce the even term component $E_k=\sum^{\lceil\frac{N}{2}\rceil-1}_{n=0}x_{2n}\left(e^{-i2\pi\frac{k}{N/2}}\right)^n$, the odd term component $O_k=\sum^{\lfloor\frac{N}{2}\rfloor-1}_{n=0}x_{2n+1}\left(e^{-i2\pi\frac{k}{N/2}}\right)^n$, and the factor $W_N=e^{-i2\pi\frac{1}{N}}$, rewriting the equation above as:

From now on, however, we will only consider the case where the sample dimensions are positive powers of two, as we are implementing the radix-2 Cooley-Tukey algorithm that is most common and easy to understand.

After restricting the dimension of the input to powers of two, the floor and ceil functions are no longer necessary, as all divisions work out to nice positive integers, and we can start constructing an implementation. Before we get into anything concrete, let's evaluate its time complexity.

Time Complexity

A crucial step here is not to make the mistake of independently running FFT to obtain each of the $n$ output coefficients, or the time complexity will be $O(n^2\log n)$! A good way to think about it is that we have a function $DFT: \R^n \rightarrow \C^n$, then with $DFT(X) = f(DFT(E), DFT(O))$, where the work function satisifies the relation: $W(DFT(X)) = 2W(DFT(E))= 2W(DFT(O))$ and $E$, $O$ being the even and odd subsequence respectively.

Recall Master Theorem:

Our approach allows us to split the main problem into to subproblems each costing half as much, and we also need an $f$ to combine the two subproblems. It should be apparent that the time complexity of $f$ is $O(n)$, since what it does is summing the two vectors produced by solving the subproblems. The corresponding equation is $T(n)=2T(n/2)+f(n)$, where $W(f(n))=O(n)$. Substituting $a=2$ and $b=2$, we can calculate the critical exponent: $E=\frac{\log 2}{\log 2}=1$; $O(n^1)=W(f(n))$. Hence, the time complexity of the radix-2 Cooley-Tukey algorithm is $O(n\log n)$.

Turning Equations into Imperative Procedures

There is still a considerable gap between the derivation and the concrete procedure we want to derive. Fortunately, we can use the master theorem to figure out more about the shape of the implementation. Consider how elements in $(6)$ map to components in the master theorem:

• $X_k$ is the $k$-th output of the $N$-point DFT; it is the first half of the $T(n)$ we want to construct.
• $X_{k+N/2}$ is the output at frequency index $k+N/2$; it is the second half of the $T(n)$ we want to construct.
• $E_k$ is the $k$-th DFT coefficient of the even-indexed input samples; it is the subproblem $T(n/2)$ we have already solved.
• $O_k$ is the $k$-th DFT coefficient of the odd-indexed input samples, the other subproblem.
• $W_N^k = e^{-i2\pi k/N}$ is the $k$-th power of $N$-th root of unity, also called the twiddle factor.
• $N$ is the size of the DFT, assumed to be a power of two.
• $k$ is the frequency index, with $0 \leq k < N/2$.

Where is $f$ though? $f$ is the function combining the subproblems, which is what these two equations are doing. Combining them gives the equation of $f$:

With all the components in the master theorem clearly defined, we are finally ready. A naive approach is recursion; however, evaluating the recursive function in a top-down manner will incur extra allocations, which is less than ideal. Instead, we can build bottom-up:

Consider the input $X=[x_0, x_1, x_2, x_3, x_4, x_5, x_6, x_7]$: its even subsequence $E$ is $[x_0, x_2, x_4, x_6]$ and its odd subsequence $O$ is $[x_1, x_3, x_5, x_7]$. We recursively break them down to subproblems until they are atomic:

# We represent the position of each component in the input vector with a zero-indexed integer:
[0, 1, 2, 3, 4, 5, 6, 7]
([0, 2, 4, 6], [1, 3, 5, 7])
([0, 4], [2, 6], [1, 5], [3, 7])
([0], [4], [2], [6], [1], [5], [3], [7])

Note:

While writing, I learned that the pattern is the bit reversal pattern. Following is an explanation, but note that the FFT implementation below doesn't use this property:

# Binary:
[0b000, 0b001, 0b010, 0b011, 0b100, 0b101, 0b110, 0b111]
([0b000, 0b010, 0b100, 0b110], [0b001, 0b011, 0b101, 0b111])
([0b000, 0b100], [0b010, 0b110], [0b001, 0b101], [0b011, 0b111])
([0b000], [0b100], [0b010], [0b110], [0b001], [0b101], [0b011], [0b111])

A nice way to understand it is that each operation sorts (and splits) the entire sequence by the discriminant bit in a stable manner. Since we put even terms before odd terms, it resulted in an ascending order. For each subsequence, we can say the following: the bits higher than the discriminant bit form the natural number sequence, while the bits lower than the discriminant bit are the same, and two subsequences differ exactly by the discriminant bit, which form an ascending sequence [0, 1]; that's essentially the pattern of the natural number sequence but mirrored on a bit level, hence we can mirror the index to get the number at that position.

That certainly looks messy. However, an amazing thing happens when you realize that the evaluation order of those sub-DFTs on the same recursion depth doesn't matter, we can thus sort them:

# This is the same process above but sorted:
[0, 1, 2, 3, 4, 5, 6, 7]
([0, 2, 4, 6], [1, 3, 5, 7])
([0, 4], [1, 5], [2, 6], [3, 7])
([0], [1], [2], [3], [4], [5], [6], [7])

I think you can now spot the pattern: given a sequence $X$ generated by interval $i$ and step $s$ starting from $p$ (denoted as $X=(i,s,p)$). We know its subsequences are $E=(2i,\frac{s}{2},p)$ and $O=(2i,\frac{s}{2},p+i)$. For a certain recursion depth, there are $i$ sequences whose starting positions are $0$ to $i-1$, respectively. We also know that the base case $s=1$ is essentially an identity function on the input vector (assuming the input vector is real, you need to extend it to complex domain with the imaginary part zeroed).

It should be obvious that you need an outerloop which mutates the step and interval as you go through different recursion depths; an inner loop that goes through each vector to be produced at the current recursion depth, and an inner inner loop that combines the even and odd DFT results. Now we can start writing some code.

C Implementation

We start by making a skeleton of the function consisting of:

• A signature compatible with our previous DFT implementation
• The three loops we just talked about
• A cool name

void fft(const double *data, double freq[2][SAMPLE_SIZE], size_t samples) {
  size_t intv = samples >> 1;
  size_t steps = 2;

  while (intv > 0) {
    // process all dft of sequences of length `steps`, there are `intv` of them
    for (size_t start = 0; start < intv; start++) {
      for (size_t i = 0; i < steps >> 1; i++) {
        // ...
      }
    }
    intv >>= 1;
    steps <<= 1;
  }
}

I used bit shift instead of division and multiplication of two in the spirit of being efficient. So far so good. Now we consider what goes into the blanks:

• We will need storage for the subsequences we have already computed and also for the bigger sequences we are computing from those subsequences. A naive approach is to store them independently on the heap; there are numerous downsides of this approach. If we look at the decomposition process (the codeblocks in the previous section), we will discover that the vectors at each recursion depth all sum up to the dimension of the input vector! And once we are done computing the current layer, we can discard the result of the layer below. Hence, we can store our intermediate results with two arrays of the dimensions of the output array.

  • We want to avoid copying in general, so instead of using concrete stack arrays, we store pointers to the current and previous results, and once we are done, we just swap the previous and current pointers and reuse the scratch buffer.

• The input vector should be copied to the array of the previous results with the imaginary part zeroed, since DFT of one sample is identity.

• Once we have done all computation, we need to copy the results back to the output array.

• The most inner loop is where the core of the implementation resides, we need to compute $W_N^k$, determine where to read from and where to write to and compute the combined results.

  • The process runs a pair basis, we calculate $X_k$ and $X_{k+N/2}$ from $E_k$ and $O_k$; the subscripts here denotes index in the sequence. You need to be very careful about what $k$ and $N$ maps to in your implementation.
  • We need to distinguish between write-back indices (determined by the parent sequence), subsequence component positions (the absolute positions of components in the entire sequence), and subsequence component indices (the relative position of components in a particular subsequence).
  • $W_N^k$ is determined solely by the parent sequence. Examine your calculation very carefully.

A possible implementation:

void fft(const double *data, double freq[2][SAMPLE_SIZE], size_t samples) {
  double a[2][SAMPLE_SIZE];
  double b[2][SAMPLE_SIZE];
  double (*prev)[2][SAMPLE_SIZE]= &a;
  double (*curr)[2][SAMPLE_SIZE]= &b;

  // DFT of one dimensional data is just identity.
  for (size_t k = 0; k < samples; k++) {
    (*prev)[0][k] = data[k];
    (*prev)[1][k] = 0;
  }

  size_t intv = samples >> 1;
  size_t steps = 2;

  while (intv > 0) {
    // process all dft of sequences of length steps, there are intv of them
    for (size_t start = 0; start < intv; start++) {
      for (size_t i = 0; i < steps >> 1; i++) {
        size_t pos = start + i * intv;
        // Example:
        // [2, 4, 6, 8]
        // start: 2
        // intv: 2
        // steps: 4
        // c_pos_a: 2
        // c_pos_b: 2 + 2
        // c_pos_a_next: 2 + intv + intv * 2
        // c_pos_b_next: 2 + intv + intv * 2
        // write_back: 2
        // write_back_next: 2 + intv * i
        size_t c_pos_a = start + i * intv * 2;
        size_t c_pos_b = start + i * intv * 2 + intv;

        // a <- start, intv * 2, step / 2
        // b <- start + intv, intv * 2, step / 2

        double ex = 2.0 * M_PI * i / steps;
        double w_re = cos(ex);
        double w_im = -sin(ex);

        double o_k_re = w_re * (*prev)[0][c_pos_b] - w_im * (*prev)[1][c_pos_b];
        double o_k_im = w_re * (*prev)[1][c_pos_b] + w_im * (*prev)[0][c_pos_b];

        size_t write_back_a = pos;
        size_t write_back_b = start + (i + (steps >> 1)) * intv;

        (*curr)[0][write_back_a] = (*prev)[0][c_pos_a] + o_k_re;
        (*curr)[1][write_back_a] = (*prev)[1][c_pos_a] + o_k_im;
        (*curr)[0][write_back_b] = (*prev)[0][c_pos_a] - o_k_re;
        (*curr)[1][write_back_b] = (*prev)[1][c_pos_a] - o_k_im;
      }
    }

    double (*prev_old)[2][SAMPLE_SIZE] = prev;
    double (*curr_old)[2][SAMPLE_SIZE] = curr;
    prev = curr_old;
    curr = prev_old;

    intv >>= 1;
    steps <<= 1;
  }

  for (size_t k = 0; k < samples; k++) {
     freq[0][k] = (*prev)[0][k];
     freq[1][k] = (*prev)[1][k];
  }
}

It should work, but there are a couple possible optimizations that you might be also be able to apply in your implementation.

Notice that in the inner loop, we repeatedly evaluate $e^{2\pi \frac{k}{N}}$, but also $e^{i\theta}=\cos\theta+i\sin\theta$, hence we can do some multiplications to save some calls to the expensive trigonometry functions, we can also lift it to the outter scope since it doesn't depend on any variables in the inner inner loop:

void fft(const double *data, double freq[2][SAMPLE_SIZE], size_t samples) {
	// ...
  while (intv > 0) {
    // ...
    // process all dft of sequences of length steps, there are intv of them
    for (size_t start = 0; start < intv; start++) {
      double ex = 2.0 * M_PI / steps;
      double w_u_re = cos(ex);
      double w_u_im = -sin(ex);

      double w_re = 1.0;
      double w_im = 0;

      for (size_t i = 0; i < steps >> 1; i++) {
				// ...
        double next_w_re = w_u_re * w_re - w_u_im * w_im;
        double next_w_im = w_u_im * w_re + w_u_re * w_im;
        w_re = next_w_re;
        w_im = next_w_im;
      }
    }
		// ...
  }
	// ...
}

We can factor out common expressions:

void fft(const double *data, double freq[2][SAMPLE_SIZE], size_t samples) {
	// ..
  while (intv > 0) {
		// ...
    // process all dft of sequences of length steps, there are intv of them
    for (size_t start = 0; start < intv; start++) {
  		// ...
      for (size_t i = 0; i < steps >> 1; i++) {
        size_t pos = start + i * intv;
        size_t c_pos_a = pos + i * intv;
        size_t c_pos_b = c_pos_a + intv;
				// ...
        size_t write_back_a = pos;
        size_t write_back_b = pos + (steps >> 1) * intv;
				// ...
      }
    }
		// ...
  }
	// ...
}

We can use a pointer for the real and imaginary part respectively to avoid indirections:

void fft(const double *data, double freq[2][SAMPLE_SIZE], size_t samples) {
  double a[2][SAMPLE_SIZE];
  double b[2][SAMPLE_SIZE];
  double *prev_re = a[0];
  double *prev_im = a[1];
  double *curr_re = b[0];
  double *curr_im = b[1];
  // DFT of one dimensional data is just identity.
  for (size_t k = 0; k < samples; k++) {
    prev_re[k] = data[k];
    prev_im[k] = 0;
  }
	// ...
  while (intv > 0) {
 		// ...
    // process all dft of sequences of length steps, there are intv of them
    for (size_t start = 0; start < intv; start++) {
   		// ...
      for (size_t i = 0; i < steps >> 1; i++) {
				// ...
        double o_k_re = w_re * prev_re[c_pos_b] - w_im * prev_im[c_pos_b];
        double o_k_im = w_re * prev_im[c_pos_b] + w_im * prev_re[c_pos_b];
        // ...
        curr_re[write_back_a] = prev_re[c_pos_a] + o_k_re;
        curr_im[write_back_a] = prev_im[c_pos_a] + o_k_im;
        curr_re[write_back_b] = prev_re[c_pos_a] - o_k_re;
        curr_im[write_back_b] = prev_im[c_pos_a] - o_k_im;
        // ...
      }
    }
    double *prev_re_old = prev_re;
    double *prev_im_old = prev_im;
    double *curr_re_old = curr_re;
    double *curr_im_old = curr_im;

    prev_re = curr_re_old;
    curr_re = prev_re_old;
    prev_im = curr_im_old;
    curr_im = prev_im_old;
		// ...
  }
	// ...
}

Lastly, we can avoid copying the final result to the output array by calculating the base 2 log and reusing the output array as an intermediate buffer:

void fft(const double *data, double freq[2][SAMPLE_SIZE], size_t samples) {
  double a[2][SAMPLE_SIZE];
  double b[2][SAMPLE_SIZE];
  
  int pof2 = 0, num = samples;
  while (num > 0) {
    num >>= 1;
    pof2 += 1;
  }

  double *prev_re;
  double *prev_im;
  double *curr_re;
  double *curr_im;

  if (pof2 % 2 == 0) {
    prev_re = a[0];
    prev_im = a[1];
    curr_re = freq[0];
    curr_im = freq[1];
  } else {
    curr_re = a[0];
    curr_im = a[1];
    prev_re = freq[0];
    prev_im = freq[1];
  }

  // ...
}

And there you have it, a reasonably performant discrete FFT implementation.

void fft(const double *data, double freq[2][SAMPLE_SIZE], size_t samples) {
  double a[2][SAMPLE_SIZE];
  double b[2][SAMPLE_SIZE];
  
  int pof2 = 0, num = samples;
  while (num > 0) {
    num >>= 1;
    pof2 += 1;
  }

  double *prev_re;
  double *prev_im;
  double *curr_re;
  double *curr_im;

  // arrange the order of prev/curr to avoid copying results
  if (pof2 % 2 == 0) {
    prev_re = a[0];
    prev_im = a[1];
    curr_re = freq[0];
    curr_im = freq[1];
  } else {
    curr_re = a[0];
    curr_im = a[1];
    prev_re = freq[0];
    prev_im = freq[1];
  }

  // DFT of one dimensional data is just identity.
  for (size_t k = 0; k < samples; k++) {
    prev_re[k] = data[k];
    prev_im[k] = 0;
  }

  size_t intv = samples >> 1;
  size_t steps = 2;

  while (intv > 0) {
    double ex = 2.0 * M_PI / steps;
    double w_u_re = cos(ex);
    double w_u_im = -sin(ex);

    // process all dft of sequences of length steps, there are intv of them
    for (size_t start = 0; start < intv; start++) {
      double w_re = 1.0;
      double w_im = 0;

      for (size_t i = 0; i < steps >> 1; i++) {
        size_t pos = start + i * intv;
        // Example:
        // [2, 4, 6, 8]
        // start: 2
        // intv: 2
        // steps: 4
        // c_pos_a: 2
        // c_pos_b: 2 + 2
        // c_pos_a_next: 2 + intv + intv * 2
        // c_pos_b_next: 2 + intv + intv * 2
        // write_back: 2
        // write_back_next: 2 + intv * i
        size_t c_pos_a = pos + i * intv;
        size_t c_pos_b = c_pos_a + intv;

        // a <- start, intv * 2, step / 2
        // b <- start + intv, intv * 2, step / 2

        // double ex = 2.0 * M_PI * i / steps;
        // double w_re = cos(ex);
        // double w_im = -sin(ex);

        double o_k_re = w_re * prev_re[c_pos_b] - w_im * prev_im[c_pos_b];
        double o_k_im = w_re * prev_im[c_pos_b] + w_im * prev_re[c_pos_b];

        size_t write_back_a = pos;
        size_t write_back_b = pos + (steps >> 1) * intv;

        curr_re[write_back_a] = prev_re[c_pos_a] + o_k_re;
        curr_im[write_back_a] = prev_im[c_pos_a] + o_k_im;
        curr_re[write_back_b] = prev_re[c_pos_a] - o_k_re;
        curr_im[write_back_b] = prev_im[c_pos_a] - o_k_im;

        // use multiplications to avoid redundent trigonometry function evaluations
        double next_w_re = w_u_re * w_re - w_u_im * w_im;
        double next_w_im = w_u_im * w_re + w_u_re * w_im;
        w_re = next_w_re;
        w_im = next_w_im;
      }
    }

    double *prev_re_old = prev_re;
    double *prev_im_old = prev_im;
    double *curr_re_old = curr_re;
    double *curr_im_old = curr_im;

    prev_re = curr_re_old;
    curr_re = prev_re_old;
    prev_im = curr_im_old;
    curr_im = prev_im_old;

    intv >>= 1;
    steps <<= 1;
  }
}

If you want to check your implementation, the input $[2.0, 3.0, 5.0, 7.0, 24.0, 12.0, 6.0, 3.0]$ should yield:

re = [62.000000, -31.192388, 15.000000, -12.807612, 12.000000, -12.807612, 15.000000, -31.192388]
im = [0.000000, 4.535534, -5.000000, 2.535534, 0.000000, -2.535534, 5.000000, -4.535534]