SIR Epidemic on a Grid

const GW = 200, GH = 150;
const S_ = 0, I_ = 1, R_ = 2;
const beta = 0.28, gamma = 0.04;
const R0 = (beta * 4 / gamma).toFixed(2);
let grid, next, history, peak, step, running, restartAt;
let off, octx, img;

function reset() {
  grid = new Uint8Array(GW * GH);
  next = new Uint8Array(GW * GH);
  history = [];
  peak = 0;
  step = 0;
  running = true;
  restartAt = 0;
  const cx = (GW / 2) | 0, cy = (GH / 2) | 0;
  for (let dy = -2; dy <= 2; dy++) for (let dx = -2; dx <= 2; dx++) {
    grid[(cy + dy) * GW + (cx + dx)] = I_;
  }
}

function init() {
  off = new OffscreenCanvas(GW, GH);
  octx = off.getContext("2d");
  img = octx.createImageData(GW, GH);
  reset();
}

function vaccinate(gx, gy) {
  const rad = 8;
  for (let dy = -rad; dy <= rad; dy++) for (let dx = -rad; dx <= rad; dx++) {
    if (dx * dx + dy * dy > rad * rad) continue;
    const x = gx + dx, y = gy + dy;
    if (x < 0 || x >= GW || y < 0 || y >= GH) continue;
    const k = y * GW + x;
    if (grid[k] === S_) grid[k] = R_;
  }
}

function stepCA() {
  let iCount = 0;
  for (let y = 0; y < GH; y++) {
    for (let x = 0; x < GW; x++) {
      const k = y * GW + x;
      const c = grid[k];
      if (c === R_) { next[k] = R_; continue; }
      if (c === I_) { next[k] = Math.random() < gamma ? R_ : I_; continue; }
      let infected = false;
      if (x > 0 && grid[k - 1] === I_ && Math.random() < beta) infected = true;
      else if (x < GW - 1 && grid[k + 1] === I_ && Math.random() < beta) infected = true;
      else if (y > 0 && grid[k - GW] === I_ && Math.random() < beta) infected = true;
      else if (y < GH - 1 && grid[k + GW] === I_ && Math.random() < beta) infected = true;
      next[k] = infected ? I_ : S_;
    }
  }
  const tmp = grid; grid = next; next = tmp;
  let sC = 0, rC = 0;
  for (let i = 0; i < grid.length; i++) {
    const v = grid[i];
    if (v === S_) sC++; else if (v === I_) iCount++; else rC++;
  }
  history.push([sC, iCount, rC]);
  if (history.length > 800) history.shift();
  if (iCount > peak) peak = iCount;
  step++;
  if (iCount === 0 && step > 20) running = false;
}

function tick({ ctx, dt, width: W, height: H, input }) {
  const chartH = Math.min(80, H * 0.18);
  const gridPxH = H - chartH;
  const cellW = W / GW;
  const cellH = gridPxH / GH;

  for (const c of input.consumeClicks()) {
    if (c.y < gridPxH) {
      const gx = (c.x / cellW) | 0, gy = (c.y / cellH) | 0;
      if (gx >= 0 && gx < GW && gy >= 0 && gy < GH) vaccinate(gx, gy);
    }
  }

  if (running) stepCA();
  else {
    if (!restartAt) restartAt = performance.now() + 1500;
    else if (performance.now() > restartAt) reset();
  }

  const d = img.data;
  for (let i = 0; i < grid.length; i++) {
    const v = grid[i], o = i * 4;
    if (v === S_) { d[o] = 40; d[o + 1] = 90; d[o + 2] = 200; }
    else if (v === I_) { d[o] = 230; d[o + 1] = 50; d[o + 2] = 50; }
    else { d[o] = 110; d[o + 1] = 110; d[o + 2] = 115; }
    d[o + 3] = 255;
  }
  octx.putImageData(img, 0, 0);
  ctx.imageSmoothingEnabled = false;
  ctx.drawImage(off, 0, 0, W, gridPxH);

  ctx.fillStyle = "#111";
  ctx.fillRect(0, gridPxH, W, chartH);
  const total = GW * GH;
  const n = history.length;
  if (n > 1) {
    const xs = W / Math.max(n, 1);
    ctx.beginPath();
    ctx.moveTo(0, gridPxH + chartH);
    for (let i = 0; i < n; i++) {
      const h = (history[i][2] / total) * chartH;
      ctx.lineTo(i * xs, gridPxH + chartH - h);
    }
    ctx.lineTo((n - 1) * xs, gridPxH + chartH);
    ctx.closePath();
    ctx.fillStyle = "#6e6e73";
    ctx.fill();

    ctx.beginPath();
    ctx.moveTo(0, gridPxH + chartH);
    for (let i = 0; i < n; i++) {
      const rH = (history[i][2] / total) * chartH;
      const iH = (history[i][1] / total) * chartH;
      ctx.lineTo(i * xs, gridPxH + chartH - rH - iH);
    }
    for (let i = n - 1; i >= 0; i--) {
      const rH = (history[i][2] / total) * chartH;
      ctx.lineTo(i * xs, gridPxH + chartH - rH);
    }
    ctx.closePath();
    ctx.fillStyle = "#e63232";
    ctx.fill();
  }

  ctx.fillStyle = "rgba(0,0,0,0.55)";
  ctx.fillRect(8, 8, 230, 60);
  ctx.fillStyle = "#fff";
  ctx.font = "12px monospace";
  ctx.fillText(`R0 ≈ ${R0}  (beta=${beta}, gamma=${gamma})`, 16, 24);
  ctx.fillText(`Peak I: ${peak}  (${((peak / total) * 100).toFixed(1)}%)`, 16, 40);
  ctx.fillText(`Step: ${step}  ·  click to vaccinate`, 16, 58);
}