Shallow Water

// 2D shallow water equations in conservative form. State per cell:
//   h    — fluid depth
//   hu   — depth * x-velocity
//   hv   — depth * y-velocity
// Lax-Friedrichs flux on a regular grid, gravity g, reflective walls.
// Click to drop a Gaussian bump on h. Side-view inset shows the cross
// section through the cursor row.

let GW, GH, cw, ch;
let h, hu, hv;          // state
let h2, hu2, hv2;       // next-step buffers
let off, octx, img, pix;
const G = 9.81;         // m/s^2
const H0 = 1.0;         // rest depth
const DX = 0.04;        // meters per cell
let cursorY = -1;       // last cursor row for cross-section

function alloc(w, h_) {
  cw = w; ch = h_;
  // grid is square cells; pick GW/GH to roughly fit canvas aspect
  const target = 120;
  GW = target;
  GH = Math.max(40, Math.round(target * h_ / w));
  const N = GW * GH;
  h  = new Float32Array(N);
  hu = new Float32Array(N);
  hv = new Float32Array(N);
  h2 = new Float32Array(N);
  hu2 = new Float32Array(N);
  hv2 = new Float32Array(N);
  for (let i = 0; i < N; i++) h[i] = H0;
  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;
}

function drop(px, py, amp) {
  const gx = (px / cw) * GW;
  const gy = (py / ch) * GH;
  const r = 6;
  for (let oy = -r; oy <= r; oy++) {
    for (let ox = -r; ox <= r; ox++) {
      const xi = Math.round(gx + ox);
      const yi = Math.round(gy + oy);
      if (xi < 1 || yi < 1 || xi >= GW - 1 || yi >= GH - 1) continue;
      const d2 = ox * ox + oy * oy;
      if (d2 > r * r) continue;
      h[yi * GW + xi] += amp * Math.exp(-d2 * 0.15);
    }
  }
}

function init({ width, height }) {
  alloc(width, height);
  // seed: a couple of drops so the first frames already have ripples
  drop(width * 0.35, height * 0.4, 0.35);
  drop(width * 0.65, height * 0.6, 0.30);
}

function step(dt) {
  // Lax-Friedrichs: u_i^{n+1} = 0.5*(u_{i+1}+u_{i-1}) - dt/(2dx)*(f(u_{i+1})-f(u_{i-1}))
  // applied dimension-split for 2D. Fluxes from conservative SWE:
  //   F = [hu, hu^2/h + 0.5 g h^2, hu*hv/h]
  //   G = [hv, hu*hv/h, hv^2/h + 0.5 g h^2]
  const dt2 = dt / (2 * DX);
  for (let y = 1; y < GH - 1; y++) {
    const row = y * GW;
    for (let x = 1; x < GW - 1; x++) {
      const i = row + x;
      const iL = i - 1, iR = i + 1, iU = i - GW, iD = i + GW;
      // fluxes east/west
      const hL = h[iL], hR = h[iR], hU = h[iU], hD = h[iD];
      const uL = hu[iL] / hL, uR = hu[iR] / hR;
      const uU = hu[iU] / hU, uD = hu[iD] / hD;
      const vL = hv[iL] / hL, vR = hv[iR] / hR;
      const vU = hv[iU] / hU, vD = hv[iD] / hD;
      const fhR = hu[iR],            fhL = hu[iL];
      const fhuR = hu[iR] * uR + 0.5 * G * hR * hR;
      const fhuL = hu[iL] * uL + 0.5 * G * hL * hL;
      const fhvR = hu[iR] * vR,      fhvL = hu[iL] * vL;
      const ghD = hv[iD],            ghU = hv[iU];
      const ghuD = hv[iD] * uD,      ghuU = hv[iU] * uU;
      const ghvD = hv[iD] * vD + 0.5 * G * hD * hD;
      const ghvU = hv[iU] * vU + 0.5 * G * hU * hU;
      // averaging from 4 neighbors keeps the scheme stable in 2D
      h2[i]  = 0.25 * (hL + hR + hU + hD) - dt2 * ((fhR - fhL) + (ghD - ghU));
      hu2[i] = 0.25 * (hu[iL] + hu[iR] + hu[iU] + hu[iD]) - dt2 * ((fhuR - fhuL) + (ghuD - ghuU));
      hv2[i] = 0.25 * (hv[iL] + hv[iR] + hv[iU] + hv[iD]) - dt2 * ((fhvR - fhvL) + (ghvD - ghvU));
      // safety: keep depth positive
      if (h2[i] < 0.05) h2[i] = 0.05;
    }
  }
  // reflective walls: copy interior, flip the normal momentum
  for (let x = 0; x < GW; x++) {
    const top = x, bot = (GH - 1) * GW + x;
    h2[top] = h2[top + GW]; hu2[top] = hu2[top + GW]; hv2[top] = -hv2[top + GW];
    h2[bot] = h2[bot - GW]; hu2[bot] = hu2[bot - GW]; hv2[bot] = -hv2[bot - GW];
  }
  for (let y = 0; y < GH; y++) {
    const lft = y * GW, rgt = y * GW + (GW - 1);
    h2[lft] = h2[lft + 1]; hu2[lft] = -hu2[lft + 1]; hv2[lft] = hv2[lft + 1];
    h2[rgt] = h2[rgt - 1]; hu2[rgt] = -hu2[rgt - 1]; hv2[rgt] = hv2[rgt - 1];
  }
  // swap
  const th = h, thu = hu, thv = hv;
  h = h2; hu = hu2; hv = hv2;
  h2 = th; hu2 = thu; hv2 = thv;
}

