Spectral Clustering on Two Moons

// Spectral clustering on two interlocking moons.
// Build a k-NN graph over the points, form the unnormalised Laplacian
// L = D - A, then power iterate M = I - L/lmax (the iterate is kept
// orthogonal to the constant vector 1 so the constant eigenvector is
// deflated away). The dominant remaining direction is the Fiedler
// vector v2; sign(v2) splits the graph along its weakest bottleneck —
// which on the two-moons set means each moon gets its own cluster,
// the case where k-means fails.

const N = 220;                // total points (half per moon)
const KNN = 8;                // k for the kNN graph (mutual-kNN below)
const ITERS_PER_FRAME = 12;   // power-iteration steps per frame
let pts;                      // Float32Array length 2N  [x0,y0,x1,y1,...]
let nbrs;                     // Int16Array length N*KNN  (neighbour indices)
let deg;                      // Float32Array length N    (weighted degree)
let v;                        // Float32Array length N    (current iterate -> Fiedler)
let tmp;                      // Float32Array length N    (scratch for L*v)
let lmax;                     // upper bound for largest eigenvalue of L
let kmCent, kmAssign;         // k-means baseline (k=2) for comparison
let iter, prevSignVec, prevV; // convergence tracking
let signFlips;                // how many entries changed sign last update
let stableIters;              // consecutive iters with low residual + no flips
let converged;
let restartIn;                // frames until automatic restart after converging
let W = 0, H = 0;

function rand() { return Math.random(); }

function makeMoons(width, height) {
  // Classic sklearn make_moons geometry, rotated to canvas pixels.
  // Moon A: ( cos t,        sin t)  for t in [0,pi]  -- top arc, opens down
  // Moon B: (1 - cos t,  -sin t + 0.5) for t in [0,pi] -- bottom arc, opens up
  pts = new Float32Array(2 * N);
  const r = Math.min(width, height) * 0.26;
  const cx = width / 2 - r * 0.5;            // shift so the pair sits centred
  const cy = height * 0.42;
  const noise = r * 0.06;
  const half = N / 2;
  for (let i = 0; i < N; i++) {
    const inFirst = i < half;
    const t = ((inFirst ? i : i - half) + 0.5) / half * Math.PI;   // 0..pi
    let ux, uy;
    if (inFirst) { ux = Math.cos(t);       uy = -Math.sin(t); }     // canvas y inverted
    else         { ux = 1 - Math.cos(t);   uy = Math.sin(t) - 0.5; }
    const x = cx + ux * r + (rand() - 0.5) * 2 * noise;
    const y = cy + uy * r + (rand() - 0.5) * 2 * noise;
    pts[2 * i] = x;
    pts[2 * i + 1] = y;
  }
}

function buildKnn() {
  nbrs = new Int16Array(N * KNN);
  deg = new Float32Array(N);
  const distRow = new Float32Array(N);
  const idx = new Int16Array(N);
  // Per-point local scale sigma_i = distance to its KNN-th neighbour
  // (Zelnik-Manor & Perona 2004 "self-tuning spectral clustering").
  const sigma = new Float32Array(N);
  for (let i = 0; i < N; i++) {
    const xi = pts[2 * i], yi = pts[2 * i + 1];
    for (let j = 0; j < N; j++) {
      const dx = pts[2 * j] - xi, dy = pts[2 * j + 1] - yi;
      distRow[j] = j === i ? Infinity : dx * dx + dy * dy;
      idx[j] = j;
    }
    // Partial selection sort for the KNN smallest distances.
    for (let k = 0; k < KNN; k++) {
      let best = k;
      for (let j = k + 1; j < N; j++) if (distRow[idx[j]] < distRow[idx[best]]) best = j;
      const t = idx[k]; idx[k] = idx[best]; idx[best] = t;
      nbrs[i * KNN + k] = idx[k];
    }
    sigma[i] = Math.sqrt(distRow[nbrs[i * KNN + KNN - 1]]);
  }
  // Symmetric-union kNN: edge if j in kNN(i) OR i in kNN(j).
  // Weight = exp(-d^2 / (sigma_i * sigma_j))  (self-tuning Gaussian).
  // Near-edges (within a moon) get weight ~1; long edges across the gap
  // between moons get exponentially small weight, so the graph stays
  // connected but the cut between moons is by far the cheapest cut.
  adjW = new Float32Array(N * N);
  const inK = new Uint8Array(N * N);
  for (let i = 0; i < N; i++) for (let k = 0; k < KNN; k++) inK[i * N + nbrs[i * KNN + k]] = 1;
  for (let i = 0; i < N; i++) {
    for (let j = i + 1; j < N; j++) {
      if (!(inK[i * N + j] || inK[j * N + i])) continue;
      const dx = pts[2 * i] - pts[2 * j], dy = pts[2 * i + 1] - pts[2 * j + 1];
      const d2 = dx * dx + dy * dy;
      const w = Math.exp(-d2 / Math.max(sigma[i] * sigma[j], 1e-9));
      adjW[i * N + j] = w; adjW[j * N + i] = w;
    }
  }
  // Weighted degree and a tighter lmax bound: lmax <= 2 * max_i d_i.
  let dmax = 0;
  for (let i = 0; i < N; i++) {
    let d = 0; const row = i * N;
    for (let j = 0; j < N; j++) d += adjW[row + j];
    deg[i] = d;
    if (d > dmax) dmax = d;
  }
  lmax = Math.max(2 * dmax, 1e-6);
}

