Central Limit Theorem

// Central Limit Theorem demo.
// Source distribution selectable: 1=uniform, 2=exponential, 3=bimodal mixture.
// Repeatedly draw samples of size n, compute means, histogram them, overlay
// the predicted Normal(mu, sigma^2/n).

let W = 0, H = 0;
let dist = 3; // start with bimodal — most striking convergence
let n = 30;
const BINS = 60;
let hist; // counts of sample means
let total = 0;
let sumMean = 0, sumMeanSq = 0;
let lastSample = null;
let lastMean = null;
let recentMeans = []; // for the "drop" animation
let accum = 0;

// Cached source-distribution histogram. Drawing 4000 srcSample() values
// every frame to estimate the PDF was the dominant cost; the source only
// changes when [1/2/3] is pressed, so we cache.
const PDF_BINS = 80;
let pdfCache = new Float64Array(PDF_BINS);
let pdfCacheMax = 0;
let pdfCacheDist = -1;

// Hit-rects for mobile touch controls (no keyboard available).
const buttons = {};

function srcSample() {
  if (dist === 1) {
    // Uniform(0, 1)
    return Math.random();
  } else if (dist === 2) {
    // Exponential(rate=1), clipped to [0, 6]
    let u = Math.random();
    if (u < 1e-9) u = 1e-9;
    return Math.min(6, -Math.log(u));
  } else {
    // Bimodal: 0.5 * N(0.2, 0.06^2) + 0.5 * N(0.8, 0.06^2)
    const mode = Math.random() < 0.5 ? 0.2 : 0.8;
    let s = 0;
    for (let i = 0; i < 6; i++) s += Math.random();
    const z = (s - 3) / Math.sqrt(0.5); // approx N(0,1)
    return mode + 0.06 * z;
  }
}

function srcStats() {
  if (dist === 1) return { mu: 0.5, sigma2: 1 / 12, lo: 0, hi: 1, label: "Uniform(0,1)" };
  if (dist === 2) return { mu: 1, sigma2: 1, lo: 0, hi: 4, label: "Exponential(1)" };
  // Bimodal mean = 0.5, variance = 0.5*((0.2-0.5)^2 + 0.06^2) + 0.5*((0.8-0.5)^2 + 0.06^2)
  const v = 0.5 * (0.09 + 0.0036) + 0.5 * (0.09 + 0.0036);
  return { mu: 0.5, sigma2: v, lo: -0.1, hi: 1.1, label: "Bimodal mix" };
}

function reset() {
  hist = new Int32Array(BINS);
  total = 0;
  sumMean = 0;
  sumMeanSq = 0;
  lastSample = null;
  lastMean = null;
  recentMeans = [];
}

// Rebuild cached source-distribution histogram from N samples.
function refreshPdfCache() {
  const st = srcStats();
  pdfCache.fill(0);
  const N = 4000;
  for (let i = 0; i < N; i++) {
    const v = srcSample();
    const bi = Math.floor(((v - st.lo) / (st.hi - st.lo)) * PDF_BINS);
    if (bi >= 0 && bi < PDF_BINS) pdfCache[bi]++;
  }
  let m = 0;
  for (let i = 0; i < PDF_BINS; i++) if (pdfCache[i] > m) m = pdfCache[i];
  pdfCacheMax = m;
  pdfCacheDist = dist;
}

function inRect(x, y, r) {
  return r && x >= r.x && x <= r.x + r.w && y >= r.y && y <= r.y + r.h;
}

function init({ width, height }) {
  W = width; H = height;
  reset();
  refreshPdfCache();
}

