Gradient Descent on Himmelblau

// Himmelblau: f(x,y) = (x^2 + y - 11)^2 + (x + y^2 - 7)^2
// Four minima, all with f = 0:
//   ( 3.000,  2.000), (-2.805,  3.131),
//   (-3.779, -3.283), ( 3.584, -1.848)
const XMIN = -6, XMAX = 6, YMIN = -6, YMAX = 6;
const MAX_AGENTS = 8;
const HUE_BY_MIN = [10, 140, 215, 50]; // orange, green, blue, yellow
const MINS = [[3, 2], [-2.805, 3.131], [-3.779, -3.283], [3.584, -1.848]];
let W = 0, H = 0;
let surfaceBuf, surfaceCtx, surfaceDirty = true;
let agents = [];
let lr = 0.005;
let mom = 0.85;

function init({ width, height }) {
  W = width; H = height;
  surfaceBuf = new OffscreenCanvas(W, H);
  surfaceCtx = surfaceBuf.getContext("2d");
  agents = [];
  surfaceDirty = true;
  // Drop 3 agents at random positions so trajectories are visible on first paint.
  for (let i = 0; i < 3; i++) {
    const wx = XMIN + Math.random() * (XMAX - XMIN);
    const wy = YMIN + Math.random() * (YMAX - YMIN);
    agents.push({ x: wx, y: wy, vx: 0, vy: 0, trail: [wx, wy], tick: 0, bin: nearestMinIdx(wx, wy) });
  }
}

function f(x, y) {
  const a = x * x + y - 11;
  const b = x + y * y - 7;
  return a * a + b * b;
}

function grad(x, y) {
  const a = x * x + y - 11;
  const b = x + y * y - 7;
  return [4 * x * a + 2 * b, 2 * a + 4 * y * b];
}

function screenToWorld(px, py) {
  return [XMIN + (px / W) * (XMAX - XMIN), YMAX - (py / H) * (YMAX - YMIN)];
}
function worldToScreen(x, y) {
  return [((x - XMIN) / (XMAX - XMIN)) * W, ((YMAX - y) / (YMAX - YMIN)) * H];
}

function nearestMinIdx(x, y) {
  let best = 0, bd = Infinity;
  for (let i = 0; i < 4; i++) {
    const dx = x - MINS[i][0], dy = y - MINS[i][1];
    const d = dx * dx + dy * dy;
    if (d < bd) { bd = d; best = i; }
  }
  return best;
}

function renderSurface() {
  // Render at the visible canvas resolution. The previous code computed at
  // 180×N and let drawImage upscale with bilinear smoothing, which produced
  // visibly blurry contour bands. Himmelblau is two squared polynomials —
  // cheap enough to evaluate per-pixel at full res, even on mobile.
  const img = surfaceCtx.createImageData(W, H);
  const d = img.data;
  const LMAX = Math.log(1 + 800);
  for (let j = 0; j < H; j++) {
    const y = YMAX - (j / H) * (YMAX - YMIN);
    for (let i = 0; i < W; i++) {
      const x = XMIN + (i / W) * (XMAX - XMIN);
      const t = Math.min(1, Math.log(1 + f(x, y)) / LMAX);
      const band = Math.abs((t * 12) % 1 - 0.5) < 0.05;
      const o = (j * W + i) * 4;
      d[o]     = band ? 235 : (40 + 200 * t) | 0;
      d[o + 1] = band ? 220 : (50 + 60 * (1 - t)) | 0;
      d[o + 2] = band ? 180 : (110 - 60 * t) | 0;
      d[o + 3] = 255;
    }
  }
  surfaceCtx.putImageData(img, 0, 0);
  surfaceDirty = false;
}

function drawMinima(ctx) {
  for (let i = 0; i < 4; i++) {
    const [sx, sy] = worldToScreen(MINS[i][0], MINS[i][1]);
    ctx.strokeStyle = `hsla(${HUE_BY_MIN[i]},80%,70%,0.85)`;
    ctx.lineWidth = 2;
    ctx.beginPath(); ctx.arc(sx, sy, 7, 0, Math.PI * 2); ctx.stroke();
    ctx.beginPath();
    ctx.moveTo(sx - 3, sy); ctx.lineTo(sx + 3, sy);
    ctx.moveTo(sx, sy - 3); ctx.lineTo(sx, sy + 3);
    ctx.stroke();
  }
}

