Fast Marching: Eikonal Equation

tap to set source

The fast marching method (Sethian, 1996) solves the eikonal equation $∥\nabla T (x) ∥ = 1/ F (x)$ , where $T (x)$ is the first-arrival time of a wavefront at point $x$ and $F (x) > 0$ is a spatially-varying propagation speed. Think of it as Dijkstra lifted from a graph to a continuous grid. Starting from a source where $T = 0$ , the algorithm maintains a narrow band of trial cells and, at every iteration, freezes the trial cell with the smallest $T$ — guaranteed correct because information propagates outward in $T$ , so a frozen value is never improved by a later neighbor. Each candidate $T$ is the upwind solution of $max (T - T_{x}, 0)^{2} + max (T - T_{y}, 0)^{2} = (h / F)^{2}$ , where $T_{x} = min (T_{i - 1, j}, T_{i + 1, j})$ over frozen neighbors (likewise $T_{y}$ ). When both axes contribute, that's a quadratic in $T$ ; when one upwind difference exceeds $h / F$ , the stencil degenerates and the answer is simply $T_{x} + h / F$ . Total cost is $O (N lo g N)$ for $N$ cells, dominated by the heap. Watch the front bend toward slow regions and accelerate through fast ones — classic refraction, no rays needed. Inspired by Gabriel Peyré's numerical-geometry visualizations.

idle

302 lines · vanilla

view source

// Fast Marching Method for the eikonal equation ||grad T|| = 1/F(x).
// Continuous-domain Dijkstra: maintain a narrow band, repeatedly extract the
// candidate with smallest T, then update its four neighbors via the upwind
// stencil:  max(T - Tx,0)^2 + max(T - Ty,0)^2 = (h/F)^2
// (Tx = min of the two horizontal neighbors that are FROZEN/known, Ty likewise.)
//
// Inspired by Gabriel Peyre's daily numerical-geometry gifs.

const GW = 140;          // grid resolution (cols)
const GH = 100;          // grid rows
// Cap heap pops per frame so the wavefront takes ~7s at 60fps to traverse the
// canvas. GW*GH = 14000 cells / (60fps * 7s) ~= 33 — round up a bit for headroom
// since some pops are stale (filtered out).
const PER_FRAME = 35;
const HOLD_FRAMES = 120; // ~2s @ 60fps to admire the finished field
const FAR = 1, BAND = 2, FROZEN = 3;

// --- Viridis colormap (sampled) ----------------------------------------------
const VIRIDIS = [
  [68, 1, 84], [72, 35, 116], [64, 67, 135], [52, 94, 141],
  [41, 120, 142], [32, 144, 140], [34, 167, 132], [68, 190, 112],
  [121, 209, 81], [189, 222, 38], [253, 231, 36],
];
function viridis(t) {
  if (t <= 0) return VIRIDIS[0];
  if (t >= 1) return VIRIDIS[VIRIDIS.length - 1];
  const s = t * (VIRIDIS.length - 1);
  const i = Math.floor(s);
  const f = s - i;
  const a = VIRIDIS[i], b = VIRIDIS[i + 1];
  return [
    a[0] + (b[0] - a[0]) * f,
    a[1] + (b[1] - a[1]) * f,
    a[2] + (b[2] - a[2]) * f,
  ];
}

// --- Min-heap of grid indices, keyed by T[idx] -------------------------------
class MinHeap {
  constructor(T) { this.T = T; this.data = []; }
  clear() { this.data.length = 0; }
  size() { return this.data.length; }
  push(i) {
    const d = this.data; d.push(i);
    let k = d.length - 1;
    const T = this.T;
    while (k > 0) {
      const p = (k - 1) >> 1;
      if (T[d[p]] <= T[d[k]]) break;
      const t = d[p]; d[p] = d[k]; d[k] = t;
      k = p;
    }
  }
  pop() {
    const d = this.data;
    const top = d[0];
    const last = d.pop();
    if (d.length) {
      d[0] = last;
      const T = this.T;
      const n = d.length;
      let k = 0;
      for (;;) {
        const l = 2 * k + 1, r = 2 * k + 2;
        let m = k;
        if (l < n && T[d[l]] < T[d[m]]) m = l;
        if (r < n && T[d[r]] < T[d[m]]) m = r;
        if (m === k) break;
        const t = d[k]; d[k] = d[m]; d[m] = t;
        k = m;
      }
    }
    return top;
  }
}

