path: root/tools/mq_editor/mq_extract.js

// MQ Extraction Algorithm
// McAulay-Quatieri sinusoidal analysis

// Extract partials from audio buffer
function extractPartials(params, stftCache) {
  const {fftSize, threshold, sampleRate, freqWeight, prominence} = params;
  const numFrames = stftCache.getNumFrames();

  const frames = [];
  for (let i = 0; i < numFrames; ++i) {
    const cachedFrame = stftCache.getFrameAtIndex(i);
    const squaredAmp = stftCache.getSquaredAmplitude(cachedFrame.time);
    const peaks = detectPeaks(squaredAmp, fftSize, sampleRate, threshold, freqWeight, prominence);
    frames.push({time: cachedFrame.time, peaks});
  }

  const partials = trackPartials(frames);

  // Second pass: extend partials leftward to recover onset frames
  expandPartialsLeft(partials, frames);

  for (const partial of partials) {
    partial.freqCurve = fitBezier(partial.times, partial.freqs);
    partial.ampCurve = fitBezier(partial.times, partial.amps);
  }

  return {partials, frames};
}

// Detect spectral peaks via local maxima + parabolic interpolation
// squaredAmp: pre-computed re*re+im*im per bin
// freqWeight: if true, weight by f before peak detection (f * Power(f))
function detectPeaks(squaredAmp, fftSize, sampleRate, thresholdDB, freqWeight, prominenceDB = 0) {
  const mag = new Float32Array(fftSize / 2);
  const binHz = sampleRate / fftSize;
  for (let i = 0; i < fftSize / 2; ++i) {
    const w = freqWeight ? (i * binHz) : 1.0;
    mag[i] = 10 * Math.log10(Math.max(squaredAmp[i] * w, 1e-20));
  }

  const peaks = [];
  for (let i = 2; i < mag.length - 2; ++i) {
    if (mag[i] > thresholdDB &&
        mag[i] > mag[i-1] && mag[i] > mag[i-2] &&
        mag[i] > mag[i+1] && mag[i] > mag[i+2]) {

      // Check prominence if requested
      if (prominenceDB > 0) {
        let minLeft = mag[i];
        for (let k = i - 1; k >= 0; --k) {
          if (mag[k] > mag[i]) break; // Found higher peak
          if (mag[k] < minLeft) minLeft = mag[k];
        }

        let minRight = mag[i];
        for (let k = i + 1; k < mag.length; ++k) {
          if (mag[k] > mag[i]) break; // Found higher peak
          if (mag[k] < minRight) minRight = mag[k];
        }

        const valley = Math.max(minLeft, minRight);
        if (mag[i] - valley < prominenceDB) continue;
      }

      // Parabolic interpolation for sub-bin accuracy
      const alpha = mag[i-1];
      const beta  = mag[i];
      const gamma = mag[i+1];
      const p = 0.5 * (alpha - gamma) / (alpha - 2*beta + gamma);

      const freq  = (i + p) * sampleRate / fftSize;
      const ampDB = beta - 0.25 * (alpha - gamma) * p;
      peaks.push({freq, amp: Math.pow(10, ampDB / 20)});
    }
  }

  return peaks;
}

