diff --git a/src/App.tsx b/src/App.tsx index 427d9c8..4ca19c2 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -48,6 +48,7 @@ const Lsystem = lazy(() => import("./pages/Lsystem")); const Bezier = lazy(() => import("./pages/Bezier")); const Percolation = lazy(() => import("./pages/Percolation")); const Epidemic = lazy(() => import("./pages/Epidemic")); +const Descent = lazy(() => import("./pages/Descent")); const NotFound = lazy(() => import("./pages/NotFound")); const App = () => ( @@ -550,6 +551,19 @@ const App = () => ( } /> + + + + + } + /> {/* ADD ALL CUSTOM ROUTES ABOVE THE CATCH-ALL "*" ROUTE */} number; // loss, roughly [0,1] range over [-1,1]^2 + grad: (x: number, y: number) => [number, number]; +} + +const gauss = (x: number, y: number, cx: number, cy: number, s: number) => + Math.exp(-((x - cx) ** 2 + (y - cy) ** 2) / (2 * s * s)); + +// Analytic gradient of a gaussian pit/bump with amplitude a +const gaussGrad = ( + x: number, y: number, cx: number, cy: number, s: number, a: number +): [number, number] => { + const g = a * gauss(x, y, cx, cy, s); + return [(-(x - cx) / (s * s)) * g, (-(y - cy) / (s * s)) * g]; +}; + +export const SURFACES: Record = { + bowl: { + label: "Tilted bowl", + f: (x, y) => 0.5 * (x * x + 1.8 * y * y) + 0.15 * x, + grad: (x, y) => [x + 0.15, 1.8 * y], + }, + twopits: { + label: "Two pits", + f: (x, y) => + 0.25 * (x * x + y * y) - + 0.55 * gauss(x, y, -0.45, -0.1, 0.28) - + 0.9 * gauss(x, y, 0.5, 0.25, 0.22), + grad: (x, y) => { + const [g1x, g1y] = gaussGrad(x, y, -0.45, -0.1, 0.28, -0.55); + const [g2x, g2y] = gaussGrad(x, y, 0.5, 0.25, 0.22, -0.9); + return [0.5 * x + g1x + g2x, 0.5 * y + g1y + g2y]; + }, + }, + ripple: { + label: "Ripples", + f: (x, y) => + 0.4 * (x * x + y * y) + 0.08 * Math.sin(9 * x) + 0.08 * Math.sin(9 * y), + grad: (x, y) => [ + 0.8 * x + 0.72 * Math.cos(9 * x), + 0.8 * y + 0.72 * Math.cos(9 * y), + ], + }, +}; + +export type OptimizerId = "sgd" | "momentum" | "adam"; + +export interface OptState { + x: number; + y: number; + vx: number; + vy: number; + mx: number; + my: number; // Adam first moment + sx: number; + sy: number; // Adam second moment + t: number; + trail: { x: number; y: number }[]; + done: boolean; +} + +export const OPTIMIZERS: Record = { + sgd: { label: "SGD", color: "#00cfff" }, + momentum: { label: "Momentum", color: "#ffe066" }, + adam: { label: "Adam", color: "#ff4da6" }, +}; + +export function newOpt(x: number, y: number): OptState { + return { x, y, vx: 0, vy: 0, mx: 0, my: 0, sx: 0, sy: 0, t: 0, trail: [{ x, y }], done: false }; +} + +const BETA1 = 0.9, BETA2 = 0.999, EPS = 1e-8; +const MU = 0.9; // momentum coefficient +const TRAIL_MAX = 900; + +export function stepOpt( + o: OptState, + id: OptimizerId, + surface: Surface, + lr: number +): void { + if (o.done) return; + const [gx, gy] = surface.grad(o.x, o.y); + + if (id === "sgd") { + o.x -= lr * gx; + o.y -= lr * gy; + } else if (id === "momentum") { + o.vx = MU * o.vx - lr * gx; + o.vy = MU * o.vy - lr * gy; + o.x += o.vx; + o.y += o.vy; + } else { + o.t++; + o.mx = BETA1 * o.mx + (1 - BETA1) * gx; + o.my = BETA1 * o.my + (1 - BETA1) * gy; + o.sx = BETA2 * o.sx + (1 - BETA2) * gx * gx; + o.sy = BETA2 * o.sy + (1 - BETA2) * gy * gy; + const mxh = o.mx / (1 - Math.pow(BETA1, o.t)); + const myh = o.my / (1 - Math.pow(BETA1, o.t)); + const sxh = o.sx / (1 - Math.pow(BETA2, o.t)); + const syh = o.sy / (1 - Math.pow(BETA2, o.t)); + // Adam works on a different lr scale; boost so races are fair to watch + const alr = lr * 3; + o.x -= (alr * mxh) / (Math.sqrt(sxh) + EPS); + o.y -= (alr * myh) / (Math.sqrt(syh) + EPS); + } + + // Clamp to the domain + o.x = Math.max(-1, Math.min(1, o.x)); + o.y = Math.max(-1, Math.min(1, o.y)); + + o.trail.push({ x: o.x, y: o.y }); + if (o.trail.length > TRAIL_MAX) o.trail.shift(); + + // Converged when the gradient is tiny + if (gx * gx + gy * gy < 1e-8) o.done = true; +} diff --git a/src/pages/Descent.tsx b/src/pages/Descent.tsx new file mode 100644 index 0000000..6e0c80f --- /dev/null +++ b/src/pages/Descent.tsx @@ -0,0 +1,248 @@ +import { useEffect, useRef, useState, useCallback } from "react"; +import { Link } from "react-router-dom"; +import { + SurfaceId, + SURFACES, + OptimizerId, + OPTIMIZERS, + OptState, + newOpt, + stepOpt, +} from "../features/descent/optimizers"; + +const SURFACE_KEYS = Object.keys(SURFACES) as SurfaceId[]; +const OPT_KEYS = Object.keys(OPTIMIZERS) as OptimizerId[]; +const HEAT_RES = 3; // px per heatmap sample + +export default function Descent() { + const canvasRef = useRef(null); + const heatRef = useRef(null); + const optsRef = useRef | null>(null); + const surfaceRef = useRef("twopits"); + const lrRef = useRef(0.02); + + const [surface, setSurface] = useState("twopits"); + const [lr, setLr] = useState(0.02); + const [steps, setSteps] = useState(0); + + // Render the loss surface to an offscreen heatmap + const buildHeatmap = useCallback(() => { + const canvas = canvasRef.current; + if (!canvas) return; + const { width: W, height: H } = canvas; + const gw = Math.max(2, Math.floor(W / HEAT_RES)); + const gh = Math.max(2, Math.floor(H / HEAT_RES)); + const off = document.createElement("canvas"); + off.width = gw; + off.height = gh; + const octx = off.getContext("2d")!; + const img = octx.createImageData(gw, gh); + const surf = SURFACES[surfaceRef.current]; + + // Find range for normalization + let min = Infinity, max = -Infinity; + const vals = new Float32Array(gw * gh); + for (let gy = 0; gy < gh; gy++) { + for (let gx = 0; gx < gw; gx++) { + const x = (gx / (gw - 1)) * 2 - 1; + const y = (gy / (gh - 1)) * 2 - 1; + const v = surf.f(x, y); + vals[gy * gw + gx] = v; + if (v < min) min = v; + if (v > max) max = v; + } + } + const span = Math.max(1e-9, max - min); + for (let i = 0; i < vals.length; i++) { + const t = (vals[i] - min) / span; // 0 = deepest + // Deep = bright green, high = near black, with contour banding + const band = Math.abs(((vals[i] - min) / span * 14) % 1 - 0.5) < 0.06 ? 0.25 : 0; + const g = Math.max(0, (1 - t) * 150 + band * 80); + img.data[i * 4] = g * 0.15; + img.data[i * 4 + 1] = g; + img.data[i * 4 + 2] = g * 0.3; + img.data[i * 4 + 3] = 255; + } + octx.putImageData(img, 0, 0); + heatRef.current = off; + }, []); + + const dropAt = useCallback((x: number, y: number) => { + optsRef.current = { + sgd: newOpt(x, y), + momentum: newOpt(x, y), + adam: newOpt(x, y), + }; + setSteps(0); + }, []); + + useEffect(() => { + const canvas = canvasRef.current; + if (!canvas) return; + const ctx = canvas.getContext("2d"); + if (!ctx) return; + + const resize = () => { + const rect = canvas.parentElement!.getBoundingClientRect(); + canvas.width = rect.width; + canvas.height = rect.height; + buildHeatmap(); + dropAt(-0.8, -0.75); + }; + resize(); + window.addEventListener("resize", resize); + + let frame = 0; + let raf = 0; + const loop = () => { + raf = requestAnimationFrame(loop); + frame++; + const { width: W, height: H } = canvas; + const opts = optsRef.current; + const surf = SURFACES[surfaceRef.current]; + + if (opts) { + for (const id of OPT_KEYS) { + stepOpt(opts[id], id, surf, lrRef.current); + } + if (frame % 10 === 0) { + setSteps((s) => s + 10); + } + } + + const toPx = (x: number) => ((x + 1) / 2) * W; + const toPy = (y: number) => ((y + 1) / 2) * H; + + // Heatmap + if (heatRef.current) { + ctx.imageSmoothingEnabled = true; + ctx.drawImage(heatRef.current, 0, 0, W, H); + } + + // Trails + heads + if (opts) { + for (const id of OPT_KEYS) { + const o = opts[id]; + const color = OPTIMIZERS[id].color; + ctx.strokeStyle = color; + ctx.lineWidth = 1.5; + ctx.beginPath(); + o.trail.forEach((p, i) => { + const px = toPx(p.x); + const py = toPy(p.y); + if (i === 0) ctx.moveTo(px, py); + else ctx.lineTo(px, py); + }); + ctx.stroke(); + ctx.beginPath(); + ctx.arc(toPx(o.x), toPy(o.y), 5, 0, Math.PI * 2); + ctx.fillStyle = color; + ctx.fill(); + ctx.strokeStyle = "#000"; + ctx.lineWidth = 1; + ctx.stroke(); + } + } + }; + raf = requestAnimationFrame(loop); + + return () => { + cancelAnimationFrame(raf); + window.removeEventListener("resize", resize); + }; + }, [buildHeatmap, dropAt]); + + const onClick = (e: React.MouseEvent) => { + const canvas = canvasRef.current!; + const rect = canvas.getBoundingClientRect(); + const x = ((e.clientX - rect.left) / rect.width) * 2 - 1; + const y = ((e.clientY - rect.top) / rect.height) * 2 - 1; + dropAt(x, y); + }; + + const handleSurface = (s: SurfaceId) => { + surfaceRef.current = s; + setSurface(s); + buildHeatmap(); + dropAt(-0.8, -0.75); + }; + + return ( +
+ {/* Header */} +
+
+ + ← home + + | + gradient descent arena +
+
+ {steps} steps + {OPT_KEYS.map((id) => ( + + ● {OPTIMIZERS[id].label} + + ))} +
+
+ + {/* Controls */} +
+
+ {SURFACE_KEYS.map((s) => ( + + ))} +
+ +
+ learning rate + { + const v = Number(e.target.value); + lrRef.current = v; + setLr(v); + }} + className="w-28 accent-primary" + /> + {lr.toFixed(3)} +
+ + +
+ + {/* Canvas */} +
+ +
+ click to drop all three optimizers · bright = low loss · on ripples, watch SGD stall in local dips while momentum coasts through +
+
+
+ ); +}