// --- State -------------------------------------------------------------------
let W, H, cellW, cellH;
let F;            // speed field, length GW*GH
let T;            // arrival time
let state;        // FAR / BAND / FROZEN
let heap;
let tMax;         // largest finalized T (for normalization)
let iter;
let done;
let holdLeft;
let srcX, srcY;
let pendingClick = null;
let bgCanvas, bgCtx;     // pre-rendered speed-field underlay
let fieldCanvas, fieldCtx; // upscaled T field
let fieldImage;          // ImageData for fieldCanvas (GW x GH)

function idx(x, y) { return y * GW + x; }

function makeSpeedField() {
  // Two overlapping sinusoidal modes — visually striking refraction pattern.
  // F stays positive: keep amplitude < 1.
  const k1x = 2 * Math.PI / GW * 2.0;
  const k1y = 2 * Math.PI / GH * 2.0;
  const k2x = 2 * Math.PI / GW * 3.0;
  const k2y = 2 * Math.PI / GH * 1.0;
  const phase = Math.random() * Math.PI * 2;
  for (let y = 0; y < GH; y++) {
    for (let x = 0; x < GW; x++) {
      const a = Math.sin(x * k1x + phase) * Math.cos(y * k1y - phase);
      const b = Math.sin(x * k2x - 0.3 * phase) * Math.sin(y * k2y + 0.7 * phase);
      // map into [0.35, 1.85] — never zero, plenty of contrast
      F[idx(x, y)] = 1.1 + 0.55 * a + 0.2 * b;
    }
  }
}

function renderSpeedBackground() {
  const img = bgCtx.createImageData(GW, GH);
  // F is in roughly [0.35, 1.85]; map to gray
  let lo = Infinity, hi = -Infinity;
  for (let i = 0; i < GW * GH; i++) {
    if (F[i] < lo) lo = F[i];
    if (F[i] > hi) hi = F[i];
  }
  const span = hi - lo || 1;
  for (let i = 0; i < GW * GH; i++) {
    const t = (F[i] - lo) / span;
    const g = (18 + t * 60) | 0; // dark gray range
    img.data[i * 4 + 0] = g;
    img.data[i * 4 + 1] = g;
    img.data[i * 4 + 2] = g + 4;
    img.data[i * 4 + 3] = 255;
  }
  bgCtx.putImageData(img, 0, 0);
}

function startMarch(sx, sy) {
  T.fill(Infinity);
  state.fill(FAR);
  heap.clear();
  tMax = 0;
  iter = 0;
  done = false;
  holdLeft = 0;
  srcX = sx;
  srcY = sy;
  const si = idx(sx, sy);
  T[si] = 0;
  state[si] = FROZEN;
  // seed band with the four neighbors
  pushNeighbors(sx, sy);
  // clear field image
  for (let i = 0; i < GW * GH; i++) {
    fieldImage.data[i * 4 + 0] = 0;
    fieldImage.data[i * 4 + 1] = 0;
    fieldImage.data[i * 4 + 2] = 0;
    fieldImage.data[i * 4 + 3] = 0;
  }
}

