Reaction-diffusion on spherical-harmonic surfaces
Fork of turing-sphere, reorganized around a geometry: an embedding of the
sphere into R^3 carried as spherical-harmonic coefficients for x, y and z.
A geometry is a .m file (geometries/*.m) defining shape(theta, phi, ...),
compiled by the same backend as the models. It is evaluated once on the
solver grid, then analysed into coefficients — which is what the solver and
the renderer both use, so the surface is exactly band-limited and can be
synthesized on any grid. Display oversampling therefore evaluates the
embedding at more points rather than subdividing it.
The models split lap_g = lap_s + dlap and iterate
Unew = (B + dt*D*dlap(Unew)) ./ (1 + dt*D*lam)
from the round-sphere answer: preconditioned Richardson with the exactly
invertible round-sphere operator as the preconditioner.
dlap is a placeholder, identically zero. The geometry is rendered but not
yet in the operator: that needs the induced metric and spheroidal transforms
(SHsph_to_spat / spat_to_SHsph), which the vendored WGSL transforms do not
implement. Written so the sphere case is bit-for-bit turing-sphere's
arithmetic, with no cancellation — asserted in the tests.
To run that loop on the GPU, the planner now unrolls counted `for` loops with
compile-time-known bounds. Loop-carried values need no special handling
because numbl gives a variable one cName per assignment; the loop variable is
bound as a per-iteration literal. Fusion survives, since numbl's inline pass
recurses into loop bodies — but it runs there with no protected names, so
compile.ts refuses a loop whose result was folded away.
Also: sphere/surface morph restored from figpack's SphereEmbedding view,
geometry and niter carried through the shared RunSpec (with a round-trip
test), and a fourth check module for the surface and the loop.
Not carried over from turing-sphere: bench/shtns, whose C transcription of
the model would have to track a step this project intends to change.
68 changed files+14702−0
.github/workflows/ci.ymladded+41−0View file
@@ -0,0 +1,41 @@
1+name: ci
2+on:
3+ push:
4+ branches: [main]
5+ pull_request:
6+
7+jobs:
8+ test:
9+ runs-on: ubuntu-latest
10+ steps:
11+ - uses: actions/checkout@v4
12+ - uses: actions/setup-node@v4
13+ with:
14+ node-version: 24
15+ cache: npm
16+ # numbl is a `file:../../numbl` dependency: we use its compiler internals
17+ # (parser, lowerer, IR, inline pass), which its published package `exports`
18+ # do not expose. Clone it where that relative path expects it. Pinned so a
19+ # change to those internals cannot silently break the build — the surface we
20+ # rely on is written down in src/mgpu/numbl.d.ts.
21+ #
22+ # numbl's own dependencies are NOT needed: the slice we import is
23+ # self-contained TypeScript, verified by building against a checkout with no
24+ # node_modules.
25+ - name: Check out numbl (sibling dependency)
26+ env:
27+ NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
28+ run: |
29+ git clone --filter=blob:none --no-checkout \
30+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
31+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
32+ # --ignore-scripts: npm runs a linked package's `prepare` script, and
33+ # numbl's is husky, which is not installed here.
34+ - run: npm ci --ignore-scripts
35+ # The checks compile MATLAB to compute shaders, so they need a GPU; the
36+ # browser suite below runs the same modules on SwiftShader if there is none.
37+ - run: npm run test:node -- --skip-without-gpu
38+ # headless Chrome + SwiftShader software WebGPU
39+ - run: npm run test:gpu
40+ env:
41+ CHROME_PATH: /usr/bin/google-chrome
.github/workflows/deploy.ymladded+59−0View file
@@ -0,0 +1,59 @@
1+name: deploy
2+on:
3+ push:
4+ branches: [main]
5+ workflow_dispatch:
6+
7+permissions:
8+ contents: read
9+ pages: write
10+ id-token: write
11+
12+concurrency:
13+ group: pages
14+ cancel-in-progress: true
15+
16+jobs:
17+ build-deploy:
18+ runs-on: ubuntu-latest
19+ environment:
20+ name: github-pages
21+ url: ${{ steps.deployment.outputs.page_url }}
22+ steps:
23+ - uses: actions/checkout@v4
24+ - uses: actions/setup-node@v4
25+ with:
26+ node-version: 24
27+ cache: npm
28+ # numbl is a `file:../../numbl` dependency: we use its compiler internals
29+ # (parser, lowerer, IR, inline pass), which its published package `exports`
30+ # do not expose. Clone it where that relative path expects it. Pinned so a
31+ # change to those internals cannot silently break the build — the surface we
32+ # rely on is written down in src/mgpu/numbl.d.ts.
33+ #
34+ # numbl's own dependencies are NOT needed: the slice we import is
35+ # self-contained TypeScript, verified by building against a checkout with no
36+ # node_modules.
37+ - name: Check out numbl (sibling dependency)
38+ env:
39+ NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
40+ run: |
41+ git clone --filter=blob:none --no-checkout \
42+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
43+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
44+ # --ignore-scripts: npm runs a linked package's `prepare` script, and
45+ # numbl's is husky, which is not installed here.
46+ - run: npm ci --ignore-scripts
47+ # The checks compile MATLAB to compute shaders, so they need a GPU; the
48+ # browser suite below runs the same modules on SwiftShader if there is none.
49+ - run: npm run test:node -- --skip-without-gpu
50+ - run: npm run build
51+ # Pages must already be enabled with "GitHub Actions" as the source; the
52+ # workflow token cannot create the site itself (`enablement: true` fails
53+ # with "Resource not accessible by integration").
54+ - uses: actions/configure-pages@v5
55+ - uses: actions/upload-pages-artifact@v3
56+ with:
57+ path: dist
58+ - id: deployment
59+ uses: actions/deploy-pages@v4
.gitignoreadded+12−0View file
@@ -0,0 +1,12 @@
1+node_modules/
2+dist/
3+*.log
4+*.png
5+
6+# bench/shtns: the upstream SHTns checkout, what configure decided, and the
7+# binaries. ./bootstrap.sh rebuilds all of it.
8+bench/shtns/shtns/
9+bench/shtns/shtns.mk
10+bench/shtns/shtbench
11+bench/shtns/shtbench_gpu
12+bench/shtns/*.o
README.mdadded+427−0View file
@@ -0,0 +1,427 @@
1+# turing-surface
2+
3+Reaction–diffusion systems (Turing patterns) on **closed surfaces given by
4+spherical-harmonic embeddings**, solved live in the browser with a spectral
5+method whose transforms run on the GPU via WebGPU.
6+
7+This is the sibling of
8+[turing-sphere](https://github.com/concept-collection/turing-sphere), which
9+solves the same systems on the round sphere. Everything there is here; what is
10+added is a *surface*.
11+
12+> [!WARNING]
13+> **The geometry is rendered, not yet solved on.** The Laplace–Beltrami
14+> operator in the models is still the round sphere's — the term that carries
15+> the shape is a placeholder that is identically zero. On anything but the
16+> sphere you are looking at the sphere's pattern painted onto that surface, not
17+> the pattern that surface would grow. Everything the correction needs in order
18+> to be dropped in — the embedding, the split of the operator, the iterative
19+> solve, the unrolled loop — is built and tested. See
20+> [The geometry is not in the operator yet](#the-geometry-is-not-in-the-operator-yet).
21+
22+## What a surface is here
23+
24+A geometry is an embedding of the sphere into R³: three scalar fields x, y, z
25+over the (θ, φ) parametrization, each carried as spherical-harmonic
26+coefficients. The unit sphere is the case where all three are pure degree-1
27+harmonics.
28+
29+You write one down as MATLAB, in [`geometries/`](geometries/):
30+
31+```matlab
32+function [gx, gy, gz] = shape(theta, phi, waist, stretch)
33+ st = sin(theta);
34+ r = 1 - waist * (st .^ 2);
35+ gx = r .* (st .* cos(phi));
36+ gy = r .* (st .* sin(phi));
37+ gz = (1 + stretch) * (r .* cos(theta));
38+end
39+```
40+
41+That is ordinary element-wise MATLAB and goes through the same compiler and the
42+same WGSL backend the models do. It is evaluated once on the solver's grid, and
43+then **analysed into coefficients**, which is the form everything downstream
44+uses. Two things follow from going through the coefficients rather than keeping
45+the pointwise values:
46+
47+- **It is exactly band-limited at lmax.** The surface has as many derivatives as
48+ the scheme needs and no aliased content the solver cannot see. What the solver
49+ and the renderer both use is the *synthesis* of the coefficients, so for a
50+ shape with sharp features the surface being solved on is not quite the one
51+ that was written down — which is the honest thing for a spectral method to do.
52+- **It can be evaluated on any grid.** The renderer draws the surface on the
53+ (possibly finer) display grid by synthesizing the same coefficients there.
54+ That is exact interpolation, not subdivision — the same argument that lets the
55+ species fields be oversampled, and it is checked directly in the tests.
56+
57+Four geometries ship: [sphere](geometries/sphere.m) (the reference case),
58+[ellipsoid](geometries/ellipsoid.m), [peanut](geometries/peanut.m) — a dumbbell
59+whose waist is a saddle — and [bumpy](geometries/bumpy.m). Each is editable in
60+the page, with its own parameters. Changing a shape does not recompile the
61+solver and does not disturb the run: the geometry is data whose shape in the
62+bindings depends only on the grid, so a swap is six buffer writes and the
63+pattern carries straight on.
64+
65+A **morph** slider blends the drawn surface back to the unit sphere. The
66+parametrization is the sphere's either way, so sweeping it shows which point
67+went where.
68+
69+## The scheme, and where the geometry enters
70+
71+It solves the N-species system
72+
73+```
74+d(u_k)/dt = D_k*lap_g(u_k) + f_k(t, u_1, ..., u_N), k = 1, ..., N
75+```
76+
77+where `lap_g` is the Laplace–Beltrami operator of the surface. On the round
78+sphere `lap_g` is diagonal in spherical-harmonic space with eigenvalues
79+`-l(l+1)`, which is what makes turing-sphere's implicit diffusion a single
80+divide. On a general surface it is not diagonal, and not even constant-
81+coefficient, so that divide has to become a solve.
82+
83+The models split the operator:
84+
85+```
86+lap_g = lap_s + dlap
87+```
88+
89+with `lap_s` the round-sphere one. `(I - dt*D*lap_s)` is still exactly
90+invertible, so the implicit step
91+
92+```
93+(I - dt*D*lap_g) Unew = B
94+```
95+
96+rearranges into a fixed point that keeps the whole geometry on the right-hand
97+side,
98+
99+```
100+Unew = (B + dt*D*dlap(Unew)) ./ (1 + dt*D*lam)
101+```
102+
103+and the loop iterates it from the round-sphere answer. That is preconditioned
104+Richardson, with the operator we can invert exactly as the preconditioner; it
105+converges while `dt*D*dlap` stays small against `(I - dt*D*lap_s)`, which is
106+what would keep the cost to a few transforms per step rather than a full
107+elliptic solve. Written out, the whole of
108+[`models/schnakenberg.m`](models/schnakenberg.m)'s step is:
109+
110+```matlab
111+function [Un, Vn, u, v] = step(U, V, lam, gx, gy, gz, a, b, D1, D2, dt, niter)
112+ u = synth(U);
113+ v = synth(V);
114+ uuv = u .* u .* v;
115+
116+ Bu = U + dt * analys(a - u + uuv);
117+ Bv = V + dt * analys(b - uuv);
118+
119+ Un = Bu ./ (1 + (dt * D1) * lam);
120+ Vn = Bv ./ (1 + (dt * D2) * lam);
121+
122+ for k = 1:niter
123+ dLu = 0 * Un; % <- the placeholder
124+ dLv = 0 * Vn;
125+ Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
126+ Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lam);
127+ end
128+end
129+```
130+
131+Written this way rather than as a residual correction on purpose: with `dlap`
132+zero, every iterate is *bit for bit* the first line, with no cancellation to
133+round differently. So the sphere case is not "close to" turing-sphere, it is
134+the same arithmetic, and the tests assert exactly that — the state after 20
135+steps is identical at 0, 1 and 4 iterations.
136+
137+### The geometry is not in the operator yet
138+
139+What belongs where `dLu` is now is `dlap = lap_g - lap_s` applied to the current
140+iterate. Getting it needs two things this repo does not have:
141+
142+1. **The induced metric**, `g_ij = ∂_i X · ∂_j X` for `X = (gx, gy, gz)`. The
143+ geometry is static and low-degree, so this is a one-off precomputation, not
144+ per-step work — but it needs θ- and φ-derivatives of the embedding.
145+2. **Surface derivatives of the field**, per iteration. In the round frame this
146+ is the spheroidal transform pair — SHTNS's `SHsph_to_spat` and
147+ `spat_to_SHsph`, i.e. `grad_s` and `div_s` — which lets the operator be
148+ written as `div_s(A grad_s f)` with `A` built from the metric, with no
149+ explicit `1/sin θ` to go singular at the poles.
150+
151+Both need Legendre *derivative* tables, which the vendored WGSL transforms under
152+[`src/sht/`](src/sht/) do not implement — they are scalar synthesis and analysis
153+only. That is the missing piece, and it is a substantial addition to the
154+transforms rather than a change to the models. Until it lands, the models take
155+`gx, gy, gz` (the surface on the grid) and `Gx, Gy, Gz` (the same surface as
156+coefficients) as arguments and do not use them, and the app says so.
157+
158+### `for` loops, unrolled
159+
160+A plan is a fixed list of GPU operations with no branching, which is what makes
161+a timestep pure command recording — one submit, no CPU in the loop. A counted
162+loop still fits: the planner
163+([`src/mgpu/plan.ts`](src/mgpu/plan.ts)) unrolls it, planning the body once per
164+iteration.
165+
166+Nothing else had to change for that, because numbl gives a variable one cName
167+for every assignment to it: the buffer an iteration writes is the buffer the
168+next one reads, which is exactly a loop-carried value. The loop variable gets no
169+buffer at all — it is bound as a derived scalar to that iteration's literal, so
170+a kernel reading `k` folds the number in.
171+
172+Two consequences worth stating:
173+
174+- **The bounds must be known when the model compiles.** `niter` is supplied as a
175+ fixed scalar rather than a tunable one, so changing it recompiles — unlike a
176+ parameter, which is a uniform. A runtime bound is refused at compile time with
177+ a source position, not silently mis-compiled, and there is a test for that.
178+- **Fusion survives.** numbl's inline pass recurses into loop bodies, so a line
179+ inside the loop is still one kernel. It runs there with no protected names,
180+ though, which means an assignment whose only visible use is later in the same
181+ body can be elided — correct for a body-local temp, wrong if something outside
182+ the loop wanted it. [`src/mgpu/compile.ts`](src/mgpu/compile.ts) snapshots what
183+ each loop body assigns before the pass and refuses the ones that escape, so
184+ that case is a compile error rather than a stale read.
185+
186+Unrolling is exactly linear in the trip count: 2 GPU ops per species per
187+iteration, asserted in the tests.
188+
189+## MATLAB, compiled to WebGPU
190+
191+Unchanged from turing-sphere, and it now compiles the geometry files too. numbl
192+parses and lowers each function for the concrete argument types of the current
193+grid; its inline pass folds single-use temps back into their consumer, so one
194+line of MATLAB becomes one expression tree; and this repo emits one WGSL compute
195+kernel per element-wise statement
196+([`src/mgpu/wgsl.ts`](src/mgpu/wgsl.ts)). `synth` / `analys` are external
197+operations whose type rules numbl learns from a `.mtoc2.js` workspace file, and
198+which the backend maps onto the spherical-harmonic pipelines. Anything it cannot
199+express is refused at compile time with a source position.
200+
201+The Schnakenberg step above compiles to 17 GPU operations at one solve
202+iteration: 4 transforms, 11 generated kernels, and 2 buffer copies feeding the
203+new state back.
204+
205+Two consequences carried over:
206+
207+- **The step is synchronous.** WebGPU's encode path is synchronous and every
208+ pipeline is built once at compile time, so a timestep is pure command
209+ recording; the only `await` in the loop is the single readback per rendered
210+ frame.
211+- **Parameters are uniforms, not constants.** Moving a slider rewrites a small
212+ buffer instead of triggering a recompile. Editing the MATLAB recompiles;
213+ changing `dt` does not. `niter` is the deliberate exception, above.
214+
215+## Provenance
216+
217+- **turing-sphere**, which this is a fork of: the solver, the transforms
218+ backend, the compilation path, the benchmarks and the analytic tests.
219+- **Transforms:** [shtns-webgpu](https://github.com/concept-collection/shtns-webgpu) —
220+ fp32 spherical harmonic transforms in WGSL compute shaders, modeled on
221+ [SHTNS](https://nschaeff.bitbucket.io/shtns/). Vendored under
222+ [`src/sht/`](src/sht/) (CECILL-2.1), including the f64 CPU reference transform
223+ used for testing.
224+- **Rendering:** three.js meshes with per-vertex colormaps, adapted from the
225+ `SphereEmbedding` view in
226+ [figpack](https://github.com/flatironinstitute/figpack)'s experimental
227+ extension package ([`src/render/`](src/render/)). That view displays a
228+ time-varying embedded geometry with fields on it, which is the same picture
229+ this draws — including its sphere/surface morph, which turing-sphere had
230+ dropped as having nothing to morph to.
231+
232+turing-sphere additionally carries a comparison against a native build of
233+upstream SHTNS ([`bench/shtns/`](https://github.com/concept-collection/turing-sphere/tree/main/bench/shtns)).
234+That is not duplicated here: the transforms are the same code, and its C-side
235+transcription of the model would have to be maintained against a step this
236+project intends to change.
237+
238+Because the algorithm is compiled to compute shaders, **WebGPU is required** —
239+there is no CPU fallback (the f64 CPU transform remains, for tests).
240+
241+## Numerics
242+
243+- Grid: Gauss–Legendre × equispaced-φ, dealiased for the cubic reactions with
244+ the `(pdeg+1)` rule: `nlat ≥ ((pdeg+1)·lmax+1)/2`, `nphi ≥ (pdeg+1)·lmax+1`
245+ (rounded up to a power of two for the GPU FFT path). At the default lmax 63
246+ that is a 128×256 grid.
247+- Spectral layout: SHTNS conventions — orthonormal + Condon–Shortley, complex
248+ coefficients for m ≥ 0, m-major ordering.
249+- fp32 transforms introduce ~1e-6 relative error per step; for pattern formation
250+ from 1e-2 seeded noise this is inconsequential. The geometry goes through one
251+ analysis/synthesis round trip and picks up the same round-off: the unit sphere
252+ comes back with radius 1 to ~2e-5 under Dawn, ~4e-4 under SwiftShader.
253+- The shipped geometries are all degree ≤ 5, far below any lmax the app offers,
254+ so band-limiting removes nothing from them. A shape you write yourself may not
255+ be so lucky — see the note in [`geometries/bumpy.m`](geometries/bumpy.m).
256+
257+## Desktop vs browser
258+
259+[`scripts/bench.ts`](scripts/bench.ts) runs the same thing the app runs — same
260+`.m`, same generated WGSL, same transforms — from Node on desktop WebGPU (Google
261+Dawn), and the app prints the command line that reproduces whatever it is
262+currently simulating:
263+
264+```
265+npm run bench -- --preset schnak-spots --geometry ellipsoid --lmax 63 --niter 1 \
266+ --steps 2000 --seed 1 --a 0.1 --b 0.9 --D1 0.0004 --D2 0.008 --dt 0.05 \
267+ --gax 1.5 --gay 1 --gaz 0.6
268+```
269+
270+Copy it from under the stats line and compare the `ms/step` it reports with the
271+app's. Both sides go through the one shared
272+[`src/bench/runSpec.ts`](src/bench/runSpec.ts) — the app formats a run into that
273+command, the benchmark parses it back — so there is no second copy of the
274+defaults for the two runs to drift apart on. Geometry parameters take a `g`
275+prefix (`--gwaist`) so a shape parameter can never collide with a model one.
276+
277+The app reports **two** numbers and only the first is comparable to the
278+benchmark: `solver` is the batch of steps alone, waited for but not read back;
279+`ms/frame` additionally carries a GPU→CPU readback per species, the
280+colormapping, and the vertex upload. Those per-frame costs are fixed and do not
281+shrink when the GPU gets faster, so on a quick GPU a frame can easily cost ten
282+times the steps inside it. That is expected and is not the solver being slower
283+in the browser.
284+
285+To attribute the gap rather than guess at it:
286+
287+```
288+node scripts/compare-perf.mjs [--lmax 63] [--steps 300]
289+```
290+
291+measures the same solver work in both — batched, nothing read back, no rendering
292+on either side — and reports each with its CPU-encoding share, the Fourier
293+stage, and the adapter. It stops you first if the two are not even the same
294+device, which is a common cause of "the browser is much slower". Both sides
295+resolve the geometry and the iteration count from the same constants, because
296+the iteration count is unrolled into the step and a mismatch would compare two
297+different amounts of work.
298+
299+The app's **Benchmark** button runs the same measurement in the page, plus the
300+**ramp** — the first third of the run against the last. GPUs downclock when
301+idle and an animation-paced loop leaves them idle most of every frame, so a
302+large ramp means the steady-state number is limited by clocks rather than work.
303+
304+### Is it really the same computation?
305+
306+```
307+node scripts/compare-env.mjs [--lmax 31] [--steps 200] [--preset schnak-spots]
308+```
309+
310+runs one identical spec on the desktop and in a real browser and compares the
311+final spectral state. The pipeline is deterministic given (model source,
312+geometry, parameters, lmax, niter, seed, steps), so the two should agree to fp32
313+round-off — not bit for bit, since GPUs differ in fused-multiply-add and other
314+latitude fp32 allows. It also reports which Fourier stage each side chose, since
315+FFT and DFT are genuinely different algorithms that round differently.
316+
317+Desktop WebGPU comes from the `webgpu` package (prebuilt Dawn, ~70 MB), an
318+optional dependency so that an unsupported platform fails the install of that
319+package alone. Its binaries need glibc 2.29+. Other flags: `--steps`,
320+`--warmup`, `--batch`, `--json`, `--help`; `DAWN_FLAGS='backend=vulkan'`
321+(`;`-separated) passes Dawn options through.
322+
323+## Tests
324+
325+There is no second implementation of the solver to diff against, so the `.m`
326+path is checked against **closed-form answers** and against **exact structural
327+properties**. Four modules, run in both environments:
328+
329+[`test/analyticChecks.ts`](test/analyticChecks.ts) — cases whose evolution is
330+known exactly, run through the whole real pipeline. All three are statements
331+about the round sphere, so all three build on the sphere geometry:
332+
333+- **A** — a linear reaction leaves every mode independent, growing by exactly
334+ `(1 + dt*c) / (1 + dt*D*l(l+1))` per step. Pins the transform round trip, the
335+ eigenvalue mapping, the IMEX update and the state feedback at once. ~2e-7 over
336+ 20 steps.
337+- **B** — a nonlinear reaction on a uniform field stays uniform, so each step is
338+ exactly the scalar ODE map. 1.5e-8 over 25 steps.
339+- **C** — a 1e-6 perturbation of the Schnakenberg fixed point follows the
340+ linearized 2×2 IMEX recurrence, and `(l=24, m=7)` is confirmed unstable.
341+ Looser (~4e-3) because fp32 keeps about four digits of a perturbation that
342+ small.
343+
344+[`test/geometryChecks.ts`](test/geometryChecks.ts) — the surface and the loop:
345+
346+- every geometry compiles and closes; the sphere has radius 1 everywhere and is
347+ **exactly degree 1** in the harmonics, which is what makes the reference case
348+ exact rather than merely accurate;
349+- the peanut matches its own closed-form radial profile at every grid point, and
350+ **the same coefficients give the same surface on a 2× grid** — the 2× Gauss
351+ latitudes share no point with the 1× ones, so agreeing there is agreeing
352+ everywhere, which is what "rendered exactly, not subdivided" means;
353+- unrolling is **exactly linear** in the trip count, and the state after 20 steps
354+ is **bit-identical** at 0, 1 and 4 iterations;
355+- a runtime loop bound is refused at compile time;
356+- swapping the surface mid-run leaves the spectral state untouched.
357+
358+[`test/modelChecks.ts`](test/modelChecks.ts) compiles every model the app offers
359+and asserts **how many kernels it compiles to**, split into the base step and
360+what one solve iteration adds. That is a fusion guard: if numbl's inline pass
361+stops folding, the results stay correct while every operator becomes its own
362+dispatch, which is invisible in the numbers.
363+
364+[`test/transformChecks.ts`](test/transformChecks.ts) compares the WGSL transforms
365+against shtns-webgpu's f64 CPU twin.
366+
367+- `npm run test:node` — under Dawn on the desktop, via `vite-node`. Needs a GPU;
368+ `--skip-without-gpu` lets a machine without one say so and move on (which is
369+ what CI does, since the browser suite covers the same modules).
370+- `npm run test:gpu` — builds and drives headless Chrome, on SwiftShader in CI.
371+ Also runs the soak. A few geometry tolerances are set by SwiftShader's fp32,
372+ which is about an order of magnitude looser than Dawn's.
373+
374+Other commands:
375+
376+- `npm run bench -- --help` — the desktop benchmark.
377+- `npm run bench:sht -- --help` — the transforms alone, no solver.
378+- `npx vite-node scripts/diagnose-sht.ts` — when the transform tests fail on a
379+ GPU, say *which* stage is wrong.
380+- `npx vite-node scripts/diagnose-leg.ts [--m 0]` — read the Legendre recurrence
381+ out of the production shader term by term.
382+- `npx vite-node scripts/longrun-node.ts [lmax]` — run to t = 100 and confirm the
383+ pattern saturates rather than decaying or diverging.
384+- `node scripts/soak.mjs [steps] [lmax]` — drive the demo for many steps,
385+ sampling JS heap and catching crashes.
386+- `node scripts/screenshot.mjs out.png [light|dark] [minSteps]` — screenshot the
387+ demo after a number of steps.
388+- `node scripts/check-live.mjs [url]` — smoke-check a deployed URL.
389+- `test.html?soak=<steps>&lmax=<n>` — solver-only soak with no rendering.
390+
391+## Development
392+
393+```
394+npm install
395+npm run dev # local dev server
396+npm run build # type-check + production build to dist/
397+```
398+
399+### The numbl dependency
400+
401+numbl is a local `file:../../numbl` dependency, so a sibling checkout of
402+[numbl](https://github.com/flatironinstitute/numbl) is required. We use its
403+compiler internals — parser, lowerer, IR, inline pass — which its package
404+`exports` map does not publish, so they are reached through the `numbl-src` path
405+alias in [`vite.config.ts`](vite.config.ts).
406+
407+The exact surface we depend on is written down in
408+[`src/mgpu/numbl.d.ts`](src/mgpu/numbl.d.ts) and TypeScript checks against
409+*that*, not against numbl's sources. This keeps this project's compiler settings
410+independent of numbl's, and means a change to one of those shapes upstream
411+breaks the build here with a clear diff rather than deep inside numbl's tree.
412+The `For` IR node is spelled out there, since the planner now walks it.
413+
414+CI clones numbl to the sibling path that the `file:` dependency expects, pinned
415+to a commit, with `--ignore-scripts` (npm runs a linked package's `prepare`
416+script, and numbl's is husky). numbl's own `node_modules` are not needed: the
417+slice we import is self-contained TypeScript.
418+
419+The `scripts/*.ts` entry points that touch the compiler go through `vite-node`,
420+so they resolve imports exactly as the browser build does. Plain `node` cannot:
421+numbl's sources import each other as `./foo.js` while the files are `.ts`.
422+
423+Deployed to GitHub Pages by `.github/workflows/deploy.yml` on push to `main`.
424+
425+## License
426+
427+CECILL-2.1 (inherited from SHTNS via shtns-webgpu, whose sources are vendored).
geometries/bumpy.madded+19−0View file
@@ -0,0 +1,19 @@
1+% Lobes around the equator, plus a gentle pear-shaped offset along z.
2+%
3+% The radial deformation sin(theta)^4 * cos(nlobe*phi) is the shape of a
4+% sectoral harmonic of order nlobe, concentrated at the equator and vanishing
5+% at both poles; the cos(theta) term breaks the north-south symmetry so the two
6+% ends differ.
7+%
8+% `nlobe` should be a whole number. A fraction makes cos(nlobe*phi) disagree
9+% with itself across the phi = 0 seam, which is not a closed surface; the
10+% analysis into coefficients will band-limit whatever results, but what comes
11+% back is not what this file says.
12+
13+function [gx, gy, gz] = shape(theta, phi, amp, nlobe, pear)
14+ st = sin(theta);
15+ r = 1 + amp * ((st .^ 4) .* cos(nlobe * phi)) + pear * cos(theta);
16+ gx = r .* (st .* cos(phi));
17+ gy = r .* (st .* sin(phi));
18+ gz = r .* cos(theta);
19+end
geometries/ellipsoid.madded+13−0View file
@@ -0,0 +1,13 @@
1+% A triaxial ellipsoid: the sphere with each axis scaled independently.
2+%
3+% Still degree-1 in the spherical harmonics — scaling a coordinate scales its
4+% coefficients — so the surface the solver sees is exactly the one written
5+% here, with nothing lost to band-limiting. The simplest geometry that is not
6+% the sphere, and the one whose curvature is easiest to reason about.
7+
8+function [gx, gy, gz] = shape(theta, phi, ax, ay, az)
9+ st = sin(theta);
10+ gx = ax * (st .* cos(phi));
11+ gy = ay * (st .* sin(phi));
12+ gz = az * cos(theta);
13+end
geometries/peanut.madded+18−0View file
@@ -0,0 +1,18 @@
1+% A dumbbell: a radial profile that pinches at the equator, stretched along z.
2+%
3+% The radius is 1 - waist*sin(theta)^2, so the poles keep radius 1 and the
4+% equator narrows to 1 - waist. Multiplying it into the unit-sphere coordinates
5+% raises the degree to 3, which is still far below any lmax the app offers, so
6+% the band-limited surface is again the one written here.
7+%
8+% Two regions of positive curvature joined by a saddle — the first geometry in
9+% this list where the Laplace-Beltrami operator differs from the round one in a
10+% way that should visibly change where a pattern wants to put its spots.
11+
12+function [gx, gy, gz] = shape(theta, phi, waist, stretch)
13+ st = sin(theta);
14+ r = 1 - waist * (st .^ 2);
15+ gx = r .* (st .* cos(phi));
16+ gy = r .* (st .* sin(phi));
17+ gz = (1 + stretch) * (r .* cos(theta));
18+end
geometries/sphere.madded+20−0View file
@@ -0,0 +1,20 @@
1+% The unit sphere.
2+%
3+% A geometry file defines one function, `shape`, giving the Cartesian
4+% coordinates of the surface over the solver's (theta, phi) grid. Both inputs
5+% are npts x 1 grid fields, as are all three outputs, and every line is
6+% element-wise MATLAB compiled to a WebGPU kernel — the same backend the models
7+% go through.
8+%
9+% The host then analyses the result into spherical-harmonic coefficients, which
10+% is the form the geometry is carried in: band-limited at lmax, evaluable on
11+% any grid. For this file that costs nothing, because x, y and z ARE degree-1
12+% harmonics — which is what makes this the case where everything downstream
13+% reduces exactly to the round-sphere solver.
14+
15+function [gx, gy, gz] = shape(theta, phi)
16+ st = sin(theta);
17+ gx = st .* cos(phi);
18+ gy = st .* sin(phi);
19+ gz = cos(theta);
20+end
index.htmladded+319−0View file
@@ -0,0 +1,319 @@
1+<!doctype html>
2+<html lang="en">
3+ <head>
4+ <meta charset="utf-8" />
5+ <meta name="viewport" content="width=device-width, initial-scale=1" />
6+ <link rel="icon" href="data:image/svg+xml,<svg xmlns=%22http://www.w3.org/2000/svg%22 viewBox=%220 0 100 100%22><circle cx=%2250%22 cy=%2250%22 r=%2245%22 fill=%22%232a7f62%22/><circle cx=%2235%22 cy=%2238%22 r=%2211%22 fill=%22%23f5d547%22/><circle cx=%2265%22 cy=%2258%22 r=%229%22 fill=%22%23f5d547%22/><circle cx=%2248%22 cy=%2274%22 r=%227%22 fill=%22%23f5d547%22/><circle cx=%2268%22 cy=%2230%22 r=%226%22 fill=%22%23f5d547%22/></svg>" />
7+ <title>turing-surface — reaction-diffusion on closed surfaces, live in the browser</title>
8+ <style>
9+ :root {
10+ --bg: #ffffff;
11+ --ink: #1f2328;
12+ --ink-2: #57606a;
13+ --line: #d0d7de;
14+ --accent: #0969da;
15+ --sphere-bg: #f4f6f8;
16+ --tok-com: #6e7781;
17+ --tok-str: #0a3069;
18+ --tok-num: #0550ae;
19+ --tok-kw: #cf222e;
20+ --tok-ext: #8250df;
21+ --warn-bg: #fff8e5;
22+ --warn-line: #e3c37a;
23+ --warn-edge: #bf8700;
24+ color-scheme: light dark;
25+ }
26+ @media (prefers-color-scheme: dark) {
27+ :root {
28+ --bg: #14171a;
29+ --ink: #e6e9ec;
30+ --ink-2: #9aa4af;
31+ --line: #333b44;
32+ --accent: #58a6ff;
33+ --sphere-bg: #14161c;
34+ --tok-com: #8b949e;
35+ --tok-str: #a5d6ff;
36+ --tok-num: #79c0ff;
37+ --tok-kw: #ff7b72;
38+ --tok-ext: #d2a8ff;
39+ --warn-bg: #2b2410;
40+ --warn-line: #6b5518;
41+ --warn-edge: #e3b341;
42+ }
43+ }
44+ body {
45+ margin: 0;
46+ background: var(--bg);
47+ color: var(--ink);
48+ font: 15px/1.5 system-ui, -apple-system, sans-serif;
49+ }
50+ main { max-width: 1100px; margin: 0 auto; padding: 20px 16px 48px; }
51+ h1 { font-size: 20px; margin: 0 0 2px; }
52+ .sub { color: var(--ink-2); margin: 0 0 12px; font-size: 13px; }
53+ .sub a { color: var(--accent); }
54+ .controls {
55+ display: flex; flex-wrap: wrap; gap: 8px 14px; align-items: center;
56+ padding: 6px 0;
57+ }
58+ /* display:flex above would otherwise override the UA's [hidden] rule */
59+ .controls[hidden] { display: none; }
60+ .controls label { color: var(--ink-2); font-size: 13px; white-space: nowrap; }
61+ select, input[type="number"], button {
62+ font: inherit; font-size: 13px;
63+ color: var(--ink); background: var(--bg);
64+ border: 1px solid var(--line); border-radius: 6px;
65+ padding: 4px 8px;
66+ }
67+ input[type="number"] { width: 6em; }
68+ input[type="range"] { width: 8em; vertical-align: middle; accent-color: var(--accent); }
69+ button { cursor: pointer; }
70+ #geomnote { margin-top: 4px; }
71+ button:hover { border-color: var(--accent); }
72+ button.primary { border-color: var(--accent); color: var(--accent); font-weight: 600; min-width: 5.5em; }
73+ #panels {
74+ display: flex; flex-wrap: wrap; gap: 14px; margin-top: 12px;
75+ }
76+ .panel {
77+ flex: 1 1 320px; min-width: 280px;
78+ border: 1px solid var(--line); border-radius: 8px; overflow: hidden;
79+ display: flex;
80+ }
81+ .sphere-box { flex: 1; aspect-ratio: 1 / 1; max-height: 70vh; position: relative; }
82+ .species-tag {
83+ position: absolute; top: 8px; left: 10px; z-index: 2;
84+ font-size: 15px; font-weight: 600; color: #fff;
85+ background: rgba(0, 0, 0, 0.45);
86+ padding: 1px 10px; border-radius: 12px;
87+ pointer-events: none;
88+ }
89+ .colorbar {
90+ display: flex; flex-direction: column; align-items: center; justify-content: center;
91+ gap: 4px; padding: 8px 4px; background: var(--sphere-bg);
92+ width: 52px; flex: none; box-sizing: border-box;
93+ }
94+ .colorbar canvas { border: 1px solid var(--line); border-radius: 2px; }
95+ .colorbar-label { font-size: 11px; color: var(--ink-2); font-variant-numeric: tabular-nums; }
96+ .stats { margin-top: 10px; font-size: 13px; color: var(--ink-2); font-variant-numeric: tabular-nums; }
97+ .stats b { color: var(--ink); font-weight: 600; }
98+ .cli {
99+ margin-top: 10px; border: 1px solid var(--line); border-radius: 8px;
100+ overflow: hidden;
101+ }
102+ .cli-head {
103+ display: flex; gap: 10px; align-items: center; justify-content: space-between;
104+ padding: 6px 10px; font-size: 12px; color: var(--ink-2);
105+ background: var(--sphere-bg); border-bottom: 1px solid var(--line);
106+ }
107+ .cli-head button { padding: 2px 10px; font-size: 12px; }
108+ #cmd {
109+ display: block; padding: 8px 10px; white-space: pre-wrap;
110+ font: 12px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
111+ color: var(--ink); user-select: all;
112+ }
113+ #blurb { margin-top: 4px; font-size: 13px; color: var(--ink-2); }
114+ .warn {
115+ margin: 0 0 14px;
116+ padding: 10px 14px;
117+ border: 1px solid var(--warn-line);
118+ border-left: 5px solid var(--warn-edge);
119+ border-radius: 6px;
120+ background: var(--warn-bg);
121+ color: var(--ink);
122+ font-size: 13.5px; line-height: 1.5;
123+ }
124+ .warn b { color: var(--warn-edge); }
125+ #err { color: #b35900; white-space: pre-wrap; font-size: 13px; }
126+ .editor {
127+ margin-top: 12px; border: 1px solid var(--line); border-radius: 8px;
128+ overflow: hidden;
129+ }
130+ .editor-head {
131+ display: flex; gap: 10px; align-items: center; justify-content: space-between;
132+ padding: 6px 10px; font-size: 12px; color: var(--ink-2);
133+ background: var(--sphere-bg); border-bottom: 1px solid var(--line);
134+ }
135+ .editor-head button { padding: 2px 10px; font-size: 12px; }
136+ /* Source and compiled-op list side by side, so the editor gets the
137+ height rather than sharing it with the list below. */
138+ .editor-body { display: flex; align-items: stretch; }
139+ .editor-code { position: relative; flex: 1 1 62%; min-width: 0; height: 34em; }
140+ /* The overlay and the textarea must agree on every metric that affects
141+ where a character lands. Keep these two rules together. */
142+ .editor-code > pre,
143+ .editor-code > textarea {
144+ margin: 0; padding: 10px 12px; border: 0;
145+ box-sizing: border-box; width: 100%; height: 100%;
146+ font: 12.5px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
147+ tab-size: 2;
148+ white-space: pre; overflow-wrap: normal;
149+ }
150+ #highlight {
151+ position: absolute; inset: 0; overflow: hidden;
152+ pointer-events: none; background: var(--bg); color: var(--ink);
153+ }
154+ #source {
155+ position: relative; z-index: 1; display: block;
156+ resize: none; overflow: auto;
157+ background: transparent; color: transparent; caret-color: var(--ink);
158+ }
159+ #source:focus { outline: none; }
160+ /* Transparent text means the selection must be see-through, or selected
161+ code would be invisible. */
162+ #source::selection { background: color-mix(in srgb, var(--accent) 28%, transparent); }
163+ #compiled {
164+ flex: 1 1 38%; min-width: 0; margin: 0; padding: 10px 12px; overflow: auto;
165+ border-left: 1px solid var(--line); background: var(--sphere-bg);
166+ font: 12px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
167+ color: var(--ink-2); white-space: pre;
168+ }
169+ @media (max-width: 860px) {
170+ .editor-body { flex-direction: column; }
171+ .editor-code { flex: none; height: 26em; }
172+ #compiled { border-left: 0; border-top: 1px solid var(--line); max-height: 12em; }
173+ }
174+ .tok-com { color: var(--tok-com); }
175+ .tok-str { color: var(--tok-str); }
176+ .tok-num { color: var(--tok-num); }
177+ .tok-kw { color: var(--tok-kw); font-weight: 600; }
178+ .tok-ext { color: var(--tok-ext); }
179+ </style>
180+ </head>
181+ <body>
182+ <main>
183+ <h1>turing-surface</h1>
184+ <p class="warn">
185+ <b>⚠ Work in progress: the geometry is drawn, not solved on.</b>
186+ The Laplace–Beltrami operator in the models is still the <em>round
187+ sphere's</em> — the term that would carry the shape is a placeholder
188+ that is identically zero. So on anything but the sphere, what you see
189+ is the sphere's pattern painted onto that surface, not the pattern that
190+ surface would actually grow. The spot size does not follow the
191+ curvature, and it will not until that placeholder is filled in.
192+ Everything else — the embedding, the mesh, the unrolled implicit solve
193+ the correction plugs into — is real and in place.
194+ </p>
195+ <p class="sub">
196+ Reaction-diffusion on closed surfaces, solved live with spherical harmonics:
197+ implicit spectral diffusion + explicit reaction (IMEX Euler), transforms on WebGPU via
198+ <a href="https://github.com/concept-collection/shtns-webgpu">shtns-webgpu</a>.
199+ The surface is an embedding of the sphere given by spherical-harmonic
200+ coefficients for x, y and z. Both the solver and the shape are the
201+ MATLAB below — <a href="https://numbl.org">numbl</a> parses and lowers
202+ them in your browser, and each line becomes a WebGPU compute kernel.
203+ Edit either and watch it change. Drag to rotate.
204+ </p>
205+ <div class="controls">
206+ <label>preset
207+ <select id="model"></select>
208+ </label>
209+ <label title="The surface, as a spherical-harmonic embedding. Rendered exactly; NOT yet in the Laplace-Beltrami operator, which is still the round-sphere one.">geometry
210+ <select id="geometry"></select>
211+ </label>
212+ <label title="Blend between the unit sphere (0) and the surface (1). Display only — the solver's grid is the sphere's parametrization either way.">morph
213+ <input type="range" id="morph" min="0" max="1" step="0.01" value="1" />
214+ </label>
215+ <label title="Iterations of the implicit diffusion solve, unrolled into the compiled step. Changing it recompiles. The correction it iterates is zero until the geometry reaches the operator, so today every value gives the same answer.">solve iters
216+ <select id="niter">
217+ <option value="0">0</option>
218+ <option value="1" selected>1</option>
219+ <option value="2">2</option>
220+ <option value="4">4</option>
221+ <option value="8">8</option>
222+ </select>
223+ </label>
224+ <label>lmax
225+ <select id="lmax">
226+ <option value="31">31</option>
227+ <option value="63" selected>63</option>
228+ <option value="127">127</option>
229+ <option value="255">255</option>
230+ </select>
231+ </label>
232+ <label title="Display only — the solution is evaluated on a finer grid for rendering; the solver and its grid are unchanged. Auto oversamples coarse grids and leaves fine ones alone.">display oversampling
233+ <select id="oversample">
234+ <option value="auto" selected>auto</option>
235+ <option value="1">off</option>
236+ <option value="2">2×</option>
237+ <option value="4">4×</option>
238+ <option value="8">8×</option>
239+ </select>
240+ </label>
241+ <label>colormap
242+ <select id="colormap"></select>
243+ </label>
244+ <button id="runpause" class="primary">Run</button>
245+ <button id="benchmark">Benchmark</button>
246+ <button id="reseed">Re-seed</button>
247+ <button id="resetview">Reset view</button>
248+ <button id="movietoggle" title="Export the run as an MP4 movie">Export movie</button>
249+ </div>
250+ <div class="controls" id="moviebar" hidden>
251+ <label title="Playback speed: simulation-time units per second of video — 1× plays one time unit per second">movie speed
252+ <select id="moviespeed">
253+ <option value="0.1">0.1×</option>
254+ <option value="0.5">0.5×</option>
255+ <option value="1">1×</option>
256+ <option value="3">3×</option>
257+ <option value="5">5×</option>
258+ <option value="10" selected>10×</option>
259+ <option value="20">20×</option>
260+ </select>
261+ </label>
262+ <label title="Rendered size of each sphere panel, in pixels — the video frame is the panels side by side plus the caption. Independent of the window size.">resolution
263+ <select id="movieres">
264+ <option value="480">480</option>
265+ <option value="640">640</option>
266+ <option value="768" selected>768</option>
267+ <option value="1080">1080</option>
268+ <option value="1440">1440</option>
269+ </select>
270+ </label>
271+ <label title="Slowly orbit the camera during the movie — one revolution per 2 minutes of video, starting from the current view (which is restored afterwards)">
272+ <input type="checkbox" id="movierotate" checked /> auto-rotate
273+ </label>
274+ <button id="movie" title="Recompute the run from t = 0 and download it as an MP4 movie">Export</button>
275+ </div>
276+ <div class="controls" id="params"></div>
277+ <div class="controls" id="geomparams"></div>
278+ <p id="geomnote" class="stats"></p>
279+ <div id="panels"></div>
280+ <p class="stats" id="stats"></p>
281+ <p class="stats" id="benchresult"></p>
282+ <div class="editor">
283+ <div class="editor-head">
284+ <span>
285+ <select id="editor-file" aria-label="file to edit"></select>
286+ <span id="editor-title"></span>
287+ </span>
288+ <span>
289+ <button id="recompile" type="button">Recompile</button>
290+ <button id="revert" type="button">Revert</button>
291+ </span>
292+ </div>
293+ <div class="editor-body">
294+ <div class="editor-code">
295+ <pre id="highlight" aria-hidden="true"></pre>
296+ <textarea
297+ id="source"
298+ spellcheck="false"
299+ autocomplete="off"
300+ autocapitalize="off"
301+ aria-label="model source (MATLAB)"
302+ ></textarea>
303+ </div>
304+ <pre id="compiled"></pre>
305+ </div>
306+ </div>
307+ <div class="cli">
308+ <div class="cli-head">
309+ <span>The same run on the desktop, through the same .m and the same kernels</span>
310+ <button id="copycmd" type="button">Copy</button>
311+ </div>
312+ <code id="cmd"></code>
313+ </div>
314+ <p id="blurb"></p>
315+ <p id="err"></p>
316+ </main>
317+ <script type="module" src="/src/main.ts"></script>
318+ </body>
319+</html>
models/allencahn.madded+26−0View file
@@ -0,0 +1,26 @@
1+% Allen-Cahn on a closed surface. One species: interfaces form and then coarsen
2+% until one domain swallows the surface.
3+%
4+% du/dt = eps2*lap_g(u) + u - u^3
5+%
6+% See models/schnakenberg.m for what the caller provides and for how the
7+% implicit solve is split between the round-sphere operator and the geometry.
8+
9+function [U, u] = init(noise)
10+ U = analys(noise);
11+ u = synth(U);
12+end
13+
14+function [Un, u] = step(U, lam, gx, gy, gz, eps2, dt, niter)
15+ u = synth(U);
16+
17+ Bu = U + dt * analys(u - u.^3);
18+ Un = Bu ./ (1 + (dt * eps2) * lam);
19+
20+ for k = 1:niter
21+ % ---- placeholder: dlap = lap_g - lap_s (see models/schnakenberg.m) ----
22+ dLu = 0 * Un;
23+ % ----------------------------------------------------------------------
24+ Un = (Bu + (dt * eps2) * dLu) ./ (1 + (dt * eps2) * lam);
25+ end
26+end
models/brusselator.madded+36−0View file
@@ -0,0 +1,36 @@
1+% Brusselator reaction-diffusion on a closed surface. Turing stripes and spots,
2+% from a smaller diffusivity contrast than Schnakenberg but a stiffer reaction.
3+%
4+% du/dt = D1*lap_g(u) + A - (B+1)*u + u^2*v
5+% dv/dt = D2*lap_g(v) + B*u - u^2*v
6+%
7+% See models/schnakenberg.m for what the caller provides and for how the
8+% implicit solve is split between the round-sphere operator and the geometry.
9+
10+function [U, V, u, v] = init(noise, A, B)
11+ U = analys(A + noise);
12+ V = analys((B / A) * ones(numel(noise), 1));
13+ u = synth(U);
14+ v = synth(V);
15+end
16+
17+function [Un, Vn, u, v] = step(U, V, lam, gx, gy, gz, A, B, D1, D2, dt, niter)
18+ u = synth(U);
19+ v = synth(V);
20+ uuv = u .* u .* v;
21+
22+ Bu = U + dt * analys(A - (B + 1) * u + uuv);
23+ Bv = V + dt * analys(B * u - uuv);
24+
25+ Un = Bu ./ (1 + (dt * D1) * lam);
26+ Vn = Bv ./ (1 + (dt * D2) * lam);
27+
28+ for k = 1:niter
29+ % ---- placeholder: dlap = lap_g - lap_s (see models/schnakenberg.m) ----
30+ dLu = 0 * Un;
31+ dLv = 0 * Vn;
32+ % ----------------------------------------------------------------------
33+ Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
34+ Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lam);
35+ end
36+end
models/schnakenberg.madded+69−0View file
@@ -0,0 +1,69 @@
1+% Schnakenberg reaction-diffusion on a closed surface.
2+%
3+% du/dt = D1*lap_g(u) + a - u + u^2*v
4+% dv/dt = D2*lap_g(v) + b - u^2*v
5+%
6+% lap_g is the Laplace-Beltrami operator of the surface, which the embedding
7+% (gx, gy, gz) determines. The reaction is explicit on the grid; diffusion is
8+% implicit, which on a general surface is no longer a divide — see `step`.
9+%
10+% Provided by the caller: synth/analys (the transforms), lam = l(l+1) per
11+% coefficient, gx/gy/gz (the surface, on the grid) and Gx/Gy/Gz (the same
12+% surface as spherical-harmonic coefficients), noise (the seeded perturbation),
13+% niter (iterations of the implicit solve), and the parameters. Grid fields are
14+% npts x 1; spectral fields are real 2 x nlm (row 1 real part, row 2
15+% imaginary), so no complex arithmetic is needed. Each function returns the new
16+% spectral state followed by the grid fields to display.
17+
18+function [U, V, u, v] = init(noise, a, b)
19+ us = a + b;
20+ vs = b / (us * us);
21+ U = analys(us + noise);
22+ V = analys(vs * ones(numel(noise), 1));
23+ u = synth(U);
24+ v = synth(V);
25+end
26+
27+function [Un, Vn, u, v] = step(U, V, lam, gx, gy, gz, a, b, D1, D2, dt, niter)
28+ u = synth(U);
29+ v = synth(V);
30+ uuv = u .* u .* v;
31+
32+ % Explicit reaction. Bu, Bv are the right-hand side of the implicit
33+ % diffusion solve (I - dt*D*lap_g) Unew = B.
34+ Bu = U + dt * analys(a - u + uuv);
35+ Bv = V + dt * analys(b - uuv);
36+
37+ % Split the surface Laplacian as lap_g = lap_s + dlap, where lap_s is the
38+ % round-sphere one. lap_s is diagonal in spherical-harmonic space with
39+ % eigenvalues -lam, so (I - dt*D*lap_s) inverts in a single divide and the
40+ % whole geometry sits inside dlap. That turns the implicit solve into
41+ %
42+ % Unew = (Bu + dt*D1*dlap(Unew)) ./ (1 + dt*D1*lam)
43+ %
44+ % which the loop below iterates from the round-sphere answer: preconditioned
45+ % Richardson, with the exactly invertible round-sphere operator as the
46+ % preconditioner. It converges while dt*D*dlap stays small against
47+ % (I - dt*D*lap_s), which is what would keep the cost to a few transforms.
48+ Un = Bu ./ (1 + (dt * D1) * lam);
49+ Vn = Bv ./ (1 + (dt * D2) * lam);
50+
51+ for k = 1:niter
52+ % ---- placeholder: dlap = lap_g - lap_s --------------------------------
53+ % The geometry enters HERE and nowhere else. What belongs here is the
54+ % Laplace-Beltrami operator of the embedding minus the round-sphere one,
55+ % which needs the induced metric (from derivatives of gx, gy, gz) and
56+ % surface derivatives of the field — transforms this project does not have
57+ % yet. See "The geometry is not in the operator yet" in the README.
58+ %
59+ % Until then dlap is identically zero. Not a stand-in that happens to be
60+ % small: it is exactly the round sphere, so this loop provably changes
61+ % nothing, every iterate equals the line above, and the scheme is bit for
62+ % bit the one turing-sphere runs. The surface is drawn, not solved on.
63+ dLu = 0 * Un;
64+ dLv = 0 * Vn;
65+ % ----------------------------------------------------------------------
66+ Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
67+ Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lam);
68+ end
69+end
package-lock.jsonadded+3209−0View file
This diff is 3,214 lines long and is not shown.
package.jsonadded+36−0View file
@@ -0,0 +1,36 @@
1+{
2+ "name": "turing-surface",
3+ "version": "0.1.0",
4+ "description": "Live reaction-diffusion (Turing patterns) on closed surfaces given by spherical-harmonic embeddings, spectral solver on WebGPU",
5+ "type": "module",
6+ "engines": {
7+ "node": ">=22.6"
8+ },
9+ "license": "CECILL-2.1",
10+ "scripts": {
11+ "dev": "vite",
12+ "build": "tsc --noEmit && vite build",
13+ "test:node": "vite-node scripts/test-node.ts",
14+ "test:gpu": "vite build && node scripts/test-gpu.mjs",
15+ "test": "npm run test:node && npm run test:gpu",
16+ "bench": "vite-node scripts/bench.ts",
17+ "bench:sht": "vite-node scripts/bench-sht.ts"
18+ },
19+ "dependencies": {
20+ "mp4-muxer": "^5.2.2",
21+ "numbl": "file:../../numbl",
22+ "three": "^0.183.0"
23+ },
24+ "optionalDependencies": {
25+ "webgpu": "^0.4.0"
26+ },
27+ "devDependencies": {
28+ "@types/node": "^26.1.1",
29+ "@types/three": "^0.185.1",
30+ "@webgpu/types": "^0.1.44",
31+ "puppeteer-core": "^23.0.0",
32+ "typescript": "^5.5.0",
33+ "vite": "^5.4.0",
34+ "vite-node": "^6.0.0"
35+ }
36+}
scripts/bench-sht.tsadded+452−0View file
@@ -0,0 +1,452 @@
1+/**
2+ * The transforms alone, on desktop WebGPU — the number to put next to upstream
3+ * SHTNS.
4+ *
5+ * npm run bench:sht -- --lmax 63 --steps 2000
6+ *
7+ * `npm run bench` measures a whole timestep of a .m model. This measures one
8+ * spectral -> grid -> spectral round trip and nothing else, which is what
9+ * bench/shtns/shtbench{,_gpu} --mode transform measures on the other side. The
10+ * solver does one of these per species per step, and profiling of the reference
11+ * implementation puts them at ~96% of its compute, so this is the comparison
12+ * that actually decides how fast the solver can be.
13+ *
14+ * The grid comes from the same rule the app uses, through the same
15+ * parseArgs/configForSpec as `npm run bench`, so --preset and --lmax mean here
16+ * exactly what they mean there. Nothing about the model is used beyond its
17+ * dealiasing degree.
18+ *
19+ * Like the solver benchmark it reports throughput (a batch of round trips
20+ * submitted together, awaited once) and latency (one per submit, for the
21+ * distribution).
22+ */
23+import { ShtPlan, requestShtDevice, describeAdapter, type ShtBinding } from '../src/sht/sht.ts';
24+import { lmIndex } from '../src/sht/layout.ts';
25+import { makeRand } from '../src/mgpu/noise.ts';
26+import { digestOf, formatDigest, relL2 } from '../src/mgpu/digest.ts';
27+import {
28+ parseArgs,
29+ configForSpec,
30+ modelForSpec,
31+ DEFAULT_LMAX,
32+ DEFAULT_SEED,
33+ DEFAULT_STEPS,
34+ DEFAULT_WARMUP,
35+ type RunSpec,
36+} from '../src/bench/runSpec.ts';
37+import { presets } from '../src/mgpu/registry.ts';
38+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
39+import { writeFileSync } from 'node:fs';
40+
41+const BENCH_SHT_COMMAND = 'npm run bench:sht --';
42+
43+const USAGE = `usage: npm run bench:sht -- [options]
44+
45+ --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
46+ --steps <n> timed round trips (default ${DEFAULT_STEPS})
47+ --warmup <n> untimed round trips first (default ${DEFAULT_WARMUP})
48+ --seed <n> seed of the initial spectrum (default ${DEFAULT_SEED})
49+ --batch <n> round trips per submit for the throughput number (default 16)
50+ --preset <key> only for its dealiasing degree, so the grid matches the
51+ solver benchmark's: ${presets.map((p) => p.key).join(' | ')}
52+ (default ${presets[0].key})
53+ --fourier <mode> auto | fft | dft (default auto)
54+ --digest after timing, re-run exactly --steps round trips from the
55+ seed and print a digest of the final spectrum
56+ --dump-state <f> like --digest, and write the spectrum to <f> as JSON, for
57+ scripts/compare-native.mjs to diff
58+ --json machine-readable output
59+ --help
60+
61+The native counterpart is
62+ bench/shtns/shtbench --mode transform --lmax <n> --steps <n> (CPU, fp64)
63+ bench/shtns/shtbench_gpu --mode transform --lmax <n> --steps <n> (CUDA, fp32)`;
64+
65+function fail(msg: string, code = 1): never {
66+ console.error(`bench:sht: ${msg}`);
67+ process.exit(code);
68+}
69+
70+// ---------------------------------------------------------------- arguments
71+const argv = process.argv.slice(2);
72+if (argv.includes('--help') || argv.includes('-h')) {
73+ console.log(USAGE);
74+ process.exit(0);
75+}
76+const wantJson = argv.includes('--json');
77+let batch = 16;
78+let fourier: 'auto' | 'fft' | 'dft' = 'auto';
79+let dumpState: string | null = null;
80+let wantDigest = false;
81+const rest: string[] = [];
82+for (let i = 0; i < argv.length; i++) {
83+ const a = argv[i];
84+ if (a === '--json') continue;
85+ if (a === '--digest') {
86+ wantDigest = true;
87+ continue;
88+ }
89+ const valued = (name: string): string | null => {
90+ if (a === `--${name}`) return argv[++i];
91+ if (a.startsWith(`--${name}=`)) return a.slice(name.length + 3);
92+ return null;
93+ };
94+ const b = valued('batch');
95+ if (b !== null) {
96+ batch = Number(b);
97+ continue;
98+ }
99+ const f = valued('fourier');
100+ if (f !== null) {
101+ if (f !== 'auto' && f !== 'fft' && f !== 'dft') fail(`--fourier must be auto|fft|dft`, 2);
102+ fourier = f;
103+ continue;
104+ }
105+ const d = valued('dump-state');
106+ if (d !== null) {
107+ dumpState = d;
108+ wantDigest = true;
109+ continue;
110+ }
111+ rest.push(a);
112+}
113+if (!Number.isInteger(batch) || batch < 1) fail('--batch must be an integer >= 1', 2);
114+
115+let spec: RunSpec;
116+try {
117+ spec = parseArgs(rest);
118+} catch (e) {
119+ fail(`${errMsg(e)}\n\n${USAGE}`, 2);
120+}
121+const cfg = configForSpec(spec);
122+
123+// -------------------------------------------------------------- the spectrum
124+/**
125+ * A seeded starting spectrum, uniform in [-1, 1). Deliberately the plainest
126+ * thing both sides can agree on bit for bit: mulberry32 only, no transcendental
127+ * functions, so a difference in the result is a difference in the transforms and
128+ * not in the input. The m = 0 imaginary parts are zeroed, since a real field has
129+ * none and the two libraries need not treat a coefficient that cannot occur
130+ * alike. Mirrors shtb_seeded_spectrum() in bench/shtns/spec.h.
131+ */
132+function seededSpectrum(lmax: number, mmax: number, nlm: number, seed: number): Float32Array {
133+ const rand = makeRand(seed);
134+ const qlm = new Float32Array(2 * nlm);
135+ for (let m = 0; m <= mmax; m++) {
136+ for (let l = m; l <= lmax; l++) {
137+ const lm = lmIndex(lmax, l, m);
138+ qlm[2 * lm] = 2 * rand() - 1;
139+ const im = 2 * rand() - 1;
140+ qlm[2 * lm + 1] = m === 0 ? 0 : im;
141+ }
142+ }
143+ return qlm;
144+}
145+
146+// ---------------------------------------------------------------------- run
147+let device: GPUDevice | null = null;
148+let plan: ShtPlan | null = null;
149+
150+try {
151+ const runtime = await installWebGpu();
152+ device = await requestShtDevice().catch((e: unknown) => {
153+ throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
154+ });
155+ const adapter = await describeAdapter(device);
156+ plan = await ShtPlan.create(device, cfg, { fourier });
157+ const nlm = plan.nlm;
158+ const npts = cfg.nlat * cfg.nphi;
159+
160+ // Two spectral buffers and one spatial one, so a round trip needs no copy:
161+ // round trips alternate direction, A -> spat -> B then B -> spat -> A.
162+ const mk = (label: string, size: number) =>
163+ device!.createBuffer({
164+ label,
165+ size,
166+ usage:
167+ GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
168+ });
169+ const qlm: [GPUBuffer, GPUBuffer] = [mk('sht-bench-qa', 8 * nlm), mk('sht-bench-qb', 8 * nlm)];
170+ const spat = mk('sht-bench-spat', 4 * npts);
171+ const readback = device.createBuffer({
172+ label: 'sht-bench-readback',
173+ size: 8 * nlm,
174+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
175+ });
176+ // Built once, at plan time — a bind group per round trip would be measuring
177+ // bind-group creation.
178+ const synth: [ShtBinding, ShtBinding] = [
179+ plan.createSynthBinding(qlm[0], spat),
180+ plan.createSynthBinding(qlm[1], spat),
181+ ];
182+ const analys: [ShtBinding, ShtBinding] = [
183+ plan.createAnalysBinding(spat, qlm[1]),
184+ plan.createAnalysBinding(spat, qlm[0]),
185+ ];
186+
187+ let cur = 0;
188+ /** Record `n` round trips into one submission. Returns nothing; the result is
189+ * in qlm[cur] once the queue has drained. */
190+ const submit = (n: number): void => {
191+ const enc = device!.createCommandEncoder({ label: 'sht-bench' });
192+ const pass = enc.beginComputePass({ label: 'sht-bench' });
193+ for (let i = 0; i < n; i++) {
194+ plan!.encodeSynthInto(pass, synth[cur]);
195+ plan!.encodeAnalysInto(pass, analys[cur]);
196+ cur ^= 1;
197+ }
198+ pass.end();
199+ device!.queue.submit([enc.finish()]);
200+ };
201+ /** `submit` in chunks, so an arbitrary round-trip count does not build one
202+ * command buffer with tens of thousands of dispatches in it. */
203+ const submitAll = (n: number, chunk = batch): void => {
204+ for (let done = 0; done < n; done += chunk) submit(Math.min(chunk, n - done));
205+ };
206+ const seed = (): void => {
207+ cur = 0;
208+ const q0 = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);
209+ device!.queue.writeBuffer(qlm[0], 0, q0 as Float32Array<ArrayBuffer>);
210+ };
211+ const readSpectrum = async (): Promise<Float32Array> => {
212+ const enc = device!.createCommandEncoder({ label: 'sht-bench-read' });
213+ enc.copyBufferToBuffer(qlm[cur], 0, readback, 0, 8 * nlm);
214+ device!.queue.submit([enc.finish()]);
215+ await readback.mapAsync(GPUMapMode.READ);
216+ const out = new Float32Array(readback.getMappedRange().slice(0));
217+ readback.unmap();
218+ return out;
219+ };
220+ const done = (): Promise<undefined> => device!.queue.onSubmittedWorkDone();
221+
222+ // What every report of this run says about itself, whether it succeeded, failed
223+ // early, or is being written to a state file for compare-native.mjs to diff.
224+ const identity = {
225+ mode: 'transform',
226+ spec: {
227+ preset: spec.preset,
228+ lmax: spec.lmax,
229+ seed: spec.seed,
230+ steps: spec.steps,
231+ warmup: spec.warmup,
232+ },
233+ backend: { library: 'shtns-webgpu (src/sht)', runtime, adapter, precision: 'fp32' },
234+ grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm },
235+ fourier: plan.fourierMode,
236+ };
237+
238+ if (!wantJson) {
239+ console.log('turing-surface bench:sht — transforms only, no solver, no rendering\n');
240+ console.log(
241+ ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${nlm.toLocaleString()}` +
242+ ` (dealiased for ${modelForSpec(spec).key}, pdeg ${modelForSpec(spec).pdeg})`,
243+ );
244+ console.log(` step 1 synthesis + 1 analysis (one round trip)`);
245+ console.log(` fourier ${plan.fourierMode.toUpperCase()} stage`);
246+ console.log(` backend WebGPU fp32${adapter ? ` — ${adapter}` : ''}\n ${runtime}`);
247+ console.log(
248+ ` run ${spec.warmup} warmup + ${spec.steps} timed round trips, seed ${spec.seed}\n`,
249+ );
250+ }
251+
252+ /*
253+ * Before timing anything: does one round trip work on this device?
254+ *
255+ * analys(synth(q)) is the identity for a band-limited q — exact Gauss
256+ * quadrature, nphi past the aliasing limit — so a single round trip should
257+ * return the input to fp32 round-off, and this is the sharpest check the
258+ * transforms are working at all here. Doing it separately means a bad result
259+ * says *which* it is: broken from the first transform, or drifted over the
260+ * thousands of iterations the timing run does. Otherwise that is a manual
261+ * bisection on --steps.
262+ */
263+ seed();
264+ const input = seededSpectrum(cfg.lmax, cfg.mmax, nlm, spec.seed);
265+ submit(1);
266+ await done();
267+ const afterOne = await readSpectrum();
268+ const firstFinite = afterOne.every((v) => Number.isFinite(v));
269+ const firstRelL2 = firstFinite ? relL2(afterOne, input) : NaN;
270+ if (!wantJson) {
271+ console.log(
272+ ` one round trip: ${
273+ firstFinite
274+ ? `back to the input to ${firstRelL2.toExponential(2)} relative L2`
275+ : 'NOT FINITE'
276+ }`,
277+ );
278+ }
279+ if (!firstFinite || !(firstRelL2 < 1e-3)) {
280+ // Report it the same way a good run reports itself, so a caller reading
281+ // --json learns what went wrong and on which device rather than having to
282+ // scrape stderr. Then say it in prose and stop: timing a transform that does
283+ // not transform is a waste of minutes.
284+ if (wantJson) {
285+ console.log(
286+ JSON.stringify(
287+ {
288+ ...identity,
289+ firstRoundTrip: { finite: firstFinite, relL2: firstRelL2 },
290+ throughput: null,
291+ latency: null,
292+ digest: null,
293+ input: null,
294+ state: { min: null, max: null, finite: firstFinite },
295+ },
296+ null,
297+ 2,
298+ ),
299+ );
300+ }
301+ const detail = firstFinite
302+ ? `it came back ${firstRelL2.toExponential(3)} away from the input, which is far\n` +
303+ ` outside fp32 round-off (~1e-7)`
304+ : `it came back with no finite values at all`;
305+ fail(
306+ `a single spectral -> grid -> spectral round trip does not round-trip on this\n` +
307+ ` device: ${detail}.\n\n` +
308+ ` That is a correctness problem in the transforms here, not a benchmarking one,\n` +
309+ ` so there is nothing worth timing yet. Adapter: ${adapter || '(unknown)'};\n` +
310+ ` Fourier stage: ${plan.fourierMode.toUpperCase()}.\n\n` +
311+ ` Worth trying, in order:\n` +
312+ ` npm run test:node the repo's own transform check against\n` +
313+ ` its f64 CPU twin, on this device\n` +
314+ ` ${BENCH_SHT_COMMAND} --fourier dft the other Fourier stage; if this works,\n` +
315+ ` the WGSL FFT is the problem\n` +
316+ ` ${BENCH_SHT_COMMAND} --lmax 15 does it depend on the grid size?`,
317+ );
318+ }
319+
320+ seed();
321+ submitAll(spec.warmup);
322+ await done();
323+
324+ // --- throughput: batches submitted together, awaited once each ---
325+ const batches = Math.max(1, Math.ceil(spec.steps / batch));
326+ const tp0 = performance.now();
327+ let stepsRun = 0;
328+ let encodeMs = 0;
329+ for (let b = 0; b < batches; b++) {
330+ const n = Math.min(batch, spec.steps - stepsRun);
331+ const e0 = performance.now();
332+ submit(n);
333+ encodeMs += performance.now() - e0;
334+ await done();
335+ stepsRun += n;
336+ }
337+ const throughputMs = (performance.now() - tp0) / stepsRun;
338+ const encodePerStep = encodeMs / stepsRun;
339+
340+ // --- latency: one round trip per submit ---
341+ const latencySteps = Math.min(spec.steps, 200);
342+ const samples = new Float64Array(latencySteps);
343+ for (let s = 0; s < latencySteps; s++) {
344+ const t0 = performance.now();
345+ submit(1);
346+ await done();
347+ samples[s] = performance.now() - t0;
348+ }
349+ const sorted = Float64Array.from(samples).sort();
350+ const q = (p: number): number => sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))];
351+ let latTotal = 0;
352+ for (const v of samples) latTotal += v;
353+ const latency = {
354+ meanMs: latTotal / samples.length,
355+ medianMs: q(0.5),
356+ p05Ms: q(0.05),
357+ p95Ms: q(0.95),
358+ minMs: sorted[0],
359+ };
360+
361+ // --- a reproducible spectrum to compare across implementations ---
362+ let digest = null;
363+ let inputDigest = null;
364+ let state: Float32Array | null = null;
365+ if (wantDigest) {
366+ inputDigest = digestOf(input, plan.fourierMode, adapter);
367+ seed();
368+ submitAll(spec.steps);
369+ await done();
370+ state = await readSpectrum();
371+ digest = digestOf(state, plan.fourierMode, adapter);
372+ }
373+
374+ const current = await readSpectrum();
375+ let finite = true;
376+ let min = Infinity;
377+ let max = -Infinity;
378+ for (const v of current) {
379+ if (!Number.isFinite(v)) finite = false;
380+ if (v < min) min = v;
381+ if (v > max) max = v;
382+ }
383+
384+ if (wantJson) {
385+ console.log(
386+ JSON.stringify(
387+ {
388+ ...identity,
389+ firstRoundTrip: { finite: firstFinite, relL2: firstRelL2 },
390+ throughput: {
391+ batch,
392+ msPerStep: throughputMs,
393+ stepsPerSec: 1000 / throughputMs,
394+ encodeMsPerStep: encodePerStep,
395+ },
396+ latency,
397+ digest,
398+ input: inputDigest,
399+ state: { min, max, finite },
400+ },
401+ null,
402+ 2,
403+ ),
404+ );
405+ } else {
406+ console.log(
407+ ` ${throughputMs.toFixed(3)} ms/round trip ` +
408+ `${(1000 / throughputMs).toFixed(1)} round trips/s (batches of ${batch})`,
409+ );
410+ console.log(` i.e. ${(throughputMs / 2).toFixed(3)} ms per single transform`);
411+ console.log(
412+ ` of which CPU command encoding: ${encodePerStep.toFixed(3)} ms/round trip ` +
413+ `(${((100 * encodePerStep) / throughputMs).toFixed(0)}% — the rest is the GPU)`,
414+ );
415+ console.log(
416+ ` one round trip per submit: ${latency.meanMs.toFixed(3)} ms mean · ` +
417+ `median ${latency.medianMs.toFixed(3)} · p05 ${latency.p05Ms.toFixed(3)} · ` +
418+ `p95 ${latency.p95Ms.toFixed(3)} · min ${latency.minMs.toFixed(3)}`,
419+ );
420+ if (!finite) console.log(' — NOT FINITE');
421+ if (digest) {
422+ console.log(`\n spectrum after ${spec.steps} round trips from seed ${spec.seed}:`);
423+ console.log(` ${formatDigest(digest)}`);
424+ }
425+ console.log(
426+ `\n The native counterpart is bench/shtns/shtbench{,_gpu} --mode transform;\n` +
427+ ` scripts/compare-native.mjs runs both and lines the numbers up.`,
428+ );
429+ }
430+
431+ if (dumpState && state && digest) {
432+ writeFileSync(
433+ dumpState,
434+ JSON.stringify({
435+ ...identity,
436+ digest,
437+ input: inputDigest,
438+ state: [...state],
439+ }),
440+ );
441+ if (!wantJson) console.log(`\n wrote ${dumpState}`);
442+ }
443+
444+ for (const b of [qlm[0], qlm[1], spat, readback]) b.destroy();
445+ plan.destroy();
446+ device.destroy();
447+ process.exit(finite ? 0 : 1);
448+} catch (e) {
449+ plan?.destroy();
450+ device?.destroy();
451+ fail(errMsg(e));
452+}
scripts/bench.tsadded+346−0View file
@@ -0,0 +1,346 @@
1+/**
2+ * Command-line benchmark: run exactly what the browser runs — the same .m
3+ * models, lowered by numbl and compiled to the same WGSL kernels, over the same
4+ * transforms — on desktop WebGPU (Google Dawn, via the optional `webgpu`
5+ * package), and report ms/step. The app prints the matching command under its
6+ * stats line; copy it and run it here for an apples-to-apples comparison.
7+ *
8+ * npm run bench -- --preset schnak-spots --lmax 63 --steps 2000 --seed 1 \
9+ * --a 0.1 --b 0.9 --D1 0.0004 --D2 0.008 --dt 0.05
10+ *
11+ * The only thing missing here is the rendering: this is the solver alone.
12+ *
13+ * Two numbers are reported, because they answer different questions:
14+ * - throughput: a batch of steps submitted together, awaited once. This is how
15+ * the app runs, and what keeping the state in GPU buffers is for.
16+ * - latency: one step per submit, each awaited. Comparable to a design that
17+ * reads back every step, and the only way to get a per-step distribution.
18+ */
19+import { requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
20+import { ModelSession } from '../src/mgpu/session.ts';
21+import { presets } from '../src/mgpu/registry.ts';
22+import { mGeometries, DEFAULT_GEOMETRY_KEY } from '../src/geom/registry.ts';
23+import {
24+ parseArgs,
25+ modelForSpec,
26+ resolvePreset,
27+ geometryForSpec,
28+ formatCommand,
29+ BENCH_COMMAND,
30+ DEFAULT_LMAX,
31+ DEFAULT_NITER,
32+ DEFAULT_SEED,
33+ DEFAULT_STEPS,
34+ DEFAULT_WARMUP,
35+ type RunSpec,
36+} from '../src/bench/runSpec.ts';
37+import { digestOf, formatDigest } from '../src/mgpu/digest.ts';
38+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
39+import { writeFileSync } from 'node:fs';
40+
41+const USAGE = `usage: ${BENCH_COMMAND} [options]
42+
43+ --preset <key> ${presets.map((p) => p.key).join(' | ')}
44+ (default ${presets[0].key})
45+ --geometry <key> ${mGeometries.map((g) => g.key).join(' | ')}
46+ (default ${DEFAULT_GEOMETRY_KEY})
47+ --lmax <n> spherical harmonic truncation (default ${DEFAULT_LMAX})
48+ --niter <n> iterations of the implicit solve, unrolled into the compiled
49+ step (default ${DEFAULT_NITER})
50+ --steps <n> timed steps (default ${DEFAULT_STEPS})
51+ --warmup <n> untimed steps first (default ${DEFAULT_WARMUP})
52+ --seed <n> initial-noise seed (default ${DEFAULT_SEED})
53+ --batch <n> steps per submit for the throughput number (default 16)
54+ --digest after timing, re-run exactly --steps steps from the seed and
55+ print a digest of the final state
56+ --dump-state <f> like --digest, and write the state to <f> as JSON, for
57+ scripts/compare-env.mjs to compare against a browser run
58+ --<param> <v> any parameter of the preset's model, e.g. --dt 0.05
59+ --g<param> <v> any parameter of the geometry, e.g. --gwaist 0.6
60+ --json machine-readable output
61+ --help
62+
63+The browser app shows the command for whatever it is currently simulating;
64+copy it from under the stats line to compare the same run here.`;
65+
66+function fail(msg: string, code = 1): never {
67+ console.error(`bench: ${msg}`);
68+ process.exit(code);
69+}
70+
71+// ---------------------------------------------------------------- arguments
72+const argv = process.argv.slice(2);
73+if (argv.includes('--help') || argv.includes('-h')) {
74+ console.log(USAGE);
75+ process.exit(0);
76+}
77+const wantJson = argv.includes('--json');
78+let batch = 16;
79+let dumpState: string | null = null;
80+let wantDigest = false;
81+const rest: string[] = [];
82+for (let i = 0; i < argv.length; i++) {
83+ const a = argv[i];
84+ if (a === '--json') continue;
85+ if (a === '--digest') {
86+ wantDigest = true;
87+ continue;
88+ }
89+ const valued = (name: string): string | null => {
90+ if (a === `--${name}`) return argv[++i];
91+ if (a.startsWith(`--${name}=`)) return a.slice(name.length + 3);
92+ return null;
93+ };
94+ const b = valued('batch');
95+ if (b !== null) {
96+ batch = Number(b);
97+ continue;
98+ }
99+ const d = valued('dump-state');
100+ if (d !== null) {
101+ dumpState = d;
102+ wantDigest = true;
103+ continue;
104+ }
105+ rest.push(a);
106+}
107+if (!Number.isInteger(batch) || batch < 1) fail(`--batch must be an integer >= 1`, 2);
108+
109+let spec: RunSpec;
110+try {
111+ spec = parseArgs(rest);
112+} catch (e) {
113+ fail(`${errMsg(e)}\n\n${USAGE}`, 2);
114+}
115+
116+// ---------------------------------------------------------------- statistics
117+interface Timing {
118+ meanMs: number;
119+ medianMs: number;
120+ p05Ms: number;
121+ p95Ms: number;
122+ minMs: number;
123+ totalMs: number;
124+ stepsPerSec: number;
125+}
126+
127+function timing(samples: Float64Array): Timing {
128+ const sorted = Float64Array.from(samples).sort();
129+ const q = (p: number): number =>
130+ sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))];
131+ let total = 0;
132+ for (const v of samples) total += v;
133+ const mean = total / samples.length;
134+ return {
135+ meanMs: mean,
136+ medianMs: q(0.5),
137+ p05Ms: q(0.05),
138+ p95Ms: q(0.95),
139+ minMs: sorted[0],
140+ totalMs: total,
141+ stepsPerSec: 1000 / mean,
142+ };
143+}
144+
145+function fieldRange(v: ArrayLike<number>): { min: number; max: number } {
146+ let min = Infinity;
147+ let max = -Infinity;
148+ for (let i = 0; i < v.length; i++) {
149+ if (v[i] < min) min = v[i];
150+ if (v[i] > max) max = v[i];
151+ }
152+ return { min, max };
153+}
154+
155+// ---------------------------------------------------------------- run
156+const model = modelForSpec(spec);
157+const { preset } = resolvePreset(spec.preset);
158+const geometry = geometryForSpec(spec);
159+
160+let device: GPUDevice | null = null;
161+let session: ModelSession | null = null;
162+
163+try {
164+ const runtime = await installWebGpu();
165+ device = await requestShtDevice().catch((e: unknown) => {
166+ throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
167+ });
168+ const adapter = await describeAdapter(device);
169+
170+ session = await ModelSession.create({
171+ device,
172+ model,
173+ params: spec.params,
174+ lmax: spec.lmax,
175+ geometry,
176+ geometryParams: spec.geometryParams,
177+ niter: spec.niter,
178+ });
179+ session.seed(spec.seed);
180+
181+ const plan = session.describe();
182+ const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
183+ const cfg = session.cfg;
184+
185+ if (!wantJson) {
186+ console.log(`turing-surface bench — solver only, no rendering\n`);
187+ console.log(` preset ${preset.label} (models/${model.key}.m: ${model.species.join(', ')})`);
188+ console.log(
189+ ` params ${model.params.map((p) => `${p.key}=${spec.params[p.key]}`).join(' ')}`,
190+ );
191+ const radius = session.geometry.radiusRange();
192+ console.log(
193+ ` geometry ${geometry.label} (geometries/${geometry.key}.m` +
194+ (geometry.params.length
195+ ? `: ${geometry.params.map((p) => `${p.key}=${spec.geometryParams[p.key]}`).join(' ')})`
196+ : ')') +
197+ ` radius ${radius.lo.toFixed(3)}–${radius.hi.toFixed(3)}`,
198+ );
199+ console.log(
200+ ` grid lmax ${cfg.lmax} · ${cfg.nlat}×${cfg.nphi} · nlm ${session.sht.nlm.toLocaleString()} · ` +
201+ `${spec.niter} solve iteration${spec.niter === 1 ? '' : 's'}`,
202+ );
203+ console.log(` compiled ${plan.step.length} GPU ops/step (${kernels} generated kernels)`);
204+ console.log(` fourier ${session.sht.fourierMode.toUpperCase()} stage`);
205+ console.log(` backend WebGPU fp32${adapter ? ` — ${adapter}` : ''}\n ${runtime}`);
206+ console.log(` run ${spec.warmup} warmup + ${spec.steps} timed steps, seed ${spec.seed}\n`);
207+ }
208+
209+ const done = (): Promise<undefined> => device!.queue.onSubmittedWorkDone();
210+
211+ session.step(spec.warmup);
212+ await done();
213+
214+ // --- throughput: batches submitted together, awaited once each ---
215+ const batches = Math.max(1, Math.ceil(spec.steps / batch));
216+ const progress = !wantJson && process.stderr.isTTY;
217+ let lastReport = performance.now();
218+ const tp0 = performance.now();
219+ let stepsRun = 0;
220+ let encodeMs = 0;
221+ for (let b = 0; b < batches; b++) {
222+ const n = Math.min(batch, spec.steps - stepsRun);
223+ const e0 = performance.now();
224+ session.step(n);
225+ encodeMs += performance.now() - e0;
226+ await done();
227+ stepsRun += n;
228+ if (progress && performance.now() - lastReport > 1000) {
229+ const so_far = (performance.now() - tp0) / stepsRun;
230+ process.stderr.write(
231+ `\r\x1b[K ${stepsRun}/${spec.steps} steps · ${so_far.toFixed(2)} ms/step`,
232+ );
233+ lastReport = performance.now();
234+ }
235+ }
236+ const throughputMs = (performance.now() - tp0) / stepsRun;
237+ const encodePerStep = encodeMs / stepsRun;
238+ if (progress) process.stderr.write('\r\x1b[K');
239+
240+ // --- latency: one step per submit, for the distribution ---
241+ const latencySteps = Math.min(spec.steps, 200);
242+ const samples = new Float64Array(latencySteps);
243+ for (let s = 0; s < latencySteps; s++) {
244+ const t0 = performance.now();
245+ session.step(1);
246+ await done();
247+ samples[s] = performance.now() - t0;
248+ }
249+ const t = timing(samples);
250+
251+ const field = await session.read(model.species[0]);
252+ const range = fieldRange(field);
253+ let finite = true;
254+ for (const v of field) if (!Number.isFinite(v)) finite = false;
255+
256+ // A reproducible state to compare across machines: exactly `--steps` steps
257+ // from the seed, separate from the timed runs above (which step a different
258+ // number of times to measure throughput and latency).
259+ let digest = null;
260+ let state: Float32Array | null = null;
261+ if (wantDigest) {
262+ session.seed(spec.seed);
263+ session.step(spec.steps);
264+ await done();
265+ state = await session.read(model.state[0]);
266+ digest = digestOf(state, session.sht.fourierMode, adapter);
267+ }
268+
269+ if (wantJson) {
270+ console.log(
271+ JSON.stringify(
272+ {
273+ command: formatCommand(spec),
274+ spec,
275+ model: model.key,
276+ backend: { adapter, runtime, precision: 'fp32' },
277+ grid: { lmax: cfg.lmax, nlat: cfg.nlat, nphi: cfg.nphi, nlm: session.sht.nlm },
278+ compiled: { opsPerStep: plan.step.length, kernels },
279+ digest,
280+ throughput: {
281+ batch,
282+ msPerStep: throughputMs,
283+ stepsPerSec: 1000 / throughputMs,
284+ encodeMsPerStep: encodePerStep,
285+ },
286+ latency: t,
287+ state: {
288+ t: session.t,
289+ steps: session.steps,
290+ species: model.species[0],
291+ min: range.min,
292+ max: range.max,
293+ contrast: range.max - range.min,
294+ finite,
295+ },
296+ },
297+ null,
298+ 2,
299+ ),
300+ );
301+ } else {
302+ console.log(
303+ ` ${throughputMs.toFixed(2)} ms/step ${(1000 / throughputMs).toFixed(1)} steps/s ` +
304+ `${(spec.params.dt * (1000 / throughputMs)).toFixed(2)} model time/s` +
305+ ` (batches of ${batch})`,
306+ );
307+ console.log(
308+ ` of which CPU command encoding: ${encodePerStep.toFixed(3)} ms/step ` +
309+ `(${((100 * encodePerStep) / throughputMs).toFixed(0)}% — the rest is the GPU)`,
310+ );
311+ console.log(
312+ ` one step per submit: ${t.meanMs.toFixed(2)} ms mean · median ${t.medianMs.toFixed(2)} · ` +
313+ `p05 ${t.p05Ms.toFixed(2)} · p95 ${t.p95Ms.toFixed(2)} · min ${t.minMs.toFixed(2)}`,
314+ );
315+ console.log(
316+ ` after ${session.steps} steps: t = ${session.t.toFixed(2)}, ` +
317+ `${model.species[0]} ∈ [${range.min.toFixed(4)}, ${range.max.toFixed(4)}] ` +
318+ `(contrast ${(range.max - range.min).toFixed(4)})${finite ? '' : ' — NOT FINITE'}`,
319+ );
320+ if (digest) {
321+ console.log(`\n state after ${spec.steps} steps from seed ${spec.seed}:`);
322+ console.log(` ${formatDigest(digest)}`);
323+ }
324+ console.log(
325+ `\n The app's stats line reports the same solver number (batched steps,\n` +
326+ ` nothing read back) plus a separate ms/frame that carries the readback\n` +
327+ ` and the rendering. Compare solver with solver.`,
328+ );
329+ }
330+
331+ if (dumpState && state && digest) {
332+ writeFileSync(
333+ dumpState,
334+ JSON.stringify({ command: formatCommand(spec), spec, digest, state: [...state] }),
335+ );
336+ if (!wantJson) console.log(`\n wrote ${dumpState}`);
337+ }
338+
339+ session.destroy();
340+ device.destroy();
341+ process.exit(finite ? 0 : 1);
342+} catch (e) {
343+ session?.destroy();
344+ device?.destroy();
345+ fail(errMsg(e));
346+}
scripts/check-live.mjsadded+55−0View file
@@ -0,0 +1,55 @@
1+/**
2+ * Smoke-check a deployed URL in headless Chrome: load it, press Run, and
3+ * confirm the solver actually advances. Usage: node scripts/check-live.mjs [url]
4+ */
5+import puppeteer from 'puppeteer-core';
6+
7+const url = process.argv[2] ?? 'https://concept-collection.github.io/turing-surface/';
8+const browser = await puppeteer.launch({
9+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
10+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
11+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
12+});
13+const page = await browser.newPage();
14+await page.setViewport({ width: 1000, height: 900 });
15+const problems = [];
16+page.on('pageerror', (e) => problems.push(`pageerror: ${e.message}`));
17+page.on('requestfailed', (r) => problems.push(`request failed: ${r.url()}`));
18+page.on('console', (m) => {
19+ if (m.type() === 'error' && !/GL Driver|favicon/.test(m.text())) {
20+ problems.push(`console error: ${m.text()}`);
21+ }
22+});
23+
24+try {
25+ await page.goto(url, { waitUntil: 'load', timeout: 60_000 });
26+ await page.waitForFunction(
27+ () => /grid/.test(document.getElementById('stats')?.textContent ?? ''),
28+ { timeout: 120_000 },
29+ );
30+ console.log('initial:', await page.$eval('#stats', (el) => el.textContent));
31+ await page.click('#runpause');
32+ await page.waitForFunction(
33+ () => {
34+ const m = document.getElementById('stats')?.textContent?.match(/\((\d+) steps\)/);
35+ return m && Number(m[1]) >= 20;
36+ },
37+ { timeout: 180_000 },
38+ );
39+ console.log('running:', await page.$eval('#stats', (el) => el.textContent));
40+ const panels = await page.$$eval('.sphere-box canvas', (els) => els.length);
41+ console.log('sphere canvases:', panels);
42+ if (problems.length) {
43+ console.log('PROBLEMS:');
44+ for (const p of new Set(problems)) console.log(' ' + p);
45+ process.exitCode = 1;
46+ } else {
47+ console.log('LIVE CHECK: PASS');
48+ }
49+} catch (e) {
50+ console.error(`LIVE CHECK FAIL: ${e.message}`);
51+ for (const p of new Set(problems)) console.error(' ' + p);
52+ process.exitCode = 1;
53+} finally {
54+ await browser.close();
55+}
scripts/compare-env.mjsadded+169−0View file
@@ -0,0 +1,169 @@
1+/**
2+ * Is the browser computing the same thing as the terminal?
3+ *
4+ * Runs one identical spec in both — same model source, parameters, lmax, seed and
5+ * step count — and compares the final spectral state. The pipeline is
6+ * deterministic given that spec (seeded PRNG, then fixed arithmetic), so the two
7+ * should agree to fp32 round-off. They will not agree bit for bit: GPUs differ in
8+ * fused-multiply-add and other latitude fp32 allows. They should agree to far
9+ * better than any real difference in what is being computed.
10+ *
11+ * Both sides build their spec through the same parseArgs, so neither can quietly
12+ * use a different default.
13+ *
14+ * node scripts/compare-env.mjs [--lmax 31] [--steps 200] [--preset schnak-spots]
15+ *
16+ * Requires `npm run build` first (it serves dist/), and desktop WebGPU for the
17+ * terminal side.
18+ */
19+import { createServer } from 'node:http';
20+import { readFile, unlink } from 'node:fs/promises';
21+import { readFileSync } from 'node:fs';
22+import { extname, join } from 'node:path';
23+import { spawnSync } from 'node:child_process';
24+import { tmpdir } from 'node:os';
25+import puppeteer from 'puppeteer-core';
26+
27+// ---- spec, defaulted small enough to be quick in a browser ----------------
28+const argv = process.argv.slice(2);
29+const flag = (name, dflt) => {
30+ const i = argv.indexOf(`--${name}`);
31+ if (i >= 0 && argv[i + 1] !== undefined) return argv[i + 1];
32+ const eq = argv.find((a) => a.startsWith(`--${name}=`));
33+ return eq ? eq.slice(name.length + 3) : dflt;
34+};
35+const lmax = flag('lmax', '31');
36+const steps = flag('steps', '200');
37+const preset = flag('preset', 'schnak-spots');
38+const seed = flag('seed', '1');
39+const tolerance = Number(flag('tolerance', '2e-3'));
40+
41+const statePath = join(tmpdir(), `turing-surface-desktop-${process.pid}.json`);
42+
43+// ---- desktop -------------------------------------------------------------
44+console.log(`comparing environments — preset ${preset}, lmax ${lmax}, ${steps} steps, seed ${seed}\n`);
45+console.log('desktop (Dawn):');
46+const bench = spawnSync(
47+ 'npx',
48+ [
49+ 'vite-node', 'scripts/bench.ts',
50+ '--preset', preset, '--lmax', lmax, '--seed', seed,
51+ '--steps', steps, '--warmup', '10',
52+ '--dump-state', statePath,
53+ ],
54+ { encoding: 'utf8' },
55+);
56+if (bench.status !== 0) {
57+ console.error(bench.stdout ?? '');
58+ console.error(bench.stderr ?? '');
59+ console.error('compare-env: the desktop run failed');
60+ process.exit(1);
61+}
62+const desktop = JSON.parse(readFileSync(statePath, 'utf8'));
63+console.log(` ${fmt(desktop.digest)}`);
64+console.log(` adapter: ${desktop.digest.adapter}`);
65+
66+// ---- browser -------------------------------------------------------------
67+const DIST = new URL('../dist/', import.meta.url).pathname;
68+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
69+const server = createServer(async (req, res) => {
70+ try {
71+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
72+ const data = await readFile(join(DIST, path));
73+ res.writeHead(200, {
74+ 'content-type': MIME[extname(path)] ?? 'application/octet-stream',
75+ });
76+ res.end(data);
77+ } catch {
78+ res.writeHead(404);
79+ res.end('not found');
80+ }
81+});
82+await new Promise((r) => server.listen(0, '127.0.0.1', r));
83+const port = server.address().port;
84+
85+const flagSets = [
86+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu', '--enable-features=Vulkan'],
87+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
88+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
89+];
90+
91+let browserState = null;
92+let lastError = '';
93+for (const args of flagSets) {
94+ let browser;
95+ try {
96+ browser = await puppeteer.launch({
97+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
98+ args,
99+ });
100+ const page = await browser.newPage();
101+ page.on('pageerror', (e) => (lastError = e.message));
102+ const url =
103+ `http://127.0.0.1:${port}/test.html?state=1&preset=${preset}` +
104+ `&lmax=${lmax}&seed=${seed}&steps=${steps}`;
105+ await page.goto(url, { waitUntil: 'load' });
106+ await page.waitForFunction(() => window.__STATE__ !== undefined, { timeout: 180000 });
107+ browserState = await page.evaluate(() => window.__STATE__);
108+ await browser.close();
109+ break;
110+ } catch (e) {
111+ lastError = e.message ?? String(e);
112+ await browser?.close();
113+ }
114+}
115+server.close();
116+await unlink(statePath).catch(() => {});
117+
118+if (!browserState) {
119+ console.error(`\ncompare-env: the browser run failed: ${lastError}`);
120+ process.exit(1);
121+}
122+
123+console.log('\nbrowser:');
124+console.log(` ${fmt(browserState.digest)}`);
125+console.log(` adapter: ${browserState.digest.adapter}`);
126+
127+// ---- compare -------------------------------------------------------------
128+const a = desktop.state;
129+const b = browserState.state;
130+if (a.length !== b.length) {
131+ console.error(`\nFAIL different state sizes: ${a.length} vs ${b.length}`);
132+ process.exit(1);
133+}
134+let num = 0;
135+let den = 0;
136+let worst = 0;
137+for (let i = 0; i < a.length; i++) {
138+ const d = a[i] - b[i];
139+ num += d * d;
140+ den += b[i] * b[i];
141+ worst = Math.max(worst, Math.abs(d));
142+}
143+const rel = Math.sqrt(num / Math.max(den, 1e-300));
144+
145+console.log('\ndifference:');
146+console.log(` relative L2 ${rel.toExponential(3)}`);
147+console.log(` worst element ${worst.toExponential(3)}`);
148+if (desktop.digest.fourier !== browserState.digest.fourier) {
149+ console.log(
150+ ` NOTE different Fourier stage (${desktop.digest.fourier} vs ` +
151+ `${browserState.digest.fourier}) — those are different algorithms, so they ` +
152+ `round differently. That alone can explain a difference in the values.`,
153+ );
154+}
155+
156+const ok = rel < tolerance;
157+console.log(
158+ `\n${ok ? 'PASS' : 'FAIL'} the two environments compute the same thing ` +
159+ `(relative L2 ${rel.toExponential(2)}, tolerance ${tolerance.toExponential(1)})`,
160+);
161+process.exit(ok ? 0 : 1);
162+
163+function fmt(d) {
164+ const g = (v) => v.toPrecision(9);
165+ return (
166+ `n=${d.n} min=${g(d.min)} max=${g(d.max)} mean=${g(d.mean)} rms=${g(d.rms)} ` +
167+ `fourier=${d.fourier}`
168+ );
169+}
scripts/compare-perf.mjsadded+177−0View file
@@ -0,0 +1,177 @@
1+/**
2+ * Why is the terminal faster than the browser?
3+ *
4+ * Measures the *same* solver work — same .m, same kernels, batched, nothing read
5+ * back, no rendering on either side — in the terminal (Dawn, in-process) and in a
6+ * real browser, and splits the result so the gap attributes itself:
7+ *
8+ * node scripts/compare-perf.mjs [--lmax 63] [--steps 300] [--preset schnak-spots]
9+ *
10+ * The browser side runs `test.html?soak=`, which has no renderer at all. So:
11+ *
12+ * - if the two agree, the solver is equally fast in the browser, and whatever
13+ * the app shows on top of this is readback, rendering, and animation pacing.
14+ * - if the browser is slower here, it is the GPU stack itself: submits crossing
15+ * into the GPU process, or Metal/Vulkan execution differing between Chrome's
16+ * Dawn and node-webgpu's.
17+ *
18+ * CPU command encoding is reported for both, because it is the one cost that can
19+ * make a fast GPU irrelevant — and it is usually *cheaper* in the browser, which
20+ * defers commands to the GPU process instead of validating them inline.
21+ *
22+ * Requires `npm run build` first, and desktop WebGPU for the terminal side.
23+ */
24+import { createServer } from 'node:http';
25+import { readFile } from 'node:fs/promises';
26+import { extname, join } from 'node:path';
27+import { spawnSync } from 'node:child_process';
28+import puppeteer from 'puppeteer-core';
29+
30+const argv = process.argv.slice(2);
31+const flag = (name, dflt) => {
32+ const i = argv.indexOf(`--${name}`);
33+ if (i >= 0 && argv[i + 1] !== undefined) return argv[i + 1];
34+ const eq = argv.find((a) => a.startsWith(`--${name}=`));
35+ return eq ? eq.slice(name.length + 3) : dflt;
36+};
37+const lmax = flag('lmax', '63');
38+const steps = flag('steps', '300');
39+const preset = flag('preset', 'schnak-spots');
40+
41+console.log(`comparing solver rate — preset ${preset}, lmax ${lmax}, ${steps} steps\n`);
42+
43+// ---- terminal ------------------------------------------------------------
44+const bench = spawnSync(
45+ 'npx',
46+ [
47+ 'vite-node', 'scripts/bench.ts', '--json',
48+ '--preset', preset, '--lmax', lmax, '--steps', steps, '--warmup', '30',
49+ ],
50+ { encoding: 'utf8', maxBuffer: 64 * 1024 * 1024 },
51+);
52+if (bench.status !== 0) {
53+ console.error(bench.stdout ?? '');
54+ console.error(bench.stderr ?? '');
55+ console.error('compare-perf: the terminal run failed');
56+ process.exit(1);
57+}
58+const desktop = JSON.parse(bench.stdout);
59+
60+// ---- browser -------------------------------------------------------------
61+const DIST = new URL('../dist/', import.meta.url).pathname;
62+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
63+const server = createServer(async (req, res) => {
64+ try {
65+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
66+ const data = await readFile(join(DIST, path));
67+ res.writeHead(200, {
68+ 'content-type': MIME[extname(path)] ?? 'application/octet-stream',
69+ });
70+ res.end(data);
71+ } catch {
72+ res.writeHead(404);
73+ res.end('not found');
74+ }
75+});
76+await new Promise((r) => server.listen(0, '127.0.0.1', r));
77+const port = server.address().port;
78+
79+const flagSets = [
80+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu', '--enable-features=Vulkan'],
81+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
82+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
83+];
84+
85+let soak = null;
86+let lastError = '';
87+for (const args of flagSets) {
88+ let browser;
89+ try {
90+ browser = await puppeteer.launch({
91+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
92+ args,
93+ });
94+ const page = await browser.newPage();
95+ page.on('pageerror', (e) => (lastError = e.message));
96+ await page.goto(`http://127.0.0.1:${port}/test.html?soak=${steps}&lmax=${lmax}`, {
97+ waitUntil: 'load',
98+ });
99+ await page.waitForFunction(() => window.__SOAK__ !== undefined, { timeout: 600000 });
100+ soak = await page.evaluate(() => window.__SOAK__);
101+ await browser.close();
102+ break;
103+ } catch (e) {
104+ lastError = e.message ?? String(e);
105+ await browser?.close();
106+ }
107+}
108+server.close();
109+
110+if (!soak) {
111+ console.error(`compare-perf: the browser run failed: ${lastError}`);
112+ process.exit(1);
113+}
114+
115+// ---- report --------------------------------------------------------------
116+const d = desktop.throughput;
117+const row = (label, total, encode, adapter, fourier) => {
118+ console.log(` ${label.padEnd(10)} ${total.toFixed(3)} ms/step` +
119+ ` encoding ${encode.toFixed(3)} ms/step (${((100 * encode) / total).toFixed(0)}%)` +
120+ ` ${fourier.toUpperCase()}`);
121+ console.log(` ${''.padEnd(10)} ${adapter}`);
122+};
123+console.log('solver only, batched, nothing read back, no rendering:\n');
124+row('terminal', d.msPerStep, d.encodeMsPerStep, desktop.backend.adapter, desktop.digest?.fourier ?? 'fft');
125+row('browser', soak.solverMsPerStep, soak.encodeMsPerStep, soak.adapter, soak.fourier);
126+
127+const ratio = soak.solverMsPerStep / d.msPerStep;
128+console.log(`\n browser / terminal = ${ratio.toFixed(2)}x`);
129+
130+// Before reading anything into the ratio: are these even the same GPU? A browser
131+// quietly falling back to a software adapter is a common cause of "the browser is
132+// much slower", and it makes the comparison meaningless rather than informative.
133+const software = (a) => /swiftshader|llvmpipe|software|basic render/i.test(a ?? '');
134+const desktopAdapter = desktop.backend.adapter ?? '';
135+if (software(soak.adapter) !== software(desktopAdapter)) {
136+ console.log(
137+ `\n STOP these are not the same device. One side is a software renderer:\n` +
138+ ` terminal: ${desktopAdapter}\n browser: ${soak.adapter}\n` +
139+ ` The ratio above compares different hardware and means nothing. If it is the\n` +
140+ ` browser that fell back, that IS the answer — check chrome://gpu for why\n` +
141+ ` (hardware acceleration disabled, or the GPU blocklisted).`,
142+ );
143+} else if (desktopAdapter && soak.adapter && desktopAdapter !== soak.adapter) {
144+ console.log(
145+ `\n NOTE the two report different adapters, which may just be different\n` +
146+ ` naming for the same GPU — but check it is not a second GPU:\n` +
147+ ` terminal: ${desktopAdapter}\n browser: ${soak.adapter}`,
148+ );
149+}
150+
151+if (desktop.digest && desktop.digest.fourier !== soak.fourier) {
152+ console.log(
153+ `\n NOTE different Fourier stage (${desktop.digest.fourier} vs ${soak.fourier}).\n` +
154+ ` Those are different algorithms with different cost — that is the difference,\n` +
155+ ` not a symptom of it.`,
156+ );
157+} else if (ratio < 1.3) {
158+ console.log(
159+ `\n The solver runs at the same rate in both. Anything the app shows beyond\n` +
160+ ` this is its readback per species, the colormapping, competing with the\n` +
161+ ` renderer for the GPU, and animation pacing — not the computation.`,
162+ );
163+} else {
164+ console.log(
165+ `\n The browser is slower at the same solver work, with no renderer involved,\n` +
166+ ` so it is the GPU stack rather than anything above it: every submit crosses\n` +
167+ ` into the GPU process, and Chrome's Dawn and node-webgpu's need not compile\n` +
168+ ` or schedule these shaders identically. Note also that an animation-paced\n` +
169+ ` page can leave the GPU in a low-power state where a continuous benchmark\n` +
170+ ` boosts it; this soak hammers it continuously, so if the app is slower than\n` +
171+ ` this number, that is a likely reason.`,
172+ );
173+}
174+console.log(
175+ `\n Correctness is a separate question: scripts/compare-env.mjs checks that the\n` +
176+ ` two environments compute the same state.`,
177+);
scripts/diagnose-leg.tsadded+206−0View file
@@ -0,0 +1,206 @@
1+/**
2+ * Where does leg_synth's recurrence go wrong?
3+ *
4+ * npx vite-node scripts/diagnose-leg.ts [--lmax 63] [--m 0]
5+ *
6+ * Follow-up to scripts/diagnose-sht.ts, which narrows a bad transform down to
7+ * one shader. This reads that shader's recurrence out term by term.
8+ *
9+ * The trick is to probe the production shader rather than a copy of it: with
10+ * qlm set to a single 1 at coefficient (l0, m) and zero everywhere else,
11+ *
12+ * fm[m][ilat] = sum_l Q_lm ytilde_l^m(theta_i) = ytilde_l0^m(theta_i)
13+ *
14+ * so one synthesis per l0 hands back exactly the recurrence value at that l, for
15+ * every latitude at once, computed by the same code the app runs. Sweeping l0
16+ * from m to lmax gives the whole sequence, and comparing with legendreRow (f64)
17+ * says which term first disagrees:
18+ *
19+ * - wrong at l = m -> the seed (amm, or sinpow_rescaled)
20+ * - wrong at l = m+1 -> a_{m+1}^m, i.e. the ab buffer as the shader reads it
21+ * - right until some l, then -> the two-at-a-time advance in the loop
22+ * growing steadily
23+ * - a constant wrong factor -> a scale error, not an instability
24+ */
25+import { ShtPlan, requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
26+import { ShtReference } from '../src/sht/reference.ts';
27+import { legendreRow } from '../src/sht/coeffs.ts';
28+import { gridForLmax, lmIndex, type ShtConfig } from '../src/sht/layout.ts';
29+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
30+
31+const argv = process.argv.slice(2);
32+if (argv.includes('--help') || argv.includes('-h')) {
33+ console.log(`usage: npx vite-node scripts/diagnose-leg.ts [options]
34+
35+ --lmax <n> spherical harmonic truncation (default 63, the app's)
36+ --m <n> the order to follow (default 0, which needs no rescaling at all
37+ and so isolates the plain recurrence)
38+ --lats <i,j> latitudes to sample (default 0,1,mid,last)
39+ --all print every l, not just the interesting ones
40+ --help`);
41+ process.exit(0);
42+}
43+const flag = (name: string, dflt: string): string => {
44+ const i = argv.indexOf(`--${name}`);
45+ if (i >= 0 && argv[i + 1] !== undefined) return argv[i + 1];
46+ const eq = argv.find((a) => a.startsWith(`--${name}=`));
47+ return eq ? eq.slice(name.length + 3) : dflt;
48+};
49+const lmax = Number(flag('lmax', '63'));
50+const m = Number(flag('m', '0'));
51+const showAll = argv.includes('--all');
52+
53+if (!Number.isInteger(m) || m < 0 || m > lmax) {
54+ console.error(`diagnose-leg: --m must be an integer in [0, lmax]`);
55+ process.exit(2);
56+}
57+
58+let device: GPUDevice | null = null;
59+let plan: ShtPlan | null = null;
60+try {
61+ const runtime = await installWebGpu();
62+ device = await requestShtDevice().catch((e: unknown) => {
63+ throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
64+ });
65+ const adapter = await describeAdapter(device);
66+
67+ const { nlat, nphi } = gridForLmax(lmax, 3);
68+ const cfg: ShtConfig = { lmax, mmax: lmax, nlat, nphi };
69+ const ref = new ShtReference(cfg);
70+ // The Fourier stage plays no part here — only fm is read.
71+ plan = await ShtPlan.create(device, cfg, { fourier: 'dft' });
72+
73+ const lats = flag('lats', '')
74+ ? flag('lats', '').split(',').map(Number)
75+ : [0, 1, nlat >> 1, nlat - 1];
76+
77+ console.log('turing-surface — following leg_synth\'s recurrence term by term\n');
78+ console.log(` device ${adapter || '(unknown)'}\n ${runtime}`);
79+ console.log(` grid lmax ${lmax} · ${nlat}×${nphi}`);
80+ console.log(` order m = ${m}${m === 0 ? ' (no rescaling: sinpow_rescaled returns 1, ny = 0)' : ''}`);
81+ console.log(
82+ ` latitudes ${lats
83+ .map((i) => `${i} (theta ${((Math.acos(ref.ct[i]) * 180) / Math.PI).toFixed(1)}°)`)
84+ .join(', ')}\n`,
85+ );
86+
87+ const fmBytes = 8 * (cfg.mmax + 1) * nlat;
88+ const stageFm = device.createBuffer({
89+ size: fmBytes,
90+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
91+ });
92+ const qlm = new Float32Array(2 * ref.nlm);
93+ const row = new Float64Array(lmax + 1);
94+
95+ /** fm[m][ilat] after a synthesis of the unit spectrum at (l0, m). */
96+ const probe = async (l0: number): Promise<Float32Array> => {
97+ qlm.fill(0);
98+ qlm[2 * lmIndex(lmax, l0, m)] = 1;
99+ device!.queue.writeBuffer(plan!.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
100+ const enc = device!.createCommandEncoder();
101+ plan!.encodeSynth(enc);
102+ enc.copyBufferToBuffer(plan!.fmBuf, 0, stageFm, 0, fmBytes);
103+ device!.queue.submit([enc.finish()]);
104+ await stageFm.mapAsync(GPUMapMode.READ);
105+ const out = new Float32Array(stageFm.getMappedRange().slice(0));
106+ stageFm.unmap();
107+ return out;
108+ };
109+
110+ const rows: { l: number; rels: number[]; ratios: number[] }[] = [];
111+ for (let l0 = m; l0 <= lmax; l0++) {
112+ const fm = await probe(l0);
113+ const rels: number[] = [];
114+ const ratios: number[] = [];
115+ for (const ilat of lats) {
116+ legendreRow(ref.coeffs, lmax, m, ref.ct[ilat], ref.st[ilat], row);
117+ const want = row[l0 - m];
118+ const got = fm[2 * (m * nlat + ilat)];
119+ rels.push(Math.abs(got - want) / Math.max(Math.abs(want), 1e-300));
120+ ratios.push(want === 0 ? NaN : got / want);
121+ }
122+ rows.push({ l: l0, rels, ratios });
123+ }
124+
125+ // Loose on purpose. A forward Legendre recurrence in fp32 loses relative
126+ // accuracy as it goes — by l = 63 a few 1e-5 is normal, and worse at the
127+ // latitudes where the terms nearly cancel. What we are hunting is a structural
128+ // error, which shows up as a ratio far from 1, not as a slow drift.
129+ const OK = 1e-2;
130+ const bad = (r: { rels: number[] }): boolean => r.rels.some((x) => !(x < OK));
131+ const firstBad = rows.find(bad);
132+
133+ const head = ` l ` + lats.map((i) => `ilat ${String(i).padStart(3)}`.padStart(14)).join('');
134+ console.log(head);
135+ console.log(` ${'-'.repeat(head.length)}`);
136+ for (const r of rows) {
137+ // every l when --all; otherwise the seed, the first step, the first failure
138+ // and its neighbours, and a tail sample — enough to see the shape
139+ const near = firstBad ? Math.abs(r.l - firstBad.l) <= 3 : false;
140+ const interesting =
141+ showAll || r.l <= m + 2 || near || r.l >= lmax - 1 || (r.l - m) % 8 === 0;
142+ if (!interesting) continue;
143+ const cells = r.rels
144+ .map((rel, k) =>
145+ (rel < OK
146+ ? `ok ${rel.toExponential(1)}`
147+ : `${r.ratios[k] > 1e3 || r.ratios[k] < -1e3 ? '' : 'x'}${r.ratios[k].toExponential(2)}`
148+ ).padStart(14),
149+ )
150+ .join('');
151+ console.log(` ${String(r.l).padStart(4)} ${cells}${bad(r) ? ' <-- wrong' : ''}`);
152+ }
153+ console.log(
154+ `\n cells are "ok <relative error>" when the term is right, and the ratio got/want\n` +
155+ ` when it is not. A few 1e-5 by l = ${lmax} is normal: an fp32 forward recurrence\n` +
156+ ` loses relative accuracy as it goes, worst where consecutive terms nearly cancel.`,
157+ );
158+
159+ if (!firstBad) {
160+ console.log(`\n Every term of the m = ${m} recurrence is right on this device.`);
161+ console.log(
162+ ` So the problem is not the recurrence itself — try another --m, or look at\n` +
163+ ` the accumulation into acc rather than the values going into it.`,
164+ );
165+ } else {
166+ const steps = Math.floor((firstBad.l - m) / 2);
167+ console.log(`\n First wrong term: l = ${firstBad.l}, which is`);
168+ if (firstBad.l === m) {
169+ console.log(
170+ ` the seed itself — amm[m] or sinpow_rescaled, before any recurrence runs.`,
171+ );
172+ } else if (firstBad.l === m + 1) {
173+ console.log(
174+ ` y1's initializer, ab[base + 1].x * ct * y0 — so a_{m+1}^m as the shader\n` +
175+ ` reads it, or the very first multiply. No loop iteration has run yet.`,
176+ );
177+ } else {
178+ console.log(
179+ ` ${steps} advance${steps === 1 ? '' : 's'} into the loop (l = m + ${firstBad.l - m}).`,
180+ );
181+ const ok = rows.filter((r) => !bad(r)).map((r) => r.l);
182+ console.log(
183+ ` Terms that are right: l = ${ok.slice(0, 10).join(', ')}` +
184+ (ok.length > 10 ? ', ...' : ''),
185+ );
186+ const parity = new Set(ok.map((l) => (l - m) % 2));
187+ if (parity.size === 1) {
188+ console.log(
189+ ` Every correct term has (l - m) % 2 == ${[...parity][0]}, and every wrong one\n` +
190+ ` the other parity. The loop carries two values per iteration — y0 for even\n` +
191+ ` offsets and y1 for odd — so one of the two is being updated wrongly while\n` +
192+ ` the other is fine.`,
193+ );
194+ }
195+ }
196+ }
197+
198+ stageFm.destroy();
199+ plan.destroy();
200+ device.destroy();
201+} catch (e) {
202+ plan?.destroy();
203+ device?.destroy();
204+ console.error(`diagnose-leg: ${errMsg(e)}`);
205+ process.exit(1);
206+}
scripts/diagnose-sht.tsadded+311−0View file
@@ -0,0 +1,311 @@
1+/**
2+ * Which half of a transform is wrong on this device?
3+ *
4+ * npx vite-node scripts/diagnose-sht.ts [--lmax 63] [--seed 12345]
5+ *
6+ * A transform is two stages, and both directions share code, so a single
7+ * pass/fail says very little:
8+ *
9+ * synthesis: qlm --[leg_synth]--> fm --[fft_synth | dft_synth]--> spat
10+ * analysis: spat --[fft_analys | dft_analys]--> fm --[leg_analys]--> qlm
11+ *
12+ * This reads the intermediate `fm` back out and compares each stage against
13+ * src/sht/reference.ts (f64, direct summation) on its own:
14+ *
15+ * - fm wrong -> the Legendre stage
16+ * - fm right but spat wrong -> the Fourier stage
17+ * - both right in DFT, wrong in FFT -> the WGSL FFT specifically
18+ *
19+ * and breaks the error down by m and by latitude, because "only high m" or "only
20+ * near the poles" points straight at the rescaled recurrence, while "every m
21+ * equally" points at indexing.
22+ *
23+ * Written for a report of `npm run test:node` failing on a GPU the transforms
24+ * have not run on before. Nothing here is a benchmark.
25+ */
26+import { ShtPlan, requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
27+import { ShtReference, randomSpectrum } from '../src/sht/reference.ts';
28+import { gridForLmax, isPowerOfTwo, type ShtConfig } from '../src/sht/layout.ts';
29+import { fftThreads } from '../src/sht/wgsl/fourier.ts';
30+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
31+
32+const argv = process.argv.slice(2);
33+if (argv.includes('--help') || argv.includes('-h')) {
34+ console.log(`usage: npx vite-node scripts/diagnose-sht.ts [options]
35+
36+ --lmax <n> spherical harmonic truncation (default 63, the app's)
37+ --seed <n> seed of the test spectrum (default 12345)
38+ --fourier <m> only test this stage: fft | dft (default: both)
39+ --help
40+
41+Compares each stage of each direction against the f64 CPU reference and says
42+which one is wrong. See the header of this file.`);
43+ process.exit(0);
44+}
45+const flag = (name: string, dflt: string): string => {
46+ const i = argv.indexOf(`--${name}`);
47+ if (i >= 0 && argv[i + 1] !== undefined) return argv[i + 1];
48+ const eq = argv.find((a) => a.startsWith(`--${name}=`));
49+ return eq ? eq.slice(name.length + 3) : dflt;
50+};
51+const lmax = Number(flag('lmax', '63'));
52+const seed = Number(flag('seed', '12345'));
53+const only = flag('fourier', '');
54+
55+/** Relative L2 of a against b, both flat. */
56+function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
57+ let num = 0;
58+ let den = 0;
59+ for (let i = 0; i < b.length; i++) {
60+ const d = (a[i] ?? NaN) - b[i];
61+ num += d * d;
62+ den += b[i] * b[i];
63+ }
64+ return Math.sqrt(num / Math.max(den, 1e-300));
65+}
66+
67+function anyNonFinite(a: ArrayLike<number>): boolean {
68+ for (let i = 0; i < a.length; i++) if (!Number.isFinite(a[i])) return true;
69+ return false;
70+}
71+
72+/** The Fourier half of a synthesis, on the host, from whatever fm it is given.
73+ * Mirrors ShtReference.synth's inner loop — so feeding it the GPU's own fm says
74+ * what the Fourier stage should have produced from the input it actually had. */
75+function fourierSynth(cfg: ShtConfig, fm: ArrayLike<number>): Float64Array {
76+ const { mmax, nlat, nphi } = cfg;
77+ const spat = new Float64Array(nlat * nphi);
78+ for (let i = 0; i < nlat; i++) {
79+ for (let j = 0; j < nphi; j++) {
80+ const phi = (2 * Math.PI * j) / nphi;
81+ let v = fm[2 * i];
82+ for (let m = 1; m <= mmax; m++) {
83+ const o = 2 * (m * nlat + i);
84+ v += 2 * (fm[o] * Math.cos(m * phi) - fm[o + 1] * Math.sin(m * phi));
85+ }
86+ spat[i * nphi + j] = v;
87+ }
88+ }
89+ return spat;
90+}
91+
92+/** Worst offender along one axis of the [m][ilat] complex fm array. */
93+function fmBreakdown(
94+ cfg: ShtConfig,
95+ got: ArrayLike<number>,
96+ want: ArrayLike<number>,
97+): { byM: { m: number; rel: number }[]; worstLat: { ilat: number; rel: number } } {
98+ const { mmax, nlat } = cfg;
99+ const byM: { m: number; rel: number }[] = [];
100+ const latNum = new Float64Array(nlat);
101+ const latDen = new Float64Array(nlat);
102+ for (let m = 0; m <= mmax; m++) {
103+ let num = 0;
104+ let den = 0;
105+ for (let i = 0; i < nlat; i++) {
106+ for (let c = 0; c < 2; c++) {
107+ const k = 2 * (m * nlat + i) + c;
108+ const d = (got[k] ?? NaN) - want[k];
109+ num += d * d;
110+ den += want[k] * want[k];
111+ latNum[i] += d * d;
112+ latDen[i] += want[k] * want[k];
113+ }
114+ }
115+ byM.push({ m, rel: Math.sqrt(num / Math.max(den, 1e-300)) });
116+ }
117+ let worstLat = { ilat: 0, rel: 0 };
118+ for (let i = 0; i < nlat; i++) {
119+ const rel = Math.sqrt(latNum[i] / Math.max(latDen[i], 1e-300));
120+ if (rel > worstLat.rel) worstLat = { ilat: i, rel };
121+ }
122+ return { byM, worstLat };
123+}
124+
125+const OK = 1e-4; // fp32 through these transforms lands near 1e-6; 1e-4 is generous
126+const verdict = (rel: number): string => (rel < OK ? 'ok ' : 'WRONG');
127+
128+let device: GPUDevice | null = null;
129+try {
130+ const runtime = await installWebGpu();
131+ device = await requestShtDevice().catch((e: unknown) => {
132+ throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
133+ });
134+ const adapter = await describeAdapter(device);
135+
136+ const { nlat, nphi } = gridForLmax(lmax, 3);
137+ const cfg: ShtConfig = { lmax, mmax: lmax, nlat, nphi };
138+ const ref = new ShtReference(cfg);
139+ const qlm = randomSpectrum(cfg, seed);
140+
141+ console.log('turing-surface — which stage of the transform is wrong?\n');
142+ console.log(` device ${adapter || '(unknown)'}\n ${runtime}`);
143+ console.log(` grid lmax ${cfg.lmax} · ${nlat}×${nphi} · nlm ${ref.nlm}`);
144+ console.log(
145+ ` limits maxComputeWorkgroupStorageSize ${device.limits.maxComputeWorkgroupStorageSize}` +
146+ `, maxComputeInvocationsPerWorkgroup ${device.limits.maxComputeInvocationsPerWorkgroup}`,
147+ );
148+ console.log(
149+ ` the FFT stage needs a power-of-two nphi (${isPowerOfTwo(nphi)}), ` +
150+ `16*nphi = ${16 * nphi} bytes of\n workgroup storage and ` +
151+ `${fftThreads(nphi)} invocations per workgroup\n`,
152+ );
153+
154+ // reference values, computed once
155+ const fmRef = ref.legendreSynth(qlm);
156+ const spatRef = ref.synth(qlm);
157+ const qlmRef = ref.analys(spatRef);
158+
159+ const modes: ('fft' | 'dft')[] =
160+ only === 'fft' || only === 'dft' ? [only] : ['fft', 'dft'];
161+ const summary: string[] = [];
162+
163+ for (const mode of modes) {
164+ let plan: ShtPlan | null = null;
165+ try {
166+ plan = await ShtPlan.create(device, cfg, { fourier: mode });
167+ } catch (e) {
168+ console.log(`${mode.toUpperCase()} stage: unavailable — ${errMsg(e)}\n`);
169+ continue;
170+ }
171+
172+ const stageFm = device.createBuffer({
173+ size: 8 * (cfg.mmax + 1) * nlat,
174+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
175+ });
176+ const stageSpat = device.createBuffer({
177+ size: 4 * nlat * nphi,
178+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
179+ });
180+ const read = async (buf: GPUBuffer): Promise<Float32Array> => {
181+ await buf.mapAsync(GPUMapMode.READ);
182+ const out = new Float32Array(buf.getMappedRange().slice(0));
183+ buf.unmap();
184+ return out;
185+ };
186+
187+ // --- synthesis, stopping to look at fm on the way through ---
188+ device.queue.writeBuffer(plan.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
189+ const enc = device.createCommandEncoder();
190+ plan.encodeSynth(enc);
191+ enc.copyBufferToBuffer(plan.fmBuf, 0, stageFm, 0, 8 * (cfg.mmax + 1) * nlat);
192+ enc.copyBufferToBuffer(plan.spatBuf, 0, stageSpat, 0, 4 * nlat * nphi);
193+ device.queue.submit([enc.finish()]);
194+ const fmGpu = await read(stageFm);
195+ const spatGpu = await read(stageSpat);
196+
197+ const legSynthRel = relL2(fmGpu, fmRef);
198+ const synthRel = relL2(spatGpu, spatRef);
199+ // The Fourier stage judged on its own input, not on the reference's: if the
200+ // Legendre stage is already wrong, comparing spat with spatRef only repeats
201+ // that. This asks whether the Fourier stage did the right thing with the fm
202+ // it was actually handed.
203+ const fourierRel = relL2(spatGpu, fourierSynth(cfg, fmGpu));
204+
205+ // --- analysis, for contrast: same two stages, opposite order ---
206+ const spatIn = Float32Array.from(spatRef);
207+ device.queue.writeBuffer(plan.spatBuf, 0, spatIn as Float32Array<ArrayBuffer>);
208+ const enc2 = device.createCommandEncoder();
209+ plan.encodeAnalys(enc2);
210+ enc2.copyBufferToBuffer(plan.fmBuf, 0, stageFm, 0, 8 * (cfg.mmax + 1) * nlat);
211+ device.queue.submit([enc2.finish()]);
212+ const fmAnalysGpu = await read(stageFm);
213+ const qlmGpu = await plan.analys(spatIn);
214+ // forward Fourier of the reference field, in the reference's own normalization
215+ const gmRef = new Float64Array(2 * (cfg.mmax + 1) * nlat);
216+ for (let i = 0; i < nlat; i++) {
217+ for (let m = 0; m <= cfg.mmax; m++) {
218+ let re = 0;
219+ let im = 0;
220+ for (let j = 0; j < nphi; j++) {
221+ const phi = (2 * Math.PI * j) / nphi;
222+ re += spatRef[i * nphi + j] * Math.cos(m * phi);
223+ im -= spatRef[i * nphi + j] * Math.sin(m * phi);
224+ }
225+ gmRef[2 * (m * nlat + i)] = re;
226+ gmRef[2 * (m * nlat + i) + 1] = im;
227+ }
228+ }
229+ const analysFourierRel = relL2(fmAnalysGpu, gmRef);
230+ const analysRel = relL2(qlmGpu, qlmRef);
231+
232+ console.log(`${mode.toUpperCase()} stage — plan chose ${plan.fourierMode.toUpperCase()}\n`);
233+ console.log(` synthesis qlm -> fm -> spat`);
234+ console.log(
235+ ` ${verdict(legSynthRel)} leg_synth fm vs f64 reference ` +
236+ `${legSynthRel.toExponential(2)}${anyNonFinite(fmGpu) ? ' (has NaN/Inf)' : ''}`,
237+ );
238+ console.log(
239+ ` ${verdict(fourierRel)} ${mode}_synth spat vs host Fourier of that fm ` +
240+ `${fourierRel.toExponential(2)}${anyNonFinite(spatGpu) ? ' (has NaN/Inf)' : ''}`,
241+ );
242+ console.log(
243+ ` ${verdict(synthRel)} end to end spat vs f64 reference ` +
244+ `${synthRel.toExponential(2)}`,
245+ );
246+ console.log(`\n analysis spat -> fm -> qlm`);
247+ console.log(
248+ ` ${verdict(analysFourierRel)} ${mode}_analys fm vs f64 reference ` +
249+ `${analysFourierRel.toExponential(2)}`,
250+ );
251+ console.log(
252+ ` ${verdict(analysRel)} end to end qlm vs f64 reference ` +
253+ `${analysRel.toExponential(2)}`,
254+ );
255+
256+ if (legSynthRel >= OK) {
257+ const { byM, worstLat } = fmBreakdown(cfg, fmGpu, fmRef);
258+ const bad = byM.filter((e) => e.rel >= OK);
259+ const good = byM.filter((e) => e.rel < OK);
260+ console.log(`\n leg_synth is wrong. Where:`);
261+ console.log(
262+ ` ${bad.length} of ${byM.length} orders m are wrong` +
263+ (good.length
264+ ? `; the ones that are right are m = ${good.slice(0, 12).map((e) => e.m).join(', ')}` +
265+ (good.length > 12 ? ', ...' : '')
266+ : ' (all of them)'),
267+ );
268+ if (bad.length) {
269+ const first = bad[0];
270+ const worst = bad.reduce((a, b) => (b.rel > a.rel ? b : a));
271+ console.log(
272+ ` lowest wrong m = ${first.m} (${first.rel.toExponential(2)}), ` +
273+ `worst m = ${worst.m} (${worst.rel.toExponential(2)})`,
274+ );
275+ }
276+ const theta = (Math.acos(ref.ct[worstLat.ilat]) * 180) / Math.PI;
277+ console.log(
278+ ` worst latitude ilat = ${worstLat.ilat} of ${nlat} ` +
279+ `(theta = ${theta.toFixed(1)}°, sin(theta) = ` +
280+ `${ref.st[worstLat.ilat].toExponential(2)}), rel ${worstLat.rel.toExponential(2)}`,
281+ );
282+ console.log(
283+ ` If only high m are wrong, or only latitudes near the poles where\n` +
284+ ` sin(theta) is small, the rescaled seed (sinpow_rescaled in\n` +
285+ ` src/sht/wgsl/common.ts) is the place to look. If every m is wrong by\n` +
286+ ` a similar amount, it is indexing or the dispatch, not the recurrence.\n` +
287+ ` Either way, follow it term by term from here:\n` +
288+ ` npx vite-node scripts/diagnose-leg.ts --lmax ${lmax} --m 0`,
289+ );
290+ }
291+ console.log();
292+
293+ summary.push(
294+ `${mode}: leg_synth ${verdict(legSynthRel).trim()}, ${mode}_synth ` +
295+ `${verdict(fourierRel).trim()}, ${mode}_analys ${verdict(analysFourierRel).trim()}, ` +
296+ `leg_analys ${verdict(analysRel).trim()}`,
297+ );
298+
299+ stageFm.destroy();
300+ stageSpat.destroy();
301+ plan.destroy();
302+ }
303+
304+ console.log('summary');
305+ for (const s of summary) console.log(` ${s}`);
306+ device.destroy();
307+} catch (e) {
308+ device?.destroy();
309+ console.error(`diagnose-sht: ${errMsg(e)}`);
310+ process.exit(1);
311+}
scripts/longrun-node.tsadded+57−0View file
@@ -0,0 +1,57 @@
1+/**
2+ * Long-run sanity check: run Schnakenberg to t = 100 on desktop WebGPU and
3+ * confirm the pattern saturates into O(1)-contrast spots rather than decaying or
4+ * blowing up. Short runs cannot tell a growing instability from a diverging one.
5+ *
6+ * vite-node scripts/longrun-node.ts [lmax]
7+ */
8+import { requestShtDevice } from '../src/sht/sht.ts';
9+import { ModelSession } from '../src/mgpu/session.ts';
10+import { mModelByKey, defaultParams } from '../src/mgpu/registry.ts';
11+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
12+
13+const lmax = Number(process.argv[2] ?? 31);
14+const model = mModelByKey('schnakenberg')!;
15+const params = defaultParams(model);
16+
17+const runtime = await installWebGpu();
18+const device = await requestShtDevice().catch((e: unknown) => {
19+ throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
20+});
21+const session = await ModelSession.create({ device, model, params, lmax });
22+session.seed(1);
23+console.log(`longrun — models/${model.key}.m at lmax ${lmax}, ${runtime}\n`);
24+
25+const nsteps = Math.round(100 / params.dt);
26+const BATCH = 50;
27+const t0 = performance.now();
28+let lo = 0;
29+let hi = 0;
30+for (let s = 0; s < nsteps; s += BATCH) {
31+ session.step(Math.min(BATCH, nsteps - s));
32+ const u = await session.read(model.species[0]);
33+ lo = Infinity;
34+ hi = -Infinity;
35+ for (const v of u) {
36+ if (v < lo) lo = v;
37+ if (v > hi) hi = v;
38+ }
39+ if (session.steps % 400 === 0) {
40+ console.log(
41+ `t=${session.t.toFixed(1).padStart(5)} u in [${lo.toFixed(4)}, ${hi.toFixed(4)}] ` +
42+ `contrast ${(hi - lo).toFixed(4)}`,
43+ );
44+ }
45+}
46+const ms = (performance.now() - t0) / nsteps;
47+
48+const contrast = hi - lo;
49+const saturated = Number.isFinite(contrast) && contrast > 0.5 && hi < 10;
50+console.log(
51+ `\n${saturated ? 'PASS' : 'FAIL'} pattern saturated: contrast ${contrast.toFixed(4)} ` +
52+ `after ${session.steps} steps (${ms.toFixed(1)} ms/step)`,
53+);
54+
55+session.destroy();
56+device.destroy();
57+process.exit(saturated ? 0 : 1);
scripts/nodeWebGpu.tsadded+71−0View file
@@ -0,0 +1,71 @@
1+/**
2+ * Desktop WebGPU for the command-line scripts, via the optional `webgpu`
3+ * package (prebuilt Google Dawn).
4+ *
5+ * Installs Dawn under the globals the transform code expects (navigator.gpu,
6+ * GPUBufferUsage, ...) so everything under src/ runs here unchanged —
7+ * including requestShtDevice(), which makes the same device request the
8+ * browser makes.
9+ */
10+
11+export const errMsg = (e: unknown): string =>
12+ e instanceof Error ? e.message : String(e);
13+
14+/**
15+ * Returns a human-readable runtime description. The import specifier is
16+ * indirect so typechecking does not require the optional package.
17+ */
18+export async function installWebGpu(): Promise<string> {
19+ const specifier = 'webgpu';
20+ let mod: {
21+ create: (flags: string[]) => GPU;
22+ globals: Record<string, unknown>;
23+ };
24+ try {
25+ mod = await import(specifier);
26+ } catch (e) {
27+ // Distinguish "not installed" from "installed but the prebuilt Dawn binary
28+ // will not load" — the second is what a machine missing a system library
29+ // looks like, and reporting it as the first sends people in circles.
30+ const detail = errMsg(e);
31+ if (/Cannot find (package|module) '?webgpu'?/.test(detail)) {
32+ throw new Error(
33+ 'desktop WebGPU needs the optional `webgpu` package (prebuilt Google Dawn):\n' +
34+ ' npm install webgpu\n' +
35+ 'It is an optionalDependency, so npm can skip it silently — `npm ls webgpu`\n' +
36+ 'says whether it is there.',
37+ );
38+ }
39+ const glibc = /GLIBC_([0-9.]+)/.exec(detail);
40+ throw new Error(
41+ `the \`webgpu\` package is installed but did not load:\n ${detail}\n` +
42+ (glibc
43+ ? `Dawn's prebuilt binary wants glibc ${glibc[1]} or newer and this host is older\n` +
44+ '(`ldd --version` says how old). No flag bridges that — use a container with a\n' +
45+ 'newer base image, or a newer host.\n'
46+ : 'That is usually the prebuilt Dawn binary missing a system library.\n'),
47+ );
48+ }
49+ Object.assign(globalThis, mod.globals);
50+ // DAWN_FLAGS is ';'-separated because individual Dawn options take
51+ // comma-separated lists, e.g. 'enable-dawn-features=allow_unsafe_apis,...'
52+ const dawnFlags = process.env.DAWN_FLAGS?.split(';').filter(Boolean) ?? [];
53+ Object.defineProperty(globalThis, 'navigator', {
54+ value: { gpu: mod.create(dawnFlags) },
55+ configurable: true,
56+ writable: true,
57+ });
58+ const { version } = await import(`${specifier}/package.json`, {
59+ with: { type: 'json' },
60+ }).then(
61+ (m) => m.default as { version: string },
62+ () => ({ version: '?' }),
63+ );
64+ return `node-webgpu ${version} (Google Dawn)`;
65+}
66+
67+/** The hint to print when Dawn loads but finds no adapter. */
68+export const NO_ADAPTER_HINT =
69+ ' Dawn reaches the GPU through Vulkan on Linux and Windows, Metal on macOS,\n' +
70+ " so a headless box may have no adapter at all. DAWN_FLAGS='backend=vulkan'\n" +
71+ ' makes it explain itself.';
scripts/screenshot.mjsadded+57−0View file
@@ -0,0 +1,57 @@
1+/** Screenshot the demo page (dist/) in headless Chrome. Usage: node scripts/screenshot.mjs out.png [light|dark] */
2+import { createServer } from 'node:http';
3+import { readFile } from 'node:fs/promises';
4+import { extname, join } from 'node:path';
5+import puppeteer from 'puppeteer-core';
6+
7+const out = process.argv[2] ?? 'demo.png';
8+const scheme = process.argv[3] ?? 'light';
9+const minSteps = Number(process.argv[4] ?? 200);
10+const DIST = new URL('../dist/', import.meta.url).pathname;
11+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
12+
13+const server = createServer(async (req, res) => {
14+ try {
15+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
16+ const data = await readFile(join(DIST, path));
17+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
18+ res.end(data);
19+ } catch {
20+ res.writeHead(404);
21+ res.end();
22+ }
23+});
24+await new Promise((r) => server.listen(0, '127.0.0.1', r));
25+const port = server.address().port;
26+
27+const browser = await puppeteer.launch({
28+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
29+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
30+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
31+});
32+const page = await browser.newPage();
33+await page.setViewport({ width: 1000, height: 900 });
34+await page.emulateMediaFeatures([{ name: 'prefers-color-scheme', value: scheme }]);
35+page.on('console', (m) => console.log(' [page]', m.text()));
36+await page.goto(`http://127.0.0.1:${port}/index.html`, { waitUntil: 'load' });
37+// the sim starts paused; wait for setup to finish, then press Run
38+await page.waitForFunction(() => /grid/.test(document.getElementById('stats')?.textContent ?? ''), { timeout: 120_000 });
39+await page.click('#runpause');
40+await page.waitForFunction(
41+ (min) => {
42+ const s = document.getElementById('stats');
43+ const e = document.getElementById('err');
44+ const m = s && s.textContent.match(/\((\d+) steps\)/);
45+ return (m && Number(m[1]) >= min) || (e && e.textContent.length > 4);
46+ },
47+ { timeout: 600_000 },
48+ minSteps,
49+);
50+await new Promise((r) => setTimeout(r, 300));
51+await page.screenshot({ path: out });
52+console.log('screenshot:', out);
53+console.log('stats:', await page.$eval('#stats', (el) => el.textContent));
54+const err = await page.$eval('#err', (el) => el.textContent);
55+if (err) console.log('err:', err);
56+await browser.close();
57+server.close();
scripts/soak.mjsadded+87−0View file
@@ -0,0 +1,87 @@
1+/**
2+ * Soak test: drive the demo page for many steps and report JS heap growth and
3+ * any crash, distinguishing a page crash from a renderer/driver death.
4+ *
5+ * Usage: node scripts/soak.mjs [steps] [lmax] [backend]
6+ * e.g. node scripts/soak.mjs 1500 63 webgpu
7+ */
8+import { createServer } from 'node:http';
9+import { readFile } from 'node:fs/promises';
10+import { extname, join } from 'node:path';
11+import puppeteer from 'puppeteer-core';
12+
13+const steps = Number(process.argv[2] ?? 1000);
14+const lmax = process.argv[3] ?? '63';
15+const backend = process.argv[4] ?? 'webgpu';
16+const DIST = new URL('../dist/', import.meta.url).pathname;
17+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
18+
19+const server = createServer(async (req, res) => {
20+ try {
21+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
22+ const data = await readFile(join(DIST, path));
23+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
24+ res.end(data);
25+ } catch {
26+ res.writeHead(404);
27+ res.end();
28+ }
29+});
30+await new Promise((r) => server.listen(0, '127.0.0.1', r));
31+const port = server.address().port;
32+
33+const browser = await puppeteer.launch({
34+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
35+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
36+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
37+});
38+const page = await browser.newPage();
39+await page.setViewport({ width: 1000, height: 900 });
40+
41+let crashed = null;
42+page.on('error', (e) => { crashed = `page crash: ${e.message}`; });
43+page.on('pageerror', (e) => { crashed = `page error: ${e.message}`; });
44+page.on('console', (m) => {
45+ const t = m.text();
46+ if (!/GL Driver Message|Failed to load resource/.test(t)) console.log(' [page]', t);
47+});
48+
49+await page.goto(`http://127.0.0.1:${port}/index.html`, { waitUntil: 'load' });
50+await page.waitForFunction(() => /grid/.test(document.getElementById('stats')?.textContent ?? ''), { timeout: 120_000 });
51+await page.select('#lmax', lmax);
52+await page.select('#backend', backend);
53+await page.waitForFunction(() => /grid/.test(document.getElementById('stats')?.textContent ?? ''), { timeout: 120_000 });
54+await page.click('#runpause');
55+
56+const readStep = () =>
57+ page.evaluate(() => {
58+ const m = document.getElementById('stats')?.textContent?.match(/\((\d+) steps\)/);
59+ return m ? Number(m[1]) : 0;
60+ });
61+const heapMB = async () => {
62+ const m = await page.metrics();
63+ return (m.JSHeapUsedSize / 1048576).toFixed(1);
64+};
65+
66+const t0 = Date.now();
67+let last = 0;
68+let stalls = 0;
69+try {
70+ while (last < steps) {
71+ await new Promise((r) => setTimeout(r, 5000));
72+ if (crashed) throw new Error(crashed);
73+ const now = await readStep();
74+ console.log(` step ${now} heap ${await heapMB()} MB (+${now - last} in 5s)`);
75+ if (now === last) {
76+ if (++stalls >= 6) throw new Error(`stalled at step ${now}`);
77+ } else stalls = 0;
78+ last = now;
79+ }
80+ console.log(`SOAK PASS: ${last} steps in ${((Date.now() - t0) / 1000).toFixed(0)}s, heap ${await heapMB()} MB`);
81+} catch (e) {
82+ console.error(`SOAK FAIL at step ${last}: ${e.message}`);
83+ process.exitCode = 1;
84+} finally {
85+ await browser.close().catch(() => {});
86+ server.close();
87+}
scripts/test-gpu.mjsadded+68−0View file
@@ -0,0 +1,68 @@
1+/**
2+ * Headless GPU test runner: serves dist/, opens test.html in headless
3+ * Chrome (falling back to the SwiftShader software WebGPU adapter when no
4+ * hardware GPU is available), and reports the suite results.
5+ *
6+ * Run after `vite build`: node scripts/test-gpu.mjs
7+ */
8+import { createServer } from 'node:http';
9+import { readFile } from 'node:fs/promises';
10+import { extname, join } from 'node:path';
11+import puppeteer from 'puppeteer-core';
12+
13+const DIST = new URL('../dist/', import.meta.url).pathname;
14+const CHROME = process.env.CHROME_PATH ?? '/usr/bin/google-chrome';
15+const MIME = {
16+ '.html': 'text/html',
17+ '.js': 'text/javascript',
18+ '.css': 'text/css',
19+ '.json': 'application/json',
20+ '.wasm': 'application/wasm',
21+};
22+
23+const server = createServer(async (req, res) => {
24+ try {
25+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
26+ const data = await readFile(join(DIST, path));
27+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
28+ res.end(data);
29+ } catch {
30+ res.writeHead(404);
31+ res.end('not found');
32+ }
33+});
34+await new Promise((r) => server.listen(0, '127.0.0.1', r));
35+const port = server.address().port;
36+
37+const flagSets = [
38+ // hardware first, then SwiftShader (software) WebGPU
39+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu', '--enable-features=Vulkan'],
40+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu', '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
41+];
42+
43+let final = null;
44+for (const flags of flagSets) {
45+ const browser = await puppeteer.launch({ executablePath: CHROME, args: flags });
46+ try {
47+ const page = await browser.newPage();
48+ page.on('console', (msg) => console.log(` [page] ${msg.text()}`));
49+ page.on('pageerror', (err) => console.log(` [pageerror] ${err.message}`));
50+ await page.goto(`http://127.0.0.1:${port}/test.html`, { waitUntil: 'load' });
51+ const results = await page.waitForFunction(() => window.__RESULTS__, { timeout: 600_000 });
52+ final = await results.jsonValue();
53+ } catch (e) {
54+ console.error(`run with flags [${flags.join(' ')}] failed: ${e.message}`);
55+ } finally {
56+ await browser.close();
57+ }
58+ if (final && !final.fatal) break;
59+ console.log('retrying with next flag set…');
60+}
61+server.close();
62+
63+if (!final || final.fatal) {
64+ console.error(`GPU tests could not run: ${final?.fatal ?? 'no results'}`);
65+ process.exit(2);
66+}
67+console.log(final.ok ? 'GPU SUITE: PASS' : 'GPU SUITE: FAIL');
68+process.exit(final.ok ? 0 : 1);
scripts/test-node.tsadded+56−0View file
@@ -0,0 +1,56 @@
1+/**
2+ * The whole suite on desktop WebGPU (Google Dawn), against the real pipeline:
3+ * MATLAB source -> numbl lowering -> generated WGSL -> GPU.
4+ *
5+ * The same four check modules run in the browser (test.html), so both GPU
6+ * stacks get the same guarantees. Run through vite-node, which is what resolves
7+ * numbl's compiler sources and the `?raw` model imports:
8+ *
9+ * npm run test:node
10+ */
11+import { requestShtDevice } from '../src/sht/sht.ts';
12+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
13+import { transformChecks } from '../test/transformChecks.ts';
14+import { analyticChecks } from '../test/analyticChecks.ts';
15+import { modelChecks } from '../test/modelChecks.ts';
16+import { geometryChecks } from '../test/geometryChecks.ts';
17+
18+let failures = 0;
19+const check = (name: string, ok: boolean, detail: string): void => {
20+ console.log(`${ok ? 'PASS' : 'FAIL'} ${name} ${detail}`);
21+ if (!ok) failures++;
22+};
23+const log = (s: string): void => console.log(s);
24+
25+/**
26+ * These checks compile MATLAB to compute shaders, so unlike the old CPU-solver
27+ * suite they need a GPU. `--skip-without-gpu` lets a runner that has none say so
28+ * and move on — CI uses it, because the browser suite runs the very same check
29+ * modules on SwiftShader. A plain local run still fails loudly, so a missing GPU
30+ * is never mistaken for a pass.
31+ */
32+const skipWithoutGpu = process.argv.includes('--skip-without-gpu');
33+
34+let runtime: string;
35+let device: GPUDevice;
36+try {
37+ runtime = await installWebGpu();
38+ device = await requestShtDevice();
39+} catch (e) {
40+ const detail = `${errMsg(e)}\n${NO_ADAPTER_HINT}`;
41+ if (skipWithoutGpu) {
42+ console.log(`SKIP no WebGPU available here, so these checks did not run.\n${detail}`);
43+ process.exit(0);
44+ }
45+ console.error(`test-node: ${detail}`);
46+ process.exit(1);
47+}
48+console.log(`turing-surface tests — ${runtime}\n`);
49+
50+await transformChecks(device, check, log);
51+await analyticChecks(device, check, log);
52+await modelChecks(device, check, log);
53+await geometryChecks(device, check, log);
54+
55+console.log(failures === 0 ? '\nAll tests passed.' : `\n${failures} failed.`);
56+process.exit(failures === 0 ? 0 : 1);
src/bench/runSpec.tsadded+194−0View file
@@ -0,0 +1,194 @@
1+/**
2+ * One run, described in a single object shared by the browser app and the
3+ * command-line benchmark. The app formats the run it is currently showing into a
4+ * `npm run bench` command; the benchmark parses that command back into the same
5+ * object and compiles the same .m with it. Neither side keeps its own copy of
6+ * the defaults, so the two runs cannot drift apart — and both execute the same
7+ * MATLAB through the same pipeline, so the comparison is like for like.
8+ */
9+import {
10+ mModels,
11+ presets,
12+ defaultParams,
13+ type MModel,
14+ type Params,
15+ type Preset,
16+} from '../mgpu/registry.ts';
17+import {
18+ mGeometries,
19+ mGeometryByKey,
20+ defaultGeometryParams,
21+ DEFAULT_GEOMETRY_KEY,
22+ type MGeometry,
23+} from '../geom/registry.ts';
24+import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
25+
26+export interface RunSpec {
27+ /** Preset key from the registry; fixes the model, params may still be edited. */
28+ preset: string;
29+ lmax: number;
30+ /** Seed of the initial noise. */
31+ seed: number;
32+ /** Timed steps (the app runs forever; the benchmark stops here). */
33+ steps: number;
34+ /** Untimed steps run first, so shader/pipeline warm-up is not measured. */
35+ warmup: number;
36+ /** Full parameter set of the preset's model, as edited. */
37+ params: Params;
38+ /** Geometry key from the geometry registry. */
39+ geometry: string;
40+ /** Full parameter set of that geometry, as edited. Written on the command
41+ * line with a `g` prefix (`--gwaist`) so a shape parameter can never
42+ * collide with a model one. */
43+ geometryParams: Params;
44+ /** Iterations of the .m's implicit solve. Structural: it is unrolled into
45+ * the compiled step, so it belongs to the spec rather than to the params. */
46+ niter: number;
47+}
48+
49+export const DEFAULT_NITER = 1;
50+
51+/** Geometry + starting parameters of a geometry key. */
52+export function resolveGeometry(key: string): { geometry: MGeometry; params: Params } {
53+ const geometry = mGeometryByKey(key);
54+ if (!geometry) {
55+ throw new Error(
56+ `unknown geometry '${key}' (have: ${mGeometries.map((g) => g.key).join(', ')})`,
57+ );
58+ }
59+ return { geometry, params: defaultGeometryParams(geometry) };
60+}
61+
62+export function geometryForSpec(spec: RunSpec): MGeometry {
63+ return resolveGeometry(spec.geometry).geometry;
64+}
65+
66+/** The command the app displays and the benchmark answers to. Goes through npm
67+ * because the benchmark runs under vite-node, which is what resolves numbl's
68+ * compiler sources and the `?raw` model imports. */
69+export const BENCH_COMMAND = 'npm run bench --';
70+export const DEFAULT_LMAX = 63;
71+export const DEFAULT_SEED = 1;
72+/** Long enough that clock ramp-up and the occasional scheduling hiccup wash
73+ * out: ~10 s of GPU stepping at lmax 63. */
74+export const DEFAULT_STEPS = 2000;
75+export const DEFAULT_WARMUP = 100;
76+
77+/** Model + starting parameters of a preset, for the app's dropdown and the
78+ * benchmark's --preset flag. */
79+export function resolvePreset(key: string): {
80+ preset: Preset;
81+ model: MModel;
82+ params: Params;
83+} {
84+ const preset = presets.find((p) => p.key === key);
85+ if (!preset) {
86+ throw new Error(
87+ `unknown preset '${key}' (have: ${presets.map((p) => p.key).join(', ')})`,
88+ );
89+ }
90+ const model = mModels.find((m) => m.key === preset.modelKey);
91+ if (!model) throw new Error(`preset '${key}' names unknown model '${preset.modelKey}'`);
92+ return { preset, model, params: { ...defaultParams(model), ...preset.params } };
93+}
94+
95+export function modelForSpec(spec: RunSpec): MModel {
96+ return resolvePreset(spec.preset).model;
97+}
98+
99+/** Transform configuration implied by the spec (same rule as the app). */
100+export function configForSpec(spec: RunSpec): ShtConfig {
101+ const { nlat, nphi } = gridForLmax(spec.lmax, modelForSpec(spec).pdeg);
102+ return { lmax: spec.lmax, mmax: spec.lmax, nlat, nphi };
103+}
104+
105+/** The command line that reproduces this run. Every knob the app exposes is
106+ * written out explicitly, so the command stays valid if a preset changes. */
107+export function formatCommand(spec: RunSpec): string {
108+ const model = modelForSpec(spec);
109+ const geometry = geometryForSpec(spec);
110+ const parts = [
111+ BENCH_COMMAND,
112+ `--preset ${spec.preset}`,
113+ `--geometry ${spec.geometry}`,
114+ `--lmax ${spec.lmax}`,
115+ `--niter ${spec.niter}`,
116+ `--steps ${spec.steps}`,
117+ `--seed ${spec.seed}`,
118+ ...model.params.map((p) => `--${p.key} ${String(spec.params[p.key])}`),
119+ ...geometry.params.map((p) => `--g${p.key} ${String(spec.geometryParams[p.key])}`),
120+ ];
121+ if (spec.warmup !== DEFAULT_WARMUP) parts.push(`--warmup ${spec.warmup}`);
122+ return parts.join(' ');
123+}
124+
125+/** Inverse of formatCommand: `--key value` or `--key=value`, in any order.
126+ * Throws with a usable message on anything it does not recognize. */
127+export function parseArgs(argv: string[]): RunSpec {
128+ const flags = new Map<string, string>();
129+ for (let i = 0; i < argv.length; i++) {
130+ const arg = argv[i];
131+ if (!arg.startsWith('--')) throw new Error(`unexpected argument '${arg}'`);
132+ const eq = arg.indexOf('=');
133+ const key = eq >= 0 ? arg.slice(2, eq) : arg.slice(2);
134+ const value = eq >= 0 ? arg.slice(eq + 1) : argv[++i];
135+ if (value === undefined) throw new Error(`--${key} needs a value`);
136+ if (!key) throw new Error(`bad option '${arg}'`);
137+ flags.set(key, value);
138+ }
139+ const take = (key: string): string | undefined => {
140+ const v = flags.get(key);
141+ flags.delete(key);
142+ return v;
143+ };
144+ const number = (key: string, dflt: number): number => {
145+ const raw = take(key);
146+ if (raw === undefined) return dflt;
147+ const v = Number(raw);
148+ if (!Number.isFinite(v)) throw new Error(`--${key} must be a number (got '${raw}')`);
149+ return v;
150+ };
151+ const count = (key: string, dflt: number, min: number): number => {
152+ const v = number(key, dflt);
153+ if (!Number.isInteger(v) || v < min) {
154+ throw new Error(`--${key} must be an integer >= ${min} (got '${v}')`);
155+ }
156+ return v;
157+ };
158+
159+ const presetKey = take('preset') ?? presets[0].key;
160+ const { model, params } = resolvePreset(presetKey);
161+ const geometryKey = take('geometry') ?? DEFAULT_GEOMETRY_KEY;
162+ const { geometry, params: geometryParams } = resolveGeometry(geometryKey);
163+ const spec: RunSpec = {
164+ preset: presetKey,
165+ lmax: count('lmax', DEFAULT_LMAX, 1),
166+ seed: number('seed', DEFAULT_SEED),
167+ steps: count('steps', DEFAULT_STEPS, 1),
168+ warmup: count('warmup', DEFAULT_WARMUP, 0),
169+ params,
170+ geometry: geometryKey,
171+ geometryParams,
172+ niter: count('niter', DEFAULT_NITER, 0),
173+ };
174+ const readInto = (into: Params, key: string, flag: string): void => {
175+ const raw = take(flag);
176+ if (raw === undefined) return;
177+ const v = Number(raw);
178+ if (!Number.isFinite(v)) throw new Error(`--${flag} must be a number (got '${raw}')`);
179+ into[key] = v;
180+ };
181+ for (const p of model.params) readInto(params, p.key, p.key);
182+ for (const p of geometry.params) readInto(geometryParams, p.key, `g${p.key}`);
183+ if (flags.size) {
184+ throw new Error(
185+ `unknown option(s): ${[...flags.keys()].map((k) => `--${k}`).join(', ')}\n` +
186+ `parameters of ${model.label}: ${model.params.map((p) => `--${p.key}`).join(' ')}\n` +
187+ `parameters of ${geometry.label}: ` +
188+ (geometry.params.length
189+ ? geometry.params.map((p) => `--g${p.key}`).join(' ')
190+ : '(none)'),
191+ );
192+ }
193+ return spec;
194+}
src/editor/codeEditor.tsadded+92−0View file
@@ -0,0 +1,92 @@
1+/**
2+ * A textarea with syntax highlighting, by overlay.
3+ *
4+ * A textarea cannot colour its own text, so the highlighted source is rendered
5+ * into a <pre> underneath and the textarea sits on top with transparent text and
6+ * a visible caret. The two must agree on every metric that affects layout —
7+ * font, line height, padding, tab size, wrapping — and their scroll offsets are
8+ * kept in sync, or the colours drift away from the characters.
9+ */
10+import { highlightMatlab } from './matlab.ts';
11+
12+export interface CodeEditorOptions {
13+ textarea: HTMLTextAreaElement;
14+ /** The <pre> behind it, holding the highlighted copy. */
15+ overlay: HTMLElement;
16+ /** Names to mark as host-provided operations. */
17+ external?: ReadonlySet<string>;
18+ /** Called on every edit. */
19+ onInput?: (value: string) => void;
20+}
21+
22+export class CodeEditor {
23+ #textarea: HTMLTextAreaElement;
24+ #overlay: HTMLElement;
25+ #external: ReadonlySet<string>;
26+
27+ constructor(opts: CodeEditorOptions) {
28+ this.#textarea = opts.textarea;
29+ this.#overlay = opts.overlay;
30+ this.#external = opts.external ?? new Set();
31+
32+ this.#textarea.addEventListener('input', () => {
33+ this.#repaint();
34+ opts.onInput?.(this.#textarea.value);
35+ });
36+ // Keep the colours under the characters while scrolling.
37+ this.#textarea.addEventListener('scroll', () => this.#syncScroll());
38+ // Tab should indent rather than leave the editor.
39+ this.#textarea.addEventListener('keydown', (e) => this.#onKeyDown(e));
40+ this.#repaint();
41+ }
42+
43+ get value(): string {
44+ return this.#textarea.value;
45+ }
46+
47+ set value(next: string) {
48+ this.#textarea.value = next;
49+ this.#repaint();
50+ }
51+
52+ focus(): void {
53+ this.#textarea.focus();
54+ }
55+
56+ /** Select a character range, scrolling it into view. */
57+ select(start: number, end: number): void {
58+ this.#textarea.focus();
59+ this.#textarea.setSelectionRange(start, end);
60+ // setSelectionRange does not always scroll; nudge the line into view.
61+ const line = this.#textarea.value.slice(0, start).split('\n').length - 1;
62+ const lineHeight = this.#textarea.scrollHeight / Math.max(1, this.#lineCount());
63+ const target = line * lineHeight - this.#textarea.clientHeight / 2;
64+ this.#textarea.scrollTop = Math.max(0, target);
65+ this.#syncScroll();
66+ }
67+
68+ #lineCount(): number {
69+ return this.#textarea.value.split('\n').length + 1; // +1 for the trailing line
70+ }
71+
72+ #onKeyDown(e: KeyboardEvent): void {
73+ if (e.key !== 'Tab' || e.ctrlKey || e.metaKey || e.altKey) return;
74+ e.preventDefault();
75+ const el = this.#textarea;
76+ const { selectionStart: s, selectionEnd: t, value } = el;
77+ el.value = `${value.slice(0, s)} ${value.slice(t)}`;
78+ el.selectionStart = el.selectionEnd = s + 2;
79+ // Let the input listener repaint and notify, as for any other edit.
80+ el.dispatchEvent(new Event('input'));
81+ }
82+
83+ #repaint(): void {
84+ this.#overlay.innerHTML = highlightMatlab(this.#textarea.value, this.#external);
85+ this.#syncScroll();
86+ }
87+
88+ #syncScroll(): void {
89+ this.#overlay.scrollTop = this.#textarea.scrollTop;
90+ this.#overlay.scrollLeft = this.#textarea.scrollLeft;
91+ }
92+}
src/editor/matlab.tsadded+213−0View file
@@ -0,0 +1,213 @@
1+/**
2+ * A small MATLAB tokenizer, for syntax highlighting the model editor.
3+ *
4+ * Only what highlighting needs — comments, literals, numbers, keywords — and
5+ * deliberately not a parser: numbl does the real parsing, and reports errors
6+ * with positions. Tokens preserve the source text exactly, character for
7+ * character, because the highlighted output is overlaid on a textarea and any
8+ * dropped or added character would shift the two out of alignment.
9+ */
10+
11+export type TokenClass = 'com' | 'str' | 'num' | 'kw' | 'ext';
12+
13+export interface Token {
14+ text: string;
15+ cls: TokenClass | null;
16+}
17+
18+const KEYWORDS = new Set([
19+ 'break', 'case', 'catch', 'classdef', 'continue', 'else', 'elseif', 'end',
20+ 'for', 'function', 'global', 'if', 'otherwise', 'parfor', 'persistent',
21+ 'return', 'spmd', 'switch', 'try', 'while',
22+]);
23+
24+const isIdentStart = (c: string): boolean => /[A-Za-z_]/.test(c);
25+const isIdent = (c: string): boolean => /[A-Za-z0-9_]/.test(c);
26+const isDigit = (c: string): boolean => c >= '0' && c <= '9';
27+
28+/**
29+ * In MATLAB `'` is both the transpose operator and the char-literal delimiter.
30+ * It opens a literal unless it directly follows something that can be
31+ * transposed — a value, a closing bracket, or another transpose.
32+ */
33+function quoteIsTranspose(src: string, at: number): boolean {
34+ for (let i = at - 1; i >= 0; i--) {
35+ const c = src[i];
36+ if (c === ' ' || c === '\t') continue;
37+ return isIdent(c) || c === ')' || c === ']' || c === '}' || c === '.' || c === "'";
38+ }
39+ return false;
40+}
41+
42+/**
43+ * Tokenize `src`. `external` names (the operations the host provides, e.g.
44+ * `synth` / `analys`) get their own class so the boundary between the model and
45+ * what it is given is visible in the editor.
46+ */
47+export function tokenizeMatlab(
48+ src: string,
49+ external: ReadonlySet<string> = new Set(),
50+): Token[] {
51+ const out: Token[] = [];
52+ const push = (text: string, cls: TokenClass | null): void => {
53+ if (!text) return;
54+ const last = out[out.length - 1];
55+ if (last && last.cls === cls) last.text += text;
56+ else out.push({ text, cls });
57+ };
58+
59+ let i = 0;
60+ let atLineStart = true;
61+ let inBlockComment = false;
62+
63+ while (i < src.length) {
64+ const c = src[i];
65+
66+ // Block comments: `%{` and `%}` each alone on their line.
67+ if (atLineStart) {
68+ const eol = src.indexOf('\n', i);
69+ const lineEnd = eol === -1 ? src.length : eol;
70+ const line = src.slice(i, lineEnd);
71+ const trimmed = line.trim();
72+ if (!inBlockComment && trimmed === '%{') inBlockComment = true;
73+ else if (inBlockComment && trimmed === '%}') {
74+ push(line, 'com');
75+ i = lineEnd;
76+ inBlockComment = false;
77+ atLineStart = false;
78+ continue;
79+ }
80+ if (inBlockComment) {
81+ push(line, 'com');
82+ i = lineEnd;
83+ atLineStart = false;
84+ continue;
85+ }
86+ }
87+
88+ if (c === '\n') {
89+ push(c, null);
90+ i++;
91+ atLineStart = true;
92+ continue;
93+ }
94+ if (c === ' ' || c === '\t') {
95+ push(c, null);
96+ i++;
97+ continue;
98+ }
99+ atLineStart = false;
100+
101+ // Line comment, including MATLAB's `%%` section markers.
102+ if (c === '%') {
103+ const eol = src.indexOf('\n', i);
104+ const end = eol === -1 ? src.length : eol;
105+ push(src.slice(i, end), 'com');
106+ i = end;
107+ continue;
108+ }
109+
110+ // Line continuation is an operator, but any trailing text is a comment.
111+ if (c === '.' && src.startsWith('...', i)) {
112+ const eol = src.indexOf('\n', i);
113+ const end = eol === -1 ? src.length : eol;
114+ push('...', null);
115+ push(src.slice(i + 3, end), 'com');
116+ i = end;
117+ continue;
118+ }
119+
120+ // Char literal (or transpose).
121+ if (c === "'") {
122+ if (quoteIsTranspose(src, i)) {
123+ push("'", null);
124+ i++;
125+ continue;
126+ }
127+ let j = i + 1;
128+ while (j < src.length && src[j] !== '\n') {
129+ if (src[j] === "'") {
130+ if (src[j + 1] === "'") j += 2; // escaped quote
131+ else {
132+ j++;
133+ break;
134+ }
135+ } else j++;
136+ }
137+ push(src.slice(i, j), 'str');
138+ i = j;
139+ continue;
140+ }
141+
142+ // Double-quoted string.
143+ if (c === '"') {
144+ let j = i + 1;
145+ while (j < src.length && src[j] !== '\n') {
146+ if (src[j] === '"') {
147+ if (src[j + 1] === '"') j += 2;
148+ else {
149+ j++;
150+ break;
151+ }
152+ } else j++;
153+ }
154+ push(src.slice(i, j), 'str');
155+ i = j;
156+ continue;
157+ }
158+
159+ // Number: 12, 1.5, .5, 1e-3, 2i
160+ if (isDigit(c) || (c === '.' && isDigit(src[i + 1]))) {
161+ let j = i;
162+ while (j < src.length && isDigit(src[j])) j++;
163+ if (src[j] === '.') {
164+ j++;
165+ while (j < src.length && isDigit(src[j])) j++;
166+ }
167+ if (src[j] === 'e' || src[j] === 'E') {
168+ let k = j + 1;
169+ if (src[k] === '+' || src[k] === '-') k++;
170+ if (isDigit(src[k])) {
171+ k++;
172+ while (k < src.length && isDigit(src[k])) k++;
173+ j = k;
174+ }
175+ }
176+ if (src[j] === 'i' || src[j] === 'j') j++;
177+ push(src.slice(i, j), 'num');
178+ i = j;
179+ continue;
180+ }
181+
182+ // Identifier / keyword / external operation.
183+ if (isIdentStart(c)) {
184+ let j = i;
185+ while (j < src.length && isIdent(src[j])) j++;
186+ const word = src.slice(i, j);
187+ push(word, KEYWORDS.has(word) ? 'kw' : external.has(word) ? 'ext' : null);
188+ i = j;
189+ continue;
190+ }
191+
192+ push(c, null);
193+ i++;
194+ }
195+
196+ return out;
197+}
198+
199+const escapeHtml = (s: string): string =>
200+ s.replace(/&/g, '&').replace(/</g, '<').replace(/>/g, '>');
201+
202+/** Highlighted HTML for `src`, safe to assign to innerHTML. */
203+export function highlightMatlab(
204+ src: string,
205+ external: ReadonlySet<string> = new Set(),
206+): string {
207+ const html = tokenizeMatlab(src, external)
208+ .map((t) => (t.cls ? `<span class="tok-${t.cls}">${escapeHtml(t.text)}</span>` : escapeHtml(t.text)))
209+ .join('');
210+ // A trailing newline keeps the last line's box height stable, so the overlay
211+ // and the textarea scroll to the same extent.
212+ return `${html}\n`;
213+}
src/geom/geometry.tsadded+215−0View file
@@ -0,0 +1,215 @@
1+/**
2+ * The surface: a .m shape file, compiled and evaluated into spherical-harmonic
3+ * coefficients.
4+ *
5+ * A geometry file is ordinary MATLAB defining one function,
6+ *
7+ * function [gx, gy, gz] = shape(theta, phi, <parameters>)
8+ *
9+ * over the solver's (theta, phi) grid — the same element-wise MATLAB the models
10+ * are written in, compiled by the same backend into the same kind of WGSL
11+ * kernel. It is evaluated once, on the CPU's behalf, and then *analysed*: the
12+ * canonical geometry this project carries is the three sets of coefficients
13+ * `X`, `Y`, `Z`, one per Cartesian component of the embedding.
14+ *
15+ * Going through the coefficients rather than keeping the pointwise values is
16+ * what makes the geometry usable by a spectral method, for two reasons:
17+ *
18+ * - it is exactly band-limited at lmax afterwards, so the surface has as many
19+ * derivatives as the scheme needs and no aliased content the solver cannot
20+ * see. `x`, `y`, `z` below are the synthesis of the coefficients, not the
21+ * raw output of the .m — the shape actually being solved on, which for a
22+ * shape with sharp features is not quite the shape that was written down.
23+ * - it can be evaluated on any grid. The renderer draws the surface on the
24+ * (possibly finer) display grid by synthesizing the same coefficients
25+ * there, which is exact interpolation rather than subdivision — the same
26+ * argument that lets the species fields be oversampled.
27+ *
28+ * The unit sphere is the case where `x`, `y`, `z` are pure degree-1 harmonics
29+ * and everything downstream reduces to turing-sphere.
30+ */
31+import { ShtPlan } from '../sht/sht.ts';
32+import type { ShtConfig } from '../sht/layout.ts';
33+import { HostBuffers, ModelPlan } from '../mgpu/plan.ts';
34+import { CompiledModel, type Binding } from '../mgpu/compile.ts';
35+import { inFunction, inFunctionAsync, inModel } from '../mgpu/errors.ts';
36+import type { ModelParams } from '../mgpu/model.ts';
37+
38+/** The function a geometry file must define. */
39+export const SHAPE_FN = 'shape';
40+
41+export interface GeometryOptions {
42+ device: GPUDevice;
43+ /** The solver's transform plan — the grid the shape is evaluated on. */
44+ sht: ShtPlan;
45+ cfg: ShtConfig;
46+ /** Geometry source (.m text). */
47+ source: string;
48+ /** Parameter names the .m may take beyond `theta` and `phi`. */
49+ paramNames: string[];
50+ params: ModelParams;
51+}
52+
53+export class Geometry {
54+ /** Coordinates on the solver grid, npts each — synthesis of the coefficients. */
55+ readonly x: Float32Array;
56+ readonly y: Float32Array;
57+ readonly z: Float32Array;
58+ /** Their spherical-harmonic coefficients, 2 x nlm each. */
59+ readonly X: Float32Array;
60+ readonly Y: Float32Array;
61+ readonly Z: Float32Array;
62+
63+ private constructor(init: {
64+ x: Float32Array; y: Float32Array; z: Float32Array;
65+ X: Float32Array; Y: Float32Array; Z: Float32Array;
66+ }) {
67+ this.x = init.x;
68+ this.y = init.y;
69+ this.z = init.z;
70+ this.X = init.X;
71+ this.Y = init.Y;
72+ this.Z = init.Z;
73+ }
74+
75+ /**
76+ * Compile the shape file, evaluate it once on the solver grid, and reduce it
77+ * to coefficients. Everything here happens at build time — a geometry never
78+ * takes part in the timestep — so it reads back through the CPU freely.
79+ */
80+ static async create(opts: GeometryOptions): Promise<Geometry> {
81+ const { device, sht, cfg, source, paramNames, params } = opts;
82+ const npts = cfg.nlat * cfg.nphi;
83+ const nlm = sht.nlm;
84+
85+ const bindings: Record<string, Binding> = {
86+ theta: { kind: 'tensor', shape: [npts, 1] },
87+ phi: { kind: 'tensor', shape: [npts, 1] },
88+ npts: { kind: 'const', value: npts },
89+ };
90+ for (const p of paramNames) bindings[p] = { kind: 'param' };
91+
92+ const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
93+ const fn = inFunction(SHAPE_FN, () => compiled.specialize(SHAPE_FN, 3));
94+ compiled.finish();
95+
96+ const host = new HostBuffers(device);
97+ host.ensure('theta', npts);
98+ host.ensure('phi', npts);
99+
100+ const plan = await inFunctionAsync(SHAPE_FN, () =>
101+ // Nothing feeds back: the three outputs are read once and the plan is
102+ // thrown away.
103+ ModelPlan.create(device, sht, { fn, feedback: [null, null, null] }, host),
104+ );
105+
106+ try {
107+ const { theta, phi } = gridAngles(sht, cfg);
108+ host.upload('theta', theta);
109+ host.upload('phi', phi);
110+ plan.setParams(params);
111+
112+ const enc = device.createCommandEncoder({ label: 'geometry-shape' });
113+ plan.encodeSteps(enc, 1);
114+ device.queue.submit([enc.finish()]);
115+
116+ const raw = await Promise.all(
117+ fn.outputs.map((out) => readBuffer(device, plan, out.name, npts)),
118+ );
119+ // Coefficients first, then back to the grid: what the solver and the
120+ // renderer both see is the band-limited surface, not the raw .m output.
121+ const [X, Y, Z] = [
122+ await sht.analys(raw[0]),
123+ await sht.analys(raw[1]),
124+ await sht.analys(raw[2]),
125+ ];
126+ const [x, y, z] = [
127+ await sht.synth(X),
128+ await sht.synth(Y),
129+ await sht.synth(Z),
130+ ];
131+ return new Geometry({ x, y, z, X, Y, Z });
132+ } finally {
133+ plan.destroy();
134+ host.destroy();
135+ }
136+ }
137+
138+ /**
139+ * The surface evaluated on another plan's grid, as interleaved xyz vertex
140+ * positions (nlat * nphi * 3) — for rendering at display resolution. Exact
141+ * interpolation: the same coefficients, more evaluation points.
142+ */
143+ async positionsOn(view: ShtPlan): Promise<Float32Array> {
144+ const [x, y, z] = [
145+ await view.synth(this.X),
146+ await view.synth(this.Y),
147+ await view.synth(this.Z),
148+ ];
149+ const out = new Float32Array(x.length * 3);
150+ for (let i = 0; i < x.length; i++) {
151+ out[3 * i] = x[i];
152+ out[3 * i + 1] = y[i];
153+ out[3 * i + 2] = z[i];
154+ }
155+ return out;
156+ }
157+
158+ /** How far the surface departs from the unit sphere, as min/max radius. */
159+ radiusRange(): { lo: number; hi: number } {
160+ let lo = Infinity;
161+ let hi = -Infinity;
162+ for (let i = 0; i < this.x.length; i++) {
163+ const r = Math.hypot(this.x[i], this.y[i], this.z[i]);
164+ if (r < lo) lo = r;
165+ if (r > hi) hi = r;
166+ }
167+ return { lo, hi };
168+ }
169+}
170+
171+/** The (theta, phi) of every grid point, flattened phi-fastest as the fields are. */
172+function gridAngles(
173+ sht: ShtPlan,
174+ cfg: ShtConfig,
175+): { theta: Float32Array; phi: Float32Array } {
176+ const { nlat, nphi } = cfg;
177+ const theta = new Float32Array(nlat * nphi);
178+ const phi = new Float32Array(nlat * nphi);
179+ for (let i = 0; i < nlat; i++) {
180+ const th = Math.acos(Math.max(-1, Math.min(1, sht.cosTheta[i])));
181+ for (let j = 0; j < nphi; j++) {
182+ theta[i * nphi + j] = th;
183+ phi[i * nphi + j] = (2 * Math.PI * j) / nphi;
184+ }
185+ }
186+ return { theta, phi };
187+}
188+
189+async function readBuffer(
190+ device: GPUDevice,
191+ plan: ModelPlan,
192+ name: string,
193+ count: number,
194+): Promise<Float32Array> {
195+ const buffer = plan.buffer(name);
196+ if (!buffer) {
197+ throw new Error(`the geometry never assigns '${name}'`);
198+ }
199+ const staging = device.createBuffer({
200+ label: `geometry-read-${name}`,
201+ size: 4 * count,
202+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
203+ });
204+ try {
205+ const enc = device.createCommandEncoder({ label: `geometry-read-${name}` });
206+ enc.copyBufferToBuffer(buffer, 0, staging, 0, 4 * count);
207+ device.queue.submit([enc.finish()]);
208+ await staging.mapAsync(GPUMapMode.READ);
209+ const out = new Float32Array(staging.getMappedRange().slice(0));
210+ staging.unmap();
211+ return out;
212+ } finally {
213+ staging.destroy();
214+ }
215+}
src/geom/registry.tsadded+85−0View file
@@ -0,0 +1,85 @@
1+/**
2+ * The available geometries: their MATLAB source, and the metadata the host owns.
3+ *
4+ * The same split as the model registry — the .m is the shape, everything around
5+ * it (parameter names, defaults, slider ranges) lives here, and the .m declares
6+ * which of them it wants by naming them as arguments.
7+ *
8+ * `sphere` is first and is not merely one entry among several: it is the case
9+ * the whole project is checked against, where the surface is exactly the unit
10+ * sphere and every result must match turing-sphere.
11+ */
12+import sphereSource from '../../geometries/sphere.m?raw';
13+import ellipsoidSource from '../../geometries/ellipsoid.m?raw';
14+import peanutSource from '../../geometries/peanut.m?raw';
15+import bumpySource from '../../geometries/bumpy.m?raw';
16+import type { ParamSpec, Params } from '../mgpu/registry.ts';
17+
18+export interface MGeometry {
19+ key: string;
20+ label: string;
21+ blurb: string;
22+ params: ParamSpec[];
23+ /** MATLAB source — the shape itself. */
24+ source: string;
25+}
26+
27+const sphere: MGeometry = {
28+ key: 'sphere',
29+ label: 'Sphere',
30+ blurb: 'The unit sphere. The reference case: identical to turing-sphere.',
31+ params: [],
32+ source: sphereSource,
33+};
34+
35+const ellipsoid: MGeometry = {
36+ key: 'ellipsoid',
37+ label: 'Ellipsoid',
38+ blurb: 'Each axis scaled independently. Exactly degree 1, so nothing is lost to band-limiting.',
39+ params: [
40+ { key: 'ax', label: 'a', value: 1.5, min: 0.2, max: 3, step: 0.05 },
41+ { key: 'ay', label: 'b', value: 1, min: 0.2, max: 3, step: 0.05 },
42+ { key: 'az', label: 'c', value: 0.6, min: 0.2, max: 3, step: 0.05 },
43+ ],
44+ source: ellipsoidSource,
45+};
46+
47+const peanut: MGeometry = {
48+ key: 'peanut',
49+ label: 'Peanut',
50+ blurb: 'A dumbbell pinched at the equator: two positively curved ends joined by a saddle.',
51+ params: [
52+ { key: 'waist', label: 'waist', value: 0.6, min: 0, max: 0.9, step: 0.05 },
53+ { key: 'stretch', label: 'stretch', value: 0.6, min: 0, max: 2, step: 0.05 },
54+ ],
55+ source: peanutSource,
56+};
57+
58+const bumpy: MGeometry = {
59+ key: 'bumpy',
60+ label: 'Bumpy',
61+ blurb: 'Equatorial lobes plus a pear-shaped offset — curvature that varies in both angles.',
62+ params: [
63+ { key: 'amp', label: 'amp', value: 0.3, min: 0, max: 0.6, step: 0.02 },
64+ { key: 'nlobe', label: 'lobes', value: 5, min: 1, max: 12, step: 1 },
65+ { key: 'pear', label: 'pear', value: 0.15, min: -0.4, max: 0.4, step: 0.05 },
66+ ],
67+ source: bumpySource,
68+};
69+
70+export const mGeometries: MGeometry[] = [sphere, ellipsoid, peanut, bumpy];
71+
72+export const mGeometryByKey = (key: string): MGeometry | undefined =>
73+ mGeometries.find((g) => g.key === key);
74+
75+export const defaultGeometryParams = (g: MGeometry): Params =>
76+ Object.fromEntries(g.params.map((p) => [p.key, p.value]));
77+
78+/** The geometry every result is checked against, and what a caller who names
79+ * none gets: the case where the solver is exact. */
80+export const SPHERE_KEY = 'sphere';
81+
82+/** What the app and the benchmark start on. Not the sphere: this project
83+ * exists for the other shapes, and opening on the reference case would hide
84+ * the one thing it adds. */
85+export const DEFAULT_GEOMETRY_KEY = 'ellipsoid';
src/main.tsadded+1021−0View file
@@ -0,0 +1,1021 @@
1+import { requestShtDevice, describeAdapter } from './sht/sht.ts';
2+import { gridForLmax } from './sht/layout.ts';
3+import { ModelSession } from './mgpu/session.ts';
4+import { mModelByKey, presets, type MModel, type Params } from './mgpu/registry.ts';
5+import { ModelCompileError, formatFailure } from './mgpu/errors.ts';
6+import { EXTERNAL_OPS } from './mgpu/externals.ts';
7+import { CodeEditor } from './editor/codeEditor.ts';
8+import {
9+ formatCommand,
10+ resolvePreset,
11+ DEFAULT_STEPS,
12+ DEFAULT_WARMUP,
13+ type RunSpec,
14+} from './bench/runSpec.ts';
15+import {
16+ mGeometries,
17+ mGeometryByKey,
18+ defaultGeometryParams,
19+ SPHERE_KEY,
20+ DEFAULT_GEOMETRY_KEY,
21+ type MGeometry,
22+} from './geom/registry.ts';
23+import {
24+ buildTopology,
25+ fillFieldValues,
26+ fillPositions,
27+ fillColors,
28+ type SphereMeshTopology,
29+} from './render/sphereMesh.ts';
30+import { SphereScene } from './render/SphereScene.ts';
31+import { Colorbar, fmtValue } from './render/colorbar.ts';
32+import { colormaps, colormapNames } from './render/colormaps.ts';
33+import { MovieRecorder } from './render/movie.ts';
34+
35+const $ = <T extends HTMLElement>(id: string): T =>
36+ document.getElementById(id) as T;
37+
38+const elModel = $<HTMLSelectElement>('model');
39+const elGeometry = $<HTMLSelectElement>('geometry');
40+const elMorph = $<HTMLInputElement>('morph');
41+const elNiter = $<HTMLSelectElement>('niter');
42+const elLmax = $<HTMLSelectElement>('lmax');
43+const elOversample = $<HTMLSelectElement>('oversample');
44+const elColormap = $<HTMLSelectElement>('colormap');
45+const elRunPause = $<HTMLButtonElement>('runpause');
46+const elBenchmark = $<HTMLButtonElement>('benchmark');
47+const elReseed = $<HTMLButtonElement>('reseed');
48+const elResetView = $<HTMLButtonElement>('resetview');
49+const elMovieToggle = $<HTMLButtonElement>('movietoggle');
50+const elMovieBar = $('moviebar');
51+const elMovieSpeed = $<HTMLSelectElement>('moviespeed');
52+const elMovieRes = $<HTMLSelectElement>('movieres');
53+const elMovieRotate = $<HTMLInputElement>('movierotate');
54+const elMovie = $<HTMLButtonElement>('movie');
55+const elParams = $('params');
56+const elGeomParams = $('geomparams');
57+const elGeomNote = $('geomnote');
58+const elPanels = $('panels');
59+const elStats = $('stats');
60+const elBenchResult = $('benchresult');
61+const elCmd = $('cmd');
62+const elCopyCmd = $<HTMLButtonElement>('copycmd');
63+const elBlurb = $('blurb');
64+const elErr = $('err');
65+const elSource = $<HTMLTextAreaElement>('source');
66+const elHighlight = $('highlight');
67+const elCompiled = $('compiled');
68+const elEditorTitle = $('editor-title');
69+const elEditorFile = $<HTMLSelectElement>('editor-file');
70+const elRecompile = $<HTMLButtonElement>('recompile');
71+const elRevert = $<HTMLButtonElement>('revert');
72+
73+for (const p of presets) {
74+ const o = document.createElement('option');
75+ o.value = p.key;
76+ o.textContent = p.label;
77+ elModel.append(o);
78+}
79+for (const g of mGeometries) {
80+ const o = document.createElement('option');
81+ o.value = g.key;
82+ o.textContent = g.label;
83+ elGeometry.append(o);
84+}
85+for (const [value, label] of [['model', 'the solver'], ['geometry', 'the surface']]) {
86+ const o = document.createElement('option');
87+ o.value = value;
88+ o.textContent = label;
89+ elEditorFile.append(o);
90+}
91+for (const name of colormapNames) {
92+ const o = document.createElement('option');
93+ o.value = name;
94+ o.textContent = name;
95+ elColormap.append(o);
96+}
97+elColormap.value = 'jet';
98+
99+/** Whichever .m is open: the solver or the surface. Both are MATLAB, compiled
100+ * by the same backend, so one editor serves both. The host-provided operations
101+ * are marked so the boundary between the file and what it is given is
102+ * visible. */
103+const editor = new CodeEditor({
104+ textarea: elSource,
105+ overlay: elHighlight,
106+ external: EXTERNAL_OPS,
107+ onInput: (value) => {
108+ if (editing === 'geometry') editedGeomSource = value;
109+ else editedSource = value;
110+ elRecompile.textContent = 'Recompile *';
111+ },
112+});
113+
114+/** Timesteps submitted per rendered frame. Nothing is read back between them,
115+ * so the batch costs one submit and one readback regardless of size. */
116+const STEPS_PER_FRAME = 4;
117+
118+/**
119+ * Steps in a solver-timing burst, and how often to run one.
120+ *
121+ * Timing the solver needs a `queue.onSubmittedWorkDone()` to know the work
122+ * finished, and in a browser that is an IPC round trip into the GPU process — a
123+ * fixed cost of a few milliseconds. Spread over one frame's four steps it would
124+ * swamp them on a fast GPU and make the solver look far slower than it is. So the
125+ * rate is measured in an occasional larger batch, where the single sync is
126+ * amortized the way the desktop benchmark amortizes its own. The state is
127+ * snapshotted and restored around the batch, so measuring never advances the
128+ * simulation — otherwise the pattern would visibly lurch forward at every
129+ * measurement.
130+ */
131+const MEASURE_BURST = 32;
132+const MEASURE_EVERY_MS = 2000;
133+
134+/**
135+ * 'auto' display oversampling targets this many render latitudes: the factor is
136+ * the smallest power of two (up to 4) that reaches it. A solver grid already
137+ * this fine gains nothing visually and is not oversampled.
138+ */
139+const AUTO_RENDER_NLAT = 256;
140+
141+/** The display oversampling factor the UI currently asks for. */
142+function resolveOversample(): number {
143+ if (elOversample.value !== 'auto') return Number(elOversample.value);
144+ const { nlat } = gridForLmax(Number(elLmax.value), model.pdeg);
145+ let os = 1;
146+ while (os < 4 && os * nlat < AUTO_RENDER_NLAT) os *= 2;
147+ return os;
148+}
149+
150+/**
151+ * Movie frame rate, and a cap on frames per movie. Playback speed comes from
152+ * the UI, in simulation-time units per second of video; the movie's length is
153+ * the run's t at that speed, and the frame count follows from it — capped by
154+ * the run's own step count (a step is at most one frame) and by
155+ * MOVIE_MAX_FRAMES to bound encode time and file size. Frame timestamps are
156+ * derived from simulation time, so a capped movie keeps its duration and
157+ * speed exactly, at a lower effective frame rate.
158+ */
159+const MOVIE_FPS = 30;
160+const MOVIE_MAX_FRAMES = 3600;
161+
162+/** Movie auto-rotation: camera revolutions per second of video. Measured in
163+ * video time, so the orbit pace on screen is the same at every export speed. */
164+const MOVIE_ROTATE_RPS = 1 / 120;
165+
166+// ---------------------------------------------------------------- state
167+let device: GPUDevice | null = null;
168+let session: ModelSession | null = null;
169+let topo: SphereMeshTopology | null = null;
170+let scenes: SphereScene[] = [];
171+let colorbars: Colorbar[] = [];
172+let valueBufs: Float32Array[] = [];
173+let colorBufs: Float32Array[] = [];
174+let ranges: { lo: number; hi: number }[] = [];
175+let resizeObs: ResizeObserver | null = null;
176+
177+const initial = resolvePreset(presets[0].key);
178+let model: MModel = mModelByKey(initial.model.key)!;
179+let params: Params = initial.params;
180+let geometry: MGeometry = mGeometryByKey(DEFAULT_GEOMETRY_KEY)!;
181+let geomParams: Params = defaultGeometryParams(geometry);
182+/** Which file the editor is showing. */
183+let editing: 'model' | 'geometry' = 'model';
184+/** Each .m as edited in the page; `null` while it matches the file. */
185+let editedSource: string | null = null;
186+let editedGeomSource: string | null = null;
187+/** Sphere (0) to surface (1). Display only; does not touch the solver. */
188+let morph = 1;
189+let seed = 1;
190+let running = false;
191+let adapterName = '';
192+let pumping = false;
193+let movieBusy = false;
194+let movieCancel = false;
195+let solverMs = 0;
196+let frameMs = 0;
197+let lastMeasure = 0;
198+let generation = 0; // bumped on every rebuild to cancel stale pumps
199+/** Surface coordinates on the render grid, interleaved xyz; null before the
200+ * first build. Kept so the morph slider can re-fill positions without
201+ * re-synthesizing. */
202+let coords: Float32Array | null = null;
203+let posBuf: Float32Array | null = null;
204+
205+const source = (): string => editedSource ?? model.source;
206+const geomSource = (): string => editedGeomSource ?? geometry.source;
207+
208+// ---------------------------------------------------------------- UI wiring
209+function buildParamInputs(): void {
210+ elParams.replaceChildren();
211+ for (const spec of model.params) {
212+ const label = document.createElement('label');
213+ label.textContent = `${spec.label} `;
214+ const input = document.createElement('input');
215+ input.type = 'number';
216+ input.min = String(spec.min);
217+ input.max = String(spec.max);
218+ input.step = String(spec.step);
219+ input.value = String(params[spec.key]);
220+ input.addEventListener('change', () => {
221+ const v = Number(input.value);
222+ if (Number.isFinite(v)) params[spec.key] = v;
223+ // Parameters are uniforms, not constants baked into the kernels, so a
224+ // change costs an upload rather than a recompile.
225+ session?.setParams(params);
226+ updateCommand();
227+ });
228+ label.append(input);
229+ elParams.append(label);
230+ }
231+}
232+
233+/**
234+ * The shape's own parameters. Unlike the model's, these are NOT uniforms: the
235+ * surface is evaluated once at build time and reduced to coefficients, so
236+ * moving one rebuilds the geometry (and with it the mesh), though not the
237+ * simulation's compiled step.
238+ */
239+function buildGeomParamInputs(): void {
240+ elGeomParams.replaceChildren();
241+ if (geometry.params.length === 0) return;
242+ const tag = document.createElement('label');
243+ tag.textContent = `${geometry.key}.m`;
244+ elGeomParams.append(tag);
245+ for (const spec of geometry.params) {
246+ const label = document.createElement('label');
247+ label.textContent = `${spec.label} `;
248+ const input = document.createElement('input');
249+ input.type = 'number';
250+ input.min = String(spec.min);
251+ input.max = String(spec.max);
252+ input.step = String(spec.step);
253+ input.value = String(geomParams[spec.key]);
254+ input.addEventListener('change', () => {
255+ const v = Number(input.value);
256+ if (Number.isFinite(v)) geomParams[spec.key] = v;
257+ viewChange = viewChange.then(() => applyGeometry());
258+ });
259+ label.append(input);
260+ elGeomParams.append(label);
261+ }
262+}
263+
264+function applyPreset(presetKey: string): void {
265+ const resolved = resolvePreset(presetKey);
266+ const next = mModelByKey(resolved.model.key);
267+ if (!next) {
268+ elErr.textContent = `No .m model for '${resolved.model.key}'`;
269+ return;
270+ }
271+ model = next;
272+ params = resolved.params;
273+ editedSource = null;
274+ buildParamInputs();
275+ elBlurb.textContent = model.blurb;
276+ showEditorFile();
277+ updateCommand();
278+}
279+
280+function applyGeometryChoice(key: string): void {
281+ const next = mGeometryByKey(key);
282+ if (!next) {
283+ elErr.textContent = `No .m geometry for '${key}'`;
284+ return;
285+ }
286+ geometry = next;
287+ geomParams = defaultGeometryParams(geometry);
288+ editedGeomSource = null;
289+ buildGeomParamInputs();
290+ showEditorFile();
291+}
292+
293+/** Load the chosen file into the editor, keeping any unsaved edit to it. */
294+function showEditorFile(): void {
295+ editing = elEditorFile.value === 'geometry' ? 'geometry' : 'model';
296+ if (editing === 'geometry') {
297+ editor.value = geomSource();
298+ elEditorTitle.textContent = `geometries/${geometry.key}.m — shape(), compiled to WebGPU`;
299+ } else {
300+ editor.value = source();
301+ elEditorTitle.textContent = `models/${model.key}.m — init() and step(), compiled to WebGPU`;
302+ }
303+}
304+
305+/** The run currently on screen, as the benchmark's RunSpec. */
306+function currentSpec(): RunSpec {
307+ return {
308+ preset: elModel.value,
309+ lmax: Number(elLmax.value),
310+ seed,
311+ steps: DEFAULT_STEPS,
312+ warmup: DEFAULT_WARMUP,
313+ params,
314+ geometry: geometry.key,
315+ geometryParams: geomParams,
316+ niter: Number(elNiter.value),
317+ };
318+}
319+
320+function updateCommand(): void {
321+ elCmd.textContent = formatCommand(currentSpec());
322+}
323+
324+elModel.addEventListener('change', () => {
325+ applyPreset(elModel.value);
326+ void rebuild();
327+});
328+elLmax.addEventListener('change', () => void rebuild());
329+// The solve iteration count is unrolled into the compiled step, so unlike a
330+// parameter it cannot be changed without recompiling.
331+elNiter.addEventListener('change', () => void rebuild());
332+// Oversampling and geometry are display-or-data changes, not code ones, so
333+// they swap things in place rather than rebuilding the run. Serialized through
334+// one chain: a rapid second change waits its turn.
335+let viewChange = Promise.resolve();
336+elOversample.addEventListener('change', () => {
337+ viewChange = viewChange.then(() => applyOversample());
338+});
339+elGeometry.addEventListener('change', () => {
340+ applyGeometryChoice(elGeometry.value);
341+ viewChange = viewChange.then(() => applyGeometry());
342+});
343+// Morph is pure rendering: no readback, no GPU work, just the vertex buffer.
344+elMorph.addEventListener('input', () => {
345+ morph = Number(elMorph.value);
346+ applyMorph();
347+});
348+elColormap.addEventListener('change', () => void draw());
349+elEditorFile.addEventListener('change', () => showEditorFile());
350+
351+function setRunning(next: boolean): void {
352+ running = next;
353+ elRunPause.textContent = running ? 'Pause' : 'Run';
354+ if (running) void pump();
355+}
356+
357+elRunPause.addEventListener('click', () => setRunning(!running));
358+elBenchmark.addEventListener('click', () => void benchmark());
359+elReseed.addEventListener('click', () => {
360+ seed = (Math.random() * 2 ** 31) >>> 0;
361+ setRunning(false);
362+ updateCommand();
363+ void reseed();
364+});
365+elResetView.addEventListener('click', () => {
366+ for (const s of scenes) s.resetCamera();
367+});
368+elMovieToggle.addEventListener('click', () => {
369+ elMovieBar.hidden = !elMovieBar.hidden;
370+});
371+elMovie.addEventListener('click', () => {
372+ if (movieBusy) movieCancel = true;
373+ else void recordMovie();
374+});
375+
376+elRecompile.addEventListener('click', () => {
377+ if (editing === 'geometry') editedGeomSource = editor.value;
378+ else editedSource = editor.value;
379+ void rebuild();
380+});
381+elRevert.addEventListener('click', () => {
382+ if (editing === 'geometry') editedGeomSource = null;
383+ else editedSource = null;
384+ showEditorFile();
385+ void rebuild();
386+});
387+
388+// The command reproduces this run's parameters on the desktop; keep it
389+// selectable even where the clipboard API is unavailable.
390+elCopyCmd.addEventListener('click', () => {
391+ const text = elCmd.textContent ?? '';
392+ const flash = (msg: string): void => {
393+ elCopyCmd.textContent = msg;
394+ setTimeout(() => (elCopyCmd.textContent = 'Copy'), 1200);
395+ };
396+ const selectCommand = (): void => {
397+ const range = document.createRange();
398+ range.selectNodeContents(elCmd);
399+ const sel = getSelection();
400+ sel?.removeAllRanges();
401+ sel?.addRange(range);
402+ flash('Selected');
403+ };
404+ if (!navigator.clipboard) return selectCommand();
405+ navigator.clipboard.writeText(text).then(() => flash('Copied'), selectCommand);
406+});
407+
408+// ---------------------------------------------------------------- setup
409+function disposeView(): void {
410+ for (const s of scenes) s.dispose();
411+ scenes = [];
412+ colorbars = [];
413+ topo = null;
414+ coords = null;
415+ posBuf = null;
416+ resizeObs?.disconnect();
417+ resizeObs = null;
418+ elPanels.replaceChildren();
419+}
420+
421+/**
422+ * Build the mesh, scenes, colorbars and per-species buffers on the current
423+ * render grid, from surface coordinates already synthesized there. Call
424+ * disposeView() first. The color ranges are kept if present, so a display-only
425+ * rebuild (an oversampling change) does not pop the shading; a full rebuild
426+ * clears `ranges` beforehand.
427+ */
428+function buildView(surface: Float32Array): void {
429+ if (!session) return;
430+ const view = session.viewSht;
431+ const { nphi } = view.cfg;
432+ const phi = new Float64Array(nphi);
433+ for (let j = 0; j < nphi; j++) phi[j] = (2 * Math.PI * j) / nphi;
434+ topo = buildTopology(view.cosTheta, phi);
435+ coords = surface;
436+ posBuf = new Float32Array(topo.numVertices * 3);
437+ fillPositions(posBuf, coords, topo, morph);
438+
439+ const sphereBg = getComputedStyle(document.documentElement)
440+ .getPropertyValue('--sphere-bg')
441+ .trim();
442+ for (let k = 0; k < model.species.length; k++) {
443+ const panel = document.createElement('div');
444+ panel.className = 'panel';
445+ const box = document.createElement('div');
446+ box.className = 'sphere-box';
447+ const tag = document.createElement('div');
448+ tag.className = 'species-tag';
449+ tag.textContent = model.species[k];
450+ box.append(tag);
451+ const side = document.createElement('div');
452+ panel.append(box, side);
453+ elPanels.append(panel);
454+
455+ const scene = new SphereScene(
456+ box,
457+ topo.numVertices,
458+ topo.indices,
459+ // Each scene owns its position buffer: three.js uploads from it, and the
460+ // morph rewrites all of them from the one shared `coords`.
461+ Float32Array.from(posBuf),
462+ sphereBg || undefined,
463+ );
464+ scene.fitCamera();
465+ scenes.push(scene);
466+ colorbars.push(new Colorbar(side));
467+ valueBufs[k] = new Float32Array(topo.numVertices);
468+ colorBufs[k] = new Float32Array(topo.numVertices * 3);
469+ if (!ranges[k]) ranges[k] = { lo: NaN, hi: NaN };
470+ }
471+ for (let k = 1; k < scenes.length; k++) scenes[0].syncCamerasWith(scenes[k]);
472+
473+ resizeObs = new ResizeObserver(() => {
474+ const boxes = elPanels.querySelectorAll<HTMLElement>('.sphere-box');
475+ boxes.forEach((box, i) => {
476+ scenes[i]?.resize(box.clientWidth, box.clientHeight);
477+ });
478+ });
479+ elPanels
480+ .querySelectorAll<HTMLElement>('.sphere-box')
481+ .forEach((box) => resizeObs!.observe(box));
482+}
483+
484+/**
485+ * Apply the UI's oversampling choice to the running session. Display-only: the
486+ * session and its state survive; only the display plan, mesh and scenes are
487+ * rebuilt, keeping the camera pose and color ranges. The pump is drained first
488+ * so no readback is in flight on the plan being replaced.
489+ */
490+async function applyOversample(): Promise<void> {
491+ if (!session) return;
492+ const gen = generation;
493+ const os = resolveOversample();
494+ if (os === session.oversample) return;
495+ const wasRunning = running;
496+ setRunning(false);
497+ while (pumping) await nextFrame();
498+ if (gen !== generation || !session) return;
499+ await session.setOversample(os);
500+ if (gen !== generation || !session) return;
501+ const surface = await session.renderPositions();
502+ if (gen !== generation || !session) return;
503+ const cam = scenes[0]?.cameraState();
504+ disposeView();
505+ buildView(surface);
506+ if (cam) for (const s of scenes) s.setCameraState(cam);
507+ await draw();
508+ updateStats();
509+ if (wasRunning) setRunning(true);
510+}
511+
512+/**
513+ * Re-evaluate the surface and swap it in. Data, not code: the compiled step is
514+ * untouched and the simulation keeps its state and its model time, so a shape
515+ * can be changed mid-run. Only the mesh is rebuilt.
516+ */
517+async function applyGeometry(): Promise<void> {
518+ if (!session) return;
519+ const gen = generation;
520+ const wasRunning = running;
521+ setRunning(false);
522+ while (pumping) await nextFrame();
523+ if (gen !== generation || !session) return;
524+ try {
525+ await session.setGeometry(geometry, geomParams, geomSource());
526+ } catch (e) {
527+ reportCompileError(e);
528+ return;
529+ }
530+ if (gen !== generation || !session) return;
531+ const surface = await session.renderPositions();
532+ if (gen !== generation || !session) return;
533+ const cam = scenes[0]?.cameraState();
534+ disposeView();
535+ buildView(surface);
536+ if (cam) for (const s of scenes) s.setCameraState(cam);
537+ elErr.textContent = '';
538+ await draw();
539+ updateGeomNote();
540+ updateStats();
541+ if (wasRunning) setRunning(true);
542+}
543+
544+/** Re-place the vertices for the current morph. No GPU work and no readback —
545+ * the surface is already on the CPU, so this is a buffer fill per panel. */
546+function applyMorph(): void {
547+ if (!topo || !coords || !posBuf) return;
548+ fillPositions(posBuf, coords, topo, morph);
549+ for (const s of scenes) s.updatePositions(posBuf);
550+}
551+
552+/** What the surface is, and the standing caveat about where it is not. */
553+function updateGeomNote(): void {
554+ if (!session) {
555+ elGeomNote.textContent = '';
556+ return;
557+ }
558+ const { lo, hi } = session.geometry.radiusRange();
559+ const isSphere = session.geometryModel.key === SPHERE_KEY;
560+ elGeomNote.innerHTML =
561+ `<b>${session.geometryModel.label}</b> — ${session.geometryModel.blurb} ` +
562+ `Radius ${lo.toFixed(3)}–${hi.toFixed(3)}. ` +
563+ (isSphere
564+ ? 'This is the round-sphere case, so the solver is exact here.'
565+ : '<b>Rendered only:</b> the Laplace–Beltrami operator in the .m is still ' +
566+ 'the round sphere\'s, so the pattern is the sphere\'s pattern painted ' +
567+ 'onto this shape.');
568+}
569+
570+/** Report a compile failure, and select the offending text in the editor. */
571+function reportCompileError(e: unknown): void {
572+ elErr.textContent = formatFailure(e, source());
573+ elCompiled.textContent = '';
574+ if (e instanceof ModelCompileError && e.start !== undefined) {
575+ editor.select(e.start, e.end ?? e.start);
576+ }
577+}
578+
579+async function rebuild(): Promise<void> {
580+ generation++;
581+ const gen = generation;
582+ setRunning(false);
583+ disposeView();
584+ session?.destroy();
585+ session = null;
586+ solverMs = 0;
587+ frameMs = 0;
588+ lastMeasure = 0;
589+ elErr.textContent = '';
590+ updateCommand();
591+ if (!device) return;
592+
593+ try {
594+ session = await ModelSession.create({
595+ device,
596+ model,
597+ params,
598+ lmax: Number(elLmax.value),
599+ source: source(),
600+ oversample: resolveOversample(),
601+ geometry,
602+ geometryParams: geomParams,
603+ geometrySource: geomSource(),
604+ niter: Number(elNiter.value),
605+ });
606+ } catch (e) {
607+ reportCompileError(e);
608+ return;
609+ }
610+ if (gen !== generation) return;
611+
612+ session.seed(seed);
613+
614+ const plan = session.describe();
615+ elCompiled.textContent =
616+ `one step compiled to ${plan.step.length} GPU operations:\n` +
617+ plan.step.map((l) => ` ${l}`).join('\n');
618+ elRecompile.textContent = 'Recompile';
619+
620+ const surface = await session.renderPositions();
621+ if (gen !== generation) return;
622+
623+ ranges = [];
624+ buildView(surface);
625+
626+ await draw();
627+ updateGeomNote();
628+ updateStats();
629+ void pump();
630+}
631+
632+async function reseed(): Promise<void> {
633+ if (!session) return;
634+ const gen = generation;
635+ session.seed(seed);
636+ if (gen !== generation) return;
637+ for (const r of ranges) {
638+ r.lo = NaN;
639+ r.hi = NaN;
640+ }
641+ await draw();
642+ updateStats();
643+}
644+
645+// ---------------------------------------------------------------- drawing
646+async function draw(): Promise<void> {
647+ if (!session || !topo) return;
648+ const gen = generation;
649+ const cmap = colormaps[elColormap.value] ?? colormaps.viridis;
650+ for (let k = 0; k < model.species.length; k++) {
651+ // The one readback per frame — the loop is otherwise entirely on the GPU.
652+ // A rebuild can land while this is in flight and destroy the buffer being
653+ // mapped, which rejects the map; that result is stale anyway, so drop it.
654+ let field: Float32Array;
655+ try {
656+ field = await session.readSpecies(k);
657+ } catch (e) {
658+ if (gen !== generation) return;
659+ throw e;
660+ }
661+ if (gen !== generation || !topo) return;
662+ fillFieldValues(valueBufs[k], field, topo);
663+ let lo = Infinity;
664+ let hi = -Infinity;
665+ for (const v of valueBufs[k]) {
666+ if (v < lo) lo = v;
667+ if (v > hi) hi = v;
668+ }
669+ // smooth the color range in both directions so the shading evolves
670+ // gently as the pattern grows (out-of-range values clamp meanwhile)
671+ const r = ranges[k];
672+ if (!Number.isFinite(r.lo)) {
673+ r.lo = lo;
674+ r.hi = hi;
675+ } else {
676+ const a = 0.15;
677+ r.lo += a * (lo - r.lo);
678+ r.hi += a * (hi - r.hi);
679+ }
680+ if (r.hi - r.lo < 1e-9) {
681+ const mid = (r.hi + r.lo) / 2;
682+ r.lo = mid - 5e-10;
683+ r.hi = mid + 5e-10;
684+ }
685+ fillColors(colorBufs[k], valueBufs[k], r.lo, r.hi, cmap);
686+ scenes[k]?.updateColors(colorBufs[k]);
687+ colorbars[k]?.update(cmap, r.lo, r.hi);
688+ }
689+}
690+
691+function updateStats(): void {
692+ if (!session) return;
693+ const { nlat, nphi } = session.cfg;
694+ const kind = `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`;
695+ const solver =
696+ solverMs > 0
697+ ? `<b>${solverMs.toFixed(2)} ms/step</b> (${(1000 / solverMs).toFixed(0)} steps/s, ` +
698+ `batch of ${MEASURE_BURST}, no readback)`
699+ : '—';
700+ const frame =
701+ frameMs > 0
702+ ? `${frameMs.toFixed(1)} ms/frame (${STEPS_PER_FRAME} steps + readback + render)`
703+ : '—';
704+ const view = session.viewSht.cfg;
705+ const render =
706+ session.oversample > 1
707+ ? ` (display ${view.nlat}×${view.nphi}, ${session.oversample}×)`
708+ : '';
709+ elStats.innerHTML =
710+ `<b>${kind}</b> · grid ${nlat}×${nphi}${render} · nlm ${session.sht.nlm.toLocaleString()} · ` +
711+ `${session.sht.fourierMode.toUpperCase()} · ${session.geometryModel.key} · ` +
712+ `${session.niter} solve iter${session.niter === 1 ? '' : 's'} · ` +
713+ `solver ${solver} · ${frame} · ` +
714+ `t = <b>${session.t.toFixed(2)}</b> (${session.steps} steps)`;
715+}
716+
717+// ---------------------------------------------------------------- sim loop
718+const nextFrame = () => new Promise<number>(requestAnimationFrame);
719+
720+async function pump(): Promise<void> {
721+ if (pumping) return;
722+ pumping = true;
723+ const gen = generation;
724+ try {
725+ while (running && session && gen === generation) {
726+ // Occasionally, a burst purely to measure the solver rate: many steps,
727+ // one sync, nothing read back — directly comparable to the desktop
728+ // benchmark's throughput number. State-preserving: the display and
729+ // model time are unaffected.
730+ if (performance.now() - lastMeasure > MEASURE_EVERY_MS) {
731+ const ms = await session.measure(MEASURE_BURST);
732+ if (gen !== generation) break;
733+ solverMs = ms;
734+ lastMeasure = performance.now();
735+ }
736+
737+ // The frame itself. No explicit sync here — draw()'s readback already
738+ // waits for the steps, so asking twice would only add a round trip.
739+ const t0 = performance.now();
740+ session.step(STEPS_PER_FRAME);
741+ await draw();
742+ if (gen !== generation) break;
743+ frameMs = frameMs === 0
744+ ? performance.now() - t0
745+ : frameMs + 0.05 * (performance.now() - t0 - frameMs);
746+ updateStats();
747+ await nextFrame();
748+ }
749+ if (gen === generation) {
750+ await draw();
751+ updateStats();
752+ }
753+ } finally {
754+ pumping = false;
755+ }
756+}
757+
758+/**
759+ * Sustained solver benchmark, in the page.
760+ *
761+ * The same measurement `npm run bench` makes: batches of steps submitted
762+ * together, waited for, never read back, with no rendering and no animation
763+ * pacing in between. That makes it directly comparable to the terminal number,
764+ * which is the only way to tell a genuinely slower browser GPU stack apart from
765+ * the costs the app adds on top.
766+ *
767+ * It also reports the ramp — the first third of the run against the last. GPUs
768+ * downclock when idle, and an animation-paced loop leaves them idle most of every
769+ * frame, so a large ramp means the app's steady-state number is limited by clocks
770+ * rather than by the work.
771+ *
772+ * These are ordinary steps: the simulation advances by them.
773+ */
774+async function benchmark(): Promise<void> {
775+ if (!session || movieBusy) return;
776+ setRunning(false);
777+ const BATCH = 32;
778+ const DURATION_MS = 2000;
779+ elBenchResult.textContent = 'benchmarking…';
780+ // A movie started mid-benchmark would replay while this loop still steps.
781+ elMovie.disabled = true;
782+ try {
783+ await nextFrame();
784+
785+ const gen = generation;
786+ const perStep: number[] = [];
787+ const t0 = performance.now();
788+ while (performance.now() - t0 < DURATION_MS) {
789+ const b0 = performance.now();
790+ session.step(BATCH);
791+ await session.sync();
792+ if (gen !== generation) return;
793+ perStep.push((performance.now() - b0) / BATCH);
794+ }
795+
796+ const mean = (xs: number[]): number => xs.reduce((a, b) => a + b, 0) / xs.length;
797+ const all = mean(perStep);
798+ const best = Math.min(...perStep);
799+ const third = Math.max(1, Math.floor(perStep.length / 3));
800+ const first = mean(perStep.slice(0, third));
801+ const last = mean(perStep.slice(-third));
802+ const steps = perStep.length * BATCH;
803+
804+ elBenchResult.innerHTML =
805+ `sustained solver: <b>${all.toFixed(2)} ms/step</b> ` +
806+ `(${(1000 / all).toFixed(0)} steps/s) · best ${best.toFixed(2)} · ` +
807+ `ramp ${(first / last).toFixed(2)}× (${first.toFixed(2)} → ${last.toFixed(2)}) · ` +
808+ `${steps} steps in batches of ${BATCH} · ` +
809+ `compare with <code>npm run bench -- --lmax ${session.cfg.lmax}</code>`;
810+ await draw();
811+ updateStats();
812+ } finally {
813+ elMovie.disabled = false;
814+ }
815+}
816+
817+// ---------------------------------------------------------------- movie
818+function saveBlob(blob: Blob, filename: string): void {
819+ const url = URL.createObjectURL(blob);
820+ const a = document.createElement('a');
821+ a.href = url;
822+ a.download = filename;
823+ a.click();
824+ setTimeout(() => URL.revokeObjectURL(url), 10_000);
825+}
826+
827+/** Submit `n` steps in bounded command buffers — a single buffer encoding
828+ * many thousands of steps can exhaust the encoder. */
829+function submitSteps(n: number): void {
830+ while (n > 0 && session) {
831+ const chunk = Math.min(512, n);
832+ session.step(chunk);
833+ n -= chunk;
834+ }
835+}
836+
837+/** While recording, lock everything that could change the run mid-replay;
838+ * the Movie button itself becomes the cancel button. */
839+function setMovieUi(on: boolean): void {
840+ const locked = [
841+ elModel, elGeometry, elMorph, elNiter, elLmax, elOversample, elColormap,
842+ elRunPause, elBenchmark, elReseed, elRecompile, elRevert, elEditorFile,
843+ elMovieSpeed, elMovieRes, elMovieRotate, elMovieToggle,
844+ ];
845+ for (const el of locked) el.disabled = on;
846+ elParams.querySelectorAll('input').forEach((input) => (input.disabled = on));
847+ elGeomParams.querySelectorAll('input').forEach((input) => (input.disabled = on));
848+ elMovie.textContent = on ? 'Cancel · 0%' : 'Export';
849+}
850+
851+/**
852+ * Recompute the run from t = 0 and download it as an MP4.
853+ *
854+ * The movie is not a recording of what already happened — it is the same
855+ * trajectory recomputed: same seed, same source, and the *current* parameters
856+ * and colormap throughout. Determinism makes this exact: after the replay the
857+ * state is where it was, so the one session is reused and the app resumes as
858+ * if nothing happened. Frames are composited from the live panels, so the
859+ * movie shows the spheres at the current camera orientation — and the replay
860+ * doubles as the progress display, since it is visible on screen.
861+ */
862+async function recordMovie(): Promise<void> {
863+ if (!session || movieBusy) return;
864+ if (session.steps === 0) {
865+ elMovie.textContent = 'run first';
866+ setTimeout(() => (elMovie.textContent = 'Export'), 1200);
867+ return;
868+ }
869+ movieBusy = true;
870+ movieCancel = false;
871+ const gen = generation;
872+ setMovieUi(true);
873+ let wasRunning = false;
874+ let total = 0;
875+ let done = 0;
876+ let seeded = false;
877+ let camBefore: ReturnType<SphereScene['cameraState']> | undefined;
878+ try {
879+ // An in-flight display-grid swap replaces the scenes whose canvases the
880+ // recorder captures, and resumes the run when it lands — let it finish.
881+ await viewChange;
882+ if (gen !== generation || !session) return;
883+ wasRunning = running;
884+ setRunning(false);
885+ while (pumping) await nextFrame(); // let an in-flight live frame drain
886+ if (gen !== generation || !session) return;
887+ total = session.steps;
888+ const speed = Number(elMovieSpeed.value) || 10;
889+ const sphere = Number(elMovieRes.value) || 768;
890+ const rotate = elMovieRotate.checked;
891+ if (rotate) camBefore = scenes[0]?.cameraState();
892+ // Render the scenes at exactly the chosen resolution for the recording —
893+ // independent of the window size — and restore afterwards.
894+ for (const s of scenes) s.captureSize(sphere);
895+ const durationS = Math.max(session.t / speed, 2 / MOVIE_FPS);
896+ const frames = Math.max(
897+ 2,
898+ Math.min(Math.round(durationS * MOVIE_FPS) + 1, total + 1, MOVIE_MAX_FRAMES),
899+ );
900+ /** The step index captured as frame `i`; both endpoints land exactly. */
901+ const stepAt = (i: number): number => Math.round((i * total) / (frames - 1));
902+
903+ const title =
904+ (presets.find((p) => p.key === elModel.value)?.label ?? model.label) +
905+ ` on ${geometry.label.toLowerCase()}` +
906+ (editedSource !== null || editedGeomSource !== null ? ' (edited)' : '');
907+ const subtitle = model.params
908+ .map((spec) => `${spec.label} ${fmtValue(params[spec.key])}`)
909+ .join(' · ');
910+ const rec = await MovieRecorder.create({
911+ panels: model.species.map((label, k) => ({ canvas: scenes[k].canvas, label })),
912+ title,
913+ subtitle,
914+ speed,
915+ fps: (frames - 1) / durationS,
916+ sphere,
917+ });
918+
919+ let finished = false;
920+ try {
921+ // Reset the color-range smoothing as a re-seed does, so the shading
922+ // evolves in the movie the way it did live.
923+ session.seed(seed);
924+ seeded = true;
925+ for (const r of ranges) {
926+ r.lo = NaN;
927+ r.hi = NaN;
928+ }
929+ const cmap = colormaps[elColormap.value] ?? colormaps.viridis;
930+ let lastVideoS = 0;
931+ for (let frame = 0; ; ) {
932+ await draw();
933+ if (gen !== generation) return;
934+ if (movieCancel) break;
935+ if (rotate) {
936+ // Advance the orbit by this frame's share of video time; siblings
937+ // follow scenes[0] through the usual camera sync.
938+ const videoS = session.t / speed;
939+ scenes[0]?.orbitBy(2 * Math.PI * MOVIE_ROTATE_RPS * (videoS - lastVideoS));
940+ lastVideoS = videoS;
941+ }
942+ for (const s of scenes) s.renderNow();
943+ await rec.addFrame(
944+ session.t,
945+ model.species.map((_, k) => ({ cmap, lo: ranges[k].lo, hi: ranges[k].hi })),
946+ );
947+ if (++frame >= frames) {
948+ finished = true;
949+ break;
950+ }
951+ const target = stepAt(frame);
952+ submitSteps(target - done);
953+ done = target;
954+ elMovie.textContent = `Cancel · ${Math.round((100 * done) / total)}%`;
955+ }
956+ if (finished) {
957+ const blob = await rec.finish();
958+ saveBlob(
959+ blob,
960+ `turing-surface-${model.key}-${geometry.key}-` +
961+ `t${session.t.toFixed(2)}-${speed}x.mp4`,
962+ );
963+ }
964+ } finally {
965+ if (!finished) rec.cancel();
966+ }
967+ } catch (e) {
968+ elErr.textContent = `movie: ${e instanceof Error ? e.message : e}`;
969+ } finally {
970+ // A cancelled replay stopped short of where the run was; step the
971+ // remainder — determinism makes this land exactly there.
972+ if (seeded && gen === generation && session) {
973+ while (done < total && gen === generation && session) {
974+ const n = Math.min(4096, total - done);
975+ submitSteps(n);
976+ done += n;
977+ elMovie.textContent = `restoring · ${Math.round((100 * done) / total)}%`;
978+ await session.sync();
979+ }
980+ await draw();
981+ updateStats();
982+ }
983+ if (gen === generation) {
984+ for (const s of scenes) s.restoreSize();
985+ }
986+ if (camBefore && gen === generation) {
987+ for (const s of scenes) s.setCameraState(camBefore);
988+ }
989+ movieBusy = false;
990+ setMovieUi(false);
991+ if (gen === generation) setRunning(wasRunning);
992+ }
993+}
994+
995+// ---------------------------------------------------------------- boot
996+async function boot(): Promise<void> {
997+ elModel.value = presets[0].key;
998+ elGeometry.value = DEFAULT_GEOMETRY_KEY;
999+ elMorph.value = String(morph);
1000+ applyGeometryChoice(DEFAULT_GEOMETRY_KEY);
1001+ applyPreset(presets[0].key);
1002+ try {
1003+ device = await requestShtDevice();
1004+ adapterName = await describeAdapter(device);
1005+ } catch (e) {
1006+ device = null;
1007+ elErr.textContent =
1008+ `WebGPU is not available (${e instanceof Error ? e.message : e}). ` +
1009+ `This demo compiles the MATLAB solver to WebGPU compute shaders, so it ` +
1010+ `needs a WebGPU-capable browser (Chrome/Edge 113+).`;
1011+ return;
1012+ }
1013+ device.lost.then((info) => {
1014+ if (info.reason !== 'destroyed') {
1015+ elErr.textContent = `WebGPU device lost: ${info.message}`;
1016+ }
1017+ });
1018+ await rebuild();
1019+}
1020+
1021+void boot();
src/mgpu/compile.tsadded+299−0View file
@@ -0,0 +1,299 @@
1+/**
2+ * MATLAB source -> numbl's JIT IR, ready for the WGSL backend.
3+ *
4+ * A model file defines ordinary MATLAB functions; the host specializes the ones
5+ * it needs (`init`, `step`) for the concrete argument types of the current grid.
6+ * This is exactly how numbl drives its own JIT — the caller supplies argument
7+ * types, and lowering fixes every type and shape from there.
8+ *
9+ * Driving it through function signatures rather than injected scope means the
10+ * .m declares what it needs: each parameter name is matched against what the
11+ * host offers, and a name the host does not provide is a compile error rather
12+ * than a silently undefined variable.
13+ *
14+ * Two numbl passes matter here:
15+ * - `specializeUserFunction` lowers one function to IR, one statement per
16+ * operation (ANF), with every node's type fixed.
17+ * - `inlinePass` then folds single-use temps back into their consumer, so a
18+ * source line like `fu = a - u + u.*u.*v` becomes ONE statement whose RHS is
19+ * an expression tree — i.e. one GPU kernel instead of four.
20+ */
21+import { parseMFile } from 'numbl-src/numbl-core/parser/index.ts';
22+import { Workspace, Lowerer, tensorDouble, scalarDouble } from 'numbl-src/numbl-core/jit/index.ts';
23+import { specializeUserFunction } from 'numbl-src/numbl-core/jit/lowering/specialize.ts';
24+import { inlinePass } from 'numbl-src/numbl-core/jit/codegen/inlinePass.ts';
25+import type { For, IRExpr, IRFunc, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
26+import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
27+import { externalOpFiles, type GridSizes } from './externals.ts';
28+import { ModelCompileError } from './errors.ts';
29+
30+/** What the host can supply for an argument the .m declares. */
31+export type Binding =
32+ /** An array, passed in a GPU buffer. */
33+ | { kind: 'tensor'; shape: number[] }
34+ /** A tunable scalar. Deliberately carries no exact value: an exact scalar
35+ * would be constant-folded into the kernels, so moving a slider would force
36+ * a recompile instead of just rewriting a uniform. */
37+ | { kind: 'param' }
38+ /** A fixed scalar, exact so array constructors reading it keep static
39+ * shapes. */
40+ | { kind: 'const'; value: number };
41+
42+const typeOf = (b: Binding): Type => {
43+ switch (b.kind) {
44+ case 'tensor':
45+ return tensorDouble(b.shape);
46+ case 'param':
47+ return scalarDouble('unknown');
48+ case 'const':
49+ // Carry the sign too: numbl's sign lattice decides, for instance,
50+ // whether sqrt() of a value can go complex.
51+ return scalarDouble(
52+ b.value > 0 ? 'positive' : b.value < 0 ? 'negative' : 'zero',
53+ b.value,
54+ );
55+ }
56+};
57+
58+/** One specialized function, as the planner consumes it. */
59+export interface CompiledFunction {
60+ name: string;
61+ /** Declared arguments, in order, with the cName each lowered to. */
62+ params: { name: string; cName: string; binding: Binding }[];
63+ /** Requested outputs, in order, with the cName holding each result. */
64+ outputs: { name: string; cName: string; ty: Type }[];
65+ /** The lowered body. Read this only after `finish()`: the inline pass
66+ * REPLACES the statement array rather than mutating it, so this is a live
67+ * view of the function rather than a snapshot. */
68+ readonly body: IRStmt[];
69+}
70+
71+/** The shape of a `function` statement in numbl's AST. */
72+interface FunctionDecl {
73+ type: 'Function';
74+ name: string;
75+ params: string[];
76+ outputs: string[];
77+}
78+
79+/**
80+ * A parsed model. Specialize the functions you need, then call `finish()` once
81+ * — the inline pass rewrites every specialization together.
82+ */
83+export class CompiledModel {
84+ #lowerer: Lowerer;
85+ #decls: Map<string, FunctionDecl>;
86+ #bindings: Record<string, Binding>;
87+
88+ constructor(
89+ source: string,
90+ bindings: Record<string, Binding>,
91+ grid: GridSizes,
92+ fileName = 'model.m',
93+ ) {
94+ const ast = parseMFile(source, fileName);
95+ const ws = new Workspace(fileName, []);
96+ ws.addFile({ name: fileName, source, ast });
97+ // synth / analys become resolvable, with their type rules.
98+ for (const f of externalOpFiles(grid)) ws.addFile(f);
99+ ws.finalize();
100+
101+ this.#bindings = bindings;
102+ this.#lowerer = new Lowerer(ws);
103+ this.#decls = new Map();
104+ for (const stmt of ast.body as { type: string }[]) {
105+ if (stmt.type === 'Function') {
106+ const fn = stmt as unknown as FunctionDecl;
107+ this.#decls.set(fn.name, fn);
108+ }
109+ }
110+ }
111+
112+ /** Names of the functions the file defines. */
113+ functionNames(): string[] {
114+ return [...this.#decls.keys()];
115+ }
116+
117+ /**
118+ * Lower `name` for the current bindings, requesting `nargout` outputs.
119+ * Every declared parameter must name something the host provides.
120+ */
121+ specialize(name: string, nargout: number): CompiledFunction {
122+ const decl = this.#decls.get(name);
123+ if (!decl) {
124+ const defined = this.functionNames();
125+ throw new ModelCompileError(
126+ `the model must define a function named '${name}'` +
127+ (defined.length
128+ ? ` (it defines ${defined.map((n) => `'${n}'`).join(', ')})`
129+ : ' (it defines no functions)'),
130+ );
131+ }
132+ if (decl.outputs.length < nargout) {
133+ throw new ModelCompileError(
134+ `'${name}' must return ${nargout} value${nargout === 1 ? '' : 's'}, ` +
135+ `but declares ${decl.outputs.length}`,
136+ );
137+ }
138+
139+ const bindings = decl.params.map((p) => {
140+ const b = this.#bindings[p];
141+ if (!b) {
142+ const offered = Object.keys(this.#bindings).join(', ');
143+ throw new ModelCompileError(
144+ `'${name}' takes an argument named '${p}', which this app does not ` +
145+ `provide. Available: ${offered}.`,
146+ );
147+ }
148+ return b;
149+ });
150+
151+ const fn: IRFunc = specializeUserFunction.call(
152+ this.#lowerer,
153+ decl,
154+ bindings.map(typeOf),
155+ undefined,
156+ undefined,
157+ undefined,
158+ nargout,
159+ undefined,
160+ );
161+
162+ return {
163+ name,
164+ params: fn.params.map((p, i) => ({
165+ name: p,
166+ cName: fn.cParams[i],
167+ binding: bindings[i],
168+ })),
169+ outputs: fn.outputs.slice(0, nargout).map((o, i) => ({
170+ name: o,
171+ cName: fn.cOutputs[i],
172+ ty: fn.outputTypes[i],
173+ })),
174+ // A getter, not a snapshot: `finish()` runs after every specialization
175+ // and swaps in a rewritten statement array.
176+ get body() {
177+ return fn.body;
178+ },
179+ };
180+ }
181+
182+ /**
183+ * Run the inline pass over everything specialized so far. It rewrites the
184+ * function bodies in place, so `CompiledFunction`s handed out earlier are
185+ * updated too.
186+ */
187+ finish(): void {
188+ // Snapshot what each loop body assigns, before the pass can rewrite it.
189+ const loops = [...this.#lowerer.specializations.values()].flatMap((fn) =>
190+ forLoops(fn.body).map((loop) => ({
191+ fn,
192+ loop,
193+ assignedBefore: assignedCNames(loop.body),
194+ })),
195+ );
196+
197+ inlinePass({ topLevelStmts: [], functions: this.#lowerer.specializations });
198+
199+ for (const { fn, loop, assignedBefore } of loops) checkLoopEscapes(fn, loop, assignedBefore);
200+ }
201+}
202+
203+/** Every `for` in a statement list, including nested ones. */
204+function forLoops(stmts: IRStmt[]): For[] {
205+ const out: For[] = [];
206+ const walk = (list: IRStmt[]): void => {
207+ for (const s of list) {
208+ if (s.kind === 'For') {
209+ out.push(s);
210+ walk(s.body);
211+ }
212+ }
213+ };
214+ walk(stmts);
215+ return out;
216+}
217+
218+/** cNames assigned anywhere in a statement list, including inside loops. */
219+function assignedCNames(stmts: IRStmt[]): Set<string> {
220+ const out = new Set<string>();
221+ const walk = (list: IRStmt[]): void => {
222+ for (const s of list) {
223+ if (s.kind === 'Assign') out.add(s.cName);
224+ else if (s.kind === 'For') walk(s.body);
225+ }
226+ };
227+ walk(stmts);
228+ return out;
229+}
230+
231+/** Call `visit` for every variable read in an expression. */
232+function forEachVarRead(e: IRExpr, visit: (cName: string) => void): void {
233+ const walk = (x: IRExpr): void => {
234+ switch (x.kind) {
235+ case 'Var':
236+ return visit(x.cName);
237+ case 'Binary':
238+ walk(x.left);
239+ walk(x.right);
240+ return;
241+ case 'Unary':
242+ walk(x.operand);
243+ return;
244+ case 'Call':
245+ x.args.forEach(walk);
246+ return;
247+ default:
248+ return;
249+ }
250+ };
251+ walk(e);
252+}
253+
254+/**
255+ * Refuse a loop whose result the inline pass folded away.
256+ *
257+ * numbl's inline pass substitutes a single-use producer into its consumer and
258+ * drops the producer. Inside a loop body it runs with no protected names — it
259+ * counts uses within that body alone — so an assignment whose only *visible*
260+ * use is later in the same body can be elided even though something outside
261+ * the loop still wants the value.
262+ *
263+ * That elision is correct for a body-local temp, which is what makes fusion
264+ * work inside the loop, and it is caught downstream in the two cases where the
265+ * value has no buffer at all: the planner already refuses a declared output
266+ * that is never assigned, and a read of a name it never allocated. The case it
267+ * would not catch is a variable assigned *before* the loop as well — there the
268+ * buffer exists, holding the pre-loop value, and the loop would silently
269+ * contribute nothing. So check all three here, in one place, against what the
270+ * body assigned before the pass ran.
271+ */
272+function checkLoopEscapes(fn: IRFunc, loop: For, assignedBefore: Set<string>): void {
273+ const assignedAfter = assignedCNames(loop.body);
274+ const elided = [...assignedBefore].filter((c) => !assignedAfter.has(c));
275+ if (elided.length === 0) return;
276+
277+ // Reads anywhere in the function outside this loop's own body.
278+ const readOutside = new Set<string>();
279+ const walk = (list: IRStmt[]): void => {
280+ for (const s of list) {
281+ if (s === (loop as IRStmt)) continue; // the loop's own body is not "outside"
282+ if (s.kind === 'Assign') forEachVarRead(s.expr, (c) => readOutside.add(c));
283+ else if (s.kind === 'For') walk(s.body);
284+ }
285+ };
286+ walk(fn.body);
287+
288+ const outputs = new Set(fn.cOutputs);
289+ const escaping = elided.filter((c) => readOutside.has(c) || outputs.has(c));
290+ if (escaping.length === 0) return;
291+
292+ const names = [...new Set(escaping)].map((c) => `'${c}'`).join(', ');
293+ throw new ModelCompileError(
294+ `inside the 'for' loop, ${names} is assigned but only read later in the ` +
295+ `same iteration, so the compiler folded the assignment into its reader — ` +
296+ `yet the value is also wanted outside the loop. Read it once outside the ` +
297+ `loop instead, or use it more than once inside it.`,
298+ );
299+}
src/mgpu/digest.tsadded+90−0View file
@@ -0,0 +1,90 @@
1+/**
2+ * A run's final state, in a form two different machines can be compared on.
3+ *
4+ * The pipeline is deterministic given (model source, parameters, lmax, seed,
5+ * steps): the perturbation comes from a seeded PRNG, and everything after it is
6+ * fixed arithmetic. So the same spec run anywhere should land on the same state —
7+ * not bit for bit, since GPUs differ in fused-multiply-add and other latitude the
8+ * fp32 rules allow, but far closer than any real difference in what is being
9+ * computed would be.
10+ *
11+ * That makes a cross-environment comparison a genuine check that the browser and
12+ * the desktop are running the same computation, rather than something that merely
13+ * looks similar.
14+ */
15+
16+export interface StateDigest {
17+ /** Element count, so a shape mismatch is caught before the values are read. */
18+ n: number;
19+ min: number;
20+ max: number;
21+ mean: number;
22+ /** Root mean square — sensitive to every element, unlike min/max. */
23+ rms: number;
24+ /** Which Fourier stage the transform plan chose. FFT and DFT are different
25+ * algorithms and round differently, so a mismatch here explains a difference
26+ * in the values rather than being a symptom of one. */
27+ fourier: 'fft' | 'dft';
28+ /** Informational: the GPU the numbers came from. */
29+ adapter: string;
30+}
31+
32+export function digestOf(
33+ values: ArrayLike<number>,
34+ fourier: 'fft' | 'dft',
35+ adapter: string,
36+): StateDigest {
37+ let min = Infinity;
38+ let max = -Infinity;
39+ let sum = 0;
40+ let sumsq = 0;
41+ for (let i = 0; i < values.length; i++) {
42+ const v = values[i];
43+ if (v < min) min = v;
44+ if (v > max) max = v;
45+ sum += v;
46+ sumsq += v * v;
47+ }
48+ const n = values.length;
49+ return {
50+ n,
51+ min,
52+ max,
53+ mean: sum / n,
54+ rms: Math.sqrt(sumsq / n),
55+ fourier,
56+ adapter,
57+ };
58+}
59+
60+/** Relative L2 difference of two states of equal length. */
61+export function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
62+ let num = 0;
63+ let den = 0;
64+ for (let i = 0; i < a.length; i++) {
65+ const d = a[i] - b[i];
66+ num += d * d;
67+ den += b[i] * b[i];
68+ }
69+ return Math.sqrt(num / Math.max(den, 1e-300));
70+}
71+
72+export function formatDigest(d: StateDigest): string {
73+ const g = (v: number): string => v.toPrecision(9);
74+ return (
75+ `n=${d.n} min=${g(d.min)} max=${g(d.max)} mean=${g(d.mean)} rms=${g(d.rms)} ` +
76+ `fourier=${d.fourier}`
77+ );
78+}
79+
80+/** Worst relative disagreement between two digests' scalar summaries. */
81+export function digestDrift(a: StateDigest, b: StateDigest): number {
82+ const rel = (x: number, y: number): number =>
83+ Math.abs(x - y) / Math.max(Math.abs(x), Math.abs(y), 1e-30);
84+ return Math.max(
85+ rel(a.min, b.min),
86+ rel(a.max, b.max),
87+ rel(a.mean, b.mean),
88+ rel(a.rms, b.rms),
89+ );
90+}
src/mgpu/errors.tsadded+105−0View file
@@ -0,0 +1,105 @@
1+/**
2+ * Compile failures, reported in coordinates of the model file the user edits.
3+ *
4+ * Failures arrive from three places, each with its own idea of position:
5+ * numbl's parser (a `position` offset), numbl's lowerer (`UnsupportedConstruct`
6+ * / `JitTypeError`, with a `span`), and this project's WGSL emitter
7+ * (`UnsupportedOnGpu`, carrying the numbl span it was given). All of them are
8+ * offsets into the whole model file — the file is parsed once, and each function
9+ * is specialized from that one AST — so they need only be turned into a line and
10+ * column for the editor.
11+ */
12+
13+/** A compile failure located in the full model source. */
14+export class ModelCompileError extends Error {
15+ /** Offset into the whole .m file, when the failure has a position. */
16+ readonly start?: number;
17+ readonly end?: number;
18+ /** Name of the model function being compiled. */
19+ readonly fn?: string;
20+
21+ constructor(
22+ message: string,
23+ opts: { start?: number; end?: number; fn?: string; cause?: unknown } = {},
24+ ) {
25+ super(message, { cause: opts.cause });
26+ this.name = 'ModelCompileError';
27+ this.start = opts.start;
28+ this.end = opts.end;
29+ this.fn = opts.fn;
30+ }
31+}
32+
33+/** Extract whatever position information an error carries. */
34+function positionOf(e: unknown): { start?: number; end?: number } {
35+ const span = (e as { span?: { start?: unknown; end?: unknown } }).span;
36+ if (span && typeof span.start === 'number') {
37+ return {
38+ start: span.start,
39+ end: typeof span.end === 'number' ? span.end : undefined,
40+ };
41+ }
42+ // numbl's parser SyntaxError reports a bare offset.
43+ const position = (e as { position?: unknown }).position;
44+ if (typeof position === 'number') return { start: position };
45+ return {};
46+}
47+
48+/** Normalize any thrown value into a located `ModelCompileError`. */
49+function asCompileError(e: unknown, fn?: string): ModelCompileError {
50+ if (e instanceof ModelCompileError) return e;
51+ const { start, end } = positionOf(e);
52+ const raw = e instanceof Error ? e.message : String(e);
53+ // numbl's parse errors read as bare token complaints out of context.
54+ const message =
55+ (e as Error)?.name === 'SyntaxError' ? `MATLAB syntax error: ${raw}` : raw;
56+ return new ModelCompileError(message, { fn, start, end, cause: e });
57+}
58+
59+/**
60+ * Run `fn`, locating any compile failure in the model file. Use for whole-file
61+ * phases (parsing) that belong to no single function.
62+ */
63+export function inModel<T>(fn: () => T): T {
64+ try {
65+ return fn();
66+ } catch (e) {
67+ throw asCompileError(e);
68+ }
69+}
70+
71+/** Run `fn`, attributing any compile failure to the model function `name`. */
72+export function inFunction<T>(name: string, fn: () => T): T {
73+ try {
74+ return fn();
75+ } catch (e) {
76+ throw asCompileError(e, name);
77+ }
78+}
79+
80+/** Async form of `inFunction`. */
81+export async function inFunctionAsync<T>(
82+ name: string,
83+ fn: () => Promise<T>,
84+): Promise<T> {
85+ try {
86+ return await fn();
87+ } catch (e) {
88+ throw asCompileError(e, name);
89+ }
90+}
91+
92+/** Render a failure for display: message, section, and 1-based line/column. */
93+export function formatFailure(e: unknown, source: string): string {
94+ const message = e instanceof Error ? e.message : String(e);
95+ if (!(e instanceof ModelCompileError)) return message;
96+ const where: string[] = [];
97+ if (e.start !== undefined && e.start <= source.length) {
98+ const before = source.slice(0, e.start);
99+ const line = before.split('\n').length;
100+ const column = e.start - before.lastIndexOf('\n');
101+ where.push(`line ${line}, column ${column}`);
102+ }
103+ if (e.fn) where.push(`in ${e.fn}()`);
104+ return where.length ? `${message} (${where.join(', ')})` : message;
105+}
src/mgpu/externals.tsadded+95−0View file
@@ -0,0 +1,95 @@
1+/**
2+ * The two spherical-harmonic transforms, as external operations the .m can
3+ * call: `synth` (spectral -> grid) and `analys` (grid -> spectral).
4+ *
5+ * numbl needs only their *type rule* in order to lower a call site. It gets
6+ * that from a `.mtoc2.js` workspace file — numbl's sanctioned extension point
7+ * for a JS-defined builtin (see `mtoc2UserFunctionsByName` in numbl's
8+ * LoweringContext). The file is evaluated in a bare CommonJS sandbox with no
9+ * imports available, so `transfer` builds numbl `Type` objects as plain
10+ * literals, and the grid sizes are baked in by the generator below (a grid
11+ * change recompiles anyway).
12+ *
13+ * The `emit`/`cBody` exports exist only because the loader's contract requires
14+ * them; we never emit C. The actual implementation is supplied by the WGSL
15+ * backend, which turns each of these calls into an ShtPlan encode.
16+ *
17+ * Spectral fields are carried as REAL 2 x nlm arrays (row 0 real part, row 1
18+ * imaginary), matching the interleaved layout the GPU buffers already use.
19+ * The IMEX update is real-linear, so no complex arithmetic is needed.
20+ */
21+
22+export interface GridSizes {
23+ /** Grid points, nlat*nphi. Grid fields are npts x 1 column vectors. */
24+ npts: number;
25+ /** Spectral coefficients. Spectral fields are 2 x nlm. */
26+ nlm: number;
27+}
28+
29+const numericType = (rows: number, cols: number): string =>
30+ `{ kind: "Numeric", elem: "double", isComplex: false, ` +
31+ `dims: [${dim(rows)}, ${dim(cols)}], shape: [${rows}, ${cols}], sign: "unknown" }`;
32+
33+// numbl's tensorDouble() canonicalizes an extent of 1 to its shared DIM_ONE
34+// singleton; mirror that so types compare equal to host-built ones.
35+const dim = (n: number): string =>
36+ n === 1 ? `{ kind: "exact", value: 1 }` : `{ kind: "exact", value: ${n} }`;
37+
38+/** Source for one transform's `.mtoc2.js`. */
39+function transformSource(
40+ name: string,
41+ inRows: number,
42+ inCols: number,
43+ outRows: number,
44+ outCols: number,
45+): string {
46+ return `
47+exports.name = ${JSON.stringify(name)};
48+
49+exports.transfer = function (argTypes, nargout) {
50+ if (argTypes.length !== 1) {
51+ throw new Error("${name} takes exactly one argument, got " + argTypes.length);
52+ }
53+ if (nargout > 1) {
54+ throw new Error("${name} returns one value, but " + nargout + " were requested");
55+ }
56+ var a = argTypes[0];
57+ if (!a || a.kind !== "Numeric" || a.isComplex) {
58+ throw new Error("${name} requires a real numeric array");
59+ }
60+ var s = a.shape;
61+ if (!s || s.length !== 2 || s[0] !== ${inRows} || s[1] !== ${inCols}) {
62+ throw new Error(
63+ "${name} requires a ${inRows}x${inCols} array, got " +
64+ (s ? s.join("x") : "unknown shape")
65+ );
66+ }
67+ return [${numericType(outRows, outCols)}];
68+};
69+
70+// Never called: this project executes the IR on WebGPU and emits no C.
71+exports.emit = function () {
72+ throw new Error("${name}: no C backend (this transform runs on WebGPU)");
73+};
74+exports.cBody = function () {
75+ return "";
76+};
77+`;
78+}
79+
80+/** Workspace files that make `synth` / `analys` resolvable during lowering. */
81+export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
82+ return [
83+ {
84+ name: 'synth.mtoc2.js',
85+ source: transformSource('synth', 2, g.nlm, g.npts, 1),
86+ },
87+ {
88+ name: 'analys.mtoc2.js',
89+ source: transformSource('analys', g.npts, 1, 2, g.nlm),
90+ },
91+ ];
92+}
93+
94+/** Names the WGSL backend must implement as GPU encodes rather than kernels. */
95+export const EXTERNAL_OPS = new Set(['synth', 'analys']);
src/mgpu/model.tsadded+353−0View file
@@ -0,0 +1,353 @@
1+/**
2+ * A .m model, compiled and running on the GPU.
3+ *
4+ * A model file is ordinary MATLAB: it defines an `init` function that builds the
5+ * initial spectral state and a `step` function that advances it one timestep.
6+ * Each is specialized for the current grid and compiled into a ModelPlan, and
7+ * both operate on the same state buffers (see HostBuffers).
8+ *
9+ * Both functions return the new state followed by the grid fields the app
10+ * renders, so their signatures say exactly what they produce:
11+ *
12+ * function [U, V, u, v] = init(noise, a, b)
13+ * function [U, V, u, v] = step(U, V, lam, a, b, D1, D2, dt)
14+ *
15+ * The host supplies the things that are precomputation rather than algorithm:
16+ * the grid, the Laplace-Beltrami eigenvalues, the seeded initial noise, and the
17+ * parameter values. Each argument is matched to the .m's declared parameter
18+ * name, so the file documents its own interface.
19+ */
20+import { ShtPlan } from '../sht/sht.ts';
21+import { lmIndex, type ShtConfig } from '../sht/layout.ts';
22+import { HostBuffers, ModelPlan } from './plan.ts';
23+import { inFunction, inFunctionAsync, inModel } from './errors.ts';
24+import { CompiledModel, type Binding } from './compile.ts';
25+
26+export interface ModelParams {
27+ [key: string]: number;
28+}
29+
30+export interface GpuModelOptions {
31+ device: GPUDevice;
32+ sht: ShtPlan;
33+ cfg: ShtConfig;
34+ /** Model source (.m text). */
35+ source: string;
36+ /** Parameter names the .m may take as arguments. */
37+ paramNames: string[];
38+ /** Spectral state names, in order (e.g. ['U', 'V']). */
39+ state: string[];
40+ /** Grid fields to render, in order (e.g. ['u', 'v']). */
41+ view: string[];
42+ /**
43+ * The surface, as the .m may ask for it: `gx`, `gy`, `gz` are the embedding's
44+ * Cartesian coordinates on the grid, and `Gx`, `Gy`, `Gz` the spherical-
45+ * harmonic coefficients they were synthesized from. Omitted for a bare unit
46+ * sphere, where the .m has no geometry to take.
47+ */
48+ geometry?: GeometryBuffers;
49+ /**
50+ * Iterations of the implicit solve the .m's `for` loop runs. A fixed scalar
51+ * rather than a tunable one: the loop is unrolled into the op sequence, so
52+ * the count is part of what compiles and changing it recompiles.
53+ */
54+ niter?: number;
55+}
56+
57+/** Host-supplied surface fields, in the layout the .m sees them. */
58+export interface GeometryBuffers {
59+ /** Grid coordinates, npts each. */
60+ x: Float32Array;
61+ y: Float32Array;
62+ z: Float32Array;
63+ /** Their spherical-harmonic coefficients, 2 x nlm each. */
64+ X: Float32Array;
65+ Y: Float32Array;
66+ Z: Float32Array;
67+}
68+
69+/** Names the .m may take for the grid coordinates and for their coefficients. */
70+export const GEOMETRY_GRID_NAMES = ['gx', 'gy', 'gz'] as const;
71+export const GEOMETRY_SPECTRAL_NAMES = ['Gx', 'Gy', 'Gz'] as const;
72+
73+/** Laplace-Beltrami eigenvalues l(l+1), duplicated across re/im so the array
74+ * matches the 2 x nlm spectral layout element for element. */
75+export function eigenvalues(cfg: ShtConfig, nlm: number): Float32Array {
76+ const lam = new Float32Array(2 * nlm);
77+ for (let m = 0; m <= cfg.mmax; m++) {
78+ for (let l = m; l <= cfg.lmax; l++) {
79+ const i = lmIndex(cfg.lmax, l, m);
80+ lam[2 * i] = l * (l + 1);
81+ lam[2 * i + 1] = l * (l + 1);
82+ }
83+ }
84+ return lam;
85+}
86+
87+export class GpuModel {
88+ readonly paramNames: string[];
89+ readonly state: string[];
90+ readonly view: string[];
91+ readonly npts: number;
92+ readonly nlm: number;
93+
94+ #device: GPUDevice;
95+ #host: HostBuffers;
96+ #initPlan: ModelPlan;
97+ #stepPlan: ModelPlan;
98+ #readback: GPUBuffer;
99+ /** Scratch holding a copy of the whole spectral state; see snapshotState. */
100+ #stash: GPUBuffer;
101+ /** Which function wrote the state most recently; see `read`. */
102+ #lastRan: 'init' | 'step' = 'init';
103+ #stashedRan: 'init' | 'step' = 'init';
104+ #destroyed = false;
105+
106+ private constructor(init: {
107+ device: GPUDevice;
108+ host: HostBuffers;
109+ initPlan: ModelPlan;
110+ stepPlan: ModelPlan;
111+ readback: GPUBuffer;
112+ stash: GPUBuffer;
113+ paramNames: string[];
114+ state: string[];
115+ view: string[];
116+ npts: number;
117+ nlm: number;
118+ }) {
119+ this.#device = init.device;
120+ this.#host = init.host;
121+ this.#initPlan = init.initPlan;
122+ this.#stepPlan = init.stepPlan;
123+ this.#readback = init.readback;
124+ this.#stash = init.stash;
125+ this.paramNames = init.paramNames;
126+ this.state = init.state;
127+ this.view = init.view;
128+ this.npts = init.npts;
129+ this.nlm = init.nlm;
130+ }
131+
132+ static async create(opts: GpuModelOptions): Promise<GpuModel> {
133+ const { device, sht, cfg, source, paramNames, state, view, geometry } = opts;
134+ const npts = cfg.nlat * cfg.nphi;
135+ const nlm = sht.nlm;
136+ const niter = opts.niter ?? 0;
137+
138+ // What the .m may ask for by parameter name. Spectral state and the
139+ // eigenvalues are 2 x nlm; the seeded perturbation is a grid field.
140+ const bindings: Record<string, Binding> = {
141+ lam: { kind: 'tensor', shape: [2, nlm] },
142+ noise: { kind: 'tensor', shape: [npts, 1] },
143+ npts: { kind: 'const', value: npts },
144+ nlm: { kind: 'const', value: nlm },
145+ niter: { kind: 'const', value: niter },
146+ };
147+ if (geometry) {
148+ for (const g of GEOMETRY_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
149+ for (const g of GEOMETRY_SPECTRAL_NAMES) bindings[g] = { kind: 'tensor', shape: [2, nlm] };
150+ }
151+ for (const s of state) bindings[s] = { kind: 'tensor', shape: [2, nlm] };
152+ for (const p of paramNames) bindings[p] = { kind: 'param' };
153+
154+ // Parsing belongs to the file, not to either function.
155+ const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
156+ // Both functions return the new state first, then the rendered grid fields.
157+ const nargout = state.length + view.length;
158+ const initFn = inFunction('init', () => compiled.specialize('init', nargout));
159+ const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
160+ compiled.finish();
161+
162+ // Only the state outputs feed back into the argument buffers; the grid
163+ // fields are read for display and then overwritten next call.
164+ const feedback = [...state, ...view.map(() => null)];
165+
166+ const host = new HostBuffers(device);
167+ // The host owns the state and the inputs it uploads, whether or not a given
168+ // function happens to take them as arguments — `init` does not read `U`, but
169+ // it writes it, and `step` reads it back.
170+ for (const s of state) host.ensure(s, 2 * nlm);
171+ host.ensure('lam', 2 * nlm);
172+ host.ensure('noise', npts);
173+ if (geometry) {
174+ for (const g of GEOMETRY_GRID_NAMES) host.ensure(g, npts);
175+ for (const g of GEOMETRY_SPECTRAL_NAMES) host.ensure(g, 2 * nlm);
176+ }
177+
178+ const initPlan = await inFunctionAsync('init', () =>
179+ ModelPlan.create(device, sht, { fn: initFn, feedback }, host),
180+ );
181+ const stepPlan = await inFunctionAsync('step', () =>
182+ ModelPlan.create(device, sht, { fn: stepFn, feedback }, host),
183+ );
184+
185+ host.upload('lam', eigenvalues(cfg, nlm));
186+ if (geometry) {
187+ host.upload('gx', geometry.x);
188+ host.upload('gy', geometry.y);
189+ host.upload('gz', geometry.z);
190+ host.upload('Gx', geometry.X);
191+ host.upload('Gy', geometry.Y);
192+ host.upload('Gz', geometry.Z);
193+ }
194+
195+ const readback = device.createBuffer({
196+ label: 'mgpu-readback',
197+ size: 4 * Math.max(npts, 2 * nlm),
198+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
199+ });
200+ const stash = device.createBuffer({
201+ label: 'mgpu-state-stash',
202+ size: 4 * state.length * 2 * nlm,
203+ usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
204+ });
205+
206+ return new GpuModel({
207+ device, host, initPlan, stepPlan, readback, stash,
208+ paramNames, state, view, npts, nlm,
209+ });
210+ }
211+
212+ setParams(params: ModelParams): void {
213+ this.#initPlan.setParams(params);
214+ this.#stepPlan.setParams(params);
215+ }
216+
217+ /**
218+ * Write a host-owned value directly — the spectral state, or one of the input
219+ * fields. Lets a test set up an exact initial condition (a single spherical-
220+ * harmonic mode, say) instead of going through `init`.
221+ */
222+ upload(name: string, data: Float32Array): void {
223+ this.#host.upload(name, data);
224+ }
225+
226+ /**
227+ * Swap the surface under a running model. The geometry is data, not code —
228+ * its shape in the bindings depends only on the grid — so changing it is six
229+ * buffer writes and needs no recompile, and the simulation carries straight
230+ * on. Only meaningful if the .m took the geometry as an argument.
231+ */
232+ uploadGeometry(geometry: GeometryBuffers): void {
233+ const fields: [string, Float32Array][] = [
234+ ['gx', geometry.x], ['gy', geometry.y], ['gz', geometry.z],
235+ ['Gx', geometry.X], ['Gy', geometry.Y], ['Gz', geometry.Z],
236+ ];
237+ for (const [name, data] of fields) {
238+ if (this.#host.get(name)) this.#host.upload(name, data);
239+ }
240+ }
241+
242+ /** Upload the seeded perturbation and run `init`. */
243+ init(noise: Float32Array): void {
244+ this.#host.upload('noise', noise);
245+ const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
246+ this.#initPlan.encodeSteps(enc, 1);
247+ this.#device.queue.submit([enc.finish()]);
248+ this.#lastRan = 'init';
249+ }
250+
251+ /**
252+ * Copy the spectral state aside, so a batch of steps can run — to be timed —
253+ * and then be undone with restoreState, leaving the simulation exactly where
254+ * it was. Only the state is stashed: the grid view fields keep whatever the
255+ * batch last wrote until a subsequent step recomputes them, so step before
256+ * reading a view after a restore.
257+ */
258+ snapshotState(): void {
259+ this.#stashedRan = this.#lastRan;
260+ this.#copyState('save');
261+ }
262+
263+ restoreState(): void {
264+ this.#copyState('restore');
265+ this.#lastRan = this.#stashedRan;
266+ }
267+
268+ #copyState(dir: 'save' | 'restore'): void {
269+ // A restore can land after a rebuild destroyed the buffers mid-await;
270+ // there is nothing left to protect, so do not submit into destroyed state.
271+ if (this.#destroyed) return;
272+ const enc = this.#device.createCommandEncoder({ label: `mgpu-state-${dir}` });
273+ let offset = 0;
274+ for (const name of this.state) {
275+ const slot = this.#host.get(name);
276+ if (!slot) throw new Error(`state '${name}' has no host buffer`);
277+ const bytes = 4 * slot.count;
278+ if (dir === 'save') {
279+ enc.copyBufferToBuffer(slot.buffer, 0, this.#stash, offset, bytes);
280+ } else {
281+ enc.copyBufferToBuffer(this.#stash, offset, slot.buffer, 0, bytes);
282+ }
283+ offset += bytes;
284+ }
285+ this.#device.queue.submit([enc.finish()]);
286+ }
287+
288+ /**
289+ * Advance `steps` timesteps. Synchronous — this only records commands and
290+ * submits them; nothing is read back and nothing is awaited.
291+ */
292+ step(steps = 1): void {
293+ const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
294+ this.#stepPlan.encodeSteps(enc, steps);
295+ this.#device.queue.submit([enc.finish()]);
296+ this.#lastRan = 'step';
297+ }
298+
299+ /**
300+ * The buffer currently holding a named value. Grid fields like `u` are
301+ * produced by both functions, into separate buffers (only the spectral state
302+ * is shared), so this resolves to whichever function ran most recently —
303+ * which is what makes the first frame show the initial state rather than an
304+ * unwritten buffer.
305+ */
306+ #locate(name: string): { buffer: GPUBuffer; count: number } | null {
307+ const [first, second] =
308+ this.#lastRan === 'init'
309+ ? [this.#initPlan, this.#stepPlan]
310+ : [this.#stepPlan, this.#initPlan];
311+ const buffer = first.buffer(name) ?? second.buffer(name);
312+ const count = first.elementCount(name) ?? second.elementCount(name);
313+ if (!buffer || count === undefined) return null;
314+ return { buffer, count };
315+ }
316+
317+ /** The GPU buffer a named value would be read from right now — for encoding
318+ * further GPU work against it (e.g. a display-grid synthesis of the state)
319+ * without a CPU round trip. */
320+ valueBuffer(name: string): GPUBuffer | null {
321+ return this.#locate(name)?.buffer ?? null;
322+ }
323+
324+ /** Read a named value back to the CPU. The only await in the whole loop. */
325+ async read(name: string): Promise<Float32Array> {
326+ const located = this.#locate(name);
327+ if (!located) {
328+ throw new Error(`read: the model has no value named '${name}'`);
329+ }
330+ const { buffer, count } = located;
331+ const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
332+ enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
333+ this.#device.queue.submit([enc.finish()]);
334+ await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
335+ const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
336+ this.#readback.unmap();
337+ return out;
338+ }
339+
340+ /** What the .m compiled to, for display. */
341+ describe(): { init: string[]; step: string[] } {
342+ return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
343+ }
344+
345+ destroy(): void {
346+ this.#destroyed = true;
347+ this.#initPlan.destroy();
348+ this.#stepPlan.destroy();
349+ this.#host.destroy();
350+ this.#readback.destroy();
351+ this.#stash.destroy();
352+ }
353+}
src/mgpu/noise.tsadded+51−0View file
@@ -0,0 +1,51 @@
1+/**
2+ * The seeded perturbation a model's `init` starts from.
3+ *
4+ * Host-side rather than in the .m, so a run is reproducible from an integer
5+ * seed and the same field can be handed to any model.
6+ */
7+
8+/**
9+ * Seeded uniform deviates in [0, 1): mulberry32.
10+ *
11+ * Integer arithmetic and one division by 2^32, so any faithful port of it
12+ * produces bit-identical values — which is what lets the native benchmark under
13+ * bench/shtns/ seed the same run.
14+ */
15+export function makeRand(seed: number): () => number {
16+ let s = seed >>> 0;
17+ return (): number => {
18+ s = (s + 0x6d2b79f5) >>> 0;
19+ let t = s;
20+ t = Math.imul(t ^ (t >>> 15), t | 1);
21+ t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
22+ return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
23+ };
24+}
25+
26+/** Seeded normal deviates: mulberry32 + Box-Muller. */
27+export function makeRandn(seed: number): () => number {
28+ const rand = makeRand(seed);
29+ let spare: number | null = null;
30+ return () => {
31+ if (spare !== null) {
32+ const v = spare;
33+ spare = null;
34+ return v;
35+ }
36+ let u = 0;
37+ while (u === 0) u = rand();
38+ const r = Math.sqrt(-2 * Math.log(u));
39+ const th = 2 * Math.PI * rand();
40+ spare = r * Math.sin(th);
41+ return r * Math.cos(th);
42+ };
43+}
44+
45+/** `amp`-scaled normal deviates, one per grid point, in index order. */
46+export function seededNoise(npts: number, amp: number, seed: number): Float32Array {
47+ const randn = makeRandn(seed);
48+ const out = new Float32Array(npts);
49+ for (let i = 0; i < npts; i++) out[i] = amp * randn();
50+ return out;
51+}
src/mgpu/numbl.d.tsadded+261−0View file
@@ -0,0 +1,261 @@
1+/**
2+ * The numbl compiler surface this project depends on.
3+ *
4+ * We reach past numbl's published entry points into its JIT internals (parser,
5+ * lowerer, IR, inline pass), which its package `exports` map does not expose.
6+ * Those imports resolve through the `numbl-src` alias in vite.config.ts; these
7+ * declarations are what TypeScript checks against.
8+ *
9+ * Declaring the surface here rather than type-checking numbl's sources
10+ * directly keeps this project's compiler settings independent of numbl's, and
11+ * pins the exact contract we rely on. If numbl changes one of these shapes,
12+ * the build breaks here with a clear diff rather than deep inside its tree.
13+ *
14+ * Only the nodes the WGSL backend actually walks are spelled out; every other
15+ * IR kind is collapsed into a catch-all so that unhandled constructs are
16+ * rejected with a message instead of being silently mis-compiled.
17+ */
18+
19+declare module 'numbl-src/numbl-core/jit/lowering/types.ts' {
20+ export type Sign =
21+ | 'positive' | 'nonneg' | 'negative' | 'nonpositive'
22+ | 'zero' | 'nonzero' | 'unknown';
23+
24+ export type DimInfo = { kind: 'exact'; value: number } | { kind: 'unknown' };
25+
26+ export type NumericExact =
27+ | number
28+ | Float64Array
29+ | { re: number; im: number }
30+ | { re: Float64Array; im: Float64Array };
31+
32+ export interface NumericType {
33+ kind: 'Numeric';
34+ elem: 'double' | 'logical' | 'char' | string;
35+ isComplex: boolean;
36+ dims: DimInfo[];
37+ /** Present iff every dim is exact. */
38+ shape?: number[];
39+ sign: Sign;
40+ exact?: NumericExact;
41+ }
42+
43+ /** Everything the WGSL backend rejects. */
44+ export interface NonNumericType {
45+ kind: 'Void' | 'Unknown' | 'String' | 'Handle' | 'Struct' | 'Class' | 'Cell';
46+ }
47+
48+ export type Type = NumericType | NonNumericType;
49+
50+ export function isMultiElement(t: NumericType): boolean;
51+ export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
52+ export function scalarDouble(sign?: Sign, exact?: number): NumericType;
53+}
54+
55+declare module 'numbl-src/numbl-core/jit/lowering/ir.ts' {
56+ import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
57+
58+ export interface Span {
59+ file: string;
60+ start: number;
61+ end: number;
62+ }
63+
64+ export interface NumLit {
65+ kind: 'NumLit';
66+ value: number;
67+ ty: Type;
68+ span: Span;
69+ }
70+ export interface Var {
71+ kind: 'Var';
72+ name: string;
73+ cName: string;
74+ ty: Type;
75+ span: Span;
76+ }
77+ export interface Binary {
78+ kind: 'Binary';
79+ builtin: string;
80+ left: IRExpr;
81+ right: IRExpr;
82+ ty: Type;
83+ span: Span;
84+ }
85+ export interface Unary {
86+ kind: 'Unary';
87+ builtin: string;
88+ operand: IRExpr;
89+ ty: Type;
90+ span: Span;
91+ }
92+ export interface Call {
93+ kind: 'Call';
94+ cName: string;
95+ name: string;
96+ args: IRExpr[];
97+ ty: Type;
98+ span: Span;
99+ }
100+ /** Any other IR expression kind — rejected by the WGSL emitter. */
101+ export interface OtherExpr {
102+ kind:
103+ | 'ImagLit' | 'StringLit' | 'TensorBuild' | 'TensorConcat' | 'CellLit'
104+ | 'CellEmpty' | 'CellIndexLoad' | 'HandleLit' | 'HandleCaptureLoad'
105+ | 'StructLit' | 'MemberLoad' | 'IndexLoad' | 'IndexSlice' | 'EndRef'
106+ | 'MakeRange';
107+ ty: Type;
108+ span: Span;
109+ }
110+
111+ export type IRExpr = NumLit | Var | Binary | Unary | Call | OtherExpr;
112+
113+ export interface Assign {
114+ kind: 'Assign';
115+ name: string;
116+ cName: string;
117+ ty: Type;
118+ expr: IRExpr;
119+ span: Span;
120+ }
121+ /**
122+ * A counted loop. The planner unrolls it, so only the fields that decide
123+ * the trip count and the loop variable's value are spelled out. `step` is
124+ * already a literal number in the IR — numbl rejects a non-literal step
125+ * during lowering — while `start` and `end` are expressions that must carry
126+ * an exact value for the planner to accept the loop.
127+ */
128+ export interface For {
129+ kind: 'For';
130+ /** Loop variable, as written in the .m. */
131+ varName: string;
132+ /** Loop variable's cName, the key the planner binds its value under. */
133+ cVar: string;
134+ start: IRExpr;
135+ step: number;
136+ end: IRExpr;
137+ body: IRStmt[];
138+ span: Span;
139+ }
140+ /** Any other IR statement kind — rejected by the planner. */
141+ export interface OtherStmt {
142+ kind:
143+ | 'ExprStmt' | 'If' | 'While' | 'ReturnFromFunction' | 'Break'
144+ | 'Continue' | 'TypeComment' | 'MemberStore' | 'MultiAssignCall'
145+ | 'IndexStore' | 'IndexSliceStore' | 'CellIndexStore';
146+ span: Span;
147+ }
148+
149+ export type IRStmt = Assign | For | OtherStmt;
150+
151+ export interface IRFunc {
152+ name: string;
153+ cName: string;
154+ /** Parameter source names. */
155+ params: string[];
156+ /** Parameter cNames, parallel to `params`. */
157+ cParams: string[];
158+ paramTypes: Type[];
159+ /** Output source names. */
160+ outputs: string[];
161+ /** Output cNames, parallel to `outputs`. */
162+ cOutputs: string[];
163+ outputTypes: Type[];
164+ body: IRStmt[];
165+ span: Span;
166+ }
167+
168+ export interface IRProgram {
169+ topLevelStmts: IRStmt[];
170+ functions: Map<string, IRFunc>;
171+ }
172+}
173+
174+declare module 'numbl-src/numbl-core/parser/index.ts' {
175+ export interface AbstractSyntaxTree {
176+ body: unknown[];
177+ }
178+ export function parseMFile(input: string, fileName?: string): AbstractSyntaxTree;
179+ export class SyntaxError extends Error {}
180+}
181+
182+declare module 'numbl-src/numbl-core/jit/index.ts' {
183+ import type { AbstractSyntaxTree } from 'numbl-src/numbl-core/parser/index.ts';
184+ import type { IRProgram, IRFunc, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
185+ import type { Type, NumericType, Sign } from 'numbl-src/numbl-core/jit/lowering/types.ts';
186+
187+ export interface WorkspaceFile {
188+ name: string;
189+ source: string;
190+ ast?: AbstractSyntaxTree;
191+ }
192+
193+ export class Workspace {
194+ constructor(mainFile: string, searchPaths?: ReadonlyArray<string>);
195+ addFile(file: WorkspaceFile): void;
196+ finalize(): void;
197+ }
198+
199+ export interface EnvEntry {
200+ cName: string;
201+ ty: Type;
202+ maybeUnassigned?: boolean;
203+ }
204+
205+ export class Lowerer {
206+ constructor(workspace: Workspace);
207+ /** Pre-bindable variable scope: seed host-provided values here. */
208+ env: Map<string, EnvEntry>;
209+ specializations: Map<string, IRFunc>;
210+ lowerProgram(ast: AbstractSyntaxTree): IRProgram;
211+ }
212+
213+ /** Thrown for MATLAB the JIT pipeline cannot lower; carries a source span. */
214+ export class UnsupportedConstruct extends Error {
215+ span?: Span;
216+ }
217+ export class JitTypeError extends Error {
218+ span?: Span;
219+ }
220+
221+ export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
222+ export function scalarDouble(sign?: Sign, exact?: number): NumericType;
223+ export function isMultiElement(t: NumericType): boolean;
224+}
225+
226+declare module 'numbl-src/numbl-core/jit/lowering/specialize.ts' {
227+ import type { Lowerer } from 'numbl-src/numbl-core/jit/index.ts';
228+ import type { IRFunc, IRExpr, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
229+ import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
230+
231+ /**
232+ * Lower one user function for a concrete argument-type signature. Called with
233+ * a `Lowerer` as `this` (numbl's own JIT does the same), so specializations
234+ * accumulate in `lowerer.specializations`.
235+ */
236+ export function specializeUserFunction(
237+ this: Lowerer,
238+ decl: unknown,
239+ argTypes: Type[],
240+ specSource?: string,
241+ definingFile?: string,
242+ preSeedOutput?: { name: string; ty: Type; initExpr: IRExpr },
243+ nargout?: number,
244+ callSiteSpan?: Span,
245+ ): IRFunc;
246+}
247+
248+declare module 'numbl-src/numbl-core/jit/codegen/inlinePass.ts' {
249+ import type { IRProgram } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
250+ /** Folds single-use ANF temps into their consumer, in place. */
251+ export function inlinePass(prog: IRProgram): void;
252+}
253+
254+declare module 'numbl-src/numbl-core/jit/builtins/index.ts' {
255+ export interface Builtin {
256+ name: string;
257+ /** Safe to evaluate one output element from one input element per slot. */
258+ elementwise?: boolean;
259+ }
260+ export function getBuiltin(name: string): Builtin | undefined;
261+}
src/mgpu/plan.tsadded+609−0View file
@@ -0,0 +1,609 @@
1+/**
2+ * Statement list -> a replayable sequence of GPU operations.
3+ *
4+ * Everything expensive happens once, here: pipeline compilation, buffer
5+ * allocation, bind-group construction. Because numbl fixes every type and
6+ * shape at lowering time, the resulting op sequence is fully static — so
7+ * `encodeStep` is pure synchronous command recording, with no allocation, no
8+ * pipeline lookup and no readback. That is what lets the whole timestep be
9+ * encoded into one submit and keeps the CPU out of the loop.
10+ */
11+import { isMultiElement, scalarDouble } from 'numbl-src/numbl-core/jit/lowering/types.ts';
12+import type { Assign, For, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
13+import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
14+import { ShtPlan, type ShtBinding } from '../sht/sht.ts';
15+import type { CompiledFunction } from './compile.ts';
16+import { EXTERNAL_OPS } from './externals.ts';
17+import {
18+ buildKernel,
19+ UnsupportedOnGpu,
20+ WORKGROUP_SIZE,
21+ type KernelInputs,
22+} from './wgsl.ts';
23+
24+const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
25+const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
26+const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);
27+
28+/**
29+ * The compile-time value of a scalar expression, if it has one. A literal
30+ * carries its own; a variable carries one when it was bound to a `const` (the
31+ * host's fixed scalars) or computed from constants, because numbl propagates
32+ * `exact` through the type lattice.
33+ */
34+const exactValue = (e: IRExpr): number | undefined => {
35+ if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
36+ return e.kind === 'NumLit' ? e.value : undefined;
37+};
38+
39+/** Cap on the iterations a `for` may unroll to. Each one is real GPU work —
40+ * its own pipelines at compile time and its own dispatches per step — so a
41+ * runaway bound should be a clear error rather than a hang. */
42+const MAX_UNROLL = 64;
43+
44+interface Slot {
45+ buffer: GPUBuffer;
46+ count: number;
47+}
48+
49+const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
50+ device.createBuffer({
51+ label,
52+ size: 4 * count,
53+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
54+ });
55+
56+/**
57+ * Buffers for host-bound variables, shared across plans.
58+ *
59+ * A model is two programs — `init` and `step` — compiled separately but
60+ * operating on the same state. `U` in the step must be the very buffer `init`
61+ * wrote, so the buffers for host bindings live here rather than inside either
62+ * plan.
63+ */
64+export class HostBuffers {
65+ #device: GPUDevice;
66+ #slots = new Map<string, Slot>();
67+
68+ constructor(device: GPUDevice) {
69+ this.#device = device;
70+ }
71+
72+ ensure(name: string, count: number): Slot {
73+ const existing = this.#slots.get(name);
74+ if (existing) {
75+ if (existing.count !== count) {
76+ throw new UnsupportedOnGpu(
77+ `'${name}' is ${existing.count} elements in one program and ` +
78+ `${count} in another`,
79+ );
80+ }
81+ return existing;
82+ }
83+ const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
84+ this.#slots.set(name, slot);
85+ return slot;
86+ }
87+
88+ get(name: string): Slot | undefined {
89+ return this.#slots.get(name);
90+ }
91+
92+ /** Upload initial data for a host binding. */
93+ upload(name: string, data: Float32Array): void {
94+ const slot = this.#slots.get(name);
95+ if (!slot) throw new Error(`upload: no buffer named '${name}'`);
96+ if (data.length !== slot.count) {
97+ throw new Error(
98+ `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
99+ );
100+ }
101+ this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
102+ }
103+
104+ destroy(): void {
105+ for (const s of this.#slots.values()) s.buffer.destroy();
106+ this.#slots.clear();
107+ }
108+}
109+
110+type Op =
111+ | {
112+ kind: 'kernel';
113+ pipeline: GPUComputePipeline;
114+ bindGroup: GPUBindGroup;
115+ count: number;
116+ label: string;
117+ /** Set when the kernel had to write to scratch because its output
118+ * aliases one of its inputs; copied back after the dispatch. */
119+ copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
120+ }
121+ | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
122+ | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
123+
124+export interface PlanSpec {
125+ /** The specialized function this plan executes. */
126+ fn: CompiledFunction;
127+ /** Output index -> host binding name to copy the result into after the run,
128+ * so the next call reads it (the new spectral state feeds the old). */
129+ feedback: (string | null)[];
130+}
131+
132+/**
133+ * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
134+ * buffers after it, then the params buffer.
135+ *
136+ * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
137+ * only contains the bindings the shader actually references — so a kernel that
138+ * happens to use no parameters (`uuv = u .* u .* v`) would drop the params
139+ * binding and no longer match the bind group. An explicit layout may carry
140+ * bindings the shader ignores.
141+ */
142+function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
143+ const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
144+ binding,
145+ visibility: GPUShaderStage.COMPUTE,
146+ buffer: { type: 'read-only-storage' },
147+ });
148+ return device.createBindGroupLayout({
149+ entries: [
150+ {
151+ binding: 0,
152+ visibility: GPUShaderStage.COMPUTE,
153+ buffer: { type: 'storage' },
154+ },
155+ ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
156+ readOnly(inputs + 1),
157+ ],
158+ });
159+}
160+
161+async function makePipeline(
162+ device: GPUDevice,
163+ code: string,
164+ label: string,
165+ bindGroupLayout: GPUBindGroupLayout,
166+): Promise<GPUComputePipeline> {
167+ device.pushErrorScope('validation');
168+ const module = device.createShaderModule({ code, label });
169+ const info = await module.getCompilationInfo();
170+ const errors = info.messages.filter((m) => m.type === 'error');
171+ if (errors.length) {
172+ throw new UnsupportedOnGpu(
173+ `generated WGSL failed to compile for '${label}':\n` +
174+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
175+ `\n--- shader ---\n${code}`,
176+ );
177+ }
178+ const pipeline = await device.createComputePipelineAsync({
179+ layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
180+ compute: { module, entryPoint: 'main' },
181+ label,
182+ });
183+ const err = await device.popErrorScope();
184+ if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
185+ return pipeline;
186+}
187+
188+/** A compiled .m step, ready to run on the GPU. */
189+export class ModelPlan {
190+ /** Scalar parameter names, in the order the params buffer expects them. */
191+ readonly paramNames: string[];
192+
193+ #device: GPUDevice;
194+ #sht: ShtPlan;
195+ #ops: Op[];
196+ #owned: GPUBuffer[];
197+ #paramBuf: GPUBuffer;
198+ #paramData: Float32Array;
199+ /** Public name -> buffer, for uploading initial state and reading results. */
200+ #byName: Map<string, Slot>;
201+
202+ private constructor(init: {
203+ device: GPUDevice;
204+ sht: ShtPlan;
205+ ops: Op[];
206+ byName: Map<string, Slot>;
207+ owned: GPUBuffer[];
208+ paramBuf: GPUBuffer;
209+ paramData: Float32Array;
210+ paramNames: string[];
211+ }) {
212+ this.#device = init.device;
213+ this.#sht = init.sht;
214+ this.#ops = init.ops;
215+ this.#byName = init.byName;
216+ this.#owned = init.owned;
217+ this.#paramBuf = init.paramBuf;
218+ this.#paramData = init.paramData;
219+ this.paramNames = init.paramNames;
220+ }
221+
222+ static async create(
223+ device: GPUDevice,
224+ sht: ShtPlan,
225+ spec: PlanSpec,
226+ host: HostBuffers,
227+ ): Promise<ModelPlan> {
228+ const { fn } = spec;
229+
230+ const slots = new Map<string, Slot>();
231+ const byName = new Map<string, Slot>();
232+ const owned: GPUBuffer[] = [];
233+ /** Scalars the .m computes from its parameters, by cName. */
234+ const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
235+
236+ const alloc = (label: string, count: number): Slot => {
237+ const buffer = makeBuffer(device, label, count);
238+ owned.push(buffer);
239+ return { buffer, count };
240+ };
241+
242+ // Arguments, bound by what the function's signature declares. Array
243+ // arguments come from the shared pool, so a value one function returns is
244+ // the same buffer the next one reads. Scalar parameters share one small
245+ // storage buffer, in signature order.
246+ const paramNames: string[] = [];
247+ const paramSlots = new Map<string, number>();
248+ for (const p of fn.params) {
249+ if (p.binding.kind === 'tensor') {
250+ const count = p.binding.shape.reduce((x, y) => x * y, 1);
251+ const slot = host.ensure(p.name, count);
252+ slots.set(p.cName, slot);
253+ byName.set(p.name, slot);
254+ } else if (p.binding.kind === 'param') {
255+ paramSlots.set(p.cName, paramNames.length);
256+ paramNames.push(p.name);
257+ }
258+ // `const` arguments are exact in the IR and fold into the kernels.
259+ }
260+ const paramData = new Float32Array(Math.max(1, paramNames.length));
261+ const paramBuf = device.createBuffer({
262+ label: 'mgpu-params',
263+ size: 4 * paramData.length,
264+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
265+ });
266+
267+ const ops: Op[] = [];
268+ for (const stmt of fn.body) {
269+ await planStatement(stmt);
270+ }
271+
272+ // Feed declared outputs back into the argument buffers they replace.
273+ fn.outputs.forEach((out, i) => {
274+ const to = spec.feedback[i];
275+ if (!to) return;
276+ const src = slots.get(out.cName);
277+ const dst = host.get(to);
278+ if (!src) {
279+ throw new UnsupportedOnGpu(
280+ `'${fn.name}' declares the output '${out.name}' but never assigns it`,
281+ );
282+ }
283+ if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
284+ if (src.count !== dst.count) {
285+ throw new UnsupportedOnGpu(
286+ `'${out.name}' (${src.count} elements) cannot feed ` +
287+ `'${to}' (${dst.count})`,
288+ );
289+ }
290+ ops.push({
291+ kind: 'copy',
292+ from: src.buffer,
293+ to: dst.buffer,
294+ bytes: 4 * src.count,
295+ label: `${out.name} -> ${to}`,
296+ });
297+ });
298+
299+ return new ModelPlan({
300+ device, sht, ops, byName, owned, paramBuf, paramData, paramNames,
301+ });
302+
303+ async function planStatement(stmt: IRStmt): Promise<void> {
304+ if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
305+ if (stmt.kind === 'For') return planFor(stmt);
306+ if (stmt.kind !== 'Assign') {
307+ throw new UnsupportedOnGpu(
308+ `a model function body may only contain assignments ` +
309+ `(found '${stmt.kind}')`,
310+ stmt.span,
311+ );
312+ }
313+ if (!isNumeric(stmt.ty)) {
314+ throw new UnsupportedOnGpu(
315+ `'${stmt.name}' is not a numeric value`,
316+ stmt.span,
317+ );
318+ }
319+ if (!isTensor(stmt.ty)) {
320+ // A scalar the model derives from its parameters (`us = a + b`). It
321+ // gets no buffer and no dispatch: the kernels that read it bind it as
322+ // a `let` in their prologue.
323+ derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
324+ return;
325+ }
326+ const count = numel(stmt.ty);
327+
328+ // Reuse the destination buffer across steps: the same cName always maps
329+ // to the same buffer, so a step allocates nothing.
330+ let dest = slots.get(stmt.cName);
331+ if (!dest) {
332+ dest = alloc(`mgpu-${stmt.name}`, count);
333+ slots.set(stmt.cName, dest);
334+ } else if (dest.count !== count) {
335+ throw new UnsupportedOnGpu(
336+ `'${stmt.name}' changes size between assignments`,
337+ stmt.span,
338+ );
339+ }
340+ byName.set(stmt.name, dest);
341+
342+ const ext = externalCall(stmt);
343+ if (ext) {
344+ const argSlot = slots.get(ext.argCName);
345+ if (!argSlot) {
346+ throw new UnsupportedOnGpu(
347+ `'${ext.name}' reads '${ext.argName}', which has no buffer`,
348+ stmt.span,
349+ );
350+ }
351+ ops.push(
352+ ext.name === 'synth'
353+ ? {
354+ kind: 'synth',
355+ binding: sht.createSynthBinding(argSlot.buffer, dest.buffer),
356+ label: `${stmt.name} = synth(${ext.argName})`,
357+ }
358+ : {
359+ kind: 'analys',
360+ binding: sht.createAnalysBinding(argSlot.buffer, dest.buffer),
361+ label: `${stmt.name} = analys(${ext.argName})`,
362+ },
363+ );
364+ return;
365+ }
366+
367+ // Element-wise kernel. Collect the distinct tensor operands and give
368+ // them dense binding slots.
369+ const tensors = new Map<string, number>();
370+ collectTensorVars(stmt.expr, (cName) => {
371+ if (!tensors.has(cName)) tensors.set(cName, tensors.size);
372+ });
373+
374+ const label = `${stmt.name} = <${count} elements, element-wise>`;
375+ const kernel = buildKernel(
376+ stmt,
377+ {
378+ tensors,
379+ params: paramSlots,
380+ scalars: derivedScalars,
381+ } satisfies KernelInputs,
382+ count,
383+ label,
384+ );
385+ const bindGroupLayout = kernelLayout(device, tensors.size);
386+ const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
387+
388+ // WebGPU forbids aliasing a writable storage binding with another
389+ // binding in the same group, so an in-place update (`u = u + 1`) writes
390+ // to scratch and copies back. Element-wise kernels only ever touch
391+ // their own index, so the copy is the only cost.
392+ const aliased = tensors.has(stmt.cName);
393+ const target = aliased ? alloc(`mgpu-${stmt.name}-scratch`, count) : dest;
394+
395+ const entries: GPUBindGroupEntry[] = [
396+ { binding: 0, resource: { buffer: target.buffer } },
397+ ];
398+ for (const [cName, i] of tensors) {
399+ const s = slots.get(cName);
400+ if (!s) {
401+ throw new UnsupportedOnGpu(
402+ `'${stmt.name}' reads a value with no buffer`,
403+ stmt.span,
404+ );
405+ }
406+ entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
407+ }
408+ entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
409+
410+ ops.push({
411+ kind: 'kernel',
412+ pipeline,
413+ bindGroup: device.createBindGroup({
414+ layout: bindGroupLayout,
415+ entries,
416+ }),
417+ count,
418+ label,
419+ copyBack: aliased
420+ ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
421+ : undefined,
422+ });
423+ }
424+
425+ /**
426+ * Unroll a counted loop into the op sequence.
427+ *
428+ * A plan is a fixed list of GPU operations with no branching, which is what
429+ * makes a timestep pure command recording. A `for` with compile-time-known
430+ * bounds still fits that: it is the same body planned once per iteration.
431+ * Nothing else changes — numbl gives a variable one cName for every
432+ * assignment to it, so the buffer an iteration writes is the buffer the
433+ * next one reads, which is exactly a loop-carried value.
434+ *
435+ * The loop variable gets no buffer either: it is bound as a derived scalar
436+ * to this iteration's literal value, so a kernel that reads `k` folds the
437+ * number in. The binding is overwritten per iteration, before that
438+ * iteration's body is planned and its WGSL emitted.
439+ */
440+ async function planFor(stmt: For): Promise<void> {
441+ const from = exactValue(stmt.start);
442+ const to = exactValue(stmt.end);
443+ if (from === undefined || to === undefined) {
444+ throw new UnsupportedOnGpu(
445+ `a 'for' loop is unrolled into the op sequence, so its bounds must ` +
446+ `be known when the model is compiled — ` +
447+ `${from === undefined ? 'the start' : 'the end'} of this one is a ` +
448+ `runtime value. Use a whole number, or a count the app supplies ` +
449+ `as a fixed argument (changing it recompiles).`,
450+ stmt.span,
451+ );
452+ }
453+ const trips = Math.floor((to - from) / stmt.step) + 1;
454+ if (!Number.isFinite(trips)) {
455+ throw new UnsupportedOnGpu(`'for ${stmt.varName}' has no finite length`, stmt.span);
456+ }
457+ if (trips > MAX_UNROLL) {
458+ throw new UnsupportedOnGpu(
459+ `'for ${stmt.varName}' would unroll to ${trips} iterations, over the ` +
460+ `limit of ${MAX_UNROLL}. Every iteration is separate GPU work, so a ` +
461+ `long loop compiles slowly and runs no faster than writing it out.`,
462+ stmt.span,
463+ );
464+ }
465+ for (let i = 0; i < trips; i++) {
466+ const value = from + i * stmt.step;
467+ derivedScalars.set(stmt.cVar, {
468+ name: stmt.varName,
469+ expr: {
470+ kind: 'NumLit',
471+ value,
472+ ty: scalarDouble(
473+ value > 0 ? 'positive' : value < 0 ? 'negative' : 'zero',
474+ value,
475+ ),
476+ span: stmt.span,
477+ },
478+ });
479+ for (const s of stmt.body) await planStatement(s);
480+ }
481+ }
482+ }
483+
484+ /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
485+ setParams(values: Record<string, number>): void {
486+ this.paramNames.forEach((name, i) => {
487+ const v = values[name];
488+ this.#paramData[i] = Number.isFinite(v) ? v : 0;
489+ });
490+ this.#device.queue.writeBuffer(
491+ this.#paramBuf,
492+ 0,
493+ this.#paramData as Float32Array<ArrayBuffer>,
494+ );
495+ }
496+
497+ /** Buffer holding the named value, or undefined if the .m never binds it. */
498+ buffer(name: string): GPUBuffer | undefined {
499+ return this.#byName.get(name)?.buffer;
500+ }
501+
502+ elementCount(name: string): number | undefined {
503+ return this.#byName.get(name)?.count;
504+ }
505+
506+ /**
507+ * Record `steps` timesteps. Synchronous: no awaits, no readback. All of the
508+ * ops share one compute pass, which WebGPU executes in submission order
509+ * with a barrier between dispatches.
510+ */
511+ encodeSteps(encoder: GPUCommandEncoder, steps: number): void {
512+ for (let s = 0; s < steps; s++) {
513+ let pass: GPUComputePassEncoder | null = null;
514+ const inPass = (): GPUComputePassEncoder => {
515+ if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
516+ return pass;
517+ };
518+ const endPass = (): void => {
519+ if (pass) {
520+ pass.end();
521+ pass = null;
522+ }
523+ };
524+ for (const op of this.#ops) {
525+ switch (op.kind) {
526+ case 'kernel': {
527+ const p = inPass();
528+ p.setPipeline(op.pipeline);
529+ p.setBindGroup(0, op.bindGroup);
530+ p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
531+ if (op.copyBack) {
532+ endPass();
533+ encoder.copyBufferToBuffer(
534+ op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
535+ );
536+ }
537+ break;
538+ }
539+ case 'synth':
540+ this.#shtInto(inPass(), op);
541+ break;
542+ case 'analys':
543+ this.#shtInto(inPass(), op);
544+ break;
545+ case 'copy':
546+ endPass();
547+ encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
548+ break;
549+ }
550+ }
551+ endPass();
552+ }
553+ }
554+
555+ #shtInto(pass: GPUComputePassEncoder, op: Op & { kind: 'synth' | 'analys' }): void {
556+ if (op.kind === 'synth') this.#sht.encodeSynthInto(pass, op.binding);
557+ else this.#sht.encodeAnalysInto(pass, op.binding);
558+ }
559+
560+ /** Human-readable op sequence — what the .m actually compiled to. */
561+ describe(): string[] {
562+ return this.#ops.map((op) => `${op.kind.padEnd(7)} ${op.label}`);
563+ }
564+
565+ destroy(): void {
566+ for (const b of this.#owned) b.destroy();
567+ this.#paramBuf.destroy();
568+ this.#owned.length = 0;
569+ }
570+}
571+
572+/** `x = synth(y)` / `x = analys(y)` -> the call's name and argument. */
573+function externalCall(
574+ stmt: Assign,
575+): { name: string; argCName: string; argName: string } | null {
576+ const e = stmt.expr;
577+ if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
578+ if (e.args.length !== 1 || e.args[0].kind !== 'Var') {
579+ throw new UnsupportedOnGpu(
580+ `'${e.name}' must be applied to a single variable`,
581+ stmt.span,
582+ );
583+ }
584+ const arg = e.args[0];
585+ return { name: e.name, argCName: arg.cName, argName: arg.name };
586+}
587+
588+function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
589+ const walk = (x: IRExpr): void => {
590+ switch (x.kind) {
591+ case 'Var':
592+ if (isTensor(x.ty)) visit(x.cName);
593+ return;
594+ case 'Binary':
595+ walk(x.left);
596+ walk(x.right);
597+ return;
598+ case 'Unary':
599+ walk(x.operand);
600+ return;
601+ case 'Call':
602+ x.args.forEach(walk);
603+ return;
604+ default:
605+ return;
606+ }
607+ };
608+ walk(e);
609+}
src/mgpu/registry.tsadded+142−0View file
@@ -0,0 +1,142 @@
1+/**
2+ * The available models: their MATLAB source, and the metadata the host owns.
3+ *
4+ * A model's *algorithm* lives in its .m file. Everything around it lives here:
5+ * the parameter names the .m may take as arguments, their defaults and slider
6+ * ranges, which grid fields to render, and the dealiasing degree. The .m
7+ * declares nothing about these — it just names the parameters it wants, and
8+ * `CompiledModel` matches each against this table.
9+ *
10+ * Naming convention, documented in each .m:
11+ * `u`, `v`, ... grid fields the model computes and the app renders
12+ * `U`, `V`, ... the corresponding spectral state (uppercase)
13+ */
14+import schnakenbergSource from '../../models/schnakenberg.m?raw';
15+import brusselatorSource from '../../models/brusselator.m?raw';
16+import allencahnSource from '../../models/allencahn.m?raw';
17+
18+export type Params = Record<string, number>;
19+
20+/** A tunable scalar the .m may take as an argument. */
21+export interface ParamSpec {
22+ key: string;
23+ label: string;
24+ value: number;
25+ min: number;
26+ max: number;
27+ step: number;
28+}
29+
30+export interface MModel {
31+ key: string;
32+ label: string;
33+ blurb: string;
34+ /** Grid fields to render, one panel each. */
35+ species: string[];
36+ /** Spectral state names the .m advances. */
37+ state: string[];
38+ params: ParamSpec[];
39+ /** Polynomial degree of the reaction in the fields, for grid dealiasing. */
40+ pdeg: number;
41+ /** Amplitude of the seeded perturbation handed to `init`. */
42+ seedAmp: number;
43+ /** MATLAB source — the algorithm itself. */
44+ source: string;
45+}
46+
47+/** Spectral state names follow the grid-field names, uppercased. */
48+const stateFor = (species: string[]): string[] => species.map((s) => s.toUpperCase());
49+
50+const schnakenberg: MModel = {
51+ key: 'schnakenberg',
52+ label: 'Schnakenberg',
53+ blurb:
54+ 'Turing spots. The homogeneous state is stable to uniform perturbations ' +
55+ 'but unstable to degrees 14 ≤ l ≤ 40, most strongly at l = 24.',
56+ species: ['u', 'v'],
57+ state: stateFor(['u', 'v']),
58+ params: [
59+ { key: 'a', label: 'a', value: 0.1, min: 0.01, max: 0.5, step: 0.01 },
60+ { key: 'b', label: 'b', value: 0.9, min: 0.1, max: 2, step: 0.05 },
61+ { key: 'D1', label: 'D₁', value: 4e-4, min: 1e-5, max: 5e-3, step: 1e-5 },
62+ { key: 'D2', label: 'D₂', value: 8e-3, min: 1e-4, max: 5e-2, step: 1e-4 },
63+ { key: 'dt', label: 'dt', value: 0.05, min: 0.005, max: 0.5, step: 0.005 },
64+ ],
65+ pdeg: 3,
66+ seedAmp: 1e-2,
67+ source: schnakenbergSource,
68+};
69+
70+const brusselator: MModel = {
71+ key: 'brusselator',
72+ label: 'Brusselator',
73+ blurb:
74+ 'Turing stripes and spots, from a smaller diffusivity contrast than ' +
75+ 'Schnakenberg but with a stiffer reaction.',
76+ species: ['u', 'v'],
77+ state: stateFor(['u', 'v']),
78+ params: [
79+ { key: 'A', label: 'A', value: 3, min: 0.5, max: 6, step: 0.1 },
80+ { key: 'B', label: 'B', value: 9, min: 1, max: 15, step: 0.25 },
81+ { key: 'D1', label: 'D₁', value: 3.33e-3, min: 1e-4, max: 2e-2, step: 1e-4 },
82+ { key: 'D2', label: 'D₂', value: 1.67e-2, min: 1e-3, max: 1e-1, step: 1e-3 },
83+ { key: 'dt', label: 'dt', value: 0.02, min: 0.002, max: 0.1, step: 0.002 },
84+ ],
85+ pdeg: 3,
86+ seedAmp: 1e-2,
87+ source: brusselatorSource,
88+};
89+
90+const allencahn: MModel = {
91+ key: 'allencahn',
92+ label: 'Allen–Cahn',
93+ blurb:
94+ 'A single species: interfaces form and then coarsen until one domain ' +
95+ 'swallows the sphere.',
96+ species: ['u'],
97+ state: stateFor(['u']),
98+ params: [
99+ { key: 'eps2', label: 'ε²', value: 1e-3, min: 1e-4, max: 1e-2, step: 1e-4 },
100+ { key: 'dt', label: 'dt', value: 0.02, min: 0.002, max: 0.2, step: 0.002 },
101+ ],
102+ pdeg: 3,
103+ seedAmp: 1e-2,
104+ source: allencahnSource,
105+};
106+
107+export const mModels: MModel[] = [schnakenberg, brusselator, allencahn];
108+
109+export const mModelByKey = (key: string): MModel | undefined =>
110+ mModels.find((m) => m.key === key);
111+
112+export const defaultParams = (m: MModel): Params =>
113+ Object.fromEntries(m.params.map((p) => [p.key, p.value]));
114+
115+/** Named parameter presets shown in the UI dropdown. The pattern length scale
116+ * goes as 1/sqrt(D), so scaling both diffusivities moves the spot size without
117+ * changing the dynamics. */
118+export interface Preset {
119+ key: string;
120+ label: string;
121+ modelKey: string;
122+ /** Overrides applied on top of the model's default parameters. */
123+ params?: Params;
124+}
125+
126+export const presets: Preset[] = [
127+ { key: 'schnak-spots', label: 'Schnakenberg — spots', modelKey: 'schnakenberg' },
128+ {
129+ key: 'schnak-coarse',
130+ label: 'Schnakenberg — coarse spots',
131+ modelKey: 'schnakenberg',
132+ params: { D1: 1e-3, D2: 2e-2 },
133+ },
134+ {
135+ key: 'schnak-fine',
136+ label: 'Schnakenberg — fine spots',
137+ modelKey: 'schnakenberg',
138+ params: { D1: 1.6e-4, D2: 3.2e-3 },
139+ },
140+ { key: 'brussel', label: 'Brusselator — stripes & spots', modelKey: 'brusselator' },
141+ { key: 'allencahn', label: 'Allen–Cahn — coarsening', modelKey: 'allencahn' },
142+];
src/mgpu/session.tsadded+302−0View file
@@ -0,0 +1,302 @@
1+/**
2+ * One running model: grid, transforms, compiled .m, seeded state.
3+ *
4+ * Everything that is not rendering. The app, the desktop benchmark and the
5+ * tests all go through this, so there is one place that decides how a model is
6+ * turned into something running on the GPU — and nothing about it is
7+ * browser-specific beyond needing a GPUDevice.
8+ */
9+import { ShtPlan } from '../sht/sht.ts';
10+import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
11+import { GpuModel, type ModelParams } from './model.ts';
12+import { seededNoise } from './noise.ts';
13+import type { MModel } from './registry.ts';
14+import { Geometry } from '../geom/geometry.ts';
15+import { mGeometryByKey, defaultGeometryParams, SPHERE_KEY, type MGeometry } from '../geom/registry.ts';
16+
17+export interface ModelSessionOptions {
18+ device: GPUDevice;
19+ model: MModel;
20+ params: ModelParams;
21+ lmax: number;
22+ /** Override the model source — the editor's working copy. */
23+ source?: string;
24+ /** Linear render oversampling: read the species fields on a grid this many
25+ * times finer than the solver's in each direction (default 1). The state is
26+ * band-limited at lmax, so the finer evaluation is exact interpolation. */
27+ oversample?: number;
28+ /** The surface to solve on. Defaults to the unit sphere. */
29+ geometry?: MGeometry;
30+ geometryParams?: ModelParams;
31+ /** Override the geometry source — the editor's working copy. */
32+ geometrySource?: string;
33+ /**
34+ * Iterations of the .m's implicit solve. Structural, not tunable: the loop
35+ * is unrolled into the op sequence, so a change recompiles.
36+ */
37+ niter?: number;
38+}
39+
40+export class ModelSession {
41+ readonly device: GPUDevice;
42+ readonly model: MModel;
43+ readonly cfg: ShtConfig;
44+ readonly sht: ShtPlan;
45+ readonly gpu: GpuModel;
46+ readonly npts: number;
47+ /** Iterations of the implicit solve compiled into the step. */
48+ readonly niter: number;
49+
50+ /** The surface being solved on, as spherical-harmonic coefficients. */
51+ #geometry: Geometry;
52+ #geometryModel: MGeometry;
53+
54+ /** Model time and step count since the last seeding. */
55+ t = 0;
56+ steps = 0;
57+
58+ #params: ModelParams;
59+ /** Display-only transforms on the oversampled grid; null at 1x. */
60+ #displaySht: ShtPlan | null;
61+ #oversample: number;
62+
63+ private constructor(init: {
64+ device: GPUDevice;
65+ model: MModel;
66+ cfg: ShtConfig;
67+ sht: ShtPlan;
68+ displaySht: ShtPlan | null;
69+ gpu: GpuModel;
70+ params: ModelParams;
71+ oversample: number;
72+ geometry: Geometry;
73+ geometryModel: MGeometry;
74+ niter: number;
75+ }) {
76+ this.device = init.device;
77+ this.model = init.model;
78+ this.cfg = init.cfg;
79+ this.sht = init.sht;
80+ this.gpu = init.gpu;
81+ this.npts = init.cfg.nlat * init.cfg.nphi;
82+ this.#oversample = init.oversample;
83+ this.#params = init.params;
84+ this.#displaySht = init.displaySht;
85+ this.#geometry = init.geometry;
86+ this.#geometryModel = init.geometryModel;
87+ this.niter = init.niter;
88+ }
89+
90+ get geometry(): Geometry {
91+ return this.#geometry;
92+ }
93+
94+ get geometryModel(): MGeometry {
95+ return this.#geometryModel;
96+ }
97+
98+ /** Linear render oversampling factor (1 = read on the solver grid). */
99+ get oversample(): number {
100+ return this.#oversample;
101+ }
102+
103+ static async create(opts: ModelSessionOptions): Promise<ModelSession> {
104+ const { device, model, params, lmax } = opts;
105+ const oversample = Math.max(1, Math.round(opts.oversample ?? 1));
106+ const niter = Math.max(0, Math.round(opts.niter ?? 1));
107+ const geometryModel = opts.geometry ?? mGeometryByKey(SPHERE_KEY)!;
108+ const geometryParams = opts.geometryParams ?? defaultGeometryParams(geometryModel);
109+ const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
110+ const cfg = { lmax, mmax: lmax, nlat, nphi };
111+ const sht = await ShtPlan.create(device, cfg);
112+ let displaySht: ShtPlan | null = null;
113+ try {
114+ // The display plan shares nothing with the solver's beyond the
115+ // coefficients copied into it per readback; its grid is the solver's
116+ // scaled by the oversampling factor, so nphi stays a power of two (the
117+ // FFT path) for power-of-two factors.
118+ if (oversample > 1) {
119+ displaySht = await ShtPlan.create(device, {
120+ lmax,
121+ mmax: lmax,
122+ nlat: oversample * nlat,
123+ nphi: oversample * nphi,
124+ });
125+ }
126+ // The surface is built before the model, because the model takes it as
127+ // an argument. It is a one-off: compiled, evaluated, read back, and its
128+ // plan discarded — nothing of it survives into the timestep but six
129+ // buffers of numbers.
130+ const geometry = await Geometry.create({
131+ device,
132+ sht,
133+ cfg,
134+ source: opts.geometrySource ?? geometryModel.source,
135+ paramNames: geometryModel.params.map((p) => p.key),
136+ params: geometryParams,
137+ });
138+ const gpu = await GpuModel.create({
139+ device,
140+ sht,
141+ cfg,
142+ source: opts.source ?? model.source,
143+ paramNames: model.params.map((p) => p.key),
144+ state: model.state,
145+ view: model.species,
146+ geometry,
147+ niter,
148+ });
149+ gpu.setParams(params);
150+ return new ModelSession({
151+ device, model, cfg, sht, displaySht, gpu, params, oversample,
152+ geometry, geometryModel, niter,
153+ });
154+ } catch (e) {
155+ // The transform plans own GPU buffers; do not leak them on a compile error.
156+ displaySht?.destroy();
157+ sht.destroy();
158+ throw e;
159+ }
160+ }
161+
162+ /**
163+ * Vertex positions for the current render grid: the surface synthesized on
164+ * `viewSht`, interleaved xyz. Exact interpolation of the same coefficients
165+ * the solver sees, so the drawn surface is the one being solved on however
166+ * finely it is sampled.
167+ */
168+ renderPositions(): Promise<Float32Array> {
169+ return this.#geometry.positionsOn(this.viewSht);
170+ }
171+
172+ /**
173+ * Change the surface in place, without recompiling or disturbing the run.
174+ * The geometry's shape in the bindings depends only on the grid, so the
175+ * compiled step does not change — only the numbers it reads. The caller
176+ * still has to rebuild the mesh from `renderPositions()`.
177+ */
178+ async setGeometry(
179+ geometryModel: MGeometry,
180+ params: ModelParams,
181+ source?: string,
182+ ): Promise<void> {
183+ const next = await Geometry.create({
184+ device: this.device,
185+ sht: this.sht,
186+ cfg: this.cfg,
187+ source: source ?? geometryModel.source,
188+ paramNames: geometryModel.params.map((p) => p.key),
189+ params,
190+ });
191+ this.#geometry = next;
192+ this.#geometryModel = geometryModel;
193+ this.gpu.uploadGeometry(next);
194+ }
195+
196+ /** The plan whose grid `readSpecies` samples on — the display plan when
197+ * oversampling, otherwise the solver's. Its cosTheta/nphi define the mesh. */
198+ get viewSht(): ShtPlan {
199+ return this.#displaySht ?? this.sht;
200+ }
201+
202+ /**
203+ * Change the display oversampling in place. Display-only: the simulation
204+ * state, time and parameters are untouched, so the run continues seamlessly
205+ * on the new render grid. The caller must not have a readSpecies in flight —
206+ * its readback maps a buffer of the plan being destroyed.
207+ */
208+ async setOversample(oversample: number): Promise<void> {
209+ const os = Math.max(1, Math.round(oversample));
210+ if (os === this.#oversample) return;
211+ const next =
212+ os > 1
213+ ? await ShtPlan.create(this.device, {
214+ lmax: this.cfg.lmax,
215+ mmax: this.cfg.mmax,
216+ nlat: os * this.cfg.nlat,
217+ nphi: os * this.cfg.nphi,
218+ })
219+ : null;
220+ const old = this.#displaySht;
221+ this.#displaySht = next;
222+ this.#oversample = os;
223+ old?.destroy();
224+ }
225+
226+ /** Run `init` from a seeded perturbation, resetting model time. */
227+ seed(seed: number): void {
228+ this.gpu.init(seededNoise(this.npts, this.model.seedAmp, seed));
229+ this.t = 0;
230+ this.steps = 0;
231+ }
232+
233+ setParams(params: ModelParams): void {
234+ this.#params = params;
235+ this.gpu.setParams(params);
236+ }
237+
238+ /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
239+ step(n = 1): void {
240+ this.gpu.step(n);
241+ this.t += n * (this.#params.dt ?? 0);
242+ this.steps += n;
243+ }
244+
245+ /**
246+ * Wait for the submitted steps to finish, without reading anything back.
247+ * This is the honest way to time the solver: a readback would add a GPU->CPU
248+ * round trip, which in a browser also crosses a process boundary and can cost
249+ * more than the steps themselves.
250+ */
251+ sync(): Promise<undefined> {
252+ return this.device.queue.onSubmittedWorkDone();
253+ }
254+
255+ /**
256+ * Time a batch of `n` steps and return ms/step, leaving the simulation
257+ * exactly where it was: the spectral state is snapshotted before the batch
258+ * and restored after, and `t`/`steps` do not advance. One sync amortized
259+ * over the batch — the same measurement the desktop benchmark makes. The
260+ * grid view fields hold the batch's output until the next real step, so
261+ * step before reading them.
262+ */
263+ async measure(n: number): Promise<number> {
264+ this.gpu.snapshotState();
265+ const t0 = performance.now();
266+ this.gpu.step(n);
267+ await this.sync();
268+ const ms = (performance.now() - t0) / n;
269+ this.gpu.restoreState();
270+ return ms;
271+ }
272+
273+ /** Read a named value (a grid field or the spectral state). */
274+ read(name: string): Promise<Float32Array> {
275+ return this.gpu.read(name);
276+ }
277+
278+ /**
279+ * Read species `k` at render resolution (`viewSht`'s grid). Without
280+ * oversampling this is the grid field the .m returned. With oversampling the
281+ * spectral state is synthesized on the finer grid instead — the same field,
282+ * since the models define each species as synth of its state, evaluated
283+ * exactly on more points.
284+ */
285+ readSpecies(k: number): Promise<Float32Array> {
286+ if (!this.#displaySht) return this.read(this.model.species[k]);
287+ const state = this.model.state[k];
288+ const buf = this.gpu.valueBuffer(state);
289+ if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
290+ return this.#displaySht.synthFrom(buf);
291+ }
292+
293+ describe(): { init: string[]; step: string[] } {
294+ return this.gpu.describe();
295+ }
296+
297+ destroy(): void {
298+ this.gpu.destroy();
299+ this.#displaySht?.destroy();
300+ this.sht.destroy();
301+ }
302+}
src/mgpu/wgsl.tsadded+372−0View file
@@ -0,0 +1,372 @@
1+/**
2+ * IR expression tree -> one WGSL compute kernel.
3+ *
4+ * This is the WebGPU counterpart of numbl's C-side fused emitter
5+ * (`codegen/emitTensorFused.ts`): for an `Assign` whose right-hand side is
6+ * purely element-wise over operands of the target's shape, emit a single
7+ * kernel that computes one output element per invocation. Because numbl's
8+ * inline pass has already folded the ANF temps back together, one source line
9+ * of MATLAB becomes one kernel.
10+ *
11+ * Everything is f32, matching the existing fp32 WebGPU transform backend.
12+ */
13+import { getBuiltin } from 'numbl-src/numbl-core/jit/builtins/index.ts';
14+import { isMultiElement } from 'numbl-src/numbl-core/jit/lowering/types.ts';
15+import type { IRExpr, Assign } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
16+import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
17+
18+/** Raised for a construct the WGSL backend cannot express. Mirrors numbl's
19+ * own decline discipline: fail at compile time with a source span, never
20+ * silently produce something that computes the wrong thing. */
21+export class UnsupportedOnGpu extends Error {
22+ readonly span?: unknown;
23+ constructor(message: string, span?: unknown) {
24+ super(message);
25+ this.name = 'UnsupportedOnGpu';
26+ this.span = span;
27+ }
28+}
29+
30+const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
31+const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
32+
33+/** Element-wise binary builtins -> WGSL infix operator. */
34+const BINARY_OPS: Record<string, string> = {
35+ plus: '+',
36+ minus: '-',
37+ times: '*',
38+ rdivide: '/',
39+ // Degenerate to element-wise when at least one side is a scalar; the
40+ // both-tensor (true matrix) case is rejected below.
41+ mtimes: '*',
42+ mrdivide: '/',
43+};
44+
45+/** Element-wise unary builtins -> WGSL prefix operator. */
46+const UNARY_OPS: Record<string, string> = { uminus: '-', uplus: '+' };
47+
48+/** Element-wise builtin calls -> WGSL builtin of the same arity. */
49+const CALL_FNS: Record<string, string> = {
50+ abs: 'abs',
51+ acos: 'acos',
52+ asin: 'asin',
53+ atan: 'atan',
54+ atan2: 'atan2',
55+ ceil: 'ceil',
56+ cos: 'cos',
57+ cosh: 'cosh',
58+ exp: 'exp',
59+ floor: 'floor',
60+ log: 'log',
61+ log2: 'log2',
62+ max: 'max',
63+ min: 'min',
64+ round: 'round',
65+ sign: 'sign',
66+ sin: 'sin',
67+ sinh: 'sinh',
68+ sqrt: 'sqrt',
69+ tan: 'tan',
70+ tanh: 'tanh',
71+};
72+
73+/** WGSL f32 literal. Must always carry a decimal point or exponent, or WGSL
74+ * infers AbstractInt and rejects the mixed-type arithmetic. */
75+function f32Lit(v: number): string {
76+ if (!Number.isFinite(v)) {
77+ throw new UnsupportedOnGpu(`cannot emit non-finite literal ${v}`);
78+ }
79+ return Number.isInteger(v) && Math.abs(v) < 1e21
80+ ? `${v}.0`
81+ : String(v).includes('e')
82+ ? `${v}f`
83+ : String(v);
84+}
85+
86+/** How a scalar or tensor operand is read inside the kernel. */
87+export interface KernelInputs {
88+ /** cName -> storage binding index, for multi-element tensor operands. */
89+ tensors: Map<string, number>;
90+ /** cName -> slot in the params storage buffer, for runtime scalars. */
91+ params: Map<string, number>;
92+ /** cName -> defining expression, for scalars the .m computes from
93+ * parameters (`us = a + b`). These have no buffer and no param slot; they
94+ * become `let` bindings in the prologue of every kernel that reads them. */
95+ scalars: Map<string, { name: string; expr: IRExpr }>;
96+}
97+
98+/** Mutable state while emitting one kernel. */
99+interface Ctx {
100+ io: KernelInputs;
101+ /** `let` lines to emit before the body, in dependency order. */
102+ prologue: string[];
103+ /** cName -> WGSL identifier, for scalars already bound in the prologue. */
104+ bound: Map<string, string>;
105+}
106+
107+/** WGSL identifier for a derived scalar. Avoids a leading underscore, which
108+ * WGSL reserves. */
109+const scalarIdent = (cName: string): string =>
110+ `s_${cName.replace(/[^A-Za-z0-9_]/g, '_')}`;
111+
112+/**
113+ * Bind a .m-derived scalar in the prologue (once), after whatever it depends
114+ * on, and return its identifier.
115+ */
116+function bindScalar(cName: string, ctx: Ctx): string {
117+ const already = ctx.bound.get(cName);
118+ if (already) return already;
119+ const def = ctx.io.scalars.get(cName)!;
120+ const ident = scalarIdent(cName);
121+ // Claim the name before emitting the RHS so a (malformed) self-reference
122+ // cannot recurse forever.
123+ ctx.bound.set(cName, ident);
124+ const rhs = emitExpr(def.expr, ctx);
125+ ctx.prologue.push(` let ${ident} = ${rhs};`);
126+ return ident;
127+}
128+
129+/**
130+ * Emit the per-element WGSL expression for `e`. `i` is the element index
131+ * variable in scope.
132+ */
133+function emitExpr(e: IRExpr, ctx: Ctx): string {
134+ const io = ctx.io;
135+ switch (e.kind) {
136+ case 'NumLit':
137+ return f32Lit(e.value);
138+
139+ case 'Var': {
140+ if (isTensor(e.ty)) {
141+ const slot = io.tensors.get(e.cName);
142+ if (slot === undefined) {
143+ throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
144+ }
145+ return `in${slot}[i]`;
146+ }
147+ // Scalar: either an exact compile-time value or a runtime parameter.
148+ if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
149+ return f32Lit(e.ty.exact);
150+ }
151+ const slot = io.params.get(e.cName);
152+ if (slot !== undefined) return `prm[${slot}]`;
153+ if (io.scalars.has(e.cName)) return bindScalar(e.cName, ctx);
154+ throw new UnsupportedOnGpu(
155+ `scalar '${e.name}' is not a constant, a parameter, or computed in ` +
156+ `this model`,
157+ e.span,
158+ );
159+ }
160+
161+ case 'Binary': {
162+ if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
163+ isTensor(e.left.ty) && isTensor(e.right.ty)) {
164+ throw new UnsupportedOnGpu(
165+ `matrix '${e.builtin === 'mtimes' ? '*' : '/'}' is not supported; ` +
166+ `use the element-wise form ('.${e.builtin === 'mtimes' ? '*' : '/'}')`,
167+ e.span,
168+ );
169+ }
170+ if (e.builtin === 'power' || e.builtin === 'mpower') {
171+ return emitPower(e.left, e.right, ctx, e.span);
172+ }
173+ const op = BINARY_OPS[e.builtin];
174+ if (!op) {
175+ throw new UnsupportedOnGpu(`operator '${e.builtin}' is not supported`, e.span);
176+ }
177+ return `(${emitExpr(e.left, ctx)} ${op} ${emitExpr(e.right, ctx)})`;
178+ }
179+
180+ case 'Unary': {
181+ const op = UNARY_OPS[e.builtin];
182+ if (!op) {
183+ throw new UnsupportedOnGpu(`unary '${e.builtin}' is not supported`, e.span);
184+ }
185+ return `(${op}${emitExpr(e.operand, ctx)})`;
186+ }
187+
188+ case 'Call': {
189+ // A shape constructor used inside an element-wise expression
190+ // contributes the same constant at every slot, so it needs no buffer.
191+ // (The shape itself is validated against the target by checkShapes.)
192+ if (e.name === 'ones') return '1.0';
193+ if (e.name === 'zeros') return '0.0';
194+
195+ const fn = CALL_FNS[e.name];
196+ const b = getBuiltin(e.name);
197+ if (!fn || !b?.elementwise) {
198+ // A call numbl resolved to another function in the file gets a mangled
199+ // specialization name; a builtin keeps its source-level name. Only the
200+ // model's entry points are compiled, so a helper is a distinct failure
201+ // from an unsupported builtin and deserves to say so.
202+ const isUserFunction = e.cName !== e.name;
203+ throw new UnsupportedOnGpu(
204+ isUserFunction
205+ ? `'${e.name}' is a function defined in this model. Only init and ` +
206+ `step are compiled — inline its body into the caller.`
207+ : `'${e.name}' cannot be evaluated element-wise on the GPU`,
208+ e.span,
209+ );
210+ }
211+ return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
212+ }
213+
214+ default:
215+ throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
216+ }
217+}
218+
219+/**
220+ * `x.^k`. WGSL's `pow` is undefined for a negative base, and these fields go
221+ * negative routinely, so expand small non-negative integer exponents into
222+ * repeated multiplication — which is also what makes `u.^2` free.
223+ */
224+function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: unknown): string {
225+ const k =
226+ exponent.kind === 'NumLit'
227+ ? exponent.value
228+ : isNumeric(exponent.ty) && typeof exponent.ty.exact === 'number'
229+ ? exponent.ty.exact
230+ : undefined;
231+ const b = emitExpr(base, ctx);
232+ if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 8) {
233+ if (k === 0) return '1.0';
234+ // bind once so a compound base expression is not re-evaluated k times
235+ return `pow_i${k}(${b})`;
236+ }
237+ if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -8) {
238+ return `(1.0 / pow_i${-k}(${b}))`;
239+ }
240+ throw new UnsupportedOnGpu(
241+ `'.^' needs a literal integer exponent in [-8, 8] (got ` +
242+ `${k === undefined ? 'a runtime value' : k}); a negative base makes ` +
243+ `WGSL's pow() undefined`,
244+ span,
245+ );
246+}
247+
248+/** Fixed-exponent power helpers, emitted only when used. */
249+function powHelpers(used: Set<number>): string {
250+ const out: string[] = [];
251+ for (const k of [...used].sort((a, b) => a - b)) {
252+ const body =
253+ k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
254+ out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
255+ }
256+ return out.join('\n');
257+}
258+
259+/**
260+ * Reject implicit expansion (broadcasting).
261+ *
262+ * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
263+ * expansion semantics — but a kernel that walks one linear index across every
264+ * operand would quietly compute the wrong thing. So every multi-element
265+ * operand must have exactly the target's shape. Scalars are fine: they are
266+ * read from the params buffer or folded in as literals.
267+ */
268+function checkShapes(e: IRExpr, target: NumericType, name: string): void {
269+ const want = target.shape;
270+ const same = (t: NumericType): boolean => {
271+ const got = t.shape;
272+ return (
273+ !!want && !!got && want.length === got.length &&
274+ want.every((d, i) => d === got[i])
275+ );
276+ };
277+ const walk = (x: IRExpr): void => {
278+ if (isNumeric(x.ty) && isMultiElement(x.ty) && !same(x.ty)) {
279+ const got = x.ty.shape?.join('x') ?? 'dynamic';
280+ throw new UnsupportedOnGpu(
281+ `'${name}' would need implicit expansion: an operand is ${got} but the ` +
282+ `result is ${want?.join('x') ?? 'dynamic'}. Expand it explicitly ` +
283+ `(the GPU kernel walks one index across every operand).`,
284+ x.span,
285+ );
286+ }
287+ switch (x.kind) {
288+ case 'Binary':
289+ walk(x.left);
290+ walk(x.right);
291+ return;
292+ case 'Unary':
293+ walk(x.operand);
294+ return;
295+ case 'Call':
296+ // A shape constructor's own arguments are sizes, not data.
297+ if (x.name !== 'ones' && x.name !== 'zeros') x.args.forEach(walk);
298+ return;
299+ default:
300+ return;
301+ }
302+ };
303+ walk(e);
304+}
305+
306+export const WORKGROUP_SIZE = 64;
307+
308+export interface Kernel {
309+ code: string;
310+ /** Number of output elements. */
311+ count: number;
312+ label: string;
313+}
314+
315+/**
316+ * Build the kernel for one element-wise `Assign`. `io` must already map every
317+ * tensor operand cName to a binding index and every runtime scalar to a
318+ * params slot; the output is binding 0 and the params buffer is the binding
319+ * after the last input.
320+ */
321+export function buildKernel(
322+ stmt: Assign,
323+ io: KernelInputs,
324+ count: number,
325+ label: string,
326+): Kernel {
327+ if (!isNumeric(stmt.ty)) {
328+ throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
329+ }
330+ if (stmt.ty.isComplex) {
331+ throw new UnsupportedOnGpu(
332+ `'${stmt.name}' is complex; the GPU backend is real-only (a spectral ` +
333+ `field is carried as a real 2 x nlm array)`,
334+ stmt.span,
335+ );
336+ }
337+
338+ checkShapes(stmt.expr, stmt.ty, stmt.name);
339+ const ctx: Ctx = { io, prologue: [], bound: new Map() };
340+ const body = emitExpr(stmt.expr, ctx);
341+
342+ // pow_iK helpers are discovered during emission; scan the result for them.
343+ const used = new Set<number>();
344+ const emitted = [...ctx.prologue, body].join('\n');
345+ for (const m of emitted.matchAll(/\bpow_i(\d+)\(/g)) used.add(Number(m[1]));
346+
347+ const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
348+ for (const [, slot] of io.tensors) {
349+ decls.push(
350+ `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
351+ );
352+ }
353+ // Params live in a read-only storage buffer rather than a uniform block:
354+ // uniform arrays would need 16-byte element stride.
355+ const prmBinding = io.tensors.size + 1;
356+ decls.push(
357+ `@group(0) @binding(${prmBinding}) var<storage, read> prm: array<f32>;`,
358+ );
359+
360+ const code = `${decls.join('\n')}
361+
362+${powHelpers(used)}
363+
364+@compute @workgroup_size(${WORKGROUP_SIZE})
365+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
366+ let i = gid.x;
367+ if (i >= ${count}u) { return; }
368+${ctx.prologue.length ? `${ctx.prologue.join('\n')}\n` : ''} out[i] = ${body};
369+}
370+`;
371+ return { code, count, label };
372+}
src/raw.d.tsadded+5−0View file
@@ -0,0 +1,5 @@
1+/** Vite's `?raw` suffix imports a file's text. Used to load .m model sources. */
2+declare module '*?raw' {
3+ const source: string;
4+ export default source;
5+}
src/render/SphereScene.tsadded+292−0View file
@@ -0,0 +1,292 @@
1+import * as THREE from 'three';
2+import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
3+
4+/**
5+ * Three.js scene wrapper: a single indexed triangle mesh with dynamic
6+ * per-vertex positions and colors, orbit controls, and optional camera
7+ * synchronization with sibling scenes. The topology is fixed by the grid; the
8+ * positions are the surface, so they change when the geometry or the morph
9+ * does, and the colors every frame.
10+ *
11+ * Rendering is on demand: the animation loop ticks every frame (it has to,
12+ * to drive OrbitControls damping), but only re-renders when the colors,
13+ * camera, or canvas size actually changed.
14+ *
15+ * Adapted from figpack's SphereEmbedding view (figpack_experimental).
16+ */
17+export class SphereScene {
18+ #scene: THREE.Scene;
19+ #camera: THREE.PerspectiveCamera;
20+ #renderer: THREE.WebGLRenderer;
21+ #controls: OrbitControls;
22+ #geometry: THREE.BufferGeometry;
23+ #mesh: THREE.Mesh;
24+ #animationId: number | null = null;
25+ #defaultCameraState: {
26+ position: THREE.Vector3;
27+ target: THREE.Vector3;
28+ } | null = null;
29+ #syncing = false;
30+ #needsRender = true;
31+ #lastW = -1;
32+ #lastH = -1;
33+
34+ constructor(
35+ container: HTMLElement,
36+ numVertices: number,
37+ indices: Uint32Array,
38+ positions: Float32Array,
39+ background = '#14161c',
40+ ) {
41+ this.#scene = new THREE.Scene();
42+ this.#scene.background = new THREE.Color(background);
43+
44+ this.#camera = new THREE.PerspectiveCamera(50, 1, 0.01, 1000);
45+
46+ this.#renderer = new THREE.WebGLRenderer({ antialias: true });
47+ this.#renderer.setPixelRatio(window.devicePixelRatio || 1);
48+ // The canvas always fills its container via CSS; resize() then only
49+ // updates the drawing buffer
50+ this.#renderer.domElement.style.width = '100%';
51+ this.#renderer.domElement.style.height = '100%';
52+ this.#renderer.domElement.style.display = 'block';
53+ container.appendChild(this.#renderer.domElement);
54+
55+ // Lighting: ambient plus a headlight attached to the camera so the
56+ // surface stays lit from the viewing direction as it is rotated
57+ this.#scene.add(new THREE.AmbientLight(0xffffff, 0.65));
58+ const headlight = new THREE.DirectionalLight(0xffffff, 1.6);
59+ headlight.position.set(0.5, 0.8, 1);
60+ this.#camera.add(headlight);
61+ this.#scene.add(this.#camera);
62+
63+ this.#geometry = new THREE.BufferGeometry();
64+ // Positions move with the morph slider, so they are dynamic too.
65+ const positionAttr = new THREE.BufferAttribute(positions, 3);
66+ positionAttr.setUsage(THREE.DynamicDrawUsage);
67+ const colorAttr = new THREE.BufferAttribute(
68+ new Float32Array(numVertices * 3),
69+ 3,
70+ );
71+ colorAttr.setUsage(THREE.DynamicDrawUsage);
72+ this.#geometry.setAttribute('position', positionAttr);
73+ this.#geometry.setAttribute('color', colorAttr);
74+ this.#geometry.setIndex(new THREE.BufferAttribute(indices, 1));
75+ this.#geometry.computeVertexNormals();
76+ this.#geometry.computeBoundingSphere();
77+
78+ const material = new THREE.MeshPhongMaterial({
79+ vertexColors: true,
80+ side: THREE.DoubleSide,
81+ shininess: 25,
82+ specular: new THREE.Color(0x222222),
83+ });
84+ this.#mesh = new THREE.Mesh(this.#geometry, material);
85+ this.#scene.add(this.#mesh);
86+
87+ this.#controls = new OrbitControls(this.#camera, this.#renderer.domElement);
88+ this.#controls.enableDamping = true;
89+ this.#controls.dampingFactor = 0.1;
90+ // Fires on user input and on every damping-tail update, so the flag stays
91+ // set until the camera has fully settled.
92+ this.#controls.addEventListener('change', () => {
93+ this.#needsRender = true;
94+ });
95+
96+ this.#animate();
97+ }
98+
99+ #animate = () => {
100+ this.#animationId = requestAnimationFrame(this.#animate);
101+ this.#controls.update();
102+ if (!this.#needsRender) return;
103+ this.#needsRender = false;
104+ this.#renderer.render(this.#scene, this.#camera);
105+ };
106+
107+ updateColors(colors: Float32Array): void {
108+ const attr = this.#geometry.getAttribute('color') as THREE.BufferAttribute;
109+ (attr.array as Float32Array).set(colors);
110+ attr.needsUpdate = true;
111+ this.#needsRender = true;
112+ }
113+
114+ /**
115+ * Move the vertices — for the sphere/surface morph. Normals have to be
116+ * recomputed with them or the shading stays that of the old shape, which is
117+ * the whole thing the eye reads a curved surface by.
118+ */
119+ updatePositions(positions: Float32Array): void {
120+ const attr = this.#geometry.getAttribute('position') as THREE.BufferAttribute;
121+ (attr.array as Float32Array).set(positions);
122+ attr.needsUpdate = true;
123+ this.#geometry.computeVertexNormals();
124+ this.#geometry.computeBoundingSphere();
125+ this.#needsRender = true;
126+ }
127+
128+ /** The renderer's canvas, for capturing frames. */
129+ get canvas(): HTMLCanvasElement {
130+ return this.#renderer.domElement;
131+ }
132+
133+ /**
134+ * Render immediately, outside the animation loop. A WebGL canvas without
135+ * preserveDrawingBuffer keeps its drawing buffer only until the browser next
136+ * composites, so a capturer must render and copy within one task.
137+ */
138+ renderNow(): void {
139+ this.#needsRender = false;
140+ this.#renderer.render(this.#scene, this.#camera);
141+ }
142+
143+ /** Mirror this scene's camera whenever the other scene's controls move. */
144+ syncCamerasWith(other: SphereScene): void {
145+ const follow = (src: SphereScene, dst: SphereScene) => {
146+ src.#controls.addEventListener('change', () => {
147+ if (dst.#syncing) return;
148+ src.#syncing = true;
149+ dst.#camera.position.copy(src.#camera.position);
150+ dst.#camera.zoom = src.#camera.zoom;
151+ dst.#camera.updateProjectionMatrix();
152+ dst.#controls.target.copy(src.#controls.target);
153+ dst.#controls.update();
154+ dst.#needsRender = true;
155+ src.#syncing = false;
156+ });
157+ };
158+ follow(this, other);
159+ follow(other, this);
160+ }
161+
162+ /** Orbit the camera about the up axis by `angle` radians, keeping the
163+ * target. Synced sibling scenes follow via their controls, as with a drag. */
164+ orbitBy(angle: number): void {
165+ const offset = this.#camera.position.clone().sub(this.#controls.target);
166+ offset.applyAxisAngle(this.#camera.up, angle);
167+ this.#camera.position.copy(this.#controls.target).add(offset);
168+ this.#controls.update();
169+ this.#needsRender = true;
170+ }
171+
172+ /** Camera pose, for carrying the view across a scene rebuild. */
173+ cameraState(): { position: THREE.Vector3; target: THREE.Vector3; zoom: number } {
174+ return {
175+ position: this.#camera.position.clone(),
176+ target: this.#controls.target.clone(),
177+ zoom: this.#camera.zoom,
178+ };
179+ }
180+
181+ setCameraState(s: {
182+ position: THREE.Vector3;
183+ target: THREE.Vector3;
184+ zoom: number;
185+ }): void {
186+ this.#camera.position.copy(s.position);
187+ this.#camera.zoom = s.zoom;
188+ this.#camera.updateProjectionMatrix();
189+ this.#controls.target.copy(s.target);
190+ this.#controls.update();
191+ this.#needsRender = true;
192+ }
193+
194+ /**
195+ * Position the camera to comfortably frame the geometry.
196+ *
197+ * The distance is generous on purpose. The bounding sphere is of the surface
198+ * currently loaded, but the camera is *kept* across a geometry change and
199+ * across the morph, so a frame that only just fits the shape at hand would
200+ * clip the next one. Leaving room means switching shapes never needs a
201+ * camera reset to see what happened.
202+ */
203+ fitCamera(): void {
204+ this.#geometry.computeBoundingSphere();
205+ const bs = this.#geometry.boundingSphere;
206+ if (!bs) return;
207+ const radius = Math.max(bs.radius, 1e-6);
208+ const distance = radius * 3.4;
209+ this.#controls.target.copy(bs.center);
210+ this.#camera.position.set(
211+ bs.center.x + distance * 0.55,
212+ bs.center.y + distance * 0.35,
213+ bs.center.z + distance * 0.75,
214+ );
215+ this.#camera.near = radius * 0.01;
216+ this.#camera.far = radius * 100;
217+ this.#camera.updateProjectionMatrix();
218+ this.#controls.update();
219+ this.#needsRender = true;
220+ this.#defaultCameraState = {
221+ position: this.#camera.position.clone(),
222+ target: this.#controls.target.clone(),
223+ };
224+ }
225+
226+ resetCamera(): void {
227+ if (this.#defaultCameraState) {
228+ this.#camera.position.copy(this.#defaultCameraState.position);
229+ this.#controls.target.copy(this.#defaultCameraState.target);
230+ this.#controls.update();
231+ this.#needsRender = true;
232+ } else {
233+ this.fitCamera();
234+ }
235+ }
236+
237+ resize(width: number, height: number): void {
238+ // Setting canvas.width clears the canvas even at the same value, which
239+ // shows as a blank flash until the next render — skip no-op resizes.
240+ if (width === this.#lastW && height === this.#lastH) return;
241+ this.#lastW = width;
242+ this.#lastH = height;
243+ this.#camera.aspect = width / Math.max(1, height);
244+ this.#camera.updateProjectionMatrix();
245+ // updateStyle=false: the canvas keeps its 100%/100% CSS sizing
246+ this.#renderer.setSize(width, height, false);
247+ // setSize clears the drawing buffer, so a re-render is required even
248+ // though nothing in the scene moved
249+ this.#needsRender = true;
250+ }
251+
252+ /**
253+ * Set the drawing buffer to an exact square pixel size, independent of the
254+ * container and devicePixelRatio — for capturing at a chosen resolution.
255+ * The canvas keeps its CSS sizing, so on screen it just rescales. Undo with
256+ * restoreSize().
257+ */
258+ captureSize(px: number): void {
259+ this.#renderer.setPixelRatio(1);
260+ this.#renderer.setSize(px, px, false);
261+ this.#camera.aspect = 1;
262+ this.#camera.updateProjectionMatrix();
263+ this.#needsRender = true;
264+ }
265+
266+ /** Return from captureSize() to the container-driven buffer size. */
267+ restoreSize(): void {
268+ this.#renderer.setPixelRatio(window.devicePixelRatio || 1);
269+ if (this.#lastW > 0 && this.#lastH > 0) {
270+ this.#renderer.setSize(this.#lastW, this.#lastH, false);
271+ this.#camera.aspect = this.#lastW / Math.max(1, this.#lastH);
272+ this.#camera.updateProjectionMatrix();
273+ }
274+ this.#needsRender = true;
275+ }
276+
277+ dispose(): void {
278+ if (this.#animationId !== null) {
279+ cancelAnimationFrame(this.#animationId);
280+ this.#animationId = null;
281+ }
282+ this.#controls.dispose();
283+ this.#geometry.dispose();
284+ (this.#mesh.material as THREE.Material).dispose();
285+ if (this.#renderer.domElement.parentNode) {
286+ this.#renderer.domElement.parentNode.removeChild(
287+ this.#renderer.domElement,
288+ );
289+ }
290+ this.#renderer.dispose();
291+ }
292+}
src/render/colorbar.tsadded+38−0View file
@@ -0,0 +1,38 @@
1+import type { ColormapFunc } from './colormaps.ts';
2+
3+/** Compact numeric label: 3 significant digits, trailing zeros trimmed. */
4+export const fmtValue = (v: number): string =>
5+ Number.isFinite(v) ? v.toPrecision(3).replace(/\.?0+$/, '') : '—';
6+
7+/** Vertical colorbar drawn on a small canvas, with min/max labels. */
8+export class Colorbar {
9+ #canvas: HTMLCanvasElement;
10+ #minLabel: HTMLElement;
11+ #maxLabel: HTMLElement;
12+
13+ constructor(container: HTMLElement) {
14+ container.classList.add('colorbar');
15+ this.#maxLabel = document.createElement('div');
16+ this.#maxLabel.className = 'colorbar-label';
17+ this.#canvas = document.createElement('canvas');
18+ this.#canvas.width = 12;
19+ this.#canvas.height = 160;
20+ this.#minLabel = document.createElement('div');
21+ this.#minLabel.className = 'colorbar-label';
22+ container.append(this.#maxLabel, this.#canvas, this.#minLabel);
23+ }
24+
25+ update(cmap: ColormapFunc, vmin: number, vmax: number): void {
26+ const ctx = this.#canvas.getContext('2d');
27+ if (!ctx) return;
28+ const h = this.#canvas.height;
29+ for (let y = 0; y < h; y++) {
30+ const t = 1 - y / (h - 1);
31+ const [r, g, b] = cmap(t);
32+ ctx.fillStyle = `rgb(${r},${g},${b})`;
33+ ctx.fillRect(0, y, this.#canvas.width, 1);
34+ }
35+ this.#maxLabel.textContent = fmtValue(vmax);
36+ this.#minLabel.textContent = fmtValue(vmin);
37+ }
38+}
src/render/colormaps.tsadded+98−0View file
@@ -0,0 +1,98 @@
1+/**
2+ * Colormaps: each maps a normalized value in [0, 1] to [r, g, b] in [0, 255].
3+ * Adapted from figpack's SphereEmbedding view (figpack_experimental).
4+ */
5+
6+export type ColormapFunc = (t: number) => [number, number, number];
7+
8+const clamp01 = (t: number) => Math.max(0, Math.min(1, t));
9+
10+// Piecewise-linear interpolation through control points (r, g, b in 0-255)
11+const makeInterpolated = (stops: [number, number, number][]): ColormapFunc => {
12+ const n = stops.length;
13+ return (t: number) => {
14+ t = clamp01(t);
15+ const x = t * (n - 1);
16+ const i = Math.min(n - 2, Math.floor(x));
17+ const f = x - i;
18+ const a = stops[i];
19+ const b = stops[i + 1];
20+ return [
21+ Math.round(a[0] + (b[0] - a[0]) * f),
22+ Math.round(a[1] + (b[1] - a[1]) * f),
23+ Math.round(a[2] + (b[2] - a[2]) * f),
24+ ];
25+ };
26+};
27+
28+// Control points sampled from matplotlib colormaps
29+const viridis = makeInterpolated([
30+ [68, 1, 84],
31+ [72, 40, 120],
32+ [62, 74, 137],
33+ [49, 104, 142],
34+ [38, 130, 142],
35+ [31, 158, 137],
36+ [53, 183, 121],
37+ [109, 205, 89],
38+ [180, 222, 44],
39+ [253, 231, 37],
40+]);
41+
42+const plasma = makeInterpolated([
43+ [13, 8, 135],
44+ [84, 2, 163],
45+ [139, 10, 165],
46+ [185, 50, 137],
47+ [219, 92, 104],
48+ [244, 136, 73],
49+ [254, 188, 43],
50+ [240, 249, 33],
51+]);
52+
53+const inferno = makeInterpolated([
54+ [0, 0, 4],
55+ [40, 11, 84],
56+ [101, 21, 110],
57+ [159, 42, 99],
58+ [212, 72, 66],
59+ [245, 125, 21],
60+ [250, 193, 39],
61+ [252, 255, 164],
62+]);
63+
64+const coolwarm = makeInterpolated([
65+ [59, 76, 192],
66+ [124, 159, 249],
67+ [192, 212, 245],
68+ [242, 242, 242],
69+ [245, 195, 157],
70+ [222, 96, 77],
71+ [180, 4, 38],
72+]);
73+
74+const jet = makeInterpolated([
75+ [0, 0, 128],
76+ [0, 0, 255],
77+ [0, 255, 255],
78+ [0, 255, 0],
79+ [255, 255, 0],
80+ [255, 0, 0],
81+ [128, 0, 0],
82+]);
83+
84+const grayscale: ColormapFunc = (t: number) => {
85+ const v = Math.round(clamp01(t) * 255);
86+ return [v, v, v];
87+};
88+
89+export const colormaps: Record<string, ColormapFunc> = {
90+ viridis,
91+ plasma,
92+ inferno,
93+ coolwarm,
94+ jet,
95+ grayscale,
96+};
97+
98+export const colormapNames = Object.keys(colormaps);
src/render/movie.tsadded+299−0View file
@@ -0,0 +1,299 @@
1+/**
2+ * MP4 recording of the live view.
3+ *
4+ * Each frame is composited from the on-screen sphere canvases, so the movie
5+ * shows what the page shows — current camera orientation, colormap, theme —
6+ * with a colorbar per species and a caption (model, parameter values, running
7+ * time). Encoding is WebCodecs H.264 muxed by mp4-muxer, entirely in the
8+ * browser, so capture runs as fast as the solver recomputes rather than at
9+ * playback speed.
10+ */
11+import { ArrayBufferTarget, Muxer } from 'mp4-muxer';
12+import type { ColormapFunc } from './colormaps.ts';
13+import { fmtValue } from './colorbar.ts';
14+
15+export interface MoviePanel {
16+ /**
17+ * The scene's WebGL canvas. It must be rendered in the same task that calls
18+ * addFrame(): without preserveDrawingBuffer the drawing buffer survives only
19+ * until the browser next composites.
20+ */
21+ canvas: HTMLCanvasElement;
22+ label: string;
23+}
24+
25+/** Per-panel colorbar state for one frame. */
26+export interface MovieBar {
27+ cmap: ColormapFunc;
28+ lo: number;
29+ hi: number;
30+}
31+
32+export interface MovieOptions {
33+ panels: MoviePanel[];
34+ /** Caption line 1, bold: the model/preset. */
35+ title: string;
36+ /** Caption line 2: the parameter values. */
37+ subtitle: string;
38+ /** Playback speed: simulation-time units per second of video. Each frame is
39+ * timestamped with its simulation time divided by this, so playback speed
40+ * is exact regardless of how many frames the caller captures. */
41+ speed: number;
42+ /** Effective frames per second, for encoder rate control only. */
43+ fps: number;
44+ /** Rendered edge of each sphere panel, px. The caller renders the scene
45+ * canvases at this size; the frame is the panels side by side plus the
46+ * caption bar. */
47+ sphere: number;
48+}
49+
50+/**
51+ * H.264 profile candidates: High, then Main, then Constrained Baseline —
52+ * Chrome's software fallback encoder supports only the last. The level covers
53+ * the frame area: 4.0 up to 1080p at 30 fps, 5.1 beyond (large exports).
54+ */
55+const h264Candidates = (pixels: number): string[] => {
56+ const level = pixels <= 1920 * 1080 ? '28' : '33';
57+ return ['avc1.6400', 'avc1.4d00', 'avc1.42e0'].map((p) => p + level);
58+};
59+
60+const even = (x: number): number => 2 * Math.floor(x / 2);
61+
62+interface Layout {
63+ /** Sphere panel edge, px. Everything else scales by u = sphere/768. */
64+ sphere: number;
65+ u: number;
66+ /** Colorbar column to the right of each sphere, like the app's. */
67+ gutter: number;
68+ captionH: number;
69+ width: number;
70+ height: number;
71+}
72+
73+/** Sized from the caller's resolution choice, clamped to what H.264 encoders
74+ * comfortably handle. Even dimensions, as 4:2:0 encoders require. */
75+const layoutFor = (nPanels: number, spherePx: number): Layout => {
76+ const sphere = even(Math.max(240, Math.min(1600, spherePx)));
77+ const u = sphere / 768;
78+ const gutter = even(Math.round(72 * u));
79+ const captionH = even(Math.round(64 * u));
80+ return {
81+ sphere,
82+ u,
83+ gutter,
84+ captionH,
85+ width: nPanels * (sphere + gutter),
86+ height: sphere + captionH,
87+ };
88+};
89+
90+export class MovieRecorder {
91+ #panels: MoviePanel[];
92+ #title: string;
93+ #subtitle: string;
94+ #speed: number;
95+ #fps: number;
96+ #lastKeyUs = 0;
97+ #layout: Layout;
98+ #canvas: HTMLCanvasElement;
99+ #ctx: CanvasRenderingContext2D;
100+ #muxer: Muxer<ArrayBufferTarget>;
101+ #encoder: VideoEncoder;
102+ #frames = 0;
103+ #error: unknown = null;
104+ // Page theme, sampled at creation so the movie matches light/dark mode.
105+ #bg: string;
106+ #ink: string;
107+ #ink2: string;
108+ #line: string;
109+ #sphereBg: string;
110+
111+ static async create(opts: MovieOptions): Promise<MovieRecorder> {
112+ if (typeof VideoEncoder === 'undefined') {
113+ throw new Error('WebCodecs is not available in this browser');
114+ }
115+ const layout = layoutFor(opts.panels.length, opts.sphere);
116+ const fps = Math.max(1, Math.round(opts.fps));
117+ const config = {
118+ width: layout.width,
119+ height: layout.height,
120+ // ~0.15 bits per pixel per frame reads as visually lossless here
121+ bitrate: Math.min(
122+ 24e6,
123+ Math.max(2e6, Math.round(layout.width * layout.height * fps * 0.15)),
124+ ),
125+ framerate: fps,
126+ };
127+ for (const codec of h264Candidates(layout.width * layout.height)) {
128+ const { supported } = await VideoEncoder.isConfigSupported({ codec, ...config });
129+ if (supported) return new MovieRecorder(opts, layout, { codec, ...config });
130+ }
131+ throw new Error('no supported H.264 encoder configuration');
132+ }
133+
134+ private constructor(opts: MovieOptions, layout: Layout, config: VideoEncoderConfig) {
135+ this.#panels = opts.panels;
136+ this.#title = opts.title;
137+ this.#subtitle = opts.subtitle;
138+ this.#speed = opts.speed;
139+ this.#fps = Math.max(1, opts.fps);
140+ this.#layout = layout;
141+
142+ const css = getComputedStyle(document.documentElement);
143+ const themeVar = (name: string, fallback: string): string =>
144+ css.getPropertyValue(name).trim() || fallback;
145+ this.#bg = themeVar('--bg', '#ffffff');
146+ this.#ink = themeVar('--ink', '#1f2328');
147+ this.#ink2 = themeVar('--ink-2', '#57606a');
148+ this.#line = themeVar('--line', '#d0d7de');
149+ this.#sphereBg = themeVar('--sphere-bg', '#f4f6f8');
150+
151+ this.#canvas = document.createElement('canvas');
152+ this.#canvas.width = layout.width;
153+ this.#canvas.height = layout.height;
154+ const ctx = this.#canvas.getContext('2d');
155+ if (!ctx) throw new Error('no 2d context for the movie canvas');
156+ this.#ctx = ctx;
157+
158+ this.#muxer = new Muxer({
159+ target: new ArrayBufferTarget(),
160+ video: {
161+ codec: 'avc',
162+ width: layout.width,
163+ height: layout.height,
164+ frameRate: Math.max(1, Math.round(opts.fps)),
165+ },
166+ fastStart: 'in-memory',
167+ });
168+ this.#encoder = new VideoEncoder({
169+ output: (chunk, meta) => this.#muxer.addVideoChunk(chunk, meta),
170+ error: (e) => (this.#error = e),
171+ });
172+ this.#encoder.configure(config);
173+ }
174+
175+ /**
176+ * Composite and encode one frame. The compositing happens synchronously, in
177+ * the caller's task; the await is only encoder backpressure, so a solver
178+ * that outruns the encoder does not pile frames up in its queue.
179+ */
180+ async addFrame(t: number, bars: MovieBar[]): Promise<void> {
181+ if (this.#error) throw this.#error;
182+ this.#compose(t, bars);
183+ const timestamp = Math.round((t / this.#speed) * 1e6);
184+ const frame = new VideoFrame(this.#canvas, {
185+ timestamp,
186+ duration: Math.round(1e6 / this.#fps),
187+ });
188+ // a keyframe every ~2 s of video keeps the file seekable without bloat
189+ const keyFrame = this.#frames === 0 || timestamp - this.#lastKeyUs >= 2e6;
190+ if (keyFrame) this.#lastKeyUs = timestamp;
191+ this.#encoder.encode(frame, { keyFrame });
192+ frame.close();
193+ this.#frames++;
194+ while (this.#encoder.encodeQueueSize > 4) {
195+ await new Promise((r) => this.#encoder.addEventListener('dequeue', r, { once: true }));
196+ }
197+ }
198+
199+ async finish(): Promise<Blob> {
200+ await this.#encoder.flush();
201+ if (this.#error) throw this.#error;
202+ this.#encoder.close();
203+ this.#muxer.finalize();
204+ return new Blob([this.#muxer.target.buffer], { type: 'video/mp4' });
205+ }
206+
207+ cancel(): void {
208+ if (this.#encoder.state !== 'closed') this.#encoder.close();
209+ }
210+
211+ // ---------------------------------------------------------------- drawing
212+ #compose(t: number, bars: MovieBar[]): void {
213+ const { sphere, gutter } = this.#layout;
214+ const ctx = this.#ctx;
215+ ctx.fillStyle = this.#bg;
216+ ctx.fillRect(0, 0, this.#layout.width, this.#layout.height);
217+ this.#panels.forEach((panel, k) => {
218+ const x = k * (sphere + gutter);
219+ ctx.drawImage(panel.canvas, x, 0, sphere, sphere);
220+ ctx.fillStyle = this.#sphereBg;
221+ ctx.fillRect(x + sphere, 0, gutter, sphere);
222+ this.#drawBar(x + sphere, bars[k]);
223+ this.#drawTag(x, panel.label);
224+ });
225+ this.#drawCaption(t);
226+ }
227+
228+ #drawBar(x0: number, bar: MovieBar): void {
229+ const { sphere, u, gutter } = this.#layout;
230+ const ctx = this.#ctx;
231+ const w = Math.round(18 * u);
232+ const h = Math.round(0.55 * sphere);
233+ const bx = Math.round(x0 + (gutter - w) / 2);
234+ const by = Math.round((sphere - h) / 2);
235+ for (let y = 0; y < h; y++) {
236+ const [r, g, b] = bar.cmap(1 - y / (h - 1));
237+ ctx.fillStyle = `rgb(${r},${g},${b})`;
238+ ctx.fillRect(bx, by + y, w, 1);
239+ }
240+ ctx.strokeStyle = this.#line;
241+ ctx.strokeRect(bx + 0.5, by + 0.5, w - 1, h - 1);
242+ ctx.fillStyle = this.#ink2;
243+ ctx.font = `${Math.round(13 * u)}px system-ui, sans-serif`;
244+ ctx.textAlign = 'center';
245+ ctx.textBaseline = 'bottom';
246+ ctx.fillText(fmtValue(bar.hi), x0 + gutter / 2, by - 6 * u);
247+ ctx.textBaseline = 'top';
248+ ctx.fillText(fmtValue(bar.lo), x0 + gutter / 2, by + h + 6 * u);
249+ }
250+
251+ /** The species name, as the app's floating tag: white on a dark pill. */
252+ #drawTag(x0: number, label: string): void {
253+ const { u } = this.#layout;
254+ const ctx = this.#ctx;
255+ const size = Math.round(20 * u);
256+ ctx.font = `600 ${size}px system-ui, sans-serif`;
257+ const tw = ctx.measureText(label).width;
258+ const padX = 12 * u;
259+ const padY = 4 * u;
260+ const bx = x0 + 12 * u;
261+ const by = 10 * u;
262+ const bh = size + 2 * padY;
263+ ctx.fillStyle = 'rgba(0, 0, 0, 0.45)';
264+ ctx.beginPath();
265+ ctx.roundRect(bx, by, tw + 2 * padX, bh, bh / 2);
266+ ctx.fill();
267+ ctx.fillStyle = '#fff';
268+ ctx.textAlign = 'left';
269+ ctx.textBaseline = 'middle';
270+ ctx.fillText(label, bx + padX, by + bh / 2 + u);
271+ }
272+
273+ #drawCaption(t: number): void {
274+ const { sphere, u, captionH, width } = this.#layout;
275+ const ctx = this.#ctx;
276+ const pad = 16 * u;
277+ ctx.strokeStyle = this.#line;
278+ ctx.beginPath();
279+ ctx.moveTo(0, sphere + 0.5);
280+ ctx.lineTo(width, sphere + 0.5);
281+ ctx.stroke();
282+ ctx.textBaseline = 'middle';
283+ ctx.fillStyle = this.#ink;
284+ ctx.font = `600 ${Math.round(20 * u)}px system-ui, sans-serif`;
285+ ctx.textAlign = 'left';
286+ ctx.fillText(this.#title, pad, sphere + captionH * 0.34);
287+ ctx.font = `${Math.round(20 * u)}px system-ui, sans-serif`;
288+ ctx.textAlign = 'right';
289+ ctx.fillText(
290+ `t = ${t.toFixed(2)} · ${this.#speed}×`,
291+ width - pad,
292+ sphere + captionH * 0.34,
293+ );
294+ ctx.fillStyle = this.#ink2;
295+ ctx.font = `${Math.round(14 * u)}px system-ui, sans-serif`;
296+ ctx.textAlign = 'left';
297+ ctx.fillText(this.#subtitle, pad, sphere + captionH * 0.74);
298+ }
299+}
src/render/sphereMesh.tsadded+227−0View file
@@ -0,0 +1,227 @@
1+/**
2+ * Mesh topology for a spherical (nlat, nphi) grid following shtns conventions:
3+ * - latitudinal grid given as cos(theta) (e.g. Gauss nodes, poles not included)
4+ * - phi equally spaced starting at 0, endpoint excluded
5+ *
6+ * The phi seam is stitched when the phi grid spans the full circle, and pole
7+ * cap vertices are added when the grid does not reach the poles, so that the
8+ * rendered surface is closed.
9+ *
10+ * Adapted from figpack's SphereEmbedding view (figpack_experimental).
11+ */
12+
13+import type { ColormapFunc } from './colormaps.ts';
14+
15+export type SphereMeshTopology = {
16+ nlat: number;
17+ nphi: number;
18+ wrapPhi: boolean;
19+ // Cap adjacent to row 0 / row nlat-1 (extra vertex appended after the grid)
20+ startCapIndex: number; // -1 if absent
21+ endCapIndex: number; // -1 if absent
22+ numVertices: number;
23+ indices: Uint32Array;
24+ // Unit-sphere positions, length numVertices * 3
25+ sphereRef: Float32Array;
26+};
27+
28+export const buildTopology = (
29+ cosTheta: Float64Array | Float32Array,
30+ phi: Float64Array | Float32Array,
31+): SphereMeshTopology => {
32+ const nlat = cosTheta.length;
33+ const nphi = phi.length;
34+
35+ // Does the phi grid span the full circle (so the seam should be stitched)?
36+ let wrapPhi = false;
37+ if (nphi >= 3) {
38+ const dphi = phi[1] - phi[0];
39+ const gap = phi[0] + 2 * Math.PI - phi[nphi - 1];
40+ wrapPhi = Math.abs(gap - dphi) < 0.25 * Math.abs(dphi);
41+ }
42+
43+ // Add pole caps where the grid does not reach the pole (|cos_theta| < 1),
44+ // only when the surface wraps in phi (otherwise there is no hole to close)
45+ const poleEps = 1e-9;
46+ const hasStartCap = wrapPhi && Math.abs(Math.abs(cosTheta[0]) - 1) > poleEps;
47+ const hasEndCap =
48+ wrapPhi && Math.abs(Math.abs(cosTheta[nlat - 1]) - 1) > poleEps;
49+
50+ const numGridVertices = nlat * nphi;
51+ let numVertices = numGridVertices;
52+ const startCapIndex = hasStartCap ? numVertices++ : -1;
53+ const endCapIndex = hasEndCap ? numVertices++ : -1;
54+
55+ const numCols = wrapPhi ? nphi : nphi - 1;
56+ let numTriangles = (nlat - 1) * numCols * 2;
57+ if (hasStartCap) numTriangles += nphi;
58+ if (hasEndCap) numTriangles += nphi;
59+
60+ const indices = new Uint32Array(numTriangles * 3);
61+ let k = 0;
62+ for (let i = 0; i < nlat - 1; i++) {
63+ for (let j = 0; j < numCols; j++) {
64+ const j2 = (j + 1) % nphi;
65+ const a = i * nphi + j;
66+ const b = i * nphi + j2;
67+ const c = (i + 1) * nphi + j;
68+ const d = (i + 1) * nphi + j2;
69+ indices[k++] = a;
70+ indices[k++] = c;
71+ indices[k++] = b;
72+ indices[k++] = b;
73+ indices[k++] = c;
74+ indices[k++] = d;
75+ }
76+ }
77+ if (hasStartCap) {
78+ for (let j = 0; j < nphi; j++) {
79+ const j2 = (j + 1) % nphi;
80+ indices[k++] = startCapIndex;
81+ indices[k++] = j;
82+ indices[k++] = j2;
83+ }
84+ }
85+ if (hasEndCap) {
86+ const rowOffset = (nlat - 1) * nphi;
87+ for (let j = 0; j < nphi; j++) {
88+ const j2 = (j + 1) % nphi;
89+ indices[k++] = rowOffset + j;
90+ indices[k++] = endCapIndex;
91+ indices[k++] = rowOffset + j2;
92+ }
93+ }
94+
95+ // Unit-sphere positions (z along the polar axis)
96+ const sphereRef = new Float32Array(numVertices * 3);
97+ for (let i = 0; i < nlat; i++) {
98+ const ct = cosTheta[i];
99+ const st = Math.sqrt(Math.max(0, 1 - ct * ct));
100+ for (let j = 0; j < nphi; j++) {
101+ const p = (i * nphi + j) * 3;
102+ sphereRef[p] = st * Math.cos(phi[j]);
103+ sphereRef[p + 1] = st * Math.sin(phi[j]);
104+ sphereRef[p + 2] = ct;
105+ }
106+ }
107+ if (hasStartCap) {
108+ const p = startCapIndex * 3;
109+ sphereRef[p + 2] = cosTheta[0] >= 0 ? 1 : -1;
110+ }
111+ if (hasEndCap) {
112+ const p = endCapIndex * 3;
113+ sphereRef[p + 2] = cosTheta[nlat - 1] >= 0 ? 1 : -1;
114+ }
115+
116+ return {
117+ nlat,
118+ nphi,
119+ wrapPhi,
120+ startCapIndex,
121+ endCapIndex,
122+ numVertices,
123+ indices,
124+ sphereRef,
125+ };
126+};
127+
128+/**
129+ * Fill the position buffer (numVertices * 3) from the surface's coordinates
130+ * (nlat * nphi * 3), interpolating toward the reference unit sphere.
131+ * morph = 1 gives the surface itself; morph = 0 pulls it back to the sphere,
132+ * which is the mesh the solver's parametrization actually lives on. Sweeping
133+ * between them shows which points went where.
134+ *
135+ * The pole caps are not on the grid, so they take the mean of the adjacent
136+ * ring — for a surface that is smooth at the pole, where the ring is a small
137+ * circle around it, that is the pole to the accuracy the ring resolves.
138+ */
139+export const fillPositions = (
140+ out: Float32Array,
141+ coords: Float32Array | Float64Array,
142+ topo: SphereMeshTopology,
143+ morph: number,
144+): void => {
145+ const { nlat, nphi, sphereRef } = topo;
146+ const n = nlat * nphi * 3;
147+ for (let p = 0; p < n; p++) {
148+ out[p] = (1 - morph) * sphereRef[p] + morph * coords[p];
149+ }
150+ const fillCap = (capIndex: number, rowIndex: number): void => {
151+ let x = 0;
152+ let y = 0;
153+ let z = 0;
154+ const rowOffset = rowIndex * nphi * 3;
155+ for (let j = 0; j < nphi; j++) {
156+ x += coords[rowOffset + j * 3];
157+ y += coords[rowOffset + j * 3 + 1];
158+ z += coords[rowOffset + j * 3 + 2];
159+ }
160+ const p = capIndex * 3;
161+ out[p] = (1 - morph) * sphereRef[p] + (morph * x) / nphi;
162+ out[p + 1] = (1 - morph) * sphereRef[p + 1] + (morph * y) / nphi;
163+ out[p + 2] = (1 - morph) * sphereRef[p + 2] + (morph * z) / nphi;
164+ };
165+ if (topo.startCapIndex >= 0) fillCap(topo.startCapIndex, 0);
166+ if (topo.endCapIndex >= 0) fillCap(topo.endCapIndex, nlat - 1);
167+};
168+
169+/**
170+ * Expand a field frame (nlat * nphi) to per-vertex values (numVertices),
171+ * with cap values averaged from the adjacent ring.
172+ */
173+export const fillFieldValues = (
174+ out: Float32Array,
175+ fieldFrame: Float32Array | Float64Array,
176+ topo: SphereMeshTopology,
177+): void => {
178+ const { nlat, nphi } = topo;
179+ const n = nlat * nphi;
180+ for (let p = 0; p < n; p++) {
181+ out[p] = fieldFrame[p];
182+ }
183+ const ringMean = (rowIndex: number) => {
184+ let sum = 0;
185+ let count = 0;
186+ for (let j = 0; j < nphi; j++) {
187+ const v = fieldFrame[rowIndex * nphi + j];
188+ if (!Number.isNaN(v)) {
189+ sum += v;
190+ count++;
191+ }
192+ }
193+ return count > 0 ? sum / count : NaN;
194+ };
195+ if (topo.startCapIndex >= 0) out[topo.startCapIndex] = ringMean(0);
196+ if (topo.endCapIndex >= 0) out[topo.endCapIndex] = ringMean(nlat - 1);
197+};
198+
199+/**
200+ * Fill the color buffer (numVertices * 3, floats in [0, 1]) from per-vertex
201+ * field values using the given colormap and range. NaN values render gray.
202+ */
203+export const fillColors = (
204+ out: Float32Array,
205+ values: Float32Array,
206+ valueMin: number,
207+ valueMax: number,
208+ cmap: ColormapFunc,
209+): void => {
210+ const span = valueMax - valueMin;
211+ const invSpan = span !== 0 ? 1 / span : 0;
212+ for (let i = 0; i < values.length; i++) {
213+ const v = values[i];
214+ const p = i * 3;
215+ if (Number.isNaN(v)) {
216+ out[p] = 0.35;
217+ out[p + 1] = 0.35;
218+ out[p + 2] = 0.35;
219+ } else {
220+ const t = span !== 0 ? (v - valueMin) * invSpan : 0.5;
221+ const [r, g, b] = cmap(t);
222+ out[p] = r / 255;
223+ out[p + 1] = g / 255;
224+ out[p + 2] = b / 255;
225+ }
226+ }
227+};
src/sht/coeffs.tsadded+85−0View file
@@ -0,0 +1,85 @@
1+/**
2+ * Recurrence coefficients for orthonormal associated Legendre functions
3+ * ytilde_l^m(theta) (spherical-harmonic normalized, Condon-Shortley phase
4+ * included), matching SHTNS legendre_precomp() with norm=sht_orthonormal:
5+ *
6+ * ytilde_m^m(theta) = amm * sin(theta)^m
7+ * ytilde_{m+1}^m = a_{m+1}^m * cos(theta) * ytilde_m^m
8+ * ytilde_l^m = a_l^m * cos(theta) * ytilde_{l-1}^m + b_l^m * ytilde_{l-2}^m
9+ *
10+ * with (cf. sht_legendre.c lines 442-447):
11+ * a_{m+1}^m = sqrt(2m+3)
12+ * a_l^m = sqrt( (2l+1)(2l-1) / ((l+m)(l-m)) )
13+ * b_l^m = -sqrt( (2l+1)/(2l-3) * ((l-1+m)(l-1-m)) / ((l+m)(l-m)) )
14+ * amm = cs^m * sqrt( 1/(4pi) * prod_{k=1..m} (2k+1)/(2k) )
15+ *
16+ * With this normalization, Y_lm(theta,phi) = ytilde_l^m(theta) e^{i m phi}
17+ * and integral |Y_lm|^2 dOmega = 1.
18+ */
19+import { lmIndex, nlmCalc } from './layout.ts';
20+
21+export interface LegendreCoeffs {
22+ /** amm[m]: seed value (includes Condon-Shortley phase (-1)^m). */
23+ amm: Float64Array;
24+ /** ab[2*lm], ab[2*lm+1] = (a_l^m, b_l^m); entries at l=m unused (0), b at l=m+1 unused (0). */
25+ ab: Float64Array;
26+}
27+
28+export function legendreCoeffs(lmax: number, mmax: number): LegendreCoeffs {
29+ const nlm = nlmCalc(lmax, mmax);
30+ const amm = new Float64Array(mmax + 1);
31+ const ab = new Float64Array(2 * nlm);
32+
33+ let t = 1.0 / (4.0 * Math.PI);
34+ amm[0] = Math.sqrt(t);
35+ for (let m = 1; m <= mmax; m++) {
36+ t *= (2 * m + 1) / (2 * m);
37+ amm[m] = -Math.sqrt(t); // (-1)^m accumulates: Condon-Shortley phase
38+ if (m % 2 === 0) amm[m] = -amm[m];
39+ }
40+
41+ for (let m = 0; m <= mmax; m++) {
42+ if (m + 1 <= lmax) {
43+ const lm = lmIndex(lmax, m + 1, m);
44+ ab[2 * lm] = Math.sqrt(2 * m + 3); // a_{m+1}^m
45+ ab[2 * lm + 1] = 0;
46+ }
47+ for (let l = m + 2; l <= lmax; l++) {
48+ const lm = lmIndex(lmax, l, m);
49+ const t1 = (l + m) * (l - m);
50+ const t2 = (l - 1 + m) * (l - 1 - m);
51+ ab[2 * lm] = Math.sqrt(((2 * l + 1) * (2 * l - 1)) / t1);
52+ ab[2 * lm + 1] = -Math.sqrt(((2 * l + 1) / (2 * l - 3)) * (t2 / t1));
53+ }
54+ }
55+ return { amm, ab };
56+}
57+
58+/**
59+ * Evaluate ytilde_l^m(theta) for l = m..lmax at one point, in f64.
60+ * ct = cos(theta), st = sin(theta). Plain (unscaled) recurrence: fine in
61+ * f64 for the moderate lmax this library targets (underflow of st^m only
62+ * matters for m of several hundred very close to the poles).
63+ */
64+export function legendreRow(
65+ coeffs: LegendreCoeffs,
66+ lmax: number,
67+ m: number,
68+ ct: number,
69+ st: number,
70+ out: Float64Array, // length lmax - m + 1
71+): void {
72+ let y0 = coeffs.amm[m] * Math.pow(st, m);
73+ out[0] = y0;
74+ if (m === lmax) return;
75+ const base = lmIndex(lmax, m, m);
76+ let y1 = coeffs.ab[2 * (base + 1)] * ct * y0;
77+ out[1] = y1;
78+ for (let l = m + 2; l <= lmax; l++) {
79+ const lm = base + (l - m);
80+ const y2 = coeffs.ab[2 * lm] * ct * y1 + coeffs.ab[2 * lm + 1] * y0;
81+ y0 = y1;
82+ y1 = y2;
83+ out[l - m] = y2;
84+ }
85+}
src/sht/gauss.tsadded+51−0View file
@@ -0,0 +1,51 @@
1+/**
2+ * Gauss-Legendre quadrature nodes and weights, computed in double
3+ * precision by Newton iteration on P_n (cf. gauss_nodes() in SHTNS
4+ * sht_legendre.c).
5+ *
6+ * Returns nodes x_i = cos(theta_i) in DECREASING order (theta increasing,
7+ * north pole first), and weights w_i for integration over x in [-1, 1]:
8+ * integral f(x) dx ~= sum_i w_i f(x_i), exact for polynomials of
9+ * degree <= 2n - 1.
10+ */
11+export function gaussNodesWeights(n: number): { x: Float64Array; w: Float64Array } {
12+ const x = new Float64Array(n);
13+ const w = new Float64Array(n);
14+ const m = (n + 1) >> 1;
15+ for (let i = 0; i < m; i++) {
16+ // initial guess (Tricomi-like), then Newton
17+ let z = Math.cos((Math.PI * (i + 0.75)) / (n + 0.5));
18+ let pp = 0;
19+ for (let iter = 0; iter < 100; iter++) {
20+ // evaluate P_n(z) and P_{n-1}(z) by recurrence
21+ let p1 = 1.0;
22+ let p2 = 0.0;
23+ for (let j = 1; j <= n; j++) {
24+ const p3 = p2;
25+ p2 = p1;
26+ p1 = ((2 * j - 1) * z * p2 - (j - 1) * p3) / j;
27+ }
28+ pp = (n * (z * p1 - p2)) / (z * z - 1.0);
29+ const dz = p1 / pp;
30+ z -= dz;
31+ if (Math.abs(dz) < 1e-15 * Math.abs(z) + 1e-300) {
32+ // one extra iteration for full convergence
33+ let q1 = 1.0, q2 = 0.0;
34+ for (let j = 1; j <= n; j++) {
35+ const q3 = q2; q2 = q1;
36+ q1 = ((2 * j - 1) * z * q2 - (j - 1) * q3) / j;
37+ }
38+ pp = (n * (z * q1 - q2)) / (z * z - 1.0);
39+ z -= q1 / pp;
40+ break;
41+ }
42+ }
43+ x[i] = z; // largest roots first => theta increasing
44+ x[n - 1 - i] = -z;
45+ const wi = 2.0 / ((1.0 - z * z) * pp * pp);
46+ w[i] = wi;
47+ w[n - 1 - i] = wi;
48+ }
49+ if (n & 1) x[m - 1] = 0.0; // exact for odd n
50+ return { x, w };
51+}
src/sht/layout.tsadded+59−0View file
@@ -0,0 +1,59 @@
1+/**
2+ * Grid and spectral layout definitions, following SHTNS conventions:
3+ *
4+ * - Spectral coefficients Q_lm are complex, stored for m >= 0 only (real
5+ * fields), interleaved [re, im], with SHTNS "m-major" ordering:
6+ * for m = 0..mmax: for l = m..lmax. Index of (l, m) is lm(l, m).
7+ * - Spatial fields are real, phi-contiguous: spat[ilat * nphi + iphi],
8+ * with ilat ordered by increasing colatitude theta (north to south)
9+ * and iphi covering [0, 2*pi) uniformly.
10+ * - Normalization: orthonormal spherical harmonics INCLUDING the
11+ * Condon-Shortley phase (SHTNS default: sht_orthonormal).
12+ * A real field is f = sum_{l,m>=0} Q_lm Y_lm + c.c.(m>0), i.e.
13+ * Q_{l,-m} = (-1)^m conj(Q_lm) is implied. m=0 coefficients must
14+ * have zero imaginary part.
15+ */
16+
17+export interface ShtConfig {
18+ lmax: number;
19+ mmax: number;
20+ nlat: number;
21+ nphi: number;
22+}
23+
24+export function nlmCalc(lmax: number, mmax: number): number {
25+ // sum over m=0..mmax of (lmax - m + 1)
26+ return (mmax + 1) * (lmax + 1) - (mmax * (mmax + 1)) / 2;
27+}
28+
29+/** Index of coefficient (l, m) in the spectral array (SHTNS LM ordering). */
30+export function lmIndex(lmax: number, l: number, m: number): number {
31+ return m * (lmax + 1) - (m * (m - 1)) / 2 + (l - m);
32+}
33+
34+export function validateConfig(cfg: ShtConfig): void {
35+ const { lmax, mmax, nlat, nphi } = cfg;
36+ if (!Number.isInteger(lmax) || lmax < 1) throw new Error(`lmax must be an integer >= 1 (got ${lmax})`);
37+ if (!Number.isInteger(mmax) || mmax < 0 || mmax > lmax)
38+ throw new Error(`mmax must be an integer in [0, lmax] (got ${mmax})`);
39+ if (!Number.isInteger(nlat) || nlat <= lmax)
40+ throw new Error(`nlat must be an integer > lmax for exact Gauss quadrature (got nlat=${nlat}, lmax=${lmax})`);
41+ if (!Number.isInteger(nphi) || nphi < 2 * mmax + 1)
42+ throw new Error(`nphi must be an integer >= 2*mmax+1 to avoid aliasing (got nphi=${nphi}, mmax=${mmax})`);
43+}
44+
45+export function isPowerOfTwo(n: number): boolean {
46+ return n > 0 && (n & (n - 1)) === 0;
47+}
48+
49+/** Grid sizes for a given lmax, dealiased for a reaction of polynomial degree
50+ * `pdeg` (the rule from websph's reference implementation):
51+ * nlat >= ((pdeg+1)*lmax+1)/2, nphi >= (pdeg+1)*lmax+1. nphi is rounded up to
52+ * a power of two to keep the GPU FFT path. */
53+export function gridForLmax(lmax: number, pdeg: number): { nlat: number; nphi: number } {
54+ const minLat = Math.max(lmax + 1, ((pdeg + 1) * lmax + 1) / 2);
55+ const nlat = 2 * Math.ceil(minLat / 2);
56+ let nphi = 1;
57+ while (nphi < (pdeg + 1) * lmax + 1) nphi *= 2;
58+ return { nlat, nphi };
59+}
src/sht/reference.tsadded+136−0View file
@@ -0,0 +1,136 @@
1+/**
2+ * Double-precision reference implementation of the scalar spherical
3+ * harmonic transform, by direct summation. Slow (O(nlat*nlm) Legendre +
4+ * O(nlat*nphi*mmax) Fourier) but simple, and serves as ground truth for
5+ * validating the fp32 WebGPU implementation.
6+ *
7+ * Conventions are identical to the GPU path (see layout.ts).
8+ */
9+import { gaussNodesWeights } from './gauss.ts';
10+import { legendreCoeffs, legendreRow, type LegendreCoeffs } from './coeffs.ts';
11+import { lmIndex, nlmCalc, validateConfig, type ShtConfig } from './layout.ts';
12+
13+export class ShtReference {
14+ readonly cfg: ShtConfig;
15+ readonly nlm: number;
16+ readonly ct: Float64Array;
17+ readonly st: Float64Array;
18+ readonly wg: Float64Array; // Gauss weights (for integral over cos(theta))
19+ readonly coeffs: LegendreCoeffs;
20+
21+ constructor(cfg: ShtConfig) {
22+ validateConfig(cfg);
23+ this.cfg = cfg;
24+ this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
25+ const { x, w } = gaussNodesWeights(cfg.nlat);
26+ this.ct = x;
27+ this.wg = w;
28+ this.st = new Float64Array(cfg.nlat);
29+ for (let i = 0; i < cfg.nlat; i++) this.st[i] = Math.sqrt(1 - x[i] * x[i]);
30+ this.coeffs = legendreCoeffs(cfg.lmax, cfg.mmax);
31+ }
32+
33+ /**
34+ * Legendre stage of the synthesis: F_m(theta_i) = sum_l Q_lm ytilde_l^m(theta_i).
35+ * Returns complex array indexed [m * nlat + ilat], interleaved re/im.
36+ */
37+ legendreSynth(qlm: ArrayLike<number>): Float64Array {
38+ const { lmax, mmax, nlat } = this.cfg;
39+ const fm = new Float64Array(2 * (mmax + 1) * nlat);
40+ const row = new Float64Array(lmax + 1);
41+ for (let i = 0; i < nlat; i++) {
42+ for (let m = 0; m <= mmax; m++) {
43+ legendreRow(this.coeffs, lmax, m, this.ct[i], this.st[i], row);
44+ let re = 0, im = 0;
45+ const base = lmIndex(lmax, m, m);
46+ for (let l = m; l <= lmax; l++) {
47+ const y = row[l - m];
48+ re += y * qlm[2 * (base + l - m)];
49+ im += y * qlm[2 * (base + l - m) + 1];
50+ }
51+ const o = 2 * (m * nlat + i);
52+ fm[o] = re;
53+ fm[o + 1] = im;
54+ }
55+ }
56+ return fm;
57+ }
58+
59+ /** Full synthesis: spectral -> spatial grid [ilat * nphi + iphi]. */
60+ synth(qlm: ArrayLike<number>): Float64Array {
61+ const { mmax, nlat, nphi } = this.cfg;
62+ const fm = this.legendreSynth(qlm);
63+ const spat = new Float64Array(nlat * nphi);
64+ for (let i = 0; i < nlat; i++) {
65+ for (let j = 0; j < nphi; j++) {
66+ const phi = (2 * Math.PI * j) / nphi;
67+ let v = fm[2 * (0 * nlat + i)]; // m=0: real part (imag must be 0)
68+ for (let m = 1; m <= mmax; m++) {
69+ const o = 2 * (m * nlat + i);
70+ const c = Math.cos(m * phi);
71+ const s = Math.sin(m * phi);
72+ v += 2 * (fm[o] * c - fm[o + 1] * s);
73+ }
74+ spat[i * nphi + j] = v;
75+ }
76+ }
77+ return spat;
78+ }
79+
80+ /** Full analysis: spatial grid -> spectral coefficients (interleaved re/im). */
81+ analys(spat: ArrayLike<number>): Float64Array {
82+ const { lmax, mmax, nlat, nphi } = this.cfg;
83+ const qlm = new Float64Array(2 * this.nlm);
84+ const row = new Float64Array(lmax + 1);
85+ // forward Fourier: G_m(theta_i) = (2*pi/nphi) * sum_j f_ij e^{-i m phi_j}
86+ const gm = new Float64Array(2 * (mmax + 1) * nlat);
87+ for (let i = 0; i < nlat; i++) {
88+ for (let m = 0; m <= mmax; m++) {
89+ let re = 0, im = 0;
90+ for (let j = 0; j < nphi; j++) {
91+ const phi = (2 * Math.PI * j) / nphi;
92+ const f = spat[i * nphi + j];
93+ re += f * Math.cos(m * phi);
94+ im -= f * Math.sin(m * phi);
95+ }
96+ const o = 2 * (m * nlat + i);
97+ const norm = (2 * Math.PI) / nphi;
98+ gm[o] = re * norm;
99+ gm[o + 1] = im * norm;
100+ }
101+ }
102+ // Legendre stage with Gauss quadrature: Q_lm = sum_i w_i ytilde_l^m(theta_i) G_m(theta_i)
103+ for (let m = 0; m <= mmax; m++) {
104+ const base = lmIndex(lmax, m, m);
105+ for (let i = 0; i < nlat; i++) {
106+ legendreRow(this.coeffs, lmax, m, this.ct[i], this.st[i], row);
107+ const o = 2 * (m * nlat + i);
108+ const wr = this.wg[i] * gm[o];
109+ const wi = this.wg[i] * gm[o + 1];
110+ for (let l = m; l <= lmax; l++) {
111+ const y = row[l - m];
112+ qlm[2 * (base + l - m)] += y * wr;
113+ qlm[2 * (base + l - m) + 1] += y * wi;
114+ }
115+ }
116+ }
117+ return qlm;
118+ }
119+}
120+
121+/** Random band-limited spectrum for testing (m=0 imaginary parts zeroed). */
122+export function randomSpectrum(cfg: ShtConfig, seed = 12345): Float32Array {
123+ const nlm = nlmCalc(cfg.lmax, cfg.mmax);
124+ const q = new Float32Array(2 * nlm);
125+ let s = seed >>> 0;
126+ const rnd = () => {
127+ // xorshift32
128+ s ^= s << 13; s >>>= 0;
129+ s ^= s >> 17;
130+ s ^= s << 5; s >>>= 0;
131+ return (s / 4294967296) * 2 - 1;
132+ };
133+ for (let k = 0; k < 2 * nlm; k++) q[k] = rnd();
134+ for (let l = 0; l <= cfg.lmax; l++) q[2 * lmIndex(cfg.lmax, l, 0) + 1] = 0; // m=0 real
135+ return q;
136+}
src/sht/sht.tsadded+513−0View file
@@ -0,0 +1,513 @@
1+/**
2+ * WebGPU spherical harmonic transform plan (scalar transforms, fp32).
3+ *
4+ * Mirrors the structure of the SHTNS CUDA backend (sht_gpu.cu):
5+ * host-side f64 precomputation of grid + recurrence coefficients, shader
6+ * source generated with sizes baked in (SHTNS uses NVRTC; WGSL is always
7+ * runtime-compiled), then per-transform: Legendre stage + Fourier stage.
8+ */
9+import { gaussNodesWeights } from './gauss.ts';
10+import { legendreCoeffs } from './coeffs.ts';
11+import { nlmCalc, validateConfig, isPowerOfTwo, type ShtConfig } from './layout.ts';
12+import { legSynthWGSL, legAnalysWGSL } from './wgsl/leg.ts';
13+import {
14+ fftSynthWGSL,
15+ fftAnalysWGSL,
16+ fftSynthRealWGSL,
17+ fftAnalysRealWGSL,
18+ dftSynthWGSL,
19+ dftAnalysWGSL,
20+ fftThreads,
21+} from './wgsl/fourier.ts';
22+
23+export type FourierMode = 'auto' | 'fft' | 'dft';
24+
25+/** The two bind groups (Legendre stage, Fourier stage) of one transform. */
26+export interface ShtBinding {
27+ readonly bgLeg: GPUBindGroup;
28+ readonly bgFour: GPUBindGroup;
29+}
30+
31+const bgEntries = (bufs: GPUBuffer[]) =>
32+ bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
33+
34+export interface ShtOptions {
35+ /** Fourier stage implementation. 'auto' picks fft when nphi is a power of two that fits in workgroup memory. */
36+ fourier?: FourierMode;
37+}
38+
39+const WG_SYNTH = 64;
40+
41+/**
42+ * Tuning knob, for A/B-ing a change without editing code. Reads globalThis
43+ * first (set it before creating a plan, as scripts/_ab.ts does), then the
44+ * environment, so `SHT_SUBGROUPS=0 npm run bench:sht` works too. `process` is
45+ * absent in the browser, where only the globalThis form applies.
46+ */
47+function tuning(name: string): unknown {
48+ const g = (globalThis as Record<string, unknown>)[name];
49+ if (g !== undefined) return g;
50+ const env = (globalThis as { process?: { env?: Record<string, string | undefined> } }).process?.env?.[name];
51+ if (env === undefined || env === '') return undefined;
52+ if (env === '1' || env === 'true') return true;
53+ if (env === '0' || env === 'false') return false;
54+ const n = Number(env);
55+ return Number.isFinite(n) ? n : env;
56+}
57+
58+/**
59+ * Workgroup size for the analysis Legendre reduction. The right answer differs
60+ * between the two reduction strategies, so it is chosen per strategy.
61+ *
62+ * Measured on an RTX PRO 6000 Blackwell (analysis, us). Shared-memory tree,
63+ * where each doubling of wgAnalys costs another barrier per l-pair:
64+ *
65+ * wgAnalys: 16 32 64 128 256
66+ * nlat=128 43.2 40.6 42.6 46.7 52.9 -> 32
67+ * nlat=256 104.0 85.6 87.8 89.9 100.0 -> 32
68+ * nlat=512 408.8 216.9 184.4 190.8 204.8 -> 64
69+ *
70+ * i.e. max(32, nlat/8). A flat 32 would be worse than the old default of 256 at
71+ * nlat=512, so it cannot be fitted on one grid. With subgroupAdd the barrier
72+ * count stops growing with wgAnalys and the picture inverts: threads in flight,
73+ * (mmax+1) * wgAnalys, becomes binding, since analysis dispatches only mmax+1
74+ * workgroups. 128 then wins at every grid (round trip, us):
75+ *
76+ * 128x256 37.8 (vs 38.3), 256x512 66.6 (vs 74.7), 512x1024 131.9 (vs 156.6)
77+ */
78+function defaultWgAnalys(nlat: number, limit: number, subgroups: boolean): number {
79+ if (subgroups) {
80+ // capped at nlat so small grids do not launch threads with no latitude to own
81+ let cap = 1;
82+ while (cap < nlat) cap *= 2;
83+ return Math.min(128, limit, cap);
84+ }
85+ const target = Math.max(32, nlat / 8);
86+ let wg = 1;
87+ while (wg < target) wg *= 2; // the tree reduction halves, so a power of two
88+ return Math.min(wg, limit);
89+}
90+
91+async function makePipeline(
92+ device: GPUDevice,
93+ code: string,
94+ entryPoint: string,
95+): Promise<GPUComputePipeline> {
96+ device.pushErrorScope('validation');
97+ const module = device.createShaderModule({ code, label: entryPoint });
98+ const info = await module.getCompilationInfo();
99+ const errors = info.messages.filter((m) => m.type === 'error');
100+ if (errors.length) {
101+ throw new Error(
102+ `WGSL compile error in ${entryPoint}:\n` +
103+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
104+ );
105+ }
106+ const pipeline = await device.createComputePipelineAsync({
107+ layout: 'auto',
108+ compute: { module, entryPoint },
109+ label: entryPoint,
110+ });
111+ const err = await device.popErrorScope();
112+ if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
113+ return pipeline;
114+}
115+
116+export class ShtPlan {
117+ readonly cfg: ShtConfig;
118+ readonly nlm: number;
119+ readonly fourierMode: 'fft' | 'dft';
120+ /** Latitudes leg_synth walks: nlat/2 when parity folding. */
121+ readonly legLat: number = 0;
122+ /** Colatitudes theta_i (f64, increasing: north to south). */
123+ readonly theta: Float64Array;
124+ readonly cosTheta: Float64Array;
125+ readonly gaussWeights: Float64Array;
126+
127+ private device: GPUDevice;
128+ private bufAb!: GPUBuffer;
129+ private bufAmm!: GPUBuffer;
130+ private bufCtstw!: GPUBuffer;
131+ private bufTrig!: GPUBuffer;
132+ /** Spectral input (synthesis) — write with queue.writeBuffer or use synth(). */
133+ readonly qlmIn!: GPUBuffer;
134+ /** Spectral output (analysis). */
135+ readonly qlmOut!: GPUBuffer;
136+ /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. COPY_SRC so the
137+ * stage boundary is observable: a transform is Legendre-then-Fourier, and
138+ * scripts/diagnose-sht.ts tells the two apart by reading this. */
139+ readonly fmBuf!: GPUBuffer;
140+ /** Spatial field [ilat*nphi + iphi], f32. */
141+ readonly spatBuf!: GPUBuffer;
142+ private stageSpat!: GPUBuffer;
143+ private stageQ!: GPUBuffer;
144+
145+ private pipeLegSynth!: GPUComputePipeline;
146+ private pipeLegAnalys!: GPUComputePipeline;
147+ private pipeFourSynth!: GPUComputePipeline;
148+ private pipeFourAnalys!: GPUComputePipeline;
149+ private bgLegSynth!: GPUBindGroup;
150+ private bgLegAnalys!: GPUBindGroup;
151+ private bgFourSynth!: GPUBindGroup;
152+ private bgFourAnalys!: GPUBindGroup;
153+
154+ private constructor(device: GPUDevice, cfg: ShtConfig, fourierMode: 'fft' | 'dft') {
155+ this.device = device;
156+ this.cfg = cfg;
157+ this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
158+ this.fourierMode = fourierMode;
159+ const { x, w } = gaussNodesWeights(cfg.nlat);
160+ this.cosTheta = x;
161+ this.gaussWeights = w;
162+ this.theta = new Float64Array(cfg.nlat);
163+ for (let i = 0; i < cfg.nlat; i++) this.theta[i] = Math.acos(x[i]);
164+ }
165+
166+ static async create(device: GPUDevice, cfg: ShtConfig, opts: ShtOptions = {}): Promise<ShtPlan> {
167+ validateConfig(cfg);
168+ const want = opts.fourier ?? 'auto';
169+ const fftFits =
170+ isPowerOfTwo(cfg.nphi) &&
171+ 16 * cfg.nphi <= device.limits.maxComputeWorkgroupStorageSize &&
172+ fftThreads(cfg.nphi) <= device.limits.maxComputeInvocationsPerWorkgroup;
173+ if (want === 'fft' && !fftFits) {
174+ throw new Error(
175+ `fourier:'fft' requires power-of-two nphi with 16*nphi <= maxComputeWorkgroupStorageSize ` +
176+ `(nphi=${cfg.nphi}, limit=${device.limits.maxComputeWorkgroupStorageSize})`,
177+ );
178+ }
179+ const mode: 'fft' | 'dft' = want === 'dft' ? 'dft' : fftFits ? 'fft' : 'dft';
180+ const plan = new ShtPlan(device, cfg, mode);
181+ await plan.init();
182+ return plan;
183+ }
184+
185+ private async init(): Promise<void> {
186+ const { lmax, mmax, nlat, nphi } = this.cfg;
187+ const dev = this.device;
188+ const self = this as {
189+ -readonly [k in keyof ShtPlan]: ShtPlan[k];
190+ };
191+
192+ // --- host precomputation (f64), then downcast to f32 for upload ---
193+ const { amm, ab } = legendreCoeffs(lmax, mmax);
194+ const ctstw = new Float32Array(3 * nlat);
195+ for (let i = 0; i < nlat; i++) {
196+ ctstw[i] = this.cosTheta[i];
197+ ctstw[nlat + i] = Math.sqrt(1 - this.cosTheta[i] * this.cosTheta[i]);
198+ ctstw[2 * nlat + i] = this.gaussWeights[i] * ((2 * Math.PI) / nphi);
199+ }
200+ // twiddle/phase table in f64 (device sin/cos is too inaccurate: ~2^-11 under Vulkan)
201+ const trig = new Float32Array(2 * nphi);
202+ for (let k = 0; k < nphi; k++) {
203+ trig[2 * k] = Math.cos((2 * Math.PI * k) / nphi);
204+ trig[2 * k + 1] = Math.sin((2 * Math.PI * k) / nphi);
205+ }
206+
207+ const mkBuf = (label: string, size: number, usage: GPUBufferUsageFlags) =>
208+ dev.createBuffer({ label, size, usage });
209+ this.bufAb = mkBuf('sht-ab', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
210+ this.bufAmm = mkBuf('sht-amm', 4 * (mmax + 1), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
211+ this.bufCtstw = mkBuf('sht-ctstw', 4 * 3 * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
212+ this.bufTrig = mkBuf('sht-trig', 8 * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
213+ self.qlmIn = mkBuf('sht-qlm-in', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
214+ self.qlmOut = mkBuf('sht-qlm-out', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
215+ self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
216+ self.spatBuf = mkBuf('sht-spat', 4 * nlat * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
217+ this.stageSpat = mkBuf('sht-stage-spat', 4 * nlat * nphi, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
218+ this.stageQ = mkBuf('sht-stage-q', 8 * this.nlm, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
219+
220+ dev.queue.writeBuffer(this.bufAb, 0, new Float32Array(ab));
221+ dev.queue.writeBuffer(this.bufAmm, 0, new Float32Array(amm));
222+ dev.queue.writeBuffer(this.bufCtstw, 0, ctstw);
223+ dev.queue.writeBuffer(this.bufTrig, 0, trig);
224+
225+ // --- shaders / pipelines ---
226+ const subgroups = tuning('SHT_SUBGROUPS') !== false && dev.features.has('subgroups');
227+ // parity folding needs an equator-symmetric grid; Gauss nodes are, if nlat is even
228+ const parity = tuning('SHT_PARITY') !== false && nlat % 2 === 0;
229+ const wgAnalys =
230+ (tuning('SHT_WG_ANALYS') as number | undefined) ??
231+ defaultWgAnalys(nlat, dev.limits.maxComputeInvocationsPerWorkgroup, subgroups);
232+ const legP = {
233+ lmax,
234+ mmax,
235+ nlat,
236+ wgSynth: WG_SYNTH,
237+ wgAnalys,
238+ subgroups,
239+ spanPairs: tuning('SHT_SPAN_PAIRS') as number | undefined,
240+ parity,
241+ };
242+ (this as { legLat: number }).legLat = parity ? nlat / 2 : nlat;
243+ const fourP = { mmax, nlat, nphi, radix: (tuning('SHT_RADIX') as number | undefined) ?? 4 };
244+ // The spatial field is real (layout.ts stores m >= 0 only), so the Fourier
245+ // stage can run an nphi/2-point complex FFT plus a recombination instead of
246+ // a full nphi-point one: half the arithmetic and half the workgroup storage.
247+ // The complex kernels remain for a future complex-valued field, and are what
248+ // SHT_REAL_FFT=0 selects.
249+ const realFft =
250+ this.fourierMode === 'fft' && nphi % 2 === 0 && tuning('SHT_REAL_FFT') !== false;
251+ const fftS = realFft ? fftSynthRealWGSL : fftSynthWGSL;
252+ const fftA = realFft ? fftAnalysRealWGSL : fftAnalysWGSL;
253+ const [pLegS, pLegA, pFourS, pFourA] = await Promise.all([
254+ makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),
255+ makePipeline(dev, legAnalysWGSL(legP), 'leg_analys'),
256+ makePipeline(
257+ dev,
258+ this.fourierMode === 'fft' ? fftS(fourP) : dftSynthWGSL(fourP),
259+ this.fourierMode === 'fft' ? 'fft_synth' : 'dft_synth',
260+ ),
261+ makePipeline(
262+ dev,
263+ this.fourierMode === 'fft' ? fftA(fourP) : dftAnalysWGSL(fourP),
264+ this.fourierMode === 'fft' ? 'fft_analys' : 'dft_analys',
265+ ),
266+ ]);
267+ this.pipeLegSynth = pLegS;
268+ this.pipeLegAnalys = pLegA;
269+ this.pipeFourSynth = pFourS;
270+ this.pipeFourAnalys = pFourA;
271+
272+ const entries = bgEntries;
273+ this.bgLegSynth = dev.createBindGroup({
274+ layout: pLegS.getBindGroupLayout(0),
275+ entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.qlmIn, this.fmBuf]),
276+ });
277+ this.bgLegAnalys = dev.createBindGroup({
278+ layout: pLegA.getBindGroupLayout(0),
279+ entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, this.qlmOut]),
280+ });
281+ this.bgFourSynth = dev.createBindGroup({
282+ layout: pFourS.getBindGroupLayout(0),
283+ entries: entries([this.fmBuf, this.spatBuf, this.bufTrig]),
284+ });
285+ this.bgFourAnalys = dev.createBindGroup({
286+ layout: pFourA.getBindGroupLayout(0),
287+ entries: entries([this.spatBuf, this.fmBuf, this.bufTrig]),
288+ });
289+ }
290+
291+ /**
292+ * Bind groups for one transform against caller-supplied spectral/spatial
293+ * buffers, so a transform can read and write buffers it does not own (the
294+ * .m-driven executor keeps a buffer per IR variable). Build these once at
295+ * plan time, not per step. `fmBuf` stays internal scratch: passes and
296+ * dispatches within a submission execute in order, so sequential transforms
297+ * can share it.
298+ */
299+ createSynthBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): ShtBinding {
300+ return {
301+ bgLeg: this.device.createBindGroup({
302+ layout: this.pipeLegSynth.getBindGroupLayout(0),
303+ entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, qlmIn, this.fmBuf]),
304+ }),
305+ bgFour: this.device.createBindGroup({
306+ layout: this.pipeFourSynth.getBindGroupLayout(0),
307+ entries: bgEntries([this.fmBuf, spatOut, this.bufTrig]),
308+ }),
309+ };
310+ }
311+
312+ createAnalysBinding(spatIn: GPUBuffer, qlmOut: GPUBuffer): ShtBinding {
313+ return {
314+ bgFour: this.device.createBindGroup({
315+ layout: this.pipeFourAnalys.getBindGroupLayout(0),
316+ entries: bgEntries([spatIn, this.fmBuf, this.bufTrig]),
317+ }),
318+ bgLeg: this.device.createBindGroup({
319+ layout: this.pipeLegAnalys.getBindGroupLayout(0),
320+ entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, qlmOut]),
321+ }),
322+ };
323+ }
324+
325+ /** Record synthesis into an existing compute pass. */
326+ encodeSynthInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
327+ const { mmax, nlat, nphi } = this.cfg;
328+ pass.setPipeline(this.pipeLegSynth);
329+ pass.setBindGroup(0, b.bgLeg);
330+ pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
331+ pass.setPipeline(this.pipeFourSynth);
332+ pass.setBindGroup(0, b.bgFour);
333+ if (this.fourierMode === 'fft') {
334+ pass.dispatchWorkgroups(nlat);
335+ } else {
336+ pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
337+ }
338+ }
339+
340+ /** Record analysis into an existing compute pass. */
341+ encodeAnalysInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
342+ const { mmax, nlat } = this.cfg;
343+ pass.setPipeline(this.pipeFourAnalys);
344+ pass.setBindGroup(0, b.bgFour);
345+ if (this.fourierMode === 'fft') {
346+ pass.dispatchWorkgroups(nlat);
347+ } else {
348+ pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
349+ }
350+ pass.setPipeline(this.pipeLegAnalys);
351+ pass.setBindGroup(0, b.bgLeg);
352+ pass.dispatchWorkgroups(mmax + 1);
353+ }
354+
355+ /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
356+ encodeSynth(encoder: GPUCommandEncoder): void {
357+ const pass = encoder.beginComputePass({ label: 'sht-synth' });
358+ this.encodeSynthInto(pass, { bgLeg: this.bgLegSynth, bgFour: this.bgFourSynth });
359+ pass.end();
360+ }
361+
362+ /**
363+ * Diagnostics: encode one stage alone, in its own pass, so a timestamp query
364+ * can measure just that kernel. The solver wants both stages in a shared pass
365+ * and should use encodeSynthInto/encodeAnalysInto; this exists because
366+ * inferring per-kernel cost by subtracting trivially-sized runs is unreliable.
367+ */
368+ encodeStage(
369+ encoder: GPUCommandEncoder,
370+ stage: 'legSynth' | 'fourSynth' | 'fourAnalys' | 'legAnalys',
371+ timestampWrites?: GPUComputePassTimestampWrites,
372+ ): void {
373+ const { mmax, nlat, nphi } = this.cfg;
374+ const fft = this.fourierMode === 'fft';
375+ const pass = encoder.beginComputePass({ label: `sht-${stage}`, timestampWrites });
376+ switch (stage) {
377+ case 'legSynth':
378+ pass.setPipeline(this.pipeLegSynth);
379+ pass.setBindGroup(0, this.bgLegSynth);
380+ pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
381+ break;
382+ case 'fourSynth':
383+ pass.setPipeline(this.pipeFourSynth);
384+ pass.setBindGroup(0, this.bgFourSynth);
385+ if (fft) pass.dispatchWorkgroups(nlat);
386+ else pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
387+ break;
388+ case 'fourAnalys':
389+ pass.setPipeline(this.pipeFourAnalys);
390+ pass.setBindGroup(0, this.bgFourAnalys);
391+ if (fft) pass.dispatchWorkgroups(nlat);
392+ else pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
393+ break;
394+ case 'legAnalys':
395+ pass.setPipeline(this.pipeLegAnalys);
396+ pass.setBindGroup(0, this.bgLegAnalys);
397+ pass.dispatchWorkgroups(mmax + 1);
398+ break;
399+ }
400+ pass.end();
401+ }
402+
403+ /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
404+ encodeAnalys(encoder: GPUCommandEncoder): void {
405+ const pass = encoder.beginComputePass({ label: 'sht-analys' });
406+ this.encodeAnalysInto(pass, { bgLeg: this.bgLegAnalys, bgFour: this.bgFourAnalys });
407+ pass.end();
408+ }
409+
410+ /**
411+ * Spectral -> spatial. qlm: interleaved [re, im], SHTNS LM ordering,
412+ * length 2*nlm. Returns the spatial field, length nlat*nphi.
413+ */
414+ async synth(qlm: Float32Array): Promise<Float32Array> {
415+ const { nlat, nphi } = this.cfg;
416+ if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
417+ this.device.queue.writeBuffer(this.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
418+ const enc = this.device.createCommandEncoder();
419+ this.encodeSynth(enc);
420+ enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
421+ this.device.queue.submit([enc.finish()]);
422+ await this.stageSpat.mapAsync(GPUMapMode.READ);
423+ const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
424+ this.stageSpat.unmap();
425+ return out;
426+ }
427+
428+ /**
429+ * Spectral -> spatial, with the coefficients read from a caller-owned GPU
430+ * buffer (interleaved [re, im], 8*nlm bytes, COPY_SRC) instead of uploaded
431+ * from the CPU. This is how a field already on the device — a model's
432+ * spectral state — is evaluated on this plan's grid, e.g. a finer display
433+ * grid than the one the coefficients were produced on.
434+ */
435+ async synthFrom(qlmSrc: GPUBuffer): Promise<Float32Array> {
436+ const { nlat, nphi } = this.cfg;
437+ const enc = this.device.createCommandEncoder({ label: 'sht-synth-from' });
438+ enc.copyBufferToBuffer(qlmSrc, 0, this.qlmIn, 0, 8 * this.nlm);
439+ this.encodeSynth(enc);
440+ enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
441+ this.device.queue.submit([enc.finish()]);
442+ await this.stageSpat.mapAsync(GPUMapMode.READ);
443+ const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
444+ this.stageSpat.unmap();
445+ return out;
446+ }
447+
448+ /** Spatial -> spectral. spat: length nlat*nphi. Returns interleaved qlm, length 2*nlm. */
449+ async analys(spat: Float32Array): Promise<Float32Array> {
450+ const { nlat, nphi } = this.cfg;
451+ if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
452+ this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
453+ const enc = this.device.createCommandEncoder();
454+ this.encodeAnalys(enc);
455+ enc.copyBufferToBuffer(this.qlmOut, 0, this.stageQ, 0, 8 * this.nlm);
456+ this.device.queue.submit([enc.finish()]);
457+ await this.stageQ.mapAsync(GPUMapMode.READ);
458+ const out = new Float32Array(this.stageQ.getMappedRange().slice(0));
459+ this.stageQ.unmap();
460+ return out;
461+ }
462+
463+ destroy(): void {
464+ for (const b of [
465+ this.bufAb, this.bufAmm, this.bufCtstw, this.bufTrig, this.qlmIn, this.qlmOut,
466+ this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ,
467+ ]) b?.destroy();
468+ }
469+}
470+
471+/** Best-effort human-readable adapter name, so it is clear which GPU (or
472+ * software rasterizer) is actually running the transforms. */
473+export async function describeAdapter(device: GPUDevice): Promise<string> {
474+ const fmt = (info: GPUAdapterInfo | undefined): string => {
475+ if (!info) return '';
476+ const parts = [info.description, info.device, info.vendor].filter(
477+ (s): s is string => !!s && s.length > 0,
478+ );
479+ const name = parts[0] ?? '';
480+ return info.architecture && !name.includes(info.architecture)
481+ ? `${name} (${info.architecture})`.trim()
482+ : name;
483+ };
484+ const own = fmt((device as GPUDevice & { adapterInfo?: GPUAdapterInfo }).adapterInfo);
485+ if (own) return own;
486+ try {
487+ const adapter = await navigator.gpu.requestAdapter();
488+ return fmt(adapter?.info);
489+ } catch {
490+ return '';
491+ }
492+}
493+
494+/** Request an adapter/device suitable for the transforms. */
495+export async function requestShtDevice(): Promise<GPUDevice> {
496+ if (!navigator.gpu) throw new Error('WebGPU is not available in this browser');
497+ const adapter = await navigator.gpu.requestAdapter();
498+ if (!adapter) throw new Error('No WebGPU adapter available');
499+ // ask for a larger workgroup storage if the adapter offers it (bigger FFTs)
500+ const wgStorage = Math.min(adapter.limits.maxComputeWorkgroupStorageSize, 32768);
501+ // `subgroups` lets the analysis reduction use subgroupAdd instead of a
502+ // shared-memory tree (2 barriers per l-pair instead of 1 + log2(wgAnalys)).
503+ // Optional: ShtPlan falls back to the tree when it is not available.
504+ const features: GPUFeatureName[] = [];
505+ if (adapter.features.has('subgroups')) features.push('subgroups');
506+ // timestamp-query is only used by the profiling scripts, but it has to be
507+ // requested at device creation, and asking costs nothing when unused.
508+ if (adapter.features.has('timestamp-query')) features.push('timestamp-query');
509+ return adapter.requestDevice({
510+ requiredFeatures: features,
511+ requiredLimits: { maxComputeWorkgroupStorageSize: wgStorage },
512+ });
513+}
src/sht/wgsl/common.tsadded+56−0View file
@@ -0,0 +1,56 @@
1+/**
2+ * Shared WGSL fragments. Shaders are generated as strings with all sizes
3+ * baked in as compile-time constants (the WGSL analog of what SHTNS does
4+ * with NVRTC on CUDA: cf. init_cuda_program() in sht_gpu.cu).
5+ *
6+ * fp32 extended-range constants: same values SHTNS injects for a
7+ * single-precision recurrence (sht_gpu.cu):
8+ * SHT_ACCURACY = 1e-15
9+ * SHT_SCALE_FACTOR = 2^56 = 7.2057594037927936e16
10+ * A per-thread integer exponent `ny` counts how many times the running
11+ * Legendre value has been multiplied by SCALE to stay in fp32 range;
12+ * contributions are only accumulated once ny == 0 (value back in normal
13+ * range and significant).
14+ */
15+
16+export const RESCALE_WGSL = /* wgsl */ `
17+const SCALE: f32 = 7.2057594e16; // rounds to exactly 2^56 in f32
18+const INV_SCALE: f32 = 1.0 / 7.2057594e16;
19+const ACCURACY: f32 = 1e-15;
20+const RESCALE_THR: f32 = ACCURACY * SCALE + 1.0; // ~73: value became significant again
21+
22+struct Seed { y0: f32, ny: i32 }
23+
24+// Seed of the recurrence: y0 ~ sin(theta)^m by binary exponentiation with
25+// rescaling (ports the HI_LLIM path of SHT/cuda_legendre.gen.cu, ~651-691).
26+// The caller multiplies by amm afterwards (|amm| is O(1)).
27+fn sinpow_rescaled(st: f32, m: u32) -> Seed {
28+ var y0: f32 = 1.0;
29+ var ny: i32 = 0;
30+ if (m > 0u) {
31+ var s: f32 = st;
32+ var lb: u32 = m;
33+ if ((lb & 1u) != 0u) { y0 = s; }
34+ var nsint: i32 = 0;
35+ lb = lb >> 1u;
36+ while (lb > 0u) {
37+ s = s * s;
38+ nsint = nsint + nsint;
39+ if (s < INV_SCALE) {
40+ nsint = nsint - 1;
41+ s = s * SCALE;
42+ }
43+ if ((lb & 1u) != 0u) {
44+ y0 = y0 * s;
45+ ny = ny + nsint;
46+ if (y0 < (ACCURACY + INV_SCALE)) {
47+ y0 = y0 * SCALE;
48+ ny = ny - 1;
49+ }
50+ }
51+ lb = lb >> 1u;
52+ }
53+ }
54+ return Seed(y0, ny);
55+}
56+`;
src/sht/wgsl/fourier.tsadded+386−0View file
@@ -0,0 +1,386 @@
1+/**
2+ * WGSL Fourier-stage kernels (the role cuFFT/VkFFT plays in SHTNS).
3+ *
4+ * Real fields, band-limited to |m| <= mmax < nphi/2:
5+ * - synthesis: assemble a Hermitian spectrum from F_m (m >= 0) and do an
6+ * inverse complex FFT along phi; take the real part.
7+ * - analysis: forward complex FFT of the (real) row; keep m = 0..mmax.
8+ *
9+ * Two implementations, selected at plan creation:
10+ * - 'fft': radix-2 Stockham in workgroup memory, one workgroup per
11+ * latitude row. Requires nphi a power of two and
12+ * 2 * 8 * nphi bytes <= maxComputeWorkgroupStorageSize.
13+ * - 'dft': direct band-limited trigonometric summation, O(nphi * mmax)
14+ * per row. Works for any nphi; also useful as a cross-check.
15+ *
16+ * All trigonometric factors come from a host-precomputed (f64 -> f32)
17+ * table trig[k] = (cos, sin)(2*pi*k/nphi): device sin/cos is only
18+ * guaranteed to ~2^-11 absolute error under Vulkan, which would dominate
19+ * the fp32 transform error.
20+ */
21+
22+export interface FourierParams {
23+ mmax: number;
24+ nlat: number;
25+ nphi: number;
26+ /** 2 or 4; radix-4 uses log4(n) barrier stages instead of log2(n). */
27+ radix?: number;
28+}
29+
30+const TRIG_BINDING = /* wgsl */ `
31+@group(0) @binding(2) var<storage, read> trig: array<vec2f>; // (cos,sin)(2*pi*k/NPHI), k < NPHI
32+`;
33+
34+/**
35+ * @param n transform length (bufA/bufB are this long)
36+ * @param scale trig-table stride multiplier: the table holds
37+ * (cos,sin)(2*pi*k/NPHI), so an n-point transform needs NPHI/n.
38+ */
39+function stockham(n: number, threads: number, sign: number, scale = 1): string {
40+ const nphi = n;
41+ const log2n = Math.log2(n);
42+ if (!Number.isInteger(log2n)) throw new Error('fft requires power-of-two nphi');
43+ // twiddle for pass with half-block ns: w = e^{sign*i*pi*j/ns} = T[j * (N/(2*ns))]^sign
44+ return /* wgsl */ `
45+var<workgroup> bufA: array<vec2f, ${nphi}>;
46+var<workgroup> bufB: array<vec2f, ${nphi}>;
47+
48+fn cmul(a: vec2f, b: vec2f) -> vec2f {
49+ return vec2f(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
50+}
51+
52+fn ld(sel: u32, i: u32) -> vec2f {
53+ if (sel == 0u) { return bufA[i]; }
54+ return bufB[i];
55+}
56+fn st_(sel: u32, i: u32, v: vec2f) {
57+ if (sel == 0u) { bufA[i] = v; } else { bufB[i] = v; }
58+}
59+
60+// radix-2 Stockham, natural order in and out; data starts in bufA (sel 0)
61+// and ends in sel = LOG2N % 2. Unnormalized: X_k = sum_j x_j e^{s*2*pi*i*jk/N}.
62+fn fft_inplace(lid: u32) {
63+ for (var p = 0u; p < ${log2n}u; p++) {
64+ workgroupBarrier();
65+ let ns = 1u << p;
66+ let sel = p & 1u;
67+ let stride = ${(nphi / 2) * scale}u >> p; // (NPHI/n) * n/(2*ns)
68+ for (var t = lid; t < ${nphi / 2}u; t += ${threads}u) {
69+ let j = t & (ns - 1u);
70+ let tw = trig[j * stride];
71+ let w = vec2f(tw.x, ${sign > 0 ? '' : '-'}tw.y);
72+ let u = ld(sel, t);
73+ let v = cmul(ld(sel, t + ${nphi / 2}u), w);
74+ let idst = 2u * (t - j) + j;
75+ st_(1u - sel, idst, u + v);
76+ st_(1u - sel, idst + ns, u - v);
77+ }
78+ }
79+ workgroupBarrier();
80+}
81+const FFT_OUT_SEL: u32 = ${log2n % 2}u;
82+`;
83+}
84+
85+/**
86+ * Radix-4 Stockham. Same interface and conventions as stockham(), but log4(n)
87+ * stages instead of log2(n) -- each stage carries a workgroupBarrier, and
88+ * barriers are what these kernels are actually bound by. When log2(n) is odd a
89+ * single radix-2 stage runs first, so n = 128 costs 1 + 3 stages rather than 7.
90+ *
91+ * Butterfly, with w = e^{s 2 pi i / 4}:
92+ * a = x0 + x2, b = x0 - x2, c = x1 + x3, d = s i (x1 - x3)
93+ * y = (a + c, b + d, a - c, b - d)
94+ */
95+function stockham4(n: number, threads: number, sign: number, scale = 1): string {
96+ const log2n = Math.log2(n);
97+ if (!Number.isInteger(log2n)) throw new Error('fft requires power-of-two nphi');
98+ const needR2 = log2n % 2 === 1;
99+ const stages4 = Math.floor(log2n / 2);
100+ const total = (needR2 ? 1 : 0) + stages4;
101+ const negY = sign > 0 ? '' : '-';
102+ return /* wgsl */ `
103+var<workgroup> bufA: array<vec2f, ${n}>;
104+var<workgroup> bufB: array<vec2f, ${n}>;
105+
106+fn cmul(a: vec2f, b: vec2f) -> vec2f {
107+ return vec2f(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
108+}
109+fn ld(sel: u32, i: u32) -> vec2f {
110+ if (sel == 0u) { return bufA[i]; }
111+ return bufB[i];
112+}
113+fn st_(sel: u32, i: u32, v: vec2f) {
114+ if (sel == 0u) { bufA[i] = v; } else { bufB[i] = v; }
115+}
116+fn tw(i: u32) -> vec2f {
117+ let t = trig[i];
118+ return vec2f(t.x, ${negY}t.y);
119+}
120+
121+fn fft_inplace(lid: u32) {
122+ var sel = 0u;
123+ var ns = 1u;
124+${
125+ needR2
126+ ? ` // leading radix-2 (ns = 1, so the twiddle is 1 and is skipped)
127+ workgroupBarrier();
128+ for (var t = lid; t < ${n / 2}u; t += ${threads}u) {
129+ let u = ld(sel, t);
130+ let v = ld(sel, t + ${n / 2}u);
131+ st_(1u - sel, 2u * t, u + v);
132+ st_(1u - sel, 2u * t + 1u, u - v);
133+ }
134+ sel = 1u - sel;
135+ ns = 2u;`
136+ : ''
137+}
138+ for (var p = 0u; p < ${stages4}u; p++) {
139+ workgroupBarrier();
140+ let s4 = ${(n * scale) / 4}u / ns; // trig unit: NPHI / (4 * ns)
141+ for (var t = lid; t < ${n / 4}u; t += ${threads}u) {
142+ let j = t & (ns - 1u);
143+ let x0 = ld(sel, t);
144+ var x1 = ld(sel, t + ${n / 4}u);
145+ var x2 = ld(sel, t + ${n / 2}u);
146+ var x3 = ld(sel, t + ${(3 * n) / 4}u);
147+ if (ns > 1u) {
148+ x1 = cmul(x1, tw(j * s4));
149+ x2 = cmul(x2, tw(2u * j * s4));
150+ x3 = cmul(x3, tw(3u * j * s4));
151+ }
152+ let a = x0 + x2;
153+ let b = x0 - x2;
154+ let c = x1 + x3;
155+ let e = x1 - x3;
156+ let d = vec2f(${sign > 0 ? '-e.y, e.x' : 'e.y, -e.x'}); // s * i * e
157+ let idst = 4u * (t - j) + j;
158+ st_(1u - sel, idst, a + c);
159+ st_(1u - sel, idst + ns, b + d);
160+ st_(1u - sel, idst + 2u * ns, a - c);
161+ st_(1u - sel, idst + 3u * ns, b - d);
162+ }
163+ sel = 1u - sel;
164+ ns = ns * 4u;
165+ }
166+ workgroupBarrier();
167+}
168+const FFT_OUT_SEL: u32 = ${total % 2}u;
169+`;
170+}
171+
172+/** Choose FFT workgroup size: enough threads for the butterflies, capped at 256. */
173+export function fftThreads(nphi: number): number {
174+ return Math.max(32, Math.min(256, nphi / 2));
175+}
176+
177+export function fftSynthWGSL(p: FourierParams): string {
178+ const T = fftThreads(p.nphi);
179+ return /* wgsl */ `
180+const MMAX: u32 = ${p.mmax}u;
181+const NLAT: u32 = ${p.nlat}u;
182+const NPHI: u32 = ${p.nphi}u;
183+@group(0) @binding(0) var<storage, read> fm: array<vec2f>;
184+@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
185+${TRIG_BINDING}
186+${stockham(p.nphi, T, +1)}
187+
188+@compute @workgroup_size(${T})
189+fn fft_synth(@builtin(local_invocation_id) lid3: vec3u,
190+ @builtin(workgroup_id) wid: vec3u) {
191+ let lid = lid3.x;
192+ let ilat = wid.x;
193+ // assemble Hermitian spectrum: X[0] = Re F_0, X[m] = F_m, X[N-m] = conj(F_m)
194+ for (var k = lid; k < NPHI; k += ${T}u) {
195+ var v = vec2f(0.0);
196+ if (k == 0u) {
197+ v = vec2f(fm[ilat].x, 0.0);
198+ } else if (k <= MMAX) {
199+ v = fm[k * NLAT + ilat];
200+ } else if (k >= NPHI - MMAX) {
201+ let c = fm[(NPHI - k) * NLAT + ilat];
202+ v = vec2f(c.x, -c.y);
203+ }
204+ bufA[k] = v;
205+ }
206+ fft_inplace(lid);
207+ for (var k = lid; k < NPHI; k += ${T}u) {
208+ spat[ilat * NPHI + k] = ld(FFT_OUT_SEL, k).x;
209+ }
210+}
211+`;
212+}
213+
214+export function fftAnalysWGSL(p: FourierParams): string {
215+ const T = fftThreads(p.nphi);
216+ return /* wgsl */ `
217+const MMAX: u32 = ${p.mmax}u;
218+const NLAT: u32 = ${p.nlat}u;
219+const NPHI: u32 = ${p.nphi}u;
220+@group(0) @binding(0) var<storage, read> spat: array<f32>;
221+@group(0) @binding(1) var<storage, read_write> fm: array<vec2f>;
222+${TRIG_BINDING}
223+${stockham(p.nphi, T, -1)}
224+
225+@compute @workgroup_size(${T})
226+fn fft_analys(@builtin(local_invocation_id) lid3: vec3u,
227+ @builtin(workgroup_id) wid: vec3u) {
228+ let lid = lid3.x;
229+ let ilat = wid.x;
230+ for (var k = lid; k < NPHI; k += ${T}u) {
231+ bufA[k] = vec2f(spat[ilat * NPHI + k], 0.0);
232+ }
233+ fft_inplace(lid);
234+ for (var m = lid; m <= MMAX; m += ${T}u) {
235+ fm[m * NLAT + ilat] = ld(FFT_OUT_SEL, m);
236+ }
237+}
238+`;
239+}
240+
241+export function dftSynthWGSL(p: FourierParams): string {
242+ return /* wgsl */ `
243+const MMAX: u32 = ${p.mmax}u;
244+const NLAT: u32 = ${p.nlat}u;
245+const NPHI: u32 = ${p.nphi}u;
246+@group(0) @binding(0) var<storage, read> fm: array<vec2f>;
247+@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
248+${TRIG_BINDING}
249+
250+@compute @workgroup_size(64)
251+fn dft_synth(@builtin(global_invocation_id) gid: vec3u) {
252+ let iphi = gid.x;
253+ let ilat = gid.y;
254+ if (iphi >= NPHI) { return; }
255+ var v: f32 = fm[ilat].x; // m = 0: real part
256+ for (var m = 1u; m <= MMAX; m++) {
257+ let w = trig[(m * iphi) % NPHI]; // e^{+i m phi}
258+ let c = fm[m * NLAT + ilat];
259+ v += 2.0 * (c.x * w.x - c.y * w.y);
260+ }
261+ spat[ilat * NPHI + iphi] = v;
262+}
263+`;
264+}
265+
266+export function dftAnalysWGSL(p: FourierParams): string {
267+ return /* wgsl */ `
268+const MMAX: u32 = ${p.mmax}u;
269+const NLAT: u32 = ${p.nlat}u;
270+const NPHI: u32 = ${p.nphi}u;
271+@group(0) @binding(0) var<storage, read> spat: array<f32>;
272+@group(0) @binding(1) var<storage, read_write> fm: array<vec2f>;
273+${TRIG_BINDING}
274+
275+@compute @workgroup_size(64)
276+fn dft_analys(@builtin(global_invocation_id) gid: vec3u) {
277+ let m = gid.x;
278+ let ilat = gid.y;
279+ if (m > MMAX) { return; }
280+ var acc = vec2f(0.0);
281+ for (var j = 0u; j < NPHI; j++) {
282+ let w = trig[(m * j) % NPHI]; // conj => e^{-i m phi}
283+ let f = spat[ilat * NPHI + j];
284+ acc += f * vec2f(w.x, -w.y);
285+ }
286+ fm[m * NLAT + ilat] = acc;
287+}
288+`;
289+}
290+
291+/**
292+ * Real-field Fourier stage: half the arithmetic and half the shared memory of
293+ * the complex path, which transforms N points to get a Hermitian result.
294+ *
295+ * A length-N real transform is an N/2-point complex FFT wrapped in a
296+ * recombination. Writing H = N/2 and taking the unnormalized conventions of the
297+ * complex kernels above (synthesis e^{+i}, analysis e^{-i}):
298+ *
299+ * synthesis Z[k] = (X[k] + conj(X[H-k])) + i e^{+2pi i k/N} (X[k] - conj(X[H-k]))
300+ * z = FFT_H^{+}(Z), then x[2m] = Re z[m], x[2m+1] = Im z[m]
301+ * analysis z[m] = x[2m] + i x[2m+1], Z = FFT_H^{-}(z)
302+ * Xe = (Z[k] + conj(Z[H-k]))/2, Xo = -i (Z[k] - conj(Z[H-k]))/2
303+ * X[k] = Xe + e^{-2pi i k/N} Xo
304+ *
305+ * The factors of 2 in the synthesis direction cancel against the 1/2 in Xe/Xo,
306+ * which is why none appear there. The complex kernels are kept: they are what a
307+ * complex-valued spatial field would use, and stockham() is shared by both.
308+ */
309+export function fftSynthRealWGSL(p: FourierParams): string {
310+ const H = p.nphi / 2;
311+ const T = fftThreads(H);
312+ return /* wgsl */ `
313+const MMAX: u32 = ${p.mmax}u;
314+const NLAT: u32 = ${p.nlat}u;
315+const NPHI: u32 = ${p.nphi}u;
316+const H: u32 = ${H}u;
317+@group(0) @binding(0) var<storage, read> fm: array<vec2f>;
318+@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
319+${TRIG_BINDING}
320+${(p.radix ?? 4) === 4 ? stockham4(H, T, +1, p.nphi / H) : stockham(H, T, +1, p.nphi / H)}
321+
322+// X[k] of the Hermitian spectrum, for 0 <= k <= H. mmax < H, so the
323+// upper-conjugate branch of the complex kernel cannot be reached here.
324+fn spec(ilat: u32, k: u32) -> vec2f {
325+ if (k == 0u) { return vec2f(fm[ilat].x, 0.0); }
326+ if (k <= MMAX) { return fm[k * NLAT + ilat]; }
327+ return vec2f(0.0);
328+}
329+
330+@compute @workgroup_size(${T})
331+fn fft_synth(@builtin(local_invocation_id) lid3: vec3u,
332+ @builtin(workgroup_id) wid: vec3u) {
333+ let lid = lid3.x;
334+ let ilat = wid.x;
335+ for (var k = lid; k < H; k += ${T}u) {
336+ let xk = spec(ilat, k);
337+ let xh = spec(ilat, H - k);
338+ let cj = vec2f(xh.x, -xh.y);
339+ let b = cmul(xk - cj, trig[k]); // e^{+2 pi i k / N}
340+ bufA[k] = (xk + cj) + vec2f(-b.y, b.x); // + i * b
341+ }
342+ fft_inplace(lid);
343+ for (var m = lid; m < H; m += ${T}u) {
344+ let z = ld(FFT_OUT_SEL, m);
345+ spat[ilat * NPHI + 2u * m] = z.x;
346+ spat[ilat * NPHI + 2u * m + 1u] = z.y;
347+ }
348+}
349+`;
350+}
351+
352+export function fftAnalysRealWGSL(p: FourierParams): string {
353+ const H = p.nphi / 2;
354+ const T = fftThreads(H);
355+ return /* wgsl */ `
356+const MMAX: u32 = ${p.mmax}u;
357+const NLAT: u32 = ${p.nlat}u;
358+const NPHI: u32 = ${p.nphi}u;
359+const H: u32 = ${H}u;
360+@group(0) @binding(0) var<storage, read> spat: array<f32>;
361+@group(0) @binding(1) var<storage, read_write> fm: array<vec2f>;
362+${TRIG_BINDING}
363+${(p.radix ?? 4) === 4 ? stockham4(H, T, -1, p.nphi / H) : stockham(H, T, -1, p.nphi / H)}
364+
365+@compute @workgroup_size(${T})
366+fn fft_analys(@builtin(local_invocation_id) lid3: vec3u,
367+ @builtin(workgroup_id) wid: vec3u) {
368+ let lid = lid3.x;
369+ let ilat = wid.x;
370+ for (var m = lid; m < H; m += ${T}u) {
371+ bufA[m] = vec2f(spat[ilat * NPHI + 2u * m], spat[ilat * NPHI + 2u * m + 1u]);
372+ }
373+ fft_inplace(lid);
374+ for (var k = lid; k <= MMAX; k += ${T}u) {
375+ let zk = ld(FFT_OUT_SEL, k);
376+ let zh = ld(FFT_OUT_SEL, (H - k) % H); // Z[H] == Z[0]
377+ let cj = vec2f(zh.x, -zh.y);
378+ let xe = 0.5 * (zk + cj);
379+ let d = 0.5 * (zk - cj);
380+ let xo = vec2f(d.y, -d.x); // -i * d
381+ let w = vec2f(trig[k].x, -trig[k].y); // e^{-2 pi i k / N}
382+ fm[k * NLAT + ilat] = xe + cmul(xo, w);
383+ }
384+}
385+`;
386+}
src/sht/wgsl/leg.tsadded+352−0View file
@@ -0,0 +1,352 @@
1+/**
2+ * WGSL Legendre-transform kernels, modeled on leg_m_kernel / ileg_m_kernel
3+ * in SHT/cuda_legendre.gen.cu (non-Ishioka fp32 path: SHTNS disables the
4+ * Ishioka recurrence for fp32 because it loses too much accuracy).
5+ *
6+ * Synthesis: F_m(theta_i) = sum_{l=m..lmax} Q_lm * ytilde_l^m(theta_i)
7+ * - one thread per latitude, one workgroup row per m (workgroup_id.y).
8+ * Analysis: Q_lm = sum_i w_i * G_m(theta_i) * ytilde_l^m(theta_i)
9+ * - one workgroup per m; threads own latitudes (strided); per-l pair
10+ * workgroup tree reduction (portable stand-in for the CUDA warp
11+ * shuffles).
12+ *
13+ * The associated Legendre functions are generated on the fly by the
14+ * standard 3-term recurrence over l (coefficients a,b precomputed on the
15+ * host in f64), with the SHTNS fp32 rescaling scheme for sin(theta)^m
16+ * underflow (see common.ts).
17+ */
18+import { RESCALE_WGSL } from './common.ts';
19+
20+export interface LegParams {
21+ lmax: number;
22+ mmax: number;
23+ nlat: number;
24+ wgSynth: number; // workgroup size for synthesis (threads over latitude)
25+ wgAnalys: number; // workgroup size for analysis (power of two)
26+ /** Use subgroup reductions in the analysis kernel (needs the `subgroups` feature). */
27+ subgroups?: boolean;
28+ /** l-pairs accumulated before the span is reduced (subgroup path only). */
29+ spanPairs?: number;
30+ /**
31+ * Fold north/south latitude pairs onto one recurrence (halves Legendre work).
32+ * Needs an equator-symmetric grid with even nlat, which the Gauss grid is.
33+ */
34+ parity?: boolean;
35+}
36+
37+const BINDINGS = /* wgsl */ `
38+@group(0) @binding(0) var<storage, read> ab: array<vec2f>; // (a_l^m, b_l^m) per lm
39+@group(0) @binding(1) var<storage, read> amm: array<f32>; // seed per m
40+@group(0) @binding(2) var<storage, read> ctstw: array<f32>; // [ct | st | w], each NLAT
41+`;
42+
43+export function legSynthWGSL(p: LegParams): string {
44+ const half = p.parity === true;
45+ return /* wgsl */ `
46+${RESCALE_WGSL}
47+const LMAX: u32 = ${p.lmax}u;
48+const NLAT: u32 = ${p.nlat}u;
49+const NLAT_2: u32 = ${p.nlat / 2}u;
50+${BINDINGS}
51+@group(0) @binding(3) var<storage, read> qlm: array<vec2f>;
52+@group(0) @binding(4) var<storage, read_write> fm: array<vec2f>; // [(m)*NLAT + ilat]
53+
54+@compute @workgroup_size(${p.wgSynth})
55+fn leg_synth(@builtin(global_invocation_id) gid: vec3u,
56+ @builtin(workgroup_id) wid: vec3u) {
57+ let ilat = gid.x;
58+ let m = wid.y;
59+ if (ilat >= ${half ? 'NLAT_2' : 'NLAT'}) { return; }
60+
61+ let ct = ctstw[ilat];
62+ let st = ctstw[NLAT + ilat];
63+ let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u; // lm index of (l=m, m)
64+
65+ var seed = sinpow_rescaled(st, m);
66+ var y0 = seed.y0 * amm[m];
67+ var ny = seed.ny;
68+ var y1: f32 = 0.0;
69+ if (m < LMAX) {
70+ y1 = ab[base + 1u].x * ct * y0;
71+ }
72+
73+${
74+ half
75+ ? ` // Parity folding: ytilde_l^m(-x) = (-1)^(l-m) ytilde_l^m(x) and the Gauss
76+ // grid is symmetric, so one recurrence serves a north/south pair. y0 always
77+ // carries even (l-m) and y1 odd, so summing them apart gives
78+ // F_m(north) = accE + accO, F_m(south) = accE - accO.
79+ var accE = vec2f(0.0);
80+ var accO = vec2f(0.0);`
81+ : ` var acc = vec2f(0.0);`
82+ }
83+ var l = m;
84+ loop {
85+ if (ny == 0) {
86+${
87+ half
88+ ? ` accE += y0 * qlm[base + (l - m)];
89+ if (l + 1u <= LMAX) {
90+ accO += y1 * qlm[base + (l + 1u - m)];
91+ }`
92+ : ` acc += y0 * qlm[base + (l - m)];
93+ if (l + 1u <= LMAX) {
94+ acc += y1 * qlm[base + (l + 1u - m)];
95+ }`
96+ }
97+ } else if (abs(y0) > RESCALE_THR) {
98+ ny += 1;
99+ y0 *= INV_SCALE;
100+ y1 *= INV_SCALE;
101+ }
102+ if (l + 2u > LMAX) { break; }
103+ // Advance (y_l, y_{l+1}) to (y_{l+2}, y_{l+3}).
104+ //
105+ // Written in exactly the shape leg_analys uses below — both coefficients
106+ // fetched unconditionally, the new y0 carried in a temporary rather than
107+ // assigned and then read back by the y1 update. The shorter form,
108+ //
109+ // let c0 = ab[base + (l + 2u - m)];
110+ // y0 = c0.x * ct * y1 + c0.y * y0;
111+ // if (l + 3u <= LMAX) { ... y1 = c1.x * ct * y0 + c1.y * y1; }
112+ //
113+ // says the same thing and is what this was, but NVIDIA's Vulkan compiler
114+ // (driver 590.48, Blackwell) mis-compiles it: c0 reads as (0, 0) on the
115+ // first iteration, so y_{l+2} comes out exactly zero and every later term
116+ // follows a different solution of the recurrence, reaching ~1e11 by l = 63.
117+ // leg_analys, doing the same arithmetic in this shape, was correct on the
118+ // same driver. See scripts/diagnose-leg.ts, which is how that was found.
119+ let a0 = ab[base + (l + 2u - m)];
120+ var a1 = vec2f(0.0);
121+ if (l + 3u <= LMAX) {
122+ a1 = ab[base + (l + 3u - m)];
123+ }
124+ let t0 = a0.x * ct * y1 + a0.y * y0;
125+ y1 = a1.x * ct * t0 + a1.y * y1;
126+ y0 = t0;
127+ l += 2u;
128+ }
129+${
130+ half
131+ ? ` fm[m * NLAT + ilat] = accE + accO;
132+ fm[m * NLAT + (NLAT - 1u - ilat)] = accE - accO;`
133+ : ` fm[m * NLAT + ilat] = acc;`
134+ }
135+}
136+`;
137+}
138+
139+export function legAnalysWGSL(p: LegParams): string {
140+ const half = p.parity === true;
141+ // parity folding leaves only the northern half of the grid to walk
142+ const K = Math.ceil((half ? p.nlat / 2 : p.nlat) / p.wgAnalys);
143+ // With subgroups, the per-l-pair reduction is one subgroupAdd plus a combine
144+ // across subgroups: 2 barriers instead of 1 + log2(wgAnalys). This is what
145+ // SHTNS's CUDA kernel does with warp shuffles. `red` then holds one partial
146+ // per subgroup; WebGPU guarantees subgroup size >= 4, so wgAnalys/4 is a safe
147+ // upper bound on how many there can be.
148+ const sg = p.subgroups === true;
149+ // Reduce once per span of l-pairs rather than once per pair. The l-loop is
150+ // serial, so its barriers are the critical path: at lmax=127 the m=0
151+ // workgroup paid 2 of them 64 times over. SHTNS amortizes the same way
152+ // (LSPAN_A = 16, or 32 for fp32), staging a whole span before reducing.
153+ // Partials for the span live in registers and are combined in one batch.
154+ const nsubMax = Math.max(1, p.wgAnalys / 4); // WebGPU guarantees subgroup size >= 4
155+ // 16 pairs = 32 l-values, which is what SHTNS uses for fp32 (LSPAN_A). Clamped
156+ // so `red` stays within 8 KB of workgroup storage, since nsubMax has to assume
157+ // the smallest legal subgroup and would otherwise oversize it badly.
158+ const pairs = sg
159+ ? Math.max(1, Math.min(p.spanPairs ?? 16, Math.floor(8192 / (nsubMax * 16))))
160+ : 1;
161+ const redLen = sg ? nsubMax * pairs : p.wgAnalys;
162+ return /* wgsl */ `${sg ? 'enable subgroups;\n' : ''}
163+${RESCALE_WGSL}
164+const LMAX: u32 = ${p.lmax}u;
165+const NLAT: u32 = ${p.nlat}u;
166+const WG: u32 = ${p.wgAnalys}u;
167+const K: u32 = ${K}u;
168+const NLAT_2: u32 = ${p.nlat / 2}u;
169+const PAIRS: u32 = ${pairs}u;
170+${BINDINGS}
171+@group(0) @binding(3) var<storage, read> fm: array<vec2f>; // [(m)*NLAT + ilat]
172+@group(0) @binding(4) var<storage, read_write> qout: array<vec2f>;
173+
174+var<workgroup> red: array<vec4f, ${redLen}>;
175+
176+@compute @workgroup_size(${p.wgAnalys})
177+fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,
178+ @builtin(workgroup_id) wid: vec3u${
179+ sg
180+ ? ',\n @builtin(subgroup_size) sgSize: u32,\n @builtin(subgroup_invocation_id) sgLane: u32'
181+ : ''
182+ }) {
183+ let lid = lid3.x;
184+ let m = wid.x;
185+ let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
186+
187+ // per-thread recurrence state for K latitudes
188+ var y0v: array<f32, ${K}>;
189+ var y1v: array<f32, ${K}>;
190+ var nyv: array<i32, ${K}>;
191+ var ctv: array<f32, ${K}>;
192+${
193+ half
194+ ? ` // Transpose of the synthesis folding: splitting the latitude sum into
195+ // hemispheres gives Q_lm = sum_north w_i * ytilde * (G_north +/- G_south),
196+ // with + for even (l-m) and - for odd -- which the loop already routes
197+ // through y0 and y1 respectively.
198+ var wpv: array<vec2f, ${K}>;
199+ var wmv: array<vec2f, ${K}>;`
200+ : ` var wfv: array<vec2f, ${K}>;`
201+ }
202+
203+ for (var k = 0u; k < K; k++) {
204+ let lat = lid + k * WG;
205+ var ct: f32 = 0.0;
206+ var st: f32 = 0.0;
207+${
208+ half
209+ ? ` var wp = vec2f(0.0);
210+ var wm = vec2f(0.0);
211+ if (lat < NLAT_2) {
212+ ct = ctstw[lat];
213+ st = ctstw[NLAT + lat];
214+ let w = ctstw[2u * NLAT + lat]; // Gauss weight (incl. 2*pi/nphi)
215+ let gN = fm[m * NLAT + lat];
216+ let gS = fm[m * NLAT + (NLAT - 1u - lat)];
217+ wp = (gN + gS) * w;
218+ wm = (gN - gS) * w;
219+ }`
220+ : ` var wf = vec2f(0.0);
221+ if (lat < NLAT) {
222+ ct = ctstw[lat];
223+ st = ctstw[NLAT + lat];
224+ wf = fm[m * NLAT + lat] * ctstw[2u * NLAT + lat]; // Gauss weight (incl. 2*pi/nphi)
225+ }`
226+ }
227+ ctv[k] = ct;
228+ let seed = sinpow_rescaled(st, m);
229+ y0v[k] = seed.y0 * amm[m];
230+ nyv[k] = seed.ny;
231+ y1v[k] = 0.0;
232+ if (m < LMAX) {
233+ y1v[k] = ab[base + 1u].x * ct * y0v[k];
234+ }
235+${half ? ' wpv[k] = wp;\n wmv[k] = wm;' : ' wfv[k] = wf;'}
236+ }
237+
238+ var l = m;
239+${
240+ sg
241+ ? ` // Accumulate up to PAIRS l-pairs into registers, then reduce the whole span
242+ // at once: 2 barriers per span instead of 2 per pair.
243+ loop {
244+ let lstart = l;
245+ var npairs = 0u;
246+ var last = false;
247+ let sub = lid / sgSize;
248+ for (var jj = 0u; jj < PAIRS; jj++) {
249+ var c0 = vec2f(0.0);
250+ var c1 = vec2f(0.0);
251+ for (var k = 0u; k < K; k++) {
252+ if (nyv[k] == 0) {
253+${
254+ half
255+ ? ` c0 += wpv[k] * y0v[k]; // even (l-m): hemispheres add
256+ c1 += wmv[k] * y1v[k]; // odd (l-m): hemispheres subtract`
257+ : ` c0 += wfv[k] * y0v[k];
258+ c1 += wfv[k] * y1v[k];`
259+ }
260+ } else if (abs(y0v[k]) > RESCALE_THR) {
261+ nyv[k] += 1;
262+ y0v[k] *= INV_SCALE;
263+ y1v[k] *= INV_SCALE;
264+ }
265+ }
266+ // subgroupAdd needs no barrier, so the per-subgroup partial can go
267+ // straight to shared memory; only the cross-subgroup combine below has
268+ // to wait, and it waits once for the whole span.
269+ let part = subgroupAdd(vec4f(c0, c1));
270+ if (sgLane == 0u) { red[sub * PAIRS + jj] = part; }
271+ npairs = jj + 1u;
272+ if (l + 2u > LMAX) { last = true; break; }
273+ let a0 = ab[base + (l + 2u - m)];
274+ var a1 = vec2f(0.0);
275+ if (l + 3u <= LMAX) {
276+ a1 = ab[base + (l + 3u - m)];
277+ }
278+ for (var k = 0u; k < K; k++) {
279+ let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
280+ y0v[k] = t0;
281+ y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
282+ }
283+ l += 2u;
284+ }
285+
286+ workgroupBarrier();
287+ if (lid == 0u) {
288+ let nsub = (WG + sgSize - 1u) / sgSize;
289+ for (var jj = 0u; jj < npairs; jj++) {
290+ var tot = vec4f(0.0);
291+ for (var i = 0u; i < nsub; i++) { tot += red[i * PAIRS + jj]; }
292+ let ll = lstart + 2u * jj;
293+ qout[base + (ll - m)] = tot.xy;
294+ if (ll + 1u <= LMAX) {
295+ qout[base + (ll + 1u - m)] = tot.zw;
296+ }
297+ }
298+ }
299+ workgroupBarrier(); // red is reused by the next span
300+
301+ if (last) { break; }
302+ }`
303+ : ` loop {
304+ var c0 = vec2f(0.0);
305+ var c1 = vec2f(0.0);
306+ for (var k = 0u; k < K; k++) {
307+ if (nyv[k] == 0) {
308+${
309+ half
310+ ? ` c0 += wpv[k] * y0v[k]; // even (l-m): hemispheres add
311+ c1 += wmv[k] * y1v[k]; // odd (l-m): hemispheres subtract`
312+ : ` c0 += wfv[k] * y0v[k];
313+ c1 += wfv[k] * y1v[k];`
314+ }
315+ } else if (abs(y0v[k]) > RESCALE_THR) {
316+ nyv[k] += 1;
317+ y0v[k] *= INV_SCALE;
318+ y1v[k] *= INV_SCALE;
319+ }
320+ }
321+ // workgroup tree reduction of (c0, c1)
322+ red[lid] = vec4f(c0, c1);
323+ workgroupBarrier();
324+ var s = WG / 2u;
325+ while (s > 0u) {
326+ if (lid < s) { red[lid] += red[lid + s]; }
327+ workgroupBarrier();
328+ s = s >> 1u;
329+ }
330+ if (lid == 0u) {
331+ qout[base + (l - m)] = red[0].xy;
332+ if (l + 1u <= LMAX) {
333+ qout[base + (l + 1u - m)] = red[0].zw;
334+ }
335+ }
336+ if (l + 2u > LMAX) { break; }
337+ let a0 = ab[base + (l + 2u - m)];
338+ var a1 = vec2f(0.0);
339+ if (l + 3u <= LMAX) {
340+ a1 = ab[base + (l + 3u - m)];
341+ }
342+ for (var k = 0u; k < K; k++) {
343+ let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
344+ y0v[k] = t0;
345+ y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
346+ }
347+ l += 2u;
348+ }`
349+ }
350+}
351+`;
352+}
test.htmladded+12−0View file
@@ -0,0 +1,12 @@
1+<!doctype html>
2+<html lang="en">
3+ <head>
4+ <meta charset="utf-8" />
5+ <title>turing-surface validation</title>
6+ </head>
7+ <body>
8+ <h1>turing-surface validation suite</h1>
9+ <pre id="log">starting…</pre>
10+ <script type="module" src="/test/test-page.ts"></script>
11+ </body>
12+</html>
test/analyticChecks.tsadded+276−0View file
@@ -0,0 +1,276 @@
1+/**
2+ * Correctness of the .m -> WGSL path, against closed-form answers.
3+ *
4+ * These replace what used to be a comparison against a second TypeScript
5+ * implementation of the same scheme. Checking against arithmetic is stronger:
6+ * two implementations agreeing only shows they share assumptions, whereas an
7+ * exact recurrence pins the result. Each test picks a case whose evolution is
8+ * known in closed form, runs it through the real pipeline — MATLAB source,
9+ * numbl lowering, generated WGSL, GPU transforms — and compares.
10+ *
11+ * A: a linear reaction makes every spherical-harmonic mode independent, with a
12+ * known growth factor per degree. Checks the transform round-trip, the
13+ * eigenvalue mapping, the IMEX update and the state feedback.
14+ * B: a nonlinear reaction on a uniform field follows the scalar ODE map
15+ * exactly. Checks that a generated kernel evaluates a nonlinear reaction.
16+ * C: a small perturbation of the Schnakenberg fixed point follows the
17+ * linearized 2x2 IMEX recurrence, and the expected mode is unstable. Checks
18+ * a real two-species model.
19+ *
20+ * Everything runs in fp32 on the GPU, so tolerances are set by fp32 round-off
21+ * (~1e-7 relative) rather than by the scheme.
22+ */
23+import { ShtPlan } from '../src/sht/sht.ts';
24+import { gridForLmax, lmIndex, nlmCalc, type ShtConfig } from '../src/sht/layout.ts';
25+import { GpuModel } from '../src/mgpu/model.ts';
26+import { mModelByKey, defaultParams, type MModel, type ParamSpec } from '../src/mgpu/registry.ts';
27+import { Geometry } from '../src/geom/geometry.ts';
28+import { mGeometryByKey, SPHERE_KEY } from '../src/geom/registry.ts';
29+import linearSource from './models/linear.m?raw';
30+import logisticSource from './models/logistic.m?raw';
31+
32+export type Check = (name: string, ok: boolean, detail: string) => void;
33+export type Log = (s: string) => void;
34+
35+const param = (key: string, value: number): ParamSpec => ({
36+ key, label: key, value, min: -1e9, max: 1e9, step: 1,
37+});
38+
39+/** A one-species test model with an arbitrary parameter list. */
40+const testModel = (key: string, source: string, params: string[]): MModel => ({
41+ key,
42+ label: key,
43+ blurb: '',
44+ species: ['u'],
45+ state: ['U'],
46+ params: params.map((p) => param(p, 0)),
47+ pdeg: 1,
48+ seedAmp: 1,
49+ source,
50+});
51+
52+/**
53+ * Every closed-form case here is a statement about the *round sphere*, so
54+ * every model here is built on the sphere geometry. The models that take a
55+ * surface still get one — a geometry is always supplied, and for the sphere it
56+ * is the degree-1 embedding, which is what makes these answers exact.
57+ */
58+async function makeModel(
59+ device: GPUDevice,
60+ model: MModel,
61+ cfg: ShtConfig,
62+ niter = 1,
63+): Promise<{ sht: ShtPlan; gpu: GpuModel }> {
64+ const sht = await ShtPlan.create(device, cfg);
65+ const geometry = await Geometry.create({
66+ device,
67+ sht,
68+ cfg,
69+ source: mGeometryByKey(SPHERE_KEY)!.source,
70+ paramNames: [],
71+ params: {},
72+ });
73+ const gpu = await GpuModel.create({
74+ device,
75+ sht,
76+ cfg,
77+ source: model.source,
78+ paramNames: model.params.map((p) => p.key),
79+ state: model.state,
80+ view: model.species,
81+ geometry,
82+ niter,
83+ });
84+ return { sht, gpu };
85+}
86+
87+export async function analyticChecks(
88+ device: GPUDevice,
89+ check: Check,
90+ log: Log,
91+): Promise<void> {
92+ // ---- A: linear reaction, exact per-mode growth factor -----------------
93+ {
94+ const lmax = 15;
95+ const { nlat, nphi } = gridForLmax(lmax, 1);
96+ const cfg = { lmax, mmax: lmax, nlat, nphi };
97+ const nlm = nlmCalc(lmax, lmax);
98+ const c = -0.3;
99+ const D = 0.01;
100+ const dt = 0.1;
101+ const nsteps = 20;
102+
103+ const model = testModel('linear', linearSource, ['c', 'D', 'dt']);
104+ const { sht, gpu } = await makeModel(device, model, cfg);
105+ gpu.setParams({ c, D, dt });
106+
107+ // A single (l, m) mode, written straight into the spectral state.
108+ const l = 5;
109+ const m = 2;
110+ const idx = lmIndex(lmax, l, m);
111+ const U0 = new Float32Array(2 * nlm);
112+ U0[2 * idx] = 0.8;
113+ U0[2 * idx + 1] = -0.35;
114+ gpu.upload('U', U0);
115+
116+ gpu.step(nsteps);
117+ const U = await gpu.read('U');
118+
119+ const g = (1 + dt * c) / (1 + dt * D * l * (l + 1));
120+ const factor = g ** nsteps;
121+ const wantRe = 0.8 * factor;
122+ const wantIm = -0.35 * factor;
123+ const errRe = Math.abs(U[2 * idx] - wantRe);
124+ const errIm = Math.abs(U[2 * idx + 1] - wantIm);
125+ check(
126+ 'A: linear reaction follows the exact per-mode recurrence',
127+ errRe < 2e-6 && errIm < 2e-6,
128+ `err (${errRe.toExponential(2)}, ${errIm.toExponential(2)}) after ${nsteps} steps`,
129+ );
130+
131+ // Nothing may leak into the other modes.
132+ let leak = 0;
133+ for (let i = 0; i < nlm; i++) {
134+ if (i === idx) continue;
135+ leak = Math.max(leak, Math.abs(U[2 * i]), Math.abs(U[2 * i + 1]));
136+ }
137+ check('A: no leakage into other modes', leak < 2e-6, `max |other| ${leak.toExponential(2)}`);
138+
139+ gpu.destroy();
140+ sht.destroy();
141+ }
142+
143+ // ---- B: nonlinear reaction on a uniform field, exact ODE map ----------
144+ {
145+ const lmax = 15;
146+ const { nlat, nphi } = gridForLmax(lmax, 3);
147+ const cfg = { lmax, mmax: lmax, nlat, nphi };
148+ const npts = nlat * nphi;
149+ const r = 0.7;
150+ const D = 0.01;
151+ const dt = 0.05;
152+ const nsteps = 25;
153+ const u0 = 0.3;
154+
155+ const model = testModel('logistic', logisticSource, ['r', 'D', 'dt']);
156+ const { sht, gpu } = await makeModel(device, model, cfg);
157+ gpu.setParams({ r, D, dt });
158+
159+ // Uniform initial field: stays uniform, and diffusion cannot touch it.
160+ const field = new Float32Array(npts).fill(u0);
161+ gpu.init(field);
162+ const Ustart = await gpu.read('U');
163+ gpu.step(nsteps);
164+ const Uend = await gpu.read('U');
165+
166+ // Read the *state*, not the `u` output: a model computes its grid fields
167+ // from the state at the START of the step (`u = synth(U)` precedes the
168+ // update), so the rendered field lags the state by one step. The l=0
169+ // coefficient of a uniform field scales linearly with its value, so the
170+ // ratio gives the value back without needing Y_00's normalization.
171+ const got = u0 * (Uend[0] / Ustart[0]);
172+
173+ let want = u0;
174+ for (let s = 0; s < nsteps; s++) want += dt * r * want * (1 - want);
175+
176+ const err = Math.abs(got - want);
177+ check(
178+ 'B: uniform nonlinear reaction follows the scalar ODE map',
179+ err < 5e-6,
180+ `${got.toFixed(7)} vs ${want.toFixed(7)}, err ${err.toExponential(2)}`,
181+ );
182+
183+ // And it must still be uniform: any structure would mean the kernel is
184+ // reading the wrong elements.
185+ const u = await gpu.read('u');
186+ let lo = Infinity;
187+ let hi = -Infinity;
188+ for (const v of u) {
189+ if (v < lo) lo = v;
190+ if (v > hi) hi = v;
191+ }
192+ check(
193+ 'B: the field stays uniform',
194+ hi - lo < 1e-6,
195+ `spread ${(hi - lo).toExponential(2)}`,
196+ );
197+
198+ gpu.destroy();
199+ sht.destroy();
200+ }
201+
202+ // ---- C: linearized Turing recurrence on a real two-species model ------
203+ {
204+ const model = mModelByKey('schnakenberg')!;
205+ const p = defaultParams(model);
206+ const lmax = 31;
207+ const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
208+ const cfg = { lmax, mmax: lmax, nlat, nphi };
209+ const nlm = nlmCalc(lmax, lmax);
210+ const npts = nlat * nphi;
211+
212+ const { sht, gpu } = await makeModel(device, model, cfg);
213+ gpu.setParams(p);
214+
215+ // Seed the exact homogeneous fixed point by handing init a zero
216+ // perturbation, then add a small single-mode bump to u only.
217+ gpu.init(new Float32Array(npts));
218+ const l = 24;
219+ const m = 7;
220+ const idx = lmIndex(lmax, l, m);
221+ const eps = 1e-6;
222+ const U0 = await gpu.read('U');
223+ const V0 = await gpu.read('V');
224+ const Upert = Float32Array.from(U0);
225+ Upert[2 * idx] += eps;
226+ gpu.upload('U', Upert);
227+ gpu.upload('V', V0);
228+
229+ const nsteps = 40;
230+ gpu.step(nsteps);
231+ const U = await gpu.read('U');
232+ const V = await gpu.read('V');
233+
234+ // Jacobian of (a - u + u^2 v, b - u^2 v) at the fixed point us = a+b,
235+ // vs = b/us^2, with diffusion applied implicitly per species.
236+ const us = p.a + p.b;
237+ const vs = p.b / (us * us);
238+ const J = [
239+ [-1 + 2 * us * vs, us * us],
240+ [-2 * us * vs, -us * us],
241+ ];
242+ const lam = l * (l + 1);
243+ const du = 1 / (1 + p.dt * p.D1 * lam);
244+ const dv = 1 / (1 + p.dt * p.D2 * lam);
245+ let cu = eps;
246+ let cv = 0;
247+ for (let s = 0; s < nsteps; s++) {
248+ const nu = (cu + p.dt * (J[0][0] * cu + J[0][1] * cv)) * du;
249+ const nv = (cv + p.dt * (J[1][0] * cu + J[1][1] * cv)) * dv;
250+ cu = nu;
251+ cv = nv;
252+ }
253+
254+ const gotU = U[2 * idx] - U0[2 * idx];
255+ const gotV = V[2 * idx] - V0[2 * idx];
256+ const relU = Math.abs(gotU - cu) / Math.max(Math.abs(cu), 1e-30);
257+ const relV = Math.abs(gotV - cv) / Math.max(Math.abs(cv), 1e-30);
258+ // Looser than A and B by design: a 1e-6 perturbation sits on a state of
259+ // order 1, so fp32 keeps only ~4 significant digits of it.
260+ check(
261+ 'C: perturbation follows the linearized 2x2 IMEX recurrence',
262+ relU < 5e-3 && relV < 5e-3,
263+ `rel err (${relU.toExponential(2)}, ${relV.toExponential(2)})`,
264+ );
265+ check(
266+ `C: the (l=${l}, m=${m}) mode is unstable`,
267+ Math.abs(cu) > eps && Math.abs(gotU) > eps,
268+ `|c_u| ${eps.toExponential(2)} -> ${Math.abs(gotU).toExponential(2)}`,
269+ );
270+
271+ log(` C: growth over ${nsteps} steps = ${(Math.abs(cu) / eps).toFixed(3)}x (predicted)`);
272+
273+ gpu.destroy();
274+ sht.destroy();
275+ }
276+}
test/geometryChecks.tsadded+308−0View file
@@ -0,0 +1,308 @@
1+/**
2+ * The two things this project adds to turing-sphere: a surface, and a `for`
3+ * loop in the compiled step.
4+ *
5+ * The surface is checked against what it is supposed to be — the sphere really
6+ * is the unit sphere and really is degree 1, a deformed shape really has the
7+ * radius profile its .m says, and the coefficients really do evaluate to the
8+ * same surface on a finer grid.
9+ *
10+ * The loop is checked for the property the whole design rests on: it is
11+ * unrolled into the fixed op sequence, so more iterations means more GPU ops —
12+ * and, while the geometry correction inside it is identically zero, the answer
13+ * must be *bit for bit* independent of how many times it runs. That is a
14+ * stronger statement than "close enough": if the placeholder were ever
15+ * something that merely rounds to zero, or if the loop were miscompiled to
16+ * read a stale buffer, these would differ in the last bits and this fails.
17+ */
18+import { ShtPlan } from '../src/sht/sht.ts';
19+import { gridForLmax, lmIndex } from '../src/sht/layout.ts';
20+import { ModelSession } from '../src/mgpu/session.ts';
21+import { mModelByKey, defaultParams } from '../src/mgpu/registry.ts';
22+import { Geometry } from '../src/geom/geometry.ts';
23+import {
24+ mGeometries,
25+ mGeometryByKey,
26+ defaultGeometryParams,
27+ SPHERE_KEY,
28+} from '../src/geom/registry.ts';
29+import { ModelCompileError } from '../src/mgpu/errors.ts';
30+import type { Check, Log } from './analyticChecks.ts';
31+
32+const LMAX = 31;
33+const STEPS = 20;
34+
35+/** Build one geometry on its own transform plan, for inspection. */
36+async function buildGeometry(device: GPUDevice, key: string) {
37+ const g = mGeometryByKey(key)!;
38+ const { nlat, nphi } = gridForLmax(LMAX, 3);
39+ const cfg = { lmax: LMAX, mmax: LMAX, nlat, nphi };
40+ const sht = await ShtPlan.create(device, cfg);
41+ const geometry = await Geometry.create({
42+ device,
43+ sht,
44+ cfg,
45+ source: g.source,
46+ paramNames: g.params.map((p) => p.key),
47+ params: defaultGeometryParams(g),
48+ });
49+ return { g, sht, cfg, geometry };
50+}
51+
52+export async function geometryChecks(
53+ device: GPUDevice,
54+ check: Check,
55+ log: Log,
56+): Promise<void> {
57+ // ---- every geometry compiles and closes ---------------------------------
58+ for (const spec of mGeometries) {
59+ const { sht, geometry } = await buildGeometry(device, spec.key);
60+ let finite = true;
61+ for (const a of [geometry.x, geometry.y, geometry.z]) {
62+ for (const v of a) if (!Number.isFinite(v)) finite = false;
63+ }
64+ const { lo, hi } = geometry.radiusRange();
65+ check(
66+ `geometry: ${spec.key}.m evaluates to a finite surface`,
67+ finite && lo > 1e-3,
68+ `radius ${lo.toFixed(4)}–${hi.toFixed(4)}`,
69+ );
70+ sht.destroy();
71+ }
72+
73+ // ---- the sphere is the unit sphere, exactly, and is degree 1 ------------
74+ {
75+ const { sht, geometry } = await buildGeometry(device, SPHERE_KEY);
76+
77+ let maxRadiusErr = 0;
78+ for (let i = 0; i < geometry.x.length; i++) {
79+ const r = Math.hypot(geometry.x[i], geometry.y[i], geometry.z[i]);
80+ maxRadiusErr = Math.max(maxRadiusErr, Math.abs(r - 1));
81+ }
82+ // Tolerance is fp32 through a full analysis/synthesis round trip, not the
83+ // geometry: the exact answer is representable, and what is measured here
84+ // is the transforms' own round-off. It is set by the loosest stack this
85+ // runs on — SwiftShader in CI is an order of magnitude worse than Dawn on
86+ // real hardware (4e-4 against 2e-5). A geometry that was actually wrong
87+ // would miss by O(1), so the slack costs nothing.
88+ check(
89+ 'geometry: sphere.m has radius 1 everywhere',
90+ maxRadiusErr < 2e-3,
91+ `max |r - 1| = ${maxRadiusErr.toExponential(2)}`,
92+ );
93+
94+ // x, y, z of the unit sphere are the three degree-1 harmonics and nothing
95+ // else, so analysing them must leave every other coefficient at zero.
96+ // This is what makes the sphere case exact rather than merely accurate:
97+ // there is no content for the band limit to throw away.
98+ const degreeOne = new Set([
99+ lmIndex(LMAX, 1, 0),
100+ lmIndex(LMAX, 1, 1),
101+ ]);
102+ let leak = 0;
103+ for (const coeffs of [geometry.X, geometry.Y, geometry.Z]) {
104+ for (let i = 0; i < coeffs.length / 2; i++) {
105+ if (degreeOne.has(i)) continue;
106+ leak = Math.max(leak, Math.abs(coeffs[2 * i]), Math.abs(coeffs[2 * i + 1]));
107+ }
108+ }
109+ check(
110+ 'geometry: sphere.m is exactly degree 1 in the harmonics',
111+ leak < 1e-3,
112+ `max |coefficient| outside l = 1 is ${leak.toExponential(2)}`,
113+ );
114+ sht.destroy();
115+ }
116+
117+ // ---- a deformed surface matches its own formula, on any grid ------------
118+ {
119+ const { g, sht, cfg, geometry } = await buildGeometry(device, 'peanut');
120+ const p = defaultGeometryParams(g);
121+
122+ // peanut.m written out: r = 1 - waist*sin(theta)^2 scales the unit sphere,
123+ // and z is then stretched, so the distance from the origin depends on
124+ // theta alone. Checking every point against this closed form checks the
125+ // whole path at once — the compiled shape kernel, the analysis into
126+ // coefficients, the synthesis back — and, because the formula has no phi
127+ // in it, that the surface really is a surface of revolution.
128+ const peanutRadius = (ct: number): number => {
129+ const st2 = Math.max(0, 1 - ct * ct);
130+ const r = 1 - p.waist * st2;
131+ return r * Math.hypot(Math.sqrt(st2), (1 + p.stretch) * ct);
132+ };
133+
134+ const onGrid = (
135+ cosTheta: Float64Array,
136+ nlat: number,
137+ nphi: number,
138+ at: (i: number) => number,
139+ ): number => {
140+ let worst = 0;
141+ for (let i = 0; i < nlat; i++) {
142+ const want = peanutRadius(cosTheta[i]);
143+ for (let j = 0; j < nphi; j++) {
144+ worst = Math.max(worst, Math.abs(at(i * nphi + j) - want));
145+ }
146+ }
147+ return worst;
148+ };
149+
150+ const coarse = onGrid(sht.cosTheta, cfg.nlat, cfg.nphi, (k) =>
151+ Math.hypot(geometry.x[k], geometry.y[k], geometry.z[k]),
152+ );
153+ check(
154+ 'geometry: peanut.m matches its own radial formula on the solver grid',
155+ coarse < 1e-3,
156+ `max |dr| = ${coarse.toExponential(2)}`,
157+ );
158+
159+ // And the same on a finer grid, from the same coefficients. This is what
160+ // "the rendered surface is the surface being solved on" means: display
161+ // oversampling evaluates the embedding at more points, it does not
162+ // subdivide or smooth it. The 2x Gauss latitudes share no point with the
163+ // 1x ones, so agreeing here is agreeing everywhere, not at samples.
164+ const fine = await ShtPlan.create(device, {
165+ lmax: cfg.lmax,
166+ mmax: cfg.mmax,
167+ nlat: 2 * cfg.nlat,
168+ nphi: 2 * cfg.nphi,
169+ });
170+ const finePos = await geometry.positionsOn(fine);
171+ const refined = onGrid(fine.cosTheta, 2 * cfg.nlat, 2 * cfg.nphi, (k) =>
172+ Math.hypot(finePos[3 * k], finePos[3 * k + 1], finePos[3 * k + 2]),
173+ );
174+ check(
175+ 'geometry: the same coefficients give the same surface on a 2x grid',
176+ refined < 1e-3,
177+ `max |dr| = ${refined.toExponential(2)} at ${2 * cfg.nlat}×${2 * cfg.nphi} points`,
178+ );
179+ fine.destroy();
180+ sht.destroy();
181+ }
182+
183+ // ---- the unrolled loop: more ops, identical answer ----------------------
184+ {
185+ const model = mModelByKey('schnakenberg')!;
186+ const params = defaultParams(model);
187+ const counts = [0, 1, 4];
188+ const ops: number[] = [];
189+ const states: Float32Array[] = [];
190+
191+ for (const niter of counts) {
192+ const session = await ModelSession.create({
193+ device, model, params, lmax: LMAX, niter,
194+ });
195+ ops.push(session.describe().step.length);
196+ session.seed(1);
197+ session.step(STEPS);
198+ states.push(await session.read('U'));
199+ session.destroy();
200+ }
201+
202+ log(` schnakenberg.m ops/step by solve iterations: ${
203+ counts.map((n, i) => `${n} -> ${ops[i]}`).join(', ')
204+ }`);
205+ check(
206+ 'loop: each solve iteration adds GPU operations',
207+ ops[0] < ops[1] && ops[1] < ops[2],
208+ `${ops.join(' < ')} ops for ${counts.join(', ')} iterations`,
209+ );
210+ // Unrolling has to be exactly linear in the trip count: the body planned
211+ // once per iteration, no more and no less. Two dispatches per species per
212+ // iteration — the placeholder line and the update that reads it.
213+ const perIteration = ops[1] - ops[0];
214+ const want = 2 * model.species.length;
215+ check(
216+ 'loop: unrolling is exactly linear in the trip count',
217+ perIteration === want && ops[2] - ops[0] === 4 * perIteration,
218+ `${perIteration} ops per iteration (expected ${want}), ` +
219+ `${ops[2] - ops[0]} for 4 iterations`,
220+ );
221+
222+ let identical = true;
223+ let worst = 0;
224+ for (let k = 1; k < states.length; k++) {
225+ if (states[k].length !== states[0].length) identical = false;
226+ for (let i = 0; i < states[0].length; i++) {
227+ if (states[k][i] !== states[0][i]) identical = false;
228+ worst = Math.max(worst, Math.abs(states[k][i] - states[0][i]));
229+ }
230+ }
231+ check(
232+ 'loop: the geometry correction is exactly zero, so the answer does not move',
233+ identical,
234+ identical
235+ ? `bit-identical after ${STEPS} steps at ${counts.join('/')} iterations`
236+ : `states differ by up to ${worst.toExponential(2)}`,
237+ );
238+ }
239+
240+ // ---- a loop whose length is not known at compile time is refused --------
241+ {
242+ const model = mModelByKey('allencahn')!;
243+ // `dt` is a tunable parameter, so it reaches the compiler with no value:
244+ // the plan cannot know how many iterations to emit.
245+ const bad = model.source.replace('for k = 1:niter', 'for k = 1:dt');
246+ let message = '';
247+ try {
248+ const session = await ModelSession.create({
249+ device, model, params: defaultParams(model), lmax: LMAX, source: bad, niter: 1,
250+ });
251+ session.destroy();
252+ } catch (e) {
253+ message = e instanceof ModelCompileError ? e.message : `wrong error type: ${e}`;
254+ }
255+ check(
256+ 'loop: a runtime loop bound is refused at compile time',
257+ message.includes('known when the model is compiled'),
258+ message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
259+ );
260+ }
261+
262+ // ---- swapping the surface leaves the simulation alone ------------------
263+ {
264+ const model = mModelByKey('schnakenberg')!;
265+ const session = await ModelSession.create({
266+ device, model, params: defaultParams(model), lmax: LMAX,
267+ });
268+ session.seed(1);
269+ session.step(STEPS);
270+ const before = await session.read('U');
271+
272+ const peanut = mGeometryByKey('peanut')!;
273+ await session.setGeometry(peanut, defaultGeometryParams(peanut));
274+ const after = await session.read('U');
275+
276+ let survived = before.length === after.length;
277+ for (let i = 0; survived && i < before.length; i++) {
278+ if (before[i] !== after[i]) survived = false;
279+ }
280+ const { lo, hi } = session.geometry.radiusRange();
281+ check(
282+ 'geometry: swapping the surface mid-run does not disturb the state',
283+ survived && session.geometryModel.key === 'peanut' && hi - lo > 0.1,
284+ survived
285+ ? `state identical, now on ${session.geometryModel.key} (radius ${lo.toFixed(3)}–${hi.toFixed(3)})`
286+ : 'state changed',
287+ );
288+ session.destroy();
289+ }
290+}
291+
292+/** Index of the entry minimizing `score`, over the first `n` entries. */
293+function argMin(
294+ xs: Float64Array | Float32Array,
295+ n: number,
296+ score: (v: number) => number,
297+): number {
298+ let best = 0;
299+ let bestScore = Infinity;
300+ for (let i = 0; i < n; i++) {
301+ const s = score(xs[i]);
302+ if (s < bestScore) {
303+ bestScore = s;
304+ best = i;
305+ }
306+ }
307+ return best;
308+}
test/modelChecks.tsadded+222−0View file
@@ -0,0 +1,222 @@
1+/**
2+ * Every model the app offers: that it compiles, what it compiles to, and that it
3+ * runs stably and produces a pattern.
4+ *
5+ * Numerical correctness of the pipeline is analyticChecks.ts's job. This file is
6+ * about the models themselves and about the compilation staying as intended — in
7+ * particular the kernel count, which is a fusion guard: numbl's lowering emits
8+ * one statement per *operator*, and its inline pass folds those back into
9+ * per-line expression trees. If that stops happening the results stay correct
10+ * but every operator becomes its own dispatch, which is invisible except here.
11+ */
12+import { ModelSession } from '../src/mgpu/session.ts';
13+import { mModels, defaultParams } from '../src/mgpu/registry.ts';
14+import {
15+ formatCommand,
16+ parseArgs,
17+ BENCH_COMMAND,
18+ type RunSpec,
19+} from '../src/bench/runSpec.ts';
20+import type { Check, Log } from './analyticChecks.ts';
21+
22+/**
23+ * Kernels each model's step compiles to outside its solve loop — one per
24+ * element-wise line, where the argument of a transform counts as its own line
25+ * (it cannot fuse into an external call).
26+ */
27+const EXPECTED_KERNELS: Record<string, number> = {
28+ schnakenberg: 7,
29+ brusselator: 7,
30+ allencahn: 3,
31+};
32+
33+/**
34+ * And what one unrolled iteration of the solve loop adds, per species: the
35+ * placeholder line that will become the geometry correction, and the update
36+ * that consumes it. Two rather than one because the correction does not fuse
37+ * into its consumer — which is right, since the operator that replaces it will
38+ * be transforms and kernels of its own, not an expression.
39+ */
40+const KERNELS_PER_ITERATION = 2;
41+
42+const LMAX = 31;
43+const STEPS = 40;
44+const NITER = 1;
45+
46+export async function modelChecks(
47+ device: GPUDevice,
48+ check: Check,
49+ log: Log,
50+): Promise<void> {
51+ check('models: registry populated', mModels.length === 3, `${mModels.length} models`);
52+
53+ // The app formats the run it is showing into a `npm run bench` command and
54+ // the benchmark parses it back. That is only worth anything if the round
55+ // trip is lossless — a knob that formatCommand forgets is a knob the desktop
56+ // run would silently take a default for, and the two runs would differ while
57+ // claiming to be the same. Every field of the spec, through both directions.
58+ {
59+ const spec: RunSpec = {
60+ preset: 'schnak-fine',
61+ lmax: 127,
62+ seed: 12345,
63+ steps: 777,
64+ warmup: 13,
65+ params: { a: 0.11, b: 0.91, D1: 5e-4, D2: 9e-3, dt: 0.04 },
66+ geometry: 'peanut',
67+ geometryParams: { waist: 0.45, stretch: 1.25 },
68+ niter: 3,
69+ };
70+ const command = formatCommand(spec);
71+ const back = parseArgs(command.slice(BENCH_COMMAND.length).trim().split(/\s+/));
72+ const same = JSON.stringify(back) === JSON.stringify(spec);
73+ check(
74+ 'runSpec: the benchmark command round-trips every field',
75+ same,
76+ same ? command.slice(BENCH_COMMAND.length + 1) : `got ${JSON.stringify(back)}`,
77+ );
78+ }
79+
80+ for (const model of mModels) {
81+ const session = await ModelSession.create({
82+ device,
83+ model,
84+ params: defaultParams(model),
85+ lmax: LMAX,
86+ niter: NITER,
87+ });
88+
89+ const plan = session.describe();
90+ const kernels = plan.step.filter((l) => l.startsWith('kernel')).length;
91+ const xforms = plan.step.filter(
92+ (l) => l.startsWith('synth') || l.startsWith('analys'),
93+ ).length;
94+ const expected =
95+ EXPECTED_KERNELS[model.key] +
96+ NITER * KERNELS_PER_ITERATION * model.species.length;
97+ log(
98+ ` ${model.key}.m -> ${plan.step.length} ops/step ` +
99+ `(${kernels} generated kernels, ${xforms} transforms, ${NITER} solve iter)`,
100+ );
101+ check(
102+ `${model.key}: element-wise lines fused into one kernel each`,
103+ kernels === expected,
104+ `${kernels} kernels (expected ${expected})`,
105+ );
106+
107+ session.seed(1);
108+ session.step(STEPS);
109+
110+ // Every rendered field must be finite and have developed some contrast.
111+ for (const field of model.species) {
112+ const values = await session.read(field);
113+ let lo = Infinity;
114+ let hi = -Infinity;
115+ let finite = true;
116+ for (const v of values) {
117+ if (!Number.isFinite(v)) finite = false;
118+ if (v < lo) lo = v;
119+ if (v > hi) hi = v;
120+ }
121+ check(
122+ `${model.key}: '${field}' is finite and patterned after ${STEPS} steps`,
123+ finite && hi - lo > 1e-6,
124+ finite
125+ ? `range [${lo.toFixed(5)}, ${hi.toFixed(5)}]`
126+ : 'contains NaN or Infinity',
127+ );
128+ }
129+
130+ session.destroy();
131+ }
132+
133+ // The oversampled readback: readSpecies must be the state synthesized on the
134+ // display grid. Comparing against the display plan's own upload path
135+ // (read the state back, synth it from the CPU) exercises the GPU-to-GPU
136+ // coefficient copy against a known-good route through the same kernels.
137+ {
138+ const model = mModels.find((m) => m.key === 'allencahn')!;
139+ const session = await ModelSession.create({
140+ device,
141+ model,
142+ params: defaultParams(model),
143+ lmax: LMAX,
144+ oversample: 2,
145+ });
146+ session.seed(1);
147+ session.step(STEPS);
148+
149+ const fine = await session.readSpecies(0);
150+ const { nlat, nphi } = session.viewSht.cfg;
151+ check(
152+ 'oversample: species field is on the 2x display grid',
153+ nlat === 2 * session.cfg.nlat &&
154+ nphi === 2 * session.cfg.nphi &&
155+ fine.length === nlat * nphi,
156+ `render ${nlat}×${nphi}, ${fine.length} values`,
157+ );
158+
159+ const qlm = await session.read('U');
160+ const expected = await session.viewSht.synth(qlm);
161+ let maxDiff = 0;
162+ for (let i = 0; i < fine.length; i++) {
163+ const d = Math.abs(fine[i] - expected[i]);
164+ if (d > maxDiff) maxDiff = d;
165+ }
166+ check(
167+ 'oversample: readSpecies matches synth of the read-back state',
168+ maxDiff <= 1e-6,
169+ `max |diff| = ${maxDiff.toExponential(2)}`,
170+ );
171+
172+ // A timing burst must be invisible: the state is snapshotted and restored
173+ // around it, and model time does not advance.
174+ const tBefore = session.t;
175+ const stepsBefore = session.steps;
176+ const ms = await session.measure(8);
177+ const after = await session.read('U');
178+ let identical = qlm.length === after.length;
179+ if (identical) {
180+ for (let i = 0; i < qlm.length; i++) {
181+ if (qlm[i] !== after[i]) {
182+ identical = false;
183+ break;
184+ }
185+ }
186+ }
187+ check(
188+ 'measure: a timing burst leaves state, t and steps untouched',
189+ identical && session.t === tBefore && session.steps === stepsBefore,
190+ identical
191+ ? `state identical, t = ${session.t.toFixed(3)}, ${ms.toFixed(3)} ms/step`
192+ : 'state changed',
193+ );
194+
195+ // Changing the oversampling in place is display-only: the state survives
196+ // and the render grid drops back to the solver's.
197+ await session.setOversample(1);
198+ const qlmAfterSwap = await session.read('U');
199+ let stateSurvived = qlmAfterSwap.length === after.length;
200+ if (stateSurvived) {
201+ for (let i = 0; i < after.length; i++) {
202+ if (qlmAfterSwap[i] !== after[i]) {
203+ stateSurvived = false;
204+ break;
205+ }
206+ }
207+ }
208+ session.step(1); // recompute the view fields on the solver grid
209+ const coarse = await session.readSpecies(0);
210+ check(
211+ 'setOversample: swaps the render grid without touching the state',
212+ stateSurvived &&
213+ session.viewSht === session.sht &&
214+ coarse.length === session.cfg.nlat * session.cfg.nphi,
215+ stateSurvived
216+ ? `state survived, render back to ${session.cfg.nlat}×${session.cfg.nphi}`
217+ : 'state changed',
218+ );
219+
220+ session.destroy();
221+ }
222+}
test/models/linear.madded+19−0View file
@@ -0,0 +1,19 @@
1+% Test model: a purely linear reaction, f(u) = c*u.
2+%
3+% Every spherical-harmonic mode then evolves independently under one IMEX Euler
4+% step, with a closed-form growth factor per degree l:
5+%
6+% U_lm^{n+1} = U_lm^n * (1 + dt*c) / (1 + dt*D*l(l+1))
7+%
8+% so a run can be checked against exact arithmetic rather than against another
9+% implementation. Used by the analytic tests; not offered in the app.
10+
11+function [U, u] = init(noise)
12+ U = analys(noise);
13+ u = synth(U);
14+end
15+
16+function [Un, u] = step(U, lam, c, D, dt)
17+ u = synth(U);
18+ Un = (U + dt * analys(c * u)) ./ (1 + (dt * D) * lam);
19+end
test/models/logistic.madded+24−0View file
@@ -0,0 +1,24 @@
1+% Test model: a nonlinear reaction, f(u) = r*u*(1 - u).
2+%
3+% Started from a *uniform* field, the state stays uniform, and diffusion does
4+% nothing to it (the l = 0 eigenvalue is zero). So every step is exactly the
5+% explicit Euler map of the scalar ODE,
6+%
7+% u^{n+1} = u^n + dt*r*u^n*(1 - u^n)
8+%
9+% which checks that the generated kernel evaluates a nonlinear reaction
10+% correctly, against arithmetic rather than another implementation. Used by the
11+% analytic tests; not offered in the app.
12+%
13+% The caller passes the initial grid field as `noise` (the name the app uses for
14+% its seeded perturbation); here the tests put an exact field there.
15+
16+function [U, u] = init(noise)
17+ U = analys(noise);
18+ u = synth(U);
19+end
20+
21+function [Un, u] = step(U, lam, r, D, dt)
22+ u = synth(U);
23+ Un = (U + dt * analys(r * u .* (1 - u))) ./ (1 + (dt * D) * lam);
24+end
test/test-page.tsadded+228−0View file
@@ -0,0 +1,228 @@
1+/**
2+ * Browser validation, in the environment the demo actually ships to.
3+ *
4+ * Runs the same four check modules as `npm run test:node` — so both GPU stacks
5+ * (Dawn on the desktop, the browser's own here) get the same guarantees — plus a
6+ * long soak that only makes sense in a page.
7+ *
8+ * Results are posted to window.__RESULTS__ for the headless runner.
9+ */
10+import { requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
11+import { ModelSession } from '../src/mgpu/session.ts';
12+import { mModels, defaultParams } from '../src/mgpu/registry.ts';
13+import { digestOf, formatDigest, type StateDigest } from '../src/mgpu/digest.ts';
14+import {
15+ parseArgs,
16+ modelForSpec,
17+ geometryForSpec,
18+ formatCommand,
19+ DEFAULT_NITER,
20+} from '../src/bench/runSpec.ts';
21+import {
22+ mGeometryByKey,
23+ defaultGeometryParams,
24+ DEFAULT_GEOMETRY_KEY,
25+} from '../src/geom/registry.ts';
26+import { transformChecks } from './transformChecks.ts';
27+import { analyticChecks } from './analyticChecks.ts';
28+import { modelChecks } from './modelChecks.ts';
29+import { geometryChecks } from './geometryChecks.ts';
30+
31+declare global {
32+ interface Window {
33+ __RESULTS__?: { ok: boolean; fatal?: string; lines: string[] };
34+ /** Set by the ?state= mode, for scripts/compare-env.mjs. */
35+ __STATE__?: { digest: StateDigest; state: number[] };
36+ /** Set by the ?soak= mode, for scripts/compare-perf.mjs. */
37+ __SOAK__?: {
38+ lmax: number;
39+ steps: number;
40+ batch: number;
41+ solverMsPerStep: number;
42+ encodeMsPerStep: number;
43+ adapter: string;
44+ fourier: 'fft' | 'dft';
45+ };
46+ }
47+}
48+
49+const logEl = document.getElementById('log')!;
50+const lines: string[] = [];
51+let failures = 0;
52+
53+function log(s: string): void {
54+ lines.push(s);
55+ logEl.textContent = lines.join('\n');
56+ console.log(s);
57+}
58+
59+function check(name: string, ok: boolean, detail: string): void {
60+ log(`${ok ? 'PASS' : 'FAIL'} ${name} ${detail}`);
61+ if (!ok) failures++;
62+}
63+
64+/**
65+ * Solver-only soak, selected with ?soak=<steps>&lmax=<n>.
66+ *
67+ * No three.js at all, so this is the browser's honest solver rate: the same
68+ * batched, no-readback measurement the desktop benchmark reports. If this number
69+ * matches the benchmark's but the app's frame cost does not, the difference is
70+ * the readback and competing with the renderer for the GPU, not the computation.
71+ */
72+async function soak(steps: number, lmax: number): Promise<void> {
73+ const device = await requestShtDevice();
74+ const model = mModels[0];
75+ // The desktop benchmark this is compared against resolves its geometry and
76+ // iteration count from the same two constants. They have to agree: the
77+ // iteration count is unrolled into the step, so a mismatch would compare
78+ // two different amounts of work and call the difference "the browser".
79+ const geometry = mGeometryByKey(DEFAULT_GEOMETRY_KEY)!;
80+ const session = await ModelSession.create({
81+ device,
82+ model,
83+ params: defaultParams(model),
84+ lmax,
85+ geometry,
86+ geometryParams: defaultGeometryParams(geometry),
87+ niter: DEFAULT_NITER,
88+ });
89+ session.seed(5);
90+ log(
91+ `soak: ${steps} steps at lmax ${lmax} ` +
92+ `(grid ${session.cfg.nlat}x${session.cfg.nphi}, ${geometry.key}, ` +
93+ `${DEFAULT_NITER} solve iter, ${session.describe().step.length} ops/step), solver only`,
94+ );
95+
96+ const BATCH = 25;
97+ // Timed separately from the sampling: `solverMs` counts only submitted steps
98+ // waited for, never read back, so it is comparable to `npm run bench`.
99+ let solverMs = 0;
100+ let solverSteps = 0;
101+ let encodeMs = 0;
102+ const t0 = performance.now();
103+ for (let s = 0; s < steps; s += BATCH) {
104+ const n = Math.min(BATCH, steps - s);
105+ const b0 = performance.now();
106+ session.step(n);
107+ // CPU-side command encoding, separated from GPU execution: in a browser each
108+ // WebGPU call crosses Blink's bindings and Dawn's validation, so on a fast
109+ // GPU the encoding can be what actually limits the step rate.
110+ const b1 = performance.now();
111+ encodeMs += b1 - b0;
112+ await session.sync();
113+ solverMs += performance.now() - b0;
114+ solverSteps += n;
115+ const u = await session.read(model.species[0]);
116+ let lo = Infinity;
117+ let hi = -Infinity;
118+ for (const v of u) {
119+ if (v < lo) lo = v;
120+ if (v > hi) hi = v;
121+ }
122+ if ((s + BATCH) % 100 === 0) {
123+ const mem = (performance as Performance & { memory?: { usedJSHeapSize: number } })
124+ .memory;
125+ log(
126+ ` step ${session.steps} u in [${lo.toFixed(4)}, ${hi.toFixed(4)}]` +
127+ (mem ? ` heap ${(mem.usedJSHeapSize / 1048576).toFixed(1)} MB` : ''),
128+ );
129+ // yield so the page stays responsive and the runner can poll
130+ await new Promise((r) => setTimeout(r, 0));
131+ }
132+ }
133+ const ms = (performance.now() - t0) / steps;
134+
135+ const final = await session.read(model.species[0]);
136+ let finite = true;
137+ for (const v of final) if (!Number.isFinite(v)) finite = false;
138+ const solverPerStep = solverMs / solverSteps;
139+ const encodePerStep = encodeMs / solverSteps;
140+ check(
141+ `soak: ${steps} steps survived`,
142+ finite,
143+ `solver ${solverPerStep.toFixed(2)} ms/step (batches of ${BATCH}, no readback), ` +
144+ `of which ${encodePerStep.toFixed(3)} ms/step CPU encoding, ` +
145+ `${ms.toFixed(2)} ms/step incl. sampling readback`,
146+ );
147+ log(
148+ ` compare 'solver' with the ms/step from \`npm run bench -- --lmax ${lmax}\`:\n` +
149+ ` same .m, same kernels, no rendering on either side.`,
150+ );
151+
152+ window.__SOAK__ = {
153+ lmax,
154+ steps,
155+ batch: BATCH,
156+ solverMsPerStep: solverPerStep,
157+ encodeMsPerStep: encodePerStep,
158+ adapter: await describeAdapter(device),
159+ fourier: session.sht.fourierMode,
160+ };
161+ session.destroy();
162+ window.__RESULTS__ = { ok: failures === 0, lines };
163+ log(failures === 0 ? 'ALL PASS' : `${failures} FAILURE(S)`);
164+}
165+
166+/**
167+ * Run one exact spec and post its final state, for scripts/compare-env.mjs to
168+ * compare against the same spec run on the desktop. Query parameters map
169+ * straight onto the benchmark's flags — `?state=1&lmax=31&steps=200` — and go
170+ * through the same parseArgs, so neither side can quietly use different
171+ * defaults.
172+ */
173+async function dumpState(q: URLSearchParams): Promise<void> {
174+ const argv: string[] = [];
175+ for (const [k, v] of q) {
176+ if (k === 'state') continue;
177+ argv.push(`--${k}`, v);
178+ }
179+ const spec = parseArgs(argv);
180+ const model = modelForSpec(spec);
181+
182+ const device = await requestShtDevice();
183+ const adapter = await describeAdapter(device);
184+ const session = await ModelSession.create({
185+ device,
186+ model,
187+ params: spec.params,
188+ lmax: spec.lmax,
189+ geometry: geometryForSpec(spec),
190+ geometryParams: spec.geometryParams,
191+ niter: spec.niter,
192+ });
193+ session.seed(spec.seed);
194+ session.step(spec.steps);
195+ await session.sync();
196+ const state = await session.read(model.state[0]);
197+ const digest = digestOf(state, session.sht.fourierMode, adapter);
198+
199+ log(`${formatCommand(spec)}\n`);
200+ log(`state after ${spec.steps} steps from seed ${spec.seed}:`);
201+ log(` ${formatDigest(digest)}`);
202+ log(` adapter: ${adapter}`);
203+ window.__STATE__ = { digest, state: [...state] };
204+ session.destroy();
205+}
206+
207+async function main(): Promise<void> {
208+ const q = new URLSearchParams(location.search);
209+ if (q.has('state')) return dumpState(q);
210+ if (q.has('soak')) {
211+ return soak(Number(q.get('soak')) || 500, Number(q.get('lmax')) || 63);
212+ }
213+ const device = await requestShtDevice();
214+
215+ await transformChecks(device, check, log);
216+ await analyticChecks(device, check, log);
217+ await modelChecks(device, check, log);
218+ await geometryChecks(device, check, log);
219+
220+ window.__RESULTS__ = { ok: failures === 0, lines };
221+ log(failures === 0 ? 'ALL PASS' : `${failures} FAILURE(S)`);
222+}
223+
224+main().catch((e) => {
225+ const msg = e instanceof Error ? `${e.message}\n${e.stack ?? ''}` : String(e);
226+ log(`fatal: ${msg}`);
227+ window.__RESULTS__ = { ok: false, fatal: msg, lines };
228+});
test/transformChecks.tsadded+55−0View file
@@ -0,0 +1,55 @@
1+/**
2+ * The WGSL spherical-harmonic transforms against the f64 CPU reference.
3+ *
4+ * This is the one place a second implementation is still the right oracle: the
5+ * transforms are vendored shtns-webgpu, and `src/sht/reference.ts` is its direct-
6+ * summation f64 twin. Everything above them (the .m models) is checked against
7+ * closed-form answers instead — see analyticChecks.ts.
8+ */
9+import { ShtPlan } from '../src/sht/sht.ts';
10+import { ShtReference, randomSpectrum } from '../src/sht/reference.ts';
11+import { gridForLmax } from '../src/sht/layout.ts';
12+import type { Check, Log } from './analyticChecks.ts';
13+
14+function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
15+ let num = 0;
16+ let den = 0;
17+ for (let i = 0; i < a.length; i++) {
18+ const d = a[i] - b[i];
19+ num += d * d;
20+ den += b[i] * b[i];
21+ }
22+ return Math.sqrt(num / Math.max(den, 1e-300));
23+}
24+
25+export async function transformChecks(
26+ device: GPUDevice,
27+ check: Check,
28+ _log: Log,
29+): Promise<void> {
30+ const lmax = 31;
31+ const { nlat, nphi } = gridForLmax(lmax, 1);
32+ const cfg = { lmax, mmax: lmax, nlat, nphi };
33+
34+ const plan = await ShtPlan.create(device, cfg);
35+ const ref = new ShtReference(cfg);
36+
37+ const q = randomSpectrum(cfg, 42);
38+ const q64 = new Float64Array(q);
39+
40+ const spatGpu = await plan.synth(new Float32Array(q64));
41+ const spatCpu = ref.synth(q64);
42+ const errSynth = relL2(spatGpu, spatCpu);
43+
44+ const qGpu = await plan.analys(new Float32Array(spatCpu));
45+ const qCpu = ref.analys(new Float64Array(spatCpu));
46+ const errAnalys = relL2(qGpu, qCpu);
47+
48+ check(
49+ 'transforms: WGSL fp32 vs f64 CPU reference',
50+ errSynth < 1e-4 && errAnalys < 1e-4,
51+ `synth ${errSynth.toExponential(2)}, analys ${errAnalys.toExponential(2)}`,
52+ );
53+
54+ plan.destroy();
55+}
tsconfig.jsonadded+15−0View file
@@ -0,0 +1,15 @@
1+{
2+ "compilerOptions": {
3+ "target": "ES2022",
4+ "module": "ESNext",
5+ "moduleResolution": "bundler",
6+ "lib": ["ES2022", "DOM", "DOM.Iterable"],
7+ "types": ["@webgpu/types", "node"],
8+ "strict": true,
9+ "noEmit": true,
10+ "allowImportingTsExtensions": true,
11+ "verbatimModuleSyntax": true,
12+ "skipLibCheck": true
13+ },
14+ "include": ["src", "test", "scripts"]
15+}
vite.config.tsadded+31−0View file
@@ -0,0 +1,31 @@
1+import { defineConfig } from 'vite';
2+import { resolve } from 'node:path';
3+
4+// numbl is a local `file:` dependency, so node_modules/numbl is a symlink to
5+// the sibling checkout. Its package `exports` map only publishes the runtime
6+// entry points, not the compiler internals we need (parser + JIT lowering), so
7+// we reach them through a path alias. (package.json's `imports` field cannot
8+// express this — Node rejects node_modules targets — and plain Node could not
9+// resolve numbl's internal `.js`->`.ts` imports anyway, which is why the GPU
10+// tests run in the browser harness rather than under `node`.)
11+const numblSrc = resolve(import.meta.dirname, 'node_modules/numbl/src');
12+
13+export default defineConfig({
14+ base: './',
15+ resolve: {
16+ alias: { 'numbl-src': numblSrc },
17+ },
18+ server: {
19+ // the alias resolves outside the project root (through the symlink)
20+ fs: { allow: [import.meta.dirname, numblSrc] },
21+ },
22+ build: {
23+ target: 'es2022',
24+ rollupOptions: {
25+ input: {
26+ main: resolve(import.meta.dirname, 'index.html'),
27+ test: resolve(import.meta.dirname, 'test.html'),
28+ },
29+ },
30+ },
31+});