Fourier transform implementations in Nim

contents/cooley_tukey/code/nim/.gitignore

@@ -0,0 +1 @@
+fft

contents/cooley_tukey/code/nim/fft.nim

@@ -0,0 +1,128 @@
+from bitops import fastLog2
+from math import PI, `^`
+from random import rand, randomize
+from sequtils import map, toSeq, zip
+from sugar import `=>`
+import complex
+# $ nimble install fftw3
+import fftw3
+
+# For some reason this isn't in the Nim fftw3 bindings, but this is the
+# current definition from the C header.
+const FFTW_FORWARD = -1
+
+type InpNumber = SomeNumber | Complex[SomeFloat]
+
+proc allClose[T: InpNumber](reference: openArray[T], calculated: openArray[T],
+    thresh = 1.0e-11): bool =
+  ## Are all the numbers in the reference equal to those that were calculated within some threshold?
+  if reference.len != calculated.len:
+    return false
+  result = true
+  for (r, c) in zip(reference, calculated):
+    if abs(r - c) > thresh:
+      result = false
+      break
+
+proc skip[T](inp: openArray[T], stride: int): seq[T] =
+  ## Like `inp[::stride]` in Python.
+  if stride < 1:
+    return toSeq(inp)
+  result = @[]
+  var i = 0
+  while i < inp.len:
+    result.add(inp[i])
+    i += stride
+
+proc nim_to_fftw(inp: Complex64): fftw_complex =
+  [inp.re, inp.im]
+
+proc fftw_to_nim(inp: fftw_complex): Complex64 =
+  complex64(inp[0], inp[1])
+
+proc nim_to_fftw(inp: openArray[Complex64]): seq[fftw_complex] =
+  inp.map(v => nim_to_fftw(v))
+
+proc fftw_to_nim(inp: openArray[fftw_complex]): seq[Complex64] =
+  inp.map(v => fftw_to_nim(v))
+
+# Anything other than 64-bit floats isn't currently supported by the binding.
+proc dft_fftw3(inp: openArray[Complex64]): seq[Complex64] =
+  ## Perform a forward Fourier transform using a binding to FFTW3.
+  var
+    fftw_inp = nim_to_fftw(inp)
+    fftw_out = newSeq[fftw_complex](inp.len)
+  let p = fftw_plan_dft_1d(cint(inp.len),
+                           addr(fftw_inp[low(fftw_inp)]),
+                           addr(fftw_out[low(fftw_out)]),
+                           FFTW_FORWARD,
+                           FFTW_ESTIMATE)
+  fftw_execute(p)
+  fftw_destroy_plan(p)
+  fftw_to_nim(fftw_out)
+
+proc dft[T: SomeFloat](x: openArray[Complex[T]]): seq[Complex[T]] =
+  ## Perform a forward Fourier transform using the naive quadratic algorithm.
+  let n = x.len
+  result = newSeq[Complex[T]](n)
+  for i in 0..n - 1:
+    for k in 0..n - 1:
+      result[i] += x[k] * exp(-complex(0.0, 2.0) * PI * float(i * k / n))
+
+proc cooley_tukey[T: SomeFloat](x: openArray[Complex[T]]): seq[Complex[T]] =
+  ## Perform a forward Fourier transform using the recursive Cooley-Tukey algorithm.
+  let n = x.len
+  if n <= 1:
+    return toSeq(x)
+  result = newSeq[Complex[T]](n)
+  let
+    even = cooley_tukey(skip(x, 2))
+    odd = cooley_tukey(skip(x[1..high(x)], 2))
+    midpoint = n div 2
+  for k in 0..midpoint - 1:
+    let exp_term = exp(-complex(0.0, 2.0) * PI * float(k / n)) * odd[k]
+    result[k] = even[k] + exp_term
+    result[k + midpoint] = even[k] - exp_term
+
+proc bit_reverse[T: InpNumber](x: openArray[T]): seq[T] =
+  let
+    n = x.len
+    l2 = fastLog2(n)
+  result = newSeq[T](n)
+  for k in 0..n - 1:
+    let s = toSeq(0..l2 - 1)
+    var b = 0
+    for i in s:
+      if (k shr i and 1) == 1:
+        b += 1 shl (l2 - 1 - i)
+    result[k] = x[b]
+    result[b] = x[k]
+
+proc iterative_cooley_tukey[T: SomeFloat](x: openArray[Complex[T]]): seq[Complex[T]] =
+  ## Perform a forward Fourier transform using the iterative Cooley-Tukey algorithm.
+  let n = x.len
+  result = bit_reverse(x)
+  for i in 1..fastLog2(n):
+    let
+      stride = 2 ^ i
+      w = exp(-complex(0.0, 2.0) * PI / float(stride))
+    for j in skip(toSeq(0..n - 1), stride):
+      var v = complex(1.0, 0.0)
+      let s = stride div 2
+      for k in 0..s - 1:
+        result[k + j + s] = result[k + j] - v * result[k + j + s]
+        result[k + j] -= result[k + j + s] - result[k + j]
+        v *= w
+
+when isMainModule:
+  randomize()
+  let
+    x = toSeq(1..64).map(i => complex64(rand(1.0), 0.0))
+    y = cooley_tukey(x)
+    z = iterative_cooley_tukey(x)
+    t = dft(x)
+    reference = dft_fftw3(x)
+  assert(bit_reverse(@[1, 2, 3, 4, 7, 11]) == @[7, 3, 11, 4, 1, 3])
+  assert(allClose(reference, y))
+  assert(allClose(reference, z))
+  assert(allClose(reference, t))

contents/cooley_tukey/cooley_tukey.md

@@ -87,6 +87,8 @@ For some reason, though, putting code to this transformation really helped me fi
 [import:15-74, lang:"asm-x64"](code/asm-x64/fft.s)
 {% sample lang="js" %}
 [import:3-15, lang:"javascript"](code/javascript/fft.js)
+{% sample lang="nim" %}
+[import:64-70, lang:"nim"](code/nim/fft.nim)
 {% endmethod %}
 
 In this function, we define `n` to be a set of integers from $$0 \rightarrow N-1$$ and arrange them to be a column.
@@ -138,6 +140,8 @@ In the end, the code looks like:
 [import:76-165, lang:"asm-x64"](code/asm-x64/fft.s)
 {% sample lang="js" %}
 [import:17-39, lang="javascript"](code/javascript/fft.js)
+{% sample lang="nim" %}
+[import:72-85, lang:"nim"](code/nim/fft.nim)
 {% endmethod %}
 
 As a side note, we are enforcing that the array must be a power of 2 for the operation to work.
@@ -249,6 +253,8 @@ Some rather impressive scratch code was submitted by Jie and can be found here:
 [import, lang:"asm-x64"](code/asm-x64/fft.s)
 {% sample lang="js" %}
 [import, lang:"javascript"](code/javascript/fft.js)
+{% sample lang="nim" %}
+[import, lang:"nim"](code/nim/fft.nim)
 {% endmethod %}
 
 <script>