727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 1function [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).
24if nargin < 5 || isempty(record)
25 record = false;
26end
27% Trajectory recorder (used only when record is true). macro_step appends each
28% committed leapfrog path to these as it grows the orbit.
29global WREC_ON WREC_X WREC_Y WREC_SEG WREC_SEGID
30WREC_ON = record;
32dim = numel(theta0);
33chain = zeros(dim, n_samples);
34traj = {};
35theta = theta0;
36for 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
51end
52WREC_ON = false;
53end
55function 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.
59rho = randn(numel(theta), 1);
60grad = grad_log_density(theta);
61logp = log_density(theta) - sum(rho.^2) / 2;
62span_accum = make_leaf_span(theta, rho, grad, logp);
64for 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
80end
81theta_select = span_accum.theta_select;
82end
84function [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.
87if depth == 0
88 [ok, span] = build_leaf(span_in, direction, dt, max_error);
89 return
90end
91[ok, left] = build_span(span_in, direction, depth - 1, dt, max_error);
92if ~ok
93 span = left;
94 return
95end
96[ok, right] = build_span(left, direction, depth - 1, dt, max_error);
97if ~ok
98 span = right;
99 return
100end
101if uturn(left, right, direction)
102 ok = false;
103 span = left;
104 return
105end
106% Sub-orbit selection is a Barker update (use_barker = true).
107span = combine(left, right, true, direction);
108ok = true;
109end
111function [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);
114if ~success
115 ok = false;
116 span = span_in;
117 return
118end
119span = make_leaf_span(theta_next, rho_next, grad_next, logp_next);
120ok = true;
121end
123function [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.
128global WREC_ON WREC_X WREC_Y WREC_SEG WREC_SEGID
129if 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;
135else
136 theta = span.theta_bk;
137 rho = span.rho_bk;
138 grad = span.grad_bk;
139 logp = span.logp_bk;
140 step = -dt;
141end
143num_steps = 1;
144for 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)];
159 end
160 return
161 end
162 num_steps = num_steps * 2;
163 step = step * 0.5;
164end
165% No step count met the tolerance.
166success = false;
167theta_next = theta;
168rho_next = rho;
169grad_next = grad;
170logp_next = logp;
171end
173function [theta, rho, grad] = leapfrog(theta, rho, grad, step, num_steps)
174half_step = 0.5 * step;
175for 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;
180end
181end
183function 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.
186half_step = 0.5 * step;
187path = zeros(numel(theta), num_steps);
188for 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;
194end
195end
197function 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.
200if num_steps == 1
201 ok = true;
202 return
203end
204ok = true;
205while 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
211 end
212end
213end
215function ok = within_tolerance(step, num_steps, theta, rho, grad, logp, max_error)
216[theta, rho, ~] = leapfrog(theta, rho, grad, step, num_steps);
217final_logp = log_density(theta) - sum(rho.^2) / 2;
218ok = abs(final_logp - logp) <= max_error;
219end
221function u = uturn(span1, span2, direction)
222if direction == 1
223 span_bk = span1;
224 span_fw = span2;
225else
226 span_bk = span2;
227 span_fw = span1;
228end
229scaled_diff = span_fw.theta_fw - span_bk.theta_bk;
230u = (dot(span_fw.rho_fw, scaled_diff) < 0) || (dot(span_bk.rho_bk, scaled_diff) < 0);
231end
233function span = combine(span_old, span_new, use_barker, direction)
234logp_old = span_old.logp;
235logp_new = span_new.logp;
236logp_total = log_sum_exp(logp_old, logp_new);
237if use_barker
238 log_denominator = logp_total;
239else
240 log_denominator = logp_old;
241end
242update = log(rand) < (logp_new - log_denominator);
243if update
244 theta_select = span_new.theta_select;
245else
246 theta_select = span_old.theta_select;
247end
248if direction == 1
249 span_bk = span_old;
250 span_fw = span_new;
251else
252 span_bk = span_new;
253 span_fw = span_old;
254end
255span = make_combined_span(span_bk, span_fw, theta_select, logp_total);
256end
258function r = log_sum_exp(x, y)
259m = max(x, y);
260r = m + log(exp(x - m) + exp(y - m));
261end
263function s = make_leaf_span(theta, rho, grad, logp)
264s = 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);
267end
269function s = make_combined_span(span_bk, span_fw, theta_select, logp_total)
270s = 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);
275end