feat(align): features and DTW
Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
This commit is contained in:
Lucas Winther
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
b614244e4d
commit
41e38cb03f
@@ -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();
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
@@ -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([]);
|
||||
});
|
||||
Reference in New Issue
Block a user