STFT Spectrogram: Five Test Signals

// Live STFT spectrogram waterfall. Synthesized audio cycles through five
// canonical test signals (linear chirp, quadratic chirp, AM two-tone, FSK
// square, noise burst) at fs=8 kHz. A 256-point Hann-windowed FFT with hop
// 64 (8 ms per column) renders magnitudes as a warm colormap.
//
// Frequency runs left->right (DC at left, Nyquist=4 kHz at right); time
// runs top->bottom with the newest spectrum at the top. The offscreen
// waterfall is a 128-bin x NROWS Uint8 image kept as a vertical ring
// buffer and blitted in two slices.

const FS = 8000;
const N = 256, HOP = 64, BINS = 128;
const NROWS = 256;
const SEG = 5.0; // seconds per signal type
const TYPES = ["linear chirp", "quadratic chirp", "AM two-tone", "FSK square", "noise burst"];

let han;          // Hann window
let buf;          // ring of N samples
let bufHead;      // next write index in buf
let samplePhase;  // accumulated audio phase (rad) for tonal generators
let fskPhase;     // phase for FSK
let sampleIdx;    // total samples generated
let nextFftAt;    // sampleIdx at which to compute next FFT column
let off, octx;    // offscreen waterfall image
let img1;         // 1-row ImageData stripe
let writeRow;     // next row to write in waterfall (newest)
let audioT;       // audio clock (seconds)
let segStart;     // audioT at which current segment began
let segIdx;       // which signal type
let re, im;       // FFT scratch
let revIdx;       // bit-reversal table

function init({ width, height }) {
  han = new Float32Array(N);
  for (let i = 0; i < N; i++) han[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / (N - 1));
  buf = new Float32Array(N);
  bufHead = 0; samplePhase = 0; fskPhase = 0; sampleIdx = 0; nextFftAt = N;
  audioT = 0; segStart = 0; segIdx = 0;
  writeRow = 0;
  off = new OffscreenCanvas(BINS, NROWS);
  octx = off.getContext("2d");
  // Pre-fill background.
  octx.fillStyle = "#08060a"; octx.fillRect(0, 0, BINS, NROWS);
  img1 = octx.createImageData(BINS, 1);
  re = new Float32Array(N); im = new Float32Array(N);
  revIdx = new Uint16Array(N);
  const bits = Math.log2(N) | 0;
  for (let i = 0; i < N; i++) {
    let r = 0, x = i;
    for (let b = 0; b < bits; b++) { r = (r << 1) | (x & 1); x >>= 1; }
    revIdx[i] = r;
  }
}

function fft() {
  // In-place iterative Cooley-Tukey on (re, im).
  for (let i = 0; i < N; i++) {
    const j = revIdx[i];
    if (j > i) { const tr = re[i]; re[i] = re[j]; re[j] = tr;
                 const ti = im[i]; im[i] = im[j]; im[j] = ti; }
  }
  for (let size = 2; size <= N; size <<= 1) {
    const half = size >> 1;
    const theta = -2 * Math.PI / size;
    const wpr = Math.cos(theta), wpi = Math.sin(theta);
    for (let k = 0; k < N; k += size) {
      let wr = 1, wi = 0;
      for (let m = 0; m < half; m++) {
        const a = k + m, b = a + half;
        const tr = wr * re[b] - wi * im[b];
        const ti = wr * im[b] + wi * re[b];
        re[b] = re[a] - tr; im[b] = im[a] - ti;
        re[a] = re[a] + tr; im[a] = im[a] + ti;
        const nwr = wr * wpr - wi * wpi;
        wi = wr * wpi + wi * wpr; wr = nwr;
      }
    }
  }
}

// Warm colormap: black -> dark red -> orange -> yellow -> white. v in [0,1].
function colormap(v, out, oi) {
  if (v < 0) v = 0; else if (v > 1) v = 1;
  let r, g, b;
  if (v < 0.25) { const t = v / 0.25; r = 60 * t; g = 0; b = 8 * (1 - t); }
  else if (v < 0.5) { const t = (v - 0.25) / 0.25; r = 60 + 175 * t; g = 30 * t; b = 0; }
  else if (v < 0.75) { const t = (v - 0.5) / 0.25; r = 235 + 20 * t; g = 30 + 165 * t; b = 0; }
  else { const t = (v - 0.75) / 0.25; r = 255; g = 195 + 60 * t; b = 255 * t; }
  out[oi] = r; out[oi + 1] = g; out[oi + 2] = b; out[oi + 3] = 255;
}

// Instantaneous frequency in Hz for the current segment at local time tau (s).
function instFreq(tau) {
  switch (segIdx) {
    case 0: return 200 + (3500 - 200) * (tau / SEG);                  // linear chirp 200 -> 3500
    case 1: return 200 + 3300 * (tau / SEG) * (tau / SEG);            // quadratic chirp
    case 2: return 0;                                                  // handled inline
    case 3: return 0;                                                  // handled inline
    case 4: return 0;                                                  // noise
  }
  return 0;
}

