feat: dataset builder

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

src/audio/mix.ts

@@ -2,27 +2,41 @@ import type { Audio } from "../core/types.ts";
 
 const PEAK_TARGET = Math.pow(10, -1 / 20);
 
-/** Sum stems (with per-stem gain in dB), then peak-normalise to -1 dBFS. */
+/**
-export function mixStems(stems: { audio: Audio; gainDb: number }[]): Audio {
+ * Sum stems (with per-stem gain in dB), then peak-normalise to -1 dBFS.
+ * Also returns each input stem as it contributed to the mix (gain and normalisation applied, padded to the
+ * mix's channel count and length), so that sum(scaled) === mix, and the applied normalisation factor.
+ */
+export function mixStemsScaled(stems: { audio: Audio; gainDb: number }[]): { mix: Audio; scaled: Audio[]; scale: number } {
   if (stems.length === 0) throw new Error("mixStems: no stems");
   const sr = stems[0]!.audio.sampleRate;
   const nch = Math.max(...stems.map((s) => s.audio.channels.length));
   const n = Math.max(...stems.map((s) => s.audio.channels[0]?.length ?? 0));
   const out = Array.from({ length: nch }, () => new Float32Array(n));
+  const scaled: Audio[] = [];
   for (const { audio, gainDb } of stems) {
     if (audio.sampleRate !== sr) throw new Error("mixStems: sample rate mismatch");
     const g = Math.pow(10, gainDb / 20);
+    const own = Array.from({ length: nch }, () => new Float32Array(n));
     for (let c = 0; c < nch; c++) {
       const src = audio.channels[c] ?? audio.channels[0]!;
       const dst = out[c]!;
-      for (let i = 0; i < src.length; i++) dst[i]! += src[i]! * g;
+      const o = own[c]!;
+      for (let i = 0; i < src.length; i++) { const v = src[i]! * g; dst[i]! += v; o[i] = v; }
     }
+    scaled.push({ sampleRate: sr, channels: own });
   }
   let peak = 0;
   for (const ch of out) for (let i = 0; i < n; i++) peak = Math.max(peak, Math.abs(ch[i]!));
+  const k = peak > 0 ? PEAK_TARGET / peak : 1;
   if (peak > 0) {
-    const k = PEAK_TARGET / peak;
     for (const ch of out) for (let i = 0; i < n; i++) ch[i]! *= k;
+    for (const s of scaled) for (const ch of s.channels) for (let i = 0; i < n; i++) ch[i]! *= k;
   }
-  return { sampleRate: sr, channels: out };
+  return { mix: { sampleRate: sr, channels: out }, scaled, scale: k };
+}
+
+/** Sum stems (with per-stem gain in dB), then peak-normalise to -1 dBFS. */
+export function mixStems(stems: { audio: Audio; gainDb: number }[]): Audio {
+  return mixStemsScaled(stems).mix;
 }

src/bench/cli.ts

@@ -0,0 +1,42 @@
+import { buildSong } from "./dataset.ts";
+import type { TierId } from "../gen/tiers.ts";
+import { TIERS } from "../gen/tiers.ts";
+
+function parseSeeds(s: string): number[] {
+  const out: number[] = [];
+  for (const part of s.split(",")) {
+    const m = /^(\d+)-(\d+)$/.exec(part);
+    if (m) for (let i = Number(m[1]); i <= Number(m[2]); i++) out.push(i);
+    else if (/^\d+$/.test(part)) out.push(Number(part));
+    else throw new Error(`bad seeds: ${s}`);
+  }
+  return out;
+}
+
+function arg(args: string[], name: string, def: string): string {
+  const i = args.indexOf(`--${name}`);
+  return i >= 0 && args[i + 1] ? args[i + 1]! : def;
+}
+
+async function main() {
+  const [cmd, ...rest] = process.argv.slice(2);
+  switch (cmd) {
+    case "gen": {
+      const tiers = arg(rest, "tiers", "t1a,t1b,t1c,t1d").split(",") as TierId[];
+      for (const t of tiers) if (!(t in TIERS)) throw new Error(`unknown tier: ${t}`);
+      const seeds = parseSeeds(arg(rest, "seeds", "1-5"));
+      const root = arg(rest, "root", "data");
+      for (const t of tiers) for (const s of seeds) {
+        const t0 = performance.now();
+        const dir = await buildSong(t, s, root);
+        console.log(`${dir} (${((performance.now() - t0) / 1000).toFixed(1)}s)`);
+      }
+      break;
+    }
+    default:
+      console.error("usage: bun run bench gen [--tiers t1a,t1b] [--seeds 1-5] [--root data]");
+      process.exit(1);
+  }
+}
+
+await main();

src/bench/dataset.ts

