Linear Regression: OLS & Residuals

// OLS linear regression with residuals. Click to add a point; the fit and
// residual sum update live. Shows m, b, R^2, SSR.

let W = 0, H = 0;
let pts = []; // {x, y} in data coords
let dragIdx = -1;
let initialized = false;
const X_MIN = 0, X_MAX = 10;
const Y_MIN = 0, Y_MAX = 10;
let pulseAge = 0;

function seedPoints() {
  pts = [];
  const trueM = 0.65, trueB = 1.4;
  for (let i = 0; i < 30; i++) {
    const x = 0.5 + Math.random() * 9;
    // gaussian noise via Box-Muller
    const u1 = Math.random(), u2 = Math.random();
    const z = Math.sqrt(-2 * Math.log(Math.max(1e-9, u1))) * Math.cos(2 * Math.PI * u2);
    const y = trueM * x + trueB + 0.9 * z;
    pts.push({ x, y: Math.max(Y_MIN + 0.1, Math.min(Y_MAX - 0.1, y)) });
  }
}

function init({ width, height }) {
  W = width; H = height;
  if (!initialized) { seedPoints(); initialized = true; }
}

function fit() {
  const n = pts.length;
  if (n < 2) return { m: 0, b: 0, ssr: 0, sst: 0, r2: 0, xbar: 0, ybar: 0 };
  let sx = 0, sy = 0;
  for (const p of pts) { sx += p.x; sy += p.y; }
  const xbar = sx / n, ybar = sy / n;
  let sxx = 0, sxy = 0, syy = 0;
  for (const p of pts) {
    const dx = p.x - xbar, dy = p.y - ybar;
    sxx += dx * dx; sxy += dx * dy; syy += dy * dy;
  }
  const m = sxx > 0 ? sxy / sxx : 0;
  const b = ybar - m * xbar;
  let ssr = 0;
  for (const p of pts) {
    const r = p.y - (m * p.x + b);
    ssr += r * r;
  }
  const sst = syy;
  const r2 = sst > 0 ? 1 - ssr / sst : 0;
  return { m, b, ssr, sst, r2, xbar, ybar };
}

function plotRect() {
  const pad = 50;
  const padR = 30;
  const padTop = 60;
  const padBot = 80;
  return { x0: pad, y0: padTop, x1: W - padR, y1: H - padBot };
}

function dataToPx(x, y, r) {
  const px = r.x0 + ((x - X_MIN) / (X_MAX - X_MIN)) * (r.x1 - r.x0);
  const py = r.y1 - ((y - Y_MIN) / (Y_MAX - Y_MIN)) * (r.y1 - r.y0);
  return [px, py];
}

