audio-mel: FFT instead of the direct DFT in MelSpectrogram

CHANGELOG.md

@@ -6,6 +6,16 @@ All notable changes to SKaiNET-audio are documented here. The format follows
 
 ## [Unreleased]
 
+### Changed
+
+- **`audio-mel`: `MelSpectrogram` uses an FFT instead of the direct DFT** — a mixed-radix FFT
+  (radix 4, 2 and 5 butterflies in closed form, 3 and larger factors generically), so Whisper's
+  400-point frame (4·4·5·5) needs no power-of-two size. Each mel filter is summed over its non-zero
+  band only. A 4 s Whisper window takes 5.8 ms instead of 38 ms on the JVM, and on an arm32
+  Android device the log-mel drops from ≈ 0.88 s to ≈ 0.2 s; the 640 KB cos/sin tables are gone.
+  Output against the `whisper.log_mel_spectrogram` golden: cosine 0.99999999 (was 0.99999998).
+  `MelSpectrogram` stays safe to share across threads (the FFT's buffers are per call).
+
 ## [0.2.1] - 2026-09-13
 
 ### Changed

audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/MelSpectrogram.kt

@@ -5,7 +5,6 @@ import sk.ainet.audio.core.FeatureExtractor
 import kotlin.math.PI
 import kotlin.math.cos
 import kotlin.math.log10
-import kotlin.math.sin
 
 /**
  * Mel front-end parameters. Presets cover the ASR models this family supports; anything else is a
@@ -45,7 +44,7 @@ public enum class MelNorm {
 }
 
 /**
- * Log-mel spectrogram, pure Kotlin (reflect-pad ±nFft/2, periodic Hann, nFft-point DFT, mel
+ * Log-mel spectrogram, pure Kotlin (reflect-pad ±nFft/2, periodic Hann, nFft-point FFT, mel
  * filterbank, then [MelNorm]). With `MelConfig.whisper()` and Whisper's shipped filterbank this is
  * verified against `whisper.log_mel_spectrogram` (cosine ≈ 1.0) — the generalized port of the
  * proven `WhisperMel`.
@@ -66,12 +65,16 @@ public class MelSpectrogram(
 
     init { require(filters.size == nMels * nBins) { "filterbank ${filters.size} != $nMels*$nBins" } }
 
+    // Each mel filter is non-zero on a narrow band of bins only: summing over [bandStart, bandEnd) skips exact zeros,
+    // so the result is the same as over all bins, at a fraction of the cost.
+    private val bandStart = IntArray(nMels) { m -> (0 until nBins).firstOrNull { filters[m * nBins + it] != 0f } ?: 0 }
+    private val bandEnd = IntArray(nMels) { m -> ((nBins - 1 downTo 0).firstOrNull { filters[m * nBins + it] != 0f } ?: -1) + 1 }
+
     // periodic Hann (denominator nFft, matching torch.hann_window(periodic=true))
     private val hann = FloatArray(nFft) { (0.5 - 0.5 * cos(2.0 * PI * it / nFft)).toFloat() }
 
-    // DFT tables: cos/sin(2π k n / nFft); sign irrelevant since we take magnitude².
-    private val cosT = Array(nBins) { k -> FloatArray(nFft) { n -> cos(2.0 * PI * k * n / nFft).toFloat() } }
-    private val sinT = Array(nBins) { k -> FloatArray(nFft) { n -> sin(2.0 * PI * k * n / nFft).toFloat() } }
+    // mixed-radix FFT: nFft = 400 (Whisper) is 4·4·5·5, so no power-of-two restriction
+    private val spectrum = PowerSpectrum(nFft)
 
     override fun extract(pcm16k: FloatArray): FeatureFrames = logMel(pcm16k)
 
@@ -84,18 +87,14 @@ public class MelSpectrogram(
         val mel = FloatArray(nMels * frames)
         val frame = FloatArray(nFft)
         val power = FloatArray(nBins)
+        val workspace = spectrum.workspace()
         for (f in 0 until frames) {
             val off = f * hop
             for (n in 0 until nFft) frame[n] = padded[off + n] * hann[n]
-            for (k in 0 until nBins) {
-                var re = 0f; var im = 0f
-                val ck = cosT[k]; val sk = sinT[k]
-                for (n in 0 until nFft) { re += frame[n] * ck[n]; im += frame[n] * sk[n] }
-                power[k] = re * re + im * im
-            }
+            spectrum.power(frame, power, workspace)
             for (m in 0 until nMels) {
                 var s = 0f; val base = m * nBins
-                for (k in 0 until nBins) s += filters[base + k] * power[k]
+                for (k in bandStart[m] until bandEnd[m]) s += filters[base + k] * power[k]
                 mel[m * frames + f] = s
             }
         }

audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/PowerSpectrum.kt

@@ -0,0 +1,128 @@
+package sk.ainet.audio.mel
+
+import kotlin.math.PI
+import kotlin.math.cos
+import kotlin.math.sin
+
+/**
+ * Power spectrum `|X_k|²`, `k = 0..n/2`, of a real frame of length [n], by a mixed-radix FFT: `n` is factored into
+ * 4s, 2s, 3s and 5s (any remaining factor is handled by a direct DFT of that size), so Whisper's 400-point frame
+ * (4·4·5·5) and NeMo's 512-point frame both run in `O(n·Σfactors)` instead of the direct DFT's `O(n²/2)`.
+ *
+ * Recursive decimation in time over a precomputed twiddle table `e^(−2πi·j/n)`; twiddles are computed in double
+ * precision. The instance is immutable and can be shared; the buffers of one transform live in a [Workspace].
+ */
+internal class PowerSpectrum(private val n: Int) {
+    init { require(n >= 1) { "n must be positive: $n" } }
+
+    private val factors: IntArray = factor(n)
+    private val wRe = FloatArray(n) { cos(-2.0 * PI * it / n).toFloat() }
+    private val wIm = FloatArray(n) { sin(-2.0 * PI * it / n).toFloat() }
+    private val maxRadix = factors.maxOrNull() ?: 1
+
+    /** The buffers of one transform: one per thread, reusable across frames. */
+    class Workspace internal constructor(n: Int, radix: Int) {
+        internal val re = FloatArray(n)
+        internal val im = FloatArray(n)
+        internal val tRe = FloatArray(radix)
+        internal val tIm = FloatArray(radix)
+    }
+
+    fun workspace(): Workspace = Workspace(n, maxRadix)
+
+    /** Writes `|X_k|²` of [frame] (length [n]) to [out] (length ≥ `n/2 + 1`). */
+    fun power(frame: FloatArray, out: FloatArray, ws: Workspace = workspace()) {
+        require(frame.size == n) { "frame has ${frame.size} samples, this transform takes $n" }
+        transform(frame, 0, 1, n, 0, 0, ws)
+        val re = ws.re; val im = ws.im
+        for (k in 0..n / 2) out[k] = re[k] * re[k] + im[k] * im[k]
+    }
+
+    /**
+     * Transform of the [len] samples `x[inOff + j·stride]` into `re/im[outOff ..< outOff + len]`: the `p` interleaved
+     * sub-sequences are transformed recursively into consecutive blocks of `m = len/p`, then combined.
+     */
+    private fun transform(x: FloatArray, inOff: Int, stride: Int, len: Int, outOff: Int, level: Int, ws: Workspace) {
+        val re = ws.re; val im = ws.im; val tRe = ws.tRe; val tIm = ws.tIm
+        if (len == 1) { re[outOff] = x[inOff]; im[outOff] = 0f; return }
+        val p = factors[level]
+        val m = len / p
+        for (q in 0 until p) transform(x, inOff + q * stride, stride * p, m, outOff + q * m, level + 1, ws)
+
+        val step = n / len                     // twiddle W_len^j = w[j·step]
+        val rot = n / p                        // W_p^j = w[j·rot]
+        for (k in 0 until m) {
+            // t_q = W_len^{q·k} · Y_q[k]
+            var wi = 0
+            for (q in 0 until p) {
+                val yr = re[outOff + q * m + k]; val yi = im[outOff + q * m + k]
+                tRe[q] = yr * wRe[wi] - yi * wIm[wi]
+                tIm[q] = yr * wIm[wi] + yi * wRe[wi]
+                wi += k * step
+            }
+            // X[k + r·m] = Σ_q t_q · W_p^{q·r}: closed forms for 2, 4 and 5, the plain sum otherwise
+            val o = outOff + k
+            when (p) {
+                2 -> {
+                    re[o] = tRe[0] + tRe[1]; im[o] = tIm[0] + tIm[1]
+                    re[o + m] = tRe[0] - tRe[1]; im[o + m] = tIm[0] - tIm[1]
+                }
+                4 -> {
+                    val a0r = tRe[0] + tRe[2]; val a0i = tIm[0] + tIm[2]
+                    val a1r = tRe[0] - tRe[2]; val a1i = tIm[0] - tIm[2]
+                    val a2r = tRe[1] + tRe[3]; val a2i = tIm[1] + tIm[3]
+                    val a3r = tRe[1] - tRe[3]; val a3i = tIm[1] - tIm[3]
+                    re[o] = a0r + a2r; im[o] = a0i + a2i
+                    re[o + m] = a1r + a3i; im[o + m] = a1i - a3r            // a1 − i·a3
+                    re[o + 2 * m] = a0r - a2r; im[o + 2 * m] = a0i - a2i
+                    re[o + 3 * m] = a1r - a3i; im[o + 3 * m] = a1i + a3r    // a1 + i·a3
+                }
+                5 -> {
+                    val s7r = tRe[1] + tRe[4]; val s7i = tIm[1] + tIm[4]
+                    val s10r = tRe[1] - tRe[4]; val s10i = tIm[1] - tIm[4]
+                    val s8r = tRe[2] + tRe[3]; val s8i = tIm[2] + tIm[3]
+                    val s9r = tRe[2] - tRe[3]; val s9i = tIm[2] - tIm[3]
+                    re[o] = tRe[0] + s7r + s8r; im[o] = tIm[0] + s7i + s8i
+                    val s5r = tRe[0] + s7r * C1 + s8r * C2; val s5i = tIm[0] + s7i * C1 + s8i * C2
+                    val s6r = s10i * S1 + s9i * S2; val s6i = -s10r * S1 - s9r * S2
+                    re[o + m] = s5r + s6r; im[o + m] = s5i + s6i                // s6 = −i·(S1·s10 + S2·s9)
+                    re[o + 4 * m] = s5r - s6r; im[o + 4 * m] = s5i - s6i
+                    val s11r = tRe[0] + s7r * C2 + s8r * C1; val s11i = tIm[0] + s7i * C2 + s8i * C1
+                    val s12r = -s10i * S2 + s9i * S1; val s12i = s10r * S2 - s9r * S1
+                    re[o + 2 * m] = s11r - s12r; im[o + 2 * m] = s11i - s12i    // s12 = i·(S2·s10 − S1·s9)
+                    re[o + 3 * m] = s11r + s12r; im[o + 3 * m] = s11i + s12i
+                }
+                else -> for (r in 0 until p) {
+                    var sr = 0f; var si = 0f
+                    var j = 0                          // (q·r mod p)·rot, advanced without a division
+                    for (q in 0 until p) {
+                        sr += tRe[q] * wRe[j] - tIm[q] * wIm[j]
+                        si += tRe[q] * wIm[j] + tIm[q] * wRe[j]
+                        j += r * rot; if (j >= n) j -= n
+                    }
+                    re[o + r * m] = sr
+                    im[o + r * m] = si
+                }
+            }
+        }
+    }
+
+    internal companion object {
+        // W_5 = e^(−2πi/5): C1 + i·(−S1), W_5² = C2 + i·(−S2)
+        private val C1 = cos(2.0 * PI / 5).toFloat()
+        private val S1 = sin(2.0 * PI / 5).toFloat()
+        private val C2 = cos(4.0 * PI / 5).toFloat()
+        private val S2 = sin(4.0 * PI / 5).toFloat()
+
+        /** `n` as a product of 4s, 2s, 3s, 5s and whatever remains (prime factors > 5, each its own stage). */
+        fun factor(n: Int): IntArray {
+            val out = ArrayList<Int>()
+            var rest = n
+            for (p in intArrayOf(4, 2, 3, 5)) while (rest % p == 0) { out += p; rest /= p }
+            var d = 7
+            while (rest > 1 && d * d <= rest) { while (rest % d == 0) { out += d; rest /= d }; d += 2 }
+            if (rest > 1) out += rest
+            return out.toIntArray()
+        }
+    }
+}

audio-mel/src/commonTest/kotlin/sk/ainet/audio/mel/PowerSpectrumTest.kt

@@ -0,0 +1,68 @@
+package sk.ainet.audio.mel
+
+import kotlin.math.PI
+import kotlin.math.abs
+import kotlin.math.cos
+import kotlin.math.sin
+import kotlin.random.Random
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertTrue
+
+class PowerSpectrumTest {
+
+    /** `|X_k|²` by the direct DFT, in double precision — the reference. */
+    private fun directPower(x: FloatArray): DoubleArray {
+        val n = x.size
+        return DoubleArray(n / 2 + 1) { k ->
+            var re = 0.0; var im = 0.0
+            for (j in 0 until n) { val a = -2.0 * PI * k * j / n; re += x[j] * cos(a); im += x[j] * sin(a) }
+            re * re + im * im
+        }
+    }
+
+    private fun assertMatchesDirect(n: Int, seed: Int = n) {
+        val rnd = Random(seed)
+        val x = FloatArray(n) { rnd.nextFloat() * 2f - 1f }
+        val want = directPower(x)
+        val got = FloatArray(n / 2 + 1).also { PowerSpectrum(n).power(x, it) }
+        val scale = want.maxOrNull()!!.coerceAtLeast(1e-12)
+        for (k in want.indices) {
+            val err = abs(got[k] - want[k]) / scale
+            assertTrue(err < 1e-5, "n=$n k=$k: fft ${got[k]} vs dft ${want[k]} (relative error $err)")
+        }
+    }
+
+    @Test fun whispersFrameSize() = assertMatchesDirect(400)          // 4·4·5·5
+    @Test fun nemosFrameSize() = assertMatchesDirect(512)             // 4·4·4·4·2
+    @Test fun everyRadixPath() { for (n in listOf(1, 2, 3, 4, 5, 6, 8, 9, 12, 15, 16, 25, 30, 60, 100, 210, 256)) assertMatchesDirect(n) }
+    @Test fun primeAndRepeatedLargeFactors() { for (n in listOf(7, 49, 97, 397)) assertMatchesDirect(n) }
+
+    @Test
+    fun aPureToneLandsInItsBin() {
+        val n = 400
+        val x = FloatArray(n) { cos(2.0 * PI * 25 * it / n).toFloat() }
+        val p = FloatArray(n / 2 + 1).also { PowerSpectrum(n).power(x, it) }
+        val peak = p.indices.maxByOrNull { p[it] }!!
+        assertTrue(peak == 25 && abs(p[25] - (n / 2.0) * (n / 2.0)) / ((n / 2.0) * (n / 2.0)) < 1e-5, "peak at $peak = ${p[peak]}")
+    }
+
+    @Test
+    fun oneInstanceServesSeveralWorkspaces() {
+        val fft = PowerSpectrum(400)
+        val x = FloatArray(400) { sin(0.1 * it).toFloat() }
+        val a = FloatArray(201).also { fft.power(x, it, fft.workspace()) }
+        val b = FloatArray(201).also { fft.power(x, it, fft.workspace()) }
+        assertContentEquals(a, b)
+    }
+
+    @Test
+    fun factorsPreferRadixFour() {
+        assertContentEquals(intArrayOf(4, 4, 5, 5), PowerSpectrum.factor(400))
+        assertContentEquals(intArrayOf(4, 4, 4, 4, 2), PowerSpectrum.factor(512))
+        assertContentEquals(intArrayOf(2, 3, 5, 7), PowerSpectrum.factor(210))
+        assertContentEquals(intArrayOf(7, 7), PowerSpectrum.factor(49))
+        assertContentEquals(intArrayOf(397), PowerSpectrum.factor(397))
+        assertContentEquals(intArrayOf(), PowerSpectrum.factor(1))
+    }
+}