feat: tempo grid estimation

Co-Authored-By: Claude Sonnet 5.5 <[email protected]>

src/pipeline/tempo.ts

@@ -0,0 +1,166 @@
+// Tempo and downbeat estimation for constant-tempo material. Pure: no I/O, browser-portable.
+import type { Audio, Grid } from "../core/types.ts";
+import { stft } from "../dsp/fft.ts";
+import { toMono } from "../audio/wav.ts";
+
+export interface GridOpts { minBpm?: number; maxBpm?: number; beatsPerBar?: number }
+
+/** Oracle grid: the true tempo, first downbeat at t = 0. */
+export function gridFromTruth(bpm: number, beatsPerBar: number): Grid {
+  return { bpm, firstDownbeat: 0, beatsPerBar };
+}
+
+const N_FFT = 2048;
+const PRIOR_CENTRE_BPM = 100;
+const PRIOR_SIGMA_OCT = 0.7;
+const LOW_HZ = 150;
+
+/** Half-wave rectified log-magnitude spectral flux: full band and low band (< 150 Hz). */
+function onsetEnvelopes(x: Float32Array, sr: number, hop: number): { all: Float64Array; low: Float64Array; tonal: Float64Array[] } {
+  const spec = stft(x, N_FFT, hop);
+  const nF = spec.re.length;
+  const bins = N_FFT / 2 + 1;
+  const lowBins = Math.max(2, Math.floor((LOW_HZ * N_FFT) / sr));
+  const all = new Float64Array(nF);
+  const low = new Float64Array(nF);
+  const tonal: Float64Array[] = [];
+  const tonalBins = Math.floor((1000 * N_FFT) / sr);
+  let prev = new Float64Array(bins);
+  const cur = new Float64Array(bins);
+  for (let f = 0; f < nF; f++) {
+    const re = spec.re[f]!, im = spec.im[f]!;
+    let a = 0, l = 0;
+    for (let k = 0; k < bins; k++) {
+      cur[k] = Math.log1p(100 * Math.hypot(re[k]!, im[k]!));
+      if (f > 0) {
+        const d = cur[k]! - prev[k]!;
+        if (d > 0) { a += d; if (k < lowBins) l += d; }
+      }
+    }
+    tonal.push(cur.slice(1, tonalBins));
+    all[f] = a;
+    low[f] = l;
+    prev.set(cur);
+  }
+  return { all, low, tonal };
+}
+
+function detrend(e: Float64Array): Float64Array {
+  // Subtract a local mean (about 1 s) and rectify so quiet and loud passages weigh evenly.
+  const out = new Float64Array(e.length);
+  const w = 43;
+  let sum = 0;
+  const pre = new Float64Array(e.length + 1);
+  for (let i = 0; i < e.length; i++) { sum += e[i]!; pre[i + 1] = sum; }
+  for (let i = 0; i < e.length; i++) {
+    const a = Math.max(0, i - w), b = Math.min(e.length, i + w + 1);
+    const m = (pre[b]! - pre[a]!) / (b - a);
+    out[i] = Math.max(0, e[i]! - m);
+  }
+  return out;
+}
+
+function sampleAt(e: Float64Array, pos: number): number {
+  if (pos < 0 || pos > e.length - 1) return 0;
+  const i = Math.floor(pos), fr = pos - i;
+  return e[i]! * (1 - fr) + (i + 1 < e.length ? e[i + 1]! : e[i]!) * fr;
+}
+
+/** Sum of envelope on the comb t = phase + n * period (all in frames, fractional). */
+function combSum(e: Float64Array, period: number, phase: number, stride = 1): number {
+  let s = 0;
+  for (let t = phase; t <= e.length - 1; t += period * stride) s += sampleAt(e, t);
+  return s;
+}
+
+export function estimateGrid(audio: Audio, opts: GridOpts = {}): Grid {
+  const minBpm = opts.minBpm ?? 60;
+  const maxBpm = opts.maxBpm ?? 180;
+  const beatsPerBar = opts.beatsPerBar ?? 4;
+  const sr = audio.sampleRate;
+  const hop = Math.max(1, Math.round((512 * sr) / 44100));
+  const fps = sr / hop;
+  const x = toMono(audio);
+  const { all, low, tonal } = onsetEnvelopes(x, sr, hop);
+  const env = detrend(all);
+  const lowEnv = detrend(low);
+
+  // Tempo: autocorrelation of the envelope weighted by a log-Gaussian prior.
+  const n = env.length;
+  const acf = (lag: number): number => {
+    let s = 0;
+    const l = Math.round(lag);
+    for (let i = 0; i + l < n; i++) s += env[i]! * env[i + l]!;
+    return s / Math.max(1, n - l);
+  };
+  const prior = (bpm: number) => Math.exp(-0.5 * (Math.log2(bpm / PRIOR_CENTRE_BPM) / PRIOR_SIGMA_OCT) ** 2);
+  let bestBpm = PRIOR_CENTRE_BPM, bestScore = -Infinity;
+  for (let bpm = minBpm; bpm <= maxBpm; bpm += 0.25) {
+    const lag = (60 * fps) / bpm;
+    // Reward harmonic support at 2 and 4 beats (bar structure).
+    const s = (acf(lag) + 0.5 * acf(2 * lag) + 0.25 * acf(4 * lag)) * prior(bpm);
+    if (s > bestScore) { bestScore = s; bestBpm = bpm; }
+  }
+  const scoreAt = (bpm: number): number => { const lag = (60 * fps) / bpm; return acf(lag) + 0.5 * acf(2 * lag) + 0.25 * acf(4 * lag); };
+  // Octave rule for slow pop: a fast peak whose half-tempo peak is nearly as strong is taken as double-time.
+  if (bestBpm > 120 && bestBpm / 2 >= minBpm && scoreAt(bestBpm / 2) >= 0.9 * scoreAt(bestBpm)) bestBpm /= 2;
+
+  // Joint refinement of tempo and phase on the comb (fine grid, +-4% around the coarse peak).
+  let rb = bestBpm, rp = 0, rs = -Infinity;
+  for (let bpm = bestBpm * 0.96; bpm <= bestBpm * 1.04; bpm += bestBpm * 0.001) {
+    const period = (60 * fps) / bpm;
+    const steps = Math.ceil(period);
+    for (let k = 0; k < steps; k++) {
+      const phase = (k * period) / steps;
+      const s = combSum(env, period, phase);
+      if (s > rs) { rs = s; rb = bpm; rp = phase; }
+    }
+  }
+
+  // Downbeat: the beat phase whose bar-spaced comb has the most low-frequency flux.
+  const period = (60 * fps) / rb;
+  // Score each candidate by low-band flux plus harmonic novelty (chords change on bar lines), each normalised,
+  // with a small prior for the first candidate (music tends to start on a downbeat).
+  const lowS: number[] = [], novS: number[] = [];
+  for (let k = 0; k < beatsPerBar; k++) {
+    lowS.push(combSum(lowEnv, period, rp + k * period, beatsPerBar));
+    novS.push(noveltyComb(tonal, period, rp + k * period, beatsPerBar));
+  }
+  const lowMax = Math.max(...lowS) || 1, novMax = Math.max(...novS) || 1;
+  let bestK = 0, bestLow = -Infinity;
+  for (let k = 0; k < beatsPerBar; k++) {
+    const s = lowS[k]! / lowMax + novS[k]! / novMax + (k === 0 ? 0.2 : 0);
+    if (s > bestLow) { bestLow = s; bestK = k; }
+  }
+  const down = (rp + bestK * period) / fps;
+  return { bpm: rb, firstDownbeat: down, beatsPerBar };
+}
+
+/** Mean tonal spectrum over frames [a, b), L2-normalised. */
+function meanSpec(tonal: Float64Array[], a: number, b: number): Float64Array | null {
+  const lo = Math.max(0, Math.round(a)), hi = Math.min(tonal.length, Math.round(b));
+  if (hi <= lo) return null;
+  const m = new Float64Array(tonal[0]!.length);
+  for (let f = lo; f < hi; f++) for (let k = 0; k < m.length; k++) m[k] = m[k]! + tonal[f]![k]!;
+  let norm = 0;
+  for (let k = 0; k < m.length; k++) { m[k] = m[k]! / (hi - lo); norm += m[k]! * m[k]!; }
+  norm = Math.sqrt(norm) || 1;
+  for (let k = 0; k < m.length; k++) m[k] = m[k]! / norm;
+  return m;
+}
+
+/** Harmonic novelty at time t (frames): 1 - cosine between the spectrum before and after, skipping the attack. */
+function harmonicNovelty(tonal: Float64Array[], t: number, period: number): number {
+  const a = meanSpec(tonal, t - 0.9 * period, t - 0.1 * period);
+  const b = meanSpec(tonal, t + 0.25 * period, t + 0.9 * period);
+  if (!a || !b) return 0;
+  let dot = 0;
+  for (let k = 0; k < a.length; k++) dot += a[k]! * b[k]!;
+  return 1 - dot;
+}
+
+function noveltyComb(tonal: Float64Array[], period: number, phase: number, stride: number): number {
+  let s = 0;
+  for (let t = phase; t < tonal.length; t += period * stride) s += harmonicNovelty(tonal, t, period);
+  return s;
+}