let adjW;        // Float32Array N*N — symmetric weighted adjacency
let edgeList;    // Int32Array of (i,j) pairs for drawing
let edgeWeight;  // Float32Array — matching weights

function compileEdgeList() {
  const pairs = [], wts = [];
  for (let i = 0; i < N; i++) {
    for (let j = i + 1; j < N; j++) {
      const w = adjW[i * N + j];
      if (w > 0) { pairs.push(i, j); wts.push(w); }
    }
  }
  edgeList = new Int32Array(pairs);
  edgeWeight = new Float32Array(wts);
}

function applyL(out, vec) {
  // (L v)_i = deg_i * v_i - sum_j w_{ij} v_j
  for (let i = 0; i < N; i++) {
    let s = 0;
    const row = i * N;
    for (let j = 0; j < N; j++) { const w = adjW[row + j]; if (w !== 0) s += w * vec[j]; }
    out[i] = deg[i] * vec[i] - s;
  }
}

function powerStep() {
  // Deflate constant: make v zero-mean.
  let m = 0; for (let i = 0; i < N; i++) m += v[i]; m /= N;
  for (let i = 0; i < N; i++) v[i] -= m;
  // tmp = L v ; v_new = v - tmp / lmax  ==  (I - L/lmax) v
  applyL(tmp, v);
  // Save previous v for residual measurement, then update.
  let resid = 0;
  for (let i = 0; i < N; i++) {
    const prev = v[i];
    const nv = prev - tmp[i] / lmax;
    prevV[i] = prev;
    v[i] = nv;
  }
  // Re-deflate constant component that crept in numerically.
  m = 0; for (let i = 0; i < N; i++) m += v[i]; m /= N;
  for (let i = 0; i < N; i++) v[i] -= m;
  // Normalise.
  let n2 = 0; for (let i = 0; i < N; i++) n2 += v[i] * v[i];
  const inv = 1 / Math.sqrt(Math.max(n2, 1e-30));
  for (let i = 0; i < N; i++) v[i] *= inv;
  // Residual: ||v - prevV|| (both unit-norm-ish). Sign-ambiguous, so take min.
  let r1 = 0, r2 = 0;
  for (let i = 0; i < N; i++) { const d1 = v[i] - prevV[i], d2 = v[i] + prevV[i]; r1 += d1 * d1; r2 += d2 * d2; }
  resid = Math.sqrt(Math.min(r1, r2));
  // Track sign flips for convergence display.
  signFlips = 0;
  for (let i = 0; i < N; i++) {
    const s = v[i] >= 0 ? 1 : -1;
    if (prevSignVec[i] !== s) signFlips++;
    prevSignVec[i] = s;
  }
  iter++;
  if (iter > 20 && signFlips === 0 && resid < 2e-4) stableIters++;
  else stableIters = 0;
  if (stableIters >= 10) converged = true;
}

function kmeansInit(width, height) {
  // k=2 Lloyd baseline for visual comparison ("k-means would slice through").
  kmCent = new Float32Array(4);
  kmCent[0] = width * 0.35; kmCent[1] = height * 0.45;
  kmCent[2] = width * 0.65; kmCent[3] = height * 0.55;
  kmAssign = new Int8Array(N);
}

function kmeansStep() {
  let sx0 = 0, sy0 = 0, c0 = 0, sx1 = 0, sy1 = 0, c1 = 0;
  for (let i = 0; i < N; i++) {
    const x = pts[2 * i], y = pts[2 * i + 1];
    const d0 = (x - kmCent[0]) ** 2 + (y - kmCent[1]) ** 2;
    const d1 = (x - kmCent[2]) ** 2 + (y - kmCent[3]) ** 2;
    if (d0 < d1) { kmAssign[i] = 0; sx0 += x; sy0 += y; c0++; }
    else { kmAssign[i] = 1; sx1 += x; sy1 += y; c1++; }
  }
  if (c0) { kmCent[0] = sx0 / c0; kmCent[1] = sy0 / c0; }
  if (c1) { kmCent[2] = sx1 / c1; kmCent[3] = sy1 / c1; }
}

function init({ width, height }) {
  W = width; H = height;
  v = new Float32Array(N);
  tmp = new Float32Array(N);
  prevV = new Float32Array(N);
  prevSignVec = new Int8Array(N);
  restartRun();
}

function restartRun() {
  // Resample moons with a fresh noise draw; reset power-iter state.
  makeMoons(W, H);
  buildKnn();
  compileEdgeList();
  let m = 0;
  for (let i = 0; i < N; i++) { v[i] = rand() - 0.5; m += v[i]; prevSignVec[i] = 0; }
  m /= N;
  for (let i = 0; i < N; i++) v[i] -= m;
  iter = 0; signFlips = N; stableIters = 0; converged = false; restartIn = 0;
  kmeansInit(W, H);
  for (let s = 0; s < 12; s++) kmeansStep();
}

