Perona-Malik: Edge-Preserving Diffusion

// Perona-Malik anisotropic diffusion vs. linear (isotropic) heat smoothing.
//
//   d_t u = div( g(|grad u|) grad u ),   g(s) = 1 / (1 + (s/K)^2)
//
// vs. the plain heat equation d_t u = Lap(u). Same noisy initial image on
// both sides; PM preserves edges, heat smears them.
//
// Inspired by Gabriel Peyre's diffusion-and-image-processing visualizations.

const GW = 200;
const GH = 150;
const N = GW * GH;

// Numerical parameters.
const K_DEFAULT = 0.12;     // edge threshold for g(s)
const K_MIN = 0.02;         // strongest edges only
const K_MAX = 0.50;         // blurs almost everything
let K = K_DEFAULT;
let K2 = K * K;
const DT = 0.18;            // step in the PDE time (dx = 1)
const SUBSTEPS = 2;         // sim steps per render frame
const RESEED_SECONDS = 10;  // re-noise + restart every ~10 s

// Two fields evolved in parallel from the same noisy seed.
let uPM;       // Perona-Malik
let uPMNext;
let uLin;      // linear heat
let uLinNext;

// Pixel buffers (one per panel).
let bufPM, bufPMCtx, bufPMImg;
let bufLin, bufLinCtx, bufLinImg;

let W, H;
let simTime;
let secondsSinceReseed;

// Last computed panel layout, kept around so the click handler at the top
// of tick() can decide whether a click landed on a panel (either one).
let lastLayout = null;

// Synthetic noisy image: a few overlapping disks with distinct intensities
// on a flat background, plus Gaussian noise. Range roughly in [0, 1].
function makeNoisyImage(target) {
  // Background level.
  const bg = 0.45;

  // Disks: (cx, cy, r, intensity). Coordinates in grid units.
  const disks = [
    { cx: 65,  cy: 55,  r: 32, v: 0.85 },
    { cx: 130, cy: 70,  r: 38, v: 0.18 },
    { cx: 95,  cy: 105, r: 26, v: 0.62 },
    { cx: 160, cy: 30,  r: 18, v: 0.95 },
    { cx: 40,  cy: 110, r: 20, v: 0.08 },
  ];

  for (let y = 0; y < GH; y++) {
    for (let x = 0; x < GW; x++) {
      let v = bg;
      // Last-disk-wins painting, so overlaps produce sharp boundaries.
      for (let d = 0; d < disks.length; d++) {
        const dx = x - disks[d].cx;
        const dy = y - disks[d].cy;
        if (dx * dx + dy * dy <= disks[d].r * disks[d].r) {
          v = disks[d].v;
        }
      }
      // Gaussian noise via Box-Muller.
      const u1 = Math.random() || 1e-9;
      const u2 = Math.random();
      const n = Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
      v += n * 0.06;
      // Soft clamp; we let the diffusion eat any residual overshoot.
      if (v < 0) v = 0;
      else if (v > 1) v = 1;
      target[y * GW + x] = v;
    }
  }
}

function reseed() {
  makeNoisyImage(uPM);
  // Copy PM's seed into the linear field so both sides see identical noise.
  uLin.set(uPM);
  simTime = 0;
  secondsSinceReseed = 0;
}

// Perona-Malik step on a Neumann-boundary grid.
//
// Flux form: u^{n+1} = u^n + dt * sum_{neighbor n} g(|du|) * du,
// where du = u_n - u_c and g(s^2) = 1/(1 + s^2/K^2). Using s^2 directly
// saves a sqrt and the result is the same.
function stepPM(dtSim) {
  for (let y = 0; y < GH; y++) {
    const y0 = y * GW;
    const ym = y === 0 ? y0 : (y - 1) * GW;       // Neumann (reflect)
    const yp = y === GH - 1 ? y0 : (y + 1) * GW;
    for (let x = 0; x < GW; x++) {
      const xm = x === 0 ? x : x - 1;
      const xp = x === GW - 1 ? x : x + 1;
      const c = uPM[y0 + x];
      const dN = uPM[ym + x] - c;
      const dS = uPM[yp + x] - c;
      const dE = uPM[y0 + xp] - c;
      const dW = uPM[y0 + xm] - c;
      const gN = 1 / (1 + (dN * dN) / K2);
      const gS = 1 / (1 + (dS * dS) / K2);
      const gE = 1 / (1 + (dE * dE) / K2);
      const gWv = 1 / (1 + (dW * dW) / K2);
      uPMNext[y0 + x] = c + dtSim * (gN * dN + gS * dS + gE * dE + gWv * dW);
    }
  }
  const tmp = uPM;
  uPM = uPMNext;
  uPMNext = tmp;
}

// Linear heat equation, same Neumann BC, same dt. With dt=0.18 we're a bit
// over the strict 1/4 explicit-Euler limit, but the field is bounded and
// the visual is the right one (edges smear away).
function stepLin(dtSim) {
  for (let y = 0; y < GH; y++) {
    const y0 = y * GW;
    const ym = y === 0 ? y0 : (y - 1) * GW;
    const yp = y === GH - 1 ? y0 : (y + 1) * GW;
    for (let x = 0; x < GW; x++) {
      const xm = x === 0 ? x : x - 1;
      const xp = x === GW - 1 ? x : x + 1;
      const c = uLin[y0 + x];
      const lap = uLin[ym + x] + uLin[yp + x] + uLin[y0 + xp] + uLin[y0 + xm] - 4 * c;
      uLinNext[y0 + x] = c + dtSim * 0.22 * lap;
    }
  }
  const tmp = uLin;
  uLin = uLinNext;
  uLinNext = tmp;
}