function genSample(tau) {
  if (segIdx === 0 || segIdx === 1) {
    const f = instFreq(tau);
    samplePhase += (2 * Math.PI * f) / FS;
    return Math.sin(samplePhase);
  }
  if (segIdx === 2) {
    const fc = 1500, fm = 60;
    samplePhase += (2 * Math.PI * fc) / FS;
    const am = 1 + 0.85 * Math.cos((2 * Math.PI * fm * tau));
    return 0.55 * am * Math.sin(samplePhase);
  }
  if (segIdx === 3) {
    // FSK at 30 baud between 800 Hz (mark) and 2400 Hz (space).
    const baud = 30;
    const bit = (Math.floor(tau * baud) & 1) === 0;
    const f = bit ? 800 : 2400;
    fskPhase += (2 * Math.PI * f) / FS;
    return Math.sin(fskPhase);
  }
  // Noise burst with slow amplitude breathing.
  const env = 0.4 + 0.6 * Math.abs(Math.sin(2 * Math.PI * 0.6 * tau));
  return env * (Math.random() * 2 - 1);
}

function pushColumn() {
  // Read N samples ending at bufHead-1 (newest), windowed.
  for (let i = 0; i < N; i++) {
    const v = buf[(bufHead + i) % N];
    re[i] = v * han[i]; im[i] = 0;
  }
  fft();
  const data = img1.data;
  // Magnitude -> dB, normalize to ~[0,1]. Floor at -60 dB, ceil at 0 dB.
  const FLOOR = -60, CEIL = -5, SPAN = CEIL - FLOOR;
  for (let k = 0; k < BINS; k++) {
    const mag = Math.sqrt(re[k] * re[k] + im[k] * im[k]) / (N * 0.5);
    const db = 20 * Math.log10(mag + 1e-9);
    let v = (db - FLOOR) / SPAN;
    colormap(v, data, k * 4);
  }
  octx.putImageData(img1, 0, writeRow);
  writeRow = (writeRow + 1) % NROWS;
}

function tick({ ctx, dt, width, height }) {
  const stepDt = Math.min(0.1, dt);
  const targetSamples = Math.floor(stepDt * FS);
  for (let s = 0; s < targetSamples; s++) {
    const tau = audioT - segStart;
    buf[bufHead] = genSample(tau);
    bufHead = (bufHead + 1) % N;
    sampleIdx++; audioT += 1 / FS;
    if (audioT - segStart >= SEG) {
      segIdx = (segIdx + 1) % TYPES.length;
      segStart = audioT;
      samplePhase = 0; fskPhase = 0;
    }
    if (sampleIdx >= nextFftAt) { pushColumn(); nextFftAt += HOP; }
  }

  // Background.
  ctx.fillStyle = "#06050a"; ctx.fillRect(0, 0, width, height);

  // Reserve a strip at the top for the label + frequency axis.
  const topPad = 30, botPad = 22, leftPad = 4, rightPad = 4;
  const ww = width - leftPad - rightPad;
  const hh = height - topPad - botPad;

  // Draw waterfall: newest row (writeRow-1) at the top.
  // Newest stripe sits in rows [0..writeRow), older stripe in [writeRow..NROWS).
  // We want oldest at bottom, newest at top -> draw [writeRow..NROWS) at top,
  // then [0..writeRow) below it.
  ctx.imageSmoothingEnabled = true;
  const partA = NROWS - writeRow; // rows starting at writeRow, count partA
  if (partA > 0) {
    const h1 = (hh * partA) / NROWS;
    ctx.drawImage(off, 0, writeRow, BINS, partA, leftPad, topPad, ww, h1);
  }
  if (writeRow > 0) {
    const h2 = (hh * writeRow) / NROWS;
    ctx.drawImage(off, 0, 0, BINS, writeRow, leftPad, topPad + (hh - h2), ww, h2);
  }

  // Label.
  ctx.fillStyle = "rgba(20,16,28,0.85)";
  ctx.fillRect(0, 0, width, topPad);
  ctx.fillStyle = "rgba(255,210,150,0.95)";
  ctx.font = "bold 13px system-ui, sans-serif";
  ctx.fillText("Signal: " + TYPES[segIdx], 10, 19);
  // Segment progress bar.
  const prog = (audioT - segStart) / SEG;
  const barX = 180, barY = 10, barW = Math.min(220, width - barX - 10), barH = 8;
  if (barW > 30) {
    ctx.strokeStyle = "rgba(255,210,150,0.5)"; ctx.lineWidth = 1;
    ctx.strokeRect(barX + 0.5, barY + 0.5, barW, barH);
    ctx.fillStyle = "rgba(255,170,90,0.8)";
    ctx.fillRect(barX + 1, barY + 1, (barW - 1) * prog, barH - 1);
  }

  // Frequency axis labels at bottom.
  ctx.fillStyle = "rgba(20,16,28,0.85)";
  ctx.fillRect(0, height - botPad, width, botPad);
  ctx.fillStyle = "rgba(200,200,220,0.85)";
  ctx.font = "10px system-ui, sans-serif";
  const ticks = [0, 1000, 2000, 3000, 4000];
  for (const f of ticks) {
    const x = leftPad + (f / (FS / 2)) * ww;
    ctx.fillRect(x, height - botPad, 1, 4);
    const lbl = f === 0 ? "0" : (f / 1000) + "k";
    ctx.fillText(lbl, x - 6, height - botPad + 16);
  }
  ctx.fillText("frequency (Hz) ->", leftPad, height - 4);
  // Time arrow on right side.
  ctx.save();
  ctx.translate(width - 12, topPad + hh / 2);
  ctx.rotate(-Math.PI / 2);
  ctx.fillStyle = "rgba(180,180,210,0.6)";
  ctx.fillText("time (older below)", -52, 0);
  ctx.restore();
}