function drawTopPanel(ctx, topY, panelH) {
  // Edges first — weight modulates alpha (Gaussian similarity).
  ctx.lineWidth = 1;
  for (let e = 0, w = 0; e < edgeList.length; e += 2, w++) {
    const i = edgeList[e], j = edgeList[e + 1];
    const a = 0.05 + 0.35 * edgeWeight[w];
    ctx.strokeStyle = `rgba(180, 200, 255, ${a})`;
    ctx.beginPath();
    ctx.moveTo(pts[2 * i], pts[2 * i + 1]);
    ctx.lineTo(pts[2 * j], pts[2 * j + 1]);
    ctx.stroke();
  }
  // Points coloured by sign(v_2); brightness ~ |v_i|.
  for (let i = 0; i < N; i++) {
    const a = 0.45 + 0.55 * Math.min(1, Math.abs(v[i]) * Math.sqrt(N));
    ctx.fillStyle = v[i] >= 0 ? `rgba(96,180,255,${a})` : `rgba(255,120,96,${a})`;
    ctx.beginPath();
    ctx.arc(pts[2 * i], pts[2 * i + 1], 3.2, 0, Math.PI * 2);
    ctx.fill();
  }
  // Dashed line: k-means decision boundary (perp. bisector of the 2 centroids).
  const mx = (kmCent[0] + kmCent[2]) * 0.5, my = (kmCent[1] + kmCent[3]) * 0.5;
  const dx = kmCent[2] - kmCent[0], dy = kmCent[3] - kmCent[1];
  const nlen = Math.hypot(dx, dy) || 1, nxp = -dy / nlen, nyp = dx / nlen;
  const L = Math.max(W, panelH);
  ctx.strokeStyle = 'rgba(255, 255, 255, 0.22)';
  ctx.setLineDash([6, 5]); ctx.lineWidth = 1;
  ctx.beginPath();
  ctx.moveTo(mx - nxp * L, my - nyp * L);
  ctx.lineTo(mx + nxp * L, my + nyp * L);
  ctx.stroke();
  ctx.setLineDash([]);
  ctx.fillStyle = 'rgba(255,255,255,0.55)';
  ctx.font = '10px system-ui, sans-serif';
  ctx.textAlign = 'left'; ctx.textBaseline = 'top';
  ctx.fillText('k-means split (cuts both moons)', 10, topY + panelH - 16);
}

function drawHeatBar(ctx, y, h) {
  // Sorted by point index. Red-to-blue heat encodes v_i / max|v|.
  const padL = 14, barW = W - 28, cellW = barW / N;
  let maxAbs = 1e-9;
  for (let i = 0; i < N; i++) if (Math.abs(v[i]) > maxAbs) maxAbs = Math.abs(v[i]);
  for (let i = 0; i < N; i++) {
    const t = v[i] / maxAbs, u = Math.abs(t);
    const r = Math.round(255 * (1 - u) + (t >= 0 ? 96 : 255) * u);
    const g = Math.round(255 * (1 - u) + (t >= 0 ? 180 : 120) * u);
    const b = Math.round(255 * (1 - u) + (t >= 0 ? 255 : 96) * u);
    ctx.fillStyle = `rgb(${r},${g},${b})`;
    ctx.fillRect(padL + i * cellW, y, Math.ceil(cellW) + 1, h);
  }
  ctx.fillStyle = 'rgba(255,255,255,0.8)';
  ctx.font = '11px system-ui, sans-serif';
  ctx.textAlign = 'left'; ctx.textBaseline = 'bottom';
  ctx.fillText('Fiedler vector v₂ (sorted by point index)', padL, y - 3);
}

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

  if (!converged) {
    for (let s = 0; s < ITERS_PER_FRAME; s++) powerStep();
    if (converged) restartIn = 180;   // ~3s on the locked-in result before resampling
  } else if (--restartIn <= 0) {
    restartRun();
  }

  // Background.
  ctx.fillStyle = '#0a0a14';
  ctx.fillRect(0, 0, W, H);

  // Layout: top panel = points + edges; bottom strip = heatmap.
  const stripH = Math.max(22, Math.min(36, Math.round(H * 0.07)));
  const stripY = H - stripH - 22;
  const topH = stripY - 8;
  drawTopPanel(ctx, 0, topH);
  drawHeatBar(ctx, stripY, stripH);

  // HUD.
  ctx.fillStyle = 'rgba(0,0,0,0.55)';
  ctx.fillRect(8, 6, 240, 38);
  ctx.fillStyle = '#e6edf3';
  ctx.font = '12px monospace';
  ctx.textAlign = 'left'; ctx.textBaseline = 'top';
  ctx.fillText(`spectral clustering  n=${N}  k=${KNN}`, 14, 12);
  ctx.fillStyle = converged ? 'rgba(126,231,135,0.95)' : 'rgba(205,217,229,0.75)';
  ctx.fillText(
    `power iter ${iter}  flips=${signFlips}${converged ? '  CONVERGED' : ''}`,
    14, 28
  );
}