Add target distribution selector (banana / Gaussian / correlated / donut) - concept-collection/walnuts-interactive

concept-collection / walnuts-interactive

Add target distribution selector (banana / Gaussian / correlated / donut)

A dropdown switches the target; the script rebuilds the density grid and a fresh chain and replies with a full 'data' event. The target is selected via a process-global the density functions dispatch on, so walnuts.m stays target-agnostic. Per-target view bounds and start points (the donut starts on the ring).

Jeremy Magland <jmagland@flatironinstitute.org> committed 1 month ago commit cd4150f86f38 parent 727e2e4 Browse files

5 changed files+219−82

README.mdmodified+9−6View file

@@ -4,9 +4,11 @@ Runs in the browser via [numbl](https://numbl.org) — no install.

44

55 ## ▶ [Open `walnuts_demo.m`](walnuts_demo.m) and click **Run**

66

7-Samples are drawn from a 2D "banana" target by **WALNUTS** (the within-orbit

8-adaptive leapfrog No-U-Turn Sampler). The figure shows the target density as a

9-heatmap with the samples scattered on top.

7+Samples are drawn from a 2D target by **WALNUTS** (the within-orbit adaptive

8+leapfrog No-U-Turn Sampler). The figure shows the target density as a heatmap

9+with the samples scattered on top. Switch between several targets (banana,

10+Gaussian, correlated Gaussian, donut) to see the sampler handle different

11+geometries.

1012

1113 WALNUTS is a NUTS variant that adapts the leapfrog step size *within* each

1214 orbit: each macro-step halves the step (and doubles the count) until the energy

@@ -15,6 +17,7 @@ U-turn.

1517

1618 Controls:

1719

20+- **Target** — switch the distribution (banana / Gaussian / correlated / donut).

1821 - **Samples** / **Leapfrog Δt** / **Max error** — change a setting to re-run the

1922 sampler with it (lower Δt or max error → more, smaller leapfrog steps).

2023 - **Resample** — a fresh chain with the current settings.

@@ -30,9 +33,9 @@ Controls:

3033 - `app/` — a single-file React app that draws the density heatmap + samples on a canvas.

3134

3235 The script sends the density grid + samples via `uihtml(..., 'Data', ...)`. The

33-controls call back: `sendToMATLAB('resample', {n, dt, maxError})` re-runs the

34-chain, and `sendToMATLAB('movie', {dt, maxError})` records a few transitions'

35-orbit trajectories (`walnuts(..., record=true)`); the script replies with

36+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

3639 `sendEventToHTMLSource`.

3740

3841 ## Credits

app/src/App.tsxmodified+63−10View file

@@ -17,8 +17,16 @@ interface WalnutsData {

1717 n: number;

1818 dt: number;

1919 maxError: number;

20+ target: string;

2021 }

2122

23+const TARGETS: { value: string; label: string }[] = [

24+ { value: "banana", label: "Banana" },

25+ { value: "gaussian", label: "Gaussian" },

26+ { value: "correlated", label: "Correlated Gaussian" },

27+ { value: "donut", label: "Donut (ring)" },

28+];

29+

2230 function isWalnutsData(d: unknown): d is WalnutsData {

2331 return (

2432 !!d &&

@@ -66,17 +74,31 @@ export function App() {

6674 const [n, setN] = useState(DEFAULT_N);

6775 const [dt, setDt] = useState(0.4);

6876 const [maxError, setMaxError] = useState(0.8);

77+ const [target, setTargetState] = useState("banana");

6978 const [busy, setBusy] = useState(false);

7079 const [movieData, setMovieData] = useState<MovieStep[] | null>(null);

7180 const [movie, setMovie] = useState<{ si: number; k: number } | null>(null);

7281

82+ // Apply a full payload (initial Data, or a `data` event after a target change).

83+ const applyData = (d: WalnutsData) => {

84+ setData(d);

85+ setN(d.n);

86+ setDt(d.dt);

87+ setMaxError(d.maxError);

88+ setTargetState(d.target);

89+ };

90+

7391 useEffect(() => {

7492 const offData = onData(d => {

93+ if (isWalnutsData(d)) applyData(d);

94+ });

95+ // New target: full fresh payload (density + samples).

96+ const offFull = onHostEvent("data", d => {

7597 if (isWalnutsData(d)) {

76- setData(d);

77- setN(d.n);

78- setDt(d.dt);

79- setMaxError(d.maxError);

98+ applyData(d);

99+ setBusy(false);

100+ setMovie(null);

101+ setMovieData(null);

80102 }

81103 });

82104 // Resample: same target, new draws.

@@ -102,6 +124,7 @@ export function App() {

102124 });

103125 return () => {

104126 offData();

127+ offFull();

105128 offSamples();

106129 offMovie();

107130 };

@@ -132,7 +155,17 @@ export function App() {

132155 if (!data || busy) return;

133156 stopMovie();

134157 setBusy(true);

135- sendToMATLAB("resample", { n: count, dt: step, maxError: err });

158+ sendToMATLAB("resample", { n: count, dt: step, maxError: err, target });

159+ };

160+

161+ // Switch the target: the script rebuilds the density + draws and replies with

162+ // a full `data` event.

163+ const changeTarget = (value: string) => {

164+ setTargetState(value);

165+ if (!data) return;

166+ stopMovie();

167+ setBusy(true);

168+ sendToMATLAB("setTarget", { target: value, n, dt, maxError });

136169 };

137170

138171 const playMovie = () => {

@@ -142,7 +175,7 @@ export function App() {

142175 }

143176 if (!data || busy) return;

144177 setBusy(true); // until the trajectory arrives

145- sendToMATLAB("movie", { dt, maxError });

178+ sendToMATLAB("movie", { dt, maxError, target });

146179 };

147180

148181 // ── derive the movie overlay for the current frame ──

@@ -194,10 +227,23 @@ export function App() {

194227 )}

195228

196229 <div style={panelStyle}>

197- <div style={{ fontWeight: 600, marginBottom: 4 }}>WALNUTS</div>

198- <div style={{ fontSize: 11, color: "#475569", marginBottom: 8 }}>

199- sampling a banana target

200- </div>

230+ <div style={{ fontWeight: 600, marginBottom: 6 }}>WALNUTS</div>

231+

232+ <label style={labelStyle}>

233+ Target

234+ <select

235+ value={target}

236+ disabled={controlsDisabled}

237+ onChange={e => changeTarget(e.target.value)}

238+ style={selectStyle}

239+ >

240+ {TARGETS.map(t => (

241+ <option key={t.value} value={t.value}>

242+ {t.label}

243+ </option>

244+ ))}

245+ </select>

246+ </label>

201247

202248 <label style={labelStyle}>

203249 Samples: <b>{n.toLocaleString()}</b>

@@ -314,6 +360,13 @@ const sliderStyle: CSSProperties = {

314360 marginTop: 2,

315361 };

316362

363+const selectStyle: CSSProperties = {

364+ width: "100%",

365+ marginTop: 2,

366+ fontSize: 11,

367+ padding: "2px 4px",

368+};

369+

317370 const btnStyle: CSSProperties = {

318371 flex: 1,

319372 padding: "4px 6px",

helpers/grad_log_density.mmodified+34−16View file

@@ -1,18 +1,36 @@

11 function g = grad_log_density(theta)

2-%GRAD_LOG_DENSITY Gradient of the banana log density (see log_density.m).

3-a = 2;

4-b = 0.2;

5-x1 = theta(1);

6-x2 = theta(2);

7-y1 = x1 / a;

8-y2 = x2 * a + a * b * (x1.^2 + a.^2);

9-d1 = y1 - 0;

10-d2 = y2 - 4;

11-% Gradient w.r.t. y of the Gaussian log density is -inv(Sigma)*d.

12-gy1 = -((4 / 3) * d1 - (2 / 3) * d2);

13-gy2 = -(-(2 / 3) * d1 + (4 / 3) * d2);

14-% Chain rule back through the twist y(x).

15-gx1 = gy1 / a + gy2 * a * b * 2 * x1;

16-gx2 = gy2 * a;

17-g = [gx1; gx2];

2+%GRAD_LOG_DENSITY Gradient of the selected target's log density (see

3+% log_density.m). Dispatches on the global WTARGET.

4+global WTARGET

5+t = WTARGET;

6+if isempty(t)

7+ t = 1;

8+end

9+x = theta(1);

10+y = theta(2);

11+switch t

12+ case 2 % standard Gaussian

13+ g = [-x; -y];

14+ case 3 % correlated Gaussian, precision = inv([1 0.8; 0.8 1])

15+ c = 1 / (1 - 0.8^2);

16+ g = [-c * (x - 0.8 * y); -c * (y - 0.8 * x)];

17+ case 4 % donut / ring

18+ r = sqrt(x.^2 + y.^2);

19+ if r < 1e-9

20+ g = [0; 0];

21+ else

22+ k = -(r - 2.5) / 0.15 / r;

23+ g = [k * x; k * y];

24+ end

25+ otherwise % banana

26+ a = 2;

27+ b = 0.2;

28+ y1 = x / a;

29+ y2 = y * a + a * b * (x.^2 + a.^2);

30+ d1 = y1;

31+ d2 = y2 - 4;

32+ gy1 = -((4 / 3) * d1 - (2 / 3) * d2);

33+ gy2 = -(-(2 / 3) * d1 + (4 / 3) * d2);

34+ g = [gy1 / a + gy2 * a * b * 2 * x; gy2 * a];

35+end

1836 end

helpers/log_density.mmodified+28−17View file

@@ -1,19 +1,30 @@

11 function lp = log_density(theta)

2-%LOG_DENSITY Unnormalized log density of the "banana" target.

3-% The standard banana from chi-feng's mcmc-demo: a correlated Gaussian

4-% (mean [0; 4], covariance [1 0.5; 0.5 1]) pulled into a banana shape by a

5-% quadratic twist. theta is a 2x1 column vector. The normalizing constant is

6-% dropped (irrelevant for sampling and for the heatmap).

7-a = 2;

8-b = 0.2;

9-x1 = theta(1);

10-x2 = theta(2);

11-% Twist x -> y, then evaluate the Gaussian at y.

12-y1 = x1 / a;

13-y2 = x2 * a + a * b * (x1.^2 + a.^2);

14-d1 = y1 - 0;

15-d2 = y2 - 4;

16-% Q = d' * inv(Sigma) * d, with inv([1 .5; .5 1]) = [4/3 -2/3; -2/3 4/3].

17-Q = (4 / 3) * (d1.^2 + d2.^2) - (4 / 3) * d1 .* d2;

18-lp = -0.5 * Q;

2+%LOG_DENSITY Unnormalized log density of the currently selected 2D target.

3+% The target is chosen by the global WTARGET (1=banana, 2=gaussian,

4+% 3=correlated, 4=donut), set by walnuts_sampler. Defaults to the banana.

5+% Normalizing constants are dropped (irrelevant for sampling and the heatmap).

6+global WTARGET

7+t = WTARGET;

8+if isempty(t)

9+ t = 1;

10+end

11+x = theta(1);

12+y = theta(2);

13+switch t

14+ case 2 % standard Gaussian N(0, I)

15+ lp = -0.5 * (x.^2 + y.^2);

16+ case 3 % correlated Gaussian, covariance [1 0.8; 0.8 1]

17+ lp = -0.5 * (x.^2 - 2 * 0.8 * x .* y + y.^2) / (1 - 0.8^2);

18+ case 4 % donut / ring: radius 2.5, variance 0.15 in the radial direction

19+ r = sqrt(x.^2 + y.^2);

20+ lp = -(r - 2.5).^2 / (2 * 0.15);

21+ otherwise % 1: banana (correlated Gaussian under a quadratic twist)

22+ a = 2;

23+ b = 0.2;

24+ y1 = x / a;

25+ y2 = y * a + a * b * (x.^2 + a.^2);

26+ d1 = y1;

27+ d2 = y2 - 4;

28+ lp = -0.5 * ((4 / 3) * (d1.^2 + d2.^2) - (4 / 3) * d1 .* d2);

29+end

1930 end

walnuts_sampler.mmodified+85−33View file

@@ -1,66 +1,117 @@

1-function walnuts_sampler(N, dt, max_error)

2-%WALNUTS_SAMPLER Interactive figure: WALNUTS sampling of a 2D banana target.

3-% WALNUTS_SAMPLER(N, DT, MAX_ERROR) draws N samples from the banana target

4-% (helpers/log_density.m) with the WALNUTS sampler (helpers/walnuts.m) and

5-% opens a figure showing the target density and the samples.

1+function walnuts_sampler(N, dt, max_error, target)

2+%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.

66 %

7-% This is the wiring: it runs the sampler, evaluates the target on a grid for

8-% the heatmap, loads the prebuilt figure app, and sends both to it. It also

9-% re-samples on request (figure -> script), wired so adding controls later is

10-% a UI-only change. Run walnuts_demo.m (which addpath's helpers/).

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/).

1112

1213 if nargin < 1 || isempty(N); N = 1000; end

1314 if nargin < 2 || isempty(dt); dt = 0.4; end

1415 if nargin < 3 || isempty(max_error); max_error = 0.8; end

16+if nargin < 4 || isempty(target); target = 'banana'; end

1517

16-burnin = 200;

17-chain = walnuts(randn(2, 1), burnin + N, dt, max_error);

18-samples = chain(:, burnin + 1:end);

19-

20-dens = density_grid();

18+set_target(target);

19+samples = run_chain(N, dt, max_error, target);

20+dens = density_grid(target);

2121

2222 html = fileread(fullfile('app', 'dist', 'index.html'));

2323 fig = figure;

2424 gl = uigridlayout(fig, [1 1], 'Padding', [0 0 0 0], ...

2525 'RowHeight', {'1x'}, 'ColumnWidth', {'1x'});

26-uihtml(gl, 'HTMLSource', html, 'Data', pack_data(samples, dens, N, dt, max_error), ...

26+uihtml(gl, 'HTMLSource', html, 'Data', pack_data(samples, dens, N, dt, max_error, target), ...

2727 'HTMLEventReceivedFcn', @(src, ev) on_event(src, ev));

2828 end

2929

3030 function on_event(src, ev)

31-% Figure -> script. The target is fixed, so the figure only ever needs new

32-% samples or a movie trajectory back.

33-% 'resample' {n, dt, maxError} -> draw a fresh chain; reply with 'samples'.

34-% 'movie' {dt, maxError} -> record a short chain's orbit trajectories;

35-% reply with 'movie'.

31+% 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'.

3636 d = ev.HTMLEventData;

37-N = 1000; dt = 0.4; max_error = 0.8;

37+N = 1000; dt = 0.4; max_error = 0.8; target = 'banana';

3838 if isstruct(d)

3939 if isfield(d, 'n'); N = max(1, round(d.n)); end

4040 if isfield(d, 'dt'); dt = d.dt; end

4141 if isfield(d, 'maxError'); max_error = d.maxError; end

42+ if isfield(d, 'target'); target = d.target; end

4243 end

44+set_target(target);

4345 switch ev.HTMLEventName

4446 case 'resample'

45- burnin = 200;

46- chain = walnuts(randn(2, 1), burnin + N, dt, max_error);

47- samples = chain(:, burnin + 1:end);

47+ samples = run_chain(N, dt, max_error, target);

4848 sendEventToHTMLSource(src, 'samples', ...

49- struct('x', samples(1, :), 'y', samples(2, :), 'n', N, 'dt', dt, 'maxError', max_error));

49+ struct('x', samples(1, :), 'y', samples(2, :), ...

50+ 'n', N, 'dt', dt, 'maxError', max_error, 'target', target));

51+ case 'setTarget'

52+ samples = run_chain(N, dt, max_error, target);

53+ dens = density_grid(target);

54+ sendEventToHTMLSource(src, 'data', ...

55+ pack_data(samples, dens, N, dt, max_error, target));

5056 case 'movie'

51- % Record the orbit-building trajectory of a few transitions for the

52- % step-by-step animation. Each traj{i} = {px, py, seg, startX/Y, selX/Y}.

5357 n_steps = 12;

54- [~, traj] = walnuts(randn(2, 1), n_steps, dt, max_error, true);

58+ [~, traj] = walnuts(target_start(target), n_steps, dt, max_error, true);

5559 sendEventToHTMLSource(src, 'movie', traj);

5660 end

5761 end

5862

59-function dens = density_grid()

60-%DENSITY_GRID Evaluate the target log density on a grid for the heatmap.

63+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).

66+global WTARGET

67+WTARGET = target_code(name);

68+end

69+

70+function c = target_code(name)

71+switch name

72+ case 'gaussian'

73+ c = 2;

74+ case 'correlated'

75+ c = 3;

76+ case 'donut'

77+ c = 4;

78+ otherwise

79+ c = 1; % banana

80+end

81+end

82+

83+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).

86+switch name

87+ case 'donut'

88+ s = [2.5; 0];

89+ case 'banana'

90+ s = [0; 1];

91+ otherwise

92+ s = [0; 0];

93+end

94+end

95+

96+function [xmin, xmax, ymin, ymax] = target_bounds(name)

97+switch name

98+ case 'banana'

99+ xmin = -6; xmax = 6; ymin = -7; ymax = 3;

100+ otherwise

101+ xmin = -4; xmax = 4; ymin = -4; ymax = 4;

102+end

103+end

104+

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+

111+function dens = density_grid(target)

112+%DENSITY_GRID Evaluate the selected target on a grid for the heatmap.

61113 % Row-major flat values: index (iy-1)*nx + ix, iy from ymin (1) to ymax (ny).

62-xmin = -6; xmax = 6;

63-ymin = -7; ymax = 3;

114+[xmin, xmax, ymin, ymax] = target_bounds(target);

64115 nx = 100; ny = 100;

65116 xs = linspace(xmin, xmax, nx);

66117 ys = linspace(ymin, ymax, ny);

@@ -76,7 +127,7 @@ dens = struct('values', values, 'nx', nx, 'ny', ny, ...

76127 'xmin', xmin, 'xmax', xmax, 'ymin', ymin, 'ymax', ymax);

77128 end

78129

79-function data = pack_data(samples, dens, N, dt, max_error)

130+function data = pack_data(samples, dens, N, dt, max_error, target)

80131 data = struct();

81132 data.type = 'walnuts';

82133 data.samples = struct('x', samples(1, :), 'y', samples(2, :));

@@ -84,4 +135,5 @@ data.density = dens;

84135 data.n = N;

85136 data.dt = dt;

86137 data.maxError = max_error;

138+data.target = target;

87139 end