// Track partials across frames (birth/death/continuation)
function trackPartials(frames) {
  const partials        = [];
  const activePartials  = [];
  const candidates      = []; // pre-birth

  const trackingRatio   = 0.05; // 5% frequency tolerance
  const minTrackingHz   = 20;
  const birthPersistence = 3;   // frames before partial is born
  const deathAge        = 5;    // frames without match before death
  const minLength       = 10;   // frames required to keep partial

  for (const frame of frames) {
    const matched = new Set();

    // Continue active partials
    for (const partial of activePartials) {
      const lastFreq = partial.freqs[partial.freqs.length - 1];
      const velocity = partial.velocity || 0;
      const predicted = lastFreq + velocity;
      
      const tol = Math.max(lastFreq * trackingRatio, minTrackingHz);
      let bestIdx = -1, bestDist = Infinity;

      for (let i = 0; i < frame.peaks.length; ++i) {
        if (matched.has(i)) continue;
        const dist = Math.abs(frame.peaks[i].freq - predicted);
        if (dist < tol && dist < bestDist) { bestDist = dist; bestIdx = i; }
      }

      if (bestIdx >= 0) {
        const pk = frame.peaks[bestIdx];
        partial.times.push(frame.time);
        partial.freqs.push(pk.freq);
        partial.amps.push(pk.amp);
        partial.age = 0;
        partial.velocity = pk.freq - lastFreq;
        matched.add(bestIdx);
      } else {
        partial.age++;
      }
    }

    // Advance candidates
    for (let i = candidates.length - 1; i >= 0; --i) {
      const cand = candidates[i];
      const lastFreq = cand.freqs[cand.freqs.length - 1];
      const velocity = cand.velocity || 0;
      const predicted = lastFreq + velocity;

      const tol = Math.max(lastFreq * trackingRatio, minTrackingHz);
      let bestIdx = -1, bestDist = Infinity;

      for (let j = 0; j < frame.peaks.length; ++j) {
        if (matched.has(j)) continue;
        const dist = Math.abs(frame.peaks[j].freq - predicted);
        if (dist < tol && dist < bestDist) { bestDist = dist; bestIdx = j; }
      }

      if (bestIdx >= 0) {
        const pk = frame.peaks[bestIdx];
        cand.times.push(frame.time);
        cand.freqs.push(pk.freq);
        cand.amps.push(pk.amp);
        cand.velocity = pk.freq - lastFreq;
        matched.add(bestIdx);
        if (cand.times.length >= birthPersistence) {
          activePartials.push(cand);
          candidates.splice(i, 1);
        }
      } else {
        candidates.splice(i, 1);
      }
    }

    // Spawn candidates from unmatched peaks
    for (let i = 0; i < frame.peaks.length; ++i) {
      if (matched.has(i)) continue;
      const pk = frame.peaks[i];
      candidates.push({
        times: [frame.time], 
        freqs: [pk.freq], 
        amps: [pk.amp], 
        age: 0,
        velocity: 0
      });
    }

    // Kill aged-out partials
    for (let i = activePartials.length - 1; i >= 0; --i) {
      if (activePartials[i].age > deathAge) {
        if (activePartials[i].times.length >= minLength) partials.push(activePartials[i]);
        activePartials.splice(i, 1);
      }
    }
  }

  // Collect remaining active partials
  for (const partial of activePartials) {
    if (partial.times.length >= minLength) partials.push(partial);
  }

  return partials;
}

// Second pass: extend each partial leftward to recover onset frames missed
// by the birthPersistence requirement in the forward pass.
function expandPartialsLeft(partials, frames) {
  const trackingRatio = 0.05;
  const minTrackingHz = 20;

  // Build time → frame index map
  const timeToIdx = new Map();
  for (let i = 0; i < frames.length; ++i) timeToIdx.set(frames[i].time, i);

  for (const partial of partials) {
    let startIdx = timeToIdx.get(partial.times[0]);
    if (startIdx == null || startIdx === 0) continue;

    for (let i = startIdx - 1; i >= 0; --i) {
      const frame = frames[i];
      const refFreq = partial.freqs[0];
      const tol = Math.max(refFreq * trackingRatio, minTrackingHz);

      let bestIdx = -1, bestDist = Infinity;
      for (let j = 0; j < frame.peaks.length; ++j) {
        const dist = Math.abs(frame.peaks[j].freq - refFreq);
        if (dist < tol && dist < bestDist) { bestDist = dist; bestIdx = j; }
      }

      if (bestIdx < 0) break;

      const pk = frame.peaks[bestIdx];
      partial.times.unshift(frame.time);
      partial.freqs.unshift(pk.freq);
      partial.amps.unshift(pk.amp);
    }
  }
}