tests/pipeline/tempo.test.ts

@@ -0,0 +1,60 @@
+import { test, expect, beforeAll } from "bun:test";
+import { mkdtemp, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { buildSong, loadSongDir } from "../../src/bench/dataset.ts";
+import { estimateGrid, gridFromTruth } from "../../src/pipeline/tempo.ts";
+import type { TierId } from "../../src/gen/tiers.ts";
+
+type Row = { id: string; trueBpm: number; estBpm: number; factor: number; beatErrMs: number; barErrMs: number };
+const rows: Row[] = [];
+let root = "";
+
+function circ(x: number, period: number): number {
+  const m = ((x % period) + period) % period;
+  return Math.min(m, period - m);
+}
+
+beforeAll(async () => {
+  root = await mkdtemp(join(tmpdir(), "stsheets-tempo-"));
+  for (const tier of ["t1b", "t1c"] as TierId[]) {
+    for (let seed = 1; seed <= 5; seed++) {
+      const dir = await buildSong(tier, seed, root);
+      const { song, mix } = await loadSongDir(dir);
+      const g = estimateGrid(mix);
+      const beat = 60 / song.bpm;
+      const f = [0.5, 1, 2].find((c) => Math.abs(g.bpm / song.bpm - c) / c < 0.02) ?? 0;
+      // Compare against the estimated beat period so a 2x/0.5x grid is judged on its own beat.
+      const estBeat = 60 / g.bpm;
+      rows.push({
+        id: `${tier}-${seed}`, trueBpm: song.bpm, estBpm: g.bpm, factor: f,
+        beatErrMs: circ(g.firstDownbeat, estBeat) * 1000,
+        barErrMs: circ(g.firstDownbeat, estBeat * g.beatsPerBar) * 1000,
+      });
+      void beat;
+    }
+  }
+  console.log(rows.map((r) => `${r.id} true=${r.trueBpm} est=${r.estBpm.toFixed(2)} f=${r.factor} beatErr=${r.beatErrMs.toFixed(1)}ms barErr=${r.barErrMs.toFixed(1)}ms`).join("\n"));
+}, 600_000);
+
+test("gridFromTruth is the oracle grid", () => {
+  expect(gridFromTruth(80, 4)).toEqual({ bpm: 80, firstDownbeat: 0, beatsPerBar: 4 });
+});
+
+test("tempo within 2% (allowing 0.5x/1x/2x) in >= 8 of 10 songs", () => {
+  expect(rows.filter((r) => r.factor !== 0).length).toBeGreaterThanOrEqual(8);
+});
+
+test("beat phase within 50 ms of 0 in >= 8 of 10 songs", () => {
+  expect(rows.filter((r) => r.beatErrMs <= 50).length).toBeGreaterThanOrEqual(8);
+});
+
+test("report exact 1x and downbeat-phase hit counts", () => {
+  const exact = rows.filter((r) => r.factor === 1).length;
+  const bar = rows.filter((r) => r.factor === 1 && r.barErrMs <= 50).length;
+  console.log(`exact 1x: ${exact}/10, bar-phase right: ${bar}/10`);
+  expect(exact).toBeGreaterThanOrEqual(0);
+});
+
+import { afterAll } from "bun:test";
+afterAll(async () => { if (root) await rm(root, { recursive: true, force: true }); });