Stable Fluids

// Jos Stam's "Stable Fluids" — semi-Lagrangian advection of velocity and
// density, with a Jacobi-iterated pressure projection that enforces
// incompressibility (div u = 0). Drag the mouse to inject smoke and
// velocity. Light upward buoyancy makes the smoke rise.
const GW = 80, GH = 60;
const N = GW * GH;
const ITERS = 12;

let u, v, u0, v0;          // velocity + scratch
let d, d0;                 // density + scratch
let p, div;                // pressure + divergence
let off, octx, img, pix;   // GW x GH backbuffer
let cw, ch;
let lastMX = -1, lastMY = -1, lastDown = false;

function clamp(a, lo, hi) { return a < lo ? lo : a > hi ? hi : a; }

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

function init({ ctx, width, height }) {
  cw = width; ch = height;
  u = new Float32Array(N); v = new Float32Array(N);
  u0 = new Float32Array(N); v0 = new Float32Array(N);
  d = new Float32Array(N); d0 = new Float32Array(N);
  p = new Float32Array(N); div = new Float32Array(N);
  off = new OffscreenCanvas(GW, GH);
  octx = off.getContext("2d");
  img = octx.createImageData(GW, GH);
  pix = img.data;
  for (let i = 3; i < pix.length; i += 4) pix[i] = 255;
  // small seed so first frame already shows a plume
  for (let y = GH - 8; y < GH - 2; y++) {
    for (let x = GW / 2 - 4; x < GW / 2 + 4; x++) {
      d[idx(x | 0, y)] = 0.9;
      v[idx(x | 0, y)] = -0.6;
    }
  }
}

// Semi-Lagrangian advection: trace each cell backwards along the velocity
// field and bilinearly sample the source field at the foot of the
// characteristic.
function advect(out, src, vx, vy, dt) {
  const dt0x = dt * GW, dt0y = dt * GH;
  for (let y = 1; y < GH - 1; y++) {
    for (let x = 1; x < GW - 1; x++) {
      const i = idx(x, y);
      let sx = x - dt0x * vx[i];
      let sy = y - dt0y * vy[i];
      sx = clamp(sx, 0.5, GW - 1.5);
      sy = clamp(sy, 0.5, GH - 1.5);
      const x0 = sx | 0, y0 = sy | 0;
      const fx = sx - x0, fy = sy - y0;
      const j = y0 * GW + x0;
      out[i] =
        (1 - fx) * ((1 - fy) * src[j]      + fy * src[j + GW]) +
              fx * ((1 - fy) * src[j + 1]  + fy * src[j + GW + 1]);
    }
  }
}

// Project velocity onto its divergence-free component. Solve
// laplacian(p) = div(u) with Jacobi, then subtract grad(p).
function project() {
  const h = 1 / Math.max(GW, GH);
  for (let y = 1; y < GH - 1; y++) {
    for (let x = 1; x < GW - 1; x++) {
      const i = idx(x, y);
      div[i] = -0.5 * h * (u[i + 1] - u[i - 1] + v[i + GW] - v[i - GW]);
      p[i] = 0;
    }
  }
  for (let k = 0; k < ITERS; k++) {
    for (let y = 1; y < GH - 1; y++) {
      for (let x = 1; x < GW - 1; x++) {
        const i = idx(x, y);
        p[i] = 0.25 * (div[i] + p[i - 1] + p[i + 1] + p[i - GW] + p[i + GW]);
      }
    }
  }
  for (let y = 1; y < GH - 1; y++) {
    for (let x = 1; x < GW - 1; x++) {
      const i = idx(x, y);
      u[i] -= 0.5 * (p[i + 1] - p[i - 1]) / h;
      v[i] -= 0.5 * (p[i + GW] - p[i - GW]) / h;
    }
  }
}

function inject(px, py, dx, dy) {
  const gx = px / cw * GW;
  const gy = py / ch * GH;
  // velocity from drag, scaled to grid units per second
  const vx = (dx / cw) * GW * 8;
  const vy = (dy / ch) * GH * 8;
  const r = 4;
  for (let oy = -r; oy <= r; oy++) {
    for (let ox = -r; ox <= r; ox++) {
      const xi = (gx + ox) | 0, yi = (gy + oy) | 0;
      if (xi < 1 || yi < 1 || xi >= GW - 1 || yi >= GH - 1) continue;
      const d2 = ox * ox + oy * oy;
      if (d2 > r * r) continue;
      const k = Math.exp(-d2 * 0.25);
      const i = idx(xi, yi);
      d[i] = Math.min(1.6, d[i] + 0.9 * k);
      u[i] += vx * k * 0.5;
      v[i] += vy * k * 0.5;
    }
  }
}

function tick({ ctx, dt, width, height, input }) {
  if (width !== cw || height !== ch) { cw = width; ch = height; }
  const stepDt = Math.min(0.033, Math.max(0.008, dt));

  if (input.mouseDown) {
    if (lastDown && lastMX >= 0) {
      inject(input.mouseX, input.mouseY, input.mouseX - lastMX, input.mouseY - lastMY);
    } else {
      inject(input.mouseX, input.mouseY, 0, 0);
    }
    lastMX = input.mouseX; lastMY = input.mouseY; lastDown = true;
  } else {
    lastDown = false; lastMX = -1; lastMY = -1;
  }

  // ambient buoyancy: density pulls velocity upward
  for (let i = 0; i < N; i++) v[i] -= d[i] * 0.6 * stepDt;

  // velocity step: self-advect, then project to incompressible
  u0.set(u); v0.set(v);
  advect(u, u0, u0, v0, stepDt);
  advect(v, v0, u0, v0, stepDt);
  // zero normal velocity at boundaries
  for (let x = 0; x < GW; x++) { u[idx(x, 0)] = 0; u[idx(x, GH - 1)] = 0; v[idx(x, 0)] = 0; v[idx(x, GH - 1)] = 0; }
  for (let y = 0; y < GH; y++) { u[idx(0, y)] = 0; u[idx(GW - 1, y)] = 0; v[idx(0, y)] = 0; v[idx(GW - 1, y)] = 0; }
  project();

  // density step: advect, then mild dissipation
  d0.set(d);
  advect(d, d0, u, v, stepDt);
  const decay = Math.exp(-0.25 * stepDt);
  for (let i = 0; i < N; i++) d[i] *= decay;

  // warm palette: dark plum -> red -> orange -> yellow -> white
  for (let i = 0, q = 0; i < N; i++, q += 4) {
    let t = d[i];
    if (t > 1) t = 1; else if (t < 0) t = 0;
    // piecewise warm ramp
    const r = t < 0.5 ? 60 + 390 * t       : 255;
    const g = t < 0.4 ? 20 + 140 * t       : t < 0.85 ? 100 + 360 * (t - 0.4) : 240 + 60 * (t - 0.85) / 0.15;
    const b = t < 0.6 ? 20 + 20 * t        : 30 + 280 * (t - 0.6);
    pix[q] = r > 255 ? 255 : r;
    pix[q + 1] = g > 255 ? 255 : g;
    pix[q + 2] = b > 255 ? 255 : b;
  }

  octx.putImageData(img, 0, 0);
  ctx.imageSmoothingEnabled = true;
  ctx.drawImage(off, 0, 0, width, height);

  ctx.fillStyle = "rgba(20, 12, 10, 0.85)";
  ctx.fillRect(8, height - 28, 220, 20);
  ctx.fillStyle = "rgba(255, 220, 180, 0.95)";
  ctx.font = "12px system-ui, sans-serif";
  ctx.fillText("drag to inject smoke", 16, height - 14);
}