// Grayscale colormap with a slight gamma so mid-tones don't go muddy.
function grayByte(v) {
  if (v < 0) v = 0;
  else if (v > 1) v = 1;
  // mild gamma to keep contrast Peyre-ish
  const g = Math.pow(v, 0.95);
  return Math.round(g * 255);
}

function paintInto(field, bufCtx, bufImg) {
  const data = bufImg.data;
  let off = 0;
  for (let i = 0; i < N; i++) {
    const g = grayByte(field[i]);
    data[off] = g;
    data[off + 1] = g;
    data[off + 2] = g;
    data[off + 3] = 255;
    off += 4;
  }
  bufCtx.putImageData(bufImg, 0, 0);
}

function init({ ctx, width, height }) {
  W = width;
  H = height;

  uPM = new Float32Array(N);
  uPMNext = new Float32Array(N);
  uLin = new Float32Array(N);
  uLinNext = new Float32Array(N);

  bufPM = new OffscreenCanvas(GW, GH);
  bufPMCtx = bufPM.getContext('2d');
  bufPMImg = bufPMCtx.createImageData(GW, GH);

  bufLin = new OffscreenCanvas(GW, GH);
  bufLinCtx = bufLin.getContext('2d');
  bufLinImg = bufLinCtx.createImageData(GW, GH);

  reseed();

  ctx.fillStyle = '#f4f3ee';
  ctx.fillRect(0, 0, W, H);
}

function tick({ ctx, dt, width, height, input }) {
  if (width !== W || height !== H) {
    W = width;
    H = height;
  }

  // Scrub K from mouseY. Top of canvas → K_MIN (only the strongest edges
  // survive), bottom → K_MAX (everything blurs). We let the diffusion keep
  // evolving with the new K — watching edge sensitivity drift mid-flight is
  // the point.
  if (input && typeof input.mouseY === 'number' && H > 0) {
    const t = Math.max(0, Math.min(1, input.mouseY / H));
    K = K_MIN + t * (K_MAX - K_MIN);
    K2 = K * K;
  }

  // Click anywhere on either panel → fresh noise on both sides.
  if (input && typeof input.consumeClicks === 'function') {
    const clicks = input.consumeClicks();
    if (clicks && clicks.length && lastLayout) {
      const { xLeft, xRight, yTop, drawW, drawH } = lastLayout;
      for (let i = 0; i < clicks.length; i++) {
        const cx = clicks[i].x;
        const cy = clicks[i].y;
        const onLeft = cx >= xLeft && cx <= xLeft + drawW && cy >= yTop && cy <= yTop + drawH;
        const onRight = cx >= xRight && cx <= xRight + drawW && cy >= yTop && cy <= yTop + drawH;
        if (onLeft || onRight) {
          reseed();
          break;
        }
      }
    }
  }

  // Advance both PDEs on the same clock.
  for (let s = 0; s < SUBSTEPS; s++) {
    stepPM(DT);
    stepLin(DT);
    simTime += DT;
  }

  secondsSinceReseed += dt;
  if (secondsSinceReseed >= RESEED_SECONDS) {
    reseed();
  }

  paintInto(uPM, bufPMCtx, bufPMImg);
  paintInto(uLin, bufLinCtx, bufLinImg);

  // Layout: two panels side by side with a thin gutter and a band of
  // page-cream above for labels.
  const labelH = 22;
  const gutter = 10;
  const margin = 12;
  const availW = W - 2 * margin - gutter;
  const panelW = Math.floor(availW / 2);
  const availH = H - labelH - 2 * margin;
  // Honor the 4:3 image aspect (GW:GH = 4:3).
  let drawH = availH;
  let drawW = Math.floor(drawH * (GW / GH));
  if (drawW > panelW) {
    drawW = panelW;
    drawH = Math.floor(drawW * (GH / GW));
  }
  const yTop = margin + labelH + Math.floor((availH - drawH) / 2);
  const xLeft = margin + Math.floor((panelW - drawW) / 2);
  const xRight = margin + panelW + gutter + Math.floor((panelW - drawW) / 2);

  // Stash for the next frame's click hit-test.
  lastLayout = { xLeft, xRight, yTop, drawW, drawH };

  // Background.
  ctx.fillStyle = '#f4f3ee';
  ctx.fillRect(0, 0, W, H);

  // Thin frames.
  ctx.imageSmoothingEnabled = true;
  ctx.imageSmoothingQuality = 'high';
  ctx.drawImage(bufPM, xLeft, yTop, drawW, drawH);
  ctx.drawImage(bufLin, xRight, yTop, drawW, drawH);

  ctx.strokeStyle = 'rgba(20, 20, 20, 0.7)';
  ctx.lineWidth = 1;
  ctx.strokeRect(xLeft - 0.5, yTop - 0.5, drawW + 1, drawH + 1);
  ctx.strokeRect(xRight - 0.5, yTop - 0.5, drawW + 1, drawH + 1);

  // Labels above each panel.
  ctx.fillStyle = 'rgba(20, 20, 20, 0.85)';
  ctx.font = '13px ui-sans-serif, system-ui, sans-serif';
  ctx.textBaseline = 'alphabetic';
  ctx.textAlign = 'center';
  ctx.fillText('Perona-Malik (edge-preserving)', xLeft + drawW / 2, yTop - 6);
  ctx.fillText('Linear heat (isotropic)', xRight + drawW / 2, yTop - 6);

  // HUD bottom-left: K and elapsed time.
  ctx.textAlign = 'left';
  ctx.fillStyle = 'rgba(20, 20, 20, 0.65)';
  ctx.font = '12px ui-sans-serif, system-ui, sans-serif';
  ctx.textBaseline = 'bottom';
  ctx.fillText('K = ' + K.toFixed(3) + '    t = ' + simTime.toFixed(1) + '    (drag Y · click to reseed)', margin, H - 6);
}