// Autodetect spread_above / spread_below from the spectrogram.
// For each (subsampled) STFT frame within the partial, measures the
// half-power (-3dB) width of the spectral peak above and below the center.
// spread = half_bandwidth / f0  (fractional).
function autodetectSpread(partial, stftCache, fftSize, sampleRate) {
  const curve = partial.freqCurve;
  if (!curve || !stftCache) return {spread_above: 0.02, spread_below: 0.02};

  const numFrames = stftCache.getNumFrames();
  const binHz    = sampleRate / fftSize;
  const halfBins = fftSize / 2;

  let sumAbove = 0, sumBelow = 0, count = 0;

  const STEP = 4;
  for (let fi = 0; fi < numFrames; fi += STEP) {
    const frame = stftCache.getFrameAtIndex(fi);
    if (!frame) continue;
    const t = frame.time;
    if (t < curve.t0 || t > curve.t3) continue;

    const f0 = evalBezier(curve, t);
    if (f0 <= 0) continue;

    const sq = frame.squaredAmplitude;
    if (!sq) continue;

    // Find peak bin in ±10% window
    const binCenter = f0 / binHz;
    const searchBins = Math.max(3, Math.round(f0 * 0.10 / binHz));
    const binLo = Math.max(1, Math.floor(binCenter - searchBins));
    const binHi = Math.min(halfBins - 2, Math.ceil(binCenter + searchBins));

    let peakBin = binLo, peakVal = sq[binLo];
    for (let b = binLo + 1; b <= binHi; ++b) {
      if (sq[b] > peakVal) { peakVal = sq[b]; peakBin = b; }
    }

    const halfPower = peakVal * 0.5;  // -3dB in power

    // Walk above peak until half-power, interpolate crossing
    let aboveBin = peakBin;
    while (aboveBin < halfBins - 1 && sq[aboveBin] > halfPower) ++aboveBin;
    const tA = aboveBin > peakBin && sq[aboveBin - 1] !== sq[aboveBin]
      ? (halfPower - sq[aboveBin - 1]) / (sq[aboveBin] - sq[aboveBin - 1])
      : 0;
    const widthAbove = (aboveBin - 1 + tA - peakBin) * binHz;

    // Walk below peak until half-power, interpolate crossing
    let belowBin = peakBin;
    while (belowBin > 1 && sq[belowBin] > halfPower) --belowBin;
    const tB = belowBin < peakBin && sq[belowBin + 1] !== sq[belowBin]
      ? (halfPower - sq[belowBin + 1]) / (sq[belowBin] - sq[belowBin + 1])
      : 0;
    const widthBelow = (peakBin - belowBin - 1 + tB) * binHz;

    sumAbove += (widthAbove / f0) * (widthAbove / f0);
    sumBelow += (widthBelow / f0) * (widthBelow / f0);
    ++count;
  }

  const spread_above = count > 0 ? Math.sqrt(sumAbove / count) : 0.01;
  const spread_below = count > 0 ? Math.sqrt(sumBelow / count) : 0.01;
  return {spread_above, spread_below};
}

// Fit cubic bezier to trajectory using least-squares for inner control points
function fitBezier(times, values) {
  const n = times.length - 1;
  const t0 = times[0], v0 = values[0];
  const t3 = times[n], v3 = values[n];
  const dt = t3 - t0;

  if (dt <= 1e-9 || n < 2) {
    // Linear fallback for too few points or zero duration
    return {t0, v0, t1: t0 + dt / 3, v1: v0 + (v3 - v0) / 3, t2: t0 + 2 * dt / 3, v2: v0 + 2 * (v3 - v0) / 3, t3, v3};
  }

  // Least squares solve for v1, v2
  // Bezier: B(u) = (1-u)^3*v0 + 3(1-u)^2*u*v1 + 3(1-u)*u^2*v2 + u^3*v3
  // Target_i = val_i - (1-u)^3*v0 - u^3*v3
  // Model_i = A_i*v1 + B_i*v2
  // A_i = 3(1-u)^2*u
  // B_i = 3(1-u)*u^2

  let sA2 = 0, sB2 = 0, sAB = 0, sAT = 0, sBT = 0;

  for (let i = 0; i <= n; ++i) {
    const u = (times[i] - t0) / dt;
    const u2 = u * u;
    const u3 = u2 * u;
    const invU = 1.0 - u;
    const invU2 = invU * invU;
    const invU3 = invU2 * invU;

    const A = 3 * invU2 * u;
    const B = 3 * invU * u2;
    const target = values[i] - (invU3 * v0 + u3 * v3);

    sA2 += A * A;
    sB2 += B * B;
    sAB += A * B;
    sAT += A * target;
    sBT += B * target;
  }

  const det = sA2 * sB2 - sAB * sAB;
  let v1, v2;

  if (Math.abs(det) < 1e-9) {
    // Fallback to simple 1/3, 2/3 heuristic if matrix is singular
    const idx1 = Math.round(n / 3);
    const idx2 = Math.round(2 * n / 3);
    v1 = values[idx1];
    v2 = values[idx2];
  } else {
    v1 = (sB2 * sAT - sAB * sBT) / det;
    v2 = (sA2 * sBT - sAB * sAT) / det;
  }

  return {t0, v0, t1: t0 + dt / 3, v1, t2: t0 + 2 * dt / 3, v2, t3, v3};
}