@@ -0,0 +1,55 @@
+import { join } from "node:path";
+import { mkdir } from "node:fs/promises";
+import type { Audio, Song } from "../core/types.ts";
+import type { TierId } from "../gen/tiers.ts";
+import { generateSong } from "../gen/generate.ts";
+import { songToGp, truthFromGp } from "../score/gp.ts";
+import { renderStems } from "../audio/render.ts";
+import { applyTrackFx, pan } from "../audio/fx.ts";
+import { mixStemsScaled } from "../audio/mix.ts";
+import { readWav, writeWav } from "../audio/wav.ts";
+import { encodeMp3 } from "../audio/mp3.ts";
+
+const SF2_PATH = "node_modules/@coderline/alphatab/dist/soundfont/sonivox.sf2";
+const SR = 44100;
+
+/**
+ * Build one song into `<root>/<tier>-<seed>/`. Idempotent: skipped when mix.mp3 exists and song.json matches.
+ * stems/<track>.wav are exactly each track's contribution to mix.wav (post-fx, post-pan, gain, normalisation),
+ * so the stems sum to the mix up to PCM16 rounding.
+ */
+export async function buildSong(tier: TierId, seed: number, root = "data"): Promise<string> {
+  const dir = join(root, `${tier}-${seed}`);
+  const song = generateSong(tier, seed);
+  const songJson = JSON.stringify(song, null, 1);
+  const jsonFile = Bun.file(join(dir, "song.json"));
+  if ((await Bun.file(join(dir, "mix.mp3")).exists()) && (await jsonFile.exists()) && (await jsonFile.text()) === songJson) {
+    return dir;
+  }
+  await mkdir(join(dir, "stems"), { recursive: true });
+  const sf2 = new Uint8Array(await Bun.file(SF2_PATH).arrayBuffer());
+  const dry = renderStems(song, sf2, SR);
+  const inputs = song.tracks.map((t) => {
+    const raw = dry.get(t.name);
+    if (!raw) throw new Error(`buildSong: no rendered stem for ${t.name}`);
+    return { audio: pan(applyTrackFx(raw, t.role, t.fx, SR), t.pan), gainDb: t.gainDb };
+  });
+  const { mix, scaled } = mixStemsScaled(inputs);
+  await Bun.write(join(dir, "truth.gp"), songToGp(song));
+  for (let i = 0; i < song.tracks.length; i++) await writeWav(join(dir, "stems", `${song.tracks[i]!.name}.wav`), scaled[i]!);
+  await writeWav(join(dir, "mix.wav"), mix);
+  await encodeMp3(join(dir, "mix.wav"), join(dir, "mix.mp3"));
+  await Bun.write(join(dir, "song.json"), songJson); // written last: marks the build complete
+  return dir;
+}
+
+export async function loadSongDir(dir: string): Promise<{
+  song: Song; truth: ReturnType<typeof truthFromGp>; mix: Audio; stems: Map<string, Audio>;
+}> {
+  const song = (await Bun.file(join(dir, "song.json")).json()) as Song;
+  const truth = truthFromGp(new Uint8Array(await Bun.file(join(dir, "truth.gp")).arrayBuffer()));
+  const mix = await readWav(join(dir, "mix.wav"));
+  const stems = new Map<string, Audio>();
+  for (const t of song.tracks) stems.set(t.name, await readWav(join(dir, "stems", `${t.name}.wav`)));
+  return { song, truth, mix, stems };
+}

tests/bench/dataset.test.ts

@@ -0,0 +1,37 @@
+import { test, expect } 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";
+
+test("buildSong writes all artefacts; loadSongDir reads them; stems sum to mix; idempotent", async () => {
+  const root = await mkdtemp(join(tmpdir(), "stsheets-ds-"));
+  try {
+    const dir = await buildSong("t1a", 1, root);
+    const files = ["song.json", "truth.gp", "mix.wav", "mix.mp3"];
+    for (const f of files) expect(await Bun.file(join(dir, f)).exists()).toBe(true);
+    const { song, truth, mix, stems } = await loadSongDir(dir);
+    for (const t of song.tracks) expect(await Bun.file(join(dir, "stems", `${t.name}.wav`)).exists()).toBe(true);
+    expect(truth.tracks[0]!.notes.length).toBe(song.tracks[0]!.notes.length);
+    expect(mix.channels.length).toBe(2);
+
+    // Oracle stems are the exact mix contributions: they sum to the mix (PCM16 rounding only).
+    expect(stems.size).toBe(song.tracks.length);
+    let maxDiff = 0;
+    for (let c = 0; c < mix.channels.length; c++) {
+      const m = mix.channels[c]!;
+      for (let i = 0; i < m.length; i++) {
+        let s = 0;
+        for (const a of stems.values()) s += a.channels[c]![i] ?? 0;
+        maxDiff = Math.max(maxDiff, Math.abs(s - m[i]!));
+      }
+    }
+    expect(maxDiff).toBeLessThan(2e-3);
+
+    const before = (await Bun.file(join(dir, "mix.mp3")).stat()).mtimeMs;
+    expect(await buildSong("t1a", 1, root)).toBe(dir);
+    expect((await Bun.file(join(dir, "mix.mp3")).stat()).mtimeMs).toBe(before);
+  } finally {
+    await rm(root, { recursive: true, force: true });
+  }
+}, 120_000);