turing-surface-cache: reaction-diffusion solutions at a chosen end time, shared through a cloud cache
Trimmed fork of turing-surface: Schnakenberg (flux form) on sphere/ellipsoid/
peanut, every setting a choice from a short discrete list, so each combination
names one exact solution. The spec's canonical JSON is SHA-256-hashed into an
object name on tempory.net; cached solutions load as the choices are browsed,
uncached ones show empty surfaces until Compute solution runs the solver
locally on WebGPU. Runs warm-start from the longest cached shorter run of the
same spec, capture intermediate end-time states, and (with a tmpbucket API
key) upload them in the background while stepping continues. Cache files are
HDF5 in turing-surface's reference-file layout plus the spec identity, so
they load in its compare mode and read with h5py. The selection is mirrored
in the URL fragment for reload and sharing.
56 changed files+13173−0
.github/workflows/ci.ymladded+48−0View file
@@ -0,0 +1,48 @@
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) and its interpreter, which its
18+ # published package `exports` do not expose. Clone it where that relative
19+ # path expects it. Pinned so a change to those internals cannot silently
20+ # break the build — the surface we rely on is written down in
21+ # src/mgpu/numbl.d.ts.
22+ #
23+ # numbl's own dependencies are NOT needed: the slice we import is
24+ # self-contained TypeScript, verified by building against a checkout with no
25+ # node_modules. One file in it is generated rather than committed, though —
26+ # the interpreter's stdlib bundle, which numbl gitignores — so a bare
27+ # checkout is missing it and `executeCode.ts` fails to resolve it. Its
28+ # generator only reads .m files off disk, so plain `node` (type-stripping,
29+ # unflagged since 22.18) runs it without installing anything.
30+ - name: Check out numbl (sibling dependency)
31+ env:
32+ NUMBL_REF: main
33+ run: |
34+ git clone --filter=blob:none --no-checkout \
35+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
36+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
37+ node "$GITHUB_WORKSPACE/../../numbl/scripts/bundle-stdlib.ts"
38+ # --ignore-scripts: npm runs a linked package's `prepare` script, and
39+ # numbl's is husky, which is not installed here.
40+ - run: npm ci --ignore-scripts
41+ - run: npm run build
42+ # End-to-end check of the miss/compute/encode and hit/decode paths in
43+ # headless Chrome on SwiftShader WebGPU, with the cloud cache mocked;
44+ # the produced .h5 is verified with h5py.
45+ - run: pip install h5py
46+ - run: node scripts/check-app.mjs
47+ env:
48+ CHROME_PATH: /usr/bin/google-chrome
.github/workflows/deploy.ymladded+62−0View file
@@ -0,0 +1,62 @@
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) and its interpreter, which its
30+ # published package `exports` do not expose. Clone it where that relative
31+ # path expects it. Pinned so a change to those internals cannot silently
32+ # break the build — the surface we rely on is written down in
33+ # src/mgpu/numbl.d.ts.
34+ #
35+ # numbl's own dependencies are NOT needed: the slice we import is
36+ # self-contained TypeScript, verified by building against a checkout with no
37+ # node_modules. One file in it is generated rather than committed, though —
38+ # the interpreter's stdlib bundle, which numbl gitignores — so a bare
39+ # checkout is missing it and `executeCode.ts` fails to resolve it. Its
40+ # generator only reads .m files off disk, so plain `node` (type-stripping,
41+ # unflagged since 22.18) runs it without installing anything.
42+ - name: Check out numbl (sibling dependency)
43+ env:
44+ NUMBL_REF: main
45+ run: |
46+ git clone --filter=blob:none --no-checkout \
47+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
48+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
49+ node "$GITHUB_WORKSPACE/../../numbl/scripts/bundle-stdlib.ts"
50+ # --ignore-scripts: npm runs a linked package's `prepare` script, and
51+ # numbl's is husky, which is not installed here.
52+ - run: npm ci --ignore-scripts
53+ - run: npm run build
54+ # Pages must already be enabled with "GitHub Actions" as the source; the
55+ # workflow token cannot create the site itself (`enablement: true` fails
56+ # with "Resource not accessible by integration").
57+ - uses: actions/configure-pages@v5
58+ - uses: actions/upload-pages-artifact@v3
59+ with:
60+ path: dist
61+ - id: deployment
62+ uses: actions/deploy-pages@v4
.github/workflows/preview.ymladded+153−0View file
@@ -0,0 +1,153 @@
1+## Comment `/preview` on a pull request to build it and publish the result to
2+## https://tempory.net/previews/<repo>/<branch>/index.html
3+##
4+## This file is meant to be copied into other repos as-is. Only the "build"
5+## block below is repo-specific; everything else is driven by `env` at the top.
6+##
7+## Two things to know before copying it:
8+##
9+## 1. `issue_comment` workflows only ever run the copy of this file on the
10+## DEFAULT branch, so it has to be merged to main before `/preview` works
11+## on any PR — including the PR that adds it.
12+## 2. It checks out and builds the PR's head commit while the R2 credentials
13+## are in scope. The `author_association` guard means only the repo owner,
14+## an org member, or a collaborator can trigger it, but one of those people
15+## asking for a preview of an untrusted fork PR would run that fork's code.
16+## Read the diff before commenting `/preview` on a fork.
17+
18+name: preview
19+
20+on:
21+ issue_comment:
22+ types: [created]
23+
24+env:
25+ R2_BUCKET: tempory
26+ R2_PREFIX: previews
27+ PREVIEW_BASE_URL: https://tempory.net/previews
28+ DIST_DIR: dist
29+ NODE_VERSION: "24"
30+
31+permissions:
32+ contents: read
33+ issues: write
34+ pull-requests: write
35+
36+concurrency:
37+ group: preview-${{ github.event.issue.number }}
38+ cancel-in-progress: true
39+
40+jobs:
41+ preview:
42+ if: >-
43+ github.event.issue.pull_request &&
44+ startsWith(github.event.comment.body, '/preview') &&
45+ contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.comment.author_association)
46+ runs-on: ubuntu-latest
47+ env:
48+ GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
49+ PR_NUMBER: ${{ github.event.issue.number }}
50+ APP_NAME: ${{ github.event.repository.name }}
51+ steps:
52+ - name: Acknowledge the command
53+ env:
54+ COMMENT_ID: ${{ github.event.comment.id }}
55+ run: |
56+ gh api --silent -X POST \
57+ "repos/$GITHUB_REPOSITORY/issues/comments/$COMMENT_ID/reactions" \
58+ -f content=eyes
59+
60+ # The comment event carries no commit info, so ask the API what the PR
61+ # currently points at. `head.repo` differs from this repo for fork PRs.
62+ - name: Resolve the PR head
63+ id: head
64+ run: |
65+ gh api "repos/$GITHUB_REPOSITORY/pulls/$PR_NUMBER" \
66+ --jq '"repo=\(.head.repo.full_name)", "sha=\(.head.sha)", "branch=\(.head.ref)"' \
67+ >> "$GITHUB_OUTPUT"
68+
69+ # Branch names may contain characters that are awkward in a URL path
70+ # (most commonly `/`, as in `feat/thing`), so flatten to one path segment.
71+ # BRANCH goes through the environment rather than `${{ }}` because on a
72+ # fork PR its value is chosen by someone outside the org.
73+ - name: Compute the deploy path
74+ id: path
75+ env:
76+ BRANCH: ${{ steps.head.outputs.branch }}
77+ run: |
78+ slug=$(printf '%s' "$BRANCH" | tr -c 'A-Za-z0-9._-' '-')
79+ echo "key=$R2_PREFIX/$APP_NAME/$slug" >> "$GITHUB_OUTPUT"
80+ echo "url=$PREVIEW_BASE_URL/$APP_NAME/$slug/index.html" >> "$GITHUB_OUTPUT"
81+
82+ - uses: actions/checkout@v4
83+ with:
84+ repository: ${{ steps.head.outputs.repo }}
85+ ref: ${{ steps.head.outputs.sha }}
86+
87+ - uses: actions/setup-node@v4
88+ with:
89+ node-version: ${{ env.NODE_VERSION }}
90+ cache: npm
91+
92+ ## ---- repo-specific build (replace this block when copying) ----------
93+ # numbl is a `file:../../numbl` dependency: we use its compiler internals
94+ # (parser, lowerer, IR, inline pass), which its published package
95+ # `exports` do not expose. Clone it where that relative path expects it,
96+ # pinned to the same ref as ci.yml and deploy.yml. One file in it is
97+ # generated rather than committed — the interpreter's stdlib bundle,
98+ # which numbl gitignores — so a bare checkout is missing it and
99+ # `executeCode.ts` fails to resolve it; its generator only reads .m files
100+ # off disk, so plain `node` runs it without installing anything.
101+ - name: Check out numbl (sibling dependency)
102+ env:
103+ NUMBL_REF: main
104+ run: |
105+ git clone --filter=blob:none --no-checkout \
106+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
107+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
108+ node "$GITHUB_WORKSPACE/../../numbl/scripts/bundle-stdlib.ts"
109+ # --ignore-scripts: npm runs a linked package's `prepare` script, and
110+ # numbl's is husky, which is not installed here.
111+ - run: npm ci --ignore-scripts
112+ - run: npm run build
113+ ## ---- end repo-specific build ----------------------------------------
114+
115+ - name: Publish to R2
116+ env:
117+ AWS_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
118+ AWS_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
119+ AWS_DEFAULT_REGION: auto
120+ R2_ENDPOINT: ${{ secrets.R2_ENDPOINT }}
121+ # R2 rejects the extra integrity checksums AWS CLI v2 adds by default.
122+ AWS_REQUEST_CHECKSUM_CALCULATION: when_required
123+ AWS_RESPONSE_CHECKSUM_VALIDATION: when_required
124+ KEY: ${{ steps.path.outputs.key }}
125+ run: |
126+ # Two passes so the entry HTML is always revalidated while the
127+ # content-hashed assets around it can be cached hard. The filters
128+ # apply to both sides of the sync, so `--delete` in each pass only
129+ # ever removes files of that same kind.
130+ aws s3 sync "$DIST_DIR" "s3://$R2_BUCKET/$KEY" \
131+ --endpoint-url "$R2_ENDPOINT" --delete --no-progress \
132+ --exclude '*.html' \
133+ --cache-control 'public, max-age=31536000, immutable'
134+ aws s3 sync "$DIST_DIR" "s3://$R2_BUCKET/$KEY" \
135+ --endpoint-url "$R2_ENDPOINT" --delete --no-progress \
136+ --exclude '*' --include '*.html' \
137+ --cache-control 'no-cache'
138+
139+ - name: Comment with the preview link
140+ env:
141+ URL: ${{ steps.path.outputs.url }}
142+ SHA: ${{ steps.head.outputs.sha }}
143+ run: |
144+ gh pr comment "$PR_NUMBER" --repo "$GITHUB_REPOSITORY" \
145+ --body "Preview of \`${SHA:0:7}\` is live: $URL"
146+
147+ - name: Report failure
148+ if: failure()
149+ env:
150+ RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
151+ run: |
152+ gh pr comment "$PR_NUMBER" --repo "$GITHUB_REPOSITORY" \
153+ --body "Preview build failed — [run log]($RUN_URL)."
.gitignoreadded+4−0View file
@@ -0,0 +1,4 @@
1+node_modules/
2+dist/
3+*.log
4+*.png
README.mdadded+146−0View file
@@ -0,0 +1,146 @@
1+# turing-surface-cache
2+
3+Reaction-diffusion systems (Turing patterns) on curved closed surfaces,
4+evaluated at a chosen end time, with the solutions shared between all visitors
5+through a cloud cache.
6+
7+This is a trimmed fork of
8+[turing-surface](https://github.com/concept-collection/turing-surface), which
9+solves the same systems live and freely tunable. What this app changes is the
10+contract: every setting is a choice from a short list, so each combination of
11+choices names exactly one solution. Solutions already in the shared cache load
12+by themselves as the choices are browsed; a combination nobody has computed
13+shows empty surfaces, and nothing runs until the user presses **Compute
14+solution**, which runs the solver locally (in the browser, on WebGPU),
15+watching the pattern form and stopping at exactly the requested time. Users
16+who hold an upload API key contribute their locally-computed solutions back,
17+so the next visitor who asks for the same combination gets it in a second
18+rather than a minute.
19+
20+## The discrete parameter space
21+
22+One model ships (Schnakenberg, in turing-surface's 6-transform flux form) on
23+three geometries (sphere, ellipsoid, peanut). The choices, defined in
24+[`src/cache/options.ts`](src/cache/options.ts):
25+
26+| setting | choices |
27+|---|---|
28+| a | 0.05, **0.1**, 0.15, 0.2 |
29+| b | 0.7, **0.9**, 1.1, 1.3 |
30+| D₁ | 1.6e-4, **4e-4**, 1e-3 |
31+| D₂ | 3.2e-3, **8e-3**, 2e-2 |
32+| dt | 0.02, **0.05**, 0.1 |
33+| geometry | sphere, **ellipsoid** (axes each 0.6/1/1.5), peanut (waist 0.4/0.6/0.8, stretch 0/0.6/1.2) |
34+| seed | **1**–5 |
35+| end time | **100**, 200, 400, 800, 1600 |
36+
37+(Defaults in bold.) The numerical-scheme settings are fixed — lmax 63, 8
38+solve iterations, seed wavelength λ = 0.5 — but are recorded in every cache
39+key, so offering them as choices later invalidates nothing.
40+
41+The whole selection is mirrored into the URL fragment, every value written
42+explicitly, so reloading returns to the same combination and a shared link
43+opens on the same spec (and, when cached, the same solution) for whoever
44+follows it. A "Reset to defaults" button puts every choice back.
45+
46+Every end time is an exact multiple of every dt, so a run to T = 800 passes
47+exactly through t = 100, 200 and 400. Those intermediate states are captured
48+as the run passes them and, for an uploading user, encoded and uploaded in
49+the background while the run continues: one long run populates four cache
50+entries, and the earlier ones are already shared before the run finishes. The same structure works in the other
51+direction: the state is Markovian in the spectral coefficients, so before
52+computing anything the app looks for the longest cached shorter run of the
53+same spec and continues from its final state, computing only the remainder.
54+Asking for T = 1600 when T = 800 is cached costs half the run, and the t = 0
55+initial state travels inside every file of the chain, so a continuation
56+writes files identical in kind to a from-scratch run.
57+
58+## How the cache works
59+
60+The page's whole state is one small spec object (model, parameters, geometry,
61+seed, end time, scheme settings, plus the app name and a format version). Its
62+canonical JSON — keys sorted at every level — is hashed with SHA-256, and the
63+hash is the object name:
64+
65+```
66+https://tempory.net/tmpbucket/turing-surface-cache/v1/schnakenberg/<sha256>.h5
67+```
68+
69+A lookup is therefore a single GET with no index or API, and a 404 means a
70+miss. The path carries the app name, format version and model in the clear so
71+that future cleanup (lifecycle rules, prefix deletes) never has to open a file
72+to know what it belongs to; the hash input includes the version string, so a
73+format change moves every object rather than silently colliding with the old
74+ones.
75+
76+Uploads go through the [tmpbucket](https://github.com/scratchrealm/tmpbucket)
77+Worker: the client presents the API key and a file name, receives a presigned
78+R2 PUT URL, and uploads directly. Only holders of the key can write; everyone
79+can read. The key is entered in the page and kept in localStorage.
80+
81+## The cache file
82+
83+Cache files are HDF5, written in the browser with
84+[h5wasm](https://github.com/usnistgov/h5wasm) and readable from Python with
85+h5py. The layout is turing-surface's reference-file layout (see
86+`docs/ellipsoid-reference-spec.md` there) extended with the cache's identity
87+at the root, so a cache file is *also* a valid reference file — it can be
88+loaded straight into turing-surface's "Compare against uploaded data" mode:
89+
90+```
91+/ attrs: app, format_version, spec_json, model, species,
92+ created_utc, adapter
93+├─ backend/ attrs: adapter, runtime, precision
94+├─ spec/ attrs: geometry, lmax, seed, steps, niter, lam3, t_end
95+│ ├─ params/ attrs: a, b, D1, D2, dt
96+│ └─ geometry_params/ attrs: the geometry's params
97+├─ grid/ attrs: lmax, mmax, nlat, nphi, nlm
98+├─ geometry/ Gx, Gy, Gz float32[2·nlm]
99+├─ initial/ U, V (spectral state at t = 0) float32[2·nlm]
100+└─ final/ U, V (at the end time) float32[2·nlm]
101+```
102+
103+`spec_json` is the exact string that was hashed into the object name, and the
104+reader verifies it matches what was asked for. The initial state is included
105+so a file fully defines its run; the coefficients are the spherical-harmonic
106+convention documented in turing-surface (orthonormal + Condon-Shortley,
107+m-major, [re, im] interleaved). At lmax 63 a file is about 90 KB.
108+
109+Note that the solver is deterministic given the spec only to fp32 round-off:
110+different GPUs round differently, so a cached solution and a local recompute
111+agree closely but not bit-for-bit. The cache stores whichever trusted user
112+computed a combination first, and the file records which adapter that was.
113+
114+## Development
115+
116+```
117+npm install
118+npm run dev # local dev server
119+npm run build # type-check + production build to dist/
120+```
121+
122+numbl is a local `file:../../numbl` dependency, exactly as in turing-surface —
123+a sibling checkout of [numbl](https://github.com/flatironinstitute/numbl) is
124+required, reached through the `numbl-src` alias in
125+[`vite.config.ts`](vite.config.ts). See turing-surface's README for the
126+details; nothing about the arrangement changed here.
127+
128+Checks:
129+
130+- `node scripts/check-app.mjs` — end-to-end in headless Chrome (SwiftShader
131+ WebGPU) with the cloud cache mocked: a miss computes locally and produces
132+ the .h5 (verified with h5py), a fresh page loads that .h5 as a hit, and a
133+ third page asking for a longer end time warm-starts from it. The pages use
134+ the `?tend=` query hook, which substitutes short test end times for the
135+ UI's list. This is what CI runs.
136+- `node scripts/check-live.mjs [url]` — smoke-check a deployed URL against
137+ the real cache.
138+- `node scripts/screenshot.mjs out.png [light|dark] [tEnd]` — screenshot
139+ after the boot-time solve.
140+
141+Deployed to GitHub Pages by `.github/workflows/deploy.yml` on push to `main`.
142+
143+## License
144+
145+CECILL-2.1 (inherited from SHTNS via shtns-webgpu, whose sources are
146+vendored under `src/sht/`).
geometries/ellipsoid.madded+8−0View file
@@ -0,0 +1,8 @@
1+% A triaxial ellipsoid: the sphere with each axis scaled independently.
2+
3+function [gx, gy, gz] = shape(theta, phi, ax, ay, az)
4+ st = sin(theta);
5+ gx = ax * (st .* cos(phi));
6+ gy = ay * (st .* sin(phi));
7+ gz = az * cos(theta);
8+end
geometries/peanut.madded+9−0View file
@@ -0,0 +1,9 @@
1+% A dumbbell: radius 1 - waist*sin(theta)^2, stretched along z.
2+
3+function [gx, gy, gz] = shape(theta, phi, waist, stretch)
4+ st = sin(theta);
5+ r = 1 - waist * (st .^ 2);
6+ gx = r .* (st .* cos(phi));
7+ gy = r .* (st .* sin(phi));
8+ gz = (1 + stretch) * (r .* cos(theta));
9+end
geometries/sphere.madded+15−0View file
@@ -0,0 +1,15 @@
1+% The unit sphere — the reference case.
2+%
3+% A geometry file defines shape(theta, phi, ...) -> gx, gy, gz: the surface
4+% over the solver's grid (all npts x 1). Unlike the models it runs once, on
5+% the CPU through numbl's interpreter, so the full MATLAB subset is
6+% available — loops, arrays, min/max, legendre, seeded randomness via
7+% rng/randn. The host analyses the result into spherical-harmonic
8+% coefficients, band-limited at lmax.
9+
10+function [gx, gy, gz] = shape(theta, phi)
11+ st = sin(theta);
12+ gx = st .* cos(phi);
13+ gy = st .* sin(phi);
14+ gz = cos(theta);
15+end
index.htmladded+146−0View file
@@ -0,0 +1,146 @@
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%232a5f7f%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-cache — reaction-diffusion solutions from a shared cloud cache</title>
8+ <style>
9+ :root {
10+ --bg: #ffffff;
11+ --ink: #1f2328;
12+ --ink-2: #57606a;
13+ --line: #d0d7de;
14+ --accent: #0969da;
15+ --ok: #1a7f37;
16+ --sphere-bg: #f4f6f8;
17+ color-scheme: light dark;
18+ }
19+ @media (prefers-color-scheme: dark) {
20+ :root {
21+ --bg: #14171a;
22+ --ink: #e6e9ec;
23+ --ink-2: #9aa4af;
24+ --line: #333b44;
25+ --accent: #58a6ff;
26+ --ok: #3fb950;
27+ --sphere-bg: #14161c;
28+ }
29+ }
30+ body {
31+ margin: 0;
32+ background: var(--bg);
33+ color: var(--ink);
34+ font: 15px/1.5 system-ui, -apple-system, sans-serif;
35+ }
36+ main { max-width: 1100px; margin: 0 auto; padding: 20px 16px 48px; }
37+ h1 { font-size: 20px; margin: 0 0 2px; }
38+ .sub { color: var(--ink-2); margin: 0 0 12px; font-size: 13px; }
39+ .sub a { color: var(--accent); }
40+ .controls {
41+ display: flex; flex-wrap: wrap; gap: 8px 14px; align-items: center;
42+ padding: 6px 0;
43+ }
44+ .controls[hidden] { display: none; }
45+ .controls label { color: var(--ink-2); font-size: 13px; white-space: nowrap; }
46+ select, input[type="password"], button {
47+ font: inherit; font-size: 13px;
48+ color: var(--ink); background: var(--bg);
49+ border: 1px solid var(--line); border-radius: 6px;
50+ padding: 4px 8px;
51+ }
52+ button { cursor: pointer; }
53+ button:hover { border-color: var(--accent); }
54+ button.primary { border-color: var(--accent); color: var(--accent); font-weight: 600; min-width: 8em; }
55+ button:disabled { opacity: 0.5; cursor: default; }
56+ #cachenote { font-size: 13px; color: var(--ink-2); }
57+ #cachenote b { color: var(--ok); font-weight: 600; }
58+ #status { margin-top: 10px; font-size: 13.5px; }
59+ #status b { font-weight: 600; }
60+ #panels { display: flex; flex-wrap: wrap; gap: 14px; margin-top: 12px; }
61+ .panel {
62+ flex: 1 1 320px; min-width: 280px;
63+ border: 1px solid var(--line); border-radius: 8px; overflow: hidden;
64+ display: flex;
65+ }
66+ .sphere-box { flex: 1; aspect-ratio: 1 / 1; max-height: 70vh; position: relative; }
67+ .species-tag {
68+ position: absolute; top: 8px; left: 10px; z-index: 2;
69+ font-size: 15px; font-weight: 600; color: #fff;
70+ background: rgba(0, 0, 0, 0.45);
71+ padding: 1px 10px; border-radius: 12px;
72+ pointer-events: none;
73+ }
74+ .colorbar {
75+ display: flex; flex-direction: column; align-items: center; justify-content: center;
76+ gap: 4px; padding: 8px 4px; background: var(--sphere-bg);
77+ width: 52px; flex: none; box-sizing: border-box;
78+ }
79+ .colorbar canvas { border: 1px solid var(--line); border-radius: 2px; }
80+ .colorbar-label { font-size: 11px; color: var(--ink-2); font-variant-numeric: tabular-nums; }
81+ .stats { margin-top: 10px; font-size: 13px; color: var(--ink-2); font-variant-numeric: tabular-nums; }
82+ .stats b { color: var(--ink); font-weight: 600; }
83+ .cloud {
84+ margin-top: 14px; border: 1px solid var(--line); border-radius: 8px;
85+ padding: 8px 12px; font-size: 13px; color: var(--ink-2);
86+ }
87+ .cloud .controls { padding: 2px 0 0; }
88+ #err { color: #b35900; white-space: pre-wrap; font-size: 13px; }
89+ </style>
90+ </head>
91+ <body>
92+ <main>
93+ <h1>turing-surface-cache</h1>
94+ <p class="sub">
95+ Reaction-diffusion (Turing patterns) on curved surfaces, evaluated at a
96+ chosen end time. Every setting is a choice from a short list, so each
97+ combination names one exact solution. Solutions already in the shared
98+ cloud cache load by themselves as you browse the choices; a
99+ combination no one has computed shows empty surfaces until you press
100+ <b>Compute solution</b>, which runs it here (with WebGPU, via
101+ <a href="https://github.com/concept-collection/turing-surface">turing-surface</a>'s
102+ spectral solver). Drag to rotate.
103+ </p>
104+ <div class="controls" id="params"></div>
105+ <div class="controls">
106+ <label title="The surface the pattern is solved on">geometry
107+ <select id="geometry"></select>
108+ </label>
109+ <span id="geomparams" class="controls" style="padding: 0"></span>
110+ </div>
111+ <div class="controls">
112+ <label title="Which random initial perturbation to start from">seed
113+ <select id="seed"></select>
114+ </label>
115+ <label title="The solution is reported at this simulation time">end time
116+ <select id="tend"></select>
117+ </label>
118+ <button id="solve" class="primary"
119+ title="Compute the selected solution in your browser. Cached solutions load by themselves as you change the selection.">Compute solution</button>
120+ <button id="stop" hidden>Stop</button>
121+ <button id="reset" title="Set every selection back to its default">Reset to defaults</button>
122+ <span id="cachenote"></span>
123+ </div>
124+ <p id="status"></p>
125+ <div id="panels"></div>
126+ <div class="controls">
127+ <button id="resetview">Reset view</button>
128+ <a id="download" hidden download>Download .h5</a>
129+ </div>
130+ <p class="stats" id="stats"></p>
131+ <div class="cloud">
132+ <span>Solutions computed here can be contributed back to the shared
133+ cache, so the next visitor gets them instantly. Contributing
134+ requires an upload API key.</span>
135+ <div class="controls">
136+ <label>upload API key
137+ <input id="apikey" type="password" autocomplete="off" placeholder="(optional)" />
138+ </label>
139+ <span id="uploadnote"></span>
140+ </div>
141+ </div>
142+ <p id="err"></p>
143+ </main>
144+ <script type="module" src="/src/main.ts"></script>
145+ </body>
146+</html>
models/schnakenberg.madded+135−0View file
@@ -0,0 +1,135 @@
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+% Explicit reaction, implicit diffusion (IMEX Euler). The implicit solve
7+% splits lap_g = lap_s + dlap: the round-sphere part lap_s is diagonal in
8+% spherical-harmonic space (eigenvalues -lam), and the loop iterates the
9+% geometric correction dlap from that exact solve. Grid fields are npts x 1;
10+% spectral fields are real 2 x nlm. See docs/richardson-iteration.md.
11+%
12+% The correction evaluates lap_g in flux form -- 7 transforms per species
13+% per iteration where the Cartesian-gradient form (Algorithm 4 of
14+% evolving_surface/notes/algos.tex) needs 12. See
15+% docs/reduced-transforms.md, and models/schnakenberg_alg4.m
16+% for the original form kept as a live reference.
17+%
18+% The flux divergence is split against the round sphere: the sphere's share
19+% of it is -jinv*lap_s(u), exact in spectral space, and only the geometry
20+% *deviation* meets r ~ 1/sin^2(theta). Without that split the concentrated
21+% division amplifies the polar roundoff of the whole flux, and since a
22+% Turing pattern is seeded by whatever is largest in its unstable band, the
23+% amplified polar noise -- static, and re-injected every step -- picks the
24+% nucleation site and grows a spot at the pole. See docs/reduced-transforms.md
25+% Sec 5.
26+
27+% The uniform steady state, perturbed by a smooth random field: chebfun's
28+% randnfun3 on the surface's bounding box, restricted to the surface by
29+% evaluating it at the grid points -- the way surfacefun seeds a run. lam3
30+% is its wavelength; the draw is seeded on the host, the sum over its
31+% Fourier modes runs on the GPU (src/mgpu/randnfun3.ts).
32+function [U, V, u, v] = init(lam3, gx, gy, gz, a, b)
33+ f = randnfun3(lam3, gx, gy, gz);
34+ us = a + b;
35+ vs = b / (us * us);
36+ [U, V] = analys(us + 0.01*f, vs * ones(numel(f), 1));
37+ [u, v] = synth(U, V);
38+end
39+
40+function [Un, Vn, u, v] = step(U, V, lam, filt, gx, gy, gz, p2, r, dp1, dq2, jinv, jhat, a, b, D1, D2, dt, niter)
41+ % Grouped transforms -- [a, b] = synth(x, y) -- are explicit batching:
42+ % output k is the transform of input k, and the whole group runs as one
43+ % batched Legendre dispatch, or as many as the device's lane width allows
44+ % (src/mgpu/plan.ts, materializeTransforms). The grouping is a promise of
45+ % independence, never of a lane width, so the same source runs anywhere.
46+ [u, v] = synth(U, V);
47+ uuv = u .* u .* v;
48+
49+ % Right-hand side of the implicit solve (I - dt*D*lap_g) Unew = B.
50+ ru = a - u + uuv;
51+ rv = b - uuv;
52+ [Ru, Rv] = analys(ru, rv);
53+ Bu = U + dt * Ru;
54+ Bv = V + dt * Rv;
55+
56+ % Preconditioned solve, then iterate the geometric correction. jhat is
57+ % the host's minimax scale over the operator's symbol eigenvalues mu(x)
58+ % -- the inverse squared principal stretches of the embedding, direction
59+ % included (src/geom/geometry.ts, Jhat): preconditioning with lam/jhat
60+ % contracts every mode and direction at rate
61+ % (muMax - muMin)/(muMax + muMin) < 1 on any surface, where the plain
62+ % lam (jhat = 1) diverges wherever mu > 2 -- docs/reduced-transforms.md
63+ % Sec 10. The answer never depends on jhat (the lamJ term added inside
64+ % dLu is the term divided back out); only the convergence rate does. On
65+ % the sphere mu = 1 and lamJ = lam.
66+ lamJ = lam ./ jhat;
67+ Un = Bu ./ (1 + (dt * D1) * lamJ);
68+ Vn = Bv ./ (1 + (dt * D2) * lamJ);
69+
70+ for k = 1:niter
71+ % dlap = lap_g - lap_s at the current iterate, in flux form
72+ % (docs/reduced-transforms.md Sec 4). The sin-weighted derivatives
73+ % A = sin(theta)*dtheta(u) and B = dphi(u) -- both smooth on the sphere,
74+ % synthesized straight from the dthetac/dphic coefficient shuffles --
75+ % are combined pointwise through the precomputed weights into two
76+ % fluxes P,Q, also smooth. The theta flux P goes back to
77+ % coefficients, through the same shuffle again, and is synthesized as
78+ % sin(theta)*dtheta(P); the phi flux Q never leaves the grid -- d/dphi
79+ % is diagonal in the Fourier index, so dphig differentiates it with two
80+ % FFT stages and no Legendre work (masking m past filt's reach). Their
81+ % sum, scaled by r, is lap_g(u). The only division by sin(theta)
82+ % anywhere is folded into the weights at precompute time.
83+ %
84+ % The weights here are the *sphere-subtracted* ones: p1 = 1 + dp1 and
85+ % q2 = 1 + dq2 (p2 is zero on the sphere already), so P,Q below are the
86+ % deviation fluxes P' = P - A, Q' = Q - B. What that leaves out is the
87+ % round sphere's own divergence, sin(theta)*dtheta(A) + dphi(B) =
88+ % -sin^2(theta)*lap_s(u), which needs no flux machinery at all: lap_s is
89+ % diagonal, so it is -lam.*Fu synthesized once (S below, riding along in
90+ % the gradient's batched synthesis) and scaled by the bounded
91+ % jinv = 1/J = r*sin^2(theta). r therefore multiplies only the deviation
92+ % -- the difference between this and multiplying the whole flux is two
93+ % orders of magnitude of polar roundoff, and it is what keeps a pattern
94+ % from nucleating at the pole (src/geom/geometry.ts, dp1/dq2/jinv).
95+ % lamJ.*Un adds back the preconditioner's -lap_s(Un)/jhat, since lam
96+ % holds +l(l+1). filt zeroes the top two degrees, where the derivative
97+ % recurrences cannot exactly represent a derivative -- and the correction
98+ % itself is projected onto the same band (algos.tex Algorithm 5 zeroes
99+ % the same coefficients): without that, each iteration replaces a bit
100+ % more of the top degrees' implicit diffusion with nothing (their fixed
101+ % point is the undiffused Bu), and the two species un-diffuse at
102+ % different rates -- a spurious Turing band at the band edge.
103+ %
104+ % The two species share each grouped call: the six gradient-and-sphere
105+ % syntheses, the two theta-flux analyses, the two divergence syntheses
106+ % and the two final analyses each run as one batched dispatch.
107+ Fu = Un .* filt;
108+ Fv = Vn .* filt;
109+ vtu = dthetac(Fu);
110+ vpu = dphic(Fu);
111+ vtv = dthetac(Fv);
112+ vpv = dphic(Fv);
113+ [Ftu, Fpu, Ftv, Fpv, Su, Sv] = synth(vtu, vpu, vtv, vpv, lam .* Fu, lam .* Fv);
114+ Pu = dp1 .* Ftu + p2 .* Fpu;
115+ Qu = p2 .* Ftu + dq2 .* Fpu;
116+ Pv = dp1 .* Ftv + p2 .* Fpv;
117+ Qv = p2 .* Ftv + dq2 .* Fpv;
118+ [PAu, PAv] = analys(Pu, Pv);
119+ Pcu = PAu .* filt;
120+ Pcv = PAv .* filt;
121+ scu = dthetac(Pcu);
122+ scv = dthetac(Pcv);
123+ [Lu, Lv] = synth(scu, scv);
124+ dQu = dphig(Qu);
125+ dQv = dphig(Qv);
126+ lapu = r .* (Lu + dQu) - jinv .* Su;
127+ lapv = r .* (Lv + dQv) - jinv .* Sv;
128+ [LAu, LAv] = analys(lapu, lapv);
129+ dLu = (LAu + lamJ .* Un) .* filt;
130+ dLv = (LAv + lamJ .* Vn) .* filt;
131+
132+ Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lamJ);
133+ Vn = (Bv + (dt * D2) * dLv) ./ (1 + (dt * D2) * lamJ);
134+ end
135+end
package-lock.jsonadded+2237−0View file
This diff is 2,242 lines long and is not shown.
package.jsonadded+27−0View file
@@ -0,0 +1,27 @@
1+{
2+ "name": "turing-surface-cache",
3+ "version": "0.1.0",
4+ "description": "Reaction-diffusion (Turing patterns) on curved surfaces at a chosen end time, with a shared cloud cache of solutions; 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+ },
14+ "dependencies": {
15+ "h5wasm": "^0.10.3",
16+ "numbl": "file:../../numbl",
17+ "three": "^0.183.0"
18+ },
19+ "devDependencies": {
20+ "@types/node": "^26.1.1",
21+ "@types/three": "^0.185.1",
22+ "@webgpu/types": "^0.1.44",
23+ "puppeteer-core": "^23.0.0",
24+ "typescript": "^5.5.0",
25+ "vite": "^5.4.0"
26+ }
27+}
scripts/check-app.mjsadded+259−0View file
@@ -0,0 +1,259 @@
1+/**
2+ * End-to-end check of the cache flow in headless Chrome (SwiftShader WebGPU),
3+ * without touching the real cloud cache:
4+ *
5+ * 1. cache miss — the page is loaded with the real tempory.net requests
6+ * intercepted to 404, the end time is set to 5, and the run computes
7+ * locally; the produced .h5 is pulled out of the download link.
8+ * 2. the .h5 is checked with Python h5py (layout, shapes, spec_json).
9+ * 3. cache hit — a fresh page, same selection, with the interception now
10+ * answering that .h5; the page must show "from the cloud cache".
11+ * 4. warm start — a fresh page asking for a longer end time, with only the
12+ * shorter run's .h5 in the "cache"; the page must resume from it and
13+ * compute only the remainder.
14+ *
15+ * The pages are opened with the ?tend= test hook so the computed runs stay
16+ * short (5 and 10 time units instead of the UI's 100+).
17+ *
18+ * Usage: node scripts/check-app.mjs
19+ */
20+import { createServer } from 'node:http';
21+import { readFile, writeFile } from 'node:fs/promises';
22+import { execFileSync } from 'node:child_process';
23+import { extname, join } from 'node:path';
24+import puppeteer from 'puppeteer-core';
25+
26+const DIST = new URL('../dist/', import.meta.url).pathname;
27+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
28+const T_END = '5';
29+const T_END_LONG = '10';
30+
31+const server = createServer(async (req, res) => {
32+ try {
33+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
34+ const data = await readFile(join(DIST, path));
35+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
36+ res.end(data);
37+ } catch {
38+ res.writeHead(404);
39+ res.end();
40+ }
41+});
42+await new Promise((r) => server.listen(0, '127.0.0.1', r));
43+const port = server.address().port;
44+
45+const browser = await puppeteer.launch({
46+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
47+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
48+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
49+});
50+
51+const problems = [];
52+function watch(page, tag) {
53+ page.on('pageerror', (e) => problems.push(`${tag} pageerror: ${e.message}`));
54+ page.on('console', (m) => {
55+ // Cache-lookup 404s log as resource errors by design; ignore them.
56+ if (
57+ m.type() === 'error' &&
58+ !/GL Driver|favicon|tempory\.net|Failed to load resource/.test(m.text())
59+ ) {
60+ problems.push(`${tag} console error: ${m.text()}`);
61+ }
62+ });
63+}
64+
65+/** Intercept tempory.net cache reads; `bytes` null => everything 404s. The
66+ * mocked responses need the CORS header the real bucket sends, or the page's
67+ * cross-origin fetch is blocked before it sees the status. */
68+const CORS = { 'access-control-allow-origin': '*' };
69+async function interceptCache(page, bytes, name) {
70+ await page.setRequestInterception(true);
71+ page.on('request', (req) => {
72+ const url = req.url();
73+ if (!url.startsWith('https://tempory.net/')) return void req.continue();
74+ if (bytes && name && url.endsWith(`/${name}`)) {
75+ if (req.method() === 'HEAD') return void req.respond({ status: 200, headers: CORS });
76+ return void req.respond({
77+ status: 200,
78+ headers: CORS,
79+ contentType: 'application/x-hdf5',
80+ body: Buffer.from(bytes),
81+ });
82+ }
83+ req.respond({ status: 404, headers: CORS, body: 'not found' });
84+ });
85+}
86+
87+/** Open the app. `tend` selects an end time through the dropdown; a `hash`
88+ * instead carries the selection in the URL fragment, exercising the
89+ * reload/share restore path. `hookOrder` sets the ?tend test list — its
90+ * first entry is the default selection, so a restore test must pass an
91+ * order whose default differs from the hash value, or a restore that
92+ * silently does nothing would still land on the right selection. */
93+async function openAndSelect(page, tend, hash = '', hookOrder = `${T_END},${T_END_LONG}`) {
94+ await page.setViewport({ width: 1100, height: 900 });
95+ // The ?tend hook replaces the end-time list with short test values.
96+ await page.goto(
97+ `http://127.0.0.1:${port}/index.html?tend=${hookOrder}${hash}`,
98+ { waitUntil: 'load' },
99+ );
100+ // Selection changes land long before the WebGPU compile finishes, so the
101+ // boot-time auto-refresh picks them up.
102+ await page.waitForSelector('#tend');
103+ if (tend !== null) await page.select('#tend', tend);
104+}
105+
106+const statusOf = (page) => page.$eval('#status', (el) => el.textContent);
107+const errOf = (page) => page.$eval('#err', (el) => el.textContent);
108+
109+try {
110+ // ---- pass 1: miss, compute locally --------------------------------------
111+ const page1 = await browser.newPage();
112+ watch(page1, 'miss:');
113+ await interceptCache(page1, null, '');
114+ await openAndSelect(page1, T_END);
115+ // A miss never computes on its own: the page must settle on empty windows
116+ // asking for the button.
117+ await page1.waitForFunction(
118+ () => /press Compute solution|failed/.test(
119+ document.getElementById('status')?.textContent ?? '') ||
120+ (document.getElementById('err')?.textContent?.length ?? 0) > 4,
121+ { timeout: 600_000 },
122+ );
123+ console.log('pass 1 idle status:', await statusOf(page1));
124+ const solveDisabled = await page1.$eval('#solve', (b) => b.disabled);
125+ if (solveDisabled) problems.push('miss: Compute solution button disabled while idle');
126+ await page1.click('#solve');
127+ // "Not uploaded" is the terminal state of a keyless local run — it appears
128+ // only after the .h5 has been encoded and the download link filled in.
129+ await page1.waitForFunction(
130+ () => /Not uploaded|failed/.test(document.getElementById('status')?.textContent ?? '') ||
131+ (document.getElementById('err')?.textContent?.length ?? 0) > 4,
132+ { timeout: 600_000 },
133+ );
134+ const s1 = await statusOf(page1);
135+ console.log('pass 1 status:', s1);
136+ const e1 = await errOf(page1);
137+ if (e1) problems.push(`miss: err: ${e1}`);
138+ if (!/computed locally/.test(s1)) problems.push(`miss: unexpected status: ${s1}`);
139+ if (!/t = 5\b/.test(s1)) problems.push(`miss: did not stop at t = 5: ${s1}`);
140+
141+ const fileName = await page1.$eval('#download', (a) => a.download);
142+ const b64 = await page1.$eval('#download', async (a) => {
143+ const buf = await (await fetch(a.href)).arrayBuffer();
144+ let out = '';
145+ const v = new Uint8Array(buf);
146+ for (let i = 0; i < v.length; i += 0x8000) {
147+ out += String.fromCharCode(...v.subarray(i, i + 0x8000));
148+ }
149+ return btoa(out);
150+ });
151+ const bytes = Buffer.from(b64, 'base64');
152+ console.log(`downloaded ${fileName}: ${bytes.length} bytes`);
153+ if (!/^[0-9a-f]{64}\.h5$/.test(fileName)) problems.push(`odd file name: ${fileName}`);
154+ // The selection is mirrored into the URL fragment on every change.
155+ if (!new URL(page1.url()).hash.includes(`tend=${T_END}`)) {
156+ problems.push(`miss: selection not in URL: ${page1.url()}`);
157+ }
158+ const h5Path = `/tmp/turing-surface-cache-check.h5`;
159+ await writeFile(h5Path, bytes);
160+ await page1.close();
161+
162+ // ---- pass 2: the .h5 itself, via h5py -----------------------------------
163+ try {
164+ const out = execFileSync('python3', ['-c', `
165+import h5py, json, sys
166+f = h5py.File('${h5Path}', 'r')
167+spec = json.loads(f.attrs['spec_json'])
168+assert f.attrs['app'] == 'turing-surface-cache', f.attrs['app']
169+assert int(f.attrs['format_version']) == 1
170+assert spec['tEnd'] == 5 and spec['model'] == 'schnakenberg', spec
171+assert int(f['spec'].attrs['steps']) == round(5 / spec['params']['dt'])
172+nlm = (spec['lmax'] + 1) * (spec['lmax'] + 2) // 2
173+for g in ('geometry/Gx', 'geometry/Gy', 'geometry/Gz', 'initial/U', 'initial/V', 'final/U', 'final/V'):
174+ d = f[g]
175+ assert d.shape == (2 * nlm,) and d.dtype.kind == 'f', (g, d.shape, d.dtype)
176+import numpy as np
177+assert np.isfinite(f['final/U'][:]).all() and np.abs(f['final/U'][:]).max() > 0
178+print('h5py check ok; species', list(f.attrs['species']), '; adapter:', f.attrs.get('adapter', '?'))
179+`], { encoding: 'utf8' });
180+ console.log('pass 2:', out.trim());
181+ } catch (e) {
182+ problems.push(`h5py check failed: ${e.stdout ?? ''}${e.stderr ?? e.message}`);
183+ }
184+
185+ // ---- pass 3: hit, load from "cache", selection restored from the URL ----
186+ const page2 = await browser.newPage();
187+ watch(page2, 'hit:');
188+ await interceptCache(page2, bytes, fileName);
189+ // No dropdown interaction: the end time arrives in the fragment, as it
190+ // would from a shared or reloaded link. The hook default is deliberately
191+ // the OTHER value, so only a working restore reaches the cached spec.
192+ await openAndSelect(page2, null, `#tend=${T_END}`, `${T_END_LONG},${T_END}`);
193+ const restored = await page2.$eval('#tend', (el) => el.value);
194+ if (restored !== T_END) problems.push(`hit: URL restore failed, tend = ${restored}`);
195+ // "from the cloud cache" is the hit's terminal status; "checking the cloud
196+ // cache…" is transient and must not satisfy the wait.
197+ await page2.waitForFunction(
198+ () => /from the cloud cache|Not uploaded|failed/.test(
199+ document.getElementById('status')?.textContent ?? '') ||
200+ (document.getElementById('err')?.textContent?.length ?? 0) > 4,
201+ { timeout: 600_000 },
202+ );
203+ const s2 = await statusOf(page2);
204+ console.log('pass 3 status:', s2);
205+ const e2 = await errOf(page2);
206+ if (e2) problems.push(`hit: err: ${e2}`);
207+ if (!/from the.*cloud cache/.test(s2)) problems.push(`hit: expected a cache hit: ${s2}`);
208+ const note = await page2.$eval('#cachenote', (el) => el.textContent);
209+ if (!/in the cloud cache/.test(note)) problems.push(`hit: cache note wrong: '${note}'`);
210+ const panels = await page2.$$eval('.sphere-box canvas', (els) => els.length);
211+ if (panels !== 2) problems.push(`hit: expected 2 sphere canvases, got ${panels}`);
212+ await page2.close();
213+
214+ // ---- pass 4: warm start from the shorter cached run ----------------------
215+ // Only the t = 5 file is in the "cache"; asking for t = 10 must resume from
216+ // it and compute just the remainder.
217+ const page3 = await browser.newPage();
218+ watch(page3, 'warm:');
219+ await interceptCache(page3, bytes, fileName);
220+ await openAndSelect(page3, T_END_LONG);
221+ await page3.waitForFunction(
222+ () => /press Compute solution|failed/.test(
223+ document.getElementById('status')?.textContent ?? '') ||
224+ (document.getElementById('err')?.textContent?.length ?? 0) > 4,
225+ { timeout: 600_000 },
226+ );
227+ await page3.click('#solve');
228+ await page3.waitForFunction(
229+ () => /Not uploaded|failed/.test(document.getElementById('status')?.textContent ?? '') ||
230+ (document.getElementById('err')?.textContent?.length ?? 0) > 4,
231+ { timeout: 600_000 },
232+ );
233+ const s3 = await statusOf(page3);
234+ console.log('pass 4 status:', s3);
235+ const e3 = await errOf(page3);
236+ if (e3) problems.push(`warm: err: ${e3}`);
237+ if (!/t = 10\b/.test(s3)) problems.push(`warm: did not stop at t = 10: ${s3}`);
238+ if (!/resumed from cached t = 5\b/.test(s3)) {
239+ problems.push(`warm: expected a resume from t = 5: ${s3}`);
240+ }
241+ const warmFile = await page3.$eval('#download', (a) => a.download);
242+ if (warmFile === fileName || !/^[0-9a-f]{64}\.h5$/.test(warmFile)) {
243+ problems.push(`warm: odd file name: ${warmFile}`);
244+ }
245+ await page3.close();
246+} catch (e) {
247+ problems.push(`fatal: ${e.message}`);
248+} finally {
249+ await browser.close();
250+ server.close();
251+}
252+
253+if (problems.length) {
254+ console.log('PROBLEMS:');
255+ for (const p of new Set(problems)) console.log(' ' + p);
256+ process.exitCode = 1;
257+} else {
258+ console.log('CHECK-APP: PASS');
259+}
scripts/check-live.mjsadded+62−0View file
@@ -0,0 +1,62 @@
1+/**
2+ * Smoke-check a deployed URL in headless Chrome: load it and wait for the
3+ * boot-time solve to finish — either from the cloud cache or computed locally.
4+ * Talks to the real cache. Usage: node scripts/check-live.mjs [url]
5+ */
6+import puppeteer from 'puppeteer-core';
7+
8+const url = process.argv[2] ?? 'https://concept-collection.github.io/turing-surface-cache/';
9+const browser = await puppeteer.launch({
10+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
11+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
12+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
13+});
14+const page = await browser.newPage();
15+await page.setViewport({ width: 1100, height: 900 });
16+const problems = [];
17+page.on('pageerror', (e) => problems.push(`pageerror: ${e.message}`));
18+page.on('requestfailed', (r) => {
19+ // The cache lookup 404s by design when the selection is not cached.
20+ if (!r.url().startsWith('https://tempory.net/')) {
21+ problems.push(`request failed: ${r.url()}`);
22+ }
23+});
24+page.on('console', (m) => {
25+ if (m.type() === 'error' && !/GL Driver|favicon|tempory\.net/.test(m.text())) {
26+ problems.push(`console error: ${m.text()}`);
27+ }
28+});
29+
30+try {
31+ await page.goto(url, { waitUntil: 'load', timeout: 60_000 });
32+ // Nothing computes without the button, so any selection is safe to idle on.
33+ await page.waitForSelector('#tend');
34+ // Terminal statuses only — "checking the cloud cache…" is transient. A
35+ // miss settles on empty windows asking for the button; nothing computes
36+ // during this check.
37+ await page.waitForFunction(
38+ () => /from the cloud cache|press Compute solution|failed/.test(
39+ document.getElementById('status')?.textContent ?? '') ||
40+ (document.getElementById('err')?.textContent?.length ?? 0) > 4,
41+ { timeout: 600_000 },
42+ );
43+ console.log('status:', await page.$eval('#status', (el) => el.textContent));
44+ const err = await page.$eval('#err', (el) => el.textContent);
45+ if (err) problems.push(`err: ${err}`);
46+ const panels = await page.$$eval('.sphere-box canvas', (els) => els.length);
47+ console.log('sphere canvases:', panels);
48+ if (panels !== 2) problems.push(`expected 2 sphere canvases, got ${panels}`);
49+ if (problems.length) {
50+ console.log('PROBLEMS:');
51+ for (const p of new Set(problems)) console.log(' ' + p);
52+ process.exitCode = 1;
53+ } else {
54+ console.log('LIVE CHECK: PASS');
55+ }
56+} catch (e) {
57+ console.error(`LIVE CHECK FAIL: ${e.message}`);
58+ for (const p of new Set(problems)) console.error(' ' + p);
59+ process.exitCode = 1;
60+} finally {
61+ await browser.close();
62+}
scripts/screenshot.mjsadded+68−0View file
@@ -0,0 +1,68 @@
1+/** Screenshot the app (dist/) in headless Chrome after the boot-time solve
2+ * finishes. Talks to the real cache.
3+ * Usage: node scripts/screenshot.mjs out.png [light|dark] [tEnd] */
4+import { createServer } from 'node:http';
5+import { readFile } from 'node:fs/promises';
6+import { extname, join } from 'node:path';
7+import puppeteer from 'puppeteer-core';
8+
9+const out = process.argv[2] ?? 'demo.png';
10+const scheme = process.argv[3] ?? 'light';
11+const tEnd = process.argv[4] ?? '100';
12+const DIST = new URL('../dist/', import.meta.url).pathname;
13+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
14+
15+const server = createServer(async (req, res) => {
16+ try {
17+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
18+ const data = await readFile(join(DIST, path));
19+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
20+ res.end(data);
21+ } catch {
22+ res.writeHead(404);
23+ res.end();
24+ }
25+});
26+await new Promise((r) => server.listen(0, '127.0.0.1', r));
27+const port = server.address().port;
28+
29+const browser = await puppeteer.launch({
30+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
31+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
32+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
33+});
34+const page = await browser.newPage();
35+await page.setViewport({ width: 1100, height: 900 });
36+await page.emulateMediaFeatures([{ name: 'prefers-color-scheme', value: scheme }]);
37+page.on('console', (m) => console.log(' [page]', m.text()));
38+// The ?tend hook accepts any end time, listed or not, so a short test run
39+// can be screenshotted too.
40+await page.goto(`http://127.0.0.1:${port}/index.html?tend=${tEnd}`, { waitUntil: 'load' });
41+await page.waitForSelector('#tend');
42+// Terminal statuses only — "checking the cloud cache…" is transient. On a
43+// miss the page settles on empty windows; press the button so the screenshot
44+// shows a pattern either way.
45+const idle = /from the cloud cache|press Compute solution|failed/;
46+await page.waitForFunction(
47+ (re) => new RegExp(re).test(document.getElementById('status')?.textContent ?? '') ||
48+ (document.getElementById('err')?.textContent?.length ?? 0) > 4,
49+ { timeout: 600_000 },
50+ idle.source,
51+);
52+if (/press Compute solution/.test(await page.$eval('#status', (el) => el.textContent))) {
53+ await page.click('#solve');
54+ await page.waitForFunction(
55+ () => /Not uploaded|Uploaded \d|from the cloud cache|failed/.test(
56+ document.getElementById('status')?.textContent ?? '') ||
57+ (document.getElementById('err')?.textContent?.length ?? 0) > 4,
58+ { timeout: 600_000 },
59+ );
60+}
61+await new Promise((r) => setTimeout(r, 300));
62+await page.screenshot({ path: out });
63+console.log('screenshot:', out);
64+console.log('status:', await page.$eval('#status', (el) => el.textContent));
65+const err = await page.$eval('#err', (el) => el.textContent);
66+if (err) console.log('err:', err);
67+await browser.close();
68+server.close();
src/cache/client.tsadded+78−0View file
@@ -0,0 +1,78 @@
1+/**
2+ * The cloud side of the cache.
3+ *
4+ * Reads are anonymous GETs against the bucket's public base URL: the object
5+ * name is a deterministic function of the spec (src/cache/spec.ts), so a
6+ * lookup is one fetch and a 404 is a miss — no index, no API.
7+ *
8+ * Writes go through the tmpbucket Worker
9+ * (https://github.com/scratchrealm/tmpbucket): present an API key and a file name,
10+ * receive a presigned R2 PUT URL, and upload directly. Only holders of the
11+ * key can write; everyone can read.
12+ */
13+import { cacheFileName, canonicalJson, type CacheSpec } from './spec.ts';
14+
15+const PUBLIC_BASE = 'https://tempory.net/tmpbucket/';
16+const WORKER_BASE = 'https://tmpbucket.figurl.workers.dev';
17+const CONTENT_TYPE = 'application/x-hdf5';
18+
19+export interface CacheLookup {
20+ fileName: string;
21+ url: string;
22+ specJson: string;
23+}
24+
25+export async function lookupFor(spec: CacheSpec): Promise<CacheLookup> {
26+ const fileName = await cacheFileName(spec);
27+ return { fileName, url: PUBLIC_BASE + fileName, specJson: canonicalJson(spec) };
28+}
29+
30+/** Fetch a cached solution; null on a miss. Throws on network failure or an
31+ * unexpected status, which are reported rather than treated as misses. */
32+export async function fetchCached(lookup: CacheLookup): Promise<Uint8Array | null> {
33+ const res = await fetch(lookup.url, { cache: 'no-store' });
34+ if (res.status === 404) return null;
35+ if (!res.ok) throw new Error(`cache read: HTTP ${res.status} for ${lookup.url}`);
36+ return new Uint8Array(await res.arrayBuffer());
37+}
38+
39+/** Upload one cache file. Resolves to its public URL. */
40+export async function uploadCacheFile(
41+ apiKey: string,
42+ fileName: string,
43+ bytes: Uint8Array,
44+): Promise<string> {
45+ const res = await fetch(`${WORKER_BASE}/api/upload-url`, {
46+ method: 'POST',
47+ headers: {
48+ Authorization: `Bearer ${apiKey}`,
49+ 'Content-Type': 'application/json',
50+ },
51+ body: JSON.stringify({ fileName, contentType: CONTENT_TYPE }),
52+ });
53+ if (res.status === 401 || res.status === 403) {
54+ throw new Error('upload not authorized — check the API key');
55+ }
56+ if (!res.ok) {
57+ let message = `HTTP ${res.status}`;
58+ try {
59+ const err = (await res.json()) as { message?: string };
60+ if (err.message) message = err.message;
61+ } catch {
62+ // keep the status-only message
63+ }
64+ throw new Error(`upload-url request failed: ${message}`);
65+ }
66+ const grant = (await res.json()) as {
67+ uploadUrl: string;
68+ uploadHeaders?: Record<string, string>;
69+ downloadUrl: string;
70+ };
71+ const put = await fetch(grant.uploadUrl, {
72+ method: 'PUT',
73+ headers: grant.uploadHeaders ?? { 'Content-Type': CONTENT_TYPE },
74+ body: bytes as unknown as BodyInit,
75+ });
76+ if (!put.ok) throw new Error(`upload PUT failed: HTTP ${put.status}`);
77+ return grant.downloadUrl;
78+}
src/cache/h5file.tsadded+231−0View file
@@ -0,0 +1,231 @@
1+/**
2+ * Cache files are HDF5, in the layout of turing-surface's reference files
3+ * (docs/ellipsoid-reference-spec.md there) extended with the cache's own
4+ * identity at the root: the canonical spec JSON that was hashed into the
5+ * object name, the app name, and the format version. A cache file is thereby
6+ * also a valid reference file — turing-surface's "Compare to reference…" mode
7+ * opens one as-is — and readable from Python with h5py.
8+ *
9+ * / attrs: app, format_version, spec_json, model, species,
10+ * created_utc, adapter
11+ * /backend attrs: adapter, runtime, precision
12+ * /spec attrs: geometry, lmax, seed, steps, niter, lam3, t_end
13+ * /spec/params attrs: a, b, D1, D2, dt
14+ * /spec/geometry_params attrs: the geometry's params
15+ * /grid attrs: lmax, mmax, nlat, nphi, nlm
16+ * /geometry Gx, Gy, Gz float32[2*nlm]
17+ * /initial one dataset per species (U, V) float32[2*nlm]
18+ * /final one dataset per species (U, V) float32[2*nlm]
19+ *
20+ * h5wasm's browser build carries the whole HDF5 library as embedded wasm
21+ * (~4 MB), so it is imported dynamically and only here: the page pays for it
22+ * on the first cache hit or upload, never on startup.
23+ */
24+import type { ShtConfig } from '../sht/layout.ts';
25+import { nlmCalc } from '../sht/layout.ts';
26+import { APP_NAME, FORMAT_VERSION, canonicalJson, stepsFor, type CacheSpec } from './spec.ts';
27+
28+export interface CacheFileData {
29+ spec: CacheSpec;
30+ grid: ShtConfig;
31+ /** Spectral state names, in order — Schnakenberg's ['U', 'V']. */
32+ species: string[];
33+ /** The band-limited surface's own coefficients, [re, im] per (l, m). */
34+ geometry: { X: Float32Array; Y: Float32Array; Z: Float32Array };
35+ /** Spectral state at t = 0 (immediately after seeding). */
36+ initial: Record<string, Float32Array>;
37+ /** The same, at the spec's end time. */
38+ final: Record<string, Float32Array>;
39+ /** Provenance: which GPU computed it. */
40+ adapter: string;
41+}
42+
43+interface H5Module {
44+ ready: Promise<unknown>;
45+ File: new (path: string, mode: string) => H5WFile;
46+ FS?: unknown;
47+}
48+
49+interface H5Attr {
50+ value: unknown;
51+}
52+
53+interface H5Obj {
54+ attrs: Record<string, H5Attr>;
55+ get(name: string): unknown;
56+ create_group(name: string): unknown;
57+ create_attribute(name: string, data: unknown, shape?: unknown, dtype?: unknown): void;
58+ create_dataset(args: { name: string; data: unknown; shape?: number[]; dtype?: string }): void;
59+}
60+
61+interface H5WFile extends H5Obj {
62+ close(): void;
63+}
64+
65+interface EmFS {
66+ writeFile(path: string, data: Uint8Array): void;
67+ readFile(path: string): Uint8Array;
68+ unlink(path: string): void;
69+}
70+
71+let scratchCounter = 0;
72+
73+async function withH5<T>(fn: (h5: H5Module, fs: EmFS) => T | Promise<T>): Promise<T> {
74+ const h5 = (await import('h5wasm')) as unknown as H5Module;
75+ const { FS } = (await h5.ready) as { FS: EmFS };
76+ return fn(h5, FS);
77+}
78+
79+const groupOf = (node: H5Obj, name: string): H5Obj => {
80+ const g = node.get(name) as H5Obj | null;
81+ if (!g || typeof g.get !== 'function') throw new Error(`no '${name}/' group`);
82+ return g;
83+};
84+
85+const coeffsOf = (group: H5Obj, groupName: string, name: string, nlm: number): Float32Array => {
86+ const v = (group.get(name) as { value?: unknown } | null)?.value;
87+ if (!(v instanceof Float32Array)) {
88+ throw new Error(`'${groupName}/${name}' is not a float32 dataset`);
89+ }
90+ if (v.length !== 2 * nlm) {
91+ throw new Error(`'${groupName}/${name}' has ${v.length} values, expected 2*nlm = ${2 * nlm}`);
92+ }
93+ return v;
94+};
95+
96+/** Serialize one solution to HDF5 bytes. */
97+export function encodeCacheFile(data: CacheFileData): Promise<Uint8Array> {
98+ return withH5((h5, FS) => {
99+ const path = `/encode-${scratchCounter++}.h5`;
100+ const file = new h5.File(path, 'w');
101+ try {
102+ const { spec, grid } = data;
103+ file.create_attribute('app', APP_NAME);
104+ file.create_attribute('format_version', FORMAT_VERSION);
105+ file.create_attribute('spec_json', canonicalJson(spec));
106+ file.create_attribute('model', spec.model);
107+ file.create_attribute('species', data.species);
108+ file.create_attribute('created_utc', new Date().toISOString());
109+ file.create_attribute('adapter', data.adapter);
110+
111+ file.create_group('backend');
112+ const backend = groupOf(file, 'backend');
113+ backend.create_attribute('adapter', data.adapter);
114+ backend.create_attribute('runtime', 'browser-webgpu');
115+ backend.create_attribute('precision', 'fp32');
116+
117+ file.create_group('spec');
118+ const specGroup = groupOf(file, 'spec');
119+ specGroup.create_attribute('geometry', spec.geometry);
120+ specGroup.create_attribute('lmax', spec.lmax);
121+ specGroup.create_attribute('seed', spec.seed);
122+ specGroup.create_attribute('steps', stepsFor(spec));
123+ specGroup.create_attribute('niter', spec.niter);
124+ specGroup.create_attribute('lam3', spec.lam3);
125+ specGroup.create_attribute('t_end', spec.tEnd);
126+ specGroup.create_group('params');
127+ const params = groupOf(specGroup, 'params');
128+ for (const [k, v] of Object.entries(spec.params)) params.create_attribute(k, v);
129+ specGroup.create_group('geometry_params');
130+ const gparams = groupOf(specGroup, 'geometry_params');
131+ for (const [k, v] of Object.entries(spec.geometryParams)) gparams.create_attribute(k, v);
132+
133+ file.create_group('grid');
134+ const gridGroup = groupOf(file, 'grid');
135+ const nlm = nlmCalc(grid.lmax, grid.mmax);
136+ gridGroup.create_attribute('lmax', grid.lmax);
137+ gridGroup.create_attribute('mmax', grid.mmax);
138+ gridGroup.create_attribute('nlat', grid.nlat);
139+ gridGroup.create_attribute('nphi', grid.nphi);
140+ gridGroup.create_attribute('nlm', nlm);
141+
142+ file.create_group('geometry');
143+ const geom = groupOf(file, 'geometry');
144+ geom.create_dataset({ name: 'Gx', data: data.geometry.X });
145+ geom.create_dataset({ name: 'Gy', data: data.geometry.Y });
146+ geom.create_dataset({ name: 'Gz', data: data.geometry.Z });
147+
148+ for (const [groupName, states] of [
149+ ['initial', data.initial],
150+ ['final', data.final],
151+ ] as const) {
152+ file.create_group(groupName);
153+ const g = groupOf(file, groupName);
154+ for (const name of data.species) {
155+ const coeffs = states[name];
156+ if (!coeffs) throw new Error(`missing ${groupName} state '${name}'`);
157+ if (coeffs.length !== 2 * nlm) {
158+ throw new Error(`${groupName}/${name}: ${coeffs.length} values, expected ${2 * nlm}`);
159+ }
160+ g.create_dataset({ name, data: coeffs });
161+ }
162+ }
163+ } finally {
164+ file.close();
165+ }
166+ const bytes = FS.readFile(path);
167+ FS.unlink(path);
168+ return bytes;
169+ });
170+}
171+
172+export interface DecodedCacheFile {
173+ /** Parsed from the file's own spec_json — the identity it was stored under. */
174+ spec: CacheSpec;
175+ species: string[];
176+ initial: Record<string, Float32Array>;
177+ final: Record<string, Float32Array>;
178+ /** Provenance, when recorded. */
179+ adapter: string;
180+ created: string;
181+}
182+
183+/**
184+ * Read cache-file bytes back. `expectSpecJson` is the canonical JSON this
185+ * client asked the cache for; a mismatch with the file's own means the object
186+ * store handed back something other than what the key promised (corruption,
187+ * or a stale format), and is an error rather than a silent wrong answer.
188+ */
189+export function decodeCacheFile(
190+ bytes: Uint8Array,
191+ expectSpecJson: string,
192+ expectSpecies: string[],
193+): Promise<DecodedCacheFile> {
194+ return withH5((h5, FS) => {
195+ const path = `/decode-${scratchCounter++}.h5`;
196+ FS.writeFile(path, bytes);
197+ const file = new h5.File(path, 'r');
198+ try {
199+ const attr = (name: string): unknown => file.attrs[name]?.value;
200+ const app = String(attr('app') ?? '');
201+ if (app !== APP_NAME) throw new Error(`not a ${APP_NAME} file (app='${app}')`);
202+ const version = Number(attr('format_version'));
203+ if (version !== FORMAT_VERSION) throw new Error(`format version ${version}, expected ${FORMAT_VERSION}`);
204+ const specJson = String(attr('spec_json') ?? '');
205+ if (specJson !== expectSpecJson) {
206+ throw new Error('file spec does not match the requested spec');
207+ }
208+ const spec = JSON.parse(specJson) as CacheSpec;
209+ const nlm = nlmCalc(spec.lmax, spec.lmax);
210+ const initialGroup = groupOf(file, 'initial');
211+ const finalGroup = groupOf(file, 'final');
212+ const initial: Record<string, Float32Array> = {};
213+ const final: Record<string, Float32Array> = {};
214+ for (const name of expectSpecies) {
215+ initial[name] = coeffsOf(initialGroup, 'initial', name, nlm);
216+ final[name] = coeffsOf(finalGroup, 'final', name, nlm);
217+ }
218+ return {
219+ spec,
220+ species: expectSpecies,
221+ initial,
222+ final,
223+ adapter: String(attr('adapter') ?? ''),
224+ created: String(attr('created_utc') ?? ''),
225+ };
226+ } finally {
227+ file.close();
228+ FS.unlink(path);
229+ }
230+ });
231+}
src/cache/options.tsadded+80−0View file
@@ -0,0 +1,80 @@
1+/**
2+ * The discrete parameter space.
3+ *
4+ * Every knob the app offers is a choice from a short enumerated list, so a run
5+ * is fully identified by a small tuple of exact values — which is what makes
6+ * the cloud cache work: the same choices always produce the same cache key
7+ * (src/cache/spec.ts), with no floating-point formatting ambiguity, because
8+ * the values *are* these list entries, never something typed or interpolated.
9+ *
10+ * The numerical-scheme knobs (lmax, niter, seed wavelength) are fixed in this
11+ * version rather than offered, but they are recorded in the spec and the cache
12+ * key all the same, so making them choices later invalidates nothing.
13+ */
14+import type { Params } from '../mgpu/registry.ts';
15+
16+export interface DiscreteChoice {
17+ key: string;
18+ label: string;
19+ /** The allowed values, in display order. */
20+ values: number[];
21+ /** Default — must be one of `values`. */
22+ value: number;
23+}
24+
25+/** Model parameter choices (Schnakenberg). */
26+export const MODEL_CHOICES: DiscreteChoice[] = [
27+ { key: 'a', label: 'a', values: [0.05, 0.1, 0.15, 0.2], value: 0.1 },
28+ { key: 'b', label: 'b', values: [0.7, 0.9, 1.1, 1.3], value: 0.9 },
29+ { key: 'D1', label: 'D₁', values: [1.6e-4, 4e-4, 1e-3], value: 4e-4 },
30+ { key: 'D2', label: 'D₂', values: [3.2e-3, 8e-3, 2e-2], value: 8e-3 },
31+ { key: 'dt', label: 'dt', values: [0.02, 0.05, 0.1], value: 0.05 },
32+];
33+
34+/** Geometry parameter choices, by geometry key. The sphere has none. */
35+export const GEOMETRY_CHOICES: Record<string, DiscreteChoice[]> = {
36+ sphere: [],
37+ ellipsoid: [
38+ { key: 'ax', label: 'a', values: [0.6, 1, 1.5], value: 1.5 },
39+ { key: 'ay', label: 'b', values: [0.6, 1, 1.5], value: 1 },
40+ { key: 'az', label: 'c', values: [0.6, 1, 1.5], value: 0.6 },
41+ ],
42+ peanut: [
43+ { key: 'waist', label: 'waist', values: [0.4, 0.6, 0.8], value: 0.6 },
44+ { key: 'stretch', label: 'stretch', values: [0, 0.6, 1.2], value: 0.6 },
45+ ],
46+};
47+
48+export const SEED_CHOICE: DiscreteChoice = {
49+ key: 'seed',
50+ label: 'seed',
51+ values: [1, 2, 3, 4, 5],
52+ value: 1,
53+};
54+
55+/**
56+ * End times, in simulation-time units. Every entry is an exact multiple of
57+ * every dt choice, so a run to any of them lands on a whole number of steps —
58+ * and a run to a later one passes exactly through the earlier ones. That is
59+ * where the along-the-way cache snapshots come from, and, in the other
60+ * direction, why a computation can warm-start from the longest cached
61+ * shorter run of the same spec.
62+ */
63+export const T_END_CHOICE: DiscreteChoice = {
64+ key: 'tEnd',
65+ label: 'end time',
66+ values: [100, 200, 400, 800, 1600],
67+ value: 100,
68+};
69+
70+/** Fixed numerical-scheme settings (recorded in every spec and cache key). */
71+export const LMAX = 63;
72+export const NITER = 8;
73+export const LAM3 = 0.5;
74+
75+export const defaultChoiceParams = (choices: DiscreteChoice[]): Params =>
76+ Object.fromEntries(choices.map((c) => [c.key, c.value]));
77+
78+/** Display formatting: exact and compact ("4e-4", "0.05"). */
79+export const fmtChoice = (v: number): string =>
80+ v !== 0 && Math.abs(v) < 0.01 ? v.toExponential() : String(v);
src/cache/spec.tsadded+73−0View file
@@ -0,0 +1,73 @@
1+/**
2+ * The cache spec: the JSON object that fully identifies one solution, and the
3+ * mapping from it to an object name in the cloud cache.
4+ *
5+ * The name is the SHA-256 of the spec's canonical JSON serialization (keys
6+ * sorted recursively, numbers as ECMAScript shortest-round-trip strings, which
7+ * the language specifies exactly). The serialization includes the app name and
8+ * a format version, so a change to what a spec means changes every hash. The
9+ * object path also carries the app name, version, and model in the clear —
10+ * `turing-surface-cache/v1/schnakenberg/<sha256>.h5` — so that future cleanup
11+ * (lifecycle rules, prefix deletes, per-model sweeps) never has to open a file
12+ * to know what it belongs to.
13+ */
14+import type { Params } from '../mgpu/registry.ts';
15+
16+export const APP_NAME = 'turing-surface-cache';
17+/** Bump together with the `v1` path segment below. */
18+export const FORMAT_VERSION = 1;
19+
20+export interface CacheSpec {
21+ app: typeof APP_NAME;
22+ formatVersion: typeof FORMAT_VERSION;
23+ model: string;
24+ /** The model's own parameters, dt included. */
25+ params: Params;
26+ geometry: string;
27+ geometryParams: Params;
28+ lmax: number;
29+ niter: number;
30+ /** Wavelength of the seeded random field. */
31+ lam3: number;
32+ seed: number;
33+ /** Physical end time; steps = tEnd / dt, which must be an integer. */
34+ tEnd: number;
35+}
36+
37+/** Timesteps from t = 0 to the spec's end time. Throws if tEnd is not an
38+ * exact multiple of dt — the discrete lists are chosen so it always is. */
39+export function stepsFor(spec: CacheSpec): number {
40+ const dt = spec.params.dt;
41+ if (!(dt > 0)) throw new Error(`bad dt ${dt}`);
42+ const steps = Math.round(spec.tEnd / dt);
43+ if (Math.abs(steps * dt - spec.tEnd) > 1e-9 * spec.tEnd) {
44+ throw new Error(`tEnd ${spec.tEnd} is not a multiple of dt ${dt}`);
45+ }
46+ return steps;
47+}
48+
49+/** JSON with every object's keys sorted, at every level. */
50+export function canonicalJson(value: unknown): string {
51+ if (value === null || typeof value !== 'object') return JSON.stringify(value);
52+ if (Array.isArray(value)) return `[${value.map(canonicalJson).join(',')}]`;
53+ const keys = Object.keys(value as Record<string, unknown>).sort();
54+ const body = keys
55+ .map((k) => `${JSON.stringify(k)}:${canonicalJson((value as Record<string, unknown>)[k])}`)
56+ .join(',');
57+ return `{${body}}`;
58+}
59+
60+export async function sha256Hex(text: string): Promise<string> {
61+ const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(text));
62+ return [...new Uint8Array(digest)].map((b) => b.toString(16).padStart(2, '0')).join('');
63+}
64+
65+/**
66+ * Object name under the bucket's upload prefix (tmpbucket prepends
67+ * `tmpbucket/`, so the public URL is
68+ * https://tempory.net/tmpbucket/<this>).
69+ */
70+export async function cacheFileName(spec: CacheSpec): Promise<string> {
71+ const hash = await sha256Hex(canonicalJson(spec));
72+ return `${APP_NAME}/v${FORMAT_VERSION}/${spec.model}/${hash}.h5`;
73+}
src/geom/geometry.tsadded+434−0View file
@@ -0,0 +1,434 @@
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. Unlike the models it is *not* compiled
10+ * to WGSL: a model's step runs every frame and must lower to a fixed sequence
11+ * of GPU dispatches, but a shape is evaluated exactly once at build time and
12+ * survives only as coefficients. So it runs through numbl's CPU interpreter
13+ * instead, which buys the full MATLAB subset — loops, arrays, reductions,
14+ * `legendre`, seeded randomness via `rng`/`randn` — and f64 evaluation, where
15+ * the step dialect is element-wise f32. The result is then *analysed*: the
16+ * canonical geometry this project carries is the three sets of coefficients
17+ * `X`, `Y`, `Z`, one per Cartesian component of the embedding.
18+ *
19+ * Going through the coefficients rather than keeping the pointwise values is
20+ * what makes the geometry usable by a spectral method, for two reasons:
21+ *
22+ * - it is exactly band-limited at lmax afterwards, so the surface has as many
23+ * derivatives as the scheme needs and no aliased content the solver cannot
24+ * see. `x`, `y`, `z` below are the synthesis of the coefficients, not the
25+ * raw output of the .m — the shape actually being solved on, which for a
26+ * shape with sharp features is not quite the shape that was written down.
27+ * - it can be evaluated on any grid. The renderer draws the surface on the
28+ * (possibly finer) display grid by synthesizing the same coefficients
29+ * there, which is exact interpolation rather than subdivision — the same
30+ * argument that lets the species fields be oversampled.
31+ *
32+ * The unit sphere is the case where `x`, `y`, `z` are pure degree-1 harmonics
33+ * and everything downstream reduces to turing-sphere.
34+ */
35+import { parseMFile, type FunctionStmt } from 'numbl-src/numbl-core/parser/index.ts';
36+import { executeCode } from 'numbl-src/numbl-core/executeCode.ts';
37+import {
38+ RuntimeTensor,
39+ isRuntimeTensor,
40+ type RuntimeValue,
41+} from 'numbl-src/numbl-core/runtime/types.ts';
42+import { ShtPlan } from '../sht/sht.ts';
43+import type { ShtConfig } from '../sht/layout.ts';
44+import type { DerivPlan } from '../sht/deriv.ts';
45+import { computeMetric, computeFluxMetric } from './metric.ts';
46+import { toolFiles } from '../tools.ts';
47+import { inFunction, inModel, ModelCompileError } from '../mgpu/errors.ts';
48+import type { ModelParams } from '../mgpu/model.ts';
49+
50+/** The function a geometry file must define. */
51+export const SHAPE_FN = 'shape';
52+
53+export interface GeometryOptions {
54+ /** The solver's transform plan — the grid the shape is evaluated on. */
55+ sht: ShtPlan;
56+ cfg: ShtConfig;
57+ /** Geometry source (.m text). */
58+ source: string;
59+ /** Parameter names the .m may take beyond `theta` and `phi`. */
60+ paramNames: string[];
61+ params: ModelParams;
62+ /** Computes the theta/phi derivatives the inverse metric quantities need. */
63+ deriv: DerivPlan;
64+}
65+
66+export class Geometry {
67+ /** Coordinates on the solver grid, npts each — synthesis of the coefficients. */
68+ readonly x: Float32Array;
69+ readonly y: Float32Array;
70+ readonly z: Float32Array;
71+ /** Their spherical-harmonic coefficients, 2 x nlm each. */
72+ readonly X: Float32Array;
73+ readonly Y: Float32Array;
74+ readonly Z: Float32Array;
75+ /**
76+ * Inverse metric quantities (src/geom/metric.ts), grid space, npts each.
77+ * Depend only on the geometry, so — like x,y,z,X,Y,Z above — these are a
78+ * one-off computed here, not per-solve-step work. Used by the Algorithm-4
79+ * (12-transform) Laplace-Beltrami path.
80+ */
81+ readonly Vtx: Float32Array;
82+ readonly Vty: Float32Array;
83+ readonly Vtz: Float32Array;
84+ readonly Vpx: Float32Array;
85+ readonly Vpy: Float32Array;
86+ readonly Vpz: Float32Array;
87+ /**
88+ * Flux-form metric weights (src/geom/metric.ts computeFluxMetric), grid
89+ * space, npts each — the six-transform Laplace-Beltrami scheme's
90+ * replacement for the six V arrays (docs/reduced-transforms.md
91+ * Sec 3). Both sets are carried so either operator formulation can run.
92+ */
93+ readonly p1: Float32Array;
94+ readonly p2: Float32Array;
95+ readonly q2: Float32Array;
96+ readonly r: Float32Array;
97+ /**
98+ * The same flux weights with the round sphere subtracted off, plus the
99+ * bounded 1/J — what lets a model evaluate lap_g without ever multiplying
100+ * the *whole* flux divergence by r ~ 1/sin^2(theta). Writing p1 = 1 + dp1,
101+ * q2 = 1 + dq2 (p2 is already a pure deviation, zero on the sphere) splits
102+ * the divergence into a round-sphere part, whose cancelling bracket
103+ * sin(theta) dtheta(A) + dphi(B) = -sin^2(theta) lap_s u is known exactly in
104+ * spectral space, and a remainder:
105+ *
106+ * lap_g u = -jinv * lap_s u + r * (sin(theta) dtheta(P') + dphi(Q'))
107+ *
108+ * with P' = dp1*A + p2*B, Q' = p2*A + dq2*B. Only the remainder meets the
109+ * concentrated division, so the polar roundoff gain drops by |P'|/|P|
110+ * instead of applying to the full flux. Subtracting 1 in f64 here is the
111+ * point: on a near-sphere dp1 is the small quantity, and forming it as an
112+ * f32 difference in the .m would lose it. See docs/reduced-transforms.md
113+ * Sec 5 and models/schnakenberg.m.
114+ */
115+ readonly dp1: Float32Array;
116+ readonly dq2: Float32Array;
117+ readonly jinv: Float32Array;
118+ /**
119+ * Preconditioner scale for the implicit solve (docs/reduced-transforms.md
120+ * Sec 10). At high degree the Richardson iteration's per-mode factor is
121+ * governed by the operator's principal symbol: in the orthonormal frame
122+ * the surface symbol matrix is S = (1/J)[[p1, p2], [p2, q2]], whose
123+ * eigenvalues mu(x) are the inverse squared principal stretches of the
124+ * embedding — the round sphere has mu = 1. Preconditioning with lam/Jhat
125+ * contracts every mode and every direction iff Jhat*mu stays in (0, 2),
126+ * so the minimax constant is the harmonic mean of the symbol extremes,
127+ *
128+ * Jhat = 2/(muMin + muMax), rate = (muMax - muMin)/(muMax + muMin) < 1.
129+ *
130+ * The direction dependence is the point: a det-based mean of the area
131+ * factor J (mu's geometric mean, exact only for conformal surfaces)
132+ * under-corrects anisotropic stretching — on the shipped ellipsoid it
133+ * leaves a band of directional high-degree modes with amplification > 1,
134+ * which inflates the pattern's spectrum at moderate niter/lmax and
135+ * diverges at larger ones. The plain scheme (Jhat = 1) diverges wherever
136+ * muMax > 2. The solve's fixed point never depends on Jhat; only the
137+ * convergence rate does.
138+ */
139+ readonly Jhat: number;
140+ /** Symbol-eigenvalue range over the grid (see Jhat), for diagnostics. */
141+ readonly muMin: number;
142+ readonly muMax: number;
143+ /** Area-factor range over the grid, for diagnostics. */
144+ readonly Jmin: number;
145+ readonly Jmax: number;
146+
147+ private constructor(init: {
148+ x: Float32Array; y: Float32Array; z: Float32Array;
149+ X: Float32Array; Y: Float32Array; Z: Float32Array;
150+ Vtx: Float32Array; Vty: Float32Array; Vtz: Float32Array;
151+ Vpx: Float32Array; Vpy: Float32Array; Vpz: Float32Array;
152+ p1: Float32Array; p2: Float32Array; q2: Float32Array; r: Float32Array;
153+ dp1: Float32Array; dq2: Float32Array; jinv: Float32Array;
154+ Jhat: number; muMin: number; muMax: number; Jmin: number; Jmax: number;
155+ }) {
156+ this.x = init.x;
157+ this.y = init.y;
158+ this.z = init.z;
159+ this.X = init.X;
160+ this.Y = init.Y;
161+ this.Z = init.Z;
162+ this.Vtx = init.Vtx;
163+ this.Vty = init.Vty;
164+ this.Vtz = init.Vtz;
165+ this.Vpx = init.Vpx;
166+ this.Vpy = init.Vpy;
167+ this.Vpz = init.Vpz;
168+ this.p1 = init.p1;
169+ this.p2 = init.p2;
170+ this.q2 = init.q2;
171+ this.r = init.r;
172+ this.dp1 = init.dp1;
173+ this.dq2 = init.dq2;
174+ this.jinv = init.jinv;
175+ this.Jhat = init.Jhat;
176+ this.muMin = init.muMin;
177+ this.muMax = init.muMax;
178+ this.Jmin = init.Jmin;
179+ this.Jmax = init.Jmax;
180+ }
181+
182+ /**
183+ * Evaluate the shape file once on the solver grid and reduce it to
184+ * coefficients. Everything here happens at build time — a geometry never
185+ * takes part in the timestep — so the .m runs on the CPU (see
186+ * `evaluateShape`) and only the analysis onward touches the GPU.
187+ */
188+ static async create(opts: GeometryOptions): Promise<Geometry> {
189+ const { sht, cfg, source, paramNames, params, deriv } = opts;
190+ const npts = cfg.nlat * cfg.nphi;
191+
192+ const { theta, phi } = gridAngles(sht, cfg);
193+ const raw = evaluateShape(source, paramNames, params, theta, phi, npts);
194+
195+ // Coefficients first, then back to the grid: what the solver and the
196+ // renderer both see is the band-limited surface, not the raw .m output.
197+ const [X, Y, Z] = [
198+ await sht.analys(raw[0]),
199+ await sht.analys(raw[1]),
200+ await sht.analys(raw[2]),
201+ ];
202+ const [x, y, z] = [
203+ await sht.synth(X),
204+ await sht.synth(Y),
205+ await sht.synth(Z),
206+ ];
207+
208+ // Inverse metric quantities (algos.tex Algorithm 2): theta/phi
209+ // derivatives of the embedding's coefficients, contracted through the
210+ // inverse first fundamental form. Depends only on the geometry, so
211+ // this is a one-off alongside x,y,z above, not per-step work.
212+ const Xt = await deriv.dtheta(X);
213+ const Xp = await deriv.dphi(X);
214+ const Yt = await deriv.dtheta(Y);
215+ const Yp = await deriv.dphi(Y);
216+ const Zt = await deriv.dtheta(Z);
217+ const Zp = await deriv.dphi(Z);
218+ const { Vtx, Vty, Vtz, Vpx, Vpy, Vpz } = computeMetric(npts, Xt, Xp, Yt, Yp, Zt, Zp);
219+
220+ // Flux-form metric weights for the six-transform scheme, built from the
221+ // *undivided* theta tangents sin(theta)*X_theta (smooth on the sphere,
222+ // unlike X_theta itself) and the same X_phi as above. Also a one-off;
223+ // the f64 combination happens on the CPU, rounded to f32 for upload.
224+ const sXtx = await deriv.sinDtheta(X);
225+ const sXty = await deriv.sinDtheta(Y);
226+ const sXtz = await deriv.sinDtheta(Z);
227+ const flux = computeFluxMetric(npts, sXtx, sXty, sXtz, Xp, Yp, Zp);
228+
229+ // The preconditioner scale — see the Jhat field comment. The symbol
230+ // matrix in the orthonormal frame is S = (1/J)[[p1,p2],[p2,q2]] with
231+ // 1/J = r sin^2(theta); its entries are the bounded quantities
232+ // g^tt, sin g^tp, sin^2 g^pp, so the eigenvalue extremes are clean to
233+ // take over the grid. det S = 1/J^2, so the area factor comes along
234+ // for free. f64 throughout.
235+ let muMin = Infinity;
236+ let muMax = 0;
237+ let Jmin = Infinity;
238+ let Jmax = 0;
239+ // The sphere-subtracted weights ride along on this loop: 1/J is already
240+ // being formed here, and dp1/dq2 want the same f64 arithmetic.
241+ const dp1 = new Float32Array(npts);
242+ const dq2 = new Float32Array(npts);
243+ const jinv = new Float32Array(npts);
244+ for (let i = 0; i < cfg.nlat; i++) {
245+ const ct = sht.cosTheta[i];
246+ const st2 = Math.max(0, 1 - ct * ct);
247+ for (let j = 0; j < cfg.nphi; j++) {
248+ const k = i * cfg.nphi + j;
249+ const invJ = flux.r[k] * st2;
250+ dp1[k] = flux.p1[k] - 1;
251+ dq2[k] = flux.q2[k] - 1;
252+ jinv[k] = invJ;
253+ const s11 = flux.p1[k] * invJ;
254+ const s12 = flux.p2[k] * invJ;
255+ const s22 = flux.q2[k] * invJ;
256+ const mean = (s11 + s22) / 2;
257+ const disc = Math.sqrt(((s11 - s22) / 2) ** 2 + s12 * s12);
258+ if (mean - disc < muMin) muMin = mean - disc;
259+ if (mean + disc > muMax) muMax = mean + disc;
260+ const J = 1 / invJ;
261+ if (J < Jmin) Jmin = J;
262+ if (J > Jmax) Jmax = J;
263+ }
264+ }
265+ const Jhat = 2 / (muMin + muMax);
266+
267+ return new Geometry({
268+ x, y, z, X, Y, Z, Vtx, Vty, Vtz, Vpx, Vpy, Vpz,
269+ p1: new Float32Array(flux.p1),
270+ p2: new Float32Array(flux.p2),
271+ q2: new Float32Array(flux.q2),
272+ r: new Float32Array(flux.r),
273+ dp1, dq2, jinv,
274+ Jhat, muMin, muMax, Jmin, Jmax,
275+ });
276+ }
277+
278+ /**
279+ * The surface evaluated on another plan's grid, as interleaved xyz vertex
280+ * positions (nlat * nphi * 3) — for rendering at display resolution. Exact
281+ * interpolation: the same coefficients, more evaluation points.
282+ */
283+ async positionsOn(view: ShtPlan): Promise<Float32Array> {
284+ const [x, y, z] = [
285+ await view.synth(this.X),
286+ await view.synth(this.Y),
287+ await view.synth(this.Z),
288+ ];
289+ const out = new Float32Array(x.length * 3);
290+ for (let i = 0; i < x.length; i++) {
291+ out[3 * i] = x[i];
292+ out[3 * i + 1] = y[i];
293+ out[3 * i + 2] = z[i];
294+ }
295+ return out;
296+ }
297+
298+ /** How far the surface departs from the unit sphere, as min/max radius. */
299+ radiusRange(): { lo: number; hi: number } {
300+ let lo = Infinity;
301+ let hi = -Infinity;
302+ for (let i = 0; i < this.x.length; i++) {
303+ const r = Math.hypot(this.x[i], this.y[i], this.z[i]);
304+ if (r < lo) lo = r;
305+ if (r > hi) hi = r;
306+ }
307+ return { lo, hi };
308+ }
309+}
310+
311+/** The (theta, phi) of every grid point, flattened phi-fastest as the fields
312+ * are — in f64, the precision the shape is evaluated at. */
313+function gridAngles(
314+ sht: ShtPlan,
315+ cfg: ShtConfig,
316+): { theta: Float64Array; phi: Float64Array } {
317+ const { nlat, nphi } = cfg;
318+ const theta = new Float64Array(nlat * nphi);
319+ const phi = new Float64Array(nlat * nphi);
320+ for (let i = 0; i < nlat; i++) {
321+ const th = Math.acos(Math.max(-1, Math.min(1, sht.cosTheta[i])));
322+ for (let j = 0; j < nphi; j++) {
323+ theta[i * nphi + j] = th;
324+ phi[i * nphi + j] = (2 * Math.PI * j) / nphi;
325+ }
326+ }
327+ return { theta, phi };
328+}
329+
330+/**
331+ * Evaluate the shape file on the grid, through numbl's CPU interpreter.
332+ *
333+ * The .m keeps the same contract it had as a compiled model: it names the
334+ * arguments it wants — `theta`, `phi`, and any of the registry's parameters —
335+ * and the host supplies them by name, so their order in the signature is the
336+ * .m's own business. A one-line driver script calls `shape` with exactly the
337+ * arguments its signature declares, with those names pre-bound in the
338+ * driver's workspace.
339+ */
340+function evaluateShape(
341+ source: string,
342+ paramNames: string[],
343+ params: ModelParams,
344+ theta: Float64Array,
345+ phi: Float64Array,
346+ npts: number,
347+): [Float32Array, Float32Array, Float32Array] {
348+ const file = `${SHAPE_FN}.m`;
349+ const ast = inModel(() => parseMFile(source, file));
350+ const fn = ast.body.find(
351+ (s): s is FunctionStmt =>
352+ s.type === 'Function' && (s as FunctionStmt).name === SHAPE_FN,
353+ );
354+ if (!fn) {
355+ throw new ModelCompileError(
356+ `the geometry defines no function named '${SHAPE_FN}'`,
357+ );
358+ }
359+ if (fn.outputs.length !== 3) {
360+ throw new ModelCompileError(
361+ `'${SHAPE_FN}' must return three outputs [gx, gy, gz], not ${fn.outputs.length}`,
362+ { fn: SHAPE_FN, start: fn.span.start, end: fn.span.end },
363+ );
364+ }
365+ const known = new Set(['theta', 'phi', ...paramNames]);
366+ for (const p of fn.params) {
367+ if (!known.has(p)) {
368+ throw new ModelCompileError(
369+ `'${SHAPE_FN}' takes an argument '${p}' that is neither the grid ` +
370+ `(theta, phi) nor one of this geometry's parameters` +
371+ (paramNames.length ? ` (${paramNames.join(', ')})` : ''),
372+ { fn: SHAPE_FN, start: fn.span.start, end: fn.span.end },
373+ );
374+ }
375+ }
376+
377+ const vars: Record<string, RuntimeValue> = {
378+ theta: new RuntimeTensor(theta, [npts, 1]),
379+ phi: new RuntimeTensor(phi, [npts, 1]),
380+ };
381+ for (const name of paramNames) {
382+ const v = params[name];
383+ // Missing parameters read as 0, as ModelPlan.setParams has it.
384+ vars[name] = Number.isFinite(v) ? v : 0;
385+ }
386+
387+ const driver = `[gx__, gy__, gz__] = ${SHAPE_FN}(${fn.params.join(', ')});`;
388+ const result = inFunction(SHAPE_FN, () =>
389+ executeCode(
390+ driver,
391+ { initialVariableValues: vars, displayResults: false, implicitCwdPath: null },
392+ [...toolFiles, { name: file, source }],
393+ 'geometry-driver.m',
394+ ),
395+ );
396+
397+ return [
398+ toGridField(result.variableValues['gx__'], fn.outputs[0], npts),
399+ toGridField(result.variableValues['gy__'], fn.outputs[1], npts),
400+ toGridField(result.variableValues['gz__'], fn.outputs[2], npts),
401+ ];
402+}
403+
404+/** One returned coordinate → npts values, rounded to the transforms' f32. */
405+function toGridField(
406+ value: RuntimeValue | undefined,
407+ name: string,
408+ npts: number,
409+): Float32Array {
410+ // A constant coordinate stays scalar in MATLAB; spread it over the grid.
411+ if (typeof value === 'number') return new Float32Array(npts).fill(value);
412+ if (value !== undefined && isRuntimeTensor(value)) {
413+ if (value.imag) {
414+ throw new ModelCompileError(
415+ `the geometry's '${name}' is complex; coordinates must be real`,
416+ { fn: SHAPE_FN },
417+ );
418+ }
419+ // A vector of npts values, either orientation. A 2-D reshape is refused
420+ // rather than reordered: the tensor's column-major layout would not match
421+ // the grid's phi-fastest rows.
422+ if (value.data.length === npts && value.shape.every((d) => d === 1 || d === npts)) {
423+ return new Float32Array(value.data);
424+ }
425+ throw new ModelCompileError(
426+ `the geometry's '${name}' is ${value.shape.join(' x ')}, but the grid ` +
427+ `wants one value per point (${npts} x 1)`,
428+ { fn: SHAPE_FN },
429+ );
430+ }
431+ throw new ModelCompileError(`the geometry's '${name}' is not numeric`, {
432+ fn: SHAPE_FN,
433+ });
434+}
src/geom/metric.tsadded+146−0View file
@@ -0,0 +1,146 @@
1+/**
2+ * Inverse metric quantities V_theta, V_phi of a surface embedding X=(x,y,z)
3+ * (evolving_surface/notes/algos.tex Algorithm 2 / SurfaceDiffOperator.
4+ * _precompute_metric_quantities, clear_denominators=False branch): six grid
5+ * scalar fields depending only on the geometry, used by the surface
6+ * Laplace-Beltrami operator (Algorithm 3) to contract a field's theta/phi
7+ * derivatives into a tangential gradient/divergence.
8+ *
9+ * g_tt = Xt.Xt, g_tp = Xt.Xp, g_pp = Xp.Xp (first fundamental form)
10+ * det = g_tt*g_pp - g_tp^2
11+ * V_theta = ( g_pp*Xt - g_tp*Xp ) / det
12+ * V_phi = ( g_tt*Xp - g_tp*Xt ) / det
13+ */
14+
15+export interface MetricFields {
16+ /** V_theta, Cartesian components, npts each. */
17+ Vtx: Float32Array;
18+ Vty: Float32Array;
19+ Vtz: Float32Array;
20+ /** V_phi, Cartesian components, npts each. */
21+ Vpx: Float32Array;
22+ Vpy: Float32Array;
23+ Vpz: Float32Array;
24+}
25+
26+/**
27+ * Xt/Xp (etc) are the theta/phi derivatives of each Cartesian embedding
28+ * component, grid space, npts each -- the tangent vectors X_theta, X_phi of
29+ * algos.tex Sec 4.1, one component per array.
30+ */
31+export function computeMetric(
32+ npts: number,
33+ Xt: Float32Array,
34+ Xp: Float32Array,
35+ Yt: Float32Array,
36+ Yp: Float32Array,
37+ Zt: Float32Array,
38+ Zp: Float32Array,
39+): MetricFields {
40+ const Vtx = new Float32Array(npts);
41+ const Vty = new Float32Array(npts);
42+ const Vtz = new Float32Array(npts);
43+ const Vpx = new Float32Array(npts);
44+ const Vpy = new Float32Array(npts);
45+ const Vpz = new Float32Array(npts);
46+
47+ for (let i = 0; i < npts; i++) {
48+ const xt = Xt[i];
49+ const xp = Xp[i];
50+ const yt = Yt[i];
51+ const yp = Yp[i];
52+ const zt = Zt[i];
53+ const zp = Zp[i];
54+
55+ const gtt = xt * xt + yt * yt + zt * zt;
56+ const gtp = xt * xp + yt * yp + zt * zp;
57+ const gpp = xp * xp + yp * yp + zp * zp;
58+ const det = gtt * gpp - gtp * gtp;
59+
60+ Vtx[i] = (gpp * xt - gtp * xp) / det;
61+ Vty[i] = (gpp * yt - gtp * yp) / det;
62+ Vtz[i] = (gpp * zt - gtp * zp) / det;
63+ Vpx[i] = (gtt * xp - gtp * xt) / det;
64+ Vpy[i] = (gtt * yp - gtp * yt) / det;
65+ Vpz[i] = (gtt * zp - gtp * zt) / det;
66+ }
67+
68+ return { Vtx, Vty, Vtz, Vpx, Vpy, Vpz };
69+}
70+
71+/**
72+ * Flux-form metric weights p1, p2, q2, r of the six-transform Laplace-Beltrami
73+ * scheme (docs/reduced-transforms.md Sec 3). Built from the
74+ * *sin-weighted* theta tangent sin(theta)*X_theta — the undivided synthesis of
75+ * the alpha shift, DerivPlan.sinDtheta — and X_phi, both smooth on the sphere:
76+ *
77+ * gtt~ = |sin(theta) X_theta|^2 (= sin^2(theta) g_tt)
78+ * gtp~ = (sin(theta) X_theta).X_phi (= sin(theta) g_tp)
79+ * gpp = |X_phi|^2
80+ * D = sqrt(gtt~ gpp - gtp~^2) (= sin^2(theta) sqrt(det g) / sin(theta)
81+ * = J sin^2(theta), with J = sqrt(det g)/sin(theta))
82+ *
83+ * p1 = gpp / D, p2 = -gtp~ / D, q2 = gtt~ / D, r = 1 / D.
84+ *
85+ * With these, for A = sin(theta) dtheta(u) and B = dphi(u), the two fluxes
86+ *
87+ * P = p1*A + p2*B, Qtilde = p2*A + q2*B
88+ *
89+ * equal sqrt(det g) g^{theta j} u_j and sin(theta) sqrt(det g) g^{phi j} u_j —
90+ * both smooth on the sphere — and Delta_Gamma u = r * (sin(theta) dtheta(P) +
91+ * dphi(Qtilde)). p1, p2, q2 are bounded (the sin^2 in D cancels against the
92+ * vanishing numerators); r ~ 1/sin^2(theta) is finite at the Gauss nodes and
93+ * is the scheme's one concentrated division (Sec 5 of the doc).
94+ *
95+ * All arithmetic is f64 (JS numbers) regardless of the input arrays' storage
96+ * type; results are rounded to f32 only on upload. That is the doc's "CPU
97+ * precompute in float64" mitigation, inherited for free.
98+ */
99+export interface FluxMetricFields {
100+ p1: Float64Array;
101+ p2: Float64Array;
102+ q2: Float64Array;
103+ r: Float64Array;
104+}
105+
106+/**
107+ * sXt* are the Cartesian components of sin(theta)*X_theta, Xp* those of
108+ * X_phi, all grid space, npts each. No sin(theta) input is needed: every
109+ * division the scheme performs is by D, which the sin-weighted inputs build
110+ * directly.
111+ */
112+export function computeFluxMetric(
113+ npts: number,
114+ sXtx: ArrayLike<number>,
115+ sXty: ArrayLike<number>,
116+ sXtz: ArrayLike<number>,
117+ Xpx: ArrayLike<number>,
118+ Xpy: ArrayLike<number>,
119+ Xpz: ArrayLike<number>,
120+): FluxMetricFields {
121+ const p1 = new Float64Array(npts);
122+ const p2 = new Float64Array(npts);
123+ const q2 = new Float64Array(npts);
124+ const r = new Float64Array(npts);
125+
126+ for (let i = 0; i < npts; i++) {
127+ const xt = sXtx[i];
128+ const yt = sXty[i];
129+ const zt = sXtz[i];
130+ const xp = Xpx[i];
131+ const yp = Xpy[i];
132+ const zp = Xpz[i];
133+
134+ const gtt = xt * xt + yt * yt + zt * zt; // sin^2 g_tt
135+ const gtp = xt * xp + yt * yp + zt * zp; // sin g_tp
136+ const gpp = xp * xp + yp * yp + zp * zp; // g_pp
137+ const D = Math.sqrt(gtt * gpp - gtp * gtp); // J sin^2(theta)
138+
139+ p1[i] = gpp / D;
140+ p2[i] = -gtp / D;
141+ q2[i] = gtt / D;
142+ r[i] = 1 / D;
143+ }
144+
145+ return { p1, p2, q2, r };
146+}
src/geom/registry.tsadded+72−0View file
@@ -0,0 +1,72 @@
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 type { ParamSpec, Params } from '../mgpu/registry.ts';
16+
17+export interface MGeometry {
18+ key: string;
19+ label: string;
20+ blurb: string;
21+ params: ParamSpec[];
22+ /** MATLAB source — the shape itself. */
23+ source: string;
24+}
25+
26+const sphere: MGeometry = {
27+ key: 'sphere',
28+ label: 'Sphere',
29+ blurb: 'The unit sphere — the reference case.',
30+ params: [],
31+ source: sphereSource,
32+};
33+
34+const ellipsoid: MGeometry = {
35+ key: 'ellipsoid',
36+ label: 'Ellipsoid',
37+ blurb: 'The sphere with each axis scaled independently.',
38+ params: [
39+ { key: 'ax', label: 'a', value: 1.5, min: 0.2, max: 3, step: 0.05 },
40+ { key: 'ay', label: 'b', value: 1, min: 0.2, max: 3, step: 0.05 },
41+ { key: 'az', label: 'c', value: 0.6, min: 0.2, max: 3, step: 0.05 },
42+ ],
43+ source: ellipsoidSource,
44+};
45+
46+const peanut: MGeometry = {
47+ key: 'peanut',
48+ label: 'Peanut',
49+ blurb: 'A dumbbell pinched at the equator.',
50+ params: [
51+ { key: 'waist', label: 'waist', value: 0.6, min: 0, max: 0.9, step: 0.05 },
52+ { key: 'stretch', label: 'stretch', value: 0.6, min: 0, max: 2, step: 0.05 },
53+ ],
54+ source: peanutSource,
55+};
56+
57+export const mGeometries: MGeometry[] = [sphere, ellipsoid, peanut];
58+
59+export const mGeometryByKey = (key: string): MGeometry | undefined =>
60+ mGeometries.find((g) => g.key === key);
61+
62+export const defaultGeometryParams = (g: MGeometry): Params =>
63+ Object.fromEntries(g.params.map((p) => [p.key, p.value]));
64+
65+/** The geometry every result is checked against, and what a caller who names
66+ * none gets: the case where the solver is exact. */
67+export const SPHERE_KEY = 'sphere';
68+
69+/** What the app and the benchmark start on. Not the sphere: this project
70+ * exists for the other shapes, and opening on the reference case would hide
71+ * the one thing it adds. */
72+export const DEFAULT_GEOMETRY_KEY = 'ellipsoid';
src/main.tsadded+997−0View file
@@ -0,0 +1,997 @@
1+/**
2+ * turing-surface-cache: reaction-diffusion solutions at a chosen end time,
3+ * from a shared cloud cache when someone has computed them before, and from
4+ * the local GPU when not.
5+ *
6+ * Every control is a choice from a short list (src/cache/options.ts), so the
7+ * page's whole state is one small spec object. Get solution hashes that spec
8+ * into a cache object name (src/cache/spec.ts) and fetches it; a 404 means
9+ * nobody has computed it, so the solver runs here — live, watching the
10+ * pattern form — and stops at exactly the requested time. A run to T passes
11+ * exactly through every smaller listed end time, so those states are captured
12+ * along the way; with an upload API key entered, all of them are contributed
13+ * back to the cache.
14+ *
15+ * The solver is turing-surface's, unchanged: the model and geometry are
16+ * MATLAB compiled (model) or interpreted (geometry) by numbl, the transforms
17+ * are WGSL compute shaders. lmax, niter and the seed wavelength are fixed in
18+ * this app (options.ts) — fewer knobs, same machinery.
19+ */
20+import { requestShtDevice, describeAdapter } from './sht/sht.ts';
21+import { ModelSession } from './mgpu/session.ts';
22+import { mModels, type MModel, type Params } from './mgpu/registry.ts';
23+import { formatFailure } from './mgpu/errors.ts';
24+import {
25+ mGeometryByKey,
26+ DEFAULT_GEOMETRY_KEY,
27+ mGeometries,
28+ type MGeometry,
29+} from './geom/registry.ts';
30+import {
31+ buildTopology,
32+ fillPositions,
33+ fillFieldValues,
34+ fillColors,
35+ type SphereMeshTopology,
36+} from './render/sphereMesh.ts';
37+import { SphereScene } from './render/SphereScene.ts';
38+import { Colorbar, floorRange } from './render/colorbar.ts';
39+import { colormaps } from './render/colormaps.ts';
40+import {
41+ MODEL_CHOICES,
42+ GEOMETRY_CHOICES,
43+ SEED_CHOICE,
44+ T_END_CHOICE,
45+ LMAX,
46+ NITER,
47+ LAM3,
48+ defaultChoiceParams,
49+ fmtChoice,
50+ type DiscreteChoice,
51+} from './cache/options.ts';
52+import { stepsFor, type CacheSpec, APP_NAME, FORMAT_VERSION } from './cache/spec.ts';
53+import { lookupFor, fetchCached, uploadCacheFile, type CacheLookup } from './cache/client.ts';
54+import { encodeCacheFile, decodeCacheFile, type DecodedCacheFile } from './cache/h5file.ts';
55+
56+const $ = <T extends HTMLElement>(id: string): T =>
57+ document.getElementById(id) as T;
58+
59+const elParams = $('params');
60+const elGeometry = $<HTMLSelectElement>('geometry');
61+const elGeomParams = $('geomparams');
62+const elSeed = $<HTMLSelectElement>('seed');
63+const elTend = $<HTMLSelectElement>('tend');
64+const elSolve = $<HTMLButtonElement>('solve');
65+const elStop = $<HTMLButtonElement>('stop');
66+const elReset = $<HTMLButtonElement>('reset');
67+const elCacheNote = $('cachenote');
68+const elStatus = $('status');
69+const elPanels = $('panels');
70+const elResetView = $<HTMLButtonElement>('resetview');
71+const elDownload = $<HTMLAnchorElement>('download');
72+const elStats = $('stats');
73+const elApiKey = $<HTMLInputElement>('apikey');
74+const elUploadNote = $('uploadnote');
75+const elErr = $('err');
76+
77+/**
78+ * Test/debug hook: `?tend=5,10` replaces the end-time list with the given
79+ * values (still cached under their own honest specs — a test end time hashes
80+ * to its own object). The headless checks use this to keep their computed
81+ * runs short; it is not part of the normal UI.
82+ */
83+{
84+ const param = new URLSearchParams(location.search).get('tend');
85+ if (param) {
86+ const values = param
87+ .split(',')
88+ .map(Number)
89+ .filter((v) => Number.isFinite(v) && v > 0);
90+ if (values.length) {
91+ T_END_CHOICE.values = values;
92+ T_END_CHOICE.value = values[0];
93+ }
94+ }
95+}
96+
97+const API_KEY_STORAGE = `${APP_NAME}:apiKey`;
98+const COLORMAP = colormaps.viridis;
99+/** Render on a 2x finer grid than the solver's; exact interpolation. */
100+const OVERSAMPLE = 2;
101+/** Cap on GPU dispatches per submission (watchdog safety; see turing-surface). */
102+const DISPATCH_BUDGET = 1000;
103+/** Steps between syncs during a computation: many small submissions queued
104+ * back to back, one wait. The readbacks and renders that pace the live view
105+ * happen per chunk, not per submission — that is what lets the run advance
106+ * at close to the solver's own rate. */
107+const CHUNK_STEPS = 32;
108+/** How often the live view renders during a computation. */
109+const RENDER_EVERY_MS = 250;
110+
111+// ---------------------------------------------------------------- state
112+const model: MModel = mModels[0];
113+let device: GPUDevice | null = null;
114+let session: ModelSession | null = null;
115+let adapterName = '';
116+/** Steps per GPU submission, sized in boot() so one submission stays under
117+ * the dispatch budget however expensive niter has made a step. */
118+let stepsPerSubmit = 4;
119+
120+/** The discrete selections, always exactly values from options.ts. */
121+let params: Params = Object.fromEntries(MODEL_CHOICES.map((c) => [c.key, c.value]));
122+let geometry: MGeometry = mGeometryByKey(DEFAULT_GEOMETRY_KEY)!;
123+let geomParams: Params = Object.fromEntries(
124+ GEOMETRY_CHOICES[DEFAULT_GEOMETRY_KEY].map((c) => [c.key, c.value]),
125+);
126+let seed = SEED_CHOICE.value;
127+let tEnd = T_END_CHOICE.value;
128+
129+// The URL fragment carries the whole selection, so a reload comes back to it
130+// and a shared link opens on the same spec (and, through refresh(), the same
131+// cached solution). Read once at startup; rewritten on every change.
132+readUrlState();
133+
134+/** What the session currently has applied (params are cheap; geometry is a
135+ * rebuild of the surface and the mesh, so it is compared before applying). */
136+let sessionGeomKey = '';
137+let sessionGeomParams: Params = {};
138+
139+let topo: SphereMeshTopology | null = null;
140+let scenes: SphereScene[] = [];
141+let colorbars: Colorbar[] = [];
142+/** The colorbar containers, hidden while the windows are empty. */
143+let colorbarEls: HTMLElement[] = [];
144+let valueBufs: Float32Array[] = [];
145+let colorBufs: Float32Array[] = [];
146+let ranges: { lo: number; hi: number }[] = [];
147+let resizeObs: ResizeObserver | null = null;
148+let coords: Float32Array | null = null;
149+let posBuf: Float32Array | null = null;
150+
151+let generation = 0;
152+let busy = false;
153+/** True while computeLocally is stepping/reading back. Every read shares one
154+ * staging buffer (GpuModel#readback), so a new solve must drain the old
155+ * loop before issuing reads of its own. */
156+let pumping = false;
157+let stopRequested = false;
158+/** Simulation time of the state on display (loadState resets session.t). */
159+let shownT: number | null = null;
160+let downloadUrl: string | null = null;
161+
162+const nextFrame = () => new Promise<number>(requestAnimationFrame);
163+
164+// ---------------------------------------------------------------- spec
165+function currentSpec(): CacheSpec {
166+ return {
167+ app: APP_NAME,
168+ formatVersion: FORMAT_VERSION,
169+ model: model.key,
170+ params: { ...params },
171+ geometry: geometry.key,
172+ geometryParams: { ...geomParams },
173+ lmax: LMAX,
174+ niter: NITER,
175+ lam3: LAM3,
176+ seed,
177+ tEnd,
178+ };
179+}
180+
181+// ---------------------------------------------------------------- URL state
182+/**
183+ * The selection lives in the URL fragment, every value written explicitly
184+ * (`#a=0.1&b=0.9&…&geometry=ellipsoid&ax=1.5&…&seed=1&tend=100`), so a link
185+ * keeps meaning the same spec even if a default changes later. The fragment
186+ * is chosen over the query string to leave `?tend` to the test hook. Values
187+ * are only accepted if they are exactly entries of the discrete lists;
188+ * anything else keeps the default.
189+ */
190+function readUrlState(): void {
191+ const hash = location.hash.replace(/^#/, '');
192+ if (!hash) return;
193+ const p = new URLSearchParams(hash);
194+ // `name` is the key as it appears in the URL; it defaults to the choice's
195+ // own key but is passed explicitly where the two differ (tEnd vs tend).
196+ const pick = (choice: DiscreteChoice, current: number, name = choice.key): number => {
197+ const raw = p.get(name);
198+ if (raw === null) return current;
199+ const v = Number(raw);
200+ return choice.values.includes(v) ? v : current;
201+ };
202+ const g = p.get('geometry');
203+ if (g && mGeometryByKey(g) && GEOMETRY_CHOICES[g]) {
204+ geometry = mGeometryByKey(g)!;
205+ geomParams = defaultChoiceParams(GEOMETRY_CHOICES[g]);
206+ }
207+ for (const c of MODEL_CHOICES) params[c.key] = pick(c, params[c.key]);
208+ for (const c of GEOMETRY_CHOICES[geometry.key]) geomParams[c.key] = pick(c, geomParams[c.key]);
209+ seed = pick(SEED_CHOICE, seed);
210+ tEnd = pick(T_END_CHOICE, tEnd, 'tend');
211+}
212+
213+function writeUrlState(): void {
214+ const p = new URLSearchParams();
215+ for (const c of MODEL_CHOICES) p.set(c.key, fmtChoice(params[c.key]));
216+ p.set('geometry', geometry.key);
217+ for (const c of GEOMETRY_CHOICES[geometry.key]) p.set(c.key, fmtChoice(geomParams[c.key]));
218+ p.set('seed', String(seed));
219+ p.set('tend', fmtChoice(tEnd));
220+ history.replaceState(null, '', `${location.pathname}${location.search}#${p.toString()}`);
221+}
222+
223+// ---------------------------------------------------------------- controls
224+/** Every select made by makeSelect, so a reset can push new values into the
225+ * ones still on the page. */
226+const boundSelects: { el: HTMLSelectElement; get: () => number }[] = [];
227+
228+function syncSelects(): void {
229+ for (const b of boundSelects) {
230+ if (b.el.isConnected) b.el.value = String(b.get());
231+ }
232+}
233+
234+function makeSelect(
235+ choice: DiscreteChoice,
236+ get: () => number,
237+ set: (v: number) => void,
238+): HTMLLabelElement {
239+ const label = document.createElement('label');
240+ label.textContent = `${choice.label} `;
241+ const select = document.createElement('select');
242+ for (const v of choice.values) {
243+ const opt = document.createElement('option');
244+ opt.value = String(v);
245+ opt.textContent = fmtChoice(v);
246+ select.append(opt);
247+ }
248+ select.value = String(get());
249+ select.addEventListener('change', () => {
250+ set(Number(select.value));
251+ onSelectionChange();
252+ });
253+ label.append(select);
254+ boundSelects.push({ el: select, get });
255+ return label;
256+}
257+
258+/** Put every selection back to its default and refresh. */
259+function resetDefaults(): void {
260+ params = defaultChoiceParams(MODEL_CHOICES);
261+ geometry = mGeometryByKey(DEFAULT_GEOMETRY_KEY)!;
262+ geomParams = defaultChoiceParams(GEOMETRY_CHOICES[DEFAULT_GEOMETRY_KEY]);
263+ seed = SEED_CHOICE.value;
264+ tEnd = T_END_CHOICE.value;
265+ elGeometry.value = geometry.key;
266+ buildGeomParamControls();
267+ elSeed.value = String(seed);
268+ elTend.value = String(tEnd);
269+ syncSelects();
270+ onSelectionChange();
271+}
272+
273+function buildControls(): void {
274+ for (const choice of MODEL_CHOICES) {
275+ elParams.append(
276+ makeSelect(choice, () => params[choice.key], (v) => (params[choice.key] = v)),
277+ );
278+ }
279+ for (const g of mGeometries) {
280+ const opt = document.createElement('option');
281+ opt.value = g.key;
282+ opt.textContent = g.label.toLowerCase();
283+ elGeometry.append(opt);
284+ }
285+ elGeometry.value = geometry.key;
286+ elGeometry.addEventListener('change', () => {
287+ geometry = mGeometryByKey(elGeometry.value)!;
288+ geomParams = Object.fromEntries(
289+ GEOMETRY_CHOICES[geometry.key].map((c) => [c.key, c.value]),
290+ );
291+ buildGeomParamControls();
292+ onSelectionChange();
293+ });
294+ buildGeomParamControls();
295+
296+ for (const v of SEED_CHOICE.values) {
297+ const opt = document.createElement('option');
298+ opt.value = String(v);
299+ opt.textContent = String(v);
300+ elSeed.append(opt);
301+ }
302+ elSeed.value = String(seed);
303+ elSeed.addEventListener('change', () => {
304+ seed = Number(elSeed.value);
305+ onSelectionChange();
306+ });
307+
308+ for (const v of T_END_CHOICE.values) {
309+ const opt = document.createElement('option');
310+ opt.value = String(v);
311+ opt.textContent = String(v);
312+ elTend.append(opt);
313+ }
314+ elTend.value = String(tEnd);
315+ elTend.addEventListener('change', () => {
316+ tEnd = Number(elTend.value);
317+ onSelectionChange();
318+ });
319+}
320+
321+function buildGeomParamControls(): void {
322+ elGeomParams.replaceChildren();
323+ for (const choice of GEOMETRY_CHOICES[geometry.key]) {
324+ elGeomParams.append(
325+ makeSelect(choice, () => geomParams[choice.key], (v) => (geomParams[choice.key] = v)),
326+ );
327+ }
328+}
329+
330+/**
331+ * A selection change refreshes the display: a cached solution loads and
332+ * shows immediately, an uncached one shows empty surfaces until the user
333+ * explicitly presses Compute solution. While a computation is running the
334+ * change touches nothing — the run keeps going and only the is-it-cached
335+ * note follows the dropdowns.
336+ *
337+ * Refreshes and button presses are chained so two flows never talk to the
338+ * session at once.
339+ */
340+let flowChain: Promise<void> = Promise.resolve();
341+function onSelectionChange(): void {
342+ writeUrlState();
343+ // During a computation the refresh is deferred until the run finishes; the
344+ // is-it-cached note should follow the dropdowns right away regardless.
345+ if (busy) void updateCacheNote();
346+ flowChain = flowChain.then(() => refresh()).catch(() => undefined);
347+}
348+
349+// The note carries a token so a slow HEAD for a superseded selection never
350+// overwrites the note for the current one.
351+let cacheNoteToken = 0;
352+async function updateCacheNote(): Promise<void> {
353+ const token = ++cacheNoteToken;
354+ elCacheNote.textContent = '';
355+ let lookup: CacheLookup;
356+ try {
357+ lookup = await lookupFor(currentSpec());
358+ } catch {
359+ return;
360+ }
361+ let present: boolean | null = null;
362+ try {
363+ const res = await fetch(lookup.url, { method: 'HEAD', cache: 'no-store' });
364+ present = res.ok ? true : res.status === 404 ? false : null;
365+ } catch {
366+ present = null;
367+ }
368+ if (token !== cacheNoteToken) return;
369+ setCacheNote(present);
370+}
371+
372+function setCacheNote(present: boolean | null): void {
373+ if (present === true) {
374+ elCacheNote.innerHTML = '<b>✓ in the cloud cache</b>';
375+ } else if (present === false) {
376+ elCacheNote.textContent = 'not cached yet';
377+ } else {
378+ elCacheNote.textContent = '';
379+ }
380+}
381+
382+// ---------------------------------------------------------------- view
383+function disposeView(): void {
384+ for (const s of scenes) s.dispose();
385+ scenes = [];
386+ colorbars = [];
387+ colorbarEls = [];
388+ topo = null;
389+ coords = null;
390+ posBuf = null;
391+ resizeObs?.disconnect();
392+ resizeObs = null;
393+ elPanels.replaceChildren();
394+}
395+
396+function buildView(surface: Float32Array): void {
397+ if (!session) return;
398+ const view = session.viewSht;
399+ const { nphi } = view.cfg;
400+ const phi = new Float64Array(nphi);
401+ for (let j = 0; j < nphi; j++) phi[j] = (2 * Math.PI * j) / nphi;
402+ topo = buildTopology(view.cosTheta, phi);
403+ coords = surface;
404+ posBuf = new Float32Array(topo.numVertices * 3);
405+ fillPositions(posBuf, coords, topo, 1);
406+
407+ const sphereBg = getComputedStyle(document.documentElement)
408+ .getPropertyValue('--sphere-bg')
409+ .trim();
410+ for (let k = 0; k < model.species.length; k++) {
411+ const panel = document.createElement('div');
412+ panel.className = 'panel';
413+ const box = document.createElement('div');
414+ box.className = 'sphere-box';
415+ const tag = document.createElement('div');
416+ tag.className = 'species-tag';
417+ tag.textContent = model.species[k];
418+ box.append(tag);
419+ const side = document.createElement('div');
420+ panel.append(box, side);
421+ elPanels.append(panel);
422+
423+ const scene = new SphereScene(
424+ box,
425+ topo.numVertices,
426+ topo.indices,
427+ Float32Array.from(posBuf),
428+ sphereBg || undefined,
429+ );
430+ scene.fitCamera();
431+ scenes.push(scene);
432+ colorbars.push(new Colorbar(side));
433+ colorbarEls.push(side);
434+ valueBufs[k] = new Float32Array(topo.numVertices);
435+ colorBufs[k] = new Float32Array(topo.numVertices * 3);
436+ ranges[k] = { lo: NaN, hi: NaN };
437+ }
438+ for (let k = 1; k < scenes.length; k++) scenes[0].syncCamerasWith(scenes[k]);
439+
440+ resizeObs = new ResizeObserver(() => {
441+ const boxes = elPanels.querySelectorAll<HTMLElement>('.sphere-box');
442+ boxes.forEach((box, i) => {
443+ scenes[i]?.resize(box.clientWidth, box.clientHeight);
444+ });
445+ });
446+ elPanels
447+ .querySelectorAll<HTMLElement>('.sphere-box')
448+ .forEach((box) => resizeObs!.observe(box));
449+}
450+
451+async function draw(): Promise<void> {
452+ if (!session || !topo) return;
453+ const gen = generation;
454+ for (let k = 0; k < model.species.length; k++) {
455+ let field: Float32Array;
456+ try {
457+ field = await session.readSpecies(k);
458+ } catch (e) {
459+ if (gen !== generation) return;
460+ throw e;
461+ }
462+ if (gen !== generation || !topo) return;
463+ fillFieldValues(valueBufs[k], field, topo);
464+ let lo = Infinity;
465+ let hi = -Infinity;
466+ for (const v of valueBufs[k]) {
467+ if (v < lo) lo = v;
468+ if (v > hi) hi = v;
469+ }
470+ // Smooth the color range in both directions so the shading evolves gently
471+ // as the pattern grows (out-of-range values clamp meanwhile).
472+ const r = ranges[k];
473+ if (!Number.isFinite(r.lo)) {
474+ r.lo = lo;
475+ r.hi = hi;
476+ } else {
477+ const a = 0.15;
478+ r.lo += a * (lo - r.lo);
479+ r.hi += a * (hi - r.hi);
480+ }
481+ const shown = floorRange(r.lo, r.hi);
482+ fillColors(colorBufs[k], valueBufs[k], shown.lo, shown.hi, COLORMAP);
483+ scenes[k]?.updateColors(colorBufs[k]);
484+ colorbars[k]?.update(COLORMAP, shown.lo, shown.hi);
485+ if (colorbarEls[k]) colorbarEls[k].style.visibility = '';
486+ }
487+}
488+
489+/** Empty windows: the selected surface with no field on it. Shown when the
490+ * selection has no cached solution and nothing has been computed yet. */
491+function clearDisplay(): void {
492+ shownT = null;
493+ elDownload.hidden = true;
494+ if (!topo) return;
495+ for (let k = 0; k < model.species.length; k++) {
496+ // NaN renders as neutral gray in fillColors — the shape without a field.
497+ valueBufs[k].fill(NaN);
498+ fillColors(colorBufs[k], valueBufs[k], 0, 1, COLORMAP);
499+ scenes[k]?.updateColors(colorBufs[k]);
500+ if (colorbarEls[k]) colorbarEls[k].style.visibility = 'hidden';
501+ }
502+ updateStats();
503+}
504+
505+function resetRanges(): void {
506+ for (const r of ranges) {
507+ r.lo = NaN;
508+ r.hi = NaN;
509+ }
510+}
511+
512+function updateStats(): void {
513+ if (!session) return;
514+ const { nlat, nphi } = session.cfg;
515+ const kind = `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`;
516+ const t = shownT !== null ? ` · showing t = <b>${fmtChoice(shownT)}</b>` : '';
517+ elStats.innerHTML =
518+ `<b>${kind}</b> · grid ${nlat}×${nphi} · lmax ${LMAX} · ` +
519+ `solve iters ${NITER}${t}`;
520+}
521+
522+// ---------------------------------------------------------------- statuses
523+function status(html: string): void {
524+ elStatus.innerHTML = html;
525+}
526+
527+function setBusy(next: boolean): void {
528+ busy = next;
529+ elSolve.disabled = next;
530+ elStop.hidden = !next;
531+}
532+
533+function offerDownload(bytes: Uint8Array, name: string): void {
534+ if (downloadUrl) URL.revokeObjectURL(downloadUrl);
535+ downloadUrl = URL.createObjectURL(new Blob([bytes as BlobPart], { type: 'application/x-hdf5' }));
536+ elDownload.href = downloadUrl;
537+ elDownload.download = name;
538+ elDownload.hidden = false;
539+}
540+
541+// ---------------------------------------------------------------- solving
542+/** Apply the current selection to the (one, reused) session: params are a
543+ * uniform upload; a geometry change re-evaluates the surface and rebuilds
544+ * the mesh, keeping the camera. */
545+async function applySelection(spec: CacheSpec): Promise<void> {
546+ if (!session) throw new Error('no session');
547+ session.setParams(spec.params);
548+ const geomChanged =
549+ spec.geometry !== sessionGeomKey ||
550+ JSON.stringify(spec.geometryParams) !== JSON.stringify(sessionGeomParams);
551+ if (!geomChanged) return;
552+ const geomModel = mGeometryByKey(spec.geometry)!;
553+ await session.setGeometry(geomModel, spec.geometryParams);
554+ sessionGeomKey = spec.geometry;
555+ sessionGeomParams = { ...spec.geometryParams };
556+ const surface = await session.renderPositions();
557+ const cam = scenes[0]?.cameraState();
558+ disposeView();
559+ buildView(surface);
560+ if (cam) for (const s of scenes) s.setCameraState(cam);
561+ // Fresh buffers render black until the first fill; show the bare surface
562+ // instead. The caller's draw or clearDisplay follows right behind.
563+ clearDisplay();
564+}
565+
566+/** Decode a fetched cache file and put it on screen. */
567+async function displayCached(
568+ bytes: Uint8Array,
569+ lookup: CacheLookup,
570+ spec: CacheSpec,
571+ gen: number,
572+): Promise<void> {
573+ if (!session) return;
574+ const decoded = await decodeCacheFile(bytes, lookup.specJson, model.state);
575+ if (gen !== generation) return;
576+ session.loadState(decoded.final);
577+ shownT = spec.tEnd;
578+ resetRanges();
579+ await draw();
580+ updateStats();
581+ const kb = (bytes.length / 1024).toFixed(0);
582+ const from = decoded.adapter ? `, computed on ${decoded.adapter}` : '';
583+ const when = decoded.created ? ` ${decoded.created.slice(0, 10)}` : '';
584+ status(
585+ `<b>t = ${fmtChoice(spec.tEnd)}</b> — from the <b>cloud cache</b> ` +
586+ `(${kb} KB${from}${when}).`,
587+ );
588+ offerDownload(bytes, lookup.fileName.split('/').pop()!);
589+}
590+
591+/**
592+ * Bring the display in line with the current selection, without ever
593+ * starting a computation: a cached solution loads and shows, an uncached one
594+ * shows empty surfaces and waits for the Compute solution button. Runs on
595+ * startup and on every selection change; a no-op while a computation is
596+ * running (the run is not disturbed — only the cache note follows).
597+ */
598+async function refresh(): Promise<void> {
599+ if (!session || busy) {
600+ void updateCacheNote();
601+ return;
602+ }
603+ generation++;
604+ const gen = generation;
605+ elErr.textContent = '';
606+ const spec = currentSpec();
607+ try {
608+ const lookup = await lookupFor(spec);
609+ status('checking the cloud cache…');
610+ let bytes: Uint8Array | null = null;
611+ let unreachable = false;
612+ try {
613+ bytes = await fetchCached(lookup);
614+ } catch {
615+ unreachable = true;
616+ }
617+ if (gen !== generation) return;
618+ await applySelection(spec);
619+ if (gen !== generation) return;
620+ if (bytes) {
621+ await displayCached(bytes, lookup, spec, gen);
622+ setCacheNote(true);
623+ return;
624+ }
625+ clearDisplay();
626+ setCacheNote(unreachable ? null : false);
627+ status(
628+ unreachable
629+ ? 'cloud cache unreachable — <b>Compute solution</b> runs it in your browser.'
630+ : `not in the cloud cache — press <b>Compute solution</b> to run it in ` +
631+ `your browser (up to ${stepsFor(spec).toLocaleString()} steps; a ` +
632+ `cached shorter run of the same settings is picked up where it left off).`,
633+ );
634+ } catch (e) {
635+ if (gen === generation) {
636+ elErr.textContent = formatFailure(e, model.source);
637+ status('failed.');
638+ }
639+ }
640+}
641+
642+/** The Compute solution button: cache lookup, then either load or compute. */
643+async function solve(): Promise<void> {
644+ if (!session || busy) return;
645+ generation++;
646+ const gen = generation;
647+ setBusy(true);
648+ // A stopped run may still be inside an await; let it see the generation
649+ // bump and finish before touching the session.
650+ while (pumping) await nextFrame();
651+ if (gen !== generation) return;
652+ stopRequested = false;
653+ elErr.textContent = '';
654+ elDownload.hidden = true;
655+ const spec = currentSpec();
656+ try {
657+ const lookup = await lookupFor(spec);
658+ status('checking the cloud cache…');
659+ let bytes: Uint8Array | null = null;
660+ try {
661+ bytes = await fetchCached(lookup);
662+ } catch (e) {
663+ // An unreachable cache degrades to computing locally, and says so.
664+ status(`cache unreachable (${e instanceof Error ? e.message : e}) — computing locally`);
665+ }
666+ if (gen !== generation) return;
667+ await applySelection(spec);
668+ if (gen !== generation) return;
669+
670+ if (bytes) {
671+ await displayCached(bytes, lookup, spec, gen);
672+ return;
673+ }
674+ await computeLocally(spec, gen);
675+ } catch (e) {
676+ if (gen === generation) {
677+ elErr.textContent = formatFailure(e, model.source);
678+ status('failed.');
679+ }
680+ } finally {
681+ if (gen === generation) setBusy(false);
682+ void updateCacheNote();
683+ }
684+}
685+
686+/** Run the solver to the spec's end time, watching the pattern form, and
687+ * capture the state at every smaller listed end time on the way. */
688+async function computeLocally(spec: CacheSpec, gen: number): Promise<void> {
689+ if (!session) return;
690+ pumping = true;
691+ try {
692+ await computeLocallyInner(spec, gen);
693+ } finally {
694+ pumping = false;
695+ }
696+}
697+
698+async function computeLocallyInner(spec: CacheSpec, gen: number): Promise<void> {
699+ if (!session) return;
700+ const steps = stepsFor(spec);
701+ const dt = spec.params.dt;
702+
703+ // Warm start: the state is Markovian in (U, V), so a cached run of the
704+ // same spec at a smaller listed end time is an exact prefix of this one.
705+ // Take the longest one there is and continue from its final state rather
706+ // than recomputing it.
707+ let warm: { tEnd: number; decoded: DecodedCacheFile } | null = null;
708+ const earlier = T_END_CHOICE.values.filter((T) => T < spec.tEnd).sort((a, b) => b - a);
709+ if (earlier.length) status('not in the cache — looking for a shorter cached run…');
710+ for (const T of earlier) {
711+ const lookup = await lookupFor({ ...spec, tEnd: T });
712+ let bytes: Uint8Array | null = null;
713+ try {
714+ bytes = await fetchCached(lookup);
715+ } catch {
716+ break; // cache unreachable: no point probing further down the ladder
717+ }
718+ if (gen !== generation) return;
719+ if (!bytes) continue;
720+ try {
721+ warm = { tEnd: T, decoded: await decodeCacheFile(bytes, lookup.specJson, model.state) };
722+ break;
723+ } catch {
724+ continue; // an unreadable candidate is skipped, not fatal
725+ }
726+ }
727+ if (gen !== generation) return;
728+
729+ let initial: Record<string, Float32Array>;
730+ if (warm) {
731+ session.loadState(warm.decoded.final);
732+ // loadState resets the clock; put it at the cached run's end so the loop
733+ // below computes only the remainder.
734+ session.steps = Math.round(warm.tEnd / dt);
735+ session.t = warm.tEnd;
736+ // The t = 0 state travels with every file of the chain, so files written
737+ // from this continuation carry the same initial state as the one resumed.
738+ initial = warm.decoded.initial;
739+ } else {
740+ status(`not in the cache — <b>computing locally</b>: seeding…`);
741+ await session.seed(spec.seed);
742+ if (gen !== generation) return;
743+ initial = await session.readState();
744+ if (gen !== generation) return;
745+ }
746+ const startSteps = session.steps;
747+
748+ // Snapshot points: every listed end time strictly between the starting
749+ // point and this run's end. The run passes through each exactly (all are
750+ // whole multiples of every dt choice).
751+ const snapshotAt = new Map<number, number>(); // step index -> tEnd value
752+ for (const T of T_END_CHOICE.values) {
753+ if (T < spec.tEnd && T > (warm?.tEnd ?? 0)) snapshotAt.set(Math.round(T / dt), T);
754+ }
755+ const snapshots: { tEnd: number; state: Record<string, Float32Array> }[] = [];
756+
757+ // Everything a cache file needs exists before the run starts, so a snapshot
758+ // is encoded and uploaded the moment it is captured, overlapping the
759+ // network with the GPU still stepping, rather than queued for the end.
760+ const geometryCoeffs = {
761+ X: session.geometry.X,
762+ Y: session.geometry.Y,
763+ Z: session.geometry.Z,
764+ };
765+ const encode = (t: number, state: Record<string, Float32Array>) =>
766+ encodeCacheFile({
767+ spec: { ...spec, tEnd: t },
768+ grid: session!.cfg,
769+ species: model.state,
770+ geometry: geometryCoeffs,
771+ initial,
772+ final: state,
773+ adapter: adapterName,
774+ });
775+ const uploadedTimes: number[] = [];
776+ const uploadErrors: string[] = [];
777+ let uploadsStarted = 0;
778+ const pendingUploads: Promise<void>[] = [];
779+ /** Encode + upload without the stepping loop waiting. A captured snapshot
780+ * is a complete solution of its own spec, so this stays valid even if the
781+ * run is stopped afterwards. */
782+ const uploadInBackground = (
783+ t: number,
784+ state: Record<string, Float32Array>,
785+ apiKey: string,
786+ preEncoded?: Uint8Array,
787+ ): void => {
788+ uploadsStarted++;
789+ pendingUploads.push(
790+ (async () => {
791+ const bytes = preEncoded ?? (await encode(t, state));
792+ const lookup = await lookupFor({ ...spec, tEnd: t });
793+ await uploadCacheFile(apiKey, lookup.fileName, bytes);
794+ uploadedTimes.push(t);
795+ })().catch((e) => {
796+ uploadErrors.push(`t = ${fmtChoice(t)}: ${e instanceof Error ? e.message : e}`);
797+ }),
798+ );
799+ };
800+
801+ shownT = null;
802+ resetRanges();
803+ const t0 = performance.now();
804+ let lastStatus = 0;
805+ let lastDraw = 0;
806+ while (session.steps < steps) {
807+ if (gen !== generation) return;
808+ if (stopRequested) {
809+ shownT = session.steps * dt;
810+ await draw();
811+ updateStats();
812+ const up = uploadedTimes.length
813+ ? ` ${uploadedTimes.length} snapshot${uploadedTimes.length > 1 ? 's' : ''} already uploaded.`
814+ : ' Nothing uploaded.';
815+ status(`stopped at t = ${(session.steps * dt).toFixed(2)}.${up}`);
816+ return;
817+ }
818+ // One chunk: up to CHUNK_STEPS steps submitted back to back (each
819+ // submission stays under the dispatch budget), then a single sync and at
820+ // most one render. Reading back and drawing after every submission is
821+ // what made the run advance at a fraction of the solver's rate — a
822+ // readback costs several times the 3-4 steps it fenced. The chunk stops
823+ // exactly at snapshot points so those states are still captured exactly.
824+ let target = Math.min(steps, session.steps + CHUNK_STEPS);
825+ for (const s of snapshotAt.keys()) {
826+ if (s > session.steps && s < target) target = s;
827+ }
828+ while (session.steps < target) {
829+ session.step(Math.min(stepsPerSubmit, target - session.steps));
830+ }
831+ // The sync bounds how far the CPU runs ahead of the GPU, and (being a
832+ // promise) yields to the event loop, which is what keeps Stop clickable.
833+ await session.sync();
834+ if (gen !== generation) return;
835+ const hit = snapshotAt.get(session.steps);
836+ if (hit !== undefined) {
837+ const state = await session.readState();
838+ if (gen !== generation) return;
839+ // With a key on hand the snapshot goes straight to the cache; without
840+ // one it is kept, in case a key is entered before the run ends.
841+ const apiKey = elApiKey.value.trim();
842+ if (apiKey) uploadInBackground(hit, state, apiKey);
843+ else snapshots.push({ tEnd: hit, state });
844+ }
845+ const now = performance.now();
846+ if (now - lastDraw > RENDER_EVERY_MS || session.steps >= steps) {
847+ lastDraw = now;
848+ await draw();
849+ if (gen !== generation) return;
850+ await nextFrame();
851+ }
852+ if (now - lastStatus > 200) {
853+ lastStatus = now;
854+ const t = session.steps * dt;
855+ const pct = ((100 * (session.steps - startSteps)) / (steps - startSteps)).toFixed(0);
856+ const rate = (session.steps - startSteps) / ((now - t0) / 1000);
857+ const from = warm ? `resumed from cached t = ${fmtChoice(warm.tEnd)} — ` : '';
858+ const up = uploadsStarted
859+ ? `, uploaded ${uploadedTimes.length}/${uploadsStarted} snapshots`
860+ : '';
861+ status(
862+ `not in the cache — <b>computing locally</b> (${from}` +
863+ `t = ${t.toFixed(2)} / ${fmtChoice(spec.tEnd)}, ${pct}%, ${rate.toFixed(0)} steps/s${up})`,
864+ );
865+ }
866+ }
867+
868+ const final = await session.readState();
869+ if (gen !== generation) return;
870+ shownT = spec.tEnd;
871+ await draw();
872+ updateStats();
873+ const secs = ((performance.now() - t0) / 1000).toFixed(1);
874+ const doneLine =
875+ `<b>t = ${fmtChoice(spec.tEnd)}</b> — computed locally in ${secs} s` +
876+ (warm ? ` (resumed from cached t = ${fmtChoice(warm.tEnd)})` : '') +
877+ `.`;
878+ status(`${doneLine} Writing the cache file…`);
879+
880+ const finalBytes = await encode(spec.tEnd, final);
881+ if (gen !== generation) return;
882+ const finalLookup = await lookupFor(spec);
883+ offerDownload(finalBytes, finalLookup.fileName.split('/').pop()!);
884+
885+ // The final solution, plus any snapshots captured before a key was entered.
886+ const apiKey = elApiKey.value.trim();
887+ if (apiKey) {
888+ uploadInBackground(spec.tEnd, final, apiKey, finalBytes);
889+ for (const snap of snapshots) uploadInBackground(snap.tEnd, snap.state, apiKey);
890+ }
891+ if (uploadsStarted === 0) {
892+ status(`${doneLine} Not uploaded (no API key).`);
893+ return;
894+ }
895+ status(`${doneLine} Uploading to the cache (${uploadedTimes.length}/${uploadsStarted})…`);
896+ await Promise.all(pendingUploads);
897+ if (gen !== generation) return;
898+
899+ if (uploadErrors.length) elErr.textContent = `upload: ${uploadErrors.join('; ')}`;
900+ const n = uploadedTimes.length;
901+ if (n > 0) {
902+ const times = [...uploadedTimes].sort((a, b) => a - b).map(fmtChoice).join(', ');
903+ const failed = uploadErrors.length ? ` (${uploadErrors.length} failed)` : '';
904+ status(
905+ `${doneLine} <b>Uploaded ${n} solution${n > 1 ? 's' : ''}</b> ` +
906+ `to the shared cache (t = ${times})${failed}.`,
907+ );
908+ } else {
909+ status(`${doneLine} Uploads failed.`);
910+ }
911+}
912+
913+// ---------------------------------------------------------------- boot
914+elSolve.addEventListener('click', () => {
915+ flowChain = flowChain.then(() => solve()).catch(() => undefined);
916+});
917+elStop.addEventListener('click', () => {
918+ stopRequested = true;
919+ setBusy(false);
920+});
921+elReset.addEventListener('click', () => resetDefaults());
922+elResetView.addEventListener('click', () => {
923+ for (const s of scenes) s.resetCamera();
924+});
925+elApiKey.addEventListener('change', () => {
926+ const key = elApiKey.value.trim();
927+ if (key) localStorage.setItem(API_KEY_STORAGE, key);
928+ else localStorage.removeItem(API_KEY_STORAGE);
929+ updateUploadNote();
930+});
931+
932+function updateUploadNote(): void {
933+ elUploadNote.textContent = elApiKey.value.trim()
934+ ? 'uploads enabled — locally computed solutions will be contributed'
935+ : '';
936+}
937+
938+async function boot(): Promise<void> {
939+ buildControls();
940+ // Written even before any change, so the address bar is always shareable.
941+ writeUrlState();
942+ elApiKey.value = localStorage.getItem(API_KEY_STORAGE) ?? '';
943+ updateUploadNote();
944+ void updateCacheNote();
945+ try {
946+ device = await requestShtDevice();
947+ adapterName = await describeAdapter(device);
948+ } catch (e) {
949+ device = null;
950+ elErr.textContent =
951+ `WebGPU is not available (${e instanceof Error ? e.message : e}). ` +
952+ `Use a WebGPU-capable browser such as Chrome or Edge.`;
953+ return;
954+ }
955+ device.lost.then((info) => {
956+ if (info.reason !== 'destroyed') {
957+ elErr.textContent = `WebGPU device lost: ${info.message}`;
958+ }
959+ });
960+
961+ status('compiling the solver…');
962+ try {
963+ session = await ModelSession.create({
964+ device,
965+ model,
966+ params,
967+ lmax: LMAX,
968+ oversample: OVERSAMPLE,
969+ geometry,
970+ geometryParams: geomParams,
971+ niter: NITER,
972+ lam3: LAM3,
973+ });
974+ } catch (e) {
975+ elErr.textContent = formatFailure(e, model.source);
976+ status('failed to compile.');
977+ return;
978+ }
979+ sessionGeomKey = geometry.key;
980+ sessionGeomParams = { ...geomParams };
981+
982+ // Never put more dispatches in one submission than the budget allows,
983+ // however expensive niter has made one step.
984+ const opsPerStep = Math.max(1, session.describe().step.length);
985+ stepsPerSubmit = Math.max(1, Math.floor(DISPATCH_BUDGET / opsPerStep));
986+
987+ const surface = await session.renderPositions();
988+ buildView(surface);
989+ clearDisplay();
990+ updateStats();
991+ // Bring up the default selection if it is cached; otherwise show empty
992+ // surfaces. Nothing is ever computed without pressing the button.
993+ flowChain = flowChain.then(() => refresh()).catch(() => undefined);
994+ await flowChain;
995+}
996+
997+void boot();
src/mgpu/compile.tsadded+304−0View file
@@ -0,0 +1,304 @@
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+ * A MultiAssignCall assigns every bound output slot. */
220+function assignedCNames(stmts: IRStmt[]): Set<string> {
221+ const out = new Set<string>();
222+ const walk = (list: IRStmt[]): void => {
223+ for (const s of list) {
224+ if (s.kind === 'Assign') out.add(s.cName);
225+ else if (s.kind === 'MultiAssignCall') {
226+ for (const o of s.outputs) if (o.binding) out.add(o.binding.cName);
227+ } else if (s.kind === 'For') walk(s.body);
228+ }
229+ };
230+ walk(stmts);
231+ return out;
232+}
233+
234+/** Call `visit` for every variable read in an expression. */
235+function forEachVarRead(e: IRExpr, visit: (cName: string) => void): void {
236+ const walk = (x: IRExpr): void => {
237+ switch (x.kind) {
238+ case 'Var':
239+ return visit(x.cName);
240+ case 'Binary':
241+ walk(x.left);
242+ walk(x.right);
243+ return;
244+ case 'Unary':
245+ walk(x.operand);
246+ return;
247+ case 'Call':
248+ x.args.forEach(walk);
249+ return;
250+ default:
251+ return;
252+ }
253+ };
254+ walk(e);
255+}
256+
257+/**
258+ * Refuse a loop whose result the inline pass folded away.
259+ *
260+ * numbl's inline pass substitutes a single-use producer into its consumer and
261+ * drops the producer. Inside a loop body it runs with no protected names — it
262+ * counts uses within that body alone — so an assignment whose only *visible*
263+ * use is later in the same body can be elided even though something outside
264+ * the loop still wants the value.
265+ *
266+ * That elision is correct for a body-local temp, which is what makes fusion
267+ * work inside the loop, and it is caught downstream in the two cases where the
268+ * value has no buffer at all: the planner already refuses a declared output
269+ * that is never assigned, and a read of a name it never allocated. The case it
270+ * would not catch is a variable assigned *before* the loop as well — there the
271+ * buffer exists, holding the pre-loop value, and the loop would silently
272+ * contribute nothing. So check all three here, in one place, against what the
273+ * body assigned before the pass ran.
274+ */
275+function checkLoopEscapes(fn: IRFunc, loop: For, assignedBefore: Set<string>): void {
276+ const assignedAfter = assignedCNames(loop.body);
277+ const elided = [...assignedBefore].filter((c) => !assignedAfter.has(c));
278+ if (elided.length === 0) return;
279+
280+ // Reads anywhere in the function outside this loop's own body.
281+ const readOutside = new Set<string>();
282+ const walk = (list: IRStmt[]): void => {
283+ for (const s of list) {
284+ if (s === (loop as IRStmt)) continue; // the loop's own body is not "outside"
285+ if (s.kind === 'Assign') forEachVarRead(s.expr, (c) => readOutside.add(c));
286+ else if (s.kind === 'MultiAssignCall') {
287+ for (const a of s.args) forEachVarRead(a, (c) => readOutside.add(c));
288+ } else if (s.kind === 'For') walk(s.body);
289+ }
290+ };
291+ walk(fn.body);
292+
293+ const outputs = new Set(fn.cOutputs);
294+ const escaping = elided.filter((c) => readOutside.has(c) || outputs.has(c));
295+ if (escaping.length === 0) return;
296+
297+ const names = [...new Set(escaping)].map((c) => `'${c}'`).join(', ');
298+ throw new ModelCompileError(
299+ `inside the 'for' loop, ${names} is assigned but only read later in the ` +
300+ `same iteration, so the compiler folded the assignment into its reader — ` +
301+ `yet the value is also wanted outside the loop. Read it once outside the ` +
302+ `loop instead, or use it more than once inside it.`,
303+ );
304+}
src/mgpu/digest.tsadded+103−0View file
@@ -0,0 +1,103 @@
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+/** Relative L-infinity (max-norm) difference of two states of equal length. */
73+export function relLinf(a: ArrayLike<number>, b: ArrayLike<number>): number {
74+ let num = 0;
75+ let den = 0;
76+ for (let i = 0; i < a.length; i++) {
77+ const d = Math.abs(a[i] - b[i]);
78+ if (d > num) num = d;
79+ const bd = Math.abs(b[i]);
80+ if (bd > den) den = bd;
81+ }
82+ return num / Math.max(den, 1e-300);
83+}
84+
85+export function formatDigest(d: StateDigest): string {
86+ const g = (v: number): string => v.toPrecision(9);
87+ return (
88+ `n=${d.n} min=${g(d.min)} max=${g(d.max)} mean=${g(d.mean)} rms=${g(d.rms)} ` +
89+ `fourier=${d.fourier}`
90+ );
91+}
92+
93+/** Worst relative disagreement between two digests' scalar summaries. */
94+export function digestDrift(a: StateDigest, b: StateDigest): number {
95+ const rel = (x: number, y: number): number =>
96+ Math.abs(x - y) / Math.max(Math.abs(x), Math.abs(y), 1e-30);
97+ return Math.max(
98+ rel(a.min, b.min),
99+ rel(a.max, b.max),
100+ rel(a.mean, b.mean),
101+ rel(a.rms, b.rms),
102+ );
103+}
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+222−0View file
@@ -0,0 +1,222 @@
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`. With `multi`, the op maps each of
39+ * N inputs to its own output — `[a, b] = synth(x, y)` — so the backend can
40+ * run the group as one batched Legendre dispatch (or split it to whatever
41+ * lane width the device supports; the syntax promises grouping intent, not
42+ * a width). */
43+function transformSource(
44+ name: string,
45+ inRows: number,
46+ inCols: number,
47+ outRows: number,
48+ outCols: number,
49+ multi = false,
50+): string {
51+ return `
52+exports.name = ${JSON.stringify(name)};
53+
54+exports.transfer = function (argTypes, nargout) {
55+ ${
56+ multi
57+ ? `if (argTypes.length < 1) {
58+ throw new Error("${name} takes at least one argument");
59+ }
60+ if (nargout > 1 && nargout !== argTypes.length) {
61+ throw new Error(
62+ "${name}: each input produces one output, so " + argTypes.length +
63+ " input(s) return " + argTypes.length + " output(s), but " +
64+ nargout + " were requested -- write [a, b] = ${name}(x, y)"
65+ );
66+ }
67+ if (nargout <= 1 && argTypes.length !== 1) {
68+ throw new Error(
69+ "${name}: " + argTypes.length + " inputs produce " + argTypes.length +
70+ " outputs -- bind each one: [a, b] = ${name}(x, y)"
71+ );
72+ }`
73+ : `if (argTypes.length !== 1) {
74+ throw new Error("${name} takes exactly one argument, got " + argTypes.length);
75+ }
76+ if (nargout > 1) {
77+ throw new Error("${name} returns one value, but " + nargout + " were requested");
78+ }`
79+ }
80+ for (var i = 0; i < argTypes.length; i++) {
81+ var a = argTypes[i];
82+ if (!a || a.kind !== "Numeric" || a.isComplex) {
83+ throw new Error("${name} requires real numeric arrays (argument " + (i + 1) + ")");
84+ }
85+ var s = a.shape;
86+ if (!s || s.length !== 2 || s[0] !== ${inRows} || s[1] !== ${inCols}) {
87+ throw new Error(
88+ "${name} requires ${inRows}x${inCols} arrays, argument " + (i + 1) +
89+ " is " + (s ? s.join("x") : "unknown shape")
90+ );
91+ }
92+ }
93+ var out = [];
94+ for (var k = 0; k < Math.max(1, nargout); k++) {
95+ out.push(${numericType(outRows, outCols)});
96+ }
97+ return out;
98+};
99+
100+// Never called: this project executes the IR on WebGPU and emits no C.
101+exports.emit = function () {
102+ throw new Error("${name}: no C backend (this transform runs on WebGPU)");
103+};
104+exports.cBody = function () {
105+ return "";
106+};
107+`;
108+}
109+
110+/**
111+ * Workspace files that make `synth` / `analys` / `dtheta` / `dphi` /
112+ * `dthetac` / `dphic` resolvable during lowering. `dtheta` and `dphi` (the
113+ * surface's first partial derivatives, coefficients -> grid — see
114+ * src/sht/deriv.ts) have exactly `synth`'s shape rule: both take spectral
115+ * coefficients and produce a grid field. `dthetac` and `dphic` are their
116+ * coefficient-space halves alone — the alpha^+/alpha^- shift and the i*m
117+ * multiply, spectral -> spectral — which the six-transform Laplace-Beltrami
118+ * scheme (docs/reduced-transforms.md) applies twice per
119+ * matvec: to the field (gradient side) and to the analysed fluxes
120+ * (divergence side, the same shift, not its transpose).
121+ */
122+export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
123+ return [
124+ {
125+ name: 'synth.mtoc2.js',
126+ source: transformSource('synth', 2, g.nlm, g.npts, 1, true),
127+ },
128+ {
129+ name: 'analys.mtoc2.js',
130+ source: transformSource('analys', g.npts, 1, 2, g.nlm, true),
131+ },
132+ {
133+ name: 'dtheta.mtoc2.js',
134+ source: transformSource('dtheta', 2, g.nlm, g.npts, 1),
135+ },
136+ {
137+ name: 'dphi.mtoc2.js',
138+ source: transformSource('dphi', 2, g.nlm, g.npts, 1),
139+ },
140+ {
141+ name: 'dthetac.mtoc2.js',
142+ source: transformSource('dthetac', 2, g.nlm, 2, g.nlm),
143+ },
144+ {
145+ name: 'dphic.mtoc2.js',
146+ source: transformSource('dphic', 2, g.nlm, 2, g.nlm),
147+ },
148+ {
149+ // Grid-space phi-derivative: two Fourier stages and an i*m multiply,
150+ // no Legendre work (d/dphi is diagonal in the Fourier index). What
151+ // lets the flux-form divergence skip the Q-flux's spherical-harmonic
152+ // analysis.
153+ name: 'dphig.mtoc2.js',
154+ source: transformSource('dphig', g.npts, 1, g.npts, 1),
155+ },
156+ {
157+ // The seeded random field a model's `init` starts from
158+ // (src/mgpu/randnfun3.ts): a wavelength and the three surface
159+ // coordinates in, one value per grid point out.
160+ name: 'randnfun3.mtoc2.js',
161+ source: randnfun3Source(g),
162+ },
163+ ];
164+}
165+
166+/** Source for `randnfun3`'s `.mtoc2.js`: `f = randnfun3(lambda, x, y, z)`. */
167+function randnfun3Source(g: GridSizes): string {
168+ return `
169+exports.name = "randnfun3";
170+
171+exports.transfer = function (argTypes, nargout) {
172+ if (argTypes.length !== 4) {
173+ throw new Error(
174+ "randnfun3 takes a wavelength and the three surface coordinates -- " +
175+ "randnfun3(lambda, gx, gy, gz) -- got " + argTypes.length + " argument(s)"
176+ );
177+ }
178+ if (nargout > 1) {
179+ throw new Error("randnfun3 returns one value, but " + nargout + " were requested");
180+ }
181+ var lam = argTypes[0];
182+ if (!lam || lam.kind !== "Numeric" || lam.isComplex) {
183+ throw new Error("randnfun3's wavelength must be a real number");
184+ }
185+ var ls = lam.shape;
186+ if (!ls || ls.length !== 2 || ls[0] !== 1 || ls[1] !== 1) {
187+ throw new Error(
188+ "randnfun3's wavelength must be a single number, not a " +
189+ (ls ? ls.join("x") : "unknown shape") + " array"
190+ );
191+ }
192+ var names = ["gx", "gy", "gz"];
193+ for (var i = 1; i < 4; i++) {
194+ var a = argTypes[i];
195+ if (!a || a.kind !== "Numeric" || a.isComplex) {
196+ throw new Error("randnfun3 requires real numeric arrays (" + names[i - 1] + ")");
197+ }
198+ var s = a.shape;
199+ if (!s || s.length !== 2 || s[0] !== ${g.npts} || s[1] !== 1) {
200+ throw new Error(
201+ "randnfun3 evaluates on the grid, so " + names[i - 1] +
202+ " must be ${g.npts}x1, not " + (s ? s.join("x") : "unknown shape")
203+ );
204+ }
205+ }
206+ return [${numericType(g.npts, 1)}];
207+};
208+
209+// Never called: this project executes the IR on WebGPU and emits no C.
210+exports.emit = function () {
211+ throw new Error("randnfun3: no C backend (this runs on WebGPU)");
212+};
213+exports.cBody = function () {
214+ return "";
215+};
216+`;
217+}
218+
219+/** Names the WGSL backend must implement as GPU encodes rather than kernels. */
220+export const EXTERNAL_OPS = new Set([
221+ 'synth', 'analys', 'dtheta', 'dphi', 'dthetac', 'dphic', 'dphig', 'randnfun3',
222+]);
src/mgpu/model.tsadded+469−0View file
@@ -0,0 +1,469 @@
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 type { DerivPlan } from '../sht/deriv.ts';
22+import { lmIndex, type ShtConfig } from '../sht/layout.ts';
23+import { HostBuffers, ModelPlan, type Randnfun3Lambda } from './plan.ts';
24+import { MODE_BUFFER } from './randnfun3.ts';
25+import { inFunction, inFunctionAsync, inModel } from './errors.ts';
26+import { CompiledModel, type Binding } from './compile.ts';
27+
28+export interface ModelParams {
29+ [key: string]: number;
30+}
31+
32+export interface GpuModelOptions {
33+ device: GPUDevice;
34+ sht: ShtPlan;
35+ cfg: ShtConfig;
36+ /** Model source (.m text). */
37+ source: string;
38+ /** Parameter names the .m may take as arguments. */
39+ paramNames: string[];
40+ /** Spectral state names, in order (e.g. ['U', 'V']). */
41+ state: string[];
42+ /** Grid fields to render, in order (e.g. ['u', 'v']). */
43+ view: string[];
44+ /**
45+ * The surface, as the .m may ask for it: `gx`, `gy`, `gz` are the embedding's
46+ * Cartesian coordinates on the grid, and `Gx`, `Gy`, `Gz` the spherical-
47+ * harmonic coefficients they were synthesized from. Omitted for a bare unit
48+ * sphere, where the .m has no geometry to take.
49+ */
50+ geometry?: GeometryBuffers;
51+ /** Computes `dtheta`/`dphi` for the .m's surface Laplace-Beltrami
52+ * correction. Omitted for a bare unit sphere, same as `geometry`. */
53+ deriv?: DerivPlan;
54+ /**
55+ * Iterations of the implicit solve the .m's `for` loop runs. A fixed scalar
56+ * rather than a tunable one: the loop is unrolled into the op sequence, so
57+ * the count is part of what compiles and changing it recompiles.
58+ */
59+ niter?: number;
60+}
61+
62+/** Host-supplied surface fields, in the layout the .m sees them. */
63+export interface GeometryBuffers {
64+ /** Grid coordinates, npts each. */
65+ x: Float32Array;
66+ y: Float32Array;
67+ z: Float32Array;
68+ /** Their spherical-harmonic coefficients, 2 x nlm each. */
69+ X: Float32Array;
70+ Y: Float32Array;
71+ Z: Float32Array;
72+ /** Inverse metric quantities (src/geom/metric.ts), grid space, npts each —
73+ * the Algorithm-4 (12-transform) Laplace-Beltrami path. */
74+ Vtx: Float32Array;
75+ Vty: Float32Array;
76+ Vtz: Float32Array;
77+ Vpx: Float32Array;
78+ Vpy: Float32Array;
79+ Vpz: Float32Array;
80+ /** Flux-form metric weights (src/geom/metric.ts computeFluxMetric), grid
81+ * space, npts each — the six-transform Laplace-Beltrami scheme of
82+ * docs/reduced-transforms.md. */
83+ p1: Float32Array;
84+ p2: Float32Array;
85+ q2: Float32Array;
86+ r: Float32Array;
87+ /** The same weights with the round sphere subtracted (Geometry.dp1/dq2/jinv)
88+ * — the sphere-split form of the flux divergence, which keeps r off the
89+ * round-sphere part of the operator. */
90+ dp1: Float32Array;
91+ dq2: Float32Array;
92+ jinv: Float32Array;
93+ /** Mean-J preconditioner scale (Geometry.Jhat): folded into every
94+ * setParams upload as the 'jhat' uniform, so a .m that takes jhat is
95+ * never left with the zero a missing parameter would default to. An
96+ * explicit jhat in the params wins (jhat: 1 pins the plain round-sphere
97+ * preconditioner, for A/B). */
98+ Jhat: number;
99+}
100+
101+/** Names the .m may take for the grid coordinates and for their coefficients. */
102+export const GEOMETRY_GRID_NAMES = ['gx', 'gy', 'gz'] as const;
103+export const GEOMETRY_SPECTRAL_NAMES = ['Gx', 'Gy', 'Gz'] as const;
104+/** Names the .m may take for the inverse metric quantities (Algorithm 4). */
105+export const METRIC_GRID_NAMES = ['Vtx', 'Vty', 'Vtz', 'Vpx', 'Vpy', 'Vpz'] as const;
106+/** Names the .m may take for the flux-form metric weights (six-transform
107+ * scheme). A model asks for whichever set its loop uses; both are uploaded. */
108+export const FLUX_METRIC_GRID_NAMES = ['p1', 'p2', 'q2', 'r', 'dp1', 'dq2', 'jinv'] as const;
109+
110+/** Laplace-Beltrami eigenvalues l(l+1), duplicated across re/im so the array
111+ * matches the 2 x nlm spectral layout element for element. */
112+export function eigenvalues(cfg: ShtConfig, nlm: number): Float32Array {
113+ const lam = new Float32Array(2 * nlm);
114+ for (let m = 0; m <= cfg.mmax; m++) {
115+ for (let l = m; l <= cfg.lmax; l++) {
116+ const i = lmIndex(cfg.lmax, l, m);
117+ lam[2 * i] = l * (l + 1);
118+ lam[2 * i + 1] = l * (l + 1);
119+ }
120+ }
121+ return lam;
122+}
123+
124+/**
125+ * 1 where l < lmax-2, else 0, duplicated across re/im like `lam`. The
126+ * theta/phi derivative recurrences (src/sht/derivCoeffs.ts) cannot exactly
127+ * represent a derivative at the top two degrees of the band limit, so the
128+ * surface Laplace-Beltrami correction filters them out wherever it
129+ * re-differentiates a field (evolving_surface/notes/algos.tex Sec 6,
130+ * "Miscellaneous implementation details").
131+ */
132+export function filterMask(cfg: ShtConfig, nlm: number): Float32Array {
133+ const filt = new Float32Array(2 * nlm);
134+ for (let m = 0; m <= cfg.mmax; m++) {
135+ for (let l = m; l <= cfg.lmax; l++) {
136+ const i = lmIndex(cfg.lmax, l, m);
137+ const keep = l < cfg.lmax - 2 ? 1 : 0;
138+ filt[2 * i] = keep;
139+ filt[2 * i + 1] = keep;
140+ }
141+ }
142+ return filt;
143+}
144+
145+export class GpuModel {
146+ readonly paramNames: string[];
147+ readonly state: string[];
148+ readonly view: string[];
149+ readonly npts: number;
150+ readonly nlm: number;
151+
152+ #device: GPUDevice;
153+ #host: HostBuffers;
154+ #initPlan: ModelPlan;
155+ #stepPlan: ModelPlan;
156+ /** Current geometry's mean-J scale; 1 with no geometry (the sphere). */
157+ #jhat = 1;
158+ #readback: GPUBuffer;
159+ /** Scratch holding a copy of the whole spectral state; see snapshotState. */
160+ #stash: GPUBuffer;
161+ /** Which function wrote the state most recently; see `read`. */
162+ #lastRan: 'init' | 'step' = 'init';
163+ #stashedRan: 'init' | 'step' = 'init';
164+ #destroyed = false;
165+
166+ private constructor(init: {
167+ device: GPUDevice;
168+ host: HostBuffers;
169+ initPlan: ModelPlan;
170+ stepPlan: ModelPlan;
171+ readback: GPUBuffer;
172+ stash: GPUBuffer;
173+ paramNames: string[];
174+ state: string[];
175+ view: string[];
176+ npts: number;
177+ nlm: number;
178+ }) {
179+ this.#device = init.device;
180+ this.#host = init.host;
181+ this.#initPlan = init.initPlan;
182+ this.#stepPlan = init.stepPlan;
183+ this.#readback = init.readback;
184+ this.#stash = init.stash;
185+ this.paramNames = init.paramNames;
186+ this.state = init.state;
187+ this.view = init.view;
188+ this.npts = init.npts;
189+ this.nlm = init.nlm;
190+ }
191+
192+ static async create(opts: GpuModelOptions): Promise<GpuModel> {
193+ const { device, sht, cfg, source, paramNames, state, view, geometry, deriv } = opts;
194+ const npts = cfg.nlat * cfg.nphi;
195+ const nlm = sht.nlm;
196+ const niter = opts.niter ?? 0;
197+
198+ // What the .m may ask for by parameter name. Spectral state, the
199+ // eigenvalues and the top-mode filter are 2 x nlm; the seeded
200+ // perturbation is a grid field.
201+ const bindings: Record<string, Binding> = {
202+ lam: { kind: 'tensor', shape: [2, nlm] },
203+ filt: { kind: 'tensor', shape: [2, nlm] },
204+ noise: { kind: 'tensor', shape: [npts, 1] },
205+ npts: { kind: 'const', value: npts },
206+ nlm: { kind: 'const', value: nlm },
207+ niter: { kind: 'const', value: niter },
208+ };
209+ if (geometry) {
210+ for (const g of GEOMETRY_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
211+ for (const g of GEOMETRY_SPECTRAL_NAMES) bindings[g] = { kind: 'tensor', shape: [2, nlm] };
212+ for (const g of METRIC_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
213+ for (const g of FLUX_METRIC_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
214+ // Mean-J preconditioner scale (Geometry.Jhat): a uniform, not a const,
215+ // so swapping the surface updates it with no recompile. The session
216+ // folds the current geometry's value into every setParams call.
217+ bindings['jhat'] = { kind: 'param' };
218+ // The wavelength of the seeded random field (src/mgpu/randnfun3.ts).
219+ // A uniform like jhat, not a const: changing it redraws the field
220+ // without recompiling the step.
221+ bindings['lam3'] = { kind: 'param' };
222+ }
223+ for (const s of state) bindings[s] = { kind: 'tensor', shape: [2, nlm] };
224+ for (const p of paramNames) bindings[p] = { kind: 'param' };
225+
226+ // Parsing belongs to the file, not to either function.
227+ const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
228+ // Both functions return the new state first, then the rendered grid fields.
229+ const nargout = state.length + view.length;
230+ const initFn = inFunction('init', () => compiled.specialize('init', nargout));
231+ const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
232+ compiled.finish();
233+
234+ // Only the state outputs feed back into the argument buffers; the grid
235+ // fields are read for display and then overwritten next call.
236+ const feedback = [...state, ...view.map(() => null)];
237+
238+ const host = new HostBuffers(device);
239+ // The host owns the state and the inputs it uploads, whether or not a given
240+ // function happens to take them as arguments — `init` does not read `U`, but
241+ // it writes it, and `step` reads it back.
242+ for (const s of state) host.ensure(s, 2 * nlm);
243+ host.ensure('lam', 2 * nlm);
244+ host.ensure('filt', 2 * nlm);
245+ host.ensure('noise', npts);
246+ if (geometry) {
247+ for (const g of GEOMETRY_GRID_NAMES) host.ensure(g, npts);
248+ for (const g of GEOMETRY_SPECTRAL_NAMES) host.ensure(g, 2 * nlm);
249+ for (const g of METRIC_GRID_NAMES) host.ensure(g, npts);
250+ for (const g of FLUX_METRIC_GRID_NAMES) host.ensure(g, npts);
251+ }
252+
253+ const initPlan = await inFunctionAsync('init', () =>
254+ ModelPlan.create(device, sht, { fn: initFn, feedback }, host, deriv),
255+ );
256+ const stepPlan = await inFunctionAsync('step', () =>
257+ ModelPlan.create(device, sht, { fn: stepFn, feedback }, host, deriv),
258+ );
259+
260+ host.upload('lam', eigenvalues(cfg, nlm));
261+ host.upload('filt', filterMask(cfg, nlm));
262+ if (geometry) {
263+ host.upload('gx', geometry.x);
264+ host.upload('gy', geometry.y);
265+ host.upload('gz', geometry.z);
266+ host.upload('Gx', geometry.X);
267+ host.upload('Gy', geometry.Y);
268+ host.upload('Gz', geometry.Z);
269+ host.upload('Vtx', geometry.Vtx);
270+ host.upload('Vty', geometry.Vty);
271+ host.upload('Vtz', geometry.Vtz);
272+ host.upload('Vpx', geometry.Vpx);
273+ host.upload('Vpy', geometry.Vpy);
274+ host.upload('Vpz', geometry.Vpz);
275+ host.upload('p1', geometry.p1);
276+ host.upload('p2', geometry.p2);
277+ host.upload('q2', geometry.q2);
278+ host.upload('r', geometry.r);
279+ host.upload('dp1', geometry.dp1);
280+ host.upload('dq2', geometry.dq2);
281+ host.upload('jinv', geometry.jinv);
282+ }
283+
284+ const readback = device.createBuffer({
285+ label: 'mgpu-readback',
286+ size: 4 * Math.max(npts, 2 * nlm),
287+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
288+ });
289+ const stash = device.createBuffer({
290+ label: 'mgpu-state-stash',
291+ size: 4 * state.length * 2 * nlm,
292+ usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
293+ });
294+
295+ const gpu = new GpuModel({
296+ device, host, initPlan, stepPlan, readback, stash,
297+ paramNames, state, view, npts, nlm,
298+ });
299+ if (geometry) gpu.#jhat = geometry.Jhat;
300+ return gpu;
301+ }
302+
303+ setParams(params: ModelParams): void {
304+ const merged = { jhat: this.#jhat, ...params };
305+ this.#initPlan.setParams(merged);
306+ this.#stepPlan.setParams(merged);
307+ }
308+
309+ /**
310+ * Write a host-owned value directly — the spectral state, or one of the input
311+ * fields. Lets a test set up an exact initial condition (a single spherical-
312+ * harmonic mode, say) instead of going through `init`.
313+ */
314+ upload(name: string, data: Float32Array): void {
315+ this.#host.upload(name, data);
316+ }
317+
318+ /**
319+ * Swap the surface under a running model. The geometry is data, not code —
320+ * its shape in the bindings depends only on the grid — so changing it is six
321+ * buffer writes and needs no recompile, and the simulation carries straight
322+ * on. Only meaningful if the .m took the geometry as an argument.
323+ */
324+ uploadGeometry(geometry: GeometryBuffers): void {
325+ const fields: [string, Float32Array][] = [
326+ ['gx', geometry.x], ['gy', geometry.y], ['gz', geometry.z],
327+ ['Gx', geometry.X], ['Gy', geometry.Y], ['Gz', geometry.Z],
328+ ['Vtx', geometry.Vtx], ['Vty', geometry.Vty], ['Vtz', geometry.Vtz],
329+ ['Vpx', geometry.Vpx], ['Vpy', geometry.Vpy], ['Vpz', geometry.Vpz],
330+ ['p1', geometry.p1], ['p2', geometry.p2],
331+ ['q2', geometry.q2], ['r', geometry.r],
332+ ['dp1', geometry.dp1], ['dq2', geometry.dq2], ['jinv', geometry.jinv],
333+ ];
334+ for (const [name, data] of fields) {
335+ if (this.#host.get(name)) this.#host.upload(name, data);
336+ }
337+ // The new surface's preconditioner scale takes effect on the next
338+ // setParams (the session re-applies its params after a swap).
339+ this.#jhat = geometry.Jhat;
340+ }
341+
342+ /** The wavelength this model's `init` asked `randnfun3` for, or null if it
343+ * seeds some other way. The session resolves it and draws the modes. */
344+ get randnfun3Lambda(): Randnfun3Lambda | null {
345+ return this.#initPlan.randnfun3Lambda;
346+ }
347+
348+ /**
349+ * Upload the seeded initial data and run `init`.
350+ *
351+ * Both inputs are optional in the sense that a .m uses one or the other:
352+ * `modes` is the random field's coefficient table for a model that calls
353+ * `randnfun3`, `noise` the plain grid field for one that takes `noise`
354+ * directly (the analytic test models inject exact initial conditions that
355+ * way). Only what the plan actually bound is uploaded.
356+ */
357+ async init(noise: Float32Array, modes: Float32Array | null): Promise<void> {
358+ if (this.#host.get('noise')) this.#host.upload('noise', noise);
359+ // Sized to the wavelength, so this may reallocate and rebind.
360+ if (modes) this.#initPlan.uploadRandnfun3Table(this.#host, modes);
361+ // Submitted in pieces: a fine seed wavelength makes the mode sum long
362+ // enough that one submission would stall the browser's compositor.
363+ await this.#initPlan.submitYielding('mgpu-init');
364+ this.#lastRan = 'init';
365+ }
366+
367+ /**
368+ * Copy the spectral state aside, so a batch of steps can run — to be timed —
369+ * and then be undone with restoreState, leaving the simulation exactly where
370+ * it was. Only the state is stashed: the grid view fields keep whatever the
371+ * batch last wrote until a subsequent step recomputes them, so step before
372+ * reading a view after a restore.
373+ */
374+ snapshotState(): void {
375+ this.#stashedRan = this.#lastRan;
376+ this.#copyState('save');
377+ }
378+
379+ restoreState(): void {
380+ this.#copyState('restore');
381+ this.#lastRan = this.#stashedRan;
382+ }
383+
384+ #copyState(dir: 'save' | 'restore'): void {
385+ // A restore can land after a rebuild destroyed the buffers mid-await;
386+ // there is nothing left to protect, so do not submit into destroyed state.
387+ if (this.#destroyed) return;
388+ const enc = this.#device.createCommandEncoder({ label: `mgpu-state-${dir}` });
389+ let offset = 0;
390+ for (const name of this.state) {
391+ const slot = this.#host.get(name);
392+ if (!slot) throw new Error(`state '${name}' has no host buffer`);
393+ const bytes = 4 * slot.count;
394+ if (dir === 'save') {
395+ enc.copyBufferToBuffer(slot.buffer, 0, this.#stash, offset, bytes);
396+ } else {
397+ enc.copyBufferToBuffer(this.#stash, offset, slot.buffer, 0, bytes);
398+ }
399+ offset += bytes;
400+ }
401+ this.#device.queue.submit([enc.finish()]);
402+ }
403+
404+ /**
405+ * Advance `steps` timesteps. Synchronous — this only records commands and
406+ * submits them; nothing is read back and nothing is awaited.
407+ */
408+ step(steps = 1): void {
409+ const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
410+ this.#stepPlan.encodeSteps(enc, steps);
411+ this.#device.queue.submit([enc.finish()]);
412+ this.#lastRan = 'step';
413+ }
414+
415+ /**
416+ * The buffer currently holding a named value. Grid fields like `u` are
417+ * produced by both functions, into separate buffers (only the spectral state
418+ * is shared), so this resolves to whichever function ran most recently —
419+ * which is what makes the first frame show the initial state rather than an
420+ * unwritten buffer.
421+ */
422+ #locate(name: string): { buffer: GPUBuffer; count: number } | null {
423+ const [first, second] =
424+ this.#lastRan === 'init'
425+ ? [this.#initPlan, this.#stepPlan]
426+ : [this.#stepPlan, this.#initPlan];
427+ const buffer = first.buffer(name) ?? second.buffer(name);
428+ const count = first.elementCount(name) ?? second.elementCount(name);
429+ if (!buffer || count === undefined) return null;
430+ return { buffer, count };
431+ }
432+
433+ /** The GPU buffer a named value would be read from right now — for encoding
434+ * further GPU work against it (e.g. a display-grid synthesis of the state)
435+ * without a CPU round trip. */
436+ valueBuffer(name: string): GPUBuffer | null {
437+ return this.#locate(name)?.buffer ?? null;
438+ }
439+
440+ /** Read a named value back to the CPU. The only await in the whole loop. */
441+ async read(name: string): Promise<Float32Array> {
442+ const located = this.#locate(name);
443+ if (!located) {
444+ throw new Error(`read: the model has no value named '${name}'`);
445+ }
446+ const { buffer, count } = located;
447+ const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
448+ enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
449+ this.#device.queue.submit([enc.finish()]);
450+ await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
451+ const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
452+ this.#readback.unmap();
453+ return out;
454+ }
455+
456+ /** What the .m compiled to, for display. */
457+ describe(): { init: string[]; step: string[] } {
458+ return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
459+ }
460+
461+ destroy(): void {
462+ this.#destroyed = true;
463+ this.#initPlan.destroy();
464+ this.#stepPlan.destroy();
465+ this.#host.destroy();
466+ this.#readback.destroy();
467+ this.#stash.destroy();
468+ }
469+}
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+366−0View file
@@ -0,0 +1,366 @@
1+/**
2+ * The numbl compiler surface this project depends on.
3+ *
4+ * We reach past numbl's published entry points into its internals — the JIT
5+ * side (parser, lowerer, IR, inline pass) that compiles the models, and the
6+ * interpreter side (executeCode, runtime values) that evaluates the
7+ * geometries — which its package `exports` map does not expose. Those imports
8+ * resolve through the `numbl-src` alias in vite.config.ts; these declarations
9+ * are what TypeScript checks against.
10+ *
11+ * Declaring the surface here rather than type-checking numbl's sources
12+ * directly keeps this project's compiler settings independent of numbl's, and
13+ * pins the exact contract we rely on. If numbl changes one of these shapes,
14+ * the build breaks here with a clear diff rather than deep inside its tree.
15+ *
16+ * Only the nodes the WGSL backend actually walks are spelled out; every other
17+ * IR kind is collapsed into a catch-all so that unhandled constructs are
18+ * rejected with a message instead of being silently mis-compiled.
19+ */
20+
21+declare module 'numbl-src/numbl-core/jit/lowering/types.ts' {
22+ export type Sign =
23+ | 'positive' | 'nonneg' | 'negative' | 'nonpositive'
24+ | 'zero' | 'nonzero' | 'unknown';
25+
26+ export type DimInfo = { kind: 'exact'; value: number } | { kind: 'unknown' };
27+
28+ export type NumericExact =
29+ | number
30+ | Float64Array
31+ | { re: number; im: number }
32+ | { re: Float64Array; im: Float64Array };
33+
34+ export interface NumericType {
35+ kind: 'Numeric';
36+ elem: 'double' | 'logical' | 'char' | string;
37+ isComplex: boolean;
38+ dims: DimInfo[];
39+ /** Present iff every dim is exact. */
40+ shape?: number[];
41+ sign: Sign;
42+ exact?: NumericExact;
43+ }
44+
45+ /** Everything the WGSL backend rejects. */
46+ export interface NonNumericType {
47+ kind: 'Void' | 'Unknown' | 'String' | 'Handle' | 'Struct' | 'Class' | 'Cell';
48+ }
49+
50+ export type Type = NumericType | NonNumericType;
51+
52+ export function isMultiElement(t: NumericType): boolean;
53+ export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
54+ export function scalarDouble(sign?: Sign, exact?: number): NumericType;
55+}
56+
57+declare module 'numbl-src/numbl-core/jit/lowering/ir.ts' {
58+ import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
59+
60+ export interface Span {
61+ file: string;
62+ start: number;
63+ end: number;
64+ }
65+
66+ export interface NumLit {
67+ kind: 'NumLit';
68+ value: number;
69+ ty: Type;
70+ span: Span;
71+ }
72+ export interface Var {
73+ kind: 'Var';
74+ name: string;
75+ cName: string;
76+ ty: Type;
77+ span: Span;
78+ }
79+ export interface Binary {
80+ kind: 'Binary';
81+ builtin: string;
82+ left: IRExpr;
83+ right: IRExpr;
84+ ty: Type;
85+ span: Span;
86+ }
87+ export interface Unary {
88+ kind: 'Unary';
89+ builtin: string;
90+ operand: IRExpr;
91+ ty: Type;
92+ span: Span;
93+ }
94+ export interface Call {
95+ kind: 'Call';
96+ cName: string;
97+ name: string;
98+ args: IRExpr[];
99+ ty: Type;
100+ span: Span;
101+ }
102+ /** Any other IR expression kind — rejected by the WGSL emitter. */
103+ export interface OtherExpr {
104+ kind:
105+ | 'ImagLit' | 'StringLit' | 'TensorBuild' | 'TensorConcat' | 'CellLit'
106+ | 'CellEmpty' | 'CellIndexLoad' | 'HandleLit' | 'HandleCaptureLoad'
107+ | 'StructLit' | 'MemberLoad' | 'IndexLoad' | 'IndexSlice' | 'EndRef'
108+ | 'MakeRange';
109+ ty: Type;
110+ span: Span;
111+ }
112+
113+ export type IRExpr = NumLit | Var | Binary | Unary | Call | OtherExpr;
114+
115+ export interface Assign {
116+ kind: 'Assign';
117+ name: string;
118+ cName: string;
119+ ty: Type;
120+ expr: IRExpr;
121+ span: Span;
122+ }
123+ /**
124+ * A counted loop. The planner unrolls it, so only the fields that decide
125+ * the trip count and the loop variable's value are spelled out. `step` is
126+ * already a literal number in the IR — numbl rejects a non-literal step
127+ * during lowering — while `start` and `end` are expressions that must carry
128+ * an exact value for the planner to accept the loop.
129+ */
130+ export interface For {
131+ kind: 'For';
132+ /** Loop variable, as written in the .m. */
133+ varName: string;
134+ /** Loop variable's cName, the key the planner binds its value under. */
135+ cVar: string;
136+ start: IRExpr;
137+ step: number;
138+ end: IRExpr;
139+ body: IRStmt[];
140+ span: Span;
141+ }
142+ /**
143+ * Multi-output call statement: `[a, b] = f(x, y)`. For `isBuiltin: true`
144+ * the builtin's `transfer(argTypes, nargout)` typed the slots during
145+ * lowering; args arrive ANF'd. The planner accepts this only for the
146+ * batched transforms (`synth`/`analys`), where output k is the transform
147+ * of argument k.
148+ */
149+ export interface MultiAssignCall {
150+ kind: 'MultiAssignCall';
151+ cName: string;
152+ name: string;
153+ isBuiltin?: boolean;
154+ args: IRExpr[];
155+ outputs: ReadonlyArray<{
156+ ty: Type;
157+ binding: { name: string; cName: string } | null;
158+ }>;
159+ span: Span;
160+ }
161+
162+ /** Any other IR statement kind — rejected by the planner. */
163+ export interface OtherStmt {
164+ kind:
165+ | 'ExprStmt' | 'If' | 'While' | 'ReturnFromFunction' | 'Break'
166+ | 'Continue' | 'TypeComment' | 'MemberStore'
167+ | 'IndexStore' | 'IndexSliceStore' | 'CellIndexStore';
168+ span: Span;
169+ }
170+
171+ export type IRStmt = Assign | For | MultiAssignCall | OtherStmt;
172+
173+ export interface IRFunc {
174+ name: string;
175+ cName: string;
176+ /** Parameter source names. */
177+ params: string[];
178+ /** Parameter cNames, parallel to `params`. */
179+ cParams: string[];
180+ paramTypes: Type[];
181+ /** Output source names. */
182+ outputs: string[];
183+ /** Output cNames, parallel to `outputs`. */
184+ cOutputs: string[];
185+ outputTypes: Type[];
186+ body: IRStmt[];
187+ span: Span;
188+ }
189+
190+ export interface IRProgram {
191+ topLevelStmts: IRStmt[];
192+ functions: Map<string, IRFunc>;
193+ }
194+}
195+
196+declare module 'numbl-src/numbl-core/parser/index.ts' {
197+ export interface ParseSpan {
198+ start: number;
199+ end: number;
200+ }
201+
202+ /** The one parse-tree node this project inspects (src/geom/geometry.ts,
203+ * finding `shape` and its argument names). */
204+ export interface FunctionStmt {
205+ type: 'Function';
206+ name: string;
207+ params: string[];
208+ outputs: string[];
209+ span: ParseSpan;
210+ }
211+
212+ /** Any other statement in a file's body — opaque to this project. Its
213+ * `type` is some other literal; narrowing to FunctionStmt goes through an
214+ * explicit type guard rather than the discriminant. */
215+ export interface OtherParseStmt {
216+ type: string;
217+ span: ParseSpan;
218+ }
219+
220+ export type Stmt = FunctionStmt | OtherParseStmt;
221+
222+ export interface AbstractSyntaxTree {
223+ body: Stmt[];
224+ }
225+ export function parseMFile(input: string, fileName?: string): AbstractSyntaxTree;
226+ export class SyntaxError extends Error {}
227+}
228+
229+declare module 'numbl-src/numbl-core/runtime/types.ts' {
230+ /** A numeric array: f64 data in column-major order, with its shape. */
231+ export class RuntimeTensor {
232+ readonly kind: 'tensor';
233+ data: Float64Array;
234+ /** Present iff the value is complex. */
235+ imag: Float64Array | undefined;
236+ shape: number[];
237+ constructor(data: Float64Array, shape: number[], imag?: Float64Array);
238+ }
239+
240+ /** Every other value kind the interpreter can hold, collapsed. */
241+ export interface OtherRuntimeValue {
242+ readonly kind: string;
243+ }
244+
245+ export type RuntimeValue =
246+ | number
247+ | boolean
248+ | string
249+ | RuntimeTensor
250+ | OtherRuntimeValue;
251+
252+ export function isRuntimeTensor(value: RuntimeValue): value is RuntimeTensor;
253+}
254+
255+declare module 'numbl-src/numbl-core/executeCode.ts' {
256+ import type { RuntimeValue } from 'numbl-src/numbl-core/runtime/types.ts';
257+
258+ export interface ExecOptions {
259+ /** Variables pre-bound in the script's workspace before it runs. */
260+ initialVariableValues?: Record<string, RuntimeValue>;
261+ displayResults?: boolean;
262+ onOutput?: (text: string) => void;
263+ /** null opts out of scanning a working directory for .m files. */
264+ implicitCwdPath?: string | null;
265+ }
266+
267+ export interface ExecWorkspaceFile {
268+ name: string;
269+ source: string;
270+ }
271+
272+ export interface ExecResult {
273+ output: string[];
274+ /** The script's workspace after it ran. */
275+ variableValues: Record<string, RuntimeValue>;
276+ }
277+
278+ /** Run a script through numbl's interpreter (with its JS-JIT), CPU-side. */
279+ export function executeCode(
280+ source: string,
281+ options?: ExecOptions,
282+ workspaceFiles?: ExecWorkspaceFile[],
283+ mainFileName?: string,
284+ ): ExecResult;
285+}
286+
287+declare module 'numbl-src/numbl-core/jit/index.ts' {
288+ import type { AbstractSyntaxTree } from 'numbl-src/numbl-core/parser/index.ts';
289+ import type { IRProgram, IRFunc, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
290+ import type { Type, NumericType, Sign } from 'numbl-src/numbl-core/jit/lowering/types.ts';
291+
292+ export interface WorkspaceFile {
293+ name: string;
294+ source: string;
295+ ast?: AbstractSyntaxTree;
296+ }
297+
298+ export class Workspace {
299+ constructor(mainFile: string, searchPaths?: ReadonlyArray<string>);
300+ addFile(file: WorkspaceFile): void;
301+ finalize(): void;
302+ }
303+
304+ export interface EnvEntry {
305+ cName: string;
306+ ty: Type;
307+ maybeUnassigned?: boolean;
308+ }
309+
310+ export class Lowerer {
311+ constructor(workspace: Workspace);
312+ /** Pre-bindable variable scope: seed host-provided values here. */
313+ env: Map<string, EnvEntry>;
314+ specializations: Map<string, IRFunc>;
315+ lowerProgram(ast: AbstractSyntaxTree): IRProgram;
316+ }
317+
318+ /** Thrown for MATLAB the JIT pipeline cannot lower; carries a source span. */
319+ export class UnsupportedConstruct extends Error {
320+ span?: Span;
321+ }
322+ export class JitTypeError extends Error {
323+ span?: Span;
324+ }
325+
326+ export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
327+ export function scalarDouble(sign?: Sign, exact?: number): NumericType;
328+ export function isMultiElement(t: NumericType): boolean;
329+}
330+
331+declare module 'numbl-src/numbl-core/jit/lowering/specialize.ts' {
332+ import type { Lowerer } from 'numbl-src/numbl-core/jit/index.ts';
333+ import type { IRFunc, IRExpr, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
334+ import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
335+
336+ /**
337+ * Lower one user function for a concrete argument-type signature. Called with
338+ * a `Lowerer` as `this` (numbl's own JIT does the same), so specializations
339+ * accumulate in `lowerer.specializations`.
340+ */
341+ export function specializeUserFunction(
342+ this: Lowerer,
343+ decl: unknown,
344+ argTypes: Type[],
345+ specSource?: string,
346+ definingFile?: string,
347+ preSeedOutput?: { name: string; ty: Type; initExpr: IRExpr },
348+ nargout?: number,
349+ callSiteSpan?: Span,
350+ ): IRFunc;
351+}
352+
353+declare module 'numbl-src/numbl-core/jit/codegen/inlinePass.ts' {
354+ import type { IRProgram } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
355+ /** Folds single-use ANF temps into their consumer, in place. */
356+ export function inlinePass(prog: IRProgram): void;
357+}
358+
359+declare module 'numbl-src/numbl-core/jit/builtins/index.ts' {
360+ export interface Builtin {
361+ name: string;
362+ /** Safe to evaluate one output element from one input element per slot. */
363+ elementwise?: boolean;
364+ }
365+ export function getBuiltin(name: string): Builtin | undefined;
366+}
src/mgpu/plan.tsadded+1166−0View file
@@ -0,0 +1,1166 @@
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 {
13+ Assign,
14+ For,
15+ IRExpr,
16+ IRStmt,
17+ MultiAssignCall,
18+} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
19+import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
20+import { ShtPlan, type ShtBinding, type ShtBatchBinding, type ShtDphigBinding } from '../sht/sht.ts';
21+import { DerivPlan, type DerivBinding } from '../sht/deriv.ts';
22+import type { CompiledFunction } from './compile.ts';
23+import { EXTERNAL_OPS } from './externals.ts';
24+import {
25+ MODE_BUFFER,
26+ INITIAL_MODES,
27+ modeTableLength,
28+ randnfun3Chunks,
29+ randnfun3WGSL,
30+} from './randnfun3.ts';
31+import {
32+ buildKernel,
33+ UnsupportedOnGpu,
34+ WORKGROUP_SIZE,
35+ type KernelInputs,
36+} from './wgsl.ts';
37+
38+const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
39+const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
40+const numel = (t: NumericType): number => (t.shape ?? []).reduce((a, b) => a * b, 1);
41+
42+/**
43+ * The compile-time value of a scalar expression, if it has one. A literal
44+ * carries its own; a variable carries one when it was bound to a `const` (the
45+ * host's fixed scalars) or computed from constants, because numbl propagates
46+ * `exact` through the type lattice.
47+ */
48+const exactValue = (e: IRExpr): number | undefined => {
49+ if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
50+ return e.kind === 'NumLit' ? e.value : undefined;
51+};
52+
53+/** Cap on the iterations a `for` may unroll to. Each one is real GPU work —
54+ * its own pipelines at compile time and its own dispatches per step — so a
55+ * runaway bound should be a clear error rather than a hang. */
56+const MAX_UNROLL = 64;
57+
58+interface Slot {
59+ buffer: GPUBuffer;
60+ count: number;
61+}
62+
63+const makeBuffer = (device: GPUDevice, label: string, count: number): GPUBuffer =>
64+ device.createBuffer({
65+ label,
66+ size: 4 * count,
67+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
68+ });
69+
70+/**
71+ * Buffers for host-bound variables, shared across plans.
72+ *
73+ * A model is two programs — `init` and `step` — compiled separately but
74+ * operating on the same state. `U` in the step must be the very buffer `init`
75+ * wrote, so the buffers for host bindings live here rather than inside either
76+ * plan.
77+ */
78+export class HostBuffers {
79+ #device: GPUDevice;
80+ #slots = new Map<string, Slot>();
81+
82+ constructor(device: GPUDevice) {
83+ this.#device = device;
84+ }
85+
86+ ensure(name: string, count: number): Slot {
87+ const existing = this.#slots.get(name);
88+ if (existing) {
89+ if (existing.count !== count) {
90+ throw new UnsupportedOnGpu(
91+ `'${name}' is ${existing.count} elements in one program and ` +
92+ `${count} in another`,
93+ );
94+ }
95+ return existing;
96+ }
97+ const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
98+ this.#slots.set(name, slot);
99+ return slot;
100+ }
101+
102+ get(name: string): Slot | undefined {
103+ return this.#slots.get(name);
104+ }
105+
106+ /**
107+ * Replace a slot's buffer with a larger one. Only for buffers whose size is
108+ * not fixed by the grid — the randnfun3 mode table, which grows with the
109+ * wavelength asked for. The caller must rebuild any bind group holding the
110+ * old buffer; it is destroyed here.
111+ */
112+ resize(name: string, count: number): Slot {
113+ const existing = this.#slots.get(name);
114+ if (!existing) throw new Error(`resize: no buffer named '${name}'`);
115+ if (count <= existing.count) return existing;
116+ existing.buffer.destroy();
117+ const slot = { buffer: makeBuffer(this.#device, `mgpu-${name}`, count), count };
118+ this.#slots.set(name, slot);
119+ return slot;
120+ }
121+
122+ /** Upload into the front of a slot, leaving any tail as it was. For a
123+ * variable-length payload in a buffer sized to its high-water mark. */
124+ uploadInto(name: string, data: Float32Array): void {
125+ const slot = this.#slots.get(name);
126+ if (!slot) throw new Error(`uploadInto: no buffer named '${name}'`);
127+ if (data.length > slot.count) {
128+ throw new Error(
129+ `uploadInto '${name}': ${data.length} elements into a ${slot.count}-element buffer`,
130+ );
131+ }
132+ this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
133+ }
134+
135+ /** Upload initial data for a host binding. */
136+ upload(name: string, data: Float32Array): void {
137+ const slot = this.#slots.get(name);
138+ if (!slot) throw new Error(`upload: no buffer named '${name}'`);
139+ if (data.length !== slot.count) {
140+ throw new Error(
141+ `upload '${name}': expected ${slot.count} elements, got ${data.length}`,
142+ );
143+ }
144+ this.#device.queue.writeBuffer(slot.buffer, 0, data as Float32Array<ArrayBuffer>);
145+ }
146+
147+ destroy(): void {
148+ for (const s of this.#slots.values()) s.buffer.destroy();
149+ this.#slots.clear();
150+ }
151+}
152+
153+type Op =
154+ | {
155+ kind: 'kernel';
156+ pipeline: GPUComputePipeline;
157+ bindGroup: GPUBindGroup;
158+ count: number;
159+ label: string;
160+ /** Set when the kernel had to write to scratch because its output
161+ * aliases one of its inputs; copied back after the dispatch. */
162+ copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
163+ /** End the submission here when run through `submitYielding`, so the
164+ * GPU is handed back between chunks of a long seed. */
165+ yieldAfter?: boolean;
166+ }
167+ | { kind: 'synth' | 'analys'; binding: ShtBinding; label: string }
168+ | { kind: 'synth-batch' | 'analys-batch'; binding: ShtBatchBinding; labels: string[] }
169+ | { kind: 'dtheta' | 'dphi'; binding: DerivBinding; label: string }
170+ | { kind: 'dthetac' | 'dphic'; bindGroup: GPUBindGroup; label: string }
171+ | { kind: 'dphig'; binding: ShtDphigBinding; label: string }
172+ | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string };
173+
174+/**
175+ * A transform op as planned, before bindings exist: `in`/`out` are the
176+ * caller-side buffers (spectral in / grid out for synth, the reverse for
177+ * analys). Kept unbound until every statement is planned so that adjacent
178+ * independent transforms of the same kind can be grouped into one batched
179+ * dispatch (ShtPlan.createSynthBatchBinding) — the Legendre recurrence is
180+ * the expensive shared part, and a batch walks it once for all lanes.
181+ */
182+interface PendingSht {
183+ pending: true;
184+ kind: 'synth' | 'analys';
185+ in: GPUBuffer;
186+ out: GPUBuffer;
187+ label: string;
188+}
189+
190+type Planned = Op | PendingSht;
191+
192+const isPending = (op: Planned): op is PendingSht => 'pending' in op;
193+
194+/**
195+ * Group maximal runs of adjacent same-kind transforms into batches of the
196+ * widest compiled lane count, and create all bindings. Only literal
197+ * adjacency in the op sequence is batched — no reordering — so the models
198+ * are written to keep batchable transforms consecutive (see the solve loops
199+ * in models/*.m). Batching changes dispatch shape only: per-lane arithmetic
200+ * is identical to the scalar kernels', so results do not depend on batchK.
201+ */
202+function materializeTransforms(planned: Planned[], sht: ShtPlan): Op[] {
203+ /** Lanes must not collide: distinct outputs, and no lane reading another's
204+ * output (repeated read-only inputs would be harmless, but WebGPU also
205+ * forbids aliasing a writable binding, so outputs are the hard rule). */
206+ const disjoint = (members: PendingSht[]): boolean => {
207+ const outs = new Set<GPUBuffer>();
208+ for (const m of members) {
209+ if (outs.has(m.out)) return false;
210+ outs.add(m.out);
211+ }
212+ return members.every((m) => !outs.has(m.in));
213+ };
214+ const bind = (m: PendingSht): Op =>
215+ m.kind === 'synth'
216+ ? { kind: 'synth', binding: sht.createSynthBinding(m.in, m.out), label: m.label }
217+ : { kind: 'analys', binding: sht.createAnalysBinding(m.in, m.out), label: m.label };
218+ const bindBatch = (members: PendingSht[]): Op =>
219+ members[0].kind === 'synth'
220+ ? {
221+ kind: 'synth-batch',
222+ binding: sht.createSynthBatchBinding(
223+ members.map((m) => ({ qlmIn: m.in, spatOut: m.out })),
224+ ),
225+ labels: members.map((m) => m.label),
226+ }
227+ : {
228+ kind: 'analys-batch',
229+ binding: sht.createAnalysBatchBinding(
230+ members.map((m) => ({ spatIn: m.in, qlmOut: m.out })),
231+ ),
232+ labels: members.map((m) => m.label),
233+ };
234+
235+ const out: Op[] = [];
236+ let i = 0;
237+ while (i < planned.length) {
238+ const op = planned[i];
239+ if (!isPending(op)) {
240+ out.push(op);
241+ i++;
242+ continue;
243+ }
244+ let j = i;
245+ while (j < planned.length) {
246+ const p = planned[j];
247+ if (!isPending(p) || p.kind !== op.kind) break;
248+ j++;
249+ }
250+ const run = planned.slice(i, j) as PendingSht[];
251+ let s = 0;
252+ while (s < run.length) {
253+ let take = 1;
254+ for (const K of [4, 2]) {
255+ if (K > sht.batchK || s + K > run.length) continue;
256+ if (disjoint(run.slice(s, s + K))) {
257+ take = K;
258+ break;
259+ }
260+ }
261+ out.push(take === 1 ? bind(run[s]) : bindBatch(run.slice(s, s + take)));
262+ s += take;
263+ }
264+ i = j;
265+ }
266+ return out;
267+}
268+
269+export interface PlanSpec {
270+ /** The specialized function this plan executes. */
271+ fn: CompiledFunction;
272+ /** Output index -> host binding name to copy the result into after the run,
273+ * so the next call reads it (the new spectral state feeds the old). */
274+ feedback: (string | null)[];
275+}
276+
277+/**
278+ * Bind group layout for a kernel: the output at 0, `inputs` read-only storage
279+ * buffers after it, then the params buffer.
280+ *
281+ * Declared explicitly rather than with `layout: 'auto'`, because an auto layout
282+ * only contains the bindings the shader actually references — so a kernel that
283+ * happens to use no parameters (`uuv = u .* u .* v`) would drop the params
284+ * binding and no longer match the bind group. An explicit layout may carry
285+ * bindings the shader ignores.
286+ */
287+function kernelLayout(device: GPUDevice, inputs: number): GPUBindGroupLayout {
288+ const readOnly = (binding: number): GPUBindGroupLayoutEntry => ({
289+ binding,
290+ visibility: GPUShaderStage.COMPUTE,
291+ buffer: { type: 'read-only-storage' },
292+ });
293+ return device.createBindGroupLayout({
294+ entries: [
295+ {
296+ binding: 0,
297+ visibility: GPUShaderStage.COMPUTE,
298+ buffer: { type: 'storage' },
299+ },
300+ ...Array.from({ length: inputs }, (_, i) => readOnly(i + 1)),
301+ readOnly(inputs + 1),
302+ ],
303+ });
304+}
305+
306+async function makePipeline(
307+ device: GPUDevice,
308+ code: string,
309+ label: string,
310+ bindGroupLayout: GPUBindGroupLayout,
311+): Promise<GPUComputePipeline> {
312+ device.pushErrorScope('validation');
313+ const module = device.createShaderModule({ code, label });
314+ const info = await module.getCompilationInfo();
315+ const errors = info.messages.filter((m) => m.type === 'error');
316+ if (errors.length) {
317+ throw new UnsupportedOnGpu(
318+ `generated WGSL failed to compile for '${label}':\n` +
319+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
320+ `\n--- shader ---\n${code}`,
321+ );
322+ }
323+ const pipeline = await device.createComputePipelineAsync({
324+ layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
325+ compute: { module, entryPoint: 'main' },
326+ label,
327+ });
328+ const err = await device.popErrorScope();
329+ if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
330+ return pipeline;
331+}
332+
333+/** A compiled .m step, ready to run on the GPU. */
334+export class ModelPlan {
335+ /** Scalar parameter names, in the order the params buffer expects them. */
336+ readonly paramNames: string[];
337+ /** The wavelength this plan's `randnfun3` call asked for, or null if it
338+ * makes none. The host draws the coefficient table from it. */
339+ readonly randnfun3Lambda: Randnfun3Lambda | null;
340+
341+ #device: GPUDevice;
342+ #sht: ShtPlan;
343+ #deriv?: DerivPlan;
344+ #ops: Op[];
345+ #owned: GPUBuffer[];
346+ #paramBuf: GPUBuffer;
347+ #paramData: Float32Array;
348+ #rebindRandnfun3: ((table: GPUBuffer) => void) | null;
349+ /** Public name -> buffer, for uploading initial state and reading results. */
350+ #byName: Map<string, Slot>;
351+
352+ private constructor(init: {
353+ device: GPUDevice;
354+ sht: ShtPlan;
355+ deriv?: DerivPlan;
356+ ops: Op[];
357+ byName: Map<string, Slot>;
358+ owned: GPUBuffer[];
359+ paramBuf: GPUBuffer;
360+ paramData: Float32Array;
361+ paramNames: string[];
362+ randnfun3Lambda: Randnfun3Lambda | null;
363+ rebindRandnfun3: ((table: GPUBuffer) => void) | null;
364+ }) {
365+ this.#device = init.device;
366+ this.#sht = init.sht;
367+ this.#deriv = init.deriv;
368+ this.#ops = init.ops;
369+ this.#byName = init.byName;
370+ this.#owned = init.owned;
371+ this.#paramBuf = init.paramBuf;
372+ this.#paramData = init.paramData;
373+ this.paramNames = init.paramNames;
374+ this.randnfun3Lambda = init.randnfun3Lambda;
375+ this.#rebindRandnfun3 = init.rebindRandnfun3;
376+ }
377+
378+ /**
379+ * Point the randnfun3 dispatch at a mode table big enough for `data`,
380+ * growing the buffer if this wavelength needs more modes than the last one,
381+ * and upload it.
382+ */
383+ uploadRandnfun3Table(host: HostBuffers, data: Float32Array): void {
384+ const slot = host.get(MODE_BUFFER);
385+ if (!slot || !this.#rebindRandnfun3) return;
386+ if (data.length > slot.count) {
387+ const max = this.#device.limits.maxStorageBufferBindingSize;
388+ if (4 * data.length > max) {
389+ throw new Error(
390+ `randnfun3: this wavelength needs a ${(4 * data.length / 1e6).toFixed(0)} MB ` +
391+ `mode table, past this device's ${(max / 1e6).toFixed(0)} MB limit ` +
392+ `on a single buffer. Use a larger lambda.`,
393+ );
394+ }
395+ this.#rebindRandnfun3(host.resize(MODE_BUFFER, data.length).buffer);
396+ }
397+ host.uploadInto(MODE_BUFFER, data);
398+ }
399+
400+ static async create(
401+ device: GPUDevice,
402+ sht: ShtPlan,
403+ spec: PlanSpec,
404+ host: HostBuffers,
405+ /** Computes dtheta/dphi — only needed if the .m calls them. */
406+ deriv?: DerivPlan,
407+ ): Promise<ModelPlan> {
408+ const { fn } = spec;
409+
410+ const slots = new Map<string, Slot>();
411+ const byName = new Map<string, Slot>();
412+ const owned: GPUBuffer[] = [];
413+ /** Scalars the .m computes from its parameters, by cName. */
414+ const derivedScalars = new Map<string, { name: string; expr: IRExpr }>();
415+
416+ const alloc = (label: string, count: number): Slot => {
417+ const buffer = makeBuffer(device, label, count);
418+ owned.push(buffer);
419+ return { buffer, count };
420+ };
421+
422+ // Arguments, bound by what the function's signature declares. Array
423+ // arguments come from the shared pool, so a value one function returns is
424+ // the same buffer the next one reads. Scalar parameters share one small
425+ // storage buffer, in signature order.
426+ const paramNames: string[] = [];
427+ const paramSlots = new Map<string, number>();
428+ for (const p of fn.params) {
429+ if (p.binding.kind === 'tensor') {
430+ const count = p.binding.shape.reduce((x, y) => x * y, 1);
431+ const slot = host.ensure(p.name, count);
432+ slots.set(p.cName, slot);
433+ byName.set(p.name, slot);
434+ } else if (p.binding.kind === 'param') {
435+ paramSlots.set(p.cName, paramNames.length);
436+ paramNames.push(p.name);
437+ }
438+ // `const` arguments are exact in the IR and fold into the kernels.
439+ }
440+ const paramData = new Float32Array(Math.max(1, paramNames.length));
441+ const paramBuf = device.createBuffer({
442+ label: 'mgpu-params',
443+ size: 4 * paramData.length,
444+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
445+ });
446+
447+ /** Set when the .m calls `randnfun3`: which wavelength it asked for, so
448+ * the host draws the coefficient table the kernel reads from exactly
449+ * that value (src/mgpu/randnfun3.ts). */
450+ let randnfun3Lambda: Randnfun3Lambda | null = null;
451+ /** Rebuilds the randnfun3 dispatch's bind group after the mode table is
452+ * reallocated for a finer wavelength. */
453+ let rebindRandnfun3: ((table: GPUBuffer) => void) | null = null;
454+
455+ const planned: Planned[] = [];
456+ for (const stmt of fn.body) {
457+ await planStatement(stmt);
458+ }
459+
460+ // Feed declared outputs back into the argument buffers they replace.
461+ fn.outputs.forEach((out, i) => {
462+ const to = spec.feedback[i];
463+ if (!to) return;
464+ const src = slots.get(out.cName);
465+ const dst = host.get(to);
466+ if (!src) {
467+ throw new UnsupportedOnGpu(
468+ `'${fn.name}' declares the output '${out.name}' but never assigns it`,
469+ );
470+ }
471+ if (!dst) throw new UnsupportedOnGpu(`'${to}' is not a host binding`);
472+ if (src.count !== dst.count) {
473+ throw new UnsupportedOnGpu(
474+ `'${out.name}' (${src.count} elements) cannot feed ` +
475+ `'${to}' (${dst.count})`,
476+ );
477+ }
478+ planned.push({
479+ kind: 'copy',
480+ from: src.buffer,
481+ to: dst.buffer,
482+ bytes: 4 * src.count,
483+ label: `${out.name} -> ${to}`,
484+ });
485+ });
486+
487+ // Group adjacent independent transforms into batched dispatches and
488+ // create every binding.
489+ const ops = materializeTransforms(planned, sht);
490+
491+ return new ModelPlan({
492+ device, sht, deriv, ops, byName, owned, paramBuf, paramData, paramNames,
493+ randnfun3Lambda, rebindRandnfun3,
494+ });
495+
496+ async function planStatement(stmt: IRStmt): Promise<void> {
497+ if (stmt.kind === 'ReturnFromFunction') return; // nothing follows it
498+ if (stmt.kind === 'For') return planFor(stmt);
499+ if (stmt.kind === 'MultiAssignCall') return planMultiTransform(stmt);
500+ if (stmt.kind !== 'Assign') {
501+ throw new UnsupportedOnGpu(
502+ `a model function body may only contain assignments ` +
503+ `(found '${stmt.kind}')`,
504+ stmt.span,
505+ );
506+ }
507+ if (!isNumeric(stmt.ty)) {
508+ throw new UnsupportedOnGpu(
509+ `'${stmt.name}' is not a numeric value`,
510+ stmt.span,
511+ );
512+ }
513+ if (!isTensor(stmt.ty)) {
514+ // A scalar the model derives from its parameters (`us = a + b`). It
515+ // gets no buffer and no dispatch: the kernels that read it bind it as
516+ // a `let` in their prologue.
517+ derivedScalars.set(stmt.cName, { name: stmt.name, expr: stmt.expr });
518+ return;
519+ }
520+ const count = numel(stmt.ty);
521+
522+ // Reuse the destination buffer across steps: the same cName always maps
523+ // to the same buffer, so a step allocates nothing.
524+ let dest = slots.get(stmt.cName);
525+ if (!dest) {
526+ dest = alloc(`mgpu-${stmt.name}`, count);
527+ slots.set(stmt.cName, dest);
528+ } else if (dest.count !== count) {
529+ throw new UnsupportedOnGpu(
530+ `'${stmt.name}' changes size between assignments`,
531+ stmt.span,
532+ );
533+ }
534+ byName.set(stmt.name, dest);
535+
536+ const ext = externalCall(stmt);
537+ if (ext) {
538+ if (ext.name === 'randnfun3') {
539+ await planRandnfun3(stmt, ext.args, dest);
540+ return;
541+ }
542+ const arg = ext.args[0] as IRExpr & { kind: 'Var' };
543+ const argSlot = slots.get(arg.cName);
544+ if (!argSlot) {
545+ throw new UnsupportedOnGpu(
546+ `'${ext.name}' reads '${arg.name}', which has no buffer`,
547+ stmt.span,
548+ );
549+ }
550+ const label = `${stmt.name} = ${ext.name}(${arg.name})`;
551+ if (ext.name === 'dphig') {
552+ // Grid -> grid, staged through the plan's fm scratch; safe even
553+ // in place, so no aliasing guard is needed.
554+ planned.push({
555+ kind: 'dphig',
556+ binding: sht.createDphigBinding(argSlot.buffer, dest.buffer),
557+ label,
558+ });
559+ return;
560+ }
561+ if (ext.name === 'synth' || ext.name === 'analys') {
562+ // Left unbound until materializeTransforms has grouped adjacent
563+ // independent transforms into batched dispatches.
564+ planned.push({
565+ pending: true,
566+ kind: ext.name,
567+ in: argSlot.buffer,
568+ out: dest.buffer,
569+ label,
570+ });
571+ } else if (
572+ ext.name === 'dtheta' || ext.name === 'dphi' ||
573+ ext.name === 'dthetac' || ext.name === 'dphic'
574+ ) {
575+ if (!deriv) {
576+ throw new UnsupportedOnGpu(
577+ `'${ext.name}' needs the surface's derivative transforms, ` +
578+ `which this plan was not given`,
579+ stmt.span,
580+ );
581+ }
582+ if (ext.name === 'dthetac' || ext.name === 'dphic') {
583+ // Coefficient-space shuffles read at l+-1 (dthetac) or in place
584+ // (dphic) and cannot alias their output: WebGPU forbids one buffer
585+ // being readable and writable storage in the same dispatch, and
586+ // there is no scratch-copy fallback here — refuse rather than
587+ // silently reroute.
588+ if (argSlot.buffer === dest.buffer) {
589+ throw new UnsupportedOnGpu(
590+ `'${stmt.name} = ${ext.name}(${arg.name})' reads and ` +
591+ `writes the same buffer; assign to a new name instead`,
592+ stmt.span,
593+ );
594+ }
595+ planned.push({
596+ kind: ext.name,
597+ bindGroup:
598+ ext.name === 'dthetac'
599+ ? deriv.createDthetacBinding(argSlot.buffer, dest.buffer)
600+ : deriv.createDphicBinding(argSlot.buffer, dest.buffer),
601+ label,
602+ });
603+ return;
604+ }
605+ planned.push(
606+ ext.name === 'dtheta'
607+ ? { kind: 'dtheta', binding: deriv.createDthetaBinding(argSlot.buffer, dest.buffer), label }
608+ : { kind: 'dphi', binding: deriv.createDphiBinding(argSlot.buffer, dest.buffer), label },
609+ );
610+ } else {
611+ throw new UnsupportedOnGpu(`unknown external op '${ext.name}'`, stmt.span);
612+ }
613+ return;
614+ }
615+
616+ // Element-wise kernel. Collect the distinct tensor operands and give
617+ // them dense binding slots.
618+ const tensors = new Map<string, number>();
619+ collectTensorVars(stmt.expr, (cName) => {
620+ if (!tensors.has(cName)) tensors.set(cName, tensors.size);
621+ });
622+
623+ const label = `${stmt.name} = <${count} elements, element-wise>`;
624+ const kernel = buildKernel(
625+ stmt,
626+ {
627+ tensors,
628+ params: paramSlots,
629+ scalars: derivedScalars,
630+ } satisfies KernelInputs,
631+ count,
632+ label,
633+ );
634+
635+ const bindGroupLayout = kernelLayout(device, tensors.size);
636+ const pipeline = await makePipeline(device, kernel.code, label, bindGroupLayout);
637+
638+ // WebGPU forbids aliasing a writable storage binding with another
639+ // binding in the same group, so an in-place update (`u = u + 1`) writes
640+ // to scratch and copies back. Element-wise kernels only ever touch
641+ // their own index, so the copy is the only cost.
642+ const aliased = tensors.has(stmt.cName);
643+ const target = aliased ? alloc(`mgpu-${stmt.name}-scratch`, count) : dest;
644+
645+ const entries: GPUBindGroupEntry[] = [
646+ { binding: 0, resource: { buffer: target.buffer } },
647+ ];
648+ for (const [cName, i] of tensors) {
649+ const s = slots.get(cName);
650+ if (!s) {
651+ throw new UnsupportedOnGpu(
652+ `'${stmt.name}' reads a value with no buffer`,
653+ stmt.span,
654+ );
655+ }
656+ entries.push({ binding: i + 1, resource: { buffer: s.buffer } });
657+ }
658+ entries.push({ binding: tensors.size + 1, resource: { buffer: paramBuf } });
659+
660+ planned.push({
661+ kind: 'kernel',
662+ pipeline,
663+ bindGroup: device.createBindGroup({
664+ layout: bindGroupLayout,
665+ entries,
666+ }),
667+ count,
668+ label,
669+ copyBack: aliased
670+ ? { from: target.buffer, to: dest.buffer, bytes: 4 * count }
671+ : undefined,
672+ });
673+ }
674+
675+ /**
676+ * `[a, b] = synth(x, y)` / `[a, b] = analys(x, y)`: an explicitly grouped
677+ * transform — output k is the transform of argument k. The group is
678+ * planned as consecutive pending transforms, which materializeTransforms
679+ * then chunks into whatever batched dispatch widths the device supports
680+ * (one x4 batch, two x2, or scalars with SHT_BATCH=0) — the syntax
681+ * promises grouping intent, never a lane width, so the same source
682+ * compiles everywhere.
683+ */
684+ function planMultiTransform(stmt: MultiAssignCall): void {
685+ if (stmt.name !== 'synth' && stmt.name !== 'analys') {
686+ throw new UnsupportedOnGpu(
687+ `'${stmt.name}' does not return multiple values here — only the ` +
688+ `transforms ('synth', 'analys') support [a, b] = op(x, y) grouping`,
689+ stmt.span,
690+ );
691+ }
692+ const kind = stmt.name;
693+ for (let i = 0; i < stmt.outputs.length; i++) {
694+ const slot = stmt.outputs[i];
695+ const arg = stmt.args[i];
696+ if (!slot.binding) {
697+ throw new UnsupportedOnGpu(
698+ `every output of '${kind}' must be bound to a name — output ` +
699+ `${i + 1} is dropped, but each input costs a transform`,
700+ stmt.span,
701+ );
702+ }
703+ if (!arg || arg.kind !== 'Var') {
704+ throw new UnsupportedOnGpu(
705+ `'${kind}' must be applied to variables (argument ${i + 1})`,
706+ stmt.span,
707+ );
708+ }
709+ const argSlot = slots.get(arg.cName);
710+ if (!argSlot) {
711+ throw new UnsupportedOnGpu(
712+ `'${kind}' reads '${arg.name}', which has no buffer`,
713+ stmt.span,
714+ );
715+ }
716+ if (!isNumeric(slot.ty) || !isTensor(slot.ty)) {
717+ throw new UnsupportedOnGpu(
718+ `'${slot.binding.name}' is not a numeric array`,
719+ stmt.span,
720+ );
721+ }
722+ const count = numel(slot.ty);
723+ let dest = slots.get(slot.binding.cName);
724+ if (!dest) {
725+ dest = alloc(`mgpu-${slot.binding.name}`, count);
726+ slots.set(slot.binding.cName, dest);
727+ } else if (dest.count !== count) {
728+ throw new UnsupportedOnGpu(
729+ `'${slot.binding.name}' changes size between assignments`,
730+ stmt.span,
731+ );
732+ }
733+ byName.set(slot.binding.name, dest);
734+ planned.push({
735+ pending: true,
736+ kind,
737+ in: argSlot.buffer,
738+ out: dest.buffer,
739+ label: `${slot.binding.name} = ${kind}(${arg.name})`,
740+ });
741+ }
742+ }
743+
744+ /**
745+ * `f = randnfun3(lambda, gx, gy, gz)`: the seeded random field, summed
746+ * over its Fourier modes at every surface point.
747+ *
748+ * One dispatch, one thread per point. The coefficient table is not an
749+ * argument — it is a host buffer this plan binds and the host refills per
750+ * seed, the way `synth` reads Legendre matrices the .m never names. What
751+ * the .m *does* choose is the wavelength, which is recorded here so the
752+ * host draws the table for exactly that value.
753+ */
754+ async function planRandnfun3(
755+ stmt: Assign,
756+ args: IRExpr[],
757+ dest: Slot,
758+ ): Promise<void> {
759+ const lam = args[0];
760+ const lambda: Randnfun3Lambda | null =
761+ lam.kind === 'NumLit'
762+ ? { kind: 'const', value: lam.value }
763+ : lam.kind === 'Var' && paramSlots.has(lam.cName)
764+ ? { kind: 'param', name: lam.name }
765+ : null;
766+ if (!lambda) {
767+ throw new UnsupportedOnGpu(
768+ `randnfun3's wavelength is drawn on the host before the step runs, ` +
769+ `so it must be a number or a model parameter — not a value ` +
770+ `computed on the GPU`,
771+ stmt.span,
772+ );
773+ }
774+ if (randnfun3Lambda && !sameLambda(randnfun3Lambda, lambda)) {
775+ throw new UnsupportedOnGpu(
776+ `this function calls randnfun3 with two different wavelengths; ` +
777+ `one coefficient table is drawn per plan, so only one is supported`,
778+ stmt.span,
779+ );
780+ }
781+ randnfun3Lambda = lambda;
782+
783+ const points = args.slice(1).map((a) => {
784+ const v = a as IRExpr & { kind: 'Var' };
785+ const slot = slots.get(v.cName);
786+ if (!slot) {
787+ throw new UnsupportedOnGpu(
788+ `randnfun3 reads '${v.name}', which has no buffer`,
789+ stmt.span,
790+ );
791+ }
792+ return { slot, name: v.name };
793+ });
794+
795+ const modes = host.ensure(MODE_BUFFER, modeTableLength(INITIAL_MODES));
796+ const label =
797+ `${stmt.name} = randnfun3(${
798+ lambda.kind === 'const' ? lambda.value : lambda.name
799+ }, ${points.map((p) => p.name).join(', ')})`;
800+
801+ const bindGroupLayout = device.createBindGroupLayout({
802+ label: 'mgpu-randnfun3',
803+ entries: [0, 1, 2, 3, 4].map((binding) => ({
804+ binding,
805+ visibility: GPUShaderStage.COMPUTE,
806+ buffer: { type: binding === 0 ? ('storage' as const) : ('read-only-storage' as const) },
807+ })),
808+ });
809+ // The table is sized to whatever wavelength is actually asked for, so a
810+ // finer one reallocates it — and with it these bind groups, which are
811+ // the only things holding the old buffer.
812+ const bind = (table: GPUBuffer): GPUBindGroup =>
813+ device.createBindGroup({
814+ layout: bindGroupLayout,
815+ entries: [
816+ { binding: 0, resource: { buffer: dest.buffer } },
817+ ...points.map((p, i) => ({
818+ binding: i + 1,
819+ resource: { buffer: p.slot.buffer },
820+ })),
821+ { binding: 4, resource: { buffer: table } },
822+ ],
823+ });
824+
825+ // One dispatch per slice of the mode table — see randnfun3Chunks. Each
826+ // reads the same table and accumulates into the same output, so they
827+ // share a bind group and differ only in their compiled slice index.
828+ const ops: (Op & { kind: 'kernel' })[] = [];
829+ for (let chunk = 0; chunk < randnfun3Chunks; chunk++) {
830+ const chunkLabel = `${label} [${chunk + 1}/${randnfun3Chunks}]`;
831+ const op = {
832+ kind: 'kernel' as const,
833+ pipeline: await makePipeline(
834+ device,
835+ randnfun3WGSL(dest.count, chunk),
836+ chunkLabel,
837+ bindGroupLayout,
838+ ),
839+ bindGroup: bind(modes.buffer),
840+ count: dest.count,
841+ label: chunkLabel,
842+ yieldAfter: true,
843+ };
844+ ops.push(op);
845+ planned.push(op);
846+ }
847+ rebindRandnfun3 = (table: GPUBuffer): void => {
848+ const group = bind(table);
849+ for (const op of ops) op.bindGroup = group;
850+ };
851+ }
852+
853+ /**
854+ * Unroll a counted loop into the op sequence.
855+ *
856+ * A plan is a fixed list of GPU operations with no branching, which is what
857+ * makes a timestep pure command recording. A `for` with compile-time-known
858+ * bounds still fits that: it is the same body planned once per iteration.
859+ * Nothing else changes — numbl gives a variable one cName for every
860+ * assignment to it, so the buffer an iteration writes is the buffer the
861+ * next one reads, which is exactly a loop-carried value.
862+ *
863+ * The loop variable gets no buffer either: it is bound as a derived scalar
864+ * to this iteration's literal value, so a kernel that reads `k` folds the
865+ * number in. The binding is overwritten per iteration, before that
866+ * iteration's body is planned and its WGSL emitted.
867+ */
868+ async function planFor(stmt: For): Promise<void> {
869+ const from = exactValue(stmt.start);
870+ const to = exactValue(stmt.end);
871+ if (from === undefined || to === undefined) {
872+ throw new UnsupportedOnGpu(
873+ `a 'for' loop is unrolled into the op sequence, so its bounds must ` +
874+ `be known when the model is compiled — ` +
875+ `${from === undefined ? 'the start' : 'the end'} of this one is a ` +
876+ `runtime value. Use a whole number, or a count the app supplies ` +
877+ `as a fixed argument (changing it recompiles).`,
878+ stmt.span,
879+ );
880+ }
881+ const trips = Math.floor((to - from) / stmt.step) + 1;
882+ if (!Number.isFinite(trips)) {
883+ throw new UnsupportedOnGpu(`'for ${stmt.varName}' has no finite length`, stmt.span);
884+ }
885+ if (trips > MAX_UNROLL) {
886+ throw new UnsupportedOnGpu(
887+ `'for ${stmt.varName}' would unroll to ${trips} iterations, over the ` +
888+ `limit of ${MAX_UNROLL}. Every iteration is separate GPU work, so a ` +
889+ `long loop compiles slowly and runs no faster than writing it out.`,
890+ stmt.span,
891+ );
892+ }
893+ for (let i = 0; i < trips; i++) {
894+ const value = from + i * stmt.step;
895+ derivedScalars.set(stmt.cVar, {
896+ name: stmt.varName,
897+ expr: {
898+ kind: 'NumLit',
899+ value,
900+ ty: scalarDouble(
901+ value > 0 ? 'positive' : value < 0 ? 'negative' : 'zero',
902+ value,
903+ ),
904+ span: stmt.span,
905+ },
906+ });
907+ for (const s of stmt.body) await planStatement(s);
908+ }
909+ }
910+ }
911+
912+ /** Upload parameter values, in `paramNames` order. Cheap — call freely. */
913+ setParams(values: Record<string, number>): void {
914+ this.paramNames.forEach((name, i) => {
915+ const v = values[name];
916+ this.#paramData[i] = Number.isFinite(v) ? v : 0;
917+ });
918+ this.#device.queue.writeBuffer(
919+ this.#paramBuf,
920+ 0,
921+ this.#paramData as Float32Array<ArrayBuffer>,
922+ );
923+ }
924+
925+ /** Buffer holding the named value, or undefined if the .m never binds it. */
926+ buffer(name: string): GPUBuffer | undefined {
927+ return this.#byName.get(name)?.buffer;
928+ }
929+
930+ elementCount(name: string): number | undefined {
931+ return this.#byName.get(name)?.count;
932+ }
933+
934+ /**
935+ * Run one pass of this plan, submitting in pieces so the GPU is not held for
936+ * the whole of it.
937+ *
938+ * For `init` only, and only because the seed field's mode sum can be huge:
939+ * at a fine wavelength the dispatches add up to tens of seconds, and a
940+ * browser's GPU process is shared with compositing, so one submission that
941+ * long stops the whole browser painting — the user's tabs included. Ops
942+ * marked `yieldAfter` (the randnfun3 chunks) end their submission and give
943+ * the queue back before the next one is recorded, which turns a freeze into
944+ * a wait. Everything else is recorded exactly as `encodeSteps` would.
945+ */
946+ async submitYielding(label: string): Promise<void> {
947+ let encoder = this.#device.createCommandEncoder({ label });
948+ let any = false;
949+ for (const group of this.#yieldGroups()) {
950+ if (any) {
951+ // Let the queue drain, then hand the event loop back, so compositing
952+ // and input get a turn between chunks.
953+ await this.#device.queue.onSubmittedWorkDone();
954+ await new Promise((r) => setTimeout(r, 0));
955+ encoder = this.#device.createCommandEncoder({ label });
956+ }
957+ this.#encodeOps(encoder, group);
958+ this.#device.queue.submit([encoder.finish()]);
959+ any = true;
960+ }
961+ if (!any) {
962+ this.#encodeOps(encoder, []);
963+ this.#device.queue.submit([encoder.finish()]);
964+ }
965+ }
966+
967+ /** The op list split at every `yieldAfter` boundary. */
968+ *#yieldGroups(): Generator<Op[]> {
969+ let group: Op[] = [];
970+ for (const op of this.#ops) {
971+ group.push(op);
972+ if (op.kind === 'kernel' && op.yieldAfter) {
973+ yield group;
974+ group = [];
975+ }
976+ }
977+ if (group.length) yield group;
978+ }
979+
980+ /**
981+ * Record `steps` timesteps. Synchronous: no awaits, no readback. All of the
982+ * ops share one compute pass, which WebGPU executes in submission order
983+ * with a barrier between dispatches.
984+ */
985+ encodeSteps(encoder: GPUCommandEncoder, steps: number): void {
986+ for (let s = 0; s < steps; s++) this.#encodeOps(encoder, this.#ops);
987+ }
988+
989+ /** Record one pass over `ops` into `encoder`. */
990+ #encodeOps(encoder: GPUCommandEncoder, ops: Op[]): void {
991+ {
992+ let pass: GPUComputePassEncoder | null = null;
993+ const inPass = (): GPUComputePassEncoder => {
994+ if (!pass) pass = encoder.beginComputePass({ label: 'mgpu-step' });
995+ return pass;
996+ };
997+ const endPass = (): void => {
998+ if (pass) {
999+ pass.end();
1000+ pass = null;
1001+ }
1002+ };
1003+ for (const op of ops) {
1004+ switch (op.kind) {
1005+ case 'kernel': {
1006+ const p = inPass();
1007+ p.setPipeline(op.pipeline);
1008+ p.setBindGroup(0, op.bindGroup);
1009+ p.dispatchWorkgroups(Math.ceil(op.count / WORKGROUP_SIZE));
1010+ if (op.copyBack) {
1011+ endPass();
1012+ encoder.copyBufferToBuffer(
1013+ op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
1014+ );
1015+ }
1016+ break;
1017+ }
1018+ case 'synth':
1019+ this.#shtInto(inPass(), op);
1020+ break;
1021+ case 'analys':
1022+ this.#shtInto(inPass(), op);
1023+ break;
1024+ case 'dtheta':
1025+ this.#derivInto(inPass(), op);
1026+ break;
1027+ case 'dphi':
1028+ this.#derivInto(inPass(), op);
1029+ break;
1030+ case 'dthetac':
1031+ this.#deriv!.encodeDthetacInto(inPass(), op.bindGroup);
1032+ break;
1033+ case 'dphic':
1034+ this.#deriv!.encodeDphicInto(inPass(), op.bindGroup);
1035+ break;
1036+ case 'dphig':
1037+ this.#sht.encodeDphigInto(inPass(), op.binding);
1038+ break;
1039+ case 'synth-batch':
1040+ this.#sht.encodeSynthBatchInto(inPass(), op.binding);
1041+ break;
1042+ case 'analys-batch':
1043+ this.#sht.encodeAnalysBatchInto(inPass(), op.binding);
1044+ break;
1045+ case 'copy':
1046+ endPass();
1047+ encoder.copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
1048+ break;
1049+ }
1050+ }
1051+ endPass();
1052+ }
1053+ }
1054+
1055+ #shtInto(pass: GPUComputePassEncoder, op: Op & { kind: 'synth' | 'analys' }): void {
1056+ if (op.kind === 'synth') this.#sht.encodeSynthInto(pass, op.binding);
1057+ else this.#sht.encodeAnalysInto(pass, op.binding);
1058+ }
1059+
1060+ #derivInto(pass: GPUComputePassEncoder, op: Op & { kind: 'dtheta' | 'dphi' }): void {
1061+ // planStatement already refused to plan a dtheta/dphi op without a
1062+ // DerivPlan, so #deriv is guaranteed set whenever an op of this kind exists.
1063+ if (op.kind === 'dtheta') this.#deriv!.encodeDthetaInto(pass, op.binding);
1064+ else this.#deriv!.encodeDphiInto(pass, op.binding);
1065+ }
1066+
1067+ /**
1068+ * Human-readable op sequence — what the .m actually compiled to. Batched
1069+ * transforms list one line per lane, annotated: the line count equals the
1070+ * logical op count regardless of the device's batch width, so op-count
1071+ * assertions in the tests are batch-invariant.
1072+ */
1073+ describe(): string[] {
1074+ return this.#ops.flatMap((op) => {
1075+ if ('labels' in op) {
1076+ const kind = op.kind === 'synth-batch' ? 'synth' : 'analys';
1077+ return op.labels.map(
1078+ (label, i) =>
1079+ `${kind.padEnd(7)} ${label} [batch lane ${i + 1}/${op.binding.size}]`,
1080+ );
1081+ }
1082+ return [`${op.kind.padEnd(7)} ${op.label}`];
1083+ });
1084+ }
1085+
1086+ destroy(): void {
1087+ for (const b of this.#owned) b.destroy();
1088+ this.#paramBuf.destroy();
1089+ this.#owned.length = 0;
1090+ }
1091+}
1092+
1093+/**
1094+ * `x = synth(y)` / `x = randnfun3(lam, gx, gy, gz)` -> the call's name and
1095+ * arguments.
1096+ *
1097+ * Every external op but `randnfun3` takes exactly one array; `randnfun3`
1098+ * takes a wavelength and the three surface coordinates. Its wavelength may
1099+ * be a literal, so arguments are returned as expressions and the caller
1100+ * decides which it needs as a buffer.
1101+ */
1102+function externalCall(
1103+ stmt: Assign,
1104+): { name: string; args: IRExpr[] } | null {
1105+ const e = stmt.expr;
1106+ if (e.kind !== 'Call' || !EXTERNAL_OPS.has(e.name)) return null;
1107+ const arity = e.name === 'randnfun3' ? 4 : 1;
1108+ if (e.args.length !== arity) {
1109+ throw new UnsupportedOnGpu(
1110+ arity === 1
1111+ ? `'${e.name}' must be applied to a single variable`
1112+ : `'${e.name}' takes ${arity} arguments, got ${e.args.length}`,
1113+ stmt.span,
1114+ );
1115+ }
1116+ // Only the wavelength may be something other than a plain variable.
1117+ for (let i = e.name === 'randnfun3' ? 1 : 0; i < e.args.length; i++) {
1118+ if (e.args[i].kind !== 'Var') {
1119+ throw new UnsupportedOnGpu(
1120+ `'${e.name}' must be applied to variables, not expressions`,
1121+ stmt.span,
1122+ );
1123+ }
1124+ }
1125+ return { name: e.name, args: e.args };
1126+}
1127+
1128+/** A `randnfun3` wavelength argument: a literal, or the parameter to read it
1129+ * from when the host fills the coefficient table. */
1130+export type Randnfun3Lambda =
1131+ | { kind: 'const'; value: number }
1132+ | { kind: 'param'; name: string };
1133+
1134+const sameLambda = (a: Randnfun3Lambda, b: Randnfun3Lambda): boolean =>
1135+ a.kind === 'const' && b.kind === 'const'
1136+ ? a.value === b.value
1137+ : a.kind === 'param' && b.kind === 'param' && a.name === b.name;
1138+
1139+/** The wavelength value a plan's `randnfun3` call resolves to. */
1140+export const resolveLambda = (
1141+ lambda: Randnfun3Lambda,
1142+ params: Record<string, number>,
1143+): number => (lambda.kind === 'const' ? lambda.value : params[lambda.name]);
1144+
1145+function collectTensorVars(e: IRExpr, visit: (cName: string) => void): void {
1146+ const walk = (x: IRExpr): void => {
1147+ switch (x.kind) {
1148+ case 'Var':
1149+ if (isTensor(x.ty)) visit(x.cName);
1150+ return;
1151+ case 'Binary':
1152+ walk(x.left);
1153+ walk(x.right);
1154+ return;
1155+ case 'Unary':
1156+ walk(x.operand);
1157+ return;
1158+ case 'Call':
1159+ x.args.forEach(walk);
1160+ return;
1161+ default:
1162+ return;
1163+ }
1164+ };
1165+ walk(e);
1166+}
src/mgpu/randnfun3.tsadded+294−0View file
@@ -0,0 +1,294 @@
1+/**
2+ * `randnfun3` — a smooth random function in 3D, evaluated at the surface.
3+ *
4+ * chebfun's randnfun3 is a random trig series on a box: a few thousand
5+ * Fourier modes with independent normal coefficients, confined to a ball for
6+ * isotropy and normalized to unit variance. Restricting it to a surface is
7+ * just evaluating it at the surface's points, which is what a model's `init`
8+ * wants for a seeded initial condition (surfacefun seeds exactly this way).
9+ *
10+ * The work splits in two, and the split is forced rather than chosen:
11+ *
12+ * - **Drawing the modes needs `randn`**, which the compiled WGSL dialect has
13+ * no counterpart for, and `sqrt(nnz)` normalization, which is a reduction.
14+ * Both are a few lines of MATLAB, so the draw lives in
15+ * `tools/randnfun3.m` and runs in numbl's interpreter — a few thousand
16+ * numbers, ~5 ms.
17+ * - **Evaluating is npts x nmodes**, ~6e7 terms at the default lambda. That
18+ * is the whole cost, and it is what this file's kernel does on the GPU.
19+ *
20+ * So the .m calls `f = randnfun3(lambda, gx, gy, gz)` — chebfun's signature,
21+ * lambda in and values out — and the coefficient table is filled in behind it
22+ * by the host, the way `synth` hides its Legendre matrices. lambda is not
23+ * decorative: the plan records which parameter the .m passed, and the host
24+ * draws the table from *that* parameter's value (src/mgpu/plan.ts,
25+ * `randnfun3Lambda`), so changing it in the .m changes the field.
26+ */
27+import { executeCode } from 'numbl-src/numbl-core/executeCode.ts';
28+import { isRuntimeTensor } from 'numbl-src/numbl-core/runtime/types.ts';
29+import { toolFiles } from '../tools.ts';
30+
31+/**
32+ * Dispatches the mode sum is split across.
33+ *
34+ * lambda is an absolute length and the mode count goes as its inverse cube,
35+ * so halving lambda costs eight times the work — there is no natural ceiling
36+ * to put on that, and nothing in the method breaks as it grows. It just gets
37+ * slower, which is the caller's business. What is *not* the caller's business
38+ * is a browser's GPU-process watchdog, which kills the device outright when a
39+ * single dispatch runs too long; a fine wavelength would otherwise turn "this
40+ * takes a while" into "device lost".
41+ *
42+ * So the sum is split into a fixed number of dispatches, each covering its own
43+ * slice of the table and accumulating into the same output. The count is fixed
44+ * at plan time (the op sequence has no runtime branching) and the slice bounds
45+ * come from the table's header, so one plan serves any wavelength. Slices that
46+ * fall past the end of a small table exit immediately, which is why a coarse
47+ * wavelength pays nothing for the split.
48+ */
49+const CHUNKS = 16;
50+
51+/** Floats the table needs for `nmodes` modes. */
52+export const modeTableLength = (nmodes: number): number =>
53+ HEADER + STRIDE * nmodes;
54+
55+/** Modes a table holds, from its header. */
56+export const modeCount = (table: Float32Array): number => table[0];
57+
58+/**
59+ * Largest table this will try to build, in f32. Not a policy about how fine a
60+ * wavelength is sensible — that is the caller's call, and a fine one is
61+ * merely slow — but the point past which the draw would fail anyway: the
62+ * host-side Float32Array alone would be 8 GB. The device's own
63+ * storage-buffer limit is checked separately, when the buffer is allocated.
64+ */
65+const MAX_TABLE_FLOATS = 2 ** 31;
66+
67+/** What the table starts at, before any seed has been drawn. Big enough for
68+ * the default wavelength on the shipped surfaces, so the common case never
69+ * reallocates. */
70+export const INITIAL_MODES = 4096;
71+
72+/** Wavelength of the seeded field when the app names none. Fine enough to
73+ * give a Turing pattern plenty to grow from, coarse enough that the draw is
74+ * ~1,400 modes rather than the ~11,500 of the slider's finest setting. */
75+export const DEFAULT_LAMBDA = 0.5;
76+
77+/** Floats before the first mode: `[nmodes, 0, 0, 0]`. The count travels in
78+ * the buffer rather than a second binding, so the kernel needs one storage
79+ * buffer and the host one write. */
80+const HEADER = 4;
81+/** Floats per mode: kx, ky, kz, real, imag. */
82+const STRIDE = 5;
83+
84+/** The name the coefficient buffer takes in the plan's HostBuffers. */
85+export const MODE_BUFFER = 'randnfun3_modes';
86+
87+/** How many dispatches `randnfun3WGSL` must be planned as. */
88+export const randnfun3Chunks = CHUNKS;
89+
90+/**
91+ * One thread per surface point, summing this chunk's slice of the modes.
92+ *
93+ * The inner loop is a dot product, a cos, a sin and two multiply-adds, over a
94+ * table small enough (~1,400 modes at the default lambda) to sit in cache for
95+ * every thread. Chunk 0 initializes the output and the rest accumulate onto
96+ * it; dispatches within one compute pass are ordered, so the reads see the
97+ * previous chunk's writes. Nothing here is per-step work: `init` runs once a
98+ * seed.
99+ */
100+export function randnfun3WGSL(npts: number, chunk: number): string {
101+ return `
102+@group(0) @binding(0) var<storage, read_write> outf: array<f32>;
103+@group(0) @binding(1) var<storage, read> px: array<f32>;
104+@group(0) @binding(2) var<storage, read> py: array<f32>;
105+@group(0) @binding(3) var<storage, read> pz: array<f32>;
106+// [nmodes, _, _, _], then kx, ky, kz, re, im per mode.
107+@group(0) @binding(4) var<storage, read> modes: array<f32>;
108+
109+@compute @workgroup_size(64)
110+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
111+ let i = gid.x;
112+ if (i >= ${npts}u) { return; }
113+ let n = u32(modes[0]);
114+ // This chunk's slice. Ceiling division, so the last slices are the short
115+ // ones and an empty slice costs a single comparison.
116+ let per = (n + ${CHUNKS}u - 1u) / ${CHUNKS}u;
117+ let lo = min(${chunk}u * per, n);
118+ let hi = min(lo + per, n);
119+ var acc = 0.0;
120+ if (lo < hi) {
121+ let x = px[i];
122+ let y = py[i];
123+ let z = pz[i];
124+ for (var m = lo; m < hi; m = m + 1u) {
125+ let b = ${HEADER}u + m * ${STRIDE}u;
126+ let t = modes[b] * x + modes[b + 1u] * y + modes[b + 2u] * z;
127+ acc = acc + modes[b + 3u] * cos(t) - modes[b + 4u] * sin(t);
128+ }
129+ }
130+${chunk === 0 ? ' outf[i] = acc;' : ' outf[i] = outf[i] + acc;'}
131+}
132+`;
133+}
134+
135+/**
136+ * Modes a wavelength will draw on a box, without drawing them: chebfun's
137+ * cube size, times the fraction its isotropy ball keeps (pi/6 of a cube,
138+ * approached from below at small m). Used to price a wavelength up front.
139+ */
140+function plannedModes(lambda: number, box: BoundingBox): number {
141+ const side = (w: number): number => 2 * Math.round((1.2 * w) / lambda + 2) + 1;
142+ const cube =
143+ side(box.x1 - box.x0) * side(box.y1 - box.y0) * side(box.z1 - box.z0);
144+ return Math.ceil((Math.PI / 6) * cube);
145+}
146+
147+/** The box a random field is drawn over: the surface's own bounding box. */
148+export interface BoundingBox {
149+ x0: number; x1: number;
150+ y0: number; y1: number;
151+ z0: number; z1: number;
152+}
153+
154+/** The bounding box of a surface, as `Geometry` holds its coordinates. */
155+export function boundingBox(
156+ x: Float32Array,
157+ y: Float32Array,
158+ z: Float32Array,
159+): BoundingBox {
160+ const box = {
161+ x0: Infinity, x1: -Infinity,
162+ y0: Infinity, y1: -Infinity,
163+ z0: Infinity, z1: -Infinity,
164+ };
165+ for (let i = 0; i < x.length; i++) {
166+ if (x[i] < box.x0) box.x0 = x[i];
167+ if (x[i] > box.x1) box.x1 = x[i];
168+ if (y[i] < box.y0) box.y0 = y[i];
169+ if (y[i] > box.y1) box.y1 = y[i];
170+ if (z[i] < box.z0) box.z0 = z[i];
171+ if (z[i] > box.z1) box.z1 = z[i];
172+ }
173+ return box;
174+}
175+
176+/**
177+ * Draw a field's modes and pack them for the GPU: `tools/randnfun3.m` run
178+ * through the interpreter, seeded, then interleaved into the buffer layout
179+ * above. Column-major out of MATLAB, interleaved on the way in.
180+ */
181+export function drawModes(
182+ lambda: number,
183+ box: BoundingBox,
184+ seed: number,
185+ /** Points the field will be summed at, for the cost budget. */
186+ npts: number,
187+): Float32Array {
188+ if (!(lambda > 0) || !Number.isFinite(lambda)) {
189+ throw new Error(`randnfun3: lambda must be a positive number, got ${lambda}`);
190+ }
191+ // The mode count follows from lambda and the box alone, so a table that
192+ // cannot be built is refused before anything is drawn. The only ceiling is
193+ // what fits: how slow a fine wavelength is, is the caller's to decide.
194+ const planned = plannedModes(lambda, box);
195+ if (modeTableLength(planned) > MAX_TABLE_FLOATS) {
196+ throw new Error(
197+ `randnfun3: lambda ${lambda} needs about ` +
198+ `${planned.toLocaleString()} Fourier modes on this surface, a ` +
199+ `${((4 * modeTableLength(planned)) / 1e9).toFixed(1)} GB table. ` +
200+ `lambda is an absolute length, so a larger surface needs more modes ` +
201+ `for the same value, and halving it costs eight times as many.`,
202+ );
203+ }
204+ const result = executeCode(
205+ 'rng(seed); [k, c] = randnfun3(lambda, [x0 x1 y0 y1 z0 z1]);',
206+ {
207+ initialVariableValues: { lambda, seed, ...box },
208+ displayResults: false,
209+ implicitCwdPath: null,
210+ },
211+ toolFiles,
212+ 'randnfun3-driver.m',
213+ );
214+ const k = result.variableValues['k'];
215+ const c = result.variableValues['c'];
216+ if (!k || !c || !isRuntimeTensor(k) || !isRuntimeTensor(c)) {
217+ throw new Error("randnfun3: tools/randnfun3.m did not return [k, c] arrays");
218+ }
219+ const nmodes = k.shape[0];
220+ const out = new Float32Array(modeTableLength(nmodes));
221+ out[0] = nmodes;
222+ for (let i = 0; i < nmodes; i++) {
223+ const b = HEADER + STRIDE * i;
224+ out[b] = k.data[i]; // kx
225+ out[b + 1] = k.data[nmodes + i]; // ky
226+ out[b + 2] = k.data[2 * nmodes + i]; // kz
227+ out[b + 3] = c.data[i]; // real
228+ out[b + 4] = c.data[nmodes + i]; // imag
229+ }
230+ return out;
231+}
232+
233+/**
234+ * `drawModes` on a worker thread, so a fine wavelength does not freeze the
235+ * page (src/mgpu/randnfun3.worker.ts).
236+ *
237+ * Falls back to drawing in place where there is no `Worker` — the node test
238+ * runner and the desktop benchmark, neither of which has an event loop it
239+ * would matter to. Failures surface as a rejection either way, so a caller
240+ * never has to know which path ran.
241+ */
242+export function drawModesAsync(
243+ lambda: number,
244+ box: BoundingBox,
245+ seed: number,
246+ npts: number,
247+): Promise<Float32Array> {
248+ if (typeof Worker === 'undefined') {
249+ try {
250+ return Promise.resolve(drawModes(lambda, box, seed, npts));
251+ } catch (e) {
252+ return Promise.reject(e instanceof Error ? e : new Error(String(e)));
253+ }
254+ }
255+ const w = drawWorker();
256+ const id = nextDrawId++;
257+ return new Promise((resolve, reject) => {
258+ pendingDraws.set(id, { resolve, reject });
259+ w.postMessage({ id, lambda, box, seed, npts });
260+ });
261+}
262+
263+let worker: Worker | null = null;
264+let nextDrawId = 1;
265+const pendingDraws = new Map<
266+ number,
267+ { resolve: (t: Float32Array) => void; reject: (e: Error) => void }
268+>();
269+
270+/** The draw worker, started on first use and kept for the session — starting
271+ * one re-parses numbl, which costs more than a coarse draw does. */
272+function drawWorker(): Worker {
273+ if (worker) return worker;
274+ worker = new Worker(new URL('./randnfun3.worker.ts', import.meta.url), {
275+ type: 'module',
276+ });
277+ worker.onmessage = (e: MessageEvent<{ id: number; table?: Float32Array; error?: string }>): void => {
278+ const waiting = pendingDraws.get(e.data.id);
279+ if (!waiting) return;
280+ pendingDraws.delete(e.data.id);
281+ if (e.data.error !== undefined) waiting.reject(new Error(e.data.error));
282+ else waiting.resolve(e.data.table!);
283+ };
284+ worker.onerror = (e: ErrorEvent): void => {
285+ // A worker that died takes every outstanding draw with it.
286+ for (const [, waiting] of pendingDraws) {
287+ waiting.reject(new Error(`randnfun3 draw worker failed: ${e.message}`));
288+ }
289+ pendingDraws.clear();
290+ worker?.terminate();
291+ worker = null;
292+ };
293+ return worker;
294+}
src/mgpu/randnfun3.worker.tsadded+40−0View file
@@ -0,0 +1,40 @@
1+/**
2+ * Drawing a seed field's Fourier modes, off the main thread.
3+ *
4+ * The draw is `tools/randnfun3.m` in numbl's interpreter, and its cost goes as
5+ * the inverse cube of the wavelength: milliseconds at the default lambda, but
6+ * ~13 s at lambda 0.01. Synchronous JS that long does not merely feel slow —
7+ * it blocks the event loop outright, so the page stops painting and the
8+ * browser offers to kill it. Nothing about the draw needs the main thread
9+ * (it touches no GPU and no DOM), so it runs here and the result is
10+ * transferred back.
11+ */
12+import { drawModes, type BoundingBox } from './randnfun3.ts';
13+
14+export interface DrawRequest {
15+ id: number;
16+ lambda: number;
17+ box: BoundingBox;
18+ seed: number;
19+ npts: number;
20+}
21+
22+export type DrawReply =
23+ | { id: number; table: Float32Array; error?: undefined }
24+ | { id: number; table?: undefined; error: string };
25+
26+self.onmessage = (e: MessageEvent<DrawRequest>): void => {
27+ const { id, lambda, box, seed, npts } = e.data;
28+ let reply: DrawReply;
29+ let transfer: Transferable[] = [];
30+ try {
31+ const table = drawModes(lambda, box, seed, npts);
32+ reply = { id, table };
33+ transfer = [table.buffer];
34+ } catch (err) {
35+ reply = { id, error: err instanceof Error ? err.message : String(err) };
36+ }
37+ (self as unknown as {
38+ postMessage: (m: DrawReply, t: Transferable[]) => void;
39+ }).postMessage(reply, transfer);
40+};
src/mgpu/registry.tsadded+83−0View file
@@ -0,0 +1,83 @@
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+ * Trimmed from turing-surface: this app ships one model (Schnakenberg, flux
11+ * form). The discrete parameter choices the app actually offers live in
12+ * src/cache/options.ts; the min/max/step here are only the numeric bounds.
13+ *
14+ * Naming convention, documented in each .m:
15+ * `u`, `v`, ... grid fields the model computes and the app renders
16+ * `U`, `V`, ... the corresponding spectral state (uppercase)
17+ */
18+import schnakenbergSource from '../../models/schnakenberg.m?raw';
19+
20+export type Params = Record<string, number>;
21+
22+/** A tunable scalar the .m may take as an argument. */
23+export interface ParamSpec {
24+ key: string;
25+ label: string;
26+ value: number;
27+ min: number;
28+ max: number;
29+ step: number;
30+ /**
31+ * This parameter is a random seed: its value picks a draw and means nothing
32+ * on its own, so the UI offers a button that jumps to another one rather
33+ * than a box to type a number into. `min`/`max` still bound what the button
34+ * picks.
35+ */
36+ reseed?: boolean;
37+}
38+
39+export interface MModel {
40+ key: string;
41+ label: string;
42+ blurb: string;
43+ /** Grid fields to render, one panel each. */
44+ species: string[];
45+ /** Spectral state names the .m advances. */
46+ state: string[];
47+ params: ParamSpec[];
48+ /** Polynomial degree of the reaction in the fields, for grid dealiasing. */
49+ pdeg: number;
50+ /** Amplitude of the seeded perturbation handed to `init`. */
51+ seedAmp: number;
52+ /** MATLAB source — the algorithm itself. */
53+ source: string;
54+}
55+
56+/** Spectral state names follow the grid-field names, uppercased. */
57+const stateFor = (species: string[]): string[] => species.map((s) => s.toUpperCase());
58+
59+const schnakenberg: MModel = {
60+ key: 'schnakenberg',
61+ label: 'Schnakenberg',
62+ blurb: 'Turing spots.',
63+ species: ['u', 'v'],
64+ state: stateFor(['u', 'v']),
65+ params: [
66+ { key: 'a', label: 'a', value: 0.1, min: 0.01, max: 0.5, step: 0.01 },
67+ { key: 'b', label: 'b', value: 0.9, min: 0.1, max: 2, step: 0.05 },
68+ { key: 'D1', label: 'D₁', value: 4e-4, min: 1e-5, max: 5e-3, step: 1e-5 },
69+ { key: 'D2', label: 'D₂', value: 8e-3, min: 1e-4, max: 5e-2, step: 1e-4 },
70+ { key: 'dt', label: 'dt', value: 0.05, min: 0.005, max: 0.5, step: 0.005 },
71+ ],
72+ pdeg: 3,
73+ seedAmp: 1e-2,
74+ source: schnakenbergSource,
75+};
76+
77+export const mModels: MModel[] = [schnakenberg];
78+
79+export const mModelByKey = (key: string): MModel | undefined =>
80+ mModels.find((m) => m.key === key);
81+
82+export const defaultParams = (m: MModel): Params =>
83+ Object.fromEntries(m.params.map((p) => [p.key, p.value]));
src/mgpu/session.tsadded+439−0View file
@@ -0,0 +1,439 @@
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 { DerivPlan } from '../sht/deriv.ts';
11+import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
12+import { GpuModel, type ModelParams } from './model.ts';
13+import { seededNoise } from './noise.ts';
14+import { boundingBox, drawModesAsync, DEFAULT_LAMBDA } from './randnfun3.ts';
15+import { resolveLambda } from './plan.ts';
16+import type { MModel } from './registry.ts';
17+import { Geometry } from '../geom/geometry.ts';
18+import { mGeometryByKey, defaultGeometryParams, SPHERE_KEY, type MGeometry } from '../geom/registry.ts';
19+
20+export interface ModelSessionOptions {
21+ device: GPUDevice;
22+ model: MModel;
23+ params: ModelParams;
24+ lmax: number;
25+ /** Override the model source — the editor's working copy. */
26+ source?: string;
27+ /** Linear render oversampling: read the species fields on a grid this many
28+ * times finer than the solver's in each direction (default 1). The state is
29+ * band-limited at lmax, so the finer evaluation is exact interpolation. */
30+ oversample?: number;
31+ /** The surface to solve on. Defaults to the unit sphere. */
32+ geometry?: MGeometry;
33+ geometryParams?: ModelParams;
34+ /** Override the geometry source — the editor's working copy. */
35+ geometrySource?: string;
36+ /**
37+ * Iterations of the .m's implicit solve. Structural, not tunable: the loop
38+ * is unrolled into the op sequence, so a change recompiles.
39+ */
40+ niter?: number;
41+ /** Wavelength of the seeded random field a model's `init` draws
42+ * (src/mgpu/randnfun3.ts). Redrawn on the next seed, never recompiled. */
43+ lam3?: number;
44+}
45+
46+export class ModelSession {
47+ readonly device: GPUDevice;
48+ readonly model: MModel;
49+ readonly cfg: ShtConfig;
50+ readonly sht: ShtPlan;
51+ readonly gpu: GpuModel;
52+ readonly npts: number;
53+ /** Iterations of the implicit solve compiled into the step. */
54+ readonly niter: number;
55+
56+ /** The surface being solved on, as spherical-harmonic coefficients. */
57+ #geometry: Geometry;
58+ #geometryModel: MGeometry;
59+ /** Computes the theta/phi derivatives a geometry's metric quantities need. */
60+ #deriv: DerivPlan;
61+
62+ /** Model time and step count since the last seeding. */
63+ t = 0;
64+ steps = 0;
65+
66+ #params: ModelParams;
67+ /** Display-only transforms on the oversampled grid; null at 1x. */
68+ #displaySht: ShtPlan | null;
69+ #oversample: number;
70+ /** Wavelength of the seeded random field, and the seed it was drawn from —
71+ * kept so changing one can redraw with the other unchanged. */
72+ #lam3: number;
73+ #seed = 1;
74+
75+ private constructor(init: {
76+ device: GPUDevice;
77+ model: MModel;
78+ cfg: ShtConfig;
79+ sht: ShtPlan;
80+ displaySht: ShtPlan | null;
81+ gpu: GpuModel;
82+ params: ModelParams;
83+ oversample: number;
84+ geometry: Geometry;
85+ geometryModel: MGeometry;
86+ deriv: DerivPlan;
87+ niter: number;
88+ lam3: number;
89+ }) {
90+ this.device = init.device;
91+ this.model = init.model;
92+ this.cfg = init.cfg;
93+ this.sht = init.sht;
94+ this.gpu = init.gpu;
95+ this.npts = init.cfg.nlat * init.cfg.nphi;
96+ this.#oversample = init.oversample;
97+ this.#params = init.params;
98+ this.#displaySht = init.displaySht;
99+ this.#geometry = init.geometry;
100+ this.#geometryModel = init.geometryModel;
101+ this.#deriv = init.deriv;
102+ this.niter = init.niter;
103+ this.#lam3 = init.lam3;
104+ }
105+
106+ get geometry(): Geometry {
107+ return this.#geometry;
108+ }
109+
110+ get geometryModel(): MGeometry {
111+ return this.#geometryModel;
112+ }
113+
114+ /** Linear render oversampling factor (1 = read on the solver grid). */
115+ get oversample(): number {
116+ return this.#oversample;
117+ }
118+
119+ static async create(opts: ModelSessionOptions): Promise<ModelSession> {
120+ const { device, model, params, lmax } = opts;
121+ const oversample = Math.max(1, Math.round(opts.oversample ?? 1));
122+ const niter = Math.max(0, Math.round(opts.niter ?? 1));
123+ const geometryModel = opts.geometry ?? mGeometryByKey(SPHERE_KEY)!;
124+ const geometryParams = opts.geometryParams ?? defaultGeometryParams(geometryModel);
125+ const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
126+ const cfg = { lmax, mmax: lmax, nlat, nphi };
127+ const sht = await ShtPlan.create(device, cfg);
128+ let displaySht: ShtPlan | null = null;
129+ let deriv: DerivPlan | null = null;
130+ try {
131+ // The display plan shares nothing with the solver's beyond the
132+ // coefficients copied into it per readback; its grid is the solver's
133+ // scaled by the oversampling factor, so nphi stays a power of two (the
134+ // FFT path) for power-of-two factors.
135+ if (oversample > 1) {
136+ displaySht = await ShtPlan.create(device, {
137+ lmax,
138+ mmax: lmax,
139+ nlat: oversample * nlat,
140+ nphi: oversample * nphi,
141+ });
142+ }
143+ // Computes the theta/phi derivatives the geometry's metric quantities
144+ // (and, per step, the surface Laplace-Beltrami correction) need.
145+ deriv = await DerivPlan.create(device, sht);
146+ // The surface is built before the model, because the model takes it as
147+ // an argument. It is a one-off: compiled, evaluated, read back, and its
148+ // plan discarded — nothing of it survives into the timestep but sixteen
149+ // buffers of numbers (the embedding, and both metric formulations built
150+ // on it: the inverse metric quantities and the flux-form weights).
151+ const geometry = await Geometry.create({
152+ sht,
153+ cfg,
154+ source: opts.geometrySource ?? geometryModel.source,
155+ paramNames: geometryModel.params.map((p) => p.key),
156+ params: geometryParams,
157+ deriv,
158+ });
159+ const gpu = await GpuModel.create({
160+ device,
161+ sht,
162+ cfg,
163+ source: opts.source ?? model.source,
164+ paramNames: model.params.map((p) => p.key),
165+ state: model.state,
166+ view: model.species,
167+ geometry,
168+ deriv,
169+ niter,
170+ });
171+ const lam3 = opts.lam3 ?? DEFAULT_LAMBDA;
172+ gpu.setParams({ lam3, ...params });
173+ return new ModelSession({
174+ device, model, cfg, sht, displaySht, gpu, params, oversample,
175+ geometry, geometryModel, deriv, niter, lam3,
176+ });
177+ } catch (e) {
178+ // The transform plans own GPU buffers; do not leak them on a compile error.
179+ deriv?.destroy();
180+ displaySht?.destroy();
181+ sht.destroy();
182+ throw e;
183+ }
184+ }
185+
186+ /**
187+ * Vertex positions for the current render grid: the surface synthesized on
188+ * `viewSht`, interleaved xyz. Exact interpolation of the same coefficients
189+ * the solver sees, so the drawn surface is the one being solved on however
190+ * finely it is sampled.
191+ */
192+ renderPositions(): Promise<Float32Array> {
193+ return this.#geometry.positionsOn(this.viewSht);
194+ }
195+
196+ /**
197+ * Change the surface in place, without recompiling or disturbing the run.
198+ * The geometry's shape in the bindings depends only on the grid, so the
199+ * compiled step does not change — only the numbers it reads. The caller
200+ * still has to rebuild the mesh from `renderPositions()`.
201+ */
202+ async setGeometry(
203+ geometryModel: MGeometry,
204+ params: ModelParams,
205+ source?: string,
206+ ): Promise<void> {
207+ const next = await Geometry.create({
208+ sht: this.sht,
209+ cfg: this.cfg,
210+ source: source ?? geometryModel.source,
211+ paramNames: geometryModel.params.map((p) => p.key),
212+ params,
213+ deriv: this.#deriv,
214+ });
215+ this.#geometry = next;
216+ this.#geometryModel = geometryModel;
217+ this.gpu.uploadGeometry(next);
218+ // The new surface brings a new preconditioner scale (GpuModel folds its
219+ // current geometry's jhat into every params upload).
220+ this.gpu.setParams(this.#params);
221+ }
222+
223+ /** The plan whose grid `readSpecies` samples on — the display plan when
224+ * oversampling, otherwise the solver's. Its cosTheta/nphi define the mesh. */
225+ get viewSht(): ShtPlan {
226+ return this.#displaySht ?? this.sht;
227+ }
228+
229+ /**
230+ * Change the display oversampling in place. Display-only: the simulation
231+ * state, time and parameters are untouched, so the run continues seamlessly
232+ * on the new render grid. The caller must not have a readSpecies in flight —
233+ * its readback maps a buffer of the plan being destroyed.
234+ */
235+ async setOversample(oversample: number): Promise<void> {
236+ const os = Math.max(1, Math.round(oversample));
237+ await this.setDisplayGrid(os * this.cfg.nlat, os * this.cfg.nphi);
238+ }
239+
240+ /**
241+ * Point the display plan at an arbitrary grid, rather than an integer
242+ * multiple of the solver's. Same contract as setOversample — display-only,
243+ * no readback may be in flight — and the same exactness argument, which does
244+ * not care about the ratio: the state is band-limited at lmax, so
245+ * synthesizing it anywhere is evaluation, not resampling. What this adds is a
246+ * grid that need not be *finer*: several sessions at different lmax can be
247+ * put on one common grid, which is what makes their fields directly
248+ * comparable point by point and lets one mesh serve all of them.
249+ */
250+ async setDisplayGrid(nlat: number, nphi: number): Promise<void> {
251+ const view = this.viewSht.cfg;
252+ if (nlat === view.nlat && nphi === view.nphi) return;
253+ const onSolverGrid = nlat === this.cfg.nlat && nphi === this.cfg.nphi;
254+ const next = onSolverGrid
255+ ? null
256+ : await ShtPlan.create(this.device, {
257+ lmax: this.cfg.lmax,
258+ mmax: this.cfg.mmax,
259+ nlat,
260+ nphi,
261+ });
262+ const old = this.#displaySht;
263+ this.#displaySht = next;
264+ this.#oversample = nlat / this.cfg.nlat;
265+ old?.destroy();
266+ }
267+
268+ /**
269+ * The coefficient table a model that calls `randnfun3` seeds from — drawn
270+ * over the current surface's bounding box, at the wavelength its own .m asked
271+ * for — or null for a model that seeds some other way. The draw is host-side
272+ * MATLAB (a few ms); the evaluation at every grid point is the GPU kernel
273+ * inside `init`.
274+ *
275+ * Separate from `seed` because the table is a function of space, not of a
276+ * grid: drawn once, it is the *same field* wherever it is evaluated, which is
277+ * how every variant of a comparison across lmax seeds from one random field
278+ * (see src/compare/sharedStart.ts).
279+ */
280+ drawSeedModes(seed: number): Promise<Float32Array | null> {
281+ const lambda = this.gpu.randnfun3Lambda;
282+ if (!lambda) return Promise.resolve(null);
283+ // Drawn on a worker: at a fine wavelength this is seconds of interpreter
284+ // time, and it must not be seconds of frozen page.
285+ return drawModesAsync(
286+ resolveLambda(lambda, this.#mergedParams()),
287+ boundingBox(this.#geometry.x, this.#geometry.y, this.#geometry.z),
288+ seed,
289+ this.npts,
290+ );
291+ }
292+
293+ /** Run `init` from a seeded perturbation, resetting model time. */
294+ async seed(seed: number): Promise<void> {
295+ const modes = await this.drawSeedModes(seed);
296+ await this.seedWith(seededNoise(this.npts, this.model.seedAmp, seed), modes);
297+ this.#seed = seed;
298+ }
299+
300+ /**
301+ * Run `init` from a caller-supplied perturbation, resetting model time.
302+ * `seed()` is this with what this session would draw for itself: the host
303+ * RNG's field on its own grid, and its own random-field table. Supplying them
304+ * instead is how several sessions on *different* grids can be started from
305+ * the same initial condition, which is the only way a comparison across lmax
306+ * compares one problem rather than two (see src/compare/sharedStart.ts).
307+ */
308+ async seedWith(noise: Float32Array, modes: Float32Array | null = null): Promise<void> {
309+ if (noise.length !== this.npts) {
310+ throw new Error(`seedWith: noise must have length ${this.npts} (got ${noise.length})`);
311+ }
312+ await this.gpu.init(noise, modes);
313+ this.t = 0;
314+ this.steps = 0;
315+ }
316+
317+ /**
318+ * Push an exact spectral state into the running model, bypassing seeded
319+ * init, and reset model time like seed() does. `gpu.step(0)` flips which
320+ * of the init/step buffer aliases a read resolves to, without otherwise
321+ * touching the state — see GpuModel's `#lastRan`.
322+ */
323+ loadState(coeffs: Record<string, Float32Array>): void {
324+ for (const name of this.model.state) {
325+ const data = coeffs[name];
326+ if (!data) throw new Error(`loadState: missing state '${name}'`);
327+ this.gpu.upload(name, data);
328+ }
329+ this.gpu.step(0);
330+ this.t = 0;
331+ this.steps = 0;
332+ }
333+
334+ /** The model's parameters plus the ones the host owns. */
335+ #mergedParams(): ModelParams {
336+ return { lam3: this.#lam3, ...this.#params };
337+ }
338+
339+ setParams(params: ModelParams): void {
340+ this.#params = params;
341+ this.gpu.setParams(this.#mergedParams());
342+ }
343+
344+ /** Wavelength of the seeded random field. Redraws on the next seed. */
345+ get lam3(): number {
346+ return this.#lam3;
347+ }
348+
349+ /**
350+ * Change the random field's wavelength. Nothing recompiles — lam3 is a
351+ * uniform and the coefficient table is host-drawn — but the field itself
352+ * only changes on the next `seed`, which is where it is drawn.
353+ */
354+ setLam3(lambda: number): void {
355+ if (lambda === this.#lam3) return;
356+ this.#lam3 = lambda;
357+ this.gpu.setParams(this.#mergedParams());
358+ }
359+
360+ /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
361+ step(n = 1): void {
362+ this.gpu.step(n);
363+ this.t += n * (this.#params.dt ?? 0);
364+ this.steps += n;
365+ }
366+
367+ /**
368+ * Wait for the submitted steps to finish, without reading anything back.
369+ * This is the honest way to time the solver: a readback would add a GPU->CPU
370+ * round trip, which in a browser also crosses a process boundary and can cost
371+ * more than the steps themselves.
372+ */
373+ sync(): Promise<undefined> {
374+ return this.device.queue.onSubmittedWorkDone();
375+ }
376+
377+ /**
378+ * Time a batch of `n` steps and return ms/step, leaving the simulation
379+ * exactly where it was: the spectral state is snapshotted before the batch
380+ * and restored after, and `t`/`steps` do not advance. One sync amortized
381+ * over the batch — the same measurement the desktop benchmark makes. The
382+ * grid view fields hold the batch's output until the next real step, so
383+ * step before reading them.
384+ */
385+ async measure(n: number): Promise<number> {
386+ this.gpu.snapshotState();
387+ const t0 = performance.now();
388+ this.gpu.step(n);
389+ await this.sync();
390+ const ms = (performance.now() - t0) / n;
391+ this.gpu.restoreState();
392+ return ms;
393+ }
394+
395+ /** Read a named value (a grid field or the spectral state). */
396+ read(name: string): Promise<Float32Array> {
397+ return this.gpu.read(name);
398+ }
399+
400+ /**
401+ * Read every state species back to the CPU, in the shape `loadState`
402+ * consumes — the capture side of that method's reload, so a caller can
403+ * hold onto the current spectral state and restore it exactly later
404+ * (e.g. a "restart to this run's initial condition" control). One at a
405+ * time, not `Promise.all`: every `read()` copies into the same shared
406+ * readback buffer (`GpuModel#readback`, model.ts:448-452), so two in
407+ * flight at once race `mapAsync` against each other's `unmap`.
408+ */
409+ async readState(): Promise<Record<string, Float32Array>> {
410+ const out: Record<string, Float32Array> = {};
411+ for (const name of this.model.state) out[name] = await this.read(name);
412+ return out;
413+ }
414+
415+ /**
416+ * Read species `k` at render resolution (`viewSht`'s grid): the spectral
417+ * state synthesized there. The models define each species as synth of its
418+ * state, so this is the field the .m returned — evaluated exactly, whatever
419+ * the grid — and it is current however the state last changed, including a
420+ * `loadState`, which runs no kernel that would write the grid-space fields.
421+ */
422+ readSpecies(k: number): Promise<Float32Array> {
423+ const state = this.model.state[k];
424+ const buf = this.gpu.valueBuffer(state);
425+ if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
426+ return this.viewSht.synthFrom(buf);
427+ }
428+
429+ describe(): { init: string[]; step: string[] } {
430+ return this.gpu.describe();
431+ }
432+
433+ destroy(): void {
434+ this.gpu.destroy();
435+ this.#deriv.destroy();
436+ this.#displaySht?.destroy();
437+ this.sht.destroy();
438+ }
439+}
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+15−0View file
@@ -0,0 +1,15 @@
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+}
6+
7+/** Vite's `import.meta.glob`, used to load every .m in tools/ at once
8+ * (src/tools.ts). Only the eager + `?raw` form this project uses is
9+ * declared — it returns each match's text, keyed by path. */
10+interface ImportMeta {
11+ glob(
12+ pattern: string,
13+ options: { query: '?raw'; eager: true; import: 'default' },
14+ ): Record<string, string>;
15+}
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+90−0View file
@@ -0,0 +1,90 @@
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+/**
8+ * Smallest span the colormap may be stretched across, relative to the field's
9+ * own magnitude.
10+ *
11+ * Set from measurement, not taste. A constant field analysed and re-synthesized
12+ * in fp32 comes back constant only to
13+ *
14+ * lmax 63: 2.9e-5 relative lmax 127: 9.6e-5 lmax 255: 2.4e-4
15+ *
16+ * and the residue is not white noise — it is concentrated in a few rings at the
17+ * poles (74x the equatorial level at lmax 63, 1710x at lmax 255), because what
18+ * survives the analysis is high-degree m = 0 content whose Legendre functions
19+ * all peak at the poles *and add in phase there*. The same round trip in f64 is
20+ * 2.4e-8 and flat, so this is fp32, not the algorithm.
21+ *
22+ * A floor of 1e-2 puts the worst of that (about 5e-4 of span at lmax 255) into
23+ * roughly 5% of the colormap rather than all of it, while the variation these
24+ * models actually carry — a few percent of the field's magnitude and up — is
25+ * left alone entirely.
26+ */
27+const RANGE_FLOOR_REL = 1e-2;
28+/** And an absolute floor, for a field whose magnitude is itself near zero. */
29+const RANGE_FLOOR_ABS = 1e-9;
30+
31+/**
32+ * Widen a value range so that a field which is uniform to numerical precision
33+ * is drawn as uniform.
34+ *
35+ * Scaling the colormap to a field's own extremes gives full contrast to
36+ * whatever variation it has — including none. Schnakenberg's `v` at t = 0 is
37+ * literally constant (`vs * ones(...)`), so its extremes are set purely by the
38+ * roundoff described above; painting that across the whole colormap produces a
39+ * vivid pole-capped picture that reads as structure, and since the residue
40+ * belongs to the grid, two runs at different lmax produce two entirely
41+ * different pictures of the same constant — which looks exactly like a broken
42+ * initial condition, and is not one.
43+ *
44+ * A floor rather than an "is this field constant?" test, so nothing ever jumps:
45+ * a real pattern growing up through the floor hands the range over from the
46+ * floor to its own data gradually, and once it is any larger than roundoff the
47+ * floor has no effect at all.
48+ */
49+export function floorRange(lo: number, hi: number): { lo: number; hi: number } {
50+ const minSpan = Math.max(
51+ RANGE_FLOOR_ABS,
52+ RANGE_FLOOR_REL * Math.max(Math.abs(lo), Math.abs(hi)),
53+ );
54+ if (hi - lo >= minSpan) return { lo, hi };
55+ const mid = (lo + hi) / 2;
56+ return { lo: mid - minSpan / 2, hi: mid + minSpan / 2 };
57+}
58+
59+/** Vertical colorbar drawn on a small canvas, with min/max labels. */
60+export class Colorbar {
61+ #canvas: HTMLCanvasElement;
62+ #minLabel: HTMLElement;
63+ #maxLabel: HTMLElement;
64+
65+ constructor(container: HTMLElement) {
66+ container.classList.add('colorbar');
67+ this.#maxLabel = document.createElement('div');
68+ this.#maxLabel.className = 'colorbar-label';
69+ this.#canvas = document.createElement('canvas');
70+ this.#canvas.width = 12;
71+ this.#canvas.height = 160;
72+ this.#minLabel = document.createElement('div');
73+ this.#minLabel.className = 'colorbar-label';
74+ container.append(this.#maxLabel, this.#canvas, this.#minLabel);
75+ }
76+
77+ update(cmap: ColormapFunc, vmin: number, vmax: number): void {
78+ const ctx = this.#canvas.getContext('2d');
79+ if (!ctx) return;
80+ const h = this.#canvas.height;
81+ for (let y = 0; y < h; y++) {
82+ const t = 1 - y / (h - 1);
83+ const [r, g, b] = cmap(t);
84+ ctx.fillStyle = `rgb(${r},${g},${b})`;
85+ ctx.fillRect(0, y, this.#canvas.width, 1);
86+ }
87+ this.#maxLabel.textContent = fmtValue(vmax);
88+ this.#minLabel.textContent = fmtValue(vmin);
89+ }
90+}
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/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/deriv.tsadded+290−0View file
@@ -0,0 +1,290 @@
1+/**
2+ * First derivatives of a scalar field, coefficients -> grid: dtheta and dphi
3+ * (evolving_surface/notes/algos.tex Algorithm 1, theta/phi branches only --
4+ * the Laplace-Beltrami operator built on these never needs the second-
5+ * derivative/curvature branches, so they are not ported).
6+ *
7+ * Both derivatives start with a shuffle in coefficient space (the theta
8+ * branch's +-1 index gather via the alpha recurrence, the phi branch's i*m
9+ * row-swap) and then reuse the *existing* Legendre+Fourier synthesis
10+ * pipeline (ShtPlan.createSynthBinding/encodeSynthInto) unchanged -- neither
11+ * derivative touches the Legendre recurrence stage itself. dtheta
12+ * additionally divides by sin(theta) on the grid afterwards.
13+ *
14+ * The two shuffles are also exposed on their own, coefficients -> coefficients
15+ * (`dthetac`, `dphic`), because the six-transform Laplace-Beltrami operator of
16+ * docs/reduced-transforms.md needs them apart from a
17+ * synthesis, and needs them twice: once on the field (steps 1-2) and once on
18+ * the two fluxes (step 5, which is the *same* alpha^+/alpha^- gather, not its
19+ * transpose). Everything above is then a composition of them:
20+ *
21+ * dphi(U) == synth(dphic(U))
22+ * dtheta(U) == synth(dthetac(U)) / sin(theta)
23+ */
24+import type { ShtPlan, ShtBinding } from './sht.ts';
25+import { derivCoeffs } from './derivCoeffs.ts';
26+import { dthetaShuffleWGSL, dphiShuffleWGSL, divideSinThetaWGSL } from './wgsl/deriv.ts';
27+
28+const WG = 64;
29+
30+async function makePipeline(
31+ device: GPUDevice,
32+ code: string,
33+ entryPoint: string,
34+): Promise<GPUComputePipeline> {
35+ device.pushErrorScope('validation');
36+ const module = device.createShaderModule({ code, label: entryPoint });
37+ const info = await module.getCompilationInfo();
38+ const errors = info.messages.filter((m) => m.type === 'error');
39+ if (errors.length) {
40+ throw new Error(
41+ `WGSL compile error in ${entryPoint}:\n` +
42+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
43+ );
44+ }
45+ const pipeline = await device.createComputePipelineAsync({
46+ layout: 'auto',
47+ compute: { module, entryPoint },
48+ label: entryPoint,
49+ });
50+ const err = await device.popErrorScope();
51+ if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
52+ return pipeline;
53+}
54+
55+/** Bindings for one dtheta/dphi call against caller-supplied buffers. */
56+export interface DerivBinding {
57+ readonly shuffle: GPUBindGroup;
58+ readonly sht: ShtBinding;
59+ /** Only present for dtheta: the post-synthesis divide by sin(theta). */
60+ readonly divide?: GPUBindGroup;
61+}
62+
63+export class DerivPlan {
64+ private device: GPUDevice;
65+ private sht: ShtPlan;
66+ private nlm: number;
67+ private npts: number;
68+
69+ private bufAPlus!: GPUBuffer;
70+ private bufAMinus!: GPUBuffer;
71+ private bufMOf!: GPUBuffer;
72+ private bufSinTheta!: GPUBuffer;
73+ /** Scratch coefficient buffer for the shuffled input to synth -- shared
74+ * sequentially like ShtPlan's fmBuf, since ops within one pass execute
75+ * in submission order. */
76+ private scratch!: GPUBuffer;
77+
78+ private pipeDtheta!: GPUComputePipeline;
79+ private pipeDphi!: GPUComputePipeline;
80+ private pipeDivide!: GPUComputePipeline;
81+
82+ private constructor(device: GPUDevice, sht: ShtPlan) {
83+ this.device = device;
84+ this.sht = sht;
85+ this.nlm = sht.nlm;
86+ this.npts = sht.cfg.nlat * sht.cfg.nphi;
87+ }
88+
89+ static async create(device: GPUDevice, sht: ShtPlan): Promise<DerivPlan> {
90+ const plan = new DerivPlan(device, sht);
91+ await plan.init();
92+ return plan;
93+ }
94+
95+ private async init(): Promise<void> {
96+ const { nlat, nphi } = this.sht.cfg;
97+ const dev = this.device;
98+
99+ const { aPlus, aMinus, mOf } = derivCoeffs(this.sht.cfg.lmax, this.sht.cfg.mmax);
100+ const sinTheta = new Float32Array(nlat);
101+ for (let i = 0; i < nlat; i++) {
102+ const ct = this.sht.cosTheta[i];
103+ sinTheta[i] = Math.sqrt(Math.max(0, 1 - ct * ct));
104+ }
105+
106+ const mk = (label: string, size: number, usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST) =>
107+ dev.createBuffer({ label, size, usage });
108+ this.bufAPlus = mk('deriv-aplus', 4 * this.nlm);
109+ this.bufAMinus = mk('deriv-aminus', 4 * this.nlm);
110+ this.bufMOf = mk('deriv-mof', 4 * this.nlm);
111+ this.bufSinTheta = mk('deriv-sintheta', 4 * nlat);
112+ this.scratch = mk(
113+ 'deriv-scratch',
114+ 8 * this.nlm,
115+ GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
116+ );
117+
118+ dev.queue.writeBuffer(this.bufAPlus, 0, new Float32Array(aPlus));
119+ dev.queue.writeBuffer(this.bufAMinus, 0, new Float32Array(aMinus));
120+ dev.queue.writeBuffer(this.bufMOf, 0, mOf as Uint32Array<ArrayBuffer>);
121+ dev.queue.writeBuffer(this.bufSinTheta, 0, sinTheta);
122+
123+ const [pDtheta, pDphi, pDivide] = await Promise.all([
124+ makePipeline(dev, dthetaShuffleWGSL({ nlm: this.nlm }), 'dtheta_shuffle'),
125+ makePipeline(dev, dphiShuffleWGSL({ nlm: this.nlm }), 'dphi_shuffle'),
126+ makePipeline(dev, divideSinThetaWGSL({ nlat, nphi }), 'divide_sin_theta'),
127+ ]);
128+ this.pipeDtheta = pDtheta;
129+ this.pipeDphi = pDphi;
130+ this.pipeDivide = pDivide;
131+ }
132+
133+ /**
134+ * Bind group for the coefficient-space half of dtheta on its own:
135+ * v_l^m = alpha^+(l-1,m) u_{l-1}^m + alpha^-(l+1,m) u_{l+1}^m, the
136+ * coefficients of sin(theta) * dtheta(u). Input and output must be
137+ * different buffers -- WebGPU forbids binding one buffer as both readable
138+ * and writable storage in a dispatch, and the gather reads l+-1 anyway.
139+ */
140+ createDthetacBinding(qlmIn: GPUBuffer, qlmOut: GPUBuffer): GPUBindGroup {
141+ return this.device.createBindGroup({
142+ layout: this.pipeDtheta.getBindGroupLayout(0),
143+ entries: [
144+ { binding: 0, resource: { buffer: this.bufAPlus } },
145+ { binding: 1, resource: { buffer: this.bufAMinus } },
146+ { binding: 2, resource: { buffer: qlmIn } },
147+ { binding: 3, resource: { buffer: qlmOut } },
148+ ],
149+ });
150+ }
151+
152+ /** Bind group for the coefficient-space half of dphi on its own:
153+ * (dphi u)_l^m = i*m*u_l^m. Same buffer restriction as dthetac. */
154+ createDphicBinding(qlmIn: GPUBuffer, qlmOut: GPUBuffer): GPUBindGroup {
155+ return this.device.createBindGroup({
156+ layout: this.pipeDphi.getBindGroupLayout(0),
157+ entries: [
158+ { binding: 0, resource: { buffer: this.bufMOf } },
159+ { binding: 1, resource: { buffer: qlmIn } },
160+ { binding: 2, resource: { buffer: qlmOut } },
161+ ],
162+ });
163+ }
164+
165+ /** Record the bare alpha^+/alpha^- shift into an existing compute pass. */
166+ encodeDthetacInto(pass: GPUComputePassEncoder, bindGroup: GPUBindGroup): void {
167+ pass.setPipeline(this.pipeDtheta);
168+ pass.setBindGroup(0, bindGroup);
169+ pass.dispatchWorkgroups(Math.ceil(this.nlm / WG));
170+ }
171+
172+ /** Record the bare i*m multiply into an existing compute pass. */
173+ encodeDphicInto(pass: GPUComputePassEncoder, bindGroup: GPUBindGroup): void {
174+ pass.setPipeline(this.pipeDphi);
175+ pass.setBindGroup(0, bindGroup);
176+ pass.dispatchWorkgroups(Math.ceil(this.nlm / WG));
177+ }
178+
179+ /** Bindings for dtheta(qlmIn) -> spatOut, against caller-owned buffers. */
180+ createDthetaBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): DerivBinding {
181+ const shuffle = this.createDthetacBinding(qlmIn, this.scratch);
182+ const sht = this.sht.createSynthBinding(this.scratch, spatOut);
183+ const divide = this.device.createBindGroup({
184+ layout: this.pipeDivide.getBindGroupLayout(0),
185+ entries: [
186+ { binding: 0, resource: { buffer: this.bufSinTheta } },
187+ { binding: 1, resource: { buffer: spatOut } },
188+ ],
189+ });
190+ return { shuffle, sht, divide };
191+ }
192+
193+ /** Bindings for dphi(qlmIn) -> spatOut, against caller-owned buffers. */
194+ createDphiBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): DerivBinding {
195+ const shuffle = this.createDphicBinding(qlmIn, this.scratch);
196+ const sht = this.sht.createSynthBinding(this.scratch, spatOut);
197+ return { shuffle, sht };
198+ }
199+
200+ /** Record dtheta into an existing compute pass. */
201+ encodeDthetaInto(pass: GPUComputePassEncoder, b: DerivBinding): void {
202+ this.encodeSinDthetaInto(pass, b);
203+ pass.setPipeline(this.pipeDivide);
204+ pass.setBindGroup(0, b.divide!);
205+ pass.dispatchWorkgroups(Math.ceil(this.npts / WG));
206+ }
207+
208+ /** Record dtheta *without* its final division: the grid values of
209+ * sin(theta) * dtheta(u), which unlike dtheta(u) itself is a smooth
210+ * function on the sphere. Takes a dtheta binding and simply stops early. */
211+ encodeSinDthetaInto(pass: GPUComputePassEncoder, b: DerivBinding): void {
212+ this.encodeDthetacInto(pass, b.shuffle);
213+ this.sht.encodeSynthInto(pass, b.sht);
214+ }
215+
216+ /** Record dphi into an existing compute pass. */
217+ encodeDphiInto(pass: GPUComputePassEncoder, b: DerivBinding): void {
218+ this.encodeDphicInto(pass, b.shuffle);
219+ this.sht.encodeSynthInto(pass, b.sht);
220+ }
221+
222+ /** CPU convenience: qlm (interleaved [re,im], length 2*nlm) -> grid field. */
223+ async dtheta(qlm: Float32Array): Promise<Float32Array> {
224+ return this.#runToGrid(qlm, 'dtheta');
225+ }
226+
227+ /** CPU convenience: sin(theta) * dtheta(u) on the grid, the undivided
228+ * synthesis of the alpha shift. What the flux-form metric precompute
229+ * (src/geom/metric.ts) is built from. */
230+ async sinDtheta(qlm: Float32Array): Promise<Float32Array> {
231+ return this.#runToGrid(qlm, 'sinDtheta');
232+ }
233+
234+ /** CPU convenience: qlm (interleaved [re,im], length 2*nlm) -> grid field. */
235+ async dphi(qlm: Float32Array): Promise<Float32Array> {
236+ return this.#runToGrid(qlm, 'dphi');
237+ }
238+
239+ async #runToGrid(
240+ qlm: Float32Array,
241+ mode: 'dtheta' | 'sinDtheta' | 'dphi',
242+ ): Promise<Float32Array> {
243+ if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
244+ const dev = this.device;
245+ const qlmIn = dev.createBuffer({
246+ label: 'deriv-qlm-in',
247+ size: 8 * this.nlm,
248+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
249+ });
250+ const spatOut = dev.createBuffer({
251+ label: 'deriv-spat-out',
252+ size: 4 * this.npts,
253+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
254+ });
255+ const stage = dev.createBuffer({
256+ label: 'deriv-stage',
257+ size: 4 * this.npts,
258+ usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
259+ });
260+ try {
261+ dev.queue.writeBuffer(qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
262+ const binding =
263+ mode === 'dphi'
264+ ? this.createDphiBinding(qlmIn, spatOut)
265+ : this.createDthetaBinding(qlmIn, spatOut);
266+ const enc = dev.createCommandEncoder({ label: 'deriv-run' });
267+ const pass = enc.beginComputePass({ label: 'deriv-run' });
268+ if (mode === 'dtheta') this.encodeDthetaInto(pass, binding);
269+ else if (mode === 'sinDtheta') this.encodeSinDthetaInto(pass, binding);
270+ else this.encodeDphiInto(pass, binding);
271+ pass.end();
272+ enc.copyBufferToBuffer(spatOut, 0, stage, 0, 4 * this.npts);
273+ dev.queue.submit([enc.finish()]);
274+ await stage.mapAsync(GPUMapMode.READ);
275+ const out = new Float32Array(stage.getMappedRange().slice(0));
276+ stage.unmap();
277+ return out;
278+ } finally {
279+ qlmIn.destroy();
280+ spatOut.destroy();
281+ stage.destroy();
282+ }
283+ }
284+
285+ destroy(): void {
286+ for (const b of [
287+ this.bufAPlus, this.bufAMinus, this.bufMOf, this.bufSinTheta, this.scratch,
288+ ]) b?.destroy();
289+ }
290+}
src/sht/derivCoeffs.tsadded+49−0View file
@@ -0,0 +1,49 @@
1+/**
2+ * Recurrence coefficients for the first theta-derivative of orthonormal
3+ * associated Legendre functions (Condon-Shortley phase included), matching
4+ * the alpha^+/alpha^- recurrence in evolving_surface/notes/algos.tex Sec 2.1:
5+ *
6+ * sin(theta) d/dtheta Y_l^m = alpha^+(l,m) Y_{l+1}^m + alpha^-(l,m) Y_{l-1}^m
7+ *
8+ * so the coefficients of sin(theta)*dtheta(u), by degree, are
9+ *
10+ * v_l^m = alpha^+(l-1,m) u_{l-1}^m + alpha^-(l+1,m) u_{l+1}^m
11+ *
12+ * dropping any term referring to a degree outside 0 <= l <= lmax. Baked to
13+ * zero at each m-block's first/last element (rather than left undefined), so
14+ * a consuming WGSL kernel needs only an in-bounds check, not a validity check.
15+ */
16+import { lmIndex, nlmCalc } from './layout.ts';
17+
18+export interface DerivCoeffs {
19+ /** aPlus[lm] = alpha^+(l-1,m) when l>m, else 0 -- multiplies u_{l-1}^m. */
20+ aPlus: Float64Array;
21+ /** aMinus[lm] = alpha^-(l+1,m) when l<lmax, else 0 -- multiplies u_{l+1}^m. */
22+ aMinus: Float64Array;
23+ /** m of the coefficient at flat index lm (the phi-derivative needs only this). */
24+ mOf: Uint32Array;
25+}
26+
27+export function alphaPlus(l: number, m: number): number {
28+ return l * Math.sqrt(((l - m + 1) * (l + m + 1)) / ((2 * l + 1) * (2 * l + 3)));
29+}
30+
31+export function alphaMinus(l: number, m: number): number {
32+ return -(l + 1) * Math.sqrt(((l - m) * (l + m)) / ((2 * l - 1) * (2 * l + 1)));
33+}
34+
35+export function derivCoeffs(lmax: number, mmax: number): DerivCoeffs {
36+ const nlm = nlmCalc(lmax, mmax);
37+ const aPlus = new Float64Array(nlm);
38+ const aMinus = new Float64Array(nlm);
39+ const mOf = new Uint32Array(nlm);
40+ for (let m = 0; m <= mmax; m++) {
41+ for (let l = m; l <= lmax; l++) {
42+ const lm = lmIndex(lmax, l, m);
43+ mOf[lm] = m;
44+ if (l - 1 >= m) aPlus[lm] = alphaPlus(l - 1, m);
45+ if (l + 1 <= lmax) aMinus[lm] = alphaMinus(l + 1, m);
46+ }
47+ }
48+ return { aPlus, aMinus, mOf };
49+}
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+208−0View file
@@ -0,0 +1,208 @@
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 { alphaPlus, alphaMinus } from './derivCoeffs.ts';
12+import { lmIndex, nlmCalc, validateConfig, type ShtConfig } from './layout.ts';
13+
14+export class ShtReference {
15+ readonly cfg: ShtConfig;
16+ readonly nlm: number;
17+ readonly ct: Float64Array;
18+ readonly st: Float64Array;
19+ readonly wg: Float64Array; // Gauss weights (for integral over cos(theta))
20+ readonly coeffs: LegendreCoeffs;
21+
22+ constructor(cfg: ShtConfig) {
23+ validateConfig(cfg);
24+ this.cfg = cfg;
25+ this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
26+ const { x, w } = gaussNodesWeights(cfg.nlat);
27+ this.ct = x;
28+ this.wg = w;
29+ this.st = new Float64Array(cfg.nlat);
30+ for (let i = 0; i < cfg.nlat; i++) this.st[i] = Math.sqrt(1 - x[i] * x[i]);
31+ this.coeffs = legendreCoeffs(cfg.lmax, cfg.mmax);
32+ }
33+
34+ /**
35+ * Legendre stage of the synthesis: F_m(theta_i) = sum_l Q_lm ytilde_l^m(theta_i).
36+ * Returns complex array indexed [m * nlat + ilat], interleaved re/im.
37+ */
38+ legendreSynth(qlm: ArrayLike<number>): Float64Array {
39+ const { lmax, mmax, nlat } = this.cfg;
40+ const fm = new Float64Array(2 * (mmax + 1) * nlat);
41+ const row = new Float64Array(lmax + 1);
42+ for (let i = 0; i < nlat; i++) {
43+ for (let m = 0; m <= mmax; m++) {
44+ legendreRow(this.coeffs, lmax, m, this.ct[i], this.st[i], row);
45+ let re = 0, im = 0;
46+ const base = lmIndex(lmax, m, m);
47+ for (let l = m; l <= lmax; l++) {
48+ const y = row[l - m];
49+ re += y * qlm[2 * (base + l - m)];
50+ im += y * qlm[2 * (base + l - m) + 1];
51+ }
52+ const o = 2 * (m * nlat + i);
53+ fm[o] = re;
54+ fm[o + 1] = im;
55+ }
56+ }
57+ return fm;
58+ }
59+
60+ /** Full synthesis: spectral -> spatial grid [ilat * nphi + iphi]. */
61+ synth(qlm: ArrayLike<number>): Float64Array {
62+ const { mmax, nlat, nphi } = this.cfg;
63+ const fm = this.legendreSynth(qlm);
64+ const spat = new Float64Array(nlat * nphi);
65+ for (let i = 0; i < nlat; i++) {
66+ for (let j = 0; j < nphi; j++) {
67+ const phi = (2 * Math.PI * j) / nphi;
68+ let v = fm[2 * (0 * nlat + i)]; // m=0: real part (imag must be 0)
69+ for (let m = 1; m <= mmax; m++) {
70+ const o = 2 * (m * nlat + i);
71+ const c = Math.cos(m * phi);
72+ const s = Math.sin(m * phi);
73+ v += 2 * (fm[o] * c - fm[o + 1] * s);
74+ }
75+ spat[i * nphi + j] = v;
76+ }
77+ }
78+ return spat;
79+ }
80+
81+ /** Full analysis: spatial grid -> spectral coefficients (interleaved re/im). */
82+ analys(spat: ArrayLike<number>): Float64Array {
83+ const { lmax, mmax, nlat, nphi } = this.cfg;
84+ const qlm = new Float64Array(2 * this.nlm);
85+ const row = new Float64Array(lmax + 1);
86+ // forward Fourier: G_m(theta_i) = (2*pi/nphi) * sum_j f_ij e^{-i m phi_j}
87+ const gm = new Float64Array(2 * (mmax + 1) * nlat);
88+ for (let i = 0; i < nlat; i++) {
89+ for (let m = 0; m <= mmax; m++) {
90+ let re = 0, im = 0;
91+ for (let j = 0; j < nphi; j++) {
92+ const phi = (2 * Math.PI * j) / nphi;
93+ const f = spat[i * nphi + j];
94+ re += f * Math.cos(m * phi);
95+ im -= f * Math.sin(m * phi);
96+ }
97+ const o = 2 * (m * nlat + i);
98+ const norm = (2 * Math.PI) / nphi;
99+ gm[o] = re * norm;
100+ gm[o + 1] = im * norm;
101+ }
102+ }
103+ // Legendre stage with Gauss quadrature: Q_lm = sum_i w_i ytilde_l^m(theta_i) G_m(theta_i)
104+ for (let m = 0; m <= mmax; m++) {
105+ const base = lmIndex(lmax, m, m);
106+ for (let i = 0; i < nlat; i++) {
107+ legendreRow(this.coeffs, lmax, m, this.ct[i], this.st[i], row);
108+ const o = 2 * (m * nlat + i);
109+ const wr = this.wg[i] * gm[o];
110+ const wi = this.wg[i] * gm[o + 1];
111+ for (let l = m; l <= lmax; l++) {
112+ const y = row[l - m];
113+ qlm[2 * (base + l - m)] += y * wr;
114+ qlm[2 * (base + l - m) + 1] += y * wi;
115+ }
116+ }
117+ }
118+ return qlm;
119+ }
120+
121+ /**
122+ * Coefficient-space theta shift, f64: v_l^m = alpha^+(l-1,m) u_{l-1}^m +
123+ * alpha^-(l+1,m) u_{l+1}^m (algos.tex eq. v_coeffs) — the coefficients of
124+ * sin(theta) * dtheta(u). The same shift serves the divergence side of the
125+ * six-transform Laplace-Beltrami scheme (step 5 of
126+ * docs/reduced-transforms.md), which is why it is exposed
127+ * apart from the synthesis.
128+ */
129+ dthetac(qlm: ArrayLike<number>): Float64Array {
130+ const { lmax, mmax } = this.cfg;
131+ const v = new Float64Array(2 * this.nlm);
132+ for (let m = 0; m <= mmax; m++) {
133+ for (let l = m; l <= lmax; l++) {
134+ const lm = lmIndex(lmax, l, m);
135+ let re = 0;
136+ let im = 0;
137+ if (l - 1 >= m) {
138+ const lm1 = lmIndex(lmax, l - 1, m);
139+ const a = alphaPlus(l - 1, m);
140+ re += a * qlm[2 * lm1];
141+ im += a * qlm[2 * lm1 + 1];
142+ }
143+ if (l + 1 <= lmax) {
144+ const lm2 = lmIndex(lmax, l + 1, m);
145+ const a = alphaMinus(l + 1, m);
146+ re += a * qlm[2 * lm2];
147+ im += a * qlm[2 * lm2 + 1];
148+ }
149+ v[2 * lm] = re;
150+ v[2 * lm + 1] = im;
151+ }
152+ }
153+ return v;
154+ }
155+
156+ /** Coefficient-space phi derivative, f64: (dphi u)_l^m = i*m*u_l^m. */
157+ dphic(qlm: ArrayLike<number>): Float64Array {
158+ const { lmax, mmax } = this.cfg;
159+ const v = new Float64Array(2 * this.nlm);
160+ for (let m = 0; m <= mmax; m++) {
161+ for (let l = m; l <= lmax; l++) {
162+ const lm = lmIndex(lmax, l, m);
163+ v[2 * lm] = -m * qlm[2 * lm + 1];
164+ v[2 * lm + 1] = m * qlm[2 * lm];
165+ }
166+ }
167+ return v;
168+ }
169+
170+ /** sin(theta) * dtheta(u) on the grid, f64: the undivided synthesis of the
171+ * theta shift. Smooth on the sphere, unlike dtheta(u) itself. */
172+ sinDtheta(qlm: ArrayLike<number>): Float64Array {
173+ return this.synth(this.dthetac(qlm));
174+ }
175+
176+ /** Theta-derivative, f64: synth(dthetac(u)) / sin(theta). */
177+ dtheta(qlm: ArrayLike<number>): Float64Array {
178+ const { nlat, nphi } = this.cfg;
179+ const grid = this.sinDtheta(qlm);
180+ for (let i = 0; i < nlat; i++) {
181+ const st = this.st[i];
182+ for (let j = 0; j < nphi; j++) grid[i * nphi + j] /= st;
183+ }
184+ return grid;
185+ }
186+
187+ /** Phi-derivative, f64: (dphi u)_l^m = i*m*u_l^m, then synthesize. */
188+ dphi(qlm: ArrayLike<number>): Float64Array {
189+ return this.synth(this.dphic(qlm));
190+ }
191+}
192+
193+/** Random band-limited spectrum for testing (m=0 imaginary parts zeroed). */
194+export function randomSpectrum(cfg: ShtConfig, seed = 12345): Float32Array {
195+ const nlm = nlmCalc(cfg.lmax, cfg.mmax);
196+ const q = new Float32Array(2 * nlm);
197+ let s = seed >>> 0;
198+ const rnd = () => {
199+ // xorshift32
200+ s ^= s << 13; s >>>= 0;
201+ s ^= s >> 17;
202+ s ^= s << 5; s >>>= 0;
203+ return (s / 4294967296) * 2 - 1;
204+ };
205+ for (let k = 0; k < 2 * nlm; k++) q[k] = rnd();
206+ for (let l = 0; l <= cfg.lmax; l++) q[2 * lmIndex(cfg.lmax, l, 0) + 1] = 0; // m=0 real
207+ return q;
208+}
src/sht/sht.tsadded+795−0View file
@@ -0,0 +1,795 @@
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 {
13+ legSynthWGSL,
14+ legAnalysWGSL,
15+ legSynthBatchWGSL,
16+ legAnalysBatchWGSL,
17+} from './wgsl/leg.ts';
18+import { fmDphiWGSL } from './wgsl/deriv.ts';
19+import {
20+ fftSynthWGSL,
21+ fftAnalysWGSL,
22+ fftSynthRealWGSL,
23+ fftAnalysRealWGSL,
24+ dftSynthWGSL,
25+ dftAnalysWGSL,
26+ fftThreads,
27+} from './wgsl/fourier.ts';
28+
29+export type FourierMode = 'auto' | 'fft' | 'dft';
30+
31+/** The two bind groups (Legendre stage, Fourier stage) of one transform. */
32+export interface ShtBinding {
33+ readonly bgLeg: GPUBindGroup;
34+ readonly bgFour: GPUBindGroup;
35+}
36+
37+/** The three bind groups of one grid-space phi-derivative (see dphig). */
38+export interface ShtDphigBinding {
39+ readonly bgFourAnalys: GPUBindGroup;
40+ readonly bgMul: GPUBindGroup;
41+ readonly bgFourSynth: GPUBindGroup;
42+}
43+
44+/**
45+ * One batched transform: K fields through a single Legendre dispatch (the
46+ * recurrence walked once, K accumulator lanes) plus K per-field Fourier
47+ * dispatches — the Fourier stage shares nothing across fields, so batching
48+ * it would save only bind-group switches.
49+ */
50+export interface ShtBatchBinding {
51+ /** Lanes in this batch — selects the pipeline compiled for that width. */
52+ readonly size: number;
53+ readonly bgLeg: GPUBindGroup;
54+ /** Per-lane Fourier bind group, lane k against fm arena k. */
55+ readonly bgFour: GPUBindGroup[];
56+}
57+
58+const bgEntries = (bufs: GPUBuffer[]) =>
59+ bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
60+
61+export interface ShtOptions {
62+ /** Fourier stage implementation. 'auto' picks fft when nphi is a power of two that fits in workgroup memory. */
63+ fourier?: FourierMode;
64+}
65+
66+const WG_SYNTH = 64;
67+
68+/**
69+ * Tuning knob, for A/B-ing a change without editing code. Reads globalThis
70+ * first (set it before creating a plan, as scripts/_ab.ts does), then the
71+ * environment, so `SHT_SUBGROUPS=0 npm run bench:sht` works too. `process` is
72+ * absent in the browser, where only the globalThis form applies.
73+ */
74+function tuning(name: string): unknown {
75+ const g = (globalThis as Record<string, unknown>)[name];
76+ if (g !== undefined) return g;
77+ const env = (globalThis as { process?: { env?: Record<string, string | undefined> } }).process?.env?.[name];
78+ if (env === undefined || env === '') return undefined;
79+ if (env === '1' || env === 'true') return true;
80+ if (env === '0' || env === 'false') return false;
81+ const n = Number(env);
82+ return Number.isFinite(n) ? n : env;
83+}
84+
85+/**
86+ * Workgroup size for the analysis Legendre reduction. The right answer differs
87+ * between the two reduction strategies, so it is chosen per strategy.
88+ *
89+ * Measured on an RTX PRO 6000 Blackwell (analysis, us). Shared-memory tree,
90+ * where each doubling of wgAnalys costs another barrier per l-pair:
91+ *
92+ * wgAnalys: 16 32 64 128 256
93+ * nlat=128 43.2 40.6 42.6 46.7 52.9 -> 32
94+ * nlat=256 104.0 85.6 87.8 89.9 100.0 -> 32
95+ * nlat=512 408.8 216.9 184.4 190.8 204.8 -> 64
96+ *
97+ * i.e. max(32, nlat/8). A flat 32 would be worse than the old default of 256 at
98+ * nlat=512, so it cannot be fitted on one grid. With subgroupAdd the barrier
99+ * count stops growing with wgAnalys and the picture inverts: threads in flight,
100+ * (mmax+1) * wgAnalys, becomes binding, since analysis dispatches only mmax+1
101+ * workgroups. 128 then wins at every grid (round trip, us):
102+ *
103+ * 128x256 37.8 (vs 38.3), 256x512 66.6 (vs 74.7), 512x1024 131.9 (vs 156.6)
104+ */
105+function defaultWgAnalys(nlat: number, limit: number, subgroups: boolean): number {
106+ if (subgroups) {
107+ // capped at nlat so small grids do not launch threads with no latitude to own
108+ let cap = 1;
109+ while (cap < nlat) cap *= 2;
110+ return Math.min(128, limit, cap);
111+ }
112+ const target = Math.max(32, nlat / 8);
113+ let wg = 1;
114+ while (wg < target) wg *= 2; // the tree reduction halves, so a power of two
115+ return Math.min(wg, limit);
116+}
117+
118+async function makePipeline(
119+ device: GPUDevice,
120+ code: string,
121+ entryPoint: string,
122+): Promise<GPUComputePipeline> {
123+ device.pushErrorScope('validation');
124+ const module = device.createShaderModule({ code, label: entryPoint });
125+ const info = await module.getCompilationInfo();
126+ const errors = info.messages.filter((m) => m.type === 'error');
127+ if (errors.length) {
128+ throw new Error(
129+ `WGSL compile error in ${entryPoint}:\n` +
130+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
131+ );
132+ }
133+ const pipeline = await device.createComputePipelineAsync({
134+ layout: 'auto',
135+ compute: { module, entryPoint },
136+ label: entryPoint,
137+ });
138+ const err = await device.popErrorScope();
139+ if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
140+ return pipeline;
141+}
142+
143+export class ShtPlan {
144+ readonly cfg: ShtConfig;
145+ readonly nlm: number;
146+ readonly fourierMode: 'fft' | 'dft';
147+ /** Latitudes leg_synth walks: nlat/2 when parity folding. */
148+ readonly legLat: number = 0;
149+ /**
150+ * Widest transform batch this plan supports: the largest even K <= 4 whose
151+ * Legendre bind group (3 tables + K caller fields + the shared fm arena)
152+ * fits the device's storage-buffer limit. K = 4 needs exactly the WebGPU
153+ * default of 8, so batching is fully available on every stack; 1 (no
154+ * batching) if SHT_BATCH is turned off. Batched and scalar transforms
155+ * compute identical per-lane arithmetic, so this only affects speed,
156+ * never results.
157+ */
158+ readonly batchK: number = 1;
159+ /** Colatitudes theta_i (f64, increasing: north to south). */
160+ readonly theta: Float64Array;
161+ readonly cosTheta: Float64Array;
162+ readonly gaussWeights: Float64Array;
163+
164+ private device: GPUDevice;
165+ private bufAb!: GPUBuffer;
166+ private bufAmm!: GPUBuffer;
167+ private bufCtstw!: GPUBuffer;
168+ private bufTrig!: GPUBuffer;
169+ /** Spectral input (synthesis) — write with queue.writeBuffer or use synth(). */
170+ readonly qlmIn!: GPUBuffer;
171+ /** Spectral output (analysis). */
172+ readonly qlmOut!: GPUBuffer;
173+ /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. COPY_SRC so the
174+ * stage boundary is observable: a transform is Legendre-then-Fourier, and
175+ * scripts/diagnose-sht.ts tells the two apart by reading this. */
176+ readonly fmBuf!: GPUBuffer;
177+ /** Spatial field [ilat*nphi + iphi], f32. */
178+ readonly spatBuf!: GPUBuffer;
179+ private stageSpat!: GPUBuffer;
180+ private stageQ!: GPUBuffer;
181+
182+ private pipeLegSynth!: GPUComputePipeline;
183+ private pipeLegAnalys!: GPUComputePipeline;
184+ private pipeFourSynth!: GPUComputePipeline;
185+ private pipeFourAnalys!: GPUComputePipeline;
186+ /** Fourier-space i*m multiply, the middle of dphig. */
187+ private pipeFmDphi!: GPUComputePipeline;
188+ /** Batched Legendre pipelines by lane count (even sizes up to batchK). */
189+ private pipeLegSynthB = new Map<number, GPUComputePipeline>();
190+ private pipeLegAnalysB = new Map<number, GPUComputePipeline>();
191+ /** One fm arena for all batch lanes (lane k at byte offset k * fmLaneBytes,
192+ * 256-aligned so the Fourier stage can bind a lane by buffer offset). A
193+ * single buffer keeps the batched Legendre bind group at 3 tables +
194+ * K fields + 1 arena — within WebGPU's default storage-buffer limit of 8
195+ * at K = 4, on every stack. */
196+ private fmArena: GPUBuffer | null = null;
197+ private fmLaneBytes = 0;
198+ private bgLegSynth!: GPUBindGroup;
199+ private bgLegAnalys!: GPUBindGroup;
200+ private bgFourSynth!: GPUBindGroup;
201+ private bgFourAnalys!: GPUBindGroup;
202+
203+ private constructor(device: GPUDevice, cfg: ShtConfig, fourierMode: 'fft' | 'dft') {
204+ this.device = device;
205+ this.cfg = cfg;
206+ this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
207+ this.fourierMode = fourierMode;
208+ const { x, w } = gaussNodesWeights(cfg.nlat);
209+ this.cosTheta = x;
210+ this.gaussWeights = w;
211+ this.theta = new Float64Array(cfg.nlat);
212+ for (let i = 0; i < cfg.nlat; i++) this.theta[i] = Math.acos(x[i]);
213+ }
214+
215+ static async create(device: GPUDevice, cfg: ShtConfig, opts: ShtOptions = {}): Promise<ShtPlan> {
216+ validateConfig(cfg);
217+ const want = opts.fourier ?? 'auto';
218+ const fftFits =
219+ isPowerOfTwo(cfg.nphi) &&
220+ 16 * cfg.nphi <= device.limits.maxComputeWorkgroupStorageSize &&
221+ fftThreads(cfg.nphi) <= device.limits.maxComputeInvocationsPerWorkgroup;
222+ if (want === 'fft' && !fftFits) {
223+ throw new Error(
224+ `fourier:'fft' requires power-of-two nphi with 16*nphi <= maxComputeWorkgroupStorageSize ` +
225+ `(nphi=${cfg.nphi}, limit=${device.limits.maxComputeWorkgroupStorageSize})`,
226+ );
227+ }
228+ const mode: 'fft' | 'dft' = want === 'dft' ? 'dft' : fftFits ? 'fft' : 'dft';
229+ const plan = new ShtPlan(device, cfg, mode);
230+ await plan.init();
231+ return plan;
232+ }
233+
234+ private async init(): Promise<void> {
235+ const { lmax, mmax, nlat, nphi } = this.cfg;
236+ const dev = this.device;
237+ const self = this as {
238+ -readonly [k in keyof ShtPlan]: ShtPlan[k];
239+ };
240+
241+ // --- host precomputation (f64), then downcast to f32 for upload ---
242+ const { amm, ab } = legendreCoeffs(lmax, mmax);
243+ const ctstw = new Float32Array(3 * nlat);
244+ for (let i = 0; i < nlat; i++) {
245+ ctstw[i] = this.cosTheta[i];
246+ ctstw[nlat + i] = Math.sqrt(1 - this.cosTheta[i] * this.cosTheta[i]);
247+ ctstw[2 * nlat + i] = this.gaussWeights[i] * ((2 * Math.PI) / nphi);
248+ }
249+ // twiddle/phase table in f64 (device sin/cos is too inaccurate: ~2^-11 under Vulkan)
250+ const trig = new Float32Array(2 * nphi);
251+ for (let k = 0; k < nphi; k++) {
252+ trig[2 * k] = Math.cos((2 * Math.PI * k) / nphi);
253+ trig[2 * k + 1] = Math.sin((2 * Math.PI * k) / nphi);
254+ }
255+
256+ const mkBuf = (label: string, size: number, usage: GPUBufferUsageFlags) =>
257+ dev.createBuffer({ label, size, usage });
258+ this.bufAb = mkBuf('sht-ab', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
259+ this.bufAmm = mkBuf('sht-amm', 4 * (mmax + 1), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
260+ this.bufCtstw = mkBuf('sht-ctstw', 4 * 3 * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
261+ this.bufTrig = mkBuf('sht-trig', 8 * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
262+ self.qlmIn = mkBuf('sht-qlm-in', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
263+ self.qlmOut = mkBuf('sht-qlm-out', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
264+ self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
265+ self.spatBuf = mkBuf('sht-spat', 4 * nlat * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
266+ this.stageSpat = mkBuf('sht-stage-spat', 4 * nlat * nphi, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
267+ this.stageQ = mkBuf('sht-stage-q', 8 * this.nlm, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
268+
269+ dev.queue.writeBuffer(this.bufAb, 0, new Float32Array(ab));
270+ dev.queue.writeBuffer(this.bufAmm, 0, new Float32Array(amm));
271+ dev.queue.writeBuffer(this.bufCtstw, 0, ctstw);
272+ dev.queue.writeBuffer(this.bufTrig, 0, trig);
273+
274+ // --- shaders / pipelines ---
275+ const subgroups = tuning('SHT_SUBGROUPS') !== false && dev.features.has('subgroups');
276+ // parity folding needs an equator-symmetric grid; Gauss nodes are, if nlat is even
277+ const parity = tuning('SHT_PARITY') !== false && nlat % 2 === 0;
278+ const wgAnalys =
279+ (tuning('SHT_WG_ANALYS') as number | undefined) ??
280+ defaultWgAnalys(nlat, dev.limits.maxComputeInvocationsPerWorkgroup, subgroups);
281+ const legP = {
282+ lmax,
283+ mmax,
284+ nlat,
285+ wgSynth: WG_SYNTH,
286+ wgAnalys,
287+ subgroups,
288+ spanPairs: tuning('SHT_SPAN_PAIRS') as number | undefined,
289+ parity,
290+ };
291+ (this as { legLat: number }).legLat = parity ? nlat / 2 : nlat;
292+ const fourP = { mmax, nlat, nphi, radix: (tuning('SHT_RADIX') as number | undefined) ?? 4 };
293+ // The spatial field is real (layout.ts stores m >= 0 only), so the Fourier
294+ // stage can run an nphi/2-point complex FFT plus a recombination instead of
295+ // a full nphi-point one: half the arithmetic and half the workgroup storage.
296+ // The complex kernels remain for a future complex-valued field, and are what
297+ // SHT_REAL_FFT=0 selects.
298+ const realFft =
299+ this.fourierMode === 'fft' && nphi % 2 === 0 && tuning('SHT_REAL_FFT') !== false;
300+ const fftS = realFft ? fftSynthRealWGSL : fftSynthWGSL;
301+ const fftA = realFft ? fftAnalysRealWGSL : fftAnalysWGSL;
302+ const [pLegS, pLegA, pFourS, pFourA, pFmDphi] = await Promise.all([
303+ makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),
304+ makePipeline(dev, legAnalysWGSL(legP), 'leg_analys'),
305+ makePipeline(
306+ dev,
307+ this.fourierMode === 'fft' ? fftS(fourP) : dftSynthWGSL(fourP),
308+ this.fourierMode === 'fft' ? 'fft_synth' : 'dft_synth',
309+ ),
310+ makePipeline(
311+ dev,
312+ this.fourierMode === 'fft' ? fftA(fourP) : dftAnalysWGSL(fourP),
313+ this.fourierMode === 'fft' ? 'fft_analys' : 'dft_analys',
314+ ),
315+ // Mirrors filterMask: content at l >= lmax-2 is filtered on the
316+ // l-space route, so the m-space route keeps m <= lmax-3.
317+ makePipeline(
318+ dev,
319+ fmDphiWGSL({ mmax, nlat, nphi, mcut: lmax - 3 }),
320+ 'fm_dphi',
321+ ),
322+ ]);
323+ this.pipeLegSynth = pLegS;
324+ this.pipeLegAnalys = pLegA;
325+ this.pipeFourSynth = pFourS;
326+ this.pipeFourAnalys = pFourA;
327+ this.pipeFmDphi = pFmDphi;
328+
329+ // --- batched Legendre pipelines ---
330+ // The widest even K <= 4 whose bind group (3 tables + K fields + the fm
331+ // arena) fits the device's storage-buffer budget — K = 4 needs 8, the
332+ // WebGPU default, so batching is fully available everywhere unless
333+ // SHT_BATCH=0 disables it (SHT_BATCH=2 caps it, for A/B).
334+ const batchTuning = tuning('SHT_BATCH');
335+ const batchWant =
336+ batchTuning === false || batchTuning === 0
337+ ? 1
338+ : typeof batchTuning === 'number'
339+ ? batchTuning
340+ : 4;
341+ const batchFit = dev.limits.maxStorageBuffersPerShaderStage - 4;
342+ const batchK = Math.min(4, Math.max(1, batchWant), 2 * Math.floor(batchFit / 2));
343+ (this as { batchK: number }).batchK = batchK;
344+ if (batchK >= 2) {
345+ // Lane stride rounded to the 256-byte offset alignment buffer bindings
346+ // require; laneElems is that stride in vec2f units for the kernels.
347+ this.fmLaneBytes = Math.ceil((8 * (mmax + 1) * nlat) / 256) * 256;
348+ const laneElems = this.fmLaneBytes / 8;
349+ this.fmArena = mkBuf('sht-fm-arena', batchK * this.fmLaneBytes, GPUBufferUsage.STORAGE);
350+ const sizes = [];
351+ for (let k = 2; k <= batchK; k += 2) sizes.push(k);
352+ const pipes = await Promise.all(
353+ sizes.flatMap((k) => [
354+ makePipeline(dev, legSynthBatchWGSL(legP, k, laneElems), `leg_synth_batch`),
355+ makePipeline(dev, legAnalysBatchWGSL(legP, k, laneElems), `leg_analys_batch`),
356+ ]),
357+ );
358+ sizes.forEach((k, i) => {
359+ this.pipeLegSynthB.set(k, pipes[2 * i]);
360+ this.pipeLegAnalysB.set(k, pipes[2 * i + 1]);
361+ });
362+ }
363+
364+ const entries = bgEntries;
365+ this.bgLegSynth = dev.createBindGroup({
366+ layout: pLegS.getBindGroupLayout(0),
367+ entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.qlmIn, this.fmBuf]),
368+ });
369+ this.bgLegAnalys = dev.createBindGroup({
370+ layout: pLegA.getBindGroupLayout(0),
371+ entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, this.qlmOut]),
372+ });
373+ this.bgFourSynth = dev.createBindGroup({
374+ layout: pFourS.getBindGroupLayout(0),
375+ entries: entries([this.fmBuf, this.spatBuf, this.bufTrig]),
376+ });
377+ this.bgFourAnalys = dev.createBindGroup({
378+ layout: pFourA.getBindGroupLayout(0),
379+ entries: entries([this.spatBuf, this.fmBuf, this.bufTrig]),
380+ });
381+ }
382+
383+ /**
384+ * Bind groups for one transform against caller-supplied spectral/spatial
385+ * buffers, so a transform can read and write buffers it does not own (the
386+ * .m-driven executor keeps a buffer per IR variable). Build these once at
387+ * plan time, not per step. `fmBuf` stays internal scratch: passes and
388+ * dispatches within a submission execute in order, so sequential transforms
389+ * can share it.
390+ */
391+ createSynthBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): ShtBinding {
392+ return {
393+ bgLeg: this.device.createBindGroup({
394+ layout: this.pipeLegSynth.getBindGroupLayout(0),
395+ entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, qlmIn, this.fmBuf]),
396+ }),
397+ bgFour: this.device.createBindGroup({
398+ layout: this.pipeFourSynth.getBindGroupLayout(0),
399+ entries: bgEntries([this.fmBuf, spatOut, this.bufTrig]),
400+ }),
401+ };
402+ }
403+
404+ createAnalysBinding(spatIn: GPUBuffer, qlmOut: GPUBuffer): ShtBinding {
405+ return {
406+ bgFour: this.device.createBindGroup({
407+ layout: this.pipeFourAnalys.getBindGroupLayout(0),
408+ entries: bgEntries([spatIn, this.fmBuf, this.bufTrig]),
409+ }),
410+ bgLeg: this.device.createBindGroup({
411+ layout: this.pipeLegAnalys.getBindGroupLayout(0),
412+ entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, qlmOut]),
413+ }),
414+ };
415+ }
416+
417+ /**
418+ * Bind groups for one batched synthesis: members.length must be a compiled
419+ * lane count (an even size <= batchK). Member outputs must be distinct
420+ * buffers; each lane gets its own fm arena, so batches compose in a pass
421+ * exactly like sequential scalar transforms do.
422+ */
423+ /** The fm arena sliced at lane k, sized as one transform's fm. */
424+ #fmLane(k: number): GPUBufferBinding {
425+ const { mmax, nlat } = this.cfg;
426+ return {
427+ buffer: this.fmArena!,
428+ offset: k * this.fmLaneBytes,
429+ size: 8 * (mmax + 1) * nlat,
430+ };
431+ }
432+
433+ createSynthBatchBinding(
434+ members: { qlmIn: GPUBuffer; spatOut: GPUBuffer }[],
435+ ): ShtBatchBinding {
436+ const K = members.length;
437+ const pipe = this.pipeLegSynthB.get(K);
438+ if (!pipe) throw new Error(`no batched synthesis pipeline for ${K} lanes`);
439+ return {
440+ size: K,
441+ bgLeg: this.device.createBindGroup({
442+ layout: pipe.getBindGroupLayout(0),
443+ entries: [
444+ ...bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, ...members.map((m) => m.qlmIn)]),
445+ { binding: 3 + K, resource: { buffer: this.fmArena! } },
446+ ],
447+ }),
448+ bgFour: members.map((m, k) =>
449+ this.device.createBindGroup({
450+ layout: this.pipeFourSynth.getBindGroupLayout(0),
451+ entries: [
452+ { binding: 0, resource: this.#fmLane(k) },
453+ { binding: 1, resource: { buffer: m.spatOut } },
454+ { binding: 2, resource: { buffer: this.bufTrig } },
455+ ],
456+ }),
457+ ),
458+ };
459+ }
460+
461+ createAnalysBatchBinding(
462+ members: { spatIn: GPUBuffer; qlmOut: GPUBuffer }[],
463+ ): ShtBatchBinding {
464+ const K = members.length;
465+ const pipe = this.pipeLegAnalysB.get(K);
466+ if (!pipe) throw new Error(`no batched analysis pipeline for ${K} lanes`);
467+ return {
468+ size: K,
469+ bgFour: members.map((m, k) =>
470+ this.device.createBindGroup({
471+ layout: this.pipeFourAnalys.getBindGroupLayout(0),
472+ entries: [
473+ { binding: 0, resource: { buffer: m.spatIn } },
474+ { binding: 1, resource: this.#fmLane(k) },
475+ { binding: 2, resource: { buffer: this.bufTrig } },
476+ ],
477+ }),
478+ ),
479+ bgLeg: this.device.createBindGroup({
480+ layout: pipe.getBindGroupLayout(0),
481+ entries: [
482+ ...bgEntries([this.bufAb, this.bufAmm, this.bufCtstw]),
483+ { binding: 3, resource: { buffer: this.fmArena! } },
484+ ...members.map((m, k) => ({
485+ binding: 4 + k,
486+ resource: { buffer: m.qlmOut },
487+ })),
488+ ],
489+ }),
490+ };
491+ }
492+
493+ /**
494+ * Bind groups for one grid-space phi-derivative, dphig: Fourier analysis
495+ * of each latitude row into fm (which truncates to m <= mmax for free),
496+ * the i*m/NPHI multiply (zeroing m past the top-degree filt's reach), and
497+ * Fourier synthesis back to the grid. No Legendre stage anywhere — this
498+ * is what lets the flux-form divergence drop the Q-flux's spherical-
499+ * harmonic analysis (docs/reduced-transforms.md Sec 5b's companion trick
500+ * in Sec 6-of-changes): d/dphi is diagonal in the Fourier index. Uses
501+ * fmBuf as scratch, sequentially like every transform in a pass.
502+ */
503+ createDphigBinding(spatIn: GPUBuffer, spatOut: GPUBuffer): ShtDphigBinding {
504+ return {
505+ bgFourAnalys: this.device.createBindGroup({
506+ layout: this.pipeFourAnalys.getBindGroupLayout(0),
507+ entries: bgEntries([spatIn, this.fmBuf, this.bufTrig]),
508+ }),
509+ bgMul: this.device.createBindGroup({
510+ layout: this.pipeFmDphi.getBindGroupLayout(0),
511+ entries: bgEntries([this.fmBuf]),
512+ }),
513+ bgFourSynth: this.device.createBindGroup({
514+ layout: this.pipeFourSynth.getBindGroupLayout(0),
515+ entries: bgEntries([this.fmBuf, spatOut, this.bufTrig]),
516+ }),
517+ };
518+ }
519+
520+ /** Record dphig into an existing compute pass: two Fourier stages and a
521+ * pointwise multiply — no Legendre work. */
522+ encodeDphigInto(pass: GPUComputePassEncoder, b: ShtDphigBinding): void {
523+ const { mmax, nlat, nphi } = this.cfg;
524+ pass.setPipeline(this.pipeFourAnalys);
525+ pass.setBindGroup(0, b.bgFourAnalys);
526+ if (this.fourierMode === 'fft') {
527+ pass.dispatchWorkgroups(nlat);
528+ } else {
529+ pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
530+ }
531+ pass.setPipeline(this.pipeFmDphi);
532+ pass.setBindGroup(0, b.bgMul);
533+ pass.dispatchWorkgroups(Math.ceil(((mmax + 1) * nlat) / 64));
534+ pass.setPipeline(this.pipeFourSynth);
535+ pass.setBindGroup(0, b.bgFourSynth);
536+ if (this.fourierMode === 'fft') {
537+ pass.dispatchWorkgroups(nlat);
538+ } else {
539+ pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
540+ }
541+ }
542+
543+ /** CPU convenience: grid field -> d/dphi of its trig interpolant, for tests. */
544+ async dphig(spat: Float32Array): Promise<Float32Array> {
545+ const { nlat, nphi } = this.cfg;
546+ if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
547+ this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
548+ const binding = this.createDphigBinding(this.spatBuf, this.spatBuf);
549+ const enc = this.device.createCommandEncoder({ label: 'sht-dphig' });
550+ const pass = enc.beginComputePass({ label: 'sht-dphig' });
551+ this.encodeDphigInto(pass, binding);
552+ pass.end();
553+ enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
554+ this.device.queue.submit([enc.finish()]);
555+ await this.stageSpat.mapAsync(GPUMapMode.READ);
556+ const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
557+ this.stageSpat.unmap();
558+ return out;
559+ }
560+
561+ /** Record a batched synthesis: one Legendre dispatch, K Fourier dispatches. */
562+ encodeSynthBatchInto(pass: GPUComputePassEncoder, b: ShtBatchBinding): void {
563+ const { mmax, nlat, nphi } = this.cfg;
564+ pass.setPipeline(this.pipeLegSynthB.get(b.size)!);
565+ pass.setBindGroup(0, b.bgLeg);
566+ pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
567+ pass.setPipeline(this.pipeFourSynth);
568+ for (const bg of b.bgFour) {
569+ pass.setBindGroup(0, bg);
570+ if (this.fourierMode === 'fft') {
571+ pass.dispatchWorkgroups(nlat);
572+ } else {
573+ pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
574+ }
575+ }
576+ }
577+
578+ /** Record a batched analysis: K Fourier dispatches, one Legendre dispatch. */
579+ encodeAnalysBatchInto(pass: GPUComputePassEncoder, b: ShtBatchBinding): void {
580+ const { mmax, nlat } = this.cfg;
581+ pass.setPipeline(this.pipeFourAnalys);
582+ for (const bg of b.bgFour) {
583+ pass.setBindGroup(0, bg);
584+ if (this.fourierMode === 'fft') {
585+ pass.dispatchWorkgroups(nlat);
586+ } else {
587+ pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
588+ }
589+ }
590+ pass.setPipeline(this.pipeLegAnalysB.get(b.size)!);
591+ pass.setBindGroup(0, b.bgLeg);
592+ pass.dispatchWorkgroups(mmax + 1);
593+ }
594+
595+ /** Record synthesis into an existing compute pass. */
596+ encodeSynthInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
597+ const { mmax, nlat, nphi } = this.cfg;
598+ pass.setPipeline(this.pipeLegSynth);
599+ pass.setBindGroup(0, b.bgLeg);
600+ pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
601+ pass.setPipeline(this.pipeFourSynth);
602+ pass.setBindGroup(0, b.bgFour);
603+ if (this.fourierMode === 'fft') {
604+ pass.dispatchWorkgroups(nlat);
605+ } else {
606+ pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
607+ }
608+ }
609+
610+ /** Record analysis into an existing compute pass. */
611+ encodeAnalysInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
612+ const { mmax, nlat } = this.cfg;
613+ pass.setPipeline(this.pipeFourAnalys);
614+ pass.setBindGroup(0, b.bgFour);
615+ if (this.fourierMode === 'fft') {
616+ pass.dispatchWorkgroups(nlat);
617+ } else {
618+ pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
619+ }
620+ pass.setPipeline(this.pipeLegAnalys);
621+ pass.setBindGroup(0, b.bgLeg);
622+ pass.dispatchWorkgroups(mmax + 1);
623+ }
624+
625+ /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
626+ encodeSynth(encoder: GPUCommandEncoder): void {
627+ const pass = encoder.beginComputePass({ label: 'sht-synth' });
628+ this.encodeSynthInto(pass, { bgLeg: this.bgLegSynth, bgFour: this.bgFourSynth });
629+ pass.end();
630+ }
631+
632+ /**
633+ * Diagnostics: encode one stage alone, in its own pass, so a timestamp query
634+ * can measure just that kernel. The solver wants both stages in a shared pass
635+ * and should use encodeSynthInto/encodeAnalysInto; this exists because
636+ * inferring per-kernel cost by subtracting trivially-sized runs is unreliable.
637+ */
638+ encodeStage(
639+ encoder: GPUCommandEncoder,
640+ stage: 'legSynth' | 'fourSynth' | 'fourAnalys' | 'legAnalys',
641+ timestampWrites?: GPUComputePassTimestampWrites,
642+ ): void {
643+ const { mmax, nlat, nphi } = this.cfg;
644+ const fft = this.fourierMode === 'fft';
645+ const pass = encoder.beginComputePass({ label: `sht-${stage}`, timestampWrites });
646+ switch (stage) {
647+ case 'legSynth':
648+ pass.setPipeline(this.pipeLegSynth);
649+ pass.setBindGroup(0, this.bgLegSynth);
650+ pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
651+ break;
652+ case 'fourSynth':
653+ pass.setPipeline(this.pipeFourSynth);
654+ pass.setBindGroup(0, this.bgFourSynth);
655+ if (fft) pass.dispatchWorkgroups(nlat);
656+ else pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
657+ break;
658+ case 'fourAnalys':
659+ pass.setPipeline(this.pipeFourAnalys);
660+ pass.setBindGroup(0, this.bgFourAnalys);
661+ if (fft) pass.dispatchWorkgroups(nlat);
662+ else pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
663+ break;
664+ case 'legAnalys':
665+ pass.setPipeline(this.pipeLegAnalys);
666+ pass.setBindGroup(0, this.bgLegAnalys);
667+ pass.dispatchWorkgroups(mmax + 1);
668+ break;
669+ }
670+ pass.end();
671+ }
672+
673+ /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
674+ encodeAnalys(encoder: GPUCommandEncoder): void {
675+ const pass = encoder.beginComputePass({ label: 'sht-analys' });
676+ this.encodeAnalysInto(pass, { bgLeg: this.bgLegAnalys, bgFour: this.bgFourAnalys });
677+ pass.end();
678+ }
679+
680+ /**
681+ * Spectral -> spatial. qlm: interleaved [re, im], SHTNS LM ordering,
682+ * length 2*nlm. Returns the spatial field, length nlat*nphi.
683+ */
684+ async synth(qlm: Float32Array): Promise<Float32Array> {
685+ const { nlat, nphi } = this.cfg;
686+ if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
687+ this.device.queue.writeBuffer(this.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
688+ const enc = this.device.createCommandEncoder();
689+ this.encodeSynth(enc);
690+ enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
691+ this.device.queue.submit([enc.finish()]);
692+ await this.stageSpat.mapAsync(GPUMapMode.READ);
693+ const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
694+ this.stageSpat.unmap();
695+ return out;
696+ }
697+
698+ /**
699+ * Spectral -> spatial, with the coefficients read from a caller-owned GPU
700+ * buffer (interleaved [re, im], 8*nlm bytes, COPY_SRC) instead of uploaded
701+ * from the CPU. This is how a field already on the device — a model's
702+ * spectral state — is evaluated on this plan's grid, e.g. a finer display
703+ * grid than the one the coefficients were produced on.
704+ */
705+ async synthFrom(qlmSrc: GPUBuffer): Promise<Float32Array> {
706+ const { nlat, nphi } = this.cfg;
707+ const enc = this.device.createCommandEncoder({ label: 'sht-synth-from' });
708+ enc.copyBufferToBuffer(qlmSrc, 0, this.qlmIn, 0, 8 * this.nlm);
709+ this.encodeSynth(enc);
710+ enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
711+ this.device.queue.submit([enc.finish()]);
712+ await this.stageSpat.mapAsync(GPUMapMode.READ);
713+ const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
714+ this.stageSpat.unmap();
715+ return out;
716+ }
717+
718+ /** Spatial -> spectral. spat: length nlat*nphi. Returns interleaved qlm, length 2*nlm. */
719+ async analys(spat: Float32Array): Promise<Float32Array> {
720+ const { nlat, nphi } = this.cfg;
721+ if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
722+ this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
723+ const enc = this.device.createCommandEncoder();
724+ this.encodeAnalys(enc);
725+ enc.copyBufferToBuffer(this.qlmOut, 0, this.stageQ, 0, 8 * this.nlm);
726+ this.device.queue.submit([enc.finish()]);
727+ await this.stageQ.mapAsync(GPUMapMode.READ);
728+ const out = new Float32Array(this.stageQ.getMappedRange().slice(0));
729+ this.stageQ.unmap();
730+ return out;
731+ }
732+
733+ destroy(): void {
734+ for (const b of [
735+ this.bufAb, this.bufAmm, this.bufCtstw, this.bufTrig, this.qlmIn, this.qlmOut,
736+ this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ, this.fmArena,
737+ ]) b?.destroy();
738+ }
739+}
740+
741+/** Best-effort human-readable adapter name, so it is clear which GPU (or
742+ * software rasterizer) is actually running the transforms. */
743+export async function describeAdapter(device: GPUDevice): Promise<string> {
744+ const fmt = (info: GPUAdapterInfo | undefined): string => {
745+ if (!info) return '';
746+ const parts = [info.description, info.device, info.vendor].filter(
747+ (s): s is string => !!s && s.length > 0,
748+ );
749+ const name = parts[0] ?? '';
750+ return info.architecture && !name.includes(info.architecture)
751+ ? `${name} (${info.architecture})`.trim()
752+ : name;
753+ };
754+ const own = fmt((device as GPUDevice & { adapterInfo?: GPUAdapterInfo }).adapterInfo);
755+ if (own) return own;
756+ try {
757+ const adapter = await navigator.gpu.requestAdapter();
758+ return fmt(adapter?.info);
759+ } catch {
760+ return '';
761+ }
762+}
763+
764+/** Request an adapter/device suitable for the transforms. */
765+export async function requestShtDevice(): Promise<GPUDevice> {
766+ if (!navigator.gpu) throw new Error('WebGPU is not available in this browser');
767+ const adapter = await navigator.gpu.requestAdapter();
768+ if (!adapter) throw new Error('No WebGPU adapter available');
769+ // ask for a larger workgroup storage if the adapter offers it (bigger FFTs)
770+ const wgStorage = Math.min(adapter.limits.maxComputeWorkgroupStorageSize, 32768);
771+ // `subgroups` lets the analysis reduction use subgroupAdd instead of a
772+ // shared-memory tree (2 barriers per l-pair instead of 1 + log2(wgAnalys)).
773+ // Optional: ShtPlan falls back to the tree when it is not available.
774+ const features: GPUFeatureName[] = [];
775+ if (adapter.features.has('subgroups')) features.push('subgroups');
776+ // timestamp-query is only used by the profiling scripts, but it has to be
777+ // requested at device creation, and asking costs nothing when unused.
778+ if (adapter.features.has('timestamp-query')) features.push('timestamp-query');
779+ // The seed field's mode table is the one buffer whose size is not fixed by
780+ // the grid — it grows with how fine a wavelength is asked for
781+ // (src/mgpu/randnfun3.ts), and a browser's default 128 MB storage-buffer
782+ // limit is well below what the adapter will actually give. Ask for the
783+ // adapter's own maximum so the wavelength is limited by the hardware rather
784+ // than by a default.
785+ const maxStorage = adapter.limits.maxStorageBufferBindingSize;
786+ const maxBuffer = adapter.limits.maxBufferSize;
787+ return adapter.requestDevice({
788+ requiredFeatures: features,
789+ requiredLimits: {
790+ maxComputeWorkgroupStorageSize: wgStorage,
791+ maxStorageBufferBindingSize: maxStorage,
792+ maxBufferSize: maxBuffer,
793+ },
794+ });
795+}
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/deriv.tsadded+142−0View file
@@ -0,0 +1,142 @@
1+/**
2+ * WGSL kernels for the coefficient-space step of the theta/phi first-
3+ * derivative algorithm (evolving_surface/notes/algos.tex, Algorithm 1, theta
4+ * and phi branches only -- the Laplace-Beltrami operator built on these never
5+ * needs the second-derivative/curvature branches).
6+ *
7+ * Both derivatives are a shuffle across nearby spectral coefficients --
8+ * independent of latitude/longitude, so neither touches the Legendre
9+ * recurrence or Fourier stages in leg.ts/fourier.ts -- followed by the
10+ * *existing*, unchanged scalar synthesis pipeline. dtheta additionally
11+ * divides the synthesized grid field by sin(theta) afterwards.
12+ */
13+
14+const WG = 64;
15+
16+export interface DerivCoeffParams {
17+ nlm: number;
18+}
19+
20+/**
21+ * v_l^m = alpha^+(l-1,m) * u_{l-1}^m + alpha^-(l+1,m) * u_{l+1}^m, the
22+ * coefficients of sin(theta) * dtheta(u) (algos.tex eq. v_coeffs). aPlus/
23+ * aMinus are precomputed zero at each m-block's boundary
24+ * (src/sht/derivCoeffs.ts), so the multiply is always mathematically
25+ * correct; the bounds checks below exist only to avoid reading past the ends
26+ * of the qlm array (an m-block-internal +-1 step never leaves the array, so
27+ * this is the only place it could).
28+ */
29+export function dthetaShuffleWGSL(p: DerivCoeffParams): string {
30+ return /* wgsl */ `
31+const NLM: u32 = ${p.nlm}u;
32+
33+@group(0) @binding(0) var<storage, read> aPlus: array<f32>;
34+@group(0) @binding(1) var<storage, read> aMinus: array<f32>;
35+@group(0) @binding(2) var<storage, read> qlmIn: array<vec2f>;
36+@group(0) @binding(3) var<storage, read_write> vOut: array<vec2f>;
37+
38+@compute @workgroup_size(${WG})
39+fn dtheta_shuffle(@builtin(global_invocation_id) gid: vec3u) {
40+ let lm = gid.x;
41+ if (lm >= NLM) { return; }
42+ var v = vec2f(0.0);
43+ if (lm > 0u) { v += aPlus[lm] * qlmIn[lm - 1u]; }
44+ if (lm + 1u < NLM) { v += aMinus[lm] * qlmIn[lm + 1u]; }
45+ vOut[lm] = v;
46+}
47+`;
48+}
49+
50+/**
51+ * (dphi u)_l^m = i*m*u_l^m: in the [re, im] row layout this swaps and
52+ * negates, re' = -m*im, im' = m*re (algos.tex eq. dYdphi).
53+ */
54+export function dphiShuffleWGSL(p: DerivCoeffParams): string {
55+ return /* wgsl */ `
56+const NLM: u32 = ${p.nlm}u;
57+
58+@group(0) @binding(0) var<storage, read> mOf: array<u32>;
59+@group(0) @binding(1) var<storage, read> qlmIn: array<vec2f>;
60+@group(0) @binding(2) var<storage, read_write> vOut: array<vec2f>;
61+
62+@compute @workgroup_size(${WG})
63+fn dphi_shuffle(@builtin(global_invocation_id) gid: vec3u) {
64+ let lm = gid.x;
65+ if (lm >= NLM) { return; }
66+ let m = f32(mOf[lm]);
67+ let c = qlmIn[lm];
68+ vOut[lm] = vec2f(-m * c.y, m * c.x);
69+}
70+`;
71+}
72+
73+export interface DivideParams {
74+ nlat: number;
75+ nphi: number;
76+}
77+
78+/**
79+ * Elementwise divide by sin(theta): the grid-space finish of Algorithm 1's
80+ * dtheta branch (dtheta(u) = synth(v_l^m) / sin(theta)). Gauss nodes never
81+ * sit at the poles, so this never divides by zero.
82+ */
83+export function divideSinThetaWGSL(p: DivideParams): string {
84+ const npts = p.nlat * p.nphi;
85+ return /* wgsl */ `
86+const NLAT: u32 = ${p.nlat}u;
87+const NPHI: u32 = ${p.nphi}u;
88+const NPTS: u32 = ${npts}u;
89+
90+@group(0) @binding(0) var<storage, read> sinTheta: array<f32>;
91+@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
92+
93+@compute @workgroup_size(${WG})
94+fn divide_sin_theta(@builtin(global_invocation_id) gid: vec3u) {
95+ let i = gid.x;
96+ if (i >= NPTS) { return; }
97+ let ilat = i / NPHI;
98+ spat[i] = spat[i] / sinTheta[ilat];
99+}
100+`;
101+}
102+
103+export interface FmDphiParams {
104+ mmax: number;
105+ nlat: number;
106+ nphi: number;
107+ /** Highest m kept; modes above are zeroed (mirrors the l-space filt). */
108+ mcut: number;
109+}
110+
111+/**
112+ * The Fourier-space middle of the grid-space phi-derivative `dphig`:
113+ * fm holds the unnormalized DFT modes of each latitude row (what the
114+ * Fourier analysis stage produces), so d/dphi is fm[m] *= i*m/NPHI --
115+ * the 1/NPHI undoes the unnormalized analysis+synthesis round trip.
116+ * Modes above MCUT are zeroed: the Fourier analysis stage already
117+ * truncated m > mmax for free, and MCUT additionally mirrors the
118+ * top-degree filt so the differentiated field carries no content the
119+ * l-space route would not have kept.
120+ */
121+export function fmDphiWGSL(p: FmDphiParams): string {
122+ const count = (p.mmax + 1) * p.nlat;
123+ return /* wgsl */ `
124+const NLAT: u32 = ${p.nlat}u;
125+const COUNT: u32 = ${count}u;
126+const MCUT: u32 = ${Math.max(0, p.mcut)}u;
127+const INV_NPHI: f32 = ${1 / p.nphi};
128+
129+@group(0) @binding(0) var<storage, read_write> fm: array<vec2f>;
130+
131+@compute @workgroup_size(${WG})
132+fn fm_dphi(@builtin(global_invocation_id) gid: vec3u) {
133+ let i = gid.x;
134+ if (i >= COUNT) { return; }
135+ let m = i / NLAT;
136+ var k: f32 = 0.0;
137+ if (m <= MCUT) { k = f32(m) * INV_NPHI; }
138+ let c = fm[i];
139+ fm[i] = vec2f(-k * c.y, k * c.x);
140+}
141+`;
142+}
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+688−0View file
@@ -0,0 +1,688 @@
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+}
353+
354+/**
355+ * Batched transforms: K independent fields through ONE walk of the Legendre
356+ * recurrence. The recurrence state (y0, y1, rescaling) depends only on
357+ * (m, theta), never on the field, so a batch shares it and pays only the
358+ * extra data fetches and accumulators per lane — the same amortization
359+ * SHTNS's GPU backend gets from batching fields. Per-lane arithmetic is
360+ * textually identical to the scalar kernels' (same operations, same order),
361+ * so a batched transform reproduces the scalar transform's results.
362+ *
363+ * K is a codegen parameter. The bind group needs 3 tables + K inputs +
364+ * K outputs storage buffers, so K = 2 (7 bindings) fits WebGPU's default
365+ * limit of 8 on every stack including SwiftShader, and K = 4 (11) needs the
366+ * raised limit requestShtDevice asks for where the adapter offers it.
367+ * The recurrence bodies are kept in exactly the shape the scalar kernels
368+ * use — see the driver-workaround comment in legSynthWGSL before
369+ * "simplifying" either copy.
370+ */
371+export function legSynthBatchWGSL(p: LegParams, K: number, laneElems: number): string {
372+ const half = p.parity === true;
373+ const lanes = Array.from({ length: K }, (_, k) => k);
374+ // K caller-owned inputs, ONE plan-owned fm arena: lane k writes at a fixed
375+ // 256-byte-aligned offset (laneElems vec2f), which is what keeps the bind
376+ // group at 3 + K + 1 storage buffers -- within WebGPU's default limit of 8
377+ // at K = 4. The Fourier stage binds the arena per lane with a buffer
378+ // offset, so it needs no changes.
379+ const bind =
380+ lanes
381+ .map((k) => `@group(0) @binding(${3 + k}) var<storage, read> qlm${k}: array<vec2f>;`)
382+ .join('\n') +
383+ `\n@group(0) @binding(${3 + K}) var<storage, read_write> fm: array<vec2f>;`;
384+ const decl = lanes
385+ .map((k) =>
386+ half
387+ ? ` var accE${k} = vec2f(0.0);\n var accO${k} = vec2f(0.0);`
388+ : ` var acc${k} = vec2f(0.0);`,
389+ )
390+ .join('\n');
391+ const accEven = lanes
392+ .map((k) => (half ? ` accE${k} += y0 * qlm${k}[i0];` : ` acc${k} += y0 * qlm${k}[i0];`))
393+ .join('\n');
394+ const accOdd = lanes
395+ .map((k) => (half ? ` accO${k} += y1 * qlm${k}[i1];` : ` acc${k} += y1 * qlm${k}[i1];`))
396+ .join('\n');
397+ const store = lanes
398+ .map((k) =>
399+ half
400+ ? ` fm[${k}u * LANE + m * NLAT + ilat] = accE${k} + accO${k};\n` +
401+ ` fm[${k}u * LANE + m * NLAT + (NLAT - 1u - ilat)] = accE${k} - accO${k};`
402+ : ` fm[${k}u * LANE + m * NLAT + ilat] = acc${k};`,
403+ )
404+ .join('\n');
405+ return /* wgsl */ `
406+${RESCALE_WGSL}
407+const LMAX: u32 = ${p.lmax}u;
408+const NLAT: u32 = ${p.nlat}u;
409+const NLAT_2: u32 = ${p.nlat / 2}u;
410+const LANE: u32 = ${laneElems}u;
411+${BINDINGS}
412+${bind}
413+
414+@compute @workgroup_size(${p.wgSynth})
415+fn leg_synth_batch(@builtin(global_invocation_id) gid: vec3u,
416+ @builtin(workgroup_id) wid: vec3u) {
417+ let ilat = gid.x;
418+ let m = wid.y;
419+ if (ilat >= ${half ? 'NLAT_2' : 'NLAT'}) { return; }
420+
421+ let ct = ctstw[ilat];
422+ let st = ctstw[NLAT + ilat];
423+ let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
424+
425+ var seed = sinpow_rescaled(st, m);
426+ var y0 = seed.y0 * amm[m];
427+ var ny = seed.ny;
428+ var y1: f32 = 0.0;
429+ if (m < LMAX) {
430+ y1 = ab[base + 1u].x * ct * y0;
431+ }
432+
433+${decl}
434+ var l = m;
435+ loop {
436+ if (ny == 0) {
437+ let i0 = base + (l - m);
438+${accEven}
439+ if (l + 1u <= LMAX) {
440+ let i1 = base + (l + 1u - m);
441+${accOdd}
442+ }
443+ } else if (abs(y0) > RESCALE_THR) {
444+ ny += 1;
445+ y0 *= INV_SCALE;
446+ y1 *= INV_SCALE;
447+ }
448+ if (l + 2u > LMAX) { break; }
449+ // Same two-coefficient, temporary-carried shape as leg_synth (see the
450+ // driver-workaround comment there).
451+ let a0 = ab[base + (l + 2u - m)];
452+ var a1 = vec2f(0.0);
453+ if (l + 3u <= LMAX) {
454+ a1 = ab[base + (l + 3u - m)];
455+ }
456+ let t0 = a0.x * ct * y1 + a0.y * y0;
457+ y1 = a1.x * ct * t0 + a1.y * y1;
458+ y0 = t0;
459+ l += 2u;
460+ }
461+${store}
462+}
463+`;
464+}
465+
466+/** Batched analysis: K spatial-Fourier fields reduced against one Legendre
467+ * recurrence walk. Structure follows legAnalysWGSL; see legSynthBatchWGSL
468+ * for the batching rationale and the binding budget. */
469+export function legAnalysBatchWGSL(p: LegParams, K: number, laneElems: number): string {
470+ const half = p.parity === true;
471+ const lanes = Array.from({ length: K }, (_, k) => k);
472+ const Kl = Math.ceil((half ? p.nlat / 2 : p.nlat) / p.wgAnalys);
473+ const sg = p.subgroups === true;
474+ const nsubMax = Math.max(1, p.wgAnalys / 4);
475+ // Same 8 KB workgroup-storage budget as the scalar kernel, now split
476+ // across K lanes, so spans shorten as K grows: barriers per unit of work
477+ // stay level.
478+ const pairs = sg
479+ ? Math.max(1, Math.min(p.spanPairs ?? 16, Math.floor(8192 / (nsubMax * 16 * K))))
480+ : 1;
481+ const redLen = (sg ? nsubMax * pairs : p.wgAnalys) * K;
482+ // ONE fm arena in (lane offsets baked, as in legSynthBatchWGSL), K
483+ // caller-owned outputs: 3 + 1 + K storage buffers.
484+ const bind =
485+ `@group(0) @binding(3) var<storage, read> fm: array<vec2f>;\n` +
486+ lanes
487+ .map((k) => `@group(0) @binding(${4 + k}) var<storage, read_write> qout${k}: array<vec2f>;`)
488+ .join('\n');
489+ const laneState = lanes
490+ .map((k) =>
491+ half
492+ ? ` var wpv${k}: array<vec2f, ${Kl}>;\n var wmv${k}: array<vec2f, ${Kl}>;`
493+ : ` var wfv${k}: array<vec2f, ${Kl}>;`,
494+ )
495+ .join('\n');
496+ const laneLoad = lanes
497+ .map((k) =>
498+ half
499+ ? ` let gN${k} = fm[${k}u * LANE + m * NLAT + lat];
500+ let gS${k} = fm[${k}u * LANE + m * NLAT + (NLAT - 1u - lat)];
501+ wp${k} = (gN${k} + gS${k}) * w;
502+ wm${k} = (gN${k} - gS${k}) * w;`
503+ : ` wf${k} = fm[${k}u * LANE + m * NLAT + lat] * w;`,
504+ )
505+ .join('\n');
506+ const laneLoadDecl = lanes
507+ .map((k) =>
508+ half ? ` var wp${k} = vec2f(0.0);\n var wm${k} = vec2f(0.0);` : ` var wf${k} = vec2f(0.0);`,
509+ )
510+ .join('\n');
511+ const laneLoadStore = lanes
512+ .map((k) => (half ? ` wpv${k}[k] = wp${k};\n wmv${k}[k] = wm${k};` : ` wfv${k}[k] = wf${k};`))
513+ .join('\n');
514+ const cDecl = lanes.map((k) => ` var c0_${k} = vec2f(0.0);\n var c1_${k} = vec2f(0.0);`).join('\n');
515+ const cAcc = lanes
516+ .map((k) =>
517+ half
518+ ? ` c0_${k} += wpv${k}[k] * y0v[k];
519+ c1_${k} += wmv${k}[k] * y1v[k];`
520+ : ` c0_${k} += wfv${k}[k] * y0v[k];
521+ c1_${k} += wfv${k}[k] * y1v[k];`,
522+ )
523+ .join('\n');
524+ return /* wgsl */ `${sg ? 'enable subgroups;\n' : ''}
525+${RESCALE_WGSL}
526+const LMAX: u32 = ${p.lmax}u;
527+const NLAT: u32 = ${p.nlat}u;
528+const WG: u32 = ${p.wgAnalys}u;
529+const K: u32 = ${Kl}u;
530+const NLAT_2: u32 = ${p.nlat / 2}u;
531+const PAIRS: u32 = ${pairs}u;
532+const NB: u32 = ${K}u;
533+const LANE: u32 = ${laneElems}u;
534+${BINDINGS}
535+${bind}
536+
537+var<workgroup> red: array<vec4f, ${redLen}>;
538+
539+@compute @workgroup_size(${p.wgAnalys})
540+fn leg_analys_batch(@builtin(local_invocation_id) lid3: vec3u,
541+ @builtin(workgroup_id) wid: vec3u${
542+ sg
543+ ? ',\n @builtin(subgroup_size) sgSize: u32,\n @builtin(subgroup_invocation_id) sgLane: u32'
544+ : ''
545+ }) {
546+ let lid = lid3.x;
547+ let m = wid.x;
548+ let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
549+
550+ var y0v: array<f32, ${Kl}>;
551+ var y1v: array<f32, ${Kl}>;
552+ var nyv: array<i32, ${Kl}>;
553+ var ctv: array<f32, ${Kl}>;
554+${laneState}
555+
556+ for (var k = 0u; k < K; k++) {
557+ let lat = lid + k * WG;
558+ var ct: f32 = 0.0;
559+ var st: f32 = 0.0;
560+ var w: f32 = 0.0;
561+${laneLoadDecl}
562+ if (lat < ${half ? 'NLAT_2' : 'NLAT'}) {
563+ ct = ctstw[lat];
564+ st = ctstw[NLAT + lat];
565+ w = ctstw[2u * NLAT + lat];
566+${laneLoad}
567+ }
568+ ctv[k] = ct;
569+ let seed = sinpow_rescaled(st, m);
570+ y0v[k] = seed.y0 * amm[m];
571+ nyv[k] = seed.ny;
572+ y1v[k] = 0.0;
573+ if (m < LMAX) {
574+ y1v[k] = ab[base + 1u].x * ct * y0v[k];
575+ }
576+${laneLoadStore}
577+ }
578+
579+ var l = m;
580+${
581+ sg
582+ ? ` loop {
583+ let lstart = l;
584+ var npairs = 0u;
585+ var last = false;
586+ let sub = lid / sgSize;
587+ for (var jj = 0u; jj < PAIRS; jj++) {
588+${cDecl}
589+ for (var k = 0u; k < K; k++) {
590+ if (nyv[k] == 0) {
591+${cAcc}
592+ } else if (abs(y0v[k]) > RESCALE_THR) {
593+ nyv[k] += 1;
594+ y0v[k] *= INV_SCALE;
595+ y1v[k] *= INV_SCALE;
596+ }
597+ }
598+${lanes
599+ .map(
600+ (k) => ` let part${k} = subgroupAdd(vec4f(c0_${k}, c1_${k}));
601+ if (sgLane == 0u) { red[(sub * PAIRS + jj) * NB + ${k}u] = part${k}; }`,
602+ )
603+ .join('\n')}
604+ npairs = jj + 1u;
605+ if (l + 2u > LMAX) { last = true; break; }
606+ let a0 = ab[base + (l + 2u - m)];
607+ var a1 = vec2f(0.0);
608+ if (l + 3u <= LMAX) {
609+ a1 = ab[base + (l + 3u - m)];
610+ }
611+ for (var k = 0u; k < K; k++) {
612+ let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
613+ y0v[k] = t0;
614+ y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
615+ }
616+ l += 2u;
617+ }
618+
619+ workgroupBarrier();
620+ if (lid == 0u) {
621+ let nsub = (WG + sgSize - 1u) / sgSize;
622+ for (var jj = 0u; jj < npairs; jj++) {
623+ let ll = lstart + 2u * jj;
624+${lanes
625+ .map(
626+ (k) => ` var tot${k} = vec4f(0.0);
627+ for (var i = 0u; i < nsub; i++) { tot${k} += red[(i * PAIRS + jj) * NB + ${k}u]; }
628+ qout${k}[base + (ll - m)] = tot${k}.xy;
629+ if (ll + 1u <= LMAX) {
630+ qout${k}[base + (ll + 1u - m)] = tot${k}.zw;
631+ }`,
632+ )
633+ .join('\n')}
634+ }
635+ }
636+ workgroupBarrier(); // red is reused by the next span
637+
638+ if (last) { break; }
639+ }`
640+ : ` loop {
641+${cDecl}
642+ for (var k = 0u; k < K; k++) {
643+ if (nyv[k] == 0) {
644+${cAcc.replace(/^ {10}/gm, ' ')}
645+ } else if (abs(y0v[k]) > RESCALE_THR) {
646+ nyv[k] += 1;
647+ y0v[k] *= INV_SCALE;
648+ y1v[k] *= INV_SCALE;
649+ }
650+ }
651+ // workgroup tree reduction, lane-strided
652+${lanes.map((k) => ` red[lid + ${k}u * WG] = vec4f(c0_${k}, c1_${k});`).join('\n')}
653+ workgroupBarrier();
654+ var s = WG / 2u;
655+ while (s > 0u) {
656+ if (lid < s) {
657+${lanes.map((k) => ` red[lid + ${k}u * WG] += red[lid + s + ${k}u * WG];`).join('\n')}
658+ }
659+ workgroupBarrier();
660+ s = s >> 1u;
661+ }
662+ if (lid == 0u) {
663+${lanes
664+ .map(
665+ (k) => ` qout${k}[base + (l - m)] = red[${k}u * WG].xy;
666+ if (l + 1u <= LMAX) {
667+ qout${k}[base + (l + 1u - m)] = red[${k}u * WG].zw;
668+ }`,
669+ )
670+ .join('\n')}
671+ }
672+ if (l + 2u > LMAX) { break; }
673+ let a0 = ab[base + (l + 2u - m)];
674+ var a1 = vec2f(0.0);
675+ if (l + 3u <= LMAX) {
676+ a1 = ab[base + (l + 3u - m)];
677+ }
678+ for (var k = 0u; k < K; k++) {
679+ let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
680+ y0v[k] = t0;
681+ y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
682+ }
683+ l += 2u;
684+ }`
685+ }
686+}
687+`;
688+}
src/tools.tsadded+27−0View file
@@ -0,0 +1,27 @@
1+/**
2+ * The shared MATLAB utilities in `tools/`, as interpreter workspace files.
3+ *
4+ * A geometry or a seeding draw is evaluated by numbl's interpreter (see
5+ * src/geom/geometry.ts), which resolves a call like `randnfunsphere(...)`
6+ * against the workspace files it is handed. Everything in `tools/` is handed
7+ * to every such run, so any .m can call any tool by name — MATLAB's own path
8+ * semantics, where the file name is the function name.
9+ *
10+ * These are *not* available to the models: a model's step compiles to WGSL,
11+ * where none of this exists.
12+ */
13+const sources = import.meta.glob('../tools/*.m', {
14+ query: '?raw',
15+ eager: true,
16+ import: 'default',
17+}) as Record<string, string>;
18+
19+export interface ToolFile {
20+ name: string;
21+ source: string;
22+}
23+
24+/** Every tool, named as MATLAB wants it (`randnfunsphere.m`). */
25+export const toolFiles: ToolFile[] = Object.entries(sources)
26+ .map(([path, source]) => ({ name: path.slice(path.lastIndexOf('/') + 1), source }))
27+ .sort((a, b) => a.name.localeCompare(b.name));
tools/randnfun3.madded+66−0View file
@@ -0,0 +1,66 @@
1+% Smooth random function in 3D — chebfun's randnfun3, as the Fourier modes
2+% it is built from rather than as a chebfun3.
3+%
4+% [K, C] = randnfun3(LAMBDA, DOM) draws a random trig series on the box
5+% DOM = [x0 x1 y0 y1 z0 z1] with maximum frequency about 2*pi/LAMBDA in
6+% each direction and standard normal distribution N(0,1) at each point.
7+% K is nmodes x 3 (angular wavenumbers) and C is nmodes x 2 (real and
8+% imaginary parts), defining
9+%
10+% f(x,y,z) = sum_j C(j,1)*cos(K(j,:)*[x;y;z]) - C(j,2)*sin(K(j,:)*[x;y;z])
11+%
12+% Seed the draw with rng(...) before calling.
13+%
14+% chebfun returns a chebfun3 and evaluates it later; this project has no
15+% such object, and the sum above is what the GPU evaluates at the surface
16+% points (src/mgpu/randnfun3.ts). Splitting it here is also what keeps the
17+% draw in MATLAB: randn has no counterpart in the compiled WGSL dialect.
18+
19+function [k, c] = randnfun3(lambda, dom)
20+ % chebfun's nonperiodic path builds a periodic function on a domain about
21+ % 20% larger and restricts it. Restriction is free when evaluating at
22+ % points, so we keep the enlarged period and never form the smaller one.
23+ m = round(1.2*(dom(2)-dom(1))/lambda + 2);
24+ n = round(1.2*(dom(4)-dom(3))/lambda + 2);
25+ p = round(1.2*(dom(6)-dom(5))/lambda + 2);
26+ m2 = 2*m+1;
27+ n2 = 2*n+1;
28+ p2 = 2*p+1;
29+ N = m2*n2*p2;
30+
31+ % chebfun draws the whole cube (column-major) before masking; drawing in
32+ % that same order keeps a seed meaning the same thing here as there.
33+ cr = randn(N, 1);
34+ ci = randn(N, 1);
35+
36+ % The cube's integer wavenumbers, -m:m x -n:n x -p:p in column-major order.
37+ i = (0:N-1).';
38+ jx = mod(i, m2) - m;
39+ jy = mod(floor(i/m2), n2) - n;
40+ jz = floor(i/(m2*n2)) - p;
41+
42+ % Confine to a ball for isotropy.
43+ keep = ((jx/m).^2 + (jy/n).^2 + (jz/p).^2) <= 1;
44+ jx = jx(keep);
45+ jy = jy(keep);
46+ jz = jz(keep);
47+ cr = cr(keep);
48+ ci = ci(keep);
49+
50+ % Normalize so the variance is 1 at each point.
51+ s = 1/sqrt(numel(cr));
52+ cr = s*cr;
53+ ci = s*ci;
54+
55+ % Angular wavenumbers on the enlarged period, which is a whole number of
56+ % wavelengths on each side.
57+ kx = 2*pi*jx/(m*lambda);
58+ ky = 2*pi*jy/(n*lambda);
59+ kz = 2*pi*jz/(p*lambda);
60+
61+ % Fold the box's origin into the phase, so evaluating is a plain sum over
62+ % cos(k.x) and sin(k.x) with no offset left to carry.
63+ ph = -(kx*dom(1) + ky*dom(3) + kz*dom(5));
64+ k = [kx, ky, kz];
65+ c = [cr.*cos(ph) - ci.*sin(ph), cr.*sin(ph) + ci.*cos(ph)];
66+end
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+30−0View file
@@ -0,0 +1,30 @@
1+import { defineConfig } from 'vite';
2+import { realpathSync } from 'node:fs';
3+import { resolve } from 'node:path';
4+
5+// numbl is a local `file:` dependency, so node_modules/numbl is a symlink to
6+// the sibling checkout. Its package `exports` map only publishes the runtime
7+// entry points, not the compiler internals we need (parser + JIT lowering), so
8+// we reach them through a path alias. (package.json's `imports` field cannot
9+// express this — Node rejects node_modules targets — and plain Node could not
10+// resolve numbl's internal `.js`->`.ts` imports anyway, which is why the GPU
11+// tests run in the browser harness rather than under `node`.)
12+// Realpath'd through the symlink: dev serves modules under their real ids, so
13+// aliasing the node_modules path would give the same file two identities (one
14+// per spelling) and run its side effects twice — the interpreter's builtin
15+// registry throws on the second.
16+const numblSrc = realpathSync(resolve(import.meta.dirname, 'node_modules/numbl/src'));
17+
18+export default defineConfig({
19+ base: './',
20+ resolve: {
21+ alias: { 'numbl-src': numblSrc },
22+ },
23+ server: {
24+ // the alias resolves outside the project root (through the symlink)
25+ fs: { allow: [import.meta.dirname, numblSrc] },
26+ },
27+ build: {
28+ target: 'es2022',
29+ },
30+});