727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 1import { useEffect, useState, type CSSProperties } from "react";
2import {
3 DensityView,
4 type Points,
5 type DensityGrid,
6 type OrbitPath,
7 type Pt,
8} from "./render/DensityView.js";
9import { onData, onHostEvent, sendToMATLAB } from "./bridge.js";
11/** Payload from the numbl script: the target density (for the heatmap) plus the
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 12 * WALNUTS samples. Mirrors what walnuts_sampler.m sends. `h` is the leapfrog
13 * step and `delta` the energy-variation tolerance for within-orbit refinement. */
15 type: "walnuts";
16 density: DensityGrid;
17 samples: Points;
18 n: number;
20 delta: number;
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 21 target: string;
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 24const TARGETS: { value: string; label: string }[] = [
25 { value: "banana", label: "Banana" },
26 { value: "gaussian", label: "Gaussian" },
27 { value: "correlated", label: "Correlated Gaussian" },
28 { value: "donut", label: "Donut (ring)" },
29];
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 31function isWalnutsData(d: unknown): d is WalnutsData {
32 return (
33 !!d &&
34 typeof d === "object" &&
35 (d as WalnutsData).type === "walnuts" &&
36 !!(d as WalnutsData).density &&
37 !!(d as WalnutsData).samples
38 );
39}
41interface SamplesEvent {
42 x: number[];
43 y: number[];
44 n: number;
46 delta: number;
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 49/** One recorded transition's orbit (from `walnuts(...)`'s orbit output). */
51 px: number[];
52 py: number[];
53 seg: number[];
54 startX: number;
55 startY: number;
56 selX: number;
57 selY: number;
58}
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 60// WALNUTS' orbit-based transition is heavier than a plain MH step, so keep the
61// sample counts modest.
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 62const SAMPLE_CHOICES = [100, 300, 1000, 3000];
63const DEFAULT_N = 1000;
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 64const DEFAULT_H = 0.8;
65const DEFAULT_DELTA = Math.log(1 / 0.66); // ≈ 0.42 (Nawaf's amin = 0.66)
66const H_MIN = 0.1;
67const H_MAX = 2.0;
68const DELTA_MIN = 0.1;
69const DELTA_MAX = 2.0;
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 70
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 71// Movie pacing: reveal a couple of orbit points per tick, then linger on the
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 72// selected draw before the next transition.
73const MOVIE_TICK_MS = 55;
74const MOVIE_REVEAL = 2;
75const MOVIE_HOLD = 14; // extra k-units to hold the selected point
77export function App() {
78 const [data, setData] = useState<WalnutsData | null>(null);
79 const [n, setN] = useState(DEFAULT_N);
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 80 const [h, setH] = useState(DEFAULT_H);
81 const [delta, setDelta] = useState(DEFAULT_DELTA);
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 82 const [target, setTargetState] = useState("banana");
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 83 const [busy, setBusy] = useState(false);
84 const [movieData, setMovieData] = useState<MovieStep[] | null>(null);
85 const [movie, setMovie] = useState<{ si: number; k: number } | null>(null);
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 87 // Apply a full payload (initial Data, or a `data` event after a target change).
88 const applyData = (d: WalnutsData) => {
89 setData(d);
90 setN(d.n);
92 setDelta(d.delta);
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 93 setTargetState(d.target);
94 };
97 const offData = onData(d => {
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 98 if (isWalnutsData(d)) applyData(d);
99 });
100 // New target: full fresh payload (density + samples).
101 const offFull = onHostEvent("data", d => {
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 103 applyData(d);
104 setBusy(false);
105 setMovie(null);
106 setMovieData(null);
108 });
109 // Resample: same target, new draws.
110 const offSamples = onHostEvent("samples", ev => {
111 const s = ev as SamplesEvent;
112 if (!s || !Array.isArray(s.x)) return;
113 setData(prev =>
114 prev
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 115 ? { ...prev, samples: { x: s.x, y: s.y }, n: s.n, h: s.h, delta: s.delta }
117 );
118 setBusy(false);
119 });
120 // Movie: an array of recorded transition orbits to animate.
121 const offMovie = onHostEvent("movie", ev => {
122 if (!Array.isArray(ev) || ev.length === 0) {
123 setBusy(false);
124 return;
125 }
126 setMovieData(ev as MovieStep[]);
127 setMovie({ si: 0, k: 0 });
128 setBusy(false);
129 });
130 return () => {
131 offData();
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 132 offFull();
134 offMovie();
135 };
136 }, []);
138 // Movie clock.
139 useEffect(() => {
140 if (!movie || !movieData) return;
141 const id = setTimeout(() => {
142 setMovie(m => {
143 if (!m) return m;
144 const step = movieData[m.si];
145 const count = step.px.length;
146 if (m.k < count + MOVIE_HOLD) return { si: m.si, k: m.k + MOVIE_REVEAL };
147 const nextSi = m.si + 1;
148 return nextSi >= movieData.length ? null : { si: nextSi, k: 0 };
149 });
150 }, MOVIE_TICK_MS);
151 return () => clearTimeout(id);
152 }, [movie, movieData]);
154 const stopMovie = () => {
155 setMovie(null);
156 setMovieData(null);
157 };
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 159 const resample = (count: number, hVal: number, deltaVal: number) => {
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 160 if (!data || busy) return;
161 stopMovie();
162 setBusy(true);
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 163 sendToMATLAB("resample", { n: count, h: hVal, delta: deltaVal, target });
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 164 };
166 // Switch the target: the script rebuilds the density + draws and replies with
167 // a full `data` event.
168 const changeTarget = (value: string) => {
169 setTargetState(value);
170 if (!data) return;
171 stopMovie();
172 setBusy(true);
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 173 sendToMATLAB("setTarget", { target: value, n, h, delta });
176 const playMovie = () => {
177 if (movie) {
178 stopMovie();
179 return;
180 }
181 if (!data || busy) return;
182 setBusy(true); // until the trajectory arrives
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 183 sendToMATLAB("movie", { h, delta, target });
186 // ── derive the movie overlay for the current frame ──
187 let cloud: Points = data ? data.samples : { x: [], y: [] };
188 let orbit: OrbitPath | null = null;
189 let start: Pt | null = null;
190 let lead: Pt | null = null;
191 let chainPts: Points | null = null;
192 let selected: Pt | null = null;
193 if (movie && movieData) {
194 const step = movieData[movie.si];
195 const count = step.px.length;
196 const k = Math.min(movie.k, count);
197 orbit = { x: step.px.slice(0, k), y: step.py.slice(0, k), seg: step.seg.slice(0, k) };
198 start = { x: step.startX, y: step.startY };
199 if (k > 0) lead = { x: step.px[k - 1], y: step.py[k - 1] };
200 const cx: number[] = [];
201 const cy: number[] = [];
202 for (let j = 0; j < movie.si; j++) {
203 cx.push(movieData[j].selX);
204 cy.push(movieData[j].selY);
205 }
206 chainPts = { x: cx, y: cy };
207 if (movie.k >= count) selected = { x: step.selX, y: step.selY };
208 cloud = data ? data.samples : { x: [], y: [] };
209 }
211 const controlsDisabled = !data || busy || !!movie;
212 const status = busy
213 ? "sampling…"
214 : movie && movieData
215 ? `movie · transition ${movie.si + 1}/${movieData.length}`
216 : `${data ? data.n.toLocaleString() : "—"} samples`;
218 return (
219 <div style={rootStyle}>
220 {data ? (
221 <DensityView
222 density={data.density}
223 samples={cloud}
224 orbit={orbit}
225 start={start}
226 lead={lead}
227 chainPts={chainPts}
228 selected={selected}
229 />
230 ) : (
231 <div style={waitingStyle}>Waiting for samples from the script…</div>
232 )}
234 <div style={panelStyle}>
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 235 <div style={{ fontWeight: 600, marginBottom: 6 }}>WALNUTS</div>
237 <label style={labelStyle}>
238 Target
239 <select
240 value={target}
241 disabled={controlsDisabled}
242 onChange={e => changeTarget(e.target.value)}
243 style={selectStyle}
244 >
245 {TARGETS.map(t => (
246 <option key={t.value} value={t.value}>
247 {t.label}
248 </option>
249 ))}
250 </select>
251 </label>
253 <label style={labelStyle}>
254 Samples: <b>{n.toLocaleString()}</b>
255 <input
256 type="range"
257 min={0}
258 max={SAMPLE_CHOICES.length - 1}
259 step={1}
260 value={Math.max(0, SAMPLE_CHOICES.indexOf(n))}
261 disabled={controlsDisabled}
262 onChange={e => setN(SAMPLE_CHOICES[Number(e.target.value)])}
263 onPointerUp={e =>
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 264 resample(SAMPLE_CHOICES[Number(e.currentTarget.value)], h, delta)
266 style={sliderStyle}
267 />
268 </label>
270 <label style={labelStyle}>
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 271 Step h: <b>{h.toFixed(2)}</b>
273 type="range"
275 max={H_MAX}
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 278 disabled={controlsDisabled}
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 279 onChange={e => setH(Number(e.target.value))}
280 onPointerUp={e => resample(n, Number(e.currentTarget.value), delta)}
282 />
283 </label>
285 <label style={labelStyle}>
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 286 Energy tol δ: <b>{delta.toFixed(2)}</b>
288 type="range"
290 max={DELTA_MAX}
291 step={0.05}
292 value={delta}
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 293 disabled={controlsDisabled}
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 294 onChange={e => setDelta(Number(e.target.value))}
295 onPointerUp={e => resample(n, h, Number(e.currentTarget.value))}
297 />
298 </label>
300 <div style={{ display: "flex", gap: 6, marginTop: 8 }}>
301 <button
302 style={btnStyle}
303 disabled={controlsDisabled}
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 304 onClick={() => resample(n, h, delta)}
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 305 title="Draw a fresh chain with these settings"
306 >
307 Resample
308 </button>
309 <button
310 style={btnStyle}
311 disabled={!data || busy}
312 onClick={playMovie}
313 title="Animate WALNUTS building orbits step by step"
314 >
315 {movie ? "■ Stop" : "▶ Movie"}
316 </button>
317 </div>
319 <div style={{ fontSize: 10, color: "#64748b", marginTop: 8 }}>
320 {status}
321 </div>
322 </div>
323 </div>
324 );
325}
327const rootStyle: CSSProperties = {
328 position: "absolute",
329 inset: 0,
330 overflow: "hidden",
331 background: "#ffffff",
332 fontFamily: "system-ui, -apple-system, Arial, sans-serif",
333};
335const waitingStyle: CSSProperties = {
336 position: "absolute",
337 inset: 0,
338 display: "flex",
339 alignItems: "center",
340 justifyContent: "center",
341 color: "#94a3b8",
342};
344const panelStyle: CSSProperties = {
345 position: "absolute",
346 top: 8,
347 left: 8,
348 width: 168,
349 padding: "8px 10px",
350 background: "rgba(255,255,255,0.92)",
351 border: "1px solid #e2e8f0",
352 borderRadius: 6,
353 boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
354 color: "#0f172a",
355};
357const labelStyle: CSSProperties = {
358 display: "block",
359 fontSize: 11,
360 marginTop: 6,
361};
363const sliderStyle: CSSProperties = {
364 width: "100%",
365 marginTop: 2,
366};
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 368const selectStyle: CSSProperties = {
369 width: "100%",
370 marginTop: 2,
371 fontSize: 11,
372 padding: "2px 4px",
373};
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 375const btnStyle: CSSProperties = {
376 flex: 1,
377 padding: "4px 6px",
378 fontSize: 11,
379 whiteSpace: "nowrap",
380 cursor: "pointer",
381 background: "#f8fafc",
382 border: "1px solid #cbd5e1",
383 borderRadius: 5,
384 color: "#0f172a",
385};