STSheets/tests/pipeline/tempo.test.ts

import { test, expect, beforeAll, afterAll } 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";
import type { Audio } from "../../src/core/types.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 { song, mix } = await loadSongDir(await buildSong(tier, seed, root));
      const g = estimateGrid(mix);
      const f = [0.5, 1, 2].find((c) => Math.abs(g.bpm / song.bpm - c) / c < 0.02) ?? 0;
      const estBeat = 60 / g.bpm; // judged on the estimated beat so a 2x/0.5x grid is compared on its own beat
      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,
      });
    }
  }
  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);

afterAll(async () => { if (root) await rm(root, { recursive: true, force: true }); });

// ---- synthetic drum/chord patterns -------------------------------------------------------------------------------
const SR = 44100;
function synth(opts: { bpm: number; bars: number; lead: number; pickupBeats?: number; kicks: number[]; snares: number[]; eighthHats: boolean }): Audio {
  const beat = 60 / opts.bpm;
  const total = opts.lead + ((opts.pickupBeats ?? 0) + opts.bars * 4) * beat + 0.5;
  const x = new Float32Array(Math.floor(total * SR));
  let seed = 12345;
  const rnd = () => ((seed = (seed * 1664525 + 1013904223) >>> 0) / 2 ** 32) * 2 - 1;
  const add = (t: number, len: number, f: (i: number, tt: number) => number) => {
    const s = Math.floor(t * SR);
    for (let i = 0; i < len * SR && s + i < x.length; i++) x[s + i]! += f(i, i / SR);
  };
  const kick = (t: number) => add(t, 0.25, (_, tt) => 0.9 * Math.sin(2 * Math.PI * (50 + 80 * Math.exp(-tt * 40)) * tt) * Math.exp(-tt * 14) + 1.2 * rnd() * Math.exp(-tt * 150));
  const snare = (t: number) => add(t, 0.15, (_, tt) => 1.5 * rnd() * Math.exp(-tt * 18));
  const hat = (t: number, a: number) => add(t, 0.04, (_, tt) => a * rnd() * Math.exp(-tt * 120));
  const t0 = opts.lead + (opts.pickupBeats ?? 0) * beat; // true first downbeat
  for (let p = 0; p < (opts.pickupBeats ?? 0); p++) hat(opts.lead + p * beat, 0.2);
  for (let b = 0; b < opts.bars; b++) {
    const barT = t0 + b * 4 * beat;
    // alternating two-chord progression: a sustained chord per bar (harmonic change on the bar line)
    const chord = b % 2 === 0 ? [220, 277.2, 329.6] : [196, 246.9, 293.7];
    add(barT, 4 * beat, (_, tt) => chord.reduce((s, f) => s + 0.04 * Math.sin(2 * Math.PI * f * tt), 0) * Math.min(1, tt * 4));
    for (const k of opts.kicks) kick(barT + k * beat);
    for (const k of opts.snares) snare(barT + k * beat);
    for (let h = 0; h < (opts.eighthHats ? 8 : 4); h++) hat(barT + (h * beat) / (opts.eighthHats ? 2 : 1), 0.12);
  }
  return { sampleRate: SR, channels: [x] };
}

test("gridFromTruth is the oracle grid", () => {
  expect(gridFromTruth(80, 4)).toEqual({ bpm: 80, firstDownbeat: 0, beatsPerBar: 4 });
});

test("leading silence: downbeat found at the true downbeat, not at the first beat after t=0", () => {
  const g = estimateGrid(synth({ bpm: 90, bars: 8, lead: 0.7, kicks: [0], snares: [], eighthHats: false }));
  expect(Math.abs(g.bpm / 90 - 1)).toBeLessThan(0.02);
  expect(circ(g.firstDownbeat - 0.7, (60 / g.bpm) * 4) * 1000).toBeLessThan(50);
});

test("one-beat pickup: bar phase is correct", () => {
  const beat = 60 / 90;
  const g = estimateGrid(synth({ bpm: 90, bars: 8, lead: 0, pickupBeats: 1, kicks: [0], snares: [], eighthHats: false }));
  expect(Math.abs(g.bpm / 90 - 1)).toBeLessThan(0.02);
  expect(circ(g.firstDownbeat - beat, (60 / g.bpm) * 4) * 1000).toBeLessThan(50);
});

test("genuine fast tempo (140 bpm kick/snare/8th-hat rock) is not halved", () => {
  const g = estimateGrid(synth({ bpm: 140, bars: 12, lead: 0, kicks: [0, 2], snares: [1, 3], eighthHats: true }));
  expect(Math.abs(g.bpm / 140 - 1)).toBeLessThan(0.02);
});

test("rejects too-short and silent input", () => {
  expect(() => estimateGrid({ sampleRate: SR, channels: [new Float32Array(100)] })).toThrow(/too short/);
  expect(() => estimateGrid({ sampleRate: SR, channels: [new Float32Array(SR * 5)] })).toThrow(/silent/);
});

// ---- tier-1 mixes ------------------------------------------------------------------------------------------------
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("exact 1x tempo and correct bar phase in >= 8 of 10 songs", () => {
  expect(rows.filter((r) => r.factor === 1).length).toBeGreaterThanOrEqual(8);
  expect(rows.filter((r) => r.factor === 1 && r.barErrMs <= 50).length).toBeGreaterThanOrEqual(8);
});