function solveQuadratic(x, y) {
  // h = 1 in grid units
  const f = F[idx(x, y)];
  const inv = 1 / f;
  // Tx = min of horizontal FROZEN neighbors
  let Tx = Infinity;
  if (x > 0      && state[idx(x - 1, y)] === FROZEN) Tx = T[idx(x - 1, y)];
  if (x < GW - 1 && state[idx(x + 1, y)] === FROZEN) {
    const v = T[idx(x + 1, y)];
    if (v < Tx) Tx = v;
  }
  let Ty = Infinity;
  if (y > 0      && state[idx(x, y - 1)] === FROZEN) Ty = T[idx(x, y - 1)];
  if (y < GH - 1 && state[idx(x, y + 1)] === FROZEN) {
    const v = T[idx(x, y + 1)];
    if (v < Ty) Ty = v;
  }
  // Solve max(T - Tx, 0)^2 + max(T - Ty, 0)^2 = inv^2
  // Try both axes
  if (!isFinite(Tx) && !isFinite(Ty)) return Infinity;
  if (!isFinite(Ty)) return Tx + inv;
  if (!isFinite(Tx)) return Ty + inv;
  const diff = Tx - Ty;
  if (Math.abs(diff) >= inv) {
    // upwind comes from a single axis
    return Math.min(Tx, Ty) + inv;
  }
  // Both axes contribute. Quadratic in T:
  //   (T - Tx)^2 + (T - Ty)^2 = inv^2
  //   2 T^2 - 2(Tx+Ty) T + (Tx^2 + Ty^2 - inv^2) = 0
  const sum = Tx + Ty;
  const disc = 2 * inv * inv - diff * diff;
  if (disc < 0) return Math.min(Tx, Ty) + inv;
  return 0.5 * (sum + Math.sqrt(disc));
}

function pushNeighbors(x, y) {
  const ns = [
    [x - 1, y], [x + 1, y], [x, y - 1], [x, y + 1],
  ];
  for (let k = 0; k < 4; k++) {
    const nx = ns[k][0], ny = ns[k][1];
    if (nx < 0 || nx >= GW || ny < 0 || ny >= GH) continue;
    const ni = idx(nx, ny);
    if (state[ni] === FROZEN) continue;
    const nt = solveQuadratic(nx, ny);
    if (nt < T[ni]) {
      T[ni] = nt;
      state[ni] = BAND;
      heap.push(ni); // duplicate entries are fine; we filter on pop
    }
  }
}

function marchStep() {
  while (heap.size()) {
    const i = heap.pop();
    if (state[i] === FROZEN) continue; // stale heap entry
    state[i] = FROZEN;
    if (T[i] > tMax) tMax = T[i];
    iter++;
    const y = (i / GW) | 0;
    const x = i - y * GW;
    // paint this pixel into the field image
    const t = tMax > 0 ? T[i] / tMax : 0;
    const tt = Math.min(1, t);
    const c = viridis(tt);
    const off = i * 4;
    fieldImage.data[off + 0] = c[0];
    fieldImage.data[off + 1] = c[1];
    fieldImage.data[off + 2] = c[2];
    fieldImage.data[off + 3] = 210;
    pushNeighbors(x, y);
    return true;
  }
  return false;
}

