OLS Linear Regression

// Interactive 2D OLS linear regression.
// Drag any point to feel the line tilt. Click empty space to add a point.
// Faint dashed line shows Huber robust fit for comparison.

let pts, W, H, dragIdx, hoverIdx;
let m, b, r2;        // OLS fit
let mH, bH;          // Huber robust fit
const HIT_R = 14;    // px hit radius for grabbing a point
const DOT_R = 6;

// Data space: x in [0,10], y in [0,15] roughly. Map to canvas with margins.
const PAD = { l: 44, r: 14, t: 14, b: 28 };

function viewport() {
  return {
    x0: PAD.l, y0: PAD.t,
    w: Math.max(40, W - PAD.l - PAD.r),
    h: Math.max(40, H - PAD.t - PAD.b),
    xMin: 0, xMax: 10, yMin: 0, yMax: 15,
  };
}
function toPix(dx, dy) {
  const v = viewport();
  return [
    v.x0 + ((dx - v.xMin) / (v.xMax - v.xMin)) * v.w,
    v.y0 + (1 - (dy - v.yMin) / (v.yMax - v.yMin)) * v.h,
  ];
}
function toData(px, py) {
  const v = viewport();
  return [
    v.xMin + ((px - v.x0) / v.w) * (v.xMax - v.xMin),
    v.yMin + (1 - (py - v.y0) / v.h) * (v.yMax - v.yMin),
  ];
}

function seed() {
  pts = [];
  // y = 0.6 x + 5 + noise, with x roughly spanning [0.5, 9.5].
  for (let i = 0; i < 15; i++) {
    const x = 0.5 + (i + Math.random() * 0.4) * 0.6;
    const noise = (Math.random() - 0.5) * 2.2;
    pts.push({ x, y: 0.6 * x + 5 + noise });
  }
  dragIdx = -1; hoverIdx = -1;
}

function init({ width, height }) {
  W = width; H = height;
  seed();
  fitOLS();
  fitHuber();
}

function fitOLS() {
  const n = pts.length;
  if (n < 2) { m = 0; b = pts.length ? pts[0].y : 0; r2 = 0; return; }
  let sx = 0, sy = 0;
  for (let i = 0; i < n; i++) { sx += pts[i].x; sy += pts[i].y; }
  const mx = sx / n, my = sy / n;
  let sxx = 0, sxy = 0, syy = 0;
  for (let i = 0; i < n; i++) {
    const dx = pts[i].x - mx, dy = pts[i].y - my;
    sxx += dx * dx; sxy += dx * dy; syy += dy * dy;
  }
  m = sxx > 1e-9 ? sxy / sxx : 0;
  b = my - m * mx;
  // R^2 = 1 - SSE/SST
  let sse = 0;
  for (let i = 0; i < n; i++) {
    const e = pts[i].y - (m * pts[i].x + b);
    sse += e * e;
  }
  r2 = syy > 1e-9 ? 1 - sse / syy : 1;
}

// IRLS Huber with k = 1.345 * sigma_MAD. ~10 iterations is plenty for 15 pts.
function fitHuber() {
  const n = pts.length;
  if (n < 2) { mH = m; bH = b; return; }
  let mm = m, bb = b;
  for (let it = 0; it < 12; it++) {
    // residuals
    const res = new Array(n);
    for (let i = 0; i < n; i++) res[i] = pts[i].y - (mm * pts[i].x + bb);
    // MAD
    const sorted = res.map(Math.abs).sort((a, c) => a - c);
    const mad = sorted[Math.floor(n / 2)] || 1e-6;
    const sigma = 1.4826 * mad + 1e-9;
    const k = 1.345 * sigma;
    // weights w_i = 1 if |r|<=k else k/|r|
    let sw = 0, swx = 0, swy = 0, swxx = 0, swxy = 0;
    for (let i = 0; i < n; i++) {
      const ar = Math.abs(res[i]);
      const w = ar <= k ? 1 : k / Math.max(ar, 1e-9);
      sw += w; swx += w * pts[i].x; swy += w * pts[i].y;
      swxx += w * pts[i].x * pts[i].x; swxy += w * pts[i].x * pts[i].y;
    }
    const det = sw * swxx - swx * swx;
    if (Math.abs(det) < 1e-9) break;
    const newM = (sw * swxy - swx * swy) / det;
    const newB = (swxx * swy - swx * swxy) / det;
    if (Math.abs(newM - mm) + Math.abs(newB - bb) < 1e-5) { mm = newM; bb = newB; break; }
    mm = newM; bb = newB;
  }
  mH = mm; bH = bb;
}