function pxToData(px, py, r) {
  const x = X_MIN + ((px - r.x0) / (r.x1 - r.x0)) * (X_MAX - X_MIN);
  const y = Y_MIN + ((r.y1 - py) / (r.y1 - r.y0)) * (Y_MAX - Y_MIN);
  return [x, y];
}

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

  if (input.justPressed("r") || input.justPressed("R")) {
    seedPoints();
    dragIdx = -1;
  }
  if (input.justPressed("c") || input.justPressed("C")) {
    pts = [];
    dragIdx = -1;
  }

  const r = plotRect();
  const mx = input.mouseX, my = input.mouseY;
  const inPlot = mx >= r.x0 && mx <= r.x1 && my >= r.y0 && my <= r.y1;

  // dragging: hold mouse on a nearby point, drag it
  if (input.mouseDown && inPlot) {
    if (dragIdx < 0) {
      let best = -1, bestD = 14 * 14;
      for (let i = 0; i < pts.length; i++) {
        const [px, py] = dataToPx(pts[i].x, pts[i].y, r);
        const d = (px - mx) * (px - mx) + (py - my) * (py - my);
        if (d < bestD) { bestD = d; best = i; }
      }
      if (best >= 0) dragIdx = best;
    }
    if (dragIdx >= 0) {
      const [dx, dy] = pxToData(mx, my, r);
      pts[dragIdx] = { x: dx, y: dy };
    }
  } else {
    dragIdx = -1;
  }

  // click (not on existing point) adds a new point
  const clicks = input.consumeClicks ? input.consumeClicks() : 0;
  if (clicks > 0 && inPlot) {
    // only add if not clicking near an existing point (so drag doesn't double-add)
    let near = false;
    for (const p of pts) {
      const [px, py] = dataToPx(p.x, p.y, r);
      if ((px - mx) * (px - mx) + (py - my) * (py - my) < 14 * 14) { near = true; break; }
    }
    if (!near) {
      const [dx, dy] = pxToData(mx, my, r);
      pts.push({ x: dx, y: dy });
      pulseAge = 0;
    }
  }
  pulseAge += dt;

  const f = fit();

  // bg
  ctx.fillStyle = "#0a0a10";
  ctx.fillRect(0, 0, W, H);

  // plot bg + grid
  ctx.fillStyle = "#13131c";
  ctx.fillRect(r.x0, r.y0, r.x1 - r.x0, r.y1 - r.y0);

  ctx.strokeStyle = "rgba(255,255,255,0.06)";
  ctx.lineWidth = 1;
  for (let gx = 0; gx <= 10; gx++) {
    const [px] = dataToPx(gx, 0, r);
    ctx.beginPath(); ctx.moveTo(px, r.y0); ctx.lineTo(px, r.y1); ctx.stroke();
  }
  for (let gy = 0; gy <= 10; gy++) {
    const [, py] = dataToPx(0, gy, r);
    ctx.beginPath(); ctx.moveTo(r.x0, py); ctx.lineTo(r.x1, py); ctx.stroke();
  }

  // axis tick labels
  ctx.fillStyle = "#667";
  ctx.font = "10px monospace";
  for (let gx = 0; gx <= 10; gx += 2) {
    const [px, py] = dataToPx(gx, 0, r);
    ctx.fillText(String(gx), px - 4, r.y1 + 14);
  }
  for (let gy = 0; gy <= 10; gy += 2) {
    const [px, py] = dataToPx(0, gy, r);
    ctx.fillText(String(gy), r.x0 - 18, py + 4);
  }

  // residual segments
  ctx.strokeStyle = "rgba(255,120,120,0.7)";
  ctx.lineWidth = 1.2;
  for (const p of pts) {
    const yhat = f.m * p.x + f.b;
    const [px, py] = dataToPx(p.x, p.y, r);
    const [, pyhat] = dataToPx(p.x, yhat, r);
    ctx.beginPath();
    ctx.moveTo(px, py);
    ctx.lineTo(px, pyhat);
    ctx.stroke();
  }

  // best-fit line clipped to plot
  ctx.strokeStyle = "rgba(120,200,255,0.95)";
  ctx.lineWidth = 2;
  const yAtXmin = f.m * X_MIN + f.b;
  const yAtXmax = f.m * X_MAX + f.b;
  const [lx0, ly0] = dataToPx(X_MIN, yAtXmin, r);
  const [lx1, ly1] = dataToPx(X_MAX, yAtXmax, r);
  ctx.save();
  ctx.beginPath();
  ctx.rect(r.x0, r.y0, r.x1 - r.x0, r.y1 - r.y0);
  ctx.clip();
  ctx.beginPath();
  ctx.moveTo(lx0, ly0); ctx.lineTo(lx1, ly1); ctx.stroke();
  ctx.restore();
  ctx.lineWidth = 1;

  // points
  for (let i = 0; i < pts.length; i++) {
    const p = pts[i];
    const [px, py] = dataToPx(p.x, p.y, r);
    const isDrag = i === dragIdx;
    ctx.fillStyle = isDrag ? "rgba(255,220,120,1)" : "rgba(230,230,240,0.95)";
    ctx.beginPath();
    ctx.arc(px, py, isDrag ? 5 : 3.5, 0, Math.PI * 2);
    ctx.fill();
  }

  // pulse on last-added point
  if (pulseAge < 0.6 && pts.length > 0) {
    const p = pts[pts.length - 1];
    const [px, py] = dataToPx(p.x, p.y, r);
    const rad = 5 + pulseAge * 30;
    ctx.strokeStyle = `rgba(120,255,150,${(1 - pulseAge / 0.6).toFixed(3)})`;
    ctx.lineWidth = 2;
    ctx.beginPath();
    ctx.arc(px, py, rad, 0, Math.PI * 2);
    ctx.stroke();
    ctx.lineWidth = 1;
  }

  // header
  ctx.fillStyle = "#e8e8f0";
  ctx.font = "bold 16px monospace";
  ctx.fillText("Linear Regression — OLS with Residuals", r.x0, 28);
  ctx.font = "11px monospace";
  ctx.fillStyle = "#aab";
  ctx.fillText("ŷ = m·x + b   ·   minimize Σ(y − ŷ)²", r.x0, 46);

  // readouts
  ctx.font = "13px monospace";
  const lineY = r.y1 + 32;
  ctx.fillStyle = "#9cf";
  ctx.fillText(`m = ${f.m.toFixed(4)}`, r.x0, lineY);
  ctx.fillStyle = "#9cf";
  ctx.fillText(`b = ${f.b.toFixed(4)}`, r.x0 + 160, lineY);
  ctx.fillStyle = "#fc6";
  ctx.fillText(`R² = ${f.r2.toFixed(4)}`, r.x0 + 320, lineY);
  ctx.fillStyle = "#f88";
  ctx.fillText(`SSR = ${f.ssr.toFixed(3)}`, r.x0 + 460, lineY);
  ctx.fillStyle = "#aab";
  ctx.fillText(`n = ${pts.length}`, r.x0 + 620, lineY);

  ctx.fillStyle = "#778";
  ctx.font = "10px monospace";
  ctx.fillText("click to add a point   ·   drag to move a point   ·   [C] clear   ·   [R] reseed", r.x0, H - 12);
}