Fast Fourier transform in J

In the previous post we discussed the discrete Fourier transform (DFT). Its J implementation was pretty straightforward, and the program looked almost identical to its mathematical definition. In this post, I want to explore how to implement fast Fourier transform (FFT), a recursively defined version of DFT that reduces algorithmic complexity from $O(n^2)$ to $O(N\log{N})$.

Math

Recall the DFT formula from the previous post

$\mathcal{F}_N\mathbf{f}[m] = \sum_{n=0}^{N-1} \mathbf{f}[n] e^{-2\pi i mn/N}, \quad m = 0, \dots , N-1$

I use a slightly different notation here to make recursive definition easier to understand.

We assume that N is a power of 2. Using basic arithmetic and properties of complex exponentials, we can derive the formula for the DFT of order 2N in terms of two DFTs of order N

$\mathcal{F}_{2N}\mathbf{f}[m]=\mathcal{F}_N\mathbf{f}_{even}+e^{-\pi im/N}\mathcal{F}_N\mathbf{f}_{odd}$ $\mathcal{F}_{2N}\mathbf{f}[m+N]=\mathcal{F}_N\mathbf{f}_{even}-e^{-\pi im/N}\mathcal{F}_N\mathbf{f}_{odd}$

where $m=0,\dots,N-1$. The base case, when $N=1$, is

$\mathcal{F}_1\mathbf{f}[0] = \mathbf{f}[0]$

i.e. the identity function.

Programming

Let’s express the FFT algorithm in words. If $N=1$ then the fft function is just the identity function (] in J). Overwise, it’s a function involving a recursive call on an input vector of half the length. Here is the same phrase written in J

fft=: something_recursive_here ` ] @. (1 = #)

If $N>1$, we split the input vector into two vectors by even and odd indices. We apply FFT recursively to each of these vectors. We multiply the first output vector by 1 and the second one by $e^{-\pi im/N}$. Then, we find the sum and difference of the result vectors, and concatenate them. That should be it.

Let’s start with the split function. To split a vector by even and odd indices in J, we can split it by groups of two elements, stack the groups on top of each other making an $N\times2$ matrix, then transpose it to get $2\times N$ matrix, which gives us the two vectors we need. Here is the implementation of the split function

spl=: |:@(_2 ]\ ])

and here is an example of applying this function to a range vector

   spl 0 1 2 3 4 5 6 7
0 2 4 6
1 3 5 7

The next step is the function $e^{-\pi im/N}$. We implement it in two steps: first, scalar $\omega=e^{\pi i/N}$, then vector $\boldsymbol\omega=\omega^{-m}$. The implementation is analogous to the one from the previous post, with the difference of $\pi i/N$ instead of $2\pi i/N$, and vector $\boldsymbol\omega$ instead of matrix $\Omega$

omg=: ^@o.@(0j1&%)
omv=: omg ^ -@i.

The new matrix $\Omega$ is build by stacking a unit vector $[1,\dots,1]$ of length N on top of the vector $\boldsymbol\omega$

Omg=: 1 ,: omv

Here is an example of its application to $N=4$

   Omg 4
1                  1              1                   1
1 0.707107j_0.707107 2.22045e_16j_1 _0.707107j_0.707107

Now we can go back to the fft function and add some code to the something_recursive_here placeholder

fft=: final_piece_here@(Omg@#@{. * $:"1)@spl ` ] @. (1 = #)

$:"1 means the fft function is recursively applied to every row of its input matrix (which is the result of the spl function above). Omg@#@{. means: take the first row of the input matrix, get the length of it, and pass it to the Omg function.

final_piece_here is then what we described above as “find the sum and difference of the result vectors, and concatenate them”. In J it’s implemented as the fork +/ , -/. Now we have all the pieces in place, and the final program is

spl=: |:@(_2 ]\ ])
omg=: ^@o.@(0j1&%)
omv=: omg ^ -@i.
Omg=: 1 ,: omv
fft=: (+/ , -/)@(Omg@#@{. * $:"1)@spl ` ] @. (1 = #)

References

Source code with examples on GitHub
Official FFT add-on on jsoftware wiki.