function nearestPoint(px, py) {
  let best = -1, bestD = HIT_R * HIT_R;
  for (let i = 0; i < pts.length; i++) {
    const [qx, qy] = toPix(pts[i].x, pts[i].y);
    const dd = (qx - px) * (qx - px) + (qy - py) * (qy - py);
    if (dd < bestD) { bestD = dd; best = i; }
  }
  return best;
}

function drawAxes(ctx) {
  const v = viewport();
  ctx.fillStyle = "#0b0e16";
  ctx.fillRect(0, 0, W, H);
  // gridlines
  ctx.strokeStyle = "rgba(120,140,180,0.10)";
  ctx.lineWidth = 1;
  ctx.beginPath();
  for (let gx = 0; gx <= 10; gx += 2) {
    const [x] = toPix(gx, 0);
    ctx.moveTo(x, v.y0); ctx.lineTo(x, v.y0 + v.h);
  }
  for (let gy = 0; gy <= 15; gy += 3) {
    const [, y] = toPix(0, gy);
    ctx.moveTo(v.x0, y); ctx.lineTo(v.x0 + v.w, y);
  }
  ctx.stroke();
  // axes
  ctx.strokeStyle = "rgba(180,200,230,0.35)";
  ctx.beginPath();
  ctx.moveTo(v.x0, v.y0 + v.h); ctx.lineTo(v.x0 + v.w, v.y0 + v.h);
  ctx.moveTo(v.x0, v.y0); ctx.lineTo(v.x0, v.y0 + v.h);
  ctx.stroke();
  // tick labels
  ctx.fillStyle = "rgba(180,200,230,0.55)";
  ctx.font = "10px monospace";
  ctx.textBaseline = "top";
  ctx.textAlign = "center";
  for (let gx = 0; gx <= 10; gx += 2) {
    const [x] = toPix(gx, 0);
    ctx.fillText(String(gx), x, v.y0 + v.h + 4);
  }
  ctx.textAlign = "right";
  ctx.textBaseline = "middle";
  for (let gy = 0; gy <= 15; gy += 3) {
    const [, y] = toPix(0, gy);
    ctx.fillText(String(gy), v.x0 - 4, y);
  }
}

function drawLine(ctx, slope, intercept, dashed, color, lw) {
  const v = viewport();
  const y0 = slope * v.xMin + intercept;
  const y1 = slope * v.xMax + intercept;
  const [p0x, p0y] = toPix(v.xMin, y0);
  const [p1x, p1y] = toPix(v.xMax, y1);
  ctx.save();
  // clip to plot rect so steep slopes don't bleed
  ctx.beginPath();
  ctx.rect(v.x0, v.y0, v.w, v.h);
  ctx.clip();
  ctx.strokeStyle = color;
  ctx.lineWidth = lw;
  ctx.setLineDash(dashed ? [6, 5] : []);
  ctx.beginPath();
  ctx.moveTo(p0x, p0y); ctx.lineTo(p1x, p1y);
  ctx.stroke();
  ctx.restore();
}

function drawResiduals(ctx) {
  ctx.save();
  ctx.strokeStyle = "rgba(255,90,106,0.35)";
  ctx.lineWidth = 1;
  ctx.setLineDash([2, 3]);
  ctx.beginPath();
  for (let i = 0; i < pts.length; i++) {
    const [px, py] = toPix(pts[i].x, pts[i].y);
    const [, fy] = toPix(pts[i].x, m * pts[i].x + b);
    ctx.moveTo(px, py); ctx.lineTo(px, fy);
  }
  ctx.stroke();
  ctx.restore();
}

