From 4ab92d4a86d7ced3cb4a8ce09398d60002f9e8e2 Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Fri, 30 Oct 2026 03:35:13 +0100 Subject: [PATCH] feat: tempo grid estimation Co-Authored-By: Claude Sonnet 5.5 --- src/pipeline/tempo.ts | 166 +++++++++++++++++++++++++++++++++++ tests/pipeline/tempo.test.ts | 60 +++++++++++++ 2 files changed, 226 insertions(+) create mode 100644 src/pipeline/tempo.ts create mode 100644 tests/pipeline/tempo.test.ts diff --git a/src/pipeline/tempo.ts b/src/pipeline/tempo.ts new file mode 100644 index 0000000..dc027c8 --- /dev/null +++ b/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; +} diff --git a/tests/pipeline/tempo.test.ts b/tests/pipeline/tempo.test.ts new file mode 100644 index 0000000..7aa837d --- /dev/null +++ b/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 }); });