FFT - Kotlin

fft.kt

@file:Suppress("LocalVariableName")

package com.scurab.neonimpl

import kotlin.math.ceil
import kotlin.math.cos
import kotlin.math.sin

/*
Implementation of https://rosettacode.org/wiki/Fast_Fourier_transform#C
into kt for floats
N*log(N)
*/

object KtFtt {
    private val PI = Math.PI.toFloat()

    private fun _fft(bf1: FloatArray, bf1Offset: Int, bf2: FloatArray, bf2Offset: Int, n: Int, step: Int) {
        if (step < n) {
            val stepOffset = step * 2
            _fft(bf2, bf2Offset, bf1, bf1Offset, n, stepOffset);
            _fft(bf2, bf2Offset + stepOffset, bf1, bf1Offset + stepOffset, n, stepOffset);

            var _i = 0
            while (_i < n) {
                var v1Re: Float
                var v1Im = -PI * _i / n

                val r = 1.0f/*exp(0f)*/
                //this could be precalculated based on fftSize
                v1Re = r * cos(v1Im)
                v1Im = r * sin(v1Im)

                val v2Re = bf2[(((_i + step)) * 2) + bf2Offset]
                val v2Im = bf2[(((_i + step) * 2) + 1) + bf2Offset]

                val rRe = (v1Re * v2Re) - (v1Im * v2Im)
                val rIm = (v1Re * v2Im) + (v1Im * v2Re)

                val bf1i1 = _i
                val bf1i2 = (_i + n)
                val bf2i1 = _i * 2;
                val bf2i2 = bf2i1 + 1;

                bf1[bf1i1 + bf1Offset] = bf2[bf2i1 + bf2Offset] + rRe;
                bf1[bf1i1 + 1 + bf1Offset] = bf2[bf2i2 + bf2Offset] + rIm;

                bf1[bf1i2 + bf1Offset] = bf2[bf2i1 + bf2Offset] - rRe;
                bf1[bf1i2 + 1 + bf1Offset] = bf2[bf2i2 + bf2Offset] - rIm;

                _i += stepOffset
            }
        }
    }

    /**
     * [input] is used as temp buffer as well, so it destroys the values in it!
     * returns forward FFT
     */
    fun fft(input: FloatArray, n: Int): FloatArray {
        val out = FloatArray(n * 2) { input[it] }
        _fft(out, 0, input, 0, n, 1);
        return out
    }
}

fun generateSignalData(segmentSize: Int, sampleRate: Int, f: (Double) -> Float): List<FloatArray> {
    val result = mutableListOf<FloatArray>()
    val segments = ceil(sampleRate / segmentSize.toDouble()).toInt()
    for (segment in 0 until segments) {
        val dataBlock = FloatArray(segmentSize)
        for (i in dataBlock.indices) {
            val x = i.toDouble() / sampleRate
            dataBlock[i] = f(x)
        }
        result.add(dataBlock)
    }
    return result
}

fun FloatArray.printData() {
    var i = 0
    var j = 0
    while (j < size) {
        println("$i => (${this[j]}, ${this[j + 1]})")
        i++
        j += 2
    }
}

fun main() {
    val data = generateSignalData(128, 128 /*4096,44100*/) {
        Math.cos(3.0 * it * 2.0 * Math.PI).toFloat()
    }.first()

    val input = FloatArray(data.size * 2) { if (it % 2 == 0) data[it / 2] else 0f }
    val result = KtFtt.fft(input, data.size)

    result.printData()
}