STSheets/src/align/features.ts

// 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, the stronger peak wins).
 * Peak times are refined to sub-frame precision by parabolic interpolation over the peak and its two neighbours.
 * 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 frames: 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 last = frames.length - 1;
    if (last >= 0 && (i - frames[last]!) / rate < minGap - 1e-9) {
      if (v > onset[frames[last]!]!) frames[last] = i;
    } else frames.push(i);
  }
  return frames.map((i) => {
    if (i === 0 || i === n - 1) return i / rate;
    const a = onset[i - 1]!, b = onset[i]!, c = onset[i + 1]!;
    const den = a - 2 * b + c;
    const d = den < 0 ? Math.max(-0.5, Math.min(0.5, (0.5 * (a - c)) / den)) : 0;
    return (i + d) / rate;
  });
}