Lattice Boltzmann (D2Q9)

// D2Q9 lattice Boltzmann BGK past a cylinder. Each cell carries 9
// populations f_i on the lattice velocities c_i. A step is:
//   collide:  f_i += (f_i^eq - f_i) / tau   (relax to local equilibrium)
//   stream:   f_i(x + c_i, t+1) <- f_i(x, t)
// Cylinder: half-way bounce-back. Inlet: Zou/He with u = (U, 0). Outlet:
// zero-gradient. nu = (tau - 1/2)/3, Re = U D / nu. |u| -> viridis ramp.

const GW = 200, GH = 70;
const N = GW * GH;
const TAU = 0.56;                  // relaxation time (nu = (tau-0.5)/3)
const U0 = 0.10;                   // inlet velocity (lattice units)
const STEPS_PER_FRAME = 2;

// D2Q9 lattice. Index 0 = rest, 1..4 = axes E,N,W,S, 5..8 = diagonals.
const CX = [0, 1, 0, -1, 0,  1, -1, -1,  1];
const CY = [0, 0, -1, 0, 1,  -1, -1,  1,  1];
const W  = [4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36];
const OPP = [0, 3, 4, 1, 2, 7, 8, 5, 6];   // opposite direction (for bounce-back)

let f, fTmp;                       // population arrays: 9 * N
let solid;                         // Uint8Array N: 1 if obstacle
let rho, ux, uy;                   // macroscopic moments (per cell)
let off, octx, img, pix;
let cyl;                           // { x, y, r } in lattice cells

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

function rebuildSolid() {
  solid.fill(0);
  const r2 = cyl.r * cyl.r;
  const x0 = Math.max(0, Math.floor(cyl.x - cyl.r - 1));
  const x1 = Math.min(GW - 1, Math.ceil(cyl.x + cyl.r + 1));
  const y0 = Math.max(0, Math.floor(cyl.y - cyl.r - 1));
  const y1 = Math.min(GH - 1, Math.ceil(cyl.y + cyl.r + 1));
  for (let y = y0; y <= y1; y++) {
    for (let x = x0; x <= x1; x++) {
      const dx = x - cyl.x, dy = y - cyl.y;
      if (dx * dx + dy * dy <= r2) solid[idx(x, y)] = 1;
    }
  }
}

function equilibrium(rh, vx, vy, out, base) {
  const u2 = 1 - 1.5 * (vx * vx + vy * vy);
  for (let i = 0; i < 9; i++) {
    const cu = 3 * (CX[i] * vx + CY[i] * vy);
    out[base + i] = W[i] * rh * (u2 + cu + 0.5 * cu * cu);
  }
}

function init({ width, height }) {
  f = new Float32Array(9 * N);
  fTmp = new Float32Array(9 * N);
  solid = new Uint8Array(N);
  rho = new Float32Array(N);
  ux = new Float32Array(N);
  uy = new Float32Array(N);

  cyl = { x: GW * 0.22, y: GH * 0.5, r: 6.5 };
  rebuildSolid();

  // Seed cells with equilibrium for rho=1, u=(U0,0) plus a tiny vertical
  // perturbation that breaks symmetry so the wake sheds without delay.
  for (let y = 0; y < GH; y++) {
    for (let x = 0; x < GW; x++) {
      const i = idx(x, y);
      const wig = 0.01 * Math.sin((y / GH) * Math.PI * 4);
      equilibrium(1, U0 + wig, 0, f, i * 9);
      rho[i] = 1; ux[i] = U0; uy[i] = 0;
    }
  }

  off = new OffscreenCanvas(GW, GH);
  octx = off.getContext("2d");
  img = octx.createImageData(GW, GH);
  pix = img.data;
  for (let q = 3; q < pix.length; q += 4) pix[q] = 255;
}