function drawPoints(ctx) {
  for (let i = 0; i < pts.length; i++) {
    const [px, py] = toPix(pts[i].x, pts[i].y);
    const active = i === dragIdx || i === hoverIdx;
    ctx.fillStyle = active ? "#ffd17a" : "#7bb8ff";
    ctx.strokeStyle = "rgba(8,10,18,0.9)";
    ctx.lineWidth = 1.5;
    ctx.beginPath();
    ctx.arc(px, py, active ? DOT_R + 1.5 : DOT_R, 0, Math.PI * 2);
    ctx.fill();
    ctx.stroke();
  }
}

function drawHUD(ctx) {
  ctx.fillStyle = "rgba(0,0,0,0.55)";
  ctx.fillRect(W - 168, 10, 156, 70);
  ctx.fillStyle = "#fff";
  ctx.font = "13px monospace";
  ctx.textAlign = "left";
  ctx.textBaseline = "alphabetic";
  ctx.fillText(`m  = ${m.toFixed(3)}`, W - 158, 30);
  ctx.fillText(`b  = ${b.toFixed(3)}`, W - 158, 48);
  ctx.fillText(`R² = ${r2.toFixed(3)}`, W - 158, 66);

  // Legend (bottom-left, inside plot)
  ctx.fillStyle = "rgba(0,0,0,0.45)";
  ctx.fillRect(PAD.l + 6, H - PAD.b - 38, 150, 30);
  ctx.fillStyle = "#7bb8ff";
  ctx.fillRect(PAD.l + 12, H - PAD.b - 30, 16, 2);
  ctx.fillStyle = "#fff";
  ctx.font = "11px system-ui, sans-serif";
  ctx.fillText("OLS fit", PAD.l + 34, H - PAD.b - 26);
  ctx.fillStyle = "rgba(255,210,140,0.85)";
  ctx.fillRect(PAD.l + 12, H - PAD.b - 18, 4, 2);
  ctx.fillRect(PAD.l + 18, H - PAD.b - 18, 4, 2);
  ctx.fillRect(PAD.l + 24, H - PAD.b - 18, 4, 2);
  ctx.fillStyle = "#fff";
  ctx.fillText("Huber robust", PAD.l + 34, H - PAD.b - 14);
}

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

  // Hover
  hoverIdx = nearestPoint(input.mouseX, input.mouseY);

  // Drag start / add point on empty click
  const clicks = input.consumeClicks();
  for (const c of clicks) {
    const hit = nearestPoint(c.x, c.y);
    if (hit < 0) {
      const v = viewport();
      if (c.x >= v.x0 && c.x <= v.x0 + v.w && c.y >= v.y0 && c.y <= v.y0 + v.h) {
        const [dx, dy] = toData(c.x, c.y);
        pts.push({ x: dx, y: dy });
      }
    }
  }

  if (input.mouseDown) {
    if (dragIdx < 0) dragIdx = nearestPoint(input.mouseX, input.mouseY);
    if (dragIdx >= 0) {
      const [dx, dy] = toData(input.mouseX, input.mouseY);
      const v = viewport();
      pts[dragIdx].x = Math.max(v.xMin, Math.min(v.xMax, dx));
      pts[dragIdx].y = Math.max(v.yMin, Math.min(v.yMax, dy));
    }
  } else {
    dragIdx = -1;
  }

  fitOLS();
  fitHuber();

  drawAxes(ctx);
  drawResiduals(ctx);
  drawLine(ctx, mH, bH, true, "rgba(255,210,140,0.75)", 1.5);
  drawLine(ctx, m, b, false, "rgba(123,184,255,0.95)", 2);
  drawPoints(ctx);
  drawHUD(ctx);
}