function shade(t) {
  // t in [-1,1] => deep blue (low) to white-cyan (high)
  const k = Math.max(-1, Math.min(1, t));
  if (k >= 0) {
    return [40 + 180 * k, 110 + 130 * k, 180 + 75 * k];
  } else {
    const a = -k;
    return [10 + 30 * (1 - a), 30 + 40 * (1 - a), 70 + 80 * (1 - a)];
  }
}

function tick({ ctx, dt, frame, width, height, input }) {
  if (width !== cw || height !== ch) alloc(width, height);

  // CFL-bounded substepping; gravity wave speed c = sqrt(g*h)
  const maxH = 1.6;
  const c = Math.sqrt(G * maxH);
  const dtMax = 0.4 * DX / c;            // CFL ~0.4
  const stepDt = Math.min(0.025, Math.max(0.005, dt));
  const sub = Math.max(1, Math.ceil(stepDt / dtMax));
  const sdt = stepDt / sub;
  for (let s = 0; s < sub; s++) step(sdt);

  // ambient: rare auto-drop if water is very flat
  if (frame % 360 === 0) {
    let e = 0;
    for (let i = 0; i < h.length; i += 31) { const d = h[i] - H0; e += d * d; }
    if (e < 0.0015) drop(Math.random() * width, Math.random() * height, 0.32);
  }

  const clicks = input.consumeClicks();
  for (const c2 of clicks) drop(c2.x, c2.y, 0.55);

  // track cursor row for cross section
  if (input.mouseY >= 0 && input.mouseY < height) cursorY = Math.floor((input.mouseY / ch) * GH);

  // paint heightmap into the small offscreen, scale up
  for (let i = 0, p = 0; i < h.length; i++, p += 4) {
    const dev = (h[i] - H0) / 0.6;
    const [r, g, b] = shade(dev);
    pix[p] = r; pix[p + 1] = g; pix[p + 2] = b;
  }
  octx.putImageData(img, 0, 0);
  ctx.imageSmoothingEnabled = true;
  ctx.drawImage(off, 0, 0, width, height);

  // side-view cross section through cursorY (or middle)
  const sy = cursorY >= 0 ? cursorY : (GH >> 1);
  const insetW = Math.min(260, Math.floor(width * 0.42));
  const insetH = 80;
  const ix = width - insetW - 12;
  const iy = 12;
  ctx.fillStyle = "rgba(8, 14, 28, 0.78)";
  ctx.fillRect(ix, iy, insetW, insetH);
  ctx.strokeStyle = "rgba(120, 170, 230, 0.55)";
  ctx.lineWidth = 1;
  ctx.strokeRect(ix + 0.5, iy + 0.5, insetW - 1, insetH - 1);
  // baseline (rest depth)
  const baseY = iy + insetH - 14;
  ctx.strokeStyle = "rgba(120, 170, 230, 0.35)";
  ctx.beginPath(); ctx.moveTo(ix + 4, baseY); ctx.lineTo(ix + insetW - 4, baseY); ctx.stroke();
  // surface curve
  ctx.strokeStyle = "rgba(180, 220, 255, 0.95)";
  ctx.beginPath();
  for (let x = 0; x < GW; x++) {
    const sxp = ix + 4 + (insetW - 8) * (x / (GW - 1));
    const val = h[sy * GW + x] - H0;
    const syp = baseY - val * 60;
    if (x === 0) ctx.moveTo(sxp, syp); else ctx.lineTo(sxp, syp);
  }
  ctx.stroke();
  // cursor marker line on the heightmap
  if (cursorY >= 0) {
    ctx.strokeStyle = "rgba(255, 240, 200, 0.35)";
    const yy = (sy + 0.5) / GH * height;
    ctx.beginPath(); ctx.moveTo(0, yy); ctx.lineTo(width, yy); ctx.stroke();
  }

  // HUD
  ctx.fillStyle = "rgba(8, 14, 28, 0.7)";
  ctx.fillRect(8, height - 28, 320, 20);
  ctx.fillStyle = "rgba(200, 220, 255, 0.95)";
  ctx.font = "12px system-ui, sans-serif";
  ctx.fillText("click to drop a stone — waves reflect, interfere", 16, height - 14);
  ctx.fillStyle = "rgba(200, 220, 255, 0.85)";
  ctx.fillText("cross section", ix + 8, iy + 14);
}