function drawAgents(ctx) {
  for (const a of agents) {
    // Trail.
    ctx.strokeStyle = `hsla(${HUE_BY_MIN[a.bin]},90%,70%,0.85)`;
    ctx.lineWidth = 1.6;
    ctx.beginPath();
    for (let i = 0; i < a.trail.length; i += 2) {
      const [sx, sy] = worldToScreen(a.trail[i], a.trail[i + 1]);
      if (i === 0) ctx.moveTo(sx, sy); else ctx.lineTo(sx, sy);
    }
    const [hx, hy] = worldToScreen(a.x, a.y);
    ctx.lineTo(hx, hy);
    ctx.stroke();
    // Head.
    ctx.fillStyle = `hsl(${HUE_BY_MIN[a.bin]},95%,75%)`;
    ctx.beginPath(); ctx.arc(hx, hy, 4, 0, Math.PI * 2); ctx.fill();
    ctx.strokeStyle = "rgba(0,0,0,0.6)";
    ctx.lineWidth = 1;
    ctx.stroke();
  }
}

// On-canvas tap targets so mobile users can adjust lr/mom (keyboard-only
// otherwise). Recomputed each frame; consulted by hit-test in `tick`.
const HUD_BTN_SIZE = 22;
let hudButtons = [];

function drawHUD(ctx) {
  const pad = 10;
  ctx.fillStyle = "rgba(0,0,0,0.6)";
  ctx.fillRect(pad, pad, 230, 110);
  ctx.fillStyle = "#fff";
  ctx.font = "13px monospace";
  ctx.textAlign = "left";
  ctx.textBaseline = "alphabetic";
  ctx.fillText(`lr   = ${lr.toFixed(4)}`, pad + 8, pad + 20);
  ctx.fillText(`mom  = ${mom.toFixed(3)}`, pad + 8, pad + 44);
  ctx.fillText(`agents = ${agents.length}/${MAX_AGENTS}`, pad + 8, pad + 68);
  ctx.fillText(`tap below to drop`, pad + 8, pad + 92);

  // Buttons: − and + for lr and mom. Position to the right of each readout.
  hudButtons.length = 0;
  const bs = HUD_BTN_SIZE;
  const lrY = pad + 8;
  const momY = pad + 32;
  const lrMinusX = pad + 150;
  const lrPlusX  = lrMinusX + bs + 4;
  const momMinusX = pad + 150;
  const momPlusX  = momMinusX + bs + 4;
  hudButtons.push({ x: lrMinusX, y: lrY, w: bs, h: bs, kind: 'lr-' });
  hudButtons.push({ x: lrPlusX,  y: lrY, w: bs, h: bs, kind: 'lr+' });
  hudButtons.push({ x: momMinusX, y: momY, w: bs, h: bs, kind: 'mom-' });
  hudButtons.push({ x: momPlusX,  y: momY, w: bs, h: bs, kind: 'mom+' });
  ctx.font = "16px monospace";
  ctx.textAlign = "center";
  ctx.textBaseline = "middle";
  for (const b of hudButtons) {
    ctx.fillStyle = "rgba(255,255,255,0.15)";
    ctx.fillRect(b.x, b.y, b.w, b.h);
    ctx.strokeStyle = "rgba(255,255,255,0.55)";
    ctx.lineWidth = 1;
    ctx.strokeRect(b.x + 0.5, b.y + 0.5, b.w - 1, b.h - 1);
    ctx.fillStyle = "#fff";
    ctx.fillText(b.kind.endsWith('+') ? '+' : '−', b.x + b.w / 2, b.y + b.h / 2 + 1);
  }
  ctx.textAlign = "left";
  ctx.textBaseline = "alphabetic";

  // Legend (which color = which minimum), bottom-left.
  const ly = H - pad - 18;
  ctx.fillStyle = "rgba(0,0,0,0.55)";
  ctx.fillRect(pad, ly, 230, 18);
  ctx.fillStyle = "#fff";
  ctx.font = "11px monospace";
  ctx.fillText("4 minima:", pad + 6, ly + 13);
  for (let i = 0; i < 4; i++) {
    const cx = pad + 78 + i * 36;
    ctx.fillStyle = `hsl(${HUE_BY_MIN[i]},90%,65%)`;
    ctx.beginPath(); ctx.arc(cx, ly + 9, 5, 0, Math.PI * 2); ctx.fill();
  }
}

