727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 1import { useEffect, useRef, type CSSProperties } from "react";
3export interface Points {
4 x: number[];
5 y: number[];
6}
8export interface Pt {
9 x: number;
10 y: number;
11}
13/** A recorded orbit path: positions x/y with a per-point macro-step id `seg`,
14 * so the polyline breaks where the orbit direction flips. */
15export interface OrbitPath {
16 x: number[];
17 y: number[];
18 seg: number[];
19}
21export interface DensityGrid {
22 /** Row-major log-density: index (iy)*nx + ix, iy from ymin (0) to ymax. */
23 values: number[];
24 nx: number;
25 ny: number;
26 xmin: number;
27 xmax: number;
28 ymin: number;
29 ymax: number;
30}
32interface DensityViewProps {
33 density: DensityGrid;
34 samples: Points;
35 // ── movie overlay (all optional) ──
36 orbit?: OrbitPath | null; // revealed orbit path of the current transition
37 start?: Pt | null; // start point of the current transition (ringed)
38 lead?: Pt | null; // current frontier point of the revealed path
39 chainPts?: Points | null; // accepted draws so far (the Markov chain)
40 selected?: Pt | null; // the just-selected draw (circled)
41}
43const DOT = "rgba(15, 23, 42, 0.5)";
44const DOT_RADIUS = 1.5;
45const MARGIN = 16;
46// Heatmap colour ramp: white (low density) → blue (high density).
47const LO: [number, number, number] = [255, 255, 255];
48const HI: [number, number, number] = [37, 99, 235];
49const ORBIT = "#f59e0b"; // amber: the leapfrog orbit path
50const START = "#334155"; // slate ring: where the transition started
51const SELECTED = "#0f172a"; // near-black: the selected draw (circled)
53/** Renders the target density as a heatmap with samples scattered on top, plus
54 * an optional orbit overlay for the step-by-step movie. Fits the density's
55 * world bounds to the canvas (aspect preserved, y up). DPR-aware; redraws on
56 * data change and resize. */
57export function DensityView({
58 density,
59 samples,
60 orbit,
61 start,
62 lead,
63 chainPts,
64 selected,
65}: DensityViewProps) {
66 const canvasRef = useRef<HTMLCanvasElement>(null);
67 const containerRef = useRef<HTMLDivElement>(null);
69 useEffect(() => {
70 const canvas = canvasRef.current;
71 const container = containerRef.current;
72 if (!canvas || !container) return;
74 const draw = () => {
75 const ctx = canvas.getContext("2d");
76 if (!ctx) return;
78 const dpr = window.devicePixelRatio || 1;
79 const cssW = container.clientWidth;
80 const cssH = container.clientHeight;
81 if (cssW === 0 || cssH === 0) return;
83 canvas.width = Math.round(cssW * dpr);
84 canvas.height = Math.round(cssH * dpr);
85 canvas.style.width = `${cssW}px`;
86 canvas.style.height = `${cssH}px`;
87 ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
88 ctx.clearRect(0, 0, cssW, cssH);
90 const { nx, ny, xmin, xmax, ymin, ymax, values } = density;
91 if (!nx || !ny) return;
93 const worldW = xmax - xmin || 1;
94 const worldH = ymax - ymin || 1;
95 const scale = Math.min(
96 (cssW - 2 * MARGIN) / worldW,
97 (cssH - 2 * MARGIN) / worldH
98 );
99 const destW = worldW * scale;
100 const destH = worldH * scale;
101 const destX = (cssW - destW) / 2;
102 const destY = (cssH - destH) / 2;
103 const toPx = (x: number) => destX + (x - xmin) * scale;
104 const toPy = (y: number) => destY + (ymax - y) * scale; // y up
106 // Heatmap: paint the grid into an offscreen nx×ny image, then scale it in.
107 let maxLp = -Infinity;
108 for (let i = 0; i < values.length; i++) {
109 if (values[i] > maxLp) maxLp = values[i];
110 }
111 const off = document.createElement("canvas");
112 off.width = nx;
113 off.height = ny;
114 const offCtx = off.getContext("2d");
115 if (offCtx) {
116 const img = offCtx.createImageData(nx, ny);
117 for (let r = 0; r < ny; r++) {
118 const iy = ny - 1 - r; // image row 0 = top = ymax
119 for (let c = 0; c < nx; c++) {
120 const t = Math.pow(Math.exp(values[iy * nx + c] - maxLp), 0.4);
121 const p = (r * nx + c) * 4;
122 img.data[p] = LO[0] + t * (HI[0] - LO[0]);
123 img.data[p + 1] = LO[1] + t * (HI[1] - LO[1]);
124 img.data[p + 2] = LO[2] + t * (HI[2] - LO[2]);
125 img.data[p + 3] = 255;
126 }
127 }
128 offCtx.putImageData(img, 0, 0);
129 ctx.imageSmoothingEnabled = true;
130 ctx.drawImage(off, destX, destY, destW, destH);
131 }
133 ctx.strokeStyle = "rgba(15,23,42,0.15)";
134 ctx.lineWidth = 1;
135 ctx.strokeRect(destX, destY, destW, destH);
137 // Sample cloud (dimmed when the movie overlay is active).
138 ctx.fillStyle = orbit || chainPts ? "rgba(15,23,42,0.18)" : DOT;
139 const n = Math.min(samples.x.length, samples.y.length);
140 for (let i = 0; i < n; i++) {
141 ctx.beginPath();
142 ctx.arc(toPx(samples.x[i]), toPy(samples.y[i]), DOT_RADIUS, 0, 2 * Math.PI);
143 ctx.fill();
144 }
146 // Accepted draws so far (the Markov chain).
147 if (chainPts) {
148 ctx.fillStyle = SELECTED;
149 const m = Math.min(chainPts.x.length, chainPts.y.length);
150 for (let i = 0; i < m; i++) {
151 ctx.beginPath();
152 ctx.arc(toPx(chainPts.x[i]), toPy(chainPts.y[i]), 2, 0, 2 * Math.PI);
153 ctx.fill();
154 }
155 }
157 // The orbit being traced: polyline (broken at segment changes) + a dot at
158 // each leapfrog step, so the discrete steps and step-size are visible.
159 if (orbit && orbit.x.length > 0) {
160 ctx.strokeStyle = ORBIT;
161 ctx.lineWidth = 1.5;
162 ctx.beginPath();
163 for (let i = 0; i < orbit.x.length; i++) {
164 const px = toPx(orbit.x[i]);
165 const py = toPy(orbit.y[i]);
166 if (i === 0 || orbit.seg[i] !== orbit.seg[i - 1]) ctx.moveTo(px, py);
167 else ctx.lineTo(px, py);
168 }
169 ctx.stroke();
170 ctx.fillStyle = ORBIT;
171 for (let i = 0; i < orbit.x.length; i++) {
172 ctx.beginPath();
173 ctx.arc(toPx(orbit.x[i]), toPy(orbit.y[i]), 1.7, 0, 2 * Math.PI);
174 ctx.fill();
175 }
176 }
178 // Start of the current transition (open ring).
179 if (start) {
180 ctx.strokeStyle = START;
181 ctx.lineWidth = 1.5;
182 ctx.beginPath();
183 ctx.arc(toPx(start.x), toPy(start.y), 4, 0, 2 * Math.PI);
184 ctx.stroke();
185 }
187 // Current frontier of the orbit.
188 if (lead) {
189 ctx.fillStyle = ORBIT;
190 ctx.beginPath();
191 ctx.arc(toPx(lead.x), toPy(lead.y), 3, 0, 2 * Math.PI);
192 ctx.fill();
193 }
195 // The selected draw: a circled black point.
196 if (selected) {
197 const cx = toPx(selected.x);
198 const cy = toPy(selected.y);
199 ctx.strokeStyle = SELECTED;
200 ctx.lineWidth = 1.5;
201 ctx.beginPath();
202 ctx.arc(cx, cy, 6.5, 0, 2 * Math.PI);
203 ctx.stroke();
204 ctx.fillStyle = SELECTED;
205 ctx.beginPath();
206 ctx.arc(cx, cy, 3.2, 0, 2 * Math.PI);
207 ctx.fill();
208 }
209 };
211 draw();
212 const ro = new ResizeObserver(draw);
213 ro.observe(container);
214 return () => ro.disconnect();
215 }, [density, samples, orbit, start, lead, chainPts, selected]);
217 return (
218 <div ref={containerRef} style={containerStyle}>
219 <canvas ref={canvasRef} style={{ display: "block" }} />
220 </div>
221 );
222}
224const containerStyle: CSSProperties = {
225 position: "absolute",
226 inset: 0,
227 overflow: "hidden",
228};