Sign in
concept-collection / turing-surface
Editable operator/solver files in the page, and a solver selector
The models now call a uniform solve(...) whose target is a one-line host-generated shim: a solver <select> in the controls picks richardson, bicgstab or gmres, and switching recompiles, like changing niter. The editor's file picker grows to include lib/dlap.m and the three solvers/*.m files, each editable with the same Recompile/Revert flow as the model and the surface; edits are passed to compilation as lib-source overrides. --solver joins the bench command (round-tripped like every other knob), and the solver-comparison tests switch solvers through the session option — the same path the UI uses — instead of rewriting model source.
Jeremy Magland <jmagland@flatironinstitute.org> committed commit 1015122f1d87 parent 59f3e22 Browse files
14 changed files+191−79
README.mdmodified+15−10View file
@@ -125,7 +125,7 @@ And a **model** is a reaction plus one solve per species — the whole of
125125 [`models/schnakenberg.m`](models/schnakenberg.m)'s step is:
126126
127127 ```matlab
128-function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, a, b, D1, D2, dt, niter)
128+function [Un, Vn, u, v] = step(U, V, lam, filt, wlm, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, a, b, D1, D2, dt, nlm, niter)
129129 u = synth(U);
130130 v = synth(V);
131131 uuv = u .* u .* v;
@@ -133,19 +133,23 @@ function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx,
133133 Bu = U + dt * analys(a - u + uuv);
134134 Bv = V + dt * analys(b - uuv);
135135
136- Un = richardson(Bu, dt * D1, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
137- Vn = richardson(Bv, dt * D2, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
136+ Un = solve(Bu, dt * D1, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
137+ Vn = solve(Bv, dt * D2, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
138138 end
139139 ```
140140
141-Trying a different solver against the same operator is a change to those two
142-call lines: every solver composes from `dlap` (the matvec is
141+`solve` is a one-line, host-generated shim (`src/mgpu/libs.ts`) forwarding to
142+whichever solver the app's **solver control** currently selects — every
143+solver composes from `dlap` (the matvec is
143144 `(1 + dtD.*lam).*x - dtD.*dlap(x)`, the preconditioner the elementwise
144-divide), and which one a model calls is part of what compiles — swapping
145-recompiles, like changing `niter` already does. The solver is written as a
146-full re-evaluation rather than an accumulated correction on purpose: where
147-`dlap` computes to zero there is no correction to mis-round, and the divide is
148-turing-sphere's arithmetic unchanged.
145+divide), and the choice is part of what compiles, so switching recompiles,
146+like changing `niter` already does. A model may equally name a solver
147+directly in those call lines. The operator and every solver are files in the
148+page's editor, next to the model and the surface — edit one and recompile,
149+same as the model. The Richardson solver is written as a full re-evaluation
150+rather than an accumulated correction on purpose: where `dlap` computes to
151+zero there is no correction to mis-round, and the divide is turing-sphere's
152+arithmetic unchanged.
149153
150154 Three solvers ship. [`solvers/bicgstab.m`](solvers/bicgstab.m) solves the
151155 same system by preconditioned BiCGSTAB — same `dlap`, same preconditioner, a
@@ -340,6 +344,7 @@ currently simulating:
340344
341345 ```
342346 npm run bench -- --preset schnak-spots --geometry ellipsoid --lmax 63 --niter 1 \
347+ --solver richardson \
343348 --steps 2000 --seed 1 --a 0.1 --b 0.9 --D1 0.0004 --D2 0.008 --dt 0.05 \
344349 --gax 1.5 --gay 1 --gaz 0.6
345350 ```
docs/richardson-iteration.mdmodified+4−3View file
@@ -68,9 +68,10 @@ for `k = 0 .. niter-1`. `solvers/richardson.m`'s `for k = 1:niter` loop *is*
6868 this: `Un^(0)` is the divide computed just before the loop, each pass
6969 computes `Un^(k+1)` from `Un^(k)`, and `dlap` — evaluated once per iteration —
7070 is its own function, `lib/dlap.m`. A model's step calls the solver once per
71-species (`Un = richardson(Bu, dt * D1, ...)`), which is where the split pays:
72-a different solver for the same operator is a different call in the model,
73-with `lib/dlap.m` untouched. The solver is written as a full re-evaluation
71+species (`Un = solve(Bu, dt * D1, ...)`, routed to the selected solvers/*.m
72+file by a host-generated shim), which is where the split pays: a different
73+solver for the same operator is a selector change — or a different call in
74+the model — with `lib/dlap.m` untouched. The solver is written as a full re-evaluation
7475 rather than an accumulated correction `δ = Un^(k+1) - Un^(k)` on purpose:
7576 where `dlap` evaluates to zero exactly, every `Un^(k)` is bit-for-bit
7677 `Un^(0)`, with no cancellation to round differently. (In practice `dlap` is a
index.htmlmodified+11−3View file
@@ -168,9 +168,10 @@
168168 Reaction-diffusion on closed surfaces, solved live with spherical
169169 harmonics on WebGPU via
170170 <a href="https://github.com/concept-collection/shtns-webgpu">shtns-webgpu</a>.
171- Both the solver and the shape are the MATLAB below, compiled in your
172- browser by <a href="https://numbl.org">numbl</a>. Edit either and watch
173- it change. Drag to rotate.
171+ The model, the shape, the diffusion operator and the solvers are all
172+ the MATLAB below, compiled in your browser by
173+ <a href="https://numbl.org">numbl</a>. Edit any of them and watch it
174+ change. Drag to rotate.
174175 </p>
175176 <div class="controls">
176177 <label>preset
@@ -182,6 +183,13 @@
182183 <label title="Blend between the sphere (0) and the surface (1). Display only.">morph
183184 <input type="range" id="morph" min="0" max="1" step="0.01" value="1" />
184185 </label>
186+ <label title="Which solver answers the models' solve(...) call — see solvers/. Changing it recompiles.">solver
187+ <select id="solver">
188+ <option value="richardson" selected>richardson</option>
189+ <option value="bicgstab">bicgstab</option>
190+ <option value="gmres">gmres</option>
191+ </select>
192+ </label>
185193 <label title="Iterations of the implicit diffusion solve. Changing it recompiles.">solve iters
186194 <select id="niter">
187195 <option value="0">0</option>
models/allencahn.mmodified+3−3View file
@@ -3,16 +3,16 @@
33 % du/dt = eps2*lap_g(u) + u - u^3
44 %
55 % Same scheme as models/schnakenberg.m: explicit reaction, then the implicit
6-% diffusion solve handed to solvers/richardson.m.
6+% diffusion solve handed to solve(...) — the solver the app's selector picks.
77
88 function [U, u] = init(noise)
99 U = analys(noise);
1010 u = synth(U);
1111 end
1212
13-function [Un, u] = step(U, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, eps2, dt, niter)
13+function [Un, u] = step(U, lam, filt, wlm, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, eps2, dt, nlm, niter)
1414 u = synth(U);
1515
1616 Bu = U + dt * analys(u - u.^3);
17- Un = richardson(Bu, dt * eps2, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
17+ Un = solve(Bu, dt * eps2, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
1818 end
models/brusselator.mmodified+4−4View file
@@ -4,7 +4,7 @@
44 % dv/dt = D2*lap_g(v) + B*u - u^2*v
55 %
66 % Same scheme as models/schnakenberg.m: explicit reaction, then the implicit
7-% diffusion solve handed to solvers/richardson.m.
7+% diffusion solve handed to solve(...) — the solver the app's selector picks.
88
99 function [U, V, u, v] = init(noise, A, B)
1010 U = analys(A + noise);
@@ -13,7 +13,7 @@ function [U, V, u, v] = init(noise, A, B)
1313 v = synth(V);
1414 end
1515
16-function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, A, B, D1, D2, dt, niter)
16+function [Un, Vn, u, v] = step(U, V, lam, filt, wlm, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, A, B, D1, D2, dt, nlm, niter)
1717 u = synth(U);
1818 v = synth(V);
1919 uuv = u .* u .* v;
@@ -21,6 +21,6 @@ function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx,
2121 Bu = U + dt * analys(A - (B + 1) * u + uuv);
2222 Bv = V + dt * analys(B * u - uuv);
2323
24- Un = richardson(Bu, dt * D1, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
25- Vn = richardson(Bv, dt * D2, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
24+ Un = solve(Bu, dt * D1, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
25+ Vn = solve(Bv, dt * D2, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
2626 end
models/schnakenberg.mmodified+8−7View file
@@ -5,10 +5,11 @@
55 %
66 % Explicit reaction, implicit diffusion (IMEX Euler): the step forms the
77 % right-hand side B of the linear system (I - dt*D*lap_g) Unew = B, and hands
8-% the solve to solvers/richardson.m — which applies the operator's geometric
9-% part through lib/dlap.m. Trying a different solver for the same operator is
10-% a change to these two call lines. Grid fields are npts x 1; spectral fields
11-% are real 2 x nlm. See docs/richardson-iteration.md.
8+% the solve to solve(...) — the solver the app's selector picks (richardson,
9+% bicgstab or gmres, see solvers/), every one applying the operator's
10+% geometric part through lib/dlap.m. A model may also name a solver directly
11+% in these two call lines. Grid fields are npts x 1; spectral fields are
12+% real 2 x nlm. See docs/richardson-iteration.md.
1213
1314 function [U, V, u, v] = init(noise, a, b)
1415 us = a + b;
@@ -19,7 +20,7 @@ function [U, V, u, v] = init(noise, a, b)
1920 v = synth(V);
2021 end
2122
22-function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, a, b, D1, D2, dt, niter)
23+function [Un, Vn, u, v] = step(U, V, lam, filt, wlm, gx, gy, gz, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, a, b, D1, D2, dt, nlm, niter)
2324 u = synth(U);
2425 v = synth(V);
2526 uuv = u .* u .* v;
@@ -29,6 +30,6 @@ function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, Vtx, Vty, Vtz, Vpx,
2930 Bv = V + dt * analys(b - uuv);
3031
3132 % The species diffuse independently, so each gets its own solve.
32- Un = richardson(Bu, dt * D1, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
33- Vn = richardson(Bv, dt * D2, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter);
33+ Un = solve(Bu, dt * D1, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
34+ Vn = solve(Bv, dt * D2, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter);
3435 end
scripts/bench.tsmodified+5−1View file
@@ -20,6 +20,7 @@ import { requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
2020 import { ModelSession } from '../src/mgpu/session.ts';
2121 import { presets } from '../src/mgpu/registry.ts';
2222 import { mGeometries, DEFAULT_GEOMETRY_KEY } from '../src/geom/registry.ts';
23+import { DEFAULT_SOLVER, solverKeys } from '../src/mgpu/libs.ts';
2324 import {
2425 parseArgs,
2526 modelForSpec,
@@ -47,6 +48,8 @@ const USAGE = `usage: ${BENCH_COMMAND} [options]
4748 --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
4849 --niter <n> iterations of the implicit solve, unrolled into the compiled
4950 step (default ${DEFAULT_NITER})
51+ --solver <key> ${solverKeys.join(' | ')} — which solver answers the
52+ models' solve(...) call (default ${DEFAULT_SOLVER})
5053 --steps <n> timed steps (default ${DEFAULT_STEPS})
5154 --warmup <n> untimed steps first (default ${DEFAULT_WARMUP})
5255 --seed <n> initial-noise seed (default ${DEFAULT_SEED})
@@ -175,6 +178,7 @@ try {
175178 geometry,
176179 geometryParams: spec.geometryParams,
177180 niter: spec.niter,
181+ solver: spec.solver,
178182 });
179183 session.seed(spec.seed);
180184
@@ -198,7 +202,7 @@ try {
198202 );
199203 console.log(
200204 ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${session.sht.nlm.toLocaleString()} · ` +
201- `${spec.niter} solve iteration${spec.niter === 1 ? '' : 's'}`,
205+ `${spec.niter} solve iteration${spec.niter === 1 ? '' : 's'} of ${spec.solver}`,
202206 );
203207 console.log(` compiled ${plan.step.length} GPU ops/step (${kernels} generated kernels)`);
204208 console.log(` fourier ${session.sht.fourierMode.toUpperCase()} stage`);
src/bench/runSpec.tsmodified+12−0View file
@@ -22,6 +22,7 @@ import {
2222 type MGeometry,
2323 } from '../geom/registry.ts';
2424 import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
25+import { DEFAULT_SOLVER, solverKeys, type SolverKey } from '../mgpu/libs.ts';
2526
2627 export interface RunSpec {
2728 /** Preset key from the registry; fixes the model, params may still be edited. */
@@ -44,6 +45,9 @@ export interface RunSpec {
4445 /** Iterations of the .m's implicit solve. Structural: it is unrolled into
4546 * the compiled step, so it belongs to the spec rather than to the params. */
4647 niter: number;
48+ /** Which solver answers the models' solve(...) call. Structural like
49+ * niter: the choice compiles into the step. */
50+ solver: SolverKey;
4751 }
4852
4953 export const DEFAULT_NITER = 1;
@@ -113,6 +117,7 @@ export function formatCommand(spec: RunSpec): string {
113117 `--geometry ${spec.geometry}`,
114118 `--lmax ${spec.lmax}`,
115119 `--niter ${spec.niter}`,
120+ `--solver ${spec.solver}`,
116121 `--steps ${spec.steps}`,
117122 `--seed ${spec.seed}`,
118123 ...model.params.map((p) => `--${p.key} ${String(spec.params[p.key])}`),
@@ -160,6 +165,12 @@ export function parseArgs(argv: string[]): RunSpec {
160165 const { model, params } = resolvePreset(presetKey);
161166 const geometryKey = take('geometry') ?? DEFAULT_GEOMETRY_KEY;
162167 const { geometry, params: geometryParams } = resolveGeometry(geometryKey);
168+ const solverRaw = take('solver') ?? DEFAULT_SOLVER;
169+ if (!(solverKeys as string[]).includes(solverRaw)) {
170+ throw new Error(
171+ `--solver must be one of ${solverKeys.join(', ')} (got '${solverRaw}')`,
172+ );
173+ }
163174 const spec: RunSpec = {
164175 preset: presetKey,
165176 lmax: count('lmax', DEFAULT_LMAX, 1),
@@ -170,6 +181,7 @@ export function parseArgs(argv: string[]): RunSpec {
170181 geometry: geometryKey,
171182 geometryParams,
172183 niter: count('niter', DEFAULT_NITER, 0),
184+ solver: solverRaw as SolverKey,
173185 };
174186 const readInto = (into: Params, key: string, flag: string): void => {
175187 const raw = take(flag);
src/main.tsmodified+35−12View file
@@ -4,6 +4,7 @@ import { ModelSession } from './mgpu/session.ts';
44 import { mModelByKey, presets, type MModel, type Params } from './mgpu/registry.ts';
55 import { ModelCompileError, formatFailure } from './mgpu/errors.ts';
66 import { EXTERNAL_OPS } from './mgpu/externals.ts';
7+import { libPath, modelLibs, type SolverKey } from './mgpu/libs.ts';
78 import { CodeEditor } from './editor/codeEditor.ts';
89 import {
910 formatCommand,
@@ -37,6 +38,7 @@ const $ = <T extends HTMLElement>(id: string): T =>
3738 const elModel = $<HTMLSelectElement>('model');
3839 const elGeometry = $<HTMLSelectElement>('geometry');
3940 const elMorph = $<HTMLInputElement>('morph');
41+const elSolver = $<HTMLSelectElement>('solver');
4042 const elNiter = $<HTMLSelectElement>('niter');
4143 const elLmax = $<HTMLSelectElement>('lmax');
4244 const elOversample = $<HTMLSelectElement>('oversample');
@@ -81,7 +83,11 @@ for (const g of mGeometries) {
8183 o.textContent = g.label;
8284 elGeometry.append(o);
8385 }
84-for (const [value, label] of [['model', 'the solver'], ['geometry', 'the surface']]) {
86+for (const [value, label] of [
87+ ['model', 'the model'],
88+ ['geometry', 'the surface'],
89+ ...modelLibs.map((f) => [`lib:${f.name}`, libPath(f.name)]),
90+]) {
8591 const o = document.createElement('option');
8692 o.value = value;
8793 o.textContent = label;
@@ -95,16 +101,17 @@ for (const name of colormapNames) {
95101 }
96102 elColormap.value = 'jet';
97103
98-/** Whichever .m is open: the solver or the surface. Both are MATLAB, compiled
99- * by the same backend, so one editor serves both. The host-provided operations
100- * are marked so the boundary between the file and what it is given is
101- * visible. */
104+/** Whichever .m is open: the model, the surface, or one of the shared
105+ * operator/solver files. All are MATLAB, compiled by the same backend, so
106+ * one editor serves them all. The host-provided operations are marked so
107+ * the boundary between the file and what it is given is visible. */
102108 const editor = new CodeEditor({
103109 textarea: elSource,
104110 overlay: elHighlight,
105111 external: new Set(EXTERNAL_OPS.keys()),
106112 onInput: (value) => {
107113 if (editing === 'geometry') editedGeomSource = value;
114+ else if (editing.startsWith('lib:')) editedLibs.set(editing.slice(4), value);
108115 else editedSource = value;
109116 elRecompile.textContent = 'Recompile *';
110117 },
@@ -210,11 +217,14 @@ let model: MModel = mModelByKey(initial.model.key)!;
210217 let params: Params = initial.params;
211218 let geometry: MGeometry = mGeometryByKey(DEFAULT_GEOMETRY_KEY)!;
212219 let geomParams: Params = defaultGeometryParams(geometry);
213-/** Which file the editor is showing. */
214-let editing: 'model' | 'geometry' = 'model';
220+/** Which file the editor is showing: the model, the surface, or a shared
221+ * file (`lib:<name>`). */
222+let editing: string = 'model';
215223 /** Each .m as edited in the page; `null` while it matches the file. */
216224 let editedSource: string | null = null;
217225 let editedGeomSource: string | null = null;
226+/** Shared-file working copies, by file name; absent while unedited. */
227+const editedLibs = new Map<string, string>();
218228 /** Sphere (0) to surface (1). Display only; does not touch the solver. */
219229 let morph = 1;
220230 let seed = 1;
@@ -235,6 +245,8 @@ let posBuf: Float32Array | null = null;
235245
236246 const source = (): string => editedSource ?? model.source;
237247 const geomSource = (): string => editedGeomSource ?? geometry.source;
248+const libSource = (name: string): string =>
249+ editedLibs.get(name) ?? modelLibs.find((f) => f.name === name)?.source ?? '';
238250
239251 // ---------------------------------------------------------------- UI wiring
240252 function buildParamInputs(): void {
@@ -323,10 +335,14 @@ function applyGeometryChoice(key: string): void {
323335
324336 /** Load the chosen file into the editor, keeping any unsaved edit to it. */
325337 function showEditorFile(): void {
326- editing = elEditorFile.value === 'geometry' ? 'geometry' : 'model';
338+ editing = elEditorFile.value;
327339 if (editing === 'geometry') {
328340 editor.value = geomSource();
329341 elEditorTitle.textContent = `geometries/${geometry.key}.m`;
342+ } else if (editing.startsWith('lib:')) {
343+ const name = editing.slice(4);
344+ editor.value = libSource(name);
345+ elEditorTitle.textContent = libPath(name);
330346 } else {
331347 editor.value = source();
332348 elEditorTitle.textContent = `models/${model.key}.m`;
@@ -345,6 +361,7 @@ function currentSpec(): RunSpec {
345361 geometry: geometry.key,
346362 geometryParams: geomParams,
347363 niter: Number(elNiter.value),
364+ solver: elSolver.value as SolverKey,
348365 };
349366 }
350367
@@ -358,8 +375,10 @@ elModel.addEventListener('change', () => {
358375 });
359376 elLmax.addEventListener('change', () => void rebuild());
360377 // The solve iteration count is unrolled into the compiled step, so unlike a
361-// parameter it cannot be changed without recompiling.
378+// parameter it cannot be changed without recompiling. The solver choice is a
379+// generated one-line shim compiled with the model, so it recompiles too.
362380 elNiter.addEventListener('change', () => void rebuild());
381+elSolver.addEventListener('change', () => void rebuild());
363382 // Oversampling and geometry are display-or-data changes, not code ones, so
364383 // they swap things in place rather than rebuilding the run. Serialized through
365384 // one chain: a rapid second change waits its turn.
@@ -406,11 +425,13 @@ elMovie.addEventListener('click', () => {
406425
407426 elRecompile.addEventListener('click', () => {
408427 if (editing === 'geometry') editedGeomSource = editor.value;
428+ else if (editing.startsWith('lib:')) editedLibs.set(editing.slice(4), editor.value);
409429 else editedSource = editor.value;
410430 void rebuild();
411431 });
412432 elRevert.addEventListener('click', () => {
413433 if (editing === 'geometry') editedGeomSource = null;
434+ else if (editing.startsWith('lib:')) editedLibs.delete(editing.slice(4));
414435 else editedSource = null;
415436 showEditorFile();
416437 void rebuild();
@@ -632,6 +653,8 @@ async function rebuild(): Promise<void> {
632653 geometryParams: geomParams,
633654 geometrySource: geomSource(),
634655 niter: Number(elNiter.value),
656+ solver: elSolver.value as SolverKey,
657+ libSources: Object.fromEntries(editedLibs),
635658 });
636659 } catch (e) {
637660 reportCompileError(e);
@@ -872,9 +895,9 @@ function submitSteps(n: number): void {
872895 * the Movie button itself becomes the cancel button. */
873896 function setMovieUi(on: boolean): void {
874897 const locked = [
875- elModel, elGeometry, elMorph, elNiter, elLmax, elOversample, elColormap,
876- elRunPause, elBenchmark, elReseed, elRecompile, elRevert, elEditorFile,
877- elMovieSpeed, elMovieRes, elMovieRotate, elMovieToggle,
898+ elModel, elGeometry, elMorph, elSolver, elNiter, elLmax, elOversample,
899+ elColormap, elRunPause, elBenchmark, elReseed, elRecompile, elRevert,
900+ elEditorFile, elMovieSpeed, elMovieRes, elMovieRotate, elMovieToggle,
878901 ];
879902 for (const el of locked) el.disabled = on;
880903 elParams.querySelectorAll('input').forEach((input) => (input.disabled = on));
src/mgpu/libs.tsmodified+36−0View file
@@ -32,3 +32,39 @@ export const solverLibs: LibFile[] = [
3232
3333 /** Everything a model may call. */
3434 export const modelLibs: LibFile[] = [...operatorLibs, ...solverLibs];
35+
36+/** Where each shared file lives on disk, for display. */
37+export const libPath = (name: string): string =>
38+ operatorLibs.some((f) => f.name === name) ? `lib/${name}` : `solvers/${name}`;
39+
40+export type SolverKey = 'richardson' | 'bicgstab' | 'gmres';
41+export const solverKeys: SolverKey[] = ['richardson', 'bicgstab', 'gmres'];
42+export const DEFAULT_SOLVER: SolverKey = 'richardson';
43+
44+/** What each solver actually takes: richardson needs no inner products, so
45+ * no weights; only gmres sizes a basis bank, so only it takes nlm. */
46+const SOLVE_FORWARD: Record<SolverKey, string> = {
47+ richardson: 'richardson(B, dtD, lam, filt, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter)',
48+ bicgstab: 'bicgstab(B, dtD, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, niter)',
49+ gmres: 'gmres(B, dtD, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter)',
50+};
51+
52+/**
53+ * The one-line shim behind the models' `solve(...)` call. The models pass
54+ * every argument any solver could want, and this host-generated file
55+ * forwards to the chosen one — so which solver runs is a compile-time choice
56+ * the app's solver control makes (swapping recompiles, like changing niter),
57+ * while the models stay identical across solvers. A model may also bypass
58+ * the shim and call a solver by name.
59+ */
60+export function solveShim(solver: SolverKey): LibFile {
61+ return {
62+ name: 'solve.m',
63+ source:
64+ `% Host-generated: forwards the models' solve(...) call to the solver\n` +
65+ `% selected in the app. See solvers/${solver}.m.\n` +
66+ `function X = solve(B, dtD, lam, filt, wlm, Vtx, Vty, Vtz, Vpx, Vpy, Vpz, nlm, niter)\n` +
67+ ` X = ${SOLVE_FORWARD[solver]};\n` +
68+ `end\n`,
69+ };
70+}
src/mgpu/model.tsmodified+10−2View file
@@ -23,7 +23,7 @@ import { lmIndex, type ShtConfig } from '../sht/layout.ts';
2323 import { HostBuffers, ModelPlan } from './plan.ts';
2424 import { inFunction, inFunctionAsync, inModel } from './errors.ts';
2525 import { CompiledModel, type Binding } from './compile.ts';
26-import { modelLibs } from './libs.ts';
26+import { DEFAULT_SOLVER, modelLibs, solveShim, type LibFile } from './libs.ts';
2727
2828 export interface ModelParams {
2929 [key: string]: number;
@@ -57,6 +57,13 @@ export interface GpuModelOptions {
5757 * the count is part of what compiles and changing it recompiles.
5858 */
5959 niter?: number;
60+ /**
61+ * The shared .m files compiled alongside the model. Defaults to the shipped
62+ * operator and solvers plus a `solve(...)` shim forwarding to richardson;
63+ * the session passes its own list to honor the app's solver selection and
64+ * any lib edits.
65+ */
66+ libs?: LibFile[];
6067 }
6168
6269 /** Host-supplied surface fields, in the layout the .m sees them. */
@@ -212,8 +219,9 @@ export class GpuModel {
212219 for (const p of paramNames) bindings[p] = { kind: 'param' };
213220
214221 // Parsing belongs to the file, not to either function.
222+ const libs = opts.libs ?? [...modelLibs, solveShim(DEFAULT_SOLVER)];
215223 const compiled = inModel(
216- () => new CompiledModel(source, bindings, { npts, nlm }, 'model.m', modelLibs),
224+ () => new CompiledModel(source, bindings, { npts, nlm }, 'model.m', libs),
217225 );
218226 // Both functions return the new state first, then the rendered grid fields.
219227 const nargout = state.length + view.length;
src/mgpu/session.tsmodified+29−1View file
@@ -10,6 +10,7 @@ import { ShtPlan } from '../sht/sht.ts';
1010 import { DerivPlan } from '../sht/deriv.ts';
1111 import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
1212 import { GpuModel, type ModelParams } from './model.ts';
13+import { DEFAULT_SOLVER, modelLibs, solveShim, type SolverKey } from './libs.ts';
1314 import { seededNoise } from './noise.ts';
1415 import type { MModel } from './registry.ts';
1516 import { Geometry } from '../geom/geometry.ts';
@@ -36,6 +37,17 @@ export interface ModelSessionOptions {
3637 * is unrolled into the op sequence, so a change recompiles.
3738 */
3839 niter?: number;
40+ /**
41+ * Which solver answers the models' `solve(...)` call (default richardson).
42+ * Structural like niter: the choice is a one-line generated shim compiled
43+ * with the model, so changing it recompiles.
44+ */
45+ solver?: SolverKey;
46+ /**
47+ * Overrides for the shared .m files, by file name (`'dlap.m'`,
48+ * `'richardson.m'`, ...) — the editor's working copies.
49+ */
50+ libSources?: Record<string, string>;
3951 }
4052
4153 export class ModelSession {
@@ -47,6 +59,8 @@ export class ModelSession {
4759 readonly npts: number;
4860 /** Iterations of the implicit solve compiled into the step. */
4961 readonly niter: number;
62+ /** The solver compiled behind the models' solve(...) call. */
63+ readonly solver: SolverKey;
5064
5165 /** The surface being solved on, as spherical-harmonic coefficients. */
5266 #geometry: Geometry;
@@ -76,6 +90,7 @@ export class ModelSession {
7690 geometryModel: MGeometry;
7791 deriv: DerivPlan;
7892 niter: number;
93+ solver: SolverKey;
7994 }) {
8095 this.device = init.device;
8196 this.model = init.model;
@@ -90,6 +105,7 @@ export class ModelSession {
90105 this.#geometryModel = init.geometryModel;
91106 this.#deriv = init.deriv;
92107 this.niter = init.niter;
108+ this.solver = init.solver;
93109 }
94110
95111 get geometry(): Geometry {
@@ -109,6 +125,17 @@ export class ModelSession {
109125 const { device, model, params, lmax } = opts;
110126 const oversample = Math.max(1, Math.round(opts.oversample ?? 1));
111127 const niter = Math.max(0, Math.round(opts.niter ?? 1));
128+ const solver = opts.solver ?? DEFAULT_SOLVER;
129+ // The shared files, with the editor's working copies substituted, plus
130+ // the shim that routes solve(...) to the chosen solver.
131+ const libs = [
132+ ...modelLibs.map((f) =>
133+ opts.libSources?.[f.name] !== undefined
134+ ? { name: f.name, source: opts.libSources[f.name] }
135+ : f,
136+ ),
137+ solveShim(solver),
138+ ];
112139 const geometryModel = opts.geometry ?? mGeometryByKey(SPHERE_KEY)!;
113140 const geometryParams = opts.geometryParams ?? defaultGeometryParams(geometryModel);
114141 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
@@ -156,11 +183,12 @@ export class ModelSession {
156183 geometry,
157184 deriv,
158185 niter,
186+ libs,
159187 });
160188 gpu.setParams(params);
161189 return new ModelSession({
162190 device, model, cfg, sht, displaySht, gpu, params, oversample,
163- geometry, geometryModel, deriv, niter,
191+ geometry, geometryModel, deriv, niter, solver,
164192 });
165193 } catch (e) {
166194 // The transform plans own GPU buffers; do not leak them on a compile error.
test/geometryChecks.tsmodified+18−33View file
@@ -43,26 +43,6 @@ import type { Check, Log } from './analyticChecks.ts';
4343 const LMAX = 31;
4444 const STEPS = 20;
4545
46-/**
47- * The shipped Schnakenberg model with its solver calls switched from
48- * richardson to another shipped solver — the same edit a user makes in the
49- * page, which is the point: same operator, same model, different solver.
50- * (gmres additionally takes `nlm`, for sizing its Krylov basis bank.)
51- * Throws if the model text drifted from what this rewrites, so the test
52- * fails loudly rather than silently comparing richardson with itself.
53- */
54-function solverSchnak(source: string, solver: 'bicgstab' | 'gmres'): string {
55- const out = source
56- .replace('step(U, V, lam, filt, ', 'step(U, V, lam, filt, wlm, ')
57- .replace('D2, dt, niter)', solver === 'gmres' ? 'D2, dt, nlm, niter)' : 'D2, dt, niter)')
58- .replaceAll('richardson(Bu, dt * D1, lam, filt, ', `${solver}(Bu, dt * D1, lam, filt, wlm, `)
59- .replaceAll('richardson(Bv, dt * D2, lam, filt, ', `${solver}(Bv, dt * D2, lam, filt, wlm, `)
60- .replaceAll('Vpz, niter);', solver === 'gmres' ? 'Vpz, nlm, niter);' : 'Vpz, niter);');
61- if (!out.includes('wlm,') || !out.includes(`${solver}(Bu`) || !out.includes(`${solver}(Bv`)) {
62- throw new Error('solverSchnak: the model source no longer matches the rewrite');
63- }
64- return out;
65-}
6646 /** The app's actual default lmax (README: "at the default lmax 63 that is a
6747 * 128x256 grid"), used for the niter/geometry sweep below and the peanut
6848 * check next to it -- the divergence they're both about is a real, lmax-
@@ -399,21 +379,26 @@ export async function geometryChecks(
399379 );
400380 }
401381
402- // ---- two solvers, one operator ------------------------------------------
403- // solvers/bicgstab.m against solvers/richardson.m on the same implicit
404- // system: a Krylov iteration converges superlinearly where the stationary
405- // one converges linearly, so at equal niter it must land much closer to the
406- // converged answer. The comparison is a ratio against the same reference,
407- // which keeps it meaningful on SwiftShader's looser fp32 too.
382+ // ---- three solvers, one operator ----------------------------------------
383+ // The Krylov solvers against richardson on the same implicit system: a
384+ // Krylov iteration converges superlinearly where the stationary one
385+ // converges linearly, so at equal niter it must land much closer to the
386+ // converged answer. Switched exactly the way the app's solver control does
387+ // — the session's `solver` option, which swaps the solve(...) shim. The
388+ // comparison is a ratio against the same reference, which keeps it
389+ // meaningful on SwiftShader's looser fp32 too.
408390 {
409391 const model = mModelByKey('schnakenberg')!;
410392 const params = defaultParams(model);
411393 const ellipsoid = mGeometryByKey('ellipsoid')!;
412- const run = async (source: string | undefined, niter: number): Promise<Float32Array> => {
394+ const run = async (
395+ solver: 'richardson' | 'bicgstab' | 'gmres',
396+ niter: number,
397+ ): Promise<Float32Array> => {
413398 const session = await ModelSession.create({
414399 device, model, params, lmax: LMAX,
415400 geometry: ellipsoid, geometryParams: defaultGeometryParams(ellipsoid),
416- niter, ...(source ? { source } : {}),
401+ niter, solver,
417402 });
418403 session.seed(1);
419404 session.step(STEPS);
@@ -421,10 +406,10 @@ export async function geometryChecks(
421406 session.destroy();
422407 return U;
423408 };
424- const ref = await run(undefined, 8); // richardson, effectively converged
425- const rich = await run(undefined, 2);
426- const bicg = await run(solverSchnak(model.source, 'bicgstab'), 2);
427- const gmres = await run(solverSchnak(model.source, 'gmres'), 2);
409+ const ref = await run('richardson', 8); // effectively converged
410+ const rich = await run('richardson', 2);
411+ const bicg = await run('bicgstab', 2);
412+ const gmres = await run('gmres', 2);
428413 const relRich = relL2(rich, ref);
429414 const relBicg = relL2(bicg, ref);
430415 const relGmres = relL2(gmres, ref);
@@ -512,7 +497,7 @@ export async function geometryChecks(
512497 const session = await ModelSession.create({
513498 device, model, params, lmax: SWEEP_LMAX,
514499 geometry: peanut, geometryParams: defaultGeometryParams(peanut),
515- niter, source: solverSchnak(model.source, solver),
500+ niter, solver,
516501 });
517502 session.seed(1);
518503 session.step(STEPS);
test/modelChecks.tsmodified+1−0View file
@@ -78,6 +78,7 @@ export async function modelChecks(
7878 geometry: 'peanut',
7979 geometryParams: { waist: 0.45, stretch: 1.25 },
8080 niter: 3,
81+ solver: 'bicgstab',
8182 };
8283 const command = formatCommand(spec);
8384 const back = parseArgs(command.slice(BENCH_COMMAND.length).trim().split(/\s+/));