function stepAgent(a, h) {
  // Adaptive sub-stepping: gradient can be enormous far from minima.
  // Threshold is intentionally permissive so most frames take 1 step —
  // we want the trajectory to be slow and legible.
  let steps = 1;
  if (h > 0) {
    const [gx0, gy0] = grad(a.x, a.y);
    const gn = Math.hypot(gx0, gy0);
    if (gn * lr > 0.8) steps = Math.min(4, 1 + ((gn * lr) / 0.8) | 0);
  }
  for (let s = 0; s < steps; s++) {
    const [gx, gy] = grad(a.x, a.y);
    a.vx = mom * a.vx - lr * gx;
    a.vy = mom * a.vy - lr * gy;
    // Clamp velocity so big slopes don't fling agents off the map.
    const vn = Math.hypot(a.vx, a.vy);
    const VMAX = 1.5;
    if (vn > VMAX) { a.vx *= VMAX / vn; a.vy *= VMAX / vn; }
    a.x += a.vx; a.y += a.vy;
    if (a.x < XMIN || a.x > XMAX || a.y < YMIN || a.y > YMAX) {
      a.x = Math.max(XMIN, Math.min(XMAX, a.x));
      a.y = Math.max(YMIN, Math.min(YMAX, a.y));
      a.vx *= -0.3; a.vy *= -0.3;
    }
  }
  // Trail: append every few frames, bounded length.
  a.tick++;
  if ((a.tick & 1) === 0) {
    a.trail.push(a.x, a.y);
    if (a.trail.length > 800) a.trail.splice(0, a.trail.length - 800);
  }
  a.bin = nearestMinIdx(a.x, a.y);
}

function tick({ ctx, dt, width, height, input }) {
  if (width !== W || height !== H) { W = width; H = height; surfaceDirty = true; }
  if (surfaceDirty) renderSurface();

  if (input.justPressed("[")) lr = Math.max(0.0005, lr / 1.15);
  if (input.justPressed("]")) lr = Math.min(0.08, lr * 1.15);
  if (input.justPressed(",")) mom = Math.max(0, mom - 0.05);
  if (input.justPressed(".")) mom = Math.min(0.98, mom + 0.05);

  for (const c of input.consumeClicks()) {
    // Hit-test HUD buttons first so mobile users can tune lr/mom by tap.
    let consumed = false;
    for (const b of hudButtons) {
      if (c.x >= b.x && c.x <= b.x + b.w && c.y >= b.y && c.y <= b.y + b.h) {
        if (b.kind === 'lr-') lr = Math.max(0.0005, lr / 1.15);
        else if (b.kind === 'lr+') lr = Math.min(0.08, lr * 1.15);
        else if (b.kind === 'mom-') mom = Math.max(0, mom - 0.05);
        else if (b.kind === 'mom+') mom = Math.min(0.98, mom + 0.05);
        consumed = true;
        break;
      }
    }
    if (consumed) continue;
    const [wx, wy] = screenToWorld(c.x, c.y);
    agents.push({ x: wx, y: wy, vx: 0, vy: 0, trail: [wx, wy], tick: 0, bin: nearestMinIdx(wx, wy) });
    while (agents.length > MAX_AGENTS) agents.shift();
  }

  for (const a of agents) stepAgent(a, dt);

  ctx.drawImage(surfaceBuf, 0, 0);
  drawMinima(ctx);
  drawAgents(ctx);
  drawHUD(ctx);
}