function pushSample() {
  const s = new Float64Array(n);
  let m = 0;
  for (let i = 0; i < n; i++) {
    const v = srcSample();
    s[i] = v;
    m += v;
  }
  m /= n;
  lastSample = s;
  lastMean = m;
  recentMeans.push({ m, age: 0 });
  if (recentMeans.length > 30) recentMeans.shift();

  const st = srcStats();
  const bi = Math.floor(((m - st.lo) / (st.hi - st.lo)) * BINS);
  if (bi >= 0 && bi < BINS) hist[bi]++;
  total++;
  sumMean += m;
  sumMeanSq += m * m;
}

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

  // keys
  if (input.justPressed("1")) { dist = 1; reset(); refreshPdfCache(); }
  if (input.justPressed("2")) { dist = 2; reset(); refreshPdfCache(); }
  if (input.justPressed("3")) { dist = 3; reset(); refreshPdfCache(); }
  if (input.justPressed("ArrowUp")) { n = Math.min(200, n + 5); reset(); }
  if (input.justPressed("ArrowDown")) { n = Math.max(2, n - 5); reset(); }
  if (input.justPressed("r") || input.justPressed("R")) reset();

  // consumeClicks() returns an array of {x,y,button}, NOT a count.
  // A click inside a button rect = touch control; otherwise = burst.
  const clicks = input.consumeClicks ? input.consumeClicks() : [];
  let burstClicks = 0;
  for (let i = 0; i < clicks.length; i++) {
    const ev = clicks[i];
    const cx = (ev && ev.x != null) ? ev.x : input.mouseX;
    const cy = (ev && ev.y != null) ? ev.y : input.mouseY;
    if (inRect(cx, cy, buttons.d1))     { dist = 1; reset(); refreshPdfCache(); continue; }
    if (inRect(cx, cy, buttons.d2))     { dist = 2; reset(); refreshPdfCache(); continue; }
    if (inRect(cx, cy, buttons.d3))     { dist = 3; reset(); refreshPdfCache(); continue; }
    if (inRect(cx, cy, buttons.nMinus)) { n = Math.max(2, n - 5); reset(); continue; }
    if (inRect(cx, cy, buttons.nPlus))  { n = Math.min(200, n + 5); reset(); continue; }
    if (inRect(cx, cy, buttons.reset))  { reset(); continue; }
    burstClicks++;
  }
  if (pdfCacheDist !== dist) refreshPdfCache();

  // auto-draw samples, or burst on click
  accum += dt;
  let toDraw = 0;
  while (accum > 0.06) { accum -= 0.06; toDraw++; }
  toDraw += burstClicks * 20;
  for (let i = 0; i < toDraw; i++) pushSample();

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

  const st = srcStats();
  const pad = 32;
  const topH = H * 0.38;
  const botY = topH + 70;
  // Leave 78px at the bottom for stats + touch-control row.
  const botH = H - botY - 78;

  // ---- top panel: source distribution + last sample ----
  ctx.fillStyle = "#13131c";
  ctx.fillRect(pad, 40, W - 2 * pad, topH - 10);

  // approximate source PDF — reused from cache (refreshed only when `dist` changes)
  const pdfBins = PDF_BINS;
  const pdf = pdfCache;
  const pdfMax = pdfCacheMax || 1;

  const topX0 = pad + 10;
  const topX1 = W - pad - 10;
  const topY0 = 50;
  const topY1 = 40 + topH - 20;
  const tw = topX1 - topX0;
  const th = topY1 - topY0;

  ctx.fillStyle = "rgba(120,180,255,0.55)";
  for (let i = 0; i < pdfBins; i++) {
    const h = (pdf[i] / Math.max(1, pdfMax)) * th;
    const x = topX0 + (i / pdfBins) * tw;
    ctx.fillRect(x, topY1 - h, tw / pdfBins - 1, h);
  }

  // mu marker
  const muX = topX0 + ((st.mu - st.lo) / (st.hi - st.lo)) * tw;
  ctx.strokeStyle = "rgba(255,200,80,0.9)";
  ctx.setLineDash([4, 4]);
  ctx.beginPath();
  ctx.moveTo(muX, topY0);
  ctx.lineTo(muX, topY1);
  ctx.stroke();
  ctx.setLineDash([]);

  // last sample points
  if (lastSample) {
    ctx.fillStyle = "rgba(255,255,255,0.9)";
    for (let i = 0; i < lastSample.length; i++) {
      const x = topX0 + ((lastSample[i] - st.lo) / (st.hi - st.lo)) * tw;
      const y = topY1 - 4 - (i % 6) * 2.5;
      ctx.beginPath();
      ctx.arc(x, y, 2.2, 0, Math.PI * 2);
      ctx.fill();
    }
    if (lastMean != null) {
      const mx = topX0 + ((lastMean - st.lo) / (st.hi - st.lo)) * tw;
      ctx.strokeStyle = "rgba(120,255,150,1)";
      ctx.lineWidth = 2;
      ctx.beginPath();
      ctx.moveTo(mx, topY0 + 4);
      ctx.lineTo(mx, topY1);
      ctx.stroke();
      ctx.lineWidth = 1;
    }
  }

  // ---- bottom panel: histogram of sample means + Normal overlay ----
  ctx.fillStyle = "#13131c";
  ctx.fillRect(pad, botY, W - 2 * pad, botH);

  const bX0 = pad + 10;
  const bX1 = W - pad - 10;
  const bY0 = botY + 10;
  const bY1 = botY + botH - 20;
  const bw = bX1 - bX0;
  const bh = bY1 - bY0;

  // x-range for means is narrower than source (SE = sigma/sqrt(n))
  const se = Math.sqrt(st.sigma2 / n);
  const meanLo = st.mu - 4 * se;
  const meanHi = st.mu + 4 * se;

  // Convert hist bins (over [lo, hi]) onto plot x-axis [meanLo, meanHi]
  let histMax = 0;
  for (let i = 0; i < BINS; i++) if (hist[i] > histMax) histMax = hist[i];

  // hist as density: count / total / binWidth
  const binW = (st.hi - st.lo) / BINS;
  const densMax = histMax / Math.max(1, total) / binW;

  // pdf of N(mu, se^2) max = 1/(sqrt(2pi)*se)
  const normMax = 1 / (Math.sqrt(2 * Math.PI) * se);
  const yMax = Math.max(densMax * 1.1, normMax * 1.1, 1e-6);

  // bars
  ctx.fillStyle = "rgba(120,255,150,0.55)";
  for (let i = 0; i < BINS; i++) {
    if (hist[i] === 0) continue;
    const xCenter = st.lo + (i + 0.5) * binW;
    if (xCenter < meanLo || xCenter > meanHi) continue;
    const dens = hist[i] / Math.max(1, total) / binW;
    const x = bX0 + ((xCenter - meanLo) / (meanHi - meanLo)) * bw;
    const wpx = (binW / (meanHi - meanLo)) * bw;
    const hpx = (dens / yMax) * bh;
    ctx.fillRect(x - wpx / 2, bY1 - hpx, Math.max(1, wpx - 0.5), hpx);
  }

  // Normal overlay
  ctx.strokeStyle = "rgba(255,200,80,0.95)";
  ctx.lineWidth = 2;
  ctx.beginPath();
  for (let px = 0; px <= bw; px += 1) {
    const xv = meanLo + (px / bw) * (meanHi - meanLo);
    const z = (xv - st.mu) / se;
    const p = Math.exp(-0.5 * z * z) / (Math.sqrt(2 * Math.PI) * se);
    const yp = bY1 - (p / yMax) * bh;
    if (px === 0) ctx.moveTo(bX0 + px, yp);
    else ctx.lineTo(bX0 + px, yp);
  }
  ctx.stroke();
  ctx.lineWidth = 1;

  // mu marker on histogram
  const muXh = bX0 + ((st.mu - meanLo) / (meanHi - meanLo)) * bw;
  ctx.strokeStyle = "rgba(255,200,80,0.5)";
  ctx.setLineDash([3, 4]);
  ctx.beginPath();
  ctx.moveTo(muXh, bY0);
  ctx.lineTo(muXh, bY1);
  ctx.stroke();
  ctx.setLineDash([]);

  // recent-mean drops
  for (let i = recentMeans.length - 1; i >= 0; i--) {
    const r = recentMeans[i];
    r.age += dt;
    if (r.age > 0.8) { recentMeans.splice(i, 1); continue; }
    const x = bX0 + ((r.m - meanLo) / (meanHi - meanLo)) * bw;
    if (x < bX0 || x > bX1) continue;
    const a = 1 - r.age / 0.8;
    ctx.fillStyle = `rgba(255,255,255,${(0.7 * a).toFixed(3)})`;
    ctx.beginPath();
    ctx.arc(x, bY0 + 10 + r.age * 60, 3, 0, Math.PI * 2);
    ctx.fill();
  }

  // ---- text ----
  ctx.fillStyle = "#e8e8f0";
  ctx.font = "bold 16px monospace";
  ctx.fillText("Central Limit Theorem", pad, 26);
  ctx.font = "12px monospace";
  ctx.fillStyle = "#aab";
  ctx.fillText(`source: ${st.label}   n = ${n}   samples = ${total}`, pad, topH + 50);

  ctx.fillStyle = "#9cf";
  ctx.font = "11px monospace";
  const empMean = total > 0 ? sumMean / total : 0;
  const empVar = total > 1 ? sumMeanSq / total - empMean * empMean : 0;
  // Stats live just above the touch-control row at H-26 (avoid overlap).
  ctx.fillText(`mean of sample means ≈ ${empMean.toFixed(4)}    (μ = ${st.mu.toFixed(4)})`, pad, H - 60);
  ctx.fillText(`var of sample means ≈ ${empVar.toFixed(5)}    (σ²/n = ${(st.sigma2 / n).toFixed(5)})`, pad, H - 44);

  // Touch / mobile controls — keyboard shortcuts mirror these.
  // Sized & laid out from current W so they re-flow on resize.
  const btnY = H - 26;
  const btnH = 20;
  const gap = 6;
  const labels = [
    { key: "d1",     text: dist === 1 ? "[1] uniform*"    : "[1] uniform" },
    { key: "d2",     text: dist === 2 ? "[2] expon*"      : "[2] expon" },
    { key: "d3",     text: dist === 3 ? "[3] bimodal*"    : "[3] bimodal" },
    { key: "nMinus", text: "n-" },
    { key: "nPlus",  text: "n+" },
    { key: "reset",  text: "[R]" },
  ];
  ctx.font = "10px monospace";
  // Measure to size buttons proportionally; fall back when ctx.measureText absent.
  const widths = labels.map(l => {
    const m = ctx.measureText ? ctx.measureText(l.text) : null;
    const tw = (m && m.width) ? m.width : l.text.length * 6.2;
    return Math.max(28, tw + 10);
  });
  // Fit row into [pad, W - pad]; shrink uniformly if needed.
  const rowW = widths.reduce((s, w) => s + w, 0) + gap * (labels.length - 1);
  const avail = W - 2 * pad;
  const shrink = rowW > avail ? avail / rowW : 1;
  let bx = pad;
  for (let i = 0; i < labels.length; i++) {
    const bw = widths[i] * shrink;
    const r = { x: bx, y: btnY, w: bw, h: btnH };
    buttons[labels[i].key] = r;
    const active = (labels[i].key === "d" + dist);
    ctx.fillStyle = active ? "#2a4055" : "#1a1a26";
    ctx.fillRect(r.x, r.y, r.w, r.h);
    ctx.strokeStyle = "#445";
    ctx.strokeRect(r.x + 0.5, r.y + 0.5, r.w - 1, r.h - 1);
    ctx.fillStyle = "#cce";
    ctx.fillText(labels[i].text, r.x + 5, r.y + 14);
    bx += bw + gap * shrink;
  }
}