diff --git a/contents/cooley_tukey/code/nim/.gitignore b/contents/cooley_tukey/code/nim/.gitignore new file mode 100644 index 000000000..e91bc4f92 --- /dev/null +++ b/contents/cooley_tukey/code/nim/.gitignore @@ -0,0 +1 @@ +fft diff --git a/contents/cooley_tukey/code/nim/fft.nim b/contents/cooley_tukey/code/nim/fft.nim new file mode 100644 index 000000000..3f93cb824 --- /dev/null +++ b/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)) diff --git a/contents/cooley_tukey/cooley_tukey.md b/contents/cooley_tukey/cooley_tukey.md index 1b7066063..4c5bc2af0 100644 --- a/contents/cooley_tukey/cooley_tukey.md +++ b/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 %}