diff --git a/contents/cooley_tukey/code/scala/fft.scala b/contents/cooley_tukey/code/scala/fft.scala new file mode 100644 index 000000000..951334828 --- /dev/null +++ b/contents/cooley_tukey/code/scala/fft.scala @@ -0,0 +1,126 @@ + +case class Complex(val real: Double, val imag: Double) { + def +(other: Complex): Complex = new Complex(this.real + other.real, this.imag + other.imag) + def -(other: Complex): Complex = new Complex(this.real - other.real, this.imag - other.imag) + + def *(other: Complex): Complex = new Complex( + this.real * other.real - this.imag * other.imag, + this.real * other.imag + this.imag * other.real) + + def *(other: Double): Complex = new Complex(this.real * other, this.imag * other) +} + +object Complex { + def fromPolar(mag: Double, phase: Double): Complex = + new Complex(mag * Math.cos(phase), mag * Math.sin(phase)) +} + +object FT { + + /** Calculates a single DFT coefficient. */ + private def coefficient(n: Int, k: Int, ftLength: Int): Complex = + Complex.fromPolar(1.0, -2.0 * Math.PI * k * n / ftLength) + + def dft(signal: IndexedSeq[Double]): IndexedSeq[Complex] = + signal.indices map { k => + val terms = for (n <- signal.indices) yield coefficient(n, k, signal.length) * signal(n) + terms reduce { _ + _ } + } + + /** Combines the transforms of the even and odd indices */ + private def mergeTransforms(evens: IndexedSeq[Complex], odds: IndexedSeq[Complex]): IndexedSeq[Complex] = { + val oddTerms = for (i <- odds.indices) yield coefficient(1, i, 2 * odds.length) * odds(i) + + val pairs = evens.zip(oddTerms) + + (pairs map { case (e, o) => e + o }) ++ (pairs map { case (e, o) => e - o }) + } + + def cooleyTukey(signal: IndexedSeq[Double]): IndexedSeq[Complex] = signal.length match { + case 2 => mergeTransforms( + Vector(new Complex(signal(0), 0)), + Vector(new Complex(signal(1), 0))) + case _ => + // Split signal into even and odd indices and call cooleyTukey recursively on each. + val evens = cooleyTukey(for (i <- 0 until signal.length by 2) yield signal(i)) + val odds = cooleyTukey(for (i <- 1 until signal.length by 2) yield signal(i)) + + mergeTransforms(evens, odds) + + } + + /** Reverses the bits in value. */ + private def reverseBits(value: Int, length: Int): Int = length match { + case 1 => value + case _ => + // Split bits in the middle. + val lowerHalf = length / 2 + + // The upper half will be longer if the number of bits is odd. + val upperHalf = length - lowerHalf + val mask = (1 << lowerHalf) - 1 + + // Reverse each half recursively and then swap them. + (reverseBits(value & mask, lowerHalf) << upperHalf) + + reverseBits(value >> lowerHalf, upperHalf) + } + + + private def log2(x: Double): Double = Math.log(x) / Math.log(2.0) + + def bitReverseIndices(signal: IndexedSeq[Double]): IndexedSeq[Double] = { + // Find the maximum number of bits needed. + val bitLength = log2(signal.length).ceil.toInt + + for (i <- signal.indices) + yield signal(reverseBits(i, bitLength)) + } + + private def butterfly(x1: Complex, x2: Complex, coeff: Complex): (Complex, Complex) = + (x1 + coeff * x2, x1 - coeff * x2) + + @scala.annotation.tailrec + private def iterativeCooleyTukeyLoop(signal: IndexedSeq[Complex], dist: Int): IndexedSeq[Complex] = { + // Distance between subsequent groups of butterflies + val stride = 2 * dist + val result = new Array[Complex](signal.length) + + for (groupStart <- 0 until signal.length by stride; i <- 0 until dist) { + val index = groupStart + i + val (r1, r2) = butterfly( + signal(index), + signal(index + dist), + coefficient(1, i, stride)) + + result(index) = r1 + result(index + dist) = r2 + } + + if (stride >= signal.length) + result.toVector + else + iterativeCooleyTukeyLoop(result.toVector, dist * 2) + } + + def iterativeCooleyTukey(signal: IndexedSeq[Double]): IndexedSeq[Complex] = + iterativeCooleyTukeyLoop(bitReverseIndices(signal) map { new Complex(_, 0.0) }, 1) +} + +object Main { + + private def approxEqual(a: IndexedSeq[Complex], b: IndexedSeq[Complex]): Boolean = { + val diffs = a.zip(b) map { case (x, y) => Math.abs(x.real - y.real + x.imag - y.imag) } + diffs map { _ < 1e-12 } reduce { _ && _ } + } + + def main(args: Array[String]): Unit = { + val signal = for (i <- 0 until 16) yield Math.random() * 2 - 1 + val x = FT.dft(signal) + val y = FT.cooleyTukey(signal) + val z = FT.iterativeCooleyTukey(signal) + + println("DFT and Cooley-Tukey approx. equal: " + approxEqual(x, y)) + println("DFT and iterative Cooley-Tukey approx. equal: " + approxEqual(x, z)) + println("Cooley-Tukey and iterative Cooley-Tukey approx. equal: " + approxEqual(y, z)) + } +} diff --git a/contents/cooley_tukey/cooley_tukey.md b/contents/cooley_tukey/cooley_tukey.md index de3feae49..38d1b665e 100644 --- a/contents/cooley_tukey/cooley_tukey.md +++ b/contents/cooley_tukey/cooley_tukey.md @@ -89,6 +89,8 @@ For some reason, though, putting code to this transformation really helped me fi [import:3-15, lang:"javascript"](code/javascript/fft.js) {% sample lang="rs" %} [import:24-37, lang:"rust"](code/rust/fft.rs) +{% sample lang="scala" %} +[import:20-33, lang:"scala"](code/scala/fft.scala) {% 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. @@ -142,6 +144,8 @@ In the end, the code looks like: [import:17-39, lang="javascript"](code/javascript/fft.js) {% sample lang="rs" %} [import:39-55, lang:"rust"](code/rust/fft.rs) +{% sample lang="scala" %} +[import:35-56, lang:"scala"](code/scala/fft.scala) {% endmethod %} As a side note, we are enforcing that the array must be a power of 2 for the operation to work. @@ -255,6 +259,8 @@ Some rather impressive scratch code was submitted by Jie and can be found here: [import, lang:"javascript"](code/javascript/fft.js) {% sample lang="rs" %} [import, lang:"rust"](code/rust/fft.rs) +{% sample lang="scala" %} +[import, lang:"scala"](code/scala/fft.scala) {% endmethod %}