Functionally idiomatic FFT - functional-programming

Question

I've written the this radix-2 FFT with the goal of making it functionally idiomatic without sacrificing too much performance:
let reverse x bits =
let rec reverse' x bits y =
match bits with
| 0 -> y
| _ -> ((y <<< 1) ||| (x &&& 1))
|> reverse' (x >>> 1) (bits - 1)
reverse' x bits 0
let radix2 (vector: Complex[]) (direction: int) =
let z = vector.Length
let depth = floor(Math.Log(double z, 2.0)) |> int
if (1 <<< depth) <> z then failwith "Vector length is not a power of 2"
// Complex roots of unity; "twiddle factors"
let unity: Complex[] =
let xpn = float direction * Math.PI / double z
Array.Parallel.init<Complex> (z/2) (fun i ->
Complex.FromPolarCoordinates(1.0, (float i) * xpn))
// Permutes elements of input vector via bit-reversal permutation
let pvec = Array.Parallel.init z (fun i -> vector.[reverse i depth])
let outerLoop (vec: Complex[]) =
let rec recLoop size =
if size <= z then
let mid, step = size / 2, z / size
let rec inrecLoop i =
if i < z then
let rec bottomLoop idx k =
if idx < i + mid then
let temp = vec.[idx + mid] * unity.[k]
vec.[idx + mid] <- (vec.[idx] - temp)
vec.[idx] <- (vec.[idx] + temp)
bottomLoop (idx + 1) (k + step)
bottomLoop i 0
inrecLoop (i + size)
inrecLoop 0
recLoop (size * 2)
recLoop 2
vec
outerLoop pvec
The outerLoop segment is the biggest nested tail-recursive mess I have ever written. I replicated the algorithm in the Wikipedia article for the Cooley-Tukey algorithm, but the only functional constructs I could think to implement using higher-order functions result in massive hits to both performance and memory efficiency. Are there other solutions that would yield the same results without resulting in massive slow-downs, while still being idiomatic?

Answer 1

I'm not an expert on how the algorithm works, so there might be a nice functional implementation, but it is worth noting that using a localised mutation is perfectly idiomatic in F#.
Your radix2 function is functional from the outside - it takes vector array as an input, never mutates it, creates a new array pvec which it then initializes (using some mutation along the way) and then returns it. This is a similar pattern to what built-in functions like Array.map use (which initializes a new array, mutates it and then returns it). This is often a sensible way of doing things, because some algorithms are better written using mutation.
In this case, it's perfectly reasonable to also use local mutable variables and loops. Doing that will make your code more readable compared to the tail-recursive version. I have not tested this, but my naive translation of your outerLoop function would just be to use three nested loops - something like this:
let mutable size = 2
while size <= z do
let mid, step = size / 2, z / size
let mutable i = 0
while i < z do
for j in 0 .. mid - 1 do
let idx, k = i + j, step * j
let temp = pvec.[idx + mid] * unity.[k]
pvec.[idx + mid] <- (pvec.[idx] - temp)
pvec.[idx] <- (pvec.[idx] + temp)
i <- i + size
size <- size * 2
This might not be exactly right (I did this just be refactoring your code), but I think it's actually more idiomatic than using complex nested tail-recursive functions in this case.