cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 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
38237c7Cite published JMLR version of the WALNUTS paperJeremy Magland 4% (Bou-Rabee, Carpenter, Kleppe, Liu, JMLR 27(113):1-64, 2026,
5% https://jmlr.org/papers/v27/25-1452.html). LOGMU / GRADLOGMU are
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 6% function handles for the target's log density and gradient. Returns the
7% selected draw THETA_TILDE, the orbit length T, and the orbit O (a cell of
8% {theta, rho} states) — O is exposed only for the figure's movie and does not
9% affect the algorithm.
10 d = length(theta);
11 rho = randn(d, 1);
12 theta_tilde = theta;
13 rho_tilde = rho;
14 logw_0 = logmu(theta) - 0.5 * norm(rho)^2;
16 O = {theta, rho};
17 logW = logw_0;
18 B = randi([0, 1], i_max, 1);
20 for i = 1:i_max
21 O_old = O;
22 logW_old = logW;
24 if B(i) == 1
25 [O_ext, logW_ext] = extend_orbit_forward(logmu, gradlogmu, O{end,1}, O{end,2}, logW(end), h, delta, 2^(i-1));
26 O = [O; O_ext];
27 logW = [logW; logW_ext];
28 else
29 [O_ext, logW_ext] = extend_orbit_backward(logmu, gradlogmu, O{1,1}, O{1,2}, logW(1), h, delta, 2^(i-1));
30 O = [O_ext; O];
31 logW = [logW_ext; logW];
32 end
34 if sub_u_turn(O_ext)
35 break;
36 end
38 logu = log(rand);
40 if logu <= logsumexp(logW_ext) - logsumexp(logW_old)
41 weights = exp(log_softmax(logW_ext));
42 cdf = cumsum(weights);
43 r = rand() * cdf(end);
44 idx = find(cdf >= r, 1, 'first');
45 theta_tilde = O_ext{idx, 1};
46 rho_tilde = O_ext{idx, 2};
47 end
49 if u_turn(O)
50 break;
52 end
54 T = length(logW) * h;
55end