// --- Entry points ------------------------------------------------------------
function init({ canvas, ctx, width, height }) {
  W = width; H = height;
  cellW = W / GW;
  cellH = H / GH;

  F = new Float32Array(GW * GH);
  T = new Float64Array(GW * GH);
  state = new Uint8Array(GW * GH);
  heap = new MinHeap(T);

  bgCanvas = new OffscreenCanvas(GW, GH);
  bgCtx = bgCanvas.getContext('2d');
  fieldCanvas = new OffscreenCanvas(GW, GH);
  fieldCtx = fieldCanvas.getContext('2d');
  fieldImage = fieldCtx.createImageData(GW, GH);

  makeSpeedField();
  renderSpeedBackground();
  startMarch((GW * 0.5) | 0, (GH * 0.5) | 0);

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

function pickFrontEdge(ctx) {
  // Draw thin white edge between FROZEN and not-FROZEN cells, in grid space.
  ctx.save();
  ctx.scale(cellW, cellH);
  ctx.strokeStyle = 'rgba(255,255,255,0.78)';
  ctx.lineWidth = 0.08;
  ctx.beginPath();
  for (let y = 0; y < GH; y++) {
    for (let x = 0; x < GW; x++) {
      const i = idx(x, y);
      if (state[i] !== FROZEN) continue;
      // right neighbor
      if (x + 1 < GW && state[idx(x + 1, y)] !== FROZEN) {
        ctx.moveTo(x + 1, y);
        ctx.lineTo(x + 1, y + 1);
      }
      // bottom neighbor
      if (y + 1 < GH && state[idx(x, y + 1)] !== FROZEN) {
        ctx.moveTo(x, y + 1);
        ctx.lineTo(x + 1, y + 1);
      }
    }
  }
  ctx.stroke();
  ctx.restore();
}

function tick({ ctx, width, height, input }) {
  if (width !== W || height !== H) {
    W = width; H = height;
    cellW = W / GW;
    cellH = H / GH;
  }

  // Click handling — schedule a restart from clicked grid cell.
  const clicks = input.consumeClicks ? input.consumeClicks() : 0;
  if (clicks > 0) {
    const gx = Math.min(GW - 1, Math.max(0, (input.mouseX / cellW) | 0));
    const gy = Math.min(GH - 1, Math.max(0, (input.mouseY / cellH) | 0));
    pendingClick = { x: gx, y: gy };
  }

  if (pendingClick) {
    startMarch(pendingClick.x, pendingClick.y);
    pendingClick = null;
  }

  // Advance the march
  if (!done) {
    let did = 0;
    for (let s = 0; s < PER_FRAME; s++) {
      if (!marchStep()) { done = true; holdLeft = HOLD_FRAMES; break; }
      did++;
    }
    if (did > 0) fieldCtx.putImageData(fieldImage, 0, 0);
  } else if (holdLeft > 0) {
    holdLeft--;
    if (holdLeft === 0) {
      // pick a new random source for a fresh loop
      let nx, ny;
      do {
        nx = ((Math.random() * 0.8 + 0.1) * GW) | 0;
        ny = ((Math.random() * 0.8 + 0.1) * GH) | 0;
      } while (nx === srcX && ny === srcY);
      // re-randomize the speed field every few resets for variety
      if (Math.random() < 0.5) {
        makeSpeedField();
        renderSpeedBackground();
      }
      startMarch(nx, ny);
    }
  }

  // --- Render ---
  ctx.fillStyle = '#06070b';
  ctx.fillRect(0, 0, W, H);

  // Speed field underlay (upscaled, faint)
  ctx.imageSmoothingEnabled = true;
  ctx.globalAlpha = 0.55;
  ctx.drawImage(bgCanvas, 0, 0, W, H);

  // T field overlay
  ctx.globalAlpha = 1.0;
  ctx.drawImage(fieldCanvas, 0, 0, W, H);

  // Wavefront edge
  pickFrontEdge(ctx);

  // Source marker
  const sx = (srcX + 0.5) * cellW;
  const sy = (srcY + 0.5) * cellH;
  ctx.fillStyle = '#ffffff';
  ctx.beginPath();
  ctx.arc(sx, sy, Math.max(2.5, Math.min(cellW, cellH) * 0.8), 0, Math.PI * 2);
  ctx.fill();
  ctx.strokeStyle = 'rgba(0,0,0,0.6)';
  ctx.lineWidth = 1;
  ctx.stroke();

  // HUD
  ctx.fillStyle = 'rgba(10,10,15,0.6)';
  ctx.fillRect(8, 8, 168, 38);
  ctx.fillStyle = 'rgba(230,232,240,0.92)';
  ctx.font = '11px ui-monospace, monospace';
  ctx.textBaseline = 'top';
  ctx.fillText(`fast marching  iter ${iter}`, 14, 13);
  ctx.fillText(
    done ? `done  T_max ${tMax.toFixed(2)}` : `band ${heap.size()}  T_max ${tMax.toFixed(2)}`,
    14, 28,
  );
}

Fast Marching: Eikonal Equation

Comments (0)