diff --git a/src/align/dtw.ts b/src/align/dtw.ts new file mode 100644 index 0000000..bd36d5a --- /dev/null +++ b/src/align/dtw.ts @@ -0,0 +1,67 @@ +// Pure (banded) dynamic time warping. Only the two latest accumulated-cost rows are kept; backpointers are one byte +// per cell (O(n·band) when banded), so a 3000×3000 unbanded run uses ≈ 9 MB. + +export interface DtwBand { center: (i: number) => number; width: number } + +const DIAG = 0, UP = 1, LEFT = 2; // from (i-1,j-1), (i-1,j), (i,j-1) + +/** Minimum-cost monotone path from (0,0) to (n-1,m-1) with steps (1,0), (0,1), (1,1). */ +export function dtw( + costAt: (i: number, j: number) => number, + n: number, + m: number, + band?: DtwBand, +): [number, number][] { + if (n <= 0 || m <= 0) return []; + const lo = new Int32Array(n); + const hi = new Int32Array(n); + const off = new Int32Array(n + 1); + for (let i = 0; i < n; i++) { + let a = 0, b = m - 1; + if (band) { + const c = Math.round(band.center(i)); + a = Math.max(0, c - band.width); + b = Math.min(m - 1, c + band.width); + if (i === 0) a = 0; + if (i === n - 1) b = m - 1; + if (a > b) { // centre outside the matrix: clamp to the nearest edge + a = b = Math.min(m - 1, Math.max(0, c)); + } + } + lo[i] = a; hi[i] = b; + off[i + 1] = off[i]! + (b - a + 1); + } + const bp = new Uint8Array(off[n]!); + let prev = new Float64Array(m).fill(Infinity); + let cur = new Float64Array(m).fill(Infinity); + for (let i = 0; i < n; i++) { + const a = lo[i]!, b = hi[i]!, base = off[i]! - a; + const pa = i > 0 ? lo[i - 1]! : 0, pb = i > 0 ? hi[i - 1]! : -1; + cur.fill(Infinity, Math.max(0, a - 1), Math.min(m, b + 2)); + for (let j = a; j <= b; j++) { + let best = Infinity, dir = DIAG; + if (i === 0 && j === 0) best = 0; + else { + if (i > 0 && j > 0 && j - 1 >= pa && j - 1 <= pb) { best = prev[j - 1]!; dir = DIAG; } + if (i > 0 && j >= pa && j <= pb && prev[j]! < best) { best = prev[j]!; dir = UP; } + if (j > a && cur[j - 1]! < best) { best = cur[j - 1]!; dir = LEFT; } + } + if (best === Infinity) { cur[j] = Infinity; continue; } + cur[j] = best + costAt(i, j); + bp[base + j] = dir; + } + // clear stale values outside this row's range so the next row never reads them + if (i > 0) prev.fill(Infinity, pa, pb + 1); + [prev, cur] = [cur, prev]; + } + if (prev[m - 1]! === Infinity) throw new Error("dtw: end cell unreachable within band"); + const path: [number, number][] = []; + let i = n - 1, j = m - 1; + for (;;) { + path.push([i, j]); + if (i === 0 && j === 0) break; + const d = bp[off[i]! - lo[i]! + j]!; + if (d === DIAG) { i--; j--; } else if (d === UP) i--; else j--; + } + return path.reverse(); +} diff --git a/src/align/features.ts b/src/align/features.ts new file mode 100644 index 0000000..2b86020 --- /dev/null +++ b/src/align/features.ts @@ -0,0 +1,132 @@ +// Pure feature extraction for alignment: 12-d chroma + spectral-flux onset strength at a chosen frame rate, onset +// picking. The spectra are computed once on a fixed 50 Hz grid and aggregated to the requested rate. +import type { Audio } from "../core/types.ts"; +import { fft, hann } from "../dsp/fft.ts"; +import { toMono } from "../dsp/mono.ts"; +import { resample } from "../dsp/resample.ts"; + +export interface Features { rate: number; onset: Float32Array; chroma: Float32Array[] } + +const SR = 22050; +const N_CHROMA = 4096; // chroma window (≈186 ms) +const N_ONSET = 1024; // onset window (≈46 ms) for sharper timing +const BASE_HOP = 441; // 50 Hz base grid +const BASE_RATE = SR / BASE_HOP; +const F_MIN = 27.5, F_MAX = 5000; + +/** Magnitude frame at `centre` (reflect-free: zero padded) written into re/im, windowed. */ +function frameMag(x: Float32Array, centre: number, n: number, w: Float64Array, re: Float64Array, im: Float64Array, mag: Float64Array): void { + const start = centre - (n >> 1); + for (let i = 0; i < n; i++) { + const j = start + i; + re[i] = j >= 0 && j < x.length ? x[j]! * w[i]! : 0; + im[i] = 0; + } + fft(re, im); + const scale = 1 / Math.sqrt(n); + for (let k = 0; k <= n >> 1; k++) mag[k] = Math.hypot(re[k]!, im[k]!) * scale; +} + +export function features(audio: Audio, rate: number): Features { + let x = toMono(audio); + if (audio.sampleRate !== SR) { + const k = audio.sampleRate / SR; + if (Number.isInteger(k)) { // fast path (e.g. 44.1k): boxcar average + decimate; the sinc resampler is ~9 s per 5 min + const y = new Float32Array(Math.floor(x.length / k)); + for (let i = 0; i < y.length; i++) { let s = 0; for (let q = 0; q < k; q++) s += x[i * k + q]!; y[i] = s / k; } + x = y; + } else x = resample(x, audio.sampleRate, SR); + } + const nOut = Math.max(1, 1 + Math.floor((x.length / SR) * rate)); + const nBase = 1 + Math.floor(x.length / BASE_HOP); + + // pitch-class map for chroma bins + const binPc = new Int8Array((N_CHROMA >> 1) + 1).fill(-1); + for (let k = 1; k < binPc.length; k++) { + const f = (k * SR) / N_CHROMA; + if (f < F_MIN || f > F_MAX) continue; + const midi = 69 + 12 * Math.log2(f / 440); + binPc[k] = (((Math.round(midi) % 12) + 12) % 12); + } + + const wc = hann(N_CHROMA), wo = hann(N_ONSET); + const re = new Float64Array(N_CHROMA), im = new Float64Array(N_CHROMA), mag = new Float64Array((N_CHROMA >> 1) + 1); + const baseChroma: Float32Array[] = []; + for (let t = 0; t < nBase; t++) { + frameMag(x, t * BASE_HOP, N_CHROMA, wc, re, im, mag); + const c = new Float32Array(12); + for (let k = 1; k < binPc.length; k++) { const pc = binPc[k]!; if (pc >= 0) c[pc]! += mag[k]! * mag[k]!; } + baseChroma.push(c); + } + + // onset strength: half-wave rectified log-magnitude flux on the short window + const baseFlux = new Float32Array(nBase); + let prevMag = new Float64Array((N_ONSET >> 1) + 1); + const re2 = new Float64Array(N_ONSET), im2 = new Float64Array(N_ONSET); + let curMag = new Float64Array((N_ONSET >> 1) + 1); + for (let t = 0; t < nBase; t++) { + frameMag(x, t * BASE_HOP, N_ONSET, wo, re2, im2, curMag); + let s = 0; + for (let k = 0; k < curMag.length; k++) { + const d = Math.log1p(100 * curMag[k]!) - Math.log1p(100 * prevMag[k]!); + if (t > 0 && d > 0) s += d; + } + baseFlux[t] = s; + [prevMag, curMag] = [curMag, prevMag]; + } + + // aggregate base grid -> output grid: mean of chroma energy, max of flux over each output frame's span + const chroma: Float32Array[] = []; + const onset = new Float32Array(nOut); + for (let o = 0; o < nOut; o++) { + const t0 = (o - 0.5) / rate * BASE_RATE, t1 = (o + 0.5) / rate * BASE_RATE; + let a = Math.max(0, Math.ceil(t0 - 1e-9)), b = Math.min(nBase - 1, Math.ceil(t1 - 1e-9) - 1); + if (b < a) a = b = Math.min(nBase - 1, Math.max(0, Math.round(o / rate * BASE_RATE))); + const c = new Float32Array(12); + let fl = 0; + for (let t = a; t <= b; t++) { + const bc = baseChroma[t]!; + for (let k = 0; k < 12; k++) c[k]! += bc[k]!; + if (baseFlux[t]! > fl) fl = baseFlux[t]!; + } + let norm = 0; + for (let k = 0; k < 12; k++) { c[k] = Math.sqrt(c[k]!); norm += c[k]! * c[k]!; } + norm = Math.sqrt(norm); + if (norm > 1e-6) for (let k = 0; k < 12; k++) c[k]! /= norm; else c.fill(0); + chroma.push(c); + onset[o] = fl; + } + + // z-normalise onset strength + let mean = 0; + for (const v of onset) mean += v; + mean /= nOut; + let varr = 0; + for (const v of onset) varr += (v - mean) ** 2; + const sd = Math.sqrt(varr / nOut) || 1; + for (let i = 0; i < nOut; i++) onset[i] = (onset[i]! - mean) / sd; + return { rate, onset, chroma }; +} + +/** Peak-pick an onset-strength curve (adaptive threshold median + k·MAD, 50 ms minimum gap). Times in seconds. */ +export function pickOnsets(onset: Float32Array, rate: number, k = 4): number[] { + const n = onset.length; + if (n === 0) return []; + const sorted = Float32Array.from(onset).sort(); + const med = sorted[n >> 1]!; + const dev = Float32Array.from(onset, (v) => Math.abs(v - med)).sort(); + const mad = dev[n >> 1]! * 1.4826; + const thr = med + k * Math.max(mad, 1e-3); + const minGap = 0.05; + const picked: number[] = []; + for (let i = 0; i < n; i++) { + const v = onset[i]!; + if (v <= thr || v < (onset[i - 1] ?? -Infinity) || v <= (onset[i + 1] ?? -Infinity)) continue; + const t = i / rate; + const last = picked.length - 1; + if (last >= 0 && t - picked[last]! < minGap - 1e-9) { + if (v > onset[Math.round(picked[last]! * rate)]!) picked[last] = t; + } else picked.push(t); + } + return picked; +} diff --git a/tests/align/dtw.test.ts b/tests/align/dtw.test.ts new file mode 100644 index 0000000..906744a --- /dev/null +++ b/tests/align/dtw.test.ts @@ -0,0 +1,45 @@ +import { test, expect } from "bun:test"; +import { dtw } from "../../src/align/dtw.ts"; + +// Ground-truth path: diagonal with a kink (j runs ahead by 10 after i=50). +function truthJ(i: number): number { return i < 50 ? i : i + 10; } +const N = 100, M = 110; +// Path cells cost 0, others grow with distance from the true path. +const cost = (i: number, j: number) => { + const t = truthJ(i); + return Math.min(1, Math.abs(j - t) * 0.2); +}; + +test("dtw recovers diagonal-with-kink path", () => { + const p = dtw(cost, N, M); + expect(p[0]).toEqual([0, 0]); + expect(p[p.length - 1]).toEqual([N - 1, M - 1]); + for (const [i, j] of p) if (i !== 49 && i !== 50) expect(Math.abs(j - truthJ(i))).toBeLessThanOrEqual(1); + for (let k = 1; k < p.length; k++) { + const di = p[k]![0] - p[k - 1]![0], dj = p[k]![1] - p[k - 1]![1]; + expect(di >= 0 && di <= 1 && dj >= 0 && dj <= 1 && di + dj >= 1).toBe(true); + } +}); + +test("banded dtw equals unbanded when the band covers the path", () => { + const full = dtw(cost, N, M); + const banded = dtw(cost, N, M, { center: (i) => truthJ(i), width: 15 }); + expect(banded).toEqual(full); +}); + +test("banded dtw only visits cells in the band", () => { + let calls = 0; + const p = dtw((i, j) => { calls++; return Math.abs(i - j) * 0.1; }, 200, 200, { center: (i) => i, width: 5 }); + expect(calls).toBeLessThanOrEqual(200 * 11); + expect(p[p.length - 1]).toEqual([199, 199]); +}); + +test("3000x3000 unbanded dtw finishes in < 5 s", () => { + const n = 3000; + const t0 = performance.now(); + const p = dtw((i, j) => Math.abs(i - j) * 0.001 + ((i * 7 + j * 13) % 5) * 0.0001, n, n); + const ms = performance.now() - t0; + console.log(`dtw 3000x3000: ${ms.toFixed(0)} ms`); + expect(p[p.length - 1]).toEqual([n - 1, n - 1]); + expect(ms).toBeLessThan(5000); +}); diff --git a/tests/align/features.test.ts b/tests/align/features.test.ts new file mode 100644 index 0000000..b13bc6f --- /dev/null +++ b/tests/align/features.test.ts @@ -0,0 +1,57 @@ +import { test, expect } from "bun:test"; +import { features, pickOnsets } from "../../src/align/features.ts"; +import type { Audio } from "../../src/core/types.ts"; + +const SR = 44100; +function tone(freqs: number[], sec: number): Audio { + const x = new Float32Array(Math.round(sec * SR)); + for (let i = 0; i < x.length; i++) for (const f of freqs) x[i]! += 0.3 * Math.sin((2 * Math.PI * f * i) / SR); + return { sampleRate: SR, channels: [x] }; +} + +test("features rate and length", () => { + const f = features(tone([440], 3), 10); + expect(f.rate).toBe(10); + expect(Math.abs(f.chroma.length - 31)).toBeLessThanOrEqual(1); + expect(f.onset.length).toBe(f.chroma.length); + expect(f.chroma[0]!.length).toBe(12); + const g = features(tone([440], 3), 50); + expect(Math.abs(g.chroma.length - 151)).toBeLessThanOrEqual(1); +}); + +test("chroma peaks at the right pitch class and is unit-norm", () => { + const f = features(tone([440], 2), 10); // A -> pitch class 9 + const c = f.chroma[10]!; + let best = 0; + for (let k = 1; k < 12; k++) if (c[k]! > c[best]!) best = k; + expect(best).toBe(9); + expect(Math.hypot(...c)).toBeCloseTo(1, 4); + const fe = features(tone([261.63, 329.63, 392], 2), 10).chroma[10]!; // C E G + const top = [...fe].map((v, k) => [v, k] as const).sort((a, b) => b[0] - a[0]).slice(0, 3).map((e) => e[1]).sort((a, b) => a - b); + expect(top).toEqual([0, 4, 7]); +}); + +test("silent frames stay zero", () => { + const f = features({ sampleRate: SR, channels: [new Float32Array(SR * 2)] }, 10); + for (const c of f.chroma) expect(c.every((v) => v === 0)).toBe(true); +}); + +test("pickOnsets finds a click train within 10 ms", () => { + const x = new Float32Array(SR * 4); + const times = [0.5, 1.2, 2.0, 2.7, 3.3]; + for (const t of times) { + const s = Math.round(t * SR); + for (let k = 0; k < 200; k++) x[s + k] = Math.sin(k * 0.8) * Math.exp(-k / 40) * 0.8; + } + const f = features({ sampleRate: SR, channels: [x] }, 50); + const got = pickOnsets(f.onset, f.rate); + expect(got.length).toBe(times.length); + times.forEach((t, k) => expect(Math.abs(got[k]! - t)).toBeLessThanOrEqual(0.01)); +}); + +test("pickOnsets enforces 50 ms min gap and ignores flat input", () => { + const o = new Float32Array(100); + o[20] = 5; o[22] = 6; o[60] = 5; + expect(pickOnsets(o, 100)).toEqual([0.22, 0.6]); + expect(pickOnsets(new Float32Array(100), 100)).toEqual([]); +});