Sign in
concept-collection / walnuts-interactive
main
Branches
main
walnuts-interactive / walnuts_sampler.m
175 lines · 5.7 KBCodeBlameHistory
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 1function walnuts_sampler(N, h, delta, target)
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 2%WALNUTS_SAMPLER Interactive figure: WALNUTS sampling of a 2D target.
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 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.
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 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/).
16if nargin < 1 || isempty(N); N = 1000; end
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 17if nargin < 2 || isempty(h); h = 0.8; end
18if nargin < 3 || isempty(delta); delta = log(1 / 0.66); end
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 19if nargin < 4 || isempty(target); target = 'banana'; end
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 22samples = run_chain(N, h, delta, target);
25html = fileread(fullfile('app', 'dist', 'index.html'));
26fig = figure;
27gl = uigridlayout(fig, [1 1], 'Padding', [0 0 0 0], ...
28 'RowHeight', {'1x'}, 'ColumnWidth', {'1x'});
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 29uihtml(gl, 'HTMLSource', html, 'Data', pack_data(samples, dens, N, h, delta, target), ...
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 30 'HTMLEventReceivedFcn', @(src, ev) on_event(src, ev));
31end
33function on_event(src, ev)
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 34% Figure -> script. The figure owns the current settings and passes them back:
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 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'.
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 38d = ev.HTMLEventData;
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 39N = 1000; h = 0.8; delta = log(1 / 0.66); target = 'banana';
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 40if isstruct(d)
41 if isfield(d, 'n'); N = max(1, round(d.n)); end
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 42 if isfield(d, 'h'); h = d.h; end
43 if isfield(d, 'delta'); delta = d.delta; end
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 44 if isfield(d, 'target'); target = d.target; end
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 47switch ev.HTMLEventName
48 case 'resample'
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 49 samples = run_chain(N, h, delta, target);
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 50 sendEventToHTMLSource(src, 'samples', ...
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 51 struct('x', samples(1, :), 'y', samples(2, :), ...
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 52 'n', N, 'h', h, 'delta', delta, 'target', target));
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 54 samples = run_chain(N, h, delta, target);
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 55 dens = density_grid(target);
56 sendEventToHTMLSource(src, 'data', ...
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 57 pack_data(samples, dens, N, h, delta, target));
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 59 traj = record_movie(h, delta, target);
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 60 sendEventToHTMLSource(src, 'movie', traj);
61end
62end
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 64function 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).
67i_max = 10;
68burnin = 200;
69theta = target_start(target);
70samples = zeros(2, N);
71total = burnin + N;
72for 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
77end
78end
80function 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.
84i_max = 10;
85d0056Movie: animate 40 transitions instead of 12Jeremy Magland 85n_steps = 40;
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 86theta = target_start(target);
87traj = {};
88for 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));
102end
103end
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 106% Select the target the density handles evaluate (a process-global so the
107% sampler stays target-agnostic).
109WTARGET = target_code(name);
110end
112function c = target_code(name)
113switch name
114 case 'gaussian'
115 c = 2;
116 case 'correlated'
117 c = 3;
118 case 'donut'
119 c = 4;
120 otherwise
121 c = 1; % banana
122end
123end
125function s = target_start(name)
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 126% A sensible interior start for each target (the ring's hole is a bad start, so
127% the donut starts on the ring).
129 case 'donut'
130 s = [2.5; 0];
131 case 'banana'
132 s = [0; 1];
133 otherwise
134 s = [0; 0];
135end
136end
138function [xmin, xmax, ymin, ymax] = target_bounds(name)
139switch name
140 case 'banana'
141 xmin = -6; xmax = 6; ymin = -7; ymax = 3;
142 otherwise
143 xmin = -4; xmax = 4; ymin = -4; ymax = 4;
144end
145end
147function dens = density_grid(target)
148%DENSITY_GRID Evaluate the selected target on a grid for the heatmap.
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 149% Row-major flat values: index (iy-1)*nx + ix, iy from ymin (1) to ymax (ny).
cd4150fAdd target distribution selector (banana / Gaussian / correlated / donut)Jeremy Magland 150[xmin, xmax, ymin, ymax] = target_bounds(target);
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 151nx = 100; ny = 100;
152xs = linspace(xmin, xmax, nx);
153ys = linspace(ymin, ymax, ny);
154values = zeros(1, nx * ny);
155k = 1;
156for iy = 1:ny
157 for ix = 1:nx
158 values(k) = log_density([xs(ix); ys(iy)]);
159 k = k + 1;
160 end
161end
162dens = struct('values', values, 'nx', nx, 'ny', ny, ...
163 'xmin', xmin, 'xmax', xmax, 'ymin', ymin, 'ymax', ymax);
164end
cccc996Use Nawaf Bou-Rabee's reference WALNUTS implementationJeremy Magland 166function data = pack_data(samples, dens, N, h, delta, target)
727e2e4Interactive WALNUTS sampling of a 2D banana targetJeremy Magland 167data = struct();
168data.type = 'walnuts';
169data.samples = struct('x', samples(1, :), 'y', samples(2, :));
170data.density = dens;
171data.n = N;
173data.delta = delta;