function step() {
  // --- collide (in place) + accumulate moments ---
  for (let i = 0; i < N; i++) {
    if (solid[i]) continue;
    const b = i * 9;
    let r = 0, mx = 0, my = 0;
    for (let k = 0; k < 9; k++) {
      const fk = f[b + k];
      r += fk; mx += CX[k] * fk; my += CY[k] * fk;
    }
    const vx = mx / r, vy = my / r;
    rho[i] = r; ux[i] = vx; uy[i] = vy;
    const u2 = 1 - 1.5 * (vx * vx + vy * vy);
    const invTau = 1 / TAU;
    for (let k = 0; k < 9; k++) {
      const cu = 3 * (CX[k] * vx + CY[k] * vy);
      const feq = W[k] * r * (u2 + cu + 0.5 * cu * cu);
      f[b + k] += (feq - f[b + k]) * invTau;
    }
  }

  // Stream into fTmp; obstacle neighbours reflect (half-way bounce-back).
  // Seed fTmp with f so any slot we never explicitly write (off-grid edges,
  // handled below by BCs) starts from a sane value rather than zero.
  fTmp.set(f);
  for (let y = 0; y < GH; y++) {
    for (let x = 0; x < GW; x++) {
      const i = idx(x, y);
      if (solid[i]) continue;
      const b = i * 9;
      for (let k = 0; k < 9; k++) {
        const nx = x + CX[k], ny = y + CY[k];
        if (nx < 0 || nx >= GW || ny < 0 || ny >= GH) continue;
        const ni = idx(nx, ny);
        if (solid[ni]) fTmp[b + OPP[k]] = f[b + k];
        else           fTmp[ni * 9 + k] = f[b + k];
      }
    }
  }

  // swap
  const t = f; f = fTmp; fTmp = t;

  // Inlet (x=0): Zou/He velocity BC, u = (U0, 0). Unknowns moving into the
  // domain are k = 1, 5, 8 (CX > 0); the rest are known after streaming.
  for (let y = 1; y < GH - 1; y++) {
    const b = idx(0, y) * 9;
    const r = (f[b+0] + f[b+2] + f[b+4] + 2 * (f[b+3] + f[b+6] + f[b+7])) / (1 - U0);
    f[b+1] = f[b+3] + (2/3) * r * U0;
    f[b+5] = f[b+7] + 0.5 * (f[b+4] - f[b+2]) + (1/6) * r * U0;
    f[b+8] = f[b+6] + 0.5 * (f[b+2] - f[b+4]) + (1/6) * r * U0;
  }
  // Outlet: zero-gradient (copy neighbour). Top/bottom: copy inner row.
  for (let y = 0; y < GH; y++) {
    const i = idx(GW - 1, y) * 9, j = idx(GW - 2, y) * 9;
    for (let k = 0; k < 9; k++) f[i + k] = f[j + k];
  }
  for (let x = 0; x < GW; x++) {
    const t0 = idx(x, 0) * 9, t1 = idx(x, 1) * 9;
    const b0 = idx(x, GH - 1) * 9, b1 = idx(x, GH - 2) * 9;
    for (let k = 0; k < 9; k++) { f[t0 + k] = f[t1 + k]; f[b0 + k] = f[b1 + k]; }
  }
}

// viridis-ish ramp (5 control colours, linear interp)
const RAMP = [
  [ 68,  1, 84], [ 59, 82,139], [ 33,144,141],
  [ 94,201, 98], [253,231, 37],
];
function colour(t, out, j) {
  if (t < 0) t = 0; else if (t > 1) t = 1;
  const s = t * (RAMP.length - 1);
  const k = s | 0, fr = s - k;
  const a = RAMP[k], c = RAMP[Math.min(RAMP.length - 1, k + 1)];
  out[j]     = a[0] + (c[0] - a[0]) * fr;
  out[j + 1] = a[1] + (c[1] - a[1]) * fr;
  out[j + 2] = a[2] + (c[2] - a[2]) * fr;
}

function tick({ ctx, width, height, input }) {
  if (input.mouseDown) {
    const gx = (input.mouseX / width) * GW;
    const gy = (input.mouseY / height) * GH;
    cyl.x = Math.max(cyl.r + 2, Math.min(GW * 0.75, gx));
    cyl.y = Math.max(cyl.r + 2, Math.min(GH - cyl.r - 2, gy));
    rebuildSolid();
  }
  for (let s = 0; s < STEPS_PER_FRAME; s++) step();

  // |u| -> viridis colormap; obstacle = dark grey
  const SCALE = 1 / (U0 * 1.6);
  for (let i = 0, q = 0; i < N; i++, q += 4) {
    if (solid[i]) { pix[q] = 28; pix[q+1] = 30; pix[q+2] = 38; continue; }
    const m = Math.sqrt(ux[i] * ux[i] + uy[i] * uy[i]);
    colour(m * SCALE, pix, q);
  }
  octx.putImageData(img, 0, 0);
  ctx.imageSmoothingEnabled = true;
  ctx.drawImage(off, 0, 0, width, height);

  // crisp cylinder outline in canvas space
  const px = (cyl.x / GW) * width, py = (cyl.y / GH) * height;
  const pr = (cyl.r / GW) * width;
  ctx.strokeStyle = "rgba(230,232,240,0.9)";
  ctx.lineWidth = 1.2;
  ctx.beginPath(); ctx.arc(px, py, pr, 0, 6.2831853); ctx.stroke();

  // HUD
  const nu = (TAU - 0.5) / 3, Re = (U0 * (2 * cyl.r)) / nu;
  ctx.fillStyle = "rgba(12,14,22,0.65)";
  ctx.fillRect(8, 8, 220, 38);
  ctx.fillStyle = "rgba(230,232,240,0.95)";
  ctx.font = "12px system-ui, sans-serif";
  ctx.fillText(`D2Q9 BGK  tau=${TAU.toFixed(2)}  U=${U0.toFixed(2)}`, 14, 22);
  ctx.fillText(`Re = U D / nu = ${Re.toFixed(0)}`, 14, 38);
  ctx.fillStyle = "rgba(200,205,215,0.85)";
  ctx.fillText("drag the obstacle", 10, height - 10);
}