Use Nawaf Bou-Rabee's reference WALNUTS implementation
Replace the JS-port-based sampler with Nawaf's reference MATLAB code (walnuts.m
+ extend_orbit_*/micro/leapfrog/u_turn/sub_u_turn/p_micro/pmf_p_micro), used as
provided; added only logsumexp/log_softmax (standard helpers it calls) and an
orbit output O on walnuts for the movie. The sampler loops his per-transition
walnuts(...), passing our 2D targets (banana/Gaussian/correlated/donut) as
function handles. Controls remapped to his parameters: step h and energy
tolerance delta. Paper: Bou-Rabee, Carpenter, Kleppe, Liu (arXiv:2506.18746).
18 changed files+378−402
.gitignoremodified+1−0View file
@@ -2,3 +2,4 @@
22 /_site/
33 /dist/
44 node_modules/
5+missing_files/
.numblignoremodified+1−0View file
@@ -10,3 +10,4 @@ app/package.json
1010 app/package-lock.json
1111 app/tsconfig.json
1212 app/vite.config.ts
13+missing_files
CLAUDE.mdmodified+31−27View file
@@ -19,42 +19,46 @@ function file). Run `walnuts_demo.m`, not the sampler directly.
1919 - **MATLAB → figure:** `uihtml(..., 'Data', struct)` / `sendEventToHTMLSource(src, name, data)`.
2020 - **Figure → MATLAB:** `sendToMATLAB(name, value)` (see `app/src/bridge.ts`), received via `HTMLEventReceivedFcn`.
2121
22-The script sends the density grid + samples once, then handles two requests
22+The script sends the density grid + samples once, then handles three requests
2323 (`walnuts_sampler.m` `on_event`):
2424
25-- `resample` `{n, dt, maxError}` → fresh chain → `samples` event. (Target is
26- fixed, so the density is not re-sent.)
27-- `movie` `{dt, maxError}` → `walnuts(..., record=true)` for a few transitions →
28- `movie` event carrying their orbit trajectories.
25+- `resample` `{n, h, delta, target}` → fresh chain → `samples` event.
26+- `setTarget` `{target, …}` → new target → full `data` event (density + samples).
27+- `movie` `{h, delta, target}` → record a few transitions' orbits → `movie` event.
2928
30-The figure's controls (samples / Δt / max error sliders, Resample) drive
31-`resample`; the **▶ Movie** button drives `movie` and animates the returned
32-orbits (`App.tsx` + `DensityView.tsx`).
29+The figure's controls (target dropdown, samples / step `h` / energy-tol `δ`
30+sliders, Resample) drive `resample`/`setTarget`; the **▶ Movie** button drives
31+`movie` and animates the returned orbits (`App.tsx` + `DensityView.tsx`).
3332
3433 ### Trajectory recording
3534
36-`walnuts.m` records the orbit path only when called with `record=true`, via
37-globals (`WREC_*`): `macro_step` appends each *committed* macro-step's leapfrog
38-path (recomputed by `leapfrog_capture`), tagged with a segment id so the figure
39-breaks the polyline at orbit direction flips. Plain sampling (`record=false`)
40-takes the fast path with no recording overhead.
35+`walnuts.m` exposes the orbit it built as a third output `O` (a cell of
36+`{theta, rho}` states) — purely for the movie; it doesn't affect the algorithm.
37+`record_movie` (in `walnuts_sampler.m`) runs a few transitions, pulls each
38+orbit's positions, and packs `px/py/seg` + the start and selected draw.
4139
4240 ## The algorithm
4341
44-`helpers/walnuts.m` is a direct port of Brian Ward's `algorithms/WALNUTS.js`
45-from chi-feng/mcmc-demo (based on Bob Carpenter's C++,
46-flatironinstitute/walnuts; paper arXiv:2506.18746). Spans carry both orbit
47-endpoints; `macro_step` does the within-orbit step halving; `build_span`
48-recurses (NUTS doubling); selection is Barker within sub-orbits, Metropolis at
49-the top. See README for full credits.
42+`helpers/walnuts.m` and its building blocks (`extend_orbit_forward/backward`,
43+`micro`, `leapfrog`, `u_turn`/`sub_u_turn`, `p_micro`/`pmf_p_micro`) are **Nawaf
44+Bou-Rabee's reference MATLAB implementation**, used as provided (paper
45+arXiv:2506.18746). It's the orbit-based formulation: sample momentum, grow the
46+orbit by doubling forward/backward, refine the leapfrog step within each step so
47+the energy variation ≤ `δ`, terminate on a U-turn, and pick a draw from the
48+orbit by its `log_softmax` weights. We added only `logsumexp`/`log_softmax`
49+(standard helpers it calls) and a third output `O` on `walnuts` for the movie.
50+`walnuts` takes the density as **function handles** (`@log_density`,
51+`@grad_log_density`), which dispatch on the global `WTARGET` — keep the algorithm
52+files target-agnostic. The verbatim originals + paper are in `missing_files/`
53+(git-ignored).
5054
5155 ## Performance
5256
53-Unlike the hitandrun kernel, WALNUTS uses **recursion + structs**, so it runs in
54-the numbl **interpreter** (it does not JS-JIT like a flat numeric loop). It's
55-still fine — banana orbits are shallow, ~1–2 ms/transition — so keep `N` modest
56-(the demo uses 1000 + 200 burn-in). Don't add a `%!numbl:assert_jit` guard here;
57-it would (correctly) fail.
57+WALNUTS here uses **function-handle args, 2-D cell arrays, and recursion**, so it
58+runs in the numbl **interpreter** (it does not JS-JIT). It's still fine —
59+~1–5 ms/transition on these 2-D targets — but heavier than a flat loop, so keep
60+`N` modest (`SAMPLE_CHOICES` tops out at 3000; the demo uses 1000 + 200 burn-in).
61+Don't add a `%!numbl:assert_jit` guard; it would (correctly) fail.
5862
5963 ## Key files
6064
@@ -63,9 +67,9 @@ it would (correctly) fail.
6367 | `app/src/App.tsx` | Reads data, hosts the info panel |
6468 | `app/src/render/DensityView.tsx` | Canvas: density heatmap + sample scatter |
6569 | `app/src/bridge.ts` | `onData` / `onHostEvent` / `sendToMATLAB` (generic) |
66-| `helpers/walnuts.m` | The WALNUTS sampler |
67-| `helpers/log_density.m` / `grad_log_density.m` | Banana target |
68-| `walnuts_sampler.m` | Runs the chain, builds the density grid, opens the figure |
70+| `helpers/walnuts.m` + building blocks | Nawaf's reference WALNUTS sampler |
71+| `helpers/log_density.m` / `grad_log_density.m` | 2D targets (banana/gaussian/correlated/donut), dispatched on global `WTARGET` |
72+| `walnuts_sampler.m` | Loops `walnuts(...)`, builds the density grid, opens the figure, handles events |
6973 | `walnuts_demo.m` | Driver: addpath + seed + call the sampler |
7074
7175 ## Local iteration
README.mdmodified+18−19View file
@@ -10,16 +10,17 @@ with the samples scattered on top. Switch between several targets (banana,
1010 Gaussian, correlated Gaussian, donut) to see the sampler handle different
1111 geometries.
1212
13-WALNUTS is a NUTS variant that adapts the leapfrog step size *within* each
14-orbit: each macro-step halves the step (and doubles the count) until the energy
15-error is within tolerance, then grows the orbit by NUTS-style doubling until a
16-U-turn.
13+WALNUTS is a NUTS variant that grows a Hamiltonian orbit by NUTS-style doubling
14+until a U-turn, but *within* each step it refines the leapfrog step size —
15+halving the step (and doubling the count) until the energy variation over the
16+step is within a tolerance δ — then picks a draw from the orbit.
1717
1818 Controls:
1919
2020 - **Target** — switch the distribution (banana / Gaussian / correlated / donut).
21-- **Samples** / **Leapfrog Δt** / **Max error** — change a setting to re-run the
22- sampler with it (lower Δt or max error → more, smaller leapfrog steps).
21+- **Samples** / **Step h** / **Energy tol δ** — change a setting to re-run the
22+ sampler with it (a larger base step `h` or smaller `δ` triggers more
23+ within-orbit refinement).
2324 - **Resample** — a fresh chain with the current settings.
2425 - **▶ Movie** — animate WALNUTS building orbits step by step: each transition
2526 traces its leapfrog path (amber), then circles the selected draw, which joins
@@ -28,26 +29,24 @@ Controls:
2829 ## How it works
2930
3031 - [`walnuts_demo.m`](walnuts_demo.m) — driver: `addpath('helpers')`, seed, call the sampler.
31-- [`walnuts_sampler.m`](walnuts_sampler.m) — runs the chain, evaluates the target on a grid, opens the figure.
32-- `helpers/` — [`walnuts.m`](helpers/walnuts.m) (the algorithm) and the target [`log_density.m`](helpers/log_density.m) / [`grad_log_density.m`](helpers/grad_log_density.m).
32+- [`walnuts_sampler.m`](walnuts_sampler.m) — loops the sampler (one transition per draw, passing the target as function handles), evaluates the target on a grid, opens the figure.
33+- `helpers/` — Nawaf Bou-Rabee's reference WALNUTS ([`walnuts.m`](helpers/walnuts.m) + `extend_orbit_*`, `micro`, `leapfrog`, `u_turn`/`sub_u_turn`, `p_micro`/`pmf_p_micro`), and the targets [`log_density.m`](helpers/log_density.m) / [`grad_log_density.m`](helpers/grad_log_density.m).
3334 - `app/` — a single-file React app that draws the density heatmap + samples on a canvas.
3435
3536 The script sends the density grid + samples via `uihtml(..., 'Data', ...)`. The
3637 controls call back: `resample` re-runs the chain, `setTarget` switches the
37-target (and rebuilds the density), and `movie` records a few transitions' orbit
38-trajectories (`walnuts(..., record=true)`); the script replies with
39-`sendEventToHTMLSource`.
38+target (and rebuilds the density), and `movie` records a few transitions'
39+orbits; the script replies with `sendEventToHTMLSource`.
4040
4141 ## Credits
4242
43-- **Algorithm:** N. Bou-Rabee, B. Carpenter, T. S. Kleppe, and S. Liu, *The
44- within-orbit adaptive leapfrog no-U-turn sampler*,
45- [arXiv:2506.18746](https://arxiv.org/abs/2506.18746) (2025). Reference C++ at
46- [flatironinstitute/walnuts](https://github.com/flatironinstitute/walnuts)
47- (Bob Carpenter).
48-- **This MATLAB port** follows Brian Ward's JavaScript implementation in
49- [WardBrian/mcmc-demo](https://github.com/WardBrian/mcmc-demo) (`algorithms/WALNUTS.js`).
50-- **Demo framework & banana target** from Chi Feng's
43+- **Algorithm & reference MATLAB implementation:** N. Bou-Rabee, B. Carpenter,
44+ T. S. Kleppe, and S. Liu, *The within-orbit adaptive leapfrog no-U-turn
45+ sampler*, [arXiv:2506.18746](https://arxiv.org/abs/2506.18746) (2025). The
46+ `helpers/` sampler (`walnuts.m` + its building blocks) is Nawaf Bou-Rabee's
47+ reference MATLAB code, used here as provided. Reference C++ at
48+ [flatironinstitute/walnuts](https://github.com/flatironinstitute/walnuts).
49+- **Target distributions** (banana, donut, …) from Chi Feng's
5150 [mcmc-demo](https://github.com/chi-feng/mcmc-demo).
5251
5352 ## Deploy
app/src/App.tsxmodified+41−36View file
@@ -9,14 +9,15 @@ import {
99 import { onData, onHostEvent, sendToMATLAB } from "./bridge.js";
1010
1111 /** Payload from the numbl script: the target density (for the heatmap) plus the
12- * WALNUTS samples. Mirrors what walnuts_sampler.m sends. */
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. */
1314 interface WalnutsData {
1415 type: "walnuts";
1516 density: DensityGrid;
1617 samples: Points;
1718 n: number;
18- dt: number;
19- maxError: number;
19+ h: number;
20+ delta: number;
2021 target: string;
2122 }
2223
@@ -41,11 +42,11 @@ interface SamplesEvent {
4142 x: number[];
4243 y: number[];
4344 n: number;
44- dt: number;
45- maxError: number;
45+ h: number;
46+ delta: number;
4647 }
4748
48-/** One recorded transition's orbit (from `walnuts(..., record=true)`). */
49+/** One recorded transition's orbit (from `walnuts(...)`'s orbit output). */
4950 interface MovieStep {
5051 px: number[];
5152 py: number[];
@@ -56,14 +57,18 @@ interface MovieStep {
5657 selY: number;
5758 }
5859
60+// WALNUTS' orbit-based transition is heavier than a plain MH step, so keep the
61+// sample counts modest.
5962 const SAMPLE_CHOICES = [100, 300, 1000, 3000];
6063 const DEFAULT_N = 1000;
61-const DT_MIN = 0.05;
62-const DT_MAX = 1.2;
63-const ERR_MIN = 0.1;
64-const ERR_MAX = 4;
65-
66-// Movie pacing: reveal a couple of leapfrog points per tick, then linger on the
64+const DEFAULT_H = 0.8;
65+const DEFAULT_DELTA = Math.log(1 / 0.66); // ≈ 0.42 (Nawaf's amin = 0.66)
66+const H_MIN = 0.1;
67+const H_MAX = 2.0;
68+const DELTA_MIN = 0.1;
69+const DELTA_MAX = 2.0;
70+
71+// Movie pacing: reveal a couple of orbit points per tick, then linger on the
6772 // selected draw before the next transition.
6873 const MOVIE_TICK_MS = 55;
6974 const MOVIE_REVEAL = 2;
@@ -72,8 +77,8 @@ const MOVIE_HOLD = 14; // extra k-units to hold the selected point
7277 export function App() {
7378 const [data, setData] = useState<WalnutsData | null>(null);
7479 const [n, setN] = useState(DEFAULT_N);
75- const [dt, setDt] = useState(0.4);
76- const [maxError, setMaxError] = useState(0.8);
80+ const [h, setH] = useState(DEFAULT_H);
81+ const [delta, setDelta] = useState(DEFAULT_DELTA);
7782 const [target, setTargetState] = useState("banana");
7883 const [busy, setBusy] = useState(false);
7984 const [movieData, setMovieData] = useState<MovieStep[] | null>(null);
@@ -83,8 +88,8 @@ export function App() {
8388 const applyData = (d: WalnutsData) => {
8489 setData(d);
8590 setN(d.n);
86- setDt(d.dt);
87- setMaxError(d.maxError);
91+ setH(d.h);
92+ setDelta(d.delta);
8893 setTargetState(d.target);
8994 };
9095
@@ -107,7 +112,7 @@ export function App() {
107112 if (!s || !Array.isArray(s.x)) return;
108113 setData(prev =>
109114 prev
110- ? { ...prev, samples: { x: s.x, y: s.y }, n: s.n, dt: s.dt, maxError: s.maxError }
115+ ? { ...prev, samples: { x: s.x, y: s.y }, n: s.n, h: s.h, delta: s.delta }
111116 : prev
112117 );
113118 setBusy(false);
@@ -151,11 +156,11 @@ export function App() {
151156 setMovieData(null);
152157 };
153158
154- const resample = (count: number, step: number, err: number) => {
159+ const resample = (count: number, hVal: number, deltaVal: number) => {
155160 if (!data || busy) return;
156161 stopMovie();
157162 setBusy(true);
158- sendToMATLAB("resample", { n: count, dt: step, maxError: err, target });
163+ sendToMATLAB("resample", { n: count, h: hVal, delta: deltaVal, target });
159164 };
160165
161166 // Switch the target: the script rebuilds the density + draws and replies with
@@ -165,7 +170,7 @@ export function App() {
165170 if (!data) return;
166171 stopMovie();
167172 setBusy(true);
168- sendToMATLAB("setTarget", { target: value, n, dt, maxError });
173+ sendToMATLAB("setTarget", { target: value, n, h, delta });
169174 };
170175
171176 const playMovie = () => {
@@ -175,7 +180,7 @@ export function App() {
175180 }
176181 if (!data || busy) return;
177182 setBusy(true); // until the trajectory arrives
178- sendToMATLAB("movie", { dt, maxError, target });
183+ sendToMATLAB("movie", { h, delta, target });
179184 };
180185
181186 // ── derive the movie overlay for the current frame ──
@@ -256,38 +261,38 @@ export function App() {
256261 disabled={controlsDisabled}
257262 onChange={e => setN(SAMPLE_CHOICES[Number(e.target.value)])}
258263 onPointerUp={e =>
259- resample(SAMPLE_CHOICES[Number(e.currentTarget.value)], dt, maxError)
264+ resample(SAMPLE_CHOICES[Number(e.currentTarget.value)], h, delta)
260265 }
261266 style={sliderStyle}
262267 />
263268 </label>
264269
265270 <label style={labelStyle}>
266- Leapfrog Δt: <b>{dt.toFixed(2)}</b>
271+ Step h: <b>{h.toFixed(2)}</b>
267272 <input
268273 type="range"
269- min={DT_MIN}
270- max={DT_MAX}
274+ min={H_MIN}
275+ max={H_MAX}
271276 step={0.05}
272- value={dt}
277+ value={h}
273278 disabled={controlsDisabled}
274- onChange={e => setDt(Number(e.target.value))}
275- onPointerUp={e => resample(n, Number(e.currentTarget.value), maxError)}
279+ onChange={e => setH(Number(e.target.value))}
280+ onPointerUp={e => resample(n, Number(e.currentTarget.value), delta)}
276281 style={sliderStyle}
277282 />
278283 </label>
279284
280285 <label style={labelStyle}>
281- Max error: <b>{maxError.toFixed(2)}</b>
286+ Energy tol δ: <b>{delta.toFixed(2)}</b>
282287 <input
283288 type="range"
284- min={ERR_MIN}
285- max={ERR_MAX}
286- step={0.1}
287- value={maxError}
289+ min={DELTA_MIN}
290+ max={DELTA_MAX}
291+ step={0.05}
292+ value={delta}
288293 disabled={controlsDisabled}
289- onChange={e => setMaxError(Number(e.target.value))}
290- onPointerUp={e => resample(n, dt, Number(e.currentTarget.value))}
294+ onChange={e => setDelta(Number(e.target.value))}
295+ onPointerUp={e => resample(n, h, Number(e.currentTarget.value))}
291296 style={sliderStyle}
292297 />
293298 </label>
@@ -296,7 +301,7 @@ export function App() {
296301 <button
297302 style={btnStyle}
298303 disabled={controlsDisabled}
299- onClick={() => resample(n, dt, maxError)}
304+ onClick={() => resample(n, h, delta)}
300305 title="Draw a fresh chain with these settings"
301306 >
302307 Resample
helpers/extend_orbit_backward.madded+26−0View file
@@ -0,0 +1,26 @@
1+function [O, logW] = extend_orbit_backward(logmu, gradlogmu, theta_a, rho_a, logw_a, h, delta, L)
2+ O = {};
3+ logW = [];
4+ theta = theta_a;
5+ rho = rho_a;
6+ logw = logw_a;
7+
8+ for iter = 1:L
9+ ell_b = micro(logmu, gradlogmu, theta, -rho, h, delta);
10+ ell = p_micro(ell_b);
11+ [theta_1, rho_1] = leapfrog(logmu, gradlogmu, theta, -rho, h * (2^(-ell)), 2^ell);
12+ rho_1 = -rho_1;
13+ ell_f = micro(logmu, gradlogmu, theta_1, rho_1, h, delta);
14+
15+ if pmf_p_micro(ell, ell_f) == 0
16+ logw = -Inf;
17+ else
18+ logw = logw + logmu(theta_1) - 0.5 * norm(rho_1)^2 - logmu(theta) + 0.5 * norm(rho)^2 + log(pmf_p_micro(ell, ell_f)) - log(pmf_p_micro(ell, ell_b));
19+ end
20+
21+ theta = theta_1;
22+ rho = rho_1;
23+ O = [{theta, rho}; O];
24+ logW = [logw; logW];
25+ end
26+end
helpers/extend_orbit_forward.madded+25−0View file
@@ -0,0 +1,25 @@
1+function [O, logW] = extend_orbit_forward(logmu, gradlogmu, theta_b, rho_b, logw_b, h, delta, L)
2+ O = {};
3+ logW = [];
4+ theta = theta_b;
5+ rho = rho_b;
6+ logw = logw_b;
7+
8+ for iter = 1:L
9+ ell_f = micro(logmu, gradlogmu, theta, rho, h, delta);
10+ ell = p_micro(ell_f);
11+ [theta_1, rho_1] = leapfrog(logmu, gradlogmu, theta, rho, h * (2^(-ell)), 2^ell);
12+ ell_b = micro(logmu, gradlogmu, theta_1, -rho_1, h, delta);
13+
14+ if pmf_p_micro(ell, ell_b) == 0
15+ logw = -Inf;
16+ else
17+ logw = logw + logmu(theta_1) - 0.5 * norm(rho_1)^2 - logmu(theta) + 0.5 * norm(rho)^2 + log(pmf_p_micro(ell, ell_b))- log(pmf_p_micro(ell, ell_f));
18+ end
19+
20+ theta = theta_1;
21+ rho = rho_1;
22+ O = [O; {theta, rho}];
23+ logW = [logW; logw];
24+ end
25+end
helpers/leapfrog.madded+17−0View file
@@ -0,0 +1,17 @@
1+function [theta, rho] = leapfrog(logmu, gradlogmu, theta, rho, h, ell)
2+ h2=h*h;
3+ for k = 1:ell
4+ g=gradlogmu(theta);
5+ rho_half = rho + (h / 2) * g;
6+ theta = theta + h * rho_half;
7+ g=gradlogmu(theta);
8+ rho = rho_half + (h / 2) * g;
9+ end
10+ % for k = 1:ell
11+ % g=gradlogmu(theta);
12+ % rho_half = rho + (h / 2) * g./(1+h2*abs(g));
13+ % theta = theta + h * rho_half;
14+ % g=gradlogmu(theta);
15+ % rho = rho_half + (h / 2) * g./(1+h2*abs(g));
16+ % end
17+end
helpers/log_softmax.madded+4−0View file
@@ -0,0 +1,4 @@
1+function s = log_softmax(x)
2+%LOG_SOFTMAX x - logsumexp(x); exp(log_softmax(x)) is a probability vector.
3+s = x - logsumexp(x);
4+end
helpers/logsumexp.madded+9−0View file
@@ -0,0 +1,9 @@
1+function s = logsumexp(x)
2+%LOGSUMEXP Numerically stable log(sum(exp(x))) over all elements of x.
3+m = max(x(:));
4+if ~isfinite(m)
5+ s = m; % all -Inf -> -Inf (and +Inf -> +Inf); avoids Inf-Inf = NaN
6+ return
7+end
8+s = m + log(sum(exp(x(:) - m)));
9+end
helpers/micro.madded+35−0View file
@@ -0,0 +1,35 @@
1+function ell = micro(logmu, gradlogmu, theta_0, rho_0, h_0, delta)
2+ ell = 0;
3+
4+ max_ell = 8; % default cap
5+
6+ while ell <= max_ell
7+ % reset initial state for this candidate ell
8+ theta = theta_0;
9+ rho = rho_0;
10+
11+ R = 2^(ell); % number of micro-steps
12+ h = h_0 / R; % micro step size
13+
14+ % initial energy
15+ H = -logmu(theta) + 0.5 * norm(rho)^2;
16+ H_max = H;
17+ H_min = H;
18+
19+ % integrate using *your* leapfrog, one step at a time
20+ for j = 1:R
21+ [theta, rho] = leapfrog(logmu, gradlogmu, theta, rho, h, 1);
22+ H = -logmu(theta) + 0.5 * norm(rho)^2;
23+ H_max = max(H_max, H);
24+ H_min = min(H_min, H);
25+ end
26+
27+ % accept this ell if energy variation is small enough
28+ if H_max - H_min <= delta
29+ return;
30+ end
31+
32+ % otherwise, refine (halve h, double R)
33+ ell = ell + 1;
34+ end
35+end
helpers/p_micro.madded+7−0View file
@@ -0,0 +1,7 @@
1+function j = p_micro(i)
2+ weights = [2/3, 1/3, 0];
3+ choices = [i, i+1, i+2];
4+ cdf = cumsum(weights);
5+ r = rand() * cdf(end);
6+ j = choices(find(cdf >= r, 1, 'first'));
7+end
helpers/pmf_p_micro.madded+8−0View file
@@ -0,0 +1,8 @@
1+function prob = pmf_p_micro(j, i)
2+ weights = [2/3, 1/3, 0];
3+ if ismember(j, [i, i+1, i+2])
4+ prob = weights(j - i + 1);
5+ else
6+ prob = 0;
7+ end
8+end
helpers/sub_u_turn.madded+7−0View file
@@ -0,0 +1,7 @@
1+function result = sub_u_turn(O)
2+ if length(O) < 2
3+ result = false;
4+ return;
5+ end
6+ result = u_turn(O) || sub_u_turn(O(1:floor(end/2))) || sub_u_turn(O(floor(end/2)+1:end));
7+end
helpers/u_turn.madded+11−0View file
@@ -0,0 +1,11 @@
1+function result = u_turn(O)
2+ theta_left = O{1,1};
3+ rho_left = O{1,2};
4+ theta_right = O{end,1};
5+ rho_right = O{end,2};
6+
7+ dot1 = dot(rho_right, (theta_right - theta_left));
8+ dot2 = dot(rho_left, (theta_right - theta_left));
9+
10+ result = (dot1 < 0) || (dot2 < 0);
11+end
helpers/walnuts.mmodified+52−273View file
@@ -1,275 +1,54 @@
1-function [chain, traj] = walnuts(theta0, n_samples, dt, max_error, record)
2-%WALNUTS Within-orbit adaptive No-U-Turn Sampler.
3-% CHAIN = WALNUTS(THETA0, N_SAMPLES, DT, MAX_ERROR) runs N_SAMPLES WALNUTS
4-% transitions starting from column vector THETA0 and returns CHAIN, a
5-% dim-by-N_SAMPLES matrix of draws. DT is the base leapfrog step and MAX_ERROR
6-% the per-macro-step energy-error tolerance that drives the within-orbit step
7-% halving.
8-%
9-% [CHAIN, TRAJ] = WALNUTS(..., true) additionally records, per transition, the
10-% leapfrog path of the orbit it builds (for the figure's step-by-step movie).
11-% TRAJ{i} has fields px/py (orbit positions), seg (per-point macro-step id, so
12-% the path breaks at direction flips), startX/startY and selX/selY. Recording
13-% adds overhead, so leave it off for plain sampling.
14-%
15-% This is a direct port of Brian Ward's JavaScript implementation in
16-% chi-feng/mcmc-demo (algorithms/WALNUTS.js), itself based on Bob Carpenter's
17-% C++ (flatironinstitute/walnuts) for the paper arXiv:2506.18746. The target
18-% density is supplied by log_density.m / grad_log_density.m on the path.
19-%
20-% A Span carries the two orbit endpoints (backward "_bk", forward "_fw"),
21-% each with position/momentum/gradient/joint-logp, plus a selected draw
22-% theta_select and the log of the summed orbit weight (logp).
23-
24-if nargin < 5 || isempty(record)
25- record = false;
26-end
27-% Trajectory recorder (used only when record is true). macro_step appends each
28-% committed leapfrog path to these as it grows the orbit.
29-global WREC_ON WREC_X WREC_Y WREC_SEG WREC_SEGID
30-WREC_ON = record;
31-
32-dim = numel(theta0);
33-chain = zeros(dim, n_samples);
34-traj = {};
35-theta = theta0;
36-for i = 1:n_samples
37- if record
38- start_pt = theta;
39- WREC_X = [];
40- WREC_Y = [];
41- WREC_SEG = [];
42- WREC_SEGID = 0;
43- end
44- theta = transition(theta, dt, max_error);
45- chain(:, i) = theta;
46- if record
47- traj{i} = struct('px', WREC_X, 'py', WREC_Y, 'seg', WREC_SEG, ...
48- 'startX', start_pt(1), 'startY', start_pt(2), ...
49- 'selX', theta(1), 'selY', theta(2));
50- end
51-end
52-WREC_ON = false;
53-end
54-
55-function theta_select = transition(theta, dt, max_error)
56-% One WALNUTS transition: sample momentum, grow the orbit by repeated doubling
57-% in random directions until a U-turn (or a failed step / max depth), choosing
58-% the draw by a Metropolis update against the growing orbit.
59-rho = randn(numel(theta), 1);
60-grad = grad_log_density(theta);
61-logp = log_density(theta) - sum(rho.^2) / 2;
62-span_accum = make_leaf_span(theta, rho, grad, logp);
63-
64-for depth = 0:11
65- if rand < 0.5
66- direction = -1;
67- else
68- direction = 1;
69- end
70- [ok, next_span] = build_span(span_accum, direction, depth, dt, max_error);
71- if ~ok
72- break
73- end
74- combined_uturn = uturn(span_accum, next_span, direction);
75- % Top-level selection is a Metropolis update (use_barker = false).
76- span_accum = combine(span_accum, next_span, false, direction);
77- if combined_uturn
78- break
79- end
80-end
81-theta_select = span_accum.theta_select;
82-end
83-
84-function [ok, span] = build_span(span_in, direction, depth, dt, max_error)
85-% Recursively build a balanced orbit of 2^depth macro-steps. Returns ok=false
86-% if any macro-step fails or a sub-orbit U-turns.
87-if depth == 0
88- [ok, span] = build_leaf(span_in, direction, dt, max_error);
89- return
90-end
91-[ok, left] = build_span(span_in, direction, depth - 1, dt, max_error);
92-if ~ok
93- span = left;
94- return
95-end
96-[ok, right] = build_span(left, direction, depth - 1, dt, max_error);
97-if ~ok
98- span = right;
99- return
100-end
101-if uturn(left, right, direction)
102- ok = false;
103- span = left;
104- return
105-end
106-% Sub-orbit selection is a Barker update (use_barker = true).
107-span = combine(left, right, true, direction);
108-ok = true;
109-end
110-
111-function [ok, span] = build_leaf(span_in, direction, dt, max_error)
112-[success, theta_next, rho_next, grad_next, logp_next] = ...
113- macro_step(span_in, direction, dt, max_error);
114-if ~success
115- ok = false;
116- span = span_in;
117- return
118-end
119-span = make_leaf_span(theta_next, rho_next, grad_next, logp_next);
120-ok = true;
121-end
122-
123-function [success, theta_next, rho_next, grad_next, logp_next] = ...
124- macro_step(span, direction, dt, max_error)
125-% Take one macro-step off the orbit's leading endpoint, adaptively halving the
126-% leapfrog step (and doubling the count) until the energy error is within
127-% tolerance, then check the choice is reversible.
128-global WREC_ON WREC_X WREC_Y WREC_SEG WREC_SEGID
129-if direction == 1
130- theta = span.theta_fw;
131- rho = span.rho_fw;
132- grad = span.grad_fw;
133- logp = span.logp_fw;
134- step = dt;
135-else
136- theta = span.theta_bk;
137- rho = span.rho_bk;
138- grad = span.grad_bk;
139- logp = span.logp_bk;
140- step = -dt;
141-end
142-
143-num_steps = 1;
144-for halvings = 0:9
145- theta_next = theta;
146- rho_next = rho;
147- grad_next = grad;
148- [theta_next, rho_next, grad_next] = leapfrog(theta_next, rho_next, grad_next, step, num_steps);
149- logp_next = log_density(theta_next) - sum(rho_next.^2) / 2;
150- if abs(logp - logp_next) <= max_error
151- success = reversible(step, num_steps, theta_next, rho_next, grad_next, logp_next, max_error);
152- if success && WREC_ON
153- % Record this committed macro-step's leapfrog path for the movie.
154- path = leapfrog_capture(theta, rho, grad, step, num_steps);
155- WREC_SEGID = WREC_SEGID + 1;
156- WREC_X = [WREC_X, theta(1), path(1, :)];
157- WREC_Y = [WREC_Y, theta(2), path(2, :)];
158- WREC_SEG = [WREC_SEG, WREC_SEGID * ones(1, num_steps + 1)];
1+function [theta_tilde, T, O] = walnuts(logmu, gradlogmu, theta, h, i_max, delta)
2+%WALNUTS One WALNUTS transition (within-orbit adaptive leapfrog NUTS).
3+% Reference MATLAB implementation aligned with the paper
4+% arXiv:2506.18746 (Bou-Rabee, Carpenter, Kleppe, Liu). LOGMU / GRADLOGMU are
5+% function handles for the target's log density and gradient. Returns the
6+% selected draw THETA_TILDE, the orbit length T, and the orbit O (a cell of
7+% {theta, rho} states) — O is exposed only for the figure's movie and does not
8+% affect the algorithm.
9+ d = length(theta);
10+ rho = randn(d, 1);
11+ theta_tilde = theta;
12+ rho_tilde = rho;
13+ logw_0 = logmu(theta) - 0.5 * norm(rho)^2;
14+
15+ O = {theta, rho};
16+ logW = logw_0;
17+ B = randi([0, 1], i_max, 1);
18+
19+ for i = 1:i_max
20+ O_old = O;
21+ logW_old = logW;
22+
23+ if B(i) == 1
24+ [O_ext, logW_ext] = extend_orbit_forward(logmu, gradlogmu, O{end,1}, O{end,2}, logW(end), h, delta, 2^(i-1));
25+ O = [O; O_ext];
26+ logW = [logW; logW_ext];
27+ else
28+ [O_ext, logW_ext] = extend_orbit_backward(logmu, gradlogmu, O{1,1}, O{1,2}, logW(1), h, delta, 2^(i-1));
29+ O = [O_ext; O];
30+ logW = [logW_ext; logW];
31+ end
32+
33+ if sub_u_turn(O_ext)
34+ break;
35+ end
36+
37+ logu = log(rand);
38+
39+ if logu <= logsumexp(logW_ext) - logsumexp(logW_old)
40+ weights = exp(log_softmax(logW_ext));
41+ cdf = cumsum(weights);
42+ r = rand() * cdf(end);
43+ idx = find(cdf >= r, 1, 'first');
44+ theta_tilde = O_ext{idx, 1};
45+ rho_tilde = O_ext{idx, 2};
46+ end
47+
48+ if u_turn(O)
49+ break;
15950 end
160- return
161- end
162- num_steps = num_steps * 2;
163- step = step * 0.5;
164-end
165-% No step count met the tolerance.
166-success = false;
167-theta_next = theta;
168-rho_next = rho;
169-grad_next = grad;
170-logp_next = logp;
171-end
172-
173-function [theta, rho, grad] = leapfrog(theta, rho, grad, step, num_steps)
174-half_step = 0.5 * step;
175-for n = 1:num_steps
176- rho = rho + half_step * grad;
177- theta = theta + step * rho;
178- grad = grad_log_density(theta);
179- rho = rho + half_step * grad;
180-end
181-end
182-
183-function path = leapfrog_capture(theta, rho, grad, step, num_steps)
184-% Re-run a committed macro-step, returning the position after each leapfrog
185-% step (dim-by-num_steps). Used only while recording the movie trajectory.
186-half_step = 0.5 * step;
187-path = zeros(numel(theta), num_steps);
188-for n = 1:num_steps
189- rho = rho + half_step * grad;
190- theta = theta + step * rho;
191- grad = grad_log_density(theta);
192- rho = rho + half_step * grad;
193- path(:, n) = theta;
194-end
195-end
196-
197-function ok = reversible(step, num_steps, theta, rho, grad, logp_next, max_error)
198-% The adaptive step count is reversible only if no coarser (doubled-step)
199-% backward integration would itself have been accepted.
200-if num_steps == 1
201- ok = true;
202- return
203-end
204-ok = true;
205-while num_steps >= 2
206- num_steps = floor(num_steps / 2);
207- step = step * 2;
208- if within_tolerance(step, num_steps, theta, -rho, grad, logp_next, max_error)
209- ok = false;
210- return
21151 end
212-end
213-end
214-
215-function ok = within_tolerance(step, num_steps, theta, rho, grad, logp, max_error)
216-[theta, rho, ~] = leapfrog(theta, rho, grad, step, num_steps);
217-final_logp = log_density(theta) - sum(rho.^2) / 2;
218-ok = abs(final_logp - logp) <= max_error;
219-end
220-
221-function u = uturn(span1, span2, direction)
222-if direction == 1
223- span_bk = span1;
224- span_fw = span2;
225-else
226- span_bk = span2;
227- span_fw = span1;
228-end
229-scaled_diff = span_fw.theta_fw - span_bk.theta_bk;
230-u = (dot(span_fw.rho_fw, scaled_diff) < 0) || (dot(span_bk.rho_bk, scaled_diff) < 0);
231-end
232-
233-function span = combine(span_old, span_new, use_barker, direction)
234-logp_old = span_old.logp;
235-logp_new = span_new.logp;
236-logp_total = log_sum_exp(logp_old, logp_new);
237-if use_barker
238- log_denominator = logp_total;
239-else
240- log_denominator = logp_old;
241-end
242-update = log(rand) < (logp_new - log_denominator);
243-if update
244- theta_select = span_new.theta_select;
245-else
246- theta_select = span_old.theta_select;
247-end
248-if direction == 1
249- span_bk = span_old;
250- span_fw = span_new;
251-else
252- span_bk = span_new;
253- span_fw = span_old;
254-end
255-span = make_combined_span(span_bk, span_fw, theta_select, logp_total);
256-end
257-
258-function r = log_sum_exp(x, y)
259-m = max(x, y);
260-r = m + log(exp(x - m) + exp(y - m));
261-end
262-
263-function s = make_leaf_span(theta, rho, grad, logp)
264-s = struct('theta_bk', theta, 'rho_bk', rho, 'grad_bk', grad, 'logp_bk', logp, ...
265- 'theta_fw', theta, 'rho_fw', rho, 'grad_fw', grad, 'logp_fw', logp, ...
266- 'theta_select', theta, 'logp', logp);
267-end
268-
269-function s = make_combined_span(span_bk, span_fw, theta_select, logp_total)
270-s = struct('theta_bk', span_bk.theta_bk, 'rho_bk', span_bk.rho_bk, ...
271- 'grad_bk', span_bk.grad_bk, 'logp_bk', span_bk.logp_bk, ...
272- 'theta_fw', span_fw.theta_fw, 'rho_fw', span_fw.rho_fw, ...
273- 'grad_fw', span_fw.grad_fw, 'logp_fw', span_fw.logp_fw, ...
274- 'theta_select', theta_select, 'logp', logp_total);
275-end
52+
53+ T = length(logW) * h;
54+end
walnuts_demo.mmodified+10−8View file
@@ -1,14 +1,16 @@
1-% WALNUTS sampling of a 2D "banana" target.
1+% WALNUTS sampling of a 2D target.
22 %
3-% WALNUTS (Within-orbit Adaptive No-U-Turn Sampler) is a NUTS variant that
4-% adapts the leapfrog step size within each orbit. This draws samples from the
5-% banana target and shows the density with the samples on top.
3+% WALNUTS (the within-orbit adaptive leapfrog No-U-Turn Sampler) builds a
4+% Hamiltonian orbit, refining the leapfrog step within the orbit so the energy
5+% stays accurate, and picks a draw from the orbit. This samples a 2D target
6+% (banana by default) and shows the density with the samples on top.
67 %
7-% The algorithm lives in helpers/walnuts.m; the target in helpers/log_density.m
8-% and helpers/grad_log_density.m. `addpath` must be the first statement.
8+% The sampler is Nawaf Bou-Rabee's reference MATLAB implementation
9+% (helpers/walnuts.m + helpers); the target is helpers/log_density.m /
10+% grad_log_density.m. `addpath` must be the first statement.
911
10-addpath('helpers'); % walnuts, log_density, grad_log_density
12+addpath('helpers'); % walnuts, extend_orbit_*, micro, leapfrog, u_turn, ...
1113
1214 rng(1); % reproducible
1315
14-walnuts_sampler(1000, 0.4, 0.8); % N samples, leapfrog dt, max energy error
16+walnuts_sampler(1000, 0.8, log(1 / 0.66)); % N samples, step h, energy tol delta
walnuts_sampler.mmodified+75−39View file
@@ -1,68 +1,110 @@
1-function walnuts_sampler(N, dt, max_error, target)
1+function walnuts_sampler(N, h, delta, target)
22 %WALNUTS_SAMPLER Interactive figure: WALNUTS sampling of a 2D target.
3-% WALNUTS_SAMPLER(N, DT, MAX_ERROR, TARGET) draws N samples from TARGET
4-% ('banana' | 'gaussian' | 'correlated' | 'donut') with the WALNUTS sampler
5-% (helpers/walnuts.m) and opens a figure showing the density and the samples.
3+% WALNUTS_SAMPLER(N, H, DELTA, TARGET) draws N samples from TARGET
4+% ('banana' | 'gaussian' | 'correlated' | 'donut') with Nawaf Bou-Rabee's
5+% WALNUTS sampler (helpers/walnuts.m) and opens a figure showing the density
6+% and the samples. H is the base leapfrog step and DELTA the per-macro-step
7+% energy-variation tolerance that drives the within-orbit step refinement.
68 %
7-% This is the wiring: it selects the target, runs the sampler, evaluates the
8-% target on a grid for the heatmap, loads the prebuilt figure app, and sends
9-% both. It also handles requests from the figure (figure -> script): change
10-% the target, resample, or record an orbit movie. Run walnuts_demo.m (which
11-% addpath's helpers/).
9+% This is the wiring: it selects the target, runs the sampler (one
10+% walnuts(...) transition per draw, passing the target as function handles),
11+% evaluates the target on a grid for the heatmap, loads the figure app, and
12+% sends both. It also handles requests from the figure (change target,
13+% resample, record an orbit movie). Run walnuts_demo.m (which addpath's
14+% helpers/).
1215
1316 if nargin < 1 || isempty(N); N = 1000; end
14-if nargin < 2 || isempty(dt); dt = 0.4; end
15-if nargin < 3 || isempty(max_error); max_error = 0.8; end
17+if nargin < 2 || isempty(h); h = 0.8; end
18+if nargin < 3 || isempty(delta); delta = log(1 / 0.66); end
1619 if nargin < 4 || isempty(target); target = 'banana'; end
1720
1821 set_target(target);
19-samples = run_chain(N, dt, max_error, target);
22+samples = run_chain(N, h, delta, target);
2023 dens = density_grid(target);
2124
2225 html = fileread(fullfile('app', 'dist', 'index.html'));
2326 fig = figure;
2427 gl = uigridlayout(fig, [1 1], 'Padding', [0 0 0 0], ...
2528 'RowHeight', {'1x'}, 'ColumnWidth', {'1x'});
26-uihtml(gl, 'HTMLSource', html, 'Data', pack_data(samples, dens, N, dt, max_error, target), ...
29+uihtml(gl, 'HTMLSource', html, 'Data', pack_data(samples, dens, N, h, delta, target), ...
2730 'HTMLEventReceivedFcn', @(src, ev) on_event(src, ev));
2831 end
2932
3033 function on_event(src, ev)
3134 % Figure -> script. The figure owns the current settings and passes them back:
32-% 'resample' {n, dt, maxError, target} -> fresh chain; reply 'samples'.
33-% 'setTarget' {target, ...} -> new target; reply full 'data'
34-% (density + samples).
35-% 'movie' {dt, maxError, target} -> record orbit trajectories; 'movie'.
35+% 'resample' {n, h, delta, target} -> fresh chain; reply 'samples'.
36+% 'setTarget' {target, ...} -> new target; reply full 'data'.
37+% 'movie' {h, delta, target} -> record orbit trajectories; 'movie'.
3638 d = ev.HTMLEventData;
37-N = 1000; dt = 0.4; max_error = 0.8; target = 'banana';
39+N = 1000; h = 0.8; delta = log(1 / 0.66); target = 'banana';
3840 if isstruct(d)
3941 if isfield(d, 'n'); N = max(1, round(d.n)); end
40- if isfield(d, 'dt'); dt = d.dt; end
41- if isfield(d, 'maxError'); max_error = d.maxError; end
42+ if isfield(d, 'h'); h = d.h; end
43+ if isfield(d, 'delta'); delta = d.delta; end
4244 if isfield(d, 'target'); target = d.target; end
4345 end
4446 set_target(target);
4547 switch ev.HTMLEventName
4648 case 'resample'
47- samples = run_chain(N, dt, max_error, target);
49+ samples = run_chain(N, h, delta, target);
4850 sendEventToHTMLSource(src, 'samples', ...
4951 struct('x', samples(1, :), 'y', samples(2, :), ...
50- 'n', N, 'dt', dt, 'maxError', max_error, 'target', target));
52+ 'n', N, 'h', h, 'delta', delta, 'target', target));
5153 case 'setTarget'
52- samples = run_chain(N, dt, max_error, target);
54+ samples = run_chain(N, h, delta, target);
5355 dens = density_grid(target);
5456 sendEventToHTMLSource(src, 'data', ...
55- pack_data(samples, dens, N, dt, max_error, target));
57+ pack_data(samples, dens, N, h, delta, target));
5658 case 'movie'
57- n_steps = 12;
58- [~, traj] = walnuts(target_start(target), n_steps, dt, max_error, true);
59+ traj = record_movie(h, delta, target);
5960 sendEventToHTMLSource(src, 'movie', traj);
6061 end
6162 end
6263
64+function samples = run_chain(N, h, delta, target)
65+% Nawaf's walnuts(...) is one transition; loop it, passing our target density as
66+% function handles (they dispatch on the global WTARGET set by set_target).
67+i_max = 10;
68+burnin = 200;
69+theta = target_start(target);
70+samples = zeros(2, N);
71+total = burnin + N;
72+for i = 1:total
73+ theta = walnuts(@log_density, @grad_log_density, theta, h, i_max, delta);
74+ if i > burnin
75+ samples(:, i - burnin) = theta;
76+ end
77+end
78+end
79+
80+function traj = record_movie(h, delta, target)
81+% Run a few transitions and capture each one's orbit (cell of {theta, rho})
82+% for the step-by-step movie: px/py are the orbit positions, seg keeps them one
83+% polyline, startX/Y is where the transition began, selX/Y the selected draw.
84+i_max = 10;
85+n_steps = 12;
86+theta = target_start(target);
87+traj = {};
88+for s = 1:n_steps
89+ start_pt = theta;
90+ [theta, ~, O] = walnuts(@log_density, @grad_log_density, theta, h, i_max, delta);
91+ np = size(O, 1);
92+ px = zeros(1, np);
93+ py = zeros(1, np);
94+ for k = 1:np
95+ p = O{k, 1};
96+ px(k) = p(1);
97+ py(k) = p(2);
98+ end
99+ traj{s} = struct('px', px, 'py', py, 'seg', ones(1, np), ...
100+ 'startX', start_pt(1), 'startY', start_pt(2), ...
101+ 'selX', theta(1), 'selY', theta(2));
102+end
103+end
104+
63105 function set_target(name)
64-% Select the target the density functions evaluate (a process-global so
65-% walnuts.m's leapfrog can stay target-agnostic).
106+% Select the target the density handles evaluate (a process-global so the
107+% sampler stays target-agnostic).
66108 global WTARGET
67109 WTARGET = target_code(name);
68110 end
@@ -81,8 +123,8 @@ end
81123 end
82124
83125 function s = target_start(name)
84-% A sensible interior starting point for each target (the ring's hole is a bad
85-% start, so the donut starts on the ring).
126+% A sensible interior start for each target (the ring's hole is a bad start, so
127+% the donut starts on the ring).
86128 switch name
87129 case 'donut'
88130 s = [2.5; 0];
@@ -102,12 +144,6 @@ switch name
102144 end
103145 end
104146
105-function samples = run_chain(N, dt, max_error, target)
106-burnin = 200;
107-chain = walnuts(target_start(target), burnin + N, dt, max_error);
108-samples = chain(:, burnin + 1:end);
109-end
110-
111147 function dens = density_grid(target)
112148 %DENSITY_GRID Evaluate the selected target on a grid for the heatmap.
113149 % Row-major flat values: index (iy-1)*nx + ix, iy from ymin (1) to ymax (ny).
@@ -127,13 +163,13 @@ dens = struct('values', values, 'nx', nx, 'ny', ny, ...
127163 'xmin', xmin, 'xmax', xmax, 'ymin', ymin, 'ymax', ymax);
128164 end
129165
130-function data = pack_data(samples, dens, N, dt, max_error, target)
166+function data = pack_data(samples, dens, N, h, delta, target)
131167 data = struct();
132168 data.type = 'walnuts';
133169 data.samples = struct('x', samples(1, :), 'y', samples(2, :));
134170 data.density = dens;
135171 data.n = N;
136-data.dt = dt;
137-data.maxError = max_error;
172+data.h = h;
173+data.delta = delta;
138174 data.target = target;
139175 end