1function [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;
15 O = {theta, rho};
16 logW = logw_0;
17 B = randi([0, 1], i_max, 1);
19 for i = 1:i_max
20 O_old = O;
21 logW_old = logW;
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
33 if sub_u_turn(O_ext)
34 break;
35 end
37 logu = log(rand);
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
48 if u_turn(O)
49 break;
50 end
51 end
53 T = length(logW) * h;
54end