turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPU
Solves N-species reaction-diffusion on the unit sphere with a first-order
IMEX Euler step: diffusion implicit and diagonal in spherical-harmonic space
(Laplace-Beltrami eigenvalues -l(l+1)), reaction explicit on the grid. Ported
from a MATLAB reference implementation whose transform dependency is just
coeffs2vals/vals2coeffs plus the grid.
Transforms run on the GPU via shtns-webgpu (vendored under src/sht/), with the
f64 CPU reference transform kept as a selectable backend and fallback.
Rendering adapts figpack's SphereEmbedding view: three.js spheres with
per-vertex colormaps, pole caps and seam stitching, cameras synced across
species.
Presets for Schnakenberg, Brusselator and Allen-Cahn, with live-editable
parameters. Verified in Node against exact linear, reaction-ODE and linearized
Turing-mode recurrences, and in headless Chrome against the f64 CPU solver.
32 changed files+5414−0
.github/workflows/ci.ymladded+21−0View file
@@ -0,0 +1,21 @@
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+ - run: npm ci
17+ - run: npm run test:node
18+ # headless Chrome + SwiftShader software WebGPU
19+ - run: npm run test:gpu
20+ env:
21+ CHROME_PATH: /usr/bin/google-chrome
.github/workflows/deploy.ymladded+36−0View file
@@ -0,0 +1,36 @@
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+ - run: npm ci
29+ - run: npm run test:node
30+ - run: npm run build
31+ - uses: actions/configure-pages@v5
32+ - uses: actions/upload-pages-artifact@v3
33+ with:
34+ path: dist
35+ - id: deployment
36+ uses: actions/deploy-pages@v4
.gitignoreadded+4−0View file
@@ -0,0 +1,4 @@
1+node_modules/
2+dist/
3+*.log
4+*.png
README.mdadded+95−0View file
@@ -0,0 +1,95 @@
1+# turing-sphere
2+
3+Reaction–diffusion systems (Turing patterns) solved **live in the browser on the
4+surface of a sphere**, using a spectral spherical-harmonic method with the
5+transforms running on the GPU via WebGPU.
6+
7+**Live demo:** <https://concept-collection.github.io/turing-sphere/>
8+
9+## What it does
10+
11+It solves the N-species system
12+
13+```
14+d(u_k)/dt = D_k*lap_s(u_k) + f_k(t, x, y, z, u_1, ..., u_N), k = 1, ..., N
15+```
16+
17+on the unit sphere, where `lap_s` is the Laplace–Beltrami operator. Diffusion is
18+treated implicitly in spherical-harmonic coefficient space, where `lap_s` is
19+diagonal with eigenvalues `-l(l+1)`; reaction is treated explicitly on the grid.
20+The two are combined with a first-order IMEX Euler step — the entire time loop is
21+
22+```
23+V_k = synth(U_k) # spectral -> grid
24+R_k = analys(f_k(t, x, y, z, V_1..V_N)) # reaction on grid -> spectral
25+U_k = (U_k + dt*R_k) / (1 + dt*D_k*l(l+1))
26+```
27+
28+You watch the patterns emerge in real time on orbitable 3D spheres (one per
29+species, cameras synced), with pause/resume, re-seeding, live parameter editing,
30+and colormap selection.
31+
32+Three presets are included:
33+
34+- **Schnakenberg** — Turing spots (unstable band 14 ≤ l ≤ 40, peak l = 24)
35+- **Brusselator** — stripes and spots from a stiffer reaction
36+- **Allen–Cahn** — a single species whose interfaces form and coarsen
37+
38+## Provenance
39+
40+This is the browser port of a MATLAB reference implementation
41+(`SphericalReactionDiffusion.m`, "websph"), which defines the solver through a
42+four-member porting boundary: `coeffs2vals`, `vals2coeffs`, `grid.lat`,
43+`grid.lon`. Profiling of the MATLAB version shows the transforms are ~96% of
44+compute, so this port swaps in:
45+
46+- **Transforms:** [shtns-webgpu](https://github.com/concept-collection/shtns-webgpu) —
47+ fp32 spherical harmonic transforms in WGSL compute shaders, modeled on
48+ [SHTNS](https://nschaeff.bitbucket.io/shtns/). Its source is vendored under
49+ [`src/sht/`](src/sht/) (CECILL-2.1), including the f64 CPU reference
50+ transform used for testing and as a no-WebGPU fallback.
51+- **Rendering:** three.js spheres with per-vertex colormaps, adapted from the
52+ `SphereEmbedding` view in
53+ [figpack](https://github.com/flatironinstitute/figpack)'s experimental
54+ extension package ([`src/render/`](src/render/)).
55+- **Solver:** [`src/solver/simulation.ts`](src/solver/simulation.ts), a direct
56+ TypeScript port of the MATLAB IMEX loop, in f64 on the coefficients with the
57+ transforms in fp32 on the GPU.
58+
59+## Numerics
60+
61+- Grid: Gauss–Legendre × equispaced-phi, dealiased for the cubic reactions with
62+ the `(pdeg+1)` rule from the reference implementation:
63+ `nlat ≥ ((pdeg+1)·lmax+1)/2`, `nphi ≥ (pdeg+1)·lmax+1` (rounded up to a power
64+ of two for the GPU FFT path). At the default lmax 63 that is a 128×256 grid.
65+- Spectral layout: SHTNS conventions — orthonormal + Condon–Shortley, complex
66+ coefficients for m ≥ 0, m-major ordering.
67+- fp32 transforms introduce ~1e-6 relative error per step (verified against the
68+ f64 CPU path); for pattern formation from 1e-2 seeded noise this is
69+ inconsequential.
70+
71+## Tests
72+
73+- `npm run test:node` — f64 solver correctness in Node: exact single-mode
74+ linear recurrence, exact uniform-state reaction ODE, and the linearized
75+ Turing-mode 2×2 IMEX recurrence (all at ~1e-12).
76+- `npm run test:gpu` — builds and drives headless Chrome: GPU-vs-CPU transform
77+ and solver cross-checks, plus a 100-step stability run.
78+- `node scripts/longrun-node.ts` — CPU run to t = 100 confirming pattern
79+ saturation.
80+- `node scripts/screenshot.mjs out.png [light|dark] [minSteps]` — screenshot
81+ the demo after a number of steps.
82+
83+## Development
84+
85+```
86+npm install
87+npm run dev # local dev server
88+npm run build # type-check + production build to dist/
89+```
90+
91+Deployed to GitHub Pages by `.github/workflows/deploy.yml` on push to `main`.
92+
93+## License
94+
95+CECILL-2.1 (inherited from SHTNS via shtns-webgpu, whose sources are vendored).
index.htmladded+124−0View file
@@ -0,0 +1,124 @@
1+<!doctype html>
2+<html lang="en">
3+ <head>
4+ <meta charset="utf-8" />
5+ <meta name="viewport" content="width=device-width, initial-scale=1" />
6+ <link rel="icon" href="data:image/svg+xml,<svg xmlns=%22http://www.w3.org/2000/svg%22 viewBox=%220 0 100 100%22><circle cx=%2250%22 cy=%2250%22 r=%2245%22 fill=%22%232a7f62%22/><circle cx=%2235%22 cy=%2238%22 r=%2211%22 fill=%22%23f5d547%22/><circle cx=%2265%22 cy=%2258%22 r=%229%22 fill=%22%23f5d547%22/><circle cx=%2248%22 cy=%2274%22 r=%227%22 fill=%22%23f5d547%22/><circle cx=%2268%22 cy=%2230%22 r=%226%22 fill=%22%23f5d547%22/></svg>" />
7+ <title>turing-sphere — reaction-diffusion on the sphere, live in the browser</title>
8+ <style>
9+ :root {
10+ --bg: #ffffff;
11+ --ink: #1f2328;
12+ --ink-2: #57606a;
13+ --line: #d0d7de;
14+ --accent: #0969da;
15+ --sphere-bg: #f4f6f8;
16+ color-scheme: light dark;
17+ }
18+ @media (prefers-color-scheme: dark) {
19+ :root {
20+ --bg: #14171a;
21+ --ink: #e6e9ec;
22+ --ink-2: #9aa4af;
23+ --line: #333b44;
24+ --accent: #58a6ff;
25+ --sphere-bg: #14161c;
26+ }
27+ }
28+ body {
29+ margin: 0;
30+ background: var(--bg);
31+ color: var(--ink);
32+ font: 15px/1.5 system-ui, -apple-system, sans-serif;
33+ }
34+ main { max-width: 1100px; margin: 0 auto; padding: 20px 16px 48px; }
35+ h1 { font-size: 20px; margin: 0 0 2px; }
36+ .sub { color: var(--ink-2); margin: 0 0 12px; font-size: 13px; }
37+ .sub a { color: var(--accent); }
38+ .controls {
39+ display: flex; flex-wrap: wrap; gap: 8px 14px; align-items: center;
40+ padding: 6px 0;
41+ }
42+ .controls label { color: var(--ink-2); font-size: 13px; white-space: nowrap; }
43+ select, input[type="number"], button {
44+ font: inherit; font-size: 13px;
45+ color: var(--ink); background: var(--bg);
46+ border: 1px solid var(--line); border-radius: 6px;
47+ padding: 4px 8px;
48+ }
49+ input[type="number"] { width: 6em; }
50+ button { cursor: pointer; }
51+ button:hover { border-color: var(--accent); }
52+ button.primary { border-color: var(--accent); color: var(--accent); font-weight: 600; min-width: 5.5em; }
53+ #panels {
54+ display: flex; flex-wrap: wrap; gap: 14px; margin-top: 12px;
55+ }
56+ .panel {
57+ flex: 1 1 320px; min-width: 280px;
58+ border: 1px solid var(--line); border-radius: 8px; overflow: hidden;
59+ display: flex;
60+ }
61+ .sphere-box { flex: 1; aspect-ratio: 1 / 1; max-height: 70vh; position: relative; }
62+ .species-tag {
63+ position: absolute; top: 8px; left: 10px; z-index: 2;
64+ font-size: 15px; font-weight: 600; color: #fff;
65+ background: rgba(0, 0, 0, 0.45);
66+ padding: 1px 10px; border-radius: 12px;
67+ pointer-events: none;
68+ }
69+ .colorbar {
70+ display: flex; flex-direction: column; align-items: center; justify-content: center;
71+ gap: 4px; padding: 8px 4px; background: var(--sphere-bg);
72+ width: 52px; flex: none; box-sizing: border-box;
73+ }
74+ .colorbar canvas { border: 1px solid var(--line); border-radius: 2px; }
75+ .colorbar-label { font-size: 11px; color: var(--ink-2); font-variant-numeric: tabular-nums; }
76+ .stats { margin-top: 10px; font-size: 13px; color: var(--ink-2); font-variant-numeric: tabular-nums; }
77+ .stats b { color: var(--ink); font-weight: 600; }
78+ #blurb { margin-top: 4px; font-size: 13px; color: var(--ink-2); }
79+ #err { color: #b35900; white-space: pre-wrap; font-size: 13px; }
80+ </style>
81+ </head>
82+ <body>
83+ <main>
84+ <h1>turing-sphere</h1>
85+ <p class="sub">
86+ Reaction-diffusion on the sphere, solved live with spherical harmonics:
87+ implicit spectral diffusion + explicit reaction (IMEX Euler), transforms on WebGPU via
88+ <a href="https://github.com/concept-collection/shtns-webgpu">shtns-webgpu</a>.
89+ Drag to rotate.
90+ </p>
91+ <div class="controls">
92+ <label>preset
93+ <select id="model"></select>
94+ </label>
95+ <label>lmax
96+ <select id="lmax">
97+ <option value="31">31</option>
98+ <option value="63" selected>63</option>
99+ <option value="127">127</option>
100+ <option value="255">255</option>
101+ </select>
102+ </label>
103+ <label>colormap
104+ <select id="colormap"></select>
105+ </label>
106+ <label>backend
107+ <select id="backend">
108+ <option value="webgpu" selected>WebGPU (fp32)</option>
109+ <option value="cpu">CPU (f64)</option>
110+ </select>
111+ </label>
112+ <button id="runpause" class="primary">Run</button>
113+ <button id="reseed">Re-seed</button>
114+ <button id="resetview">Reset view</button>
115+ </div>
116+ <div class="controls" id="params"></div>
117+ <div id="panels"></div>
118+ <p class="stats" id="stats"></p>
119+ <p id="blurb"></p>
120+ <p id="err"></p>
121+ </main>
122+ <script type="module" src="/src/main.ts"></script>
123+ </body>
124+</html>
package-lock.jsonadded+2185−0View file
This diff is 2,190 lines long and is not shown.
package.jsonadded+25−0View file
@@ -0,0 +1,25 @@
1+{
2+ "name": "turing-sphere",
3+ "version": "0.1.0",
4+ "description": "Live reaction-diffusion (Turing patterns) on the sphere, spectral spherical-harmonic solver on WebGPU",
5+ "type": "module",
6+ "license": "CECILL-2.1",
7+ "scripts": {
8+ "dev": "vite",
9+ "build": "tsc --noEmit && vite build",
10+ "test:node": "node scripts/test-node.ts",
11+ "test:gpu": "vite build && node scripts/test-gpu.mjs",
12+ "test": "npm run test:node && npm run test:gpu"
13+ },
14+ "dependencies": {
15+ "three": "^0.183.0"
16+ },
17+ "devDependencies": {
18+ "@types/node": "^26.1.1",
19+ "@types/three": "^0.185.1",
20+ "@webgpu/types": "^0.1.44",
21+ "puppeteer-core": "^23.0.0",
22+ "typescript": "^5.5.0",
23+ "vite": "^5.4.0"
24+ }
25+}
scripts/longrun-node.tsadded+43−0View file
@@ -0,0 +1,43 @@
1+/**
2+ * Long-run sanity check (CPU f64, lmax 31): run Schnakenberg to t = 100 and
3+ * confirm the pattern saturates into O(1)-contrast spots rather than decaying
4+ * or blowing up. Run: node scripts/longrun-node.ts
5+ */
6+import { CpuBackend } from '../src/solver/backend.ts';
7+import { Simulation, gridForLmax } from '../src/solver/simulation.ts';
8+import { models, defaultParams } from '../src/solver/models.ts';
9+
10+const schnak = models[0];
11+const params = defaultParams(schnak);
12+const lmax = 31;
13+const { nlat, nphi } = gridForLmax(lmax, schnak.pdeg);
14+const backend = new CpuBackend({ lmax, mmax: lmax, nlat, nphi });
15+const sim = new Simulation(backend, schnak, params);
16+await sim.init(1);
17+
18+const nsteps = Math.round(100 / params.dt);
19+const t0 = performance.now();
20+for (let s = 0; s < nsteps; s++) {
21+ await sim.step();
22+ if ((s + 1) % 400 === 0) {
23+ let lo = Infinity, hi = -Infinity;
24+ for (const v of sim.V[0]) {
25+ if (v < lo) lo = v;
26+ if (v > hi) hi = v;
27+ }
28+ console.log(
29+ `t=${sim.t.toFixed(1).padStart(5)} u in [${lo.toFixed(4)}, ${hi.toFixed(4)}] ` +
30+ `contrast ${(hi - lo).toFixed(4)}`,
31+ );
32+ }
33+}
34+console.log(`${((performance.now() - t0) / nsteps).toFixed(1)} ms/step CPU`);
35+
36+let lo = Infinity, hi = -Infinity;
37+for (const v of sim.V[0]) {
38+ if (v < lo) lo = v;
39+ if (v > hi) hi = v;
40+}
41+const ok = Number.isFinite(lo) && hi - lo > 0.3 && hi - lo < 5;
42+console.log(ok ? 'PASS: saturated O(1) pattern' : 'FAIL: no saturated pattern');
43+process.exit(ok ? 0 : 1);
scripts/screenshot.mjsadded+57−0View file
@@ -0,0 +1,57 @@
1+/** Screenshot the demo page (dist/) in headless Chrome. Usage: node scripts/screenshot.mjs out.png [light|dark] */
2+import { createServer } from 'node:http';
3+import { readFile } from 'node:fs/promises';
4+import { extname, join } from 'node:path';
5+import puppeteer from 'puppeteer-core';
6+
7+const out = process.argv[2] ?? 'demo.png';
8+const scheme = process.argv[3] ?? 'light';
9+const minSteps = Number(process.argv[4] ?? 200);
10+const DIST = new URL('../dist/', import.meta.url).pathname;
11+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
12+
13+const server = createServer(async (req, res) => {
14+ try {
15+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
16+ const data = await readFile(join(DIST, path));
17+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
18+ res.end(data);
19+ } catch {
20+ res.writeHead(404);
21+ res.end();
22+ }
23+});
24+await new Promise((r) => server.listen(0, '127.0.0.1', r));
25+const port = server.address().port;
26+
27+const browser = await puppeteer.launch({
28+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
29+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
30+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
31+});
32+const page = await browser.newPage();
33+await page.setViewport({ width: 1000, height: 900 });
34+await page.emulateMediaFeatures([{ name: 'prefers-color-scheme', value: scheme }]);
35+page.on('console', (m) => console.log(' [page]', m.text()));
36+await page.goto(`http://127.0.0.1:${port}/index.html`, { waitUntil: 'load' });
37+// the sim starts paused; wait for setup to finish, then press Run
38+await page.waitForFunction(() => /grid/.test(document.getElementById('stats')?.textContent ?? ''), { timeout: 120_000 });
39+await page.click('#runpause');
40+await page.waitForFunction(
41+ (min) => {
42+ const s = document.getElementById('stats');
43+ const e = document.getElementById('err');
44+ const m = s && s.textContent.match(/\((\d+) steps\)/);
45+ return (m && Number(m[1]) >= min) || (e && e.textContent.length > 4);
46+ },
47+ { timeout: 600_000 },
48+ minSteps,
49+);
50+await new Promise((r) => setTimeout(r, 300));
51+await page.screenshot({ path: out });
52+console.log('screenshot:', out);
53+console.log('stats:', await page.$eval('#stats', (el) => el.textContent));
54+const err = await page.$eval('#err', (el) => el.textContent);
55+if (err) console.log('err:', err);
56+await browser.close();
57+server.close();
scripts/soak.mjsadded+87−0View file
@@ -0,0 +1,87 @@
1+/**
2+ * Soak test: drive the demo page for many steps and report JS heap growth and
3+ * any crash, distinguishing a page crash from a renderer/driver death.
4+ *
5+ * Usage: node scripts/soak.mjs [steps] [lmax] [backend]
6+ * e.g. node scripts/soak.mjs 1500 63 webgpu
7+ */
8+import { createServer } from 'node:http';
9+import { readFile } from 'node:fs/promises';
10+import { extname, join } from 'node:path';
11+import puppeteer from 'puppeteer-core';
12+
13+const steps = Number(process.argv[2] ?? 1000);
14+const lmax = process.argv[3] ?? '63';
15+const backend = process.argv[4] ?? 'webgpu';
16+const DIST = new URL('../dist/', import.meta.url).pathname;
17+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
18+
19+const server = createServer(async (req, res) => {
20+ try {
21+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
22+ const data = await readFile(join(DIST, path));
23+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
24+ res.end(data);
25+ } catch {
26+ res.writeHead(404);
27+ res.end();
28+ }
29+});
30+await new Promise((r) => server.listen(0, '127.0.0.1', r));
31+const port = server.address().port;
32+
33+const browser = await puppeteer.launch({
34+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
35+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
36+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
37+});
38+const page = await browser.newPage();
39+await page.setViewport({ width: 1000, height: 900 });
40+
41+let crashed = null;
42+page.on('error', (e) => { crashed = `page crash: ${e.message}`; });
43+page.on('pageerror', (e) => { crashed = `page error: ${e.message}`; });
44+page.on('console', (m) => {
45+ const t = m.text();
46+ if (!/GL Driver Message|Failed to load resource/.test(t)) console.log(' [page]', t);
47+});
48+
49+await page.goto(`http://127.0.0.1:${port}/index.html`, { waitUntil: 'load' });
50+await page.waitForFunction(() => /grid/.test(document.getElementById('stats')?.textContent ?? ''), { timeout: 120_000 });
51+await page.select('#lmax', lmax);
52+await page.select('#backend', backend);
53+await page.waitForFunction(() => /grid/.test(document.getElementById('stats')?.textContent ?? ''), { timeout: 120_000 });
54+await page.click('#runpause');
55+
56+const readStep = () =>
57+ page.evaluate(() => {
58+ const m = document.getElementById('stats')?.textContent?.match(/\((\d+) steps\)/);
59+ return m ? Number(m[1]) : 0;
60+ });
61+const heapMB = async () => {
62+ const m = await page.metrics();
63+ return (m.JSHeapUsedSize / 1048576).toFixed(1);
64+};
65+
66+const t0 = Date.now();
67+let last = 0;
68+let stalls = 0;
69+try {
70+ while (last < steps) {
71+ await new Promise((r) => setTimeout(r, 5000));
72+ if (crashed) throw new Error(crashed);
73+ const now = await readStep();
74+ console.log(` step ${now} heap ${await heapMB()} MB (+${now - last} in 5s)`);
75+ if (now === last) {
76+ if (++stalls >= 6) throw new Error(`stalled at step ${now}`);
77+ } else stalls = 0;
78+ last = now;
79+ }
80+ console.log(`SOAK PASS: ${last} steps in ${((Date.now() - t0) / 1000).toFixed(0)}s, heap ${await heapMB()} MB`);
81+} catch (e) {
82+ console.error(`SOAK FAIL at step ${last}: ${e.message}`);
83+ process.exitCode = 1;
84+} finally {
85+ await browser.close().catch(() => {});
86+ server.close();
87+}
scripts/test-gpu.mjsadded+68−0View file
@@ -0,0 +1,68 @@
1+/**
2+ * Headless GPU test runner: serves dist/, opens test.html in headless
3+ * Chrome (falling back to the SwiftShader software WebGPU adapter when no
4+ * hardware GPU is available), and reports the suite results.
5+ *
6+ * Run after `vite build`: node scripts/test-gpu.mjs
7+ */
8+import { createServer } from 'node:http';
9+import { readFile } from 'node:fs/promises';
10+import { extname, join } from 'node:path';
11+import puppeteer from 'puppeteer-core';
12+
13+const DIST = new URL('../dist/', import.meta.url).pathname;
14+const CHROME = process.env.CHROME_PATH ?? '/usr/bin/google-chrome';
15+const MIME = {
16+ '.html': 'text/html',
17+ '.js': 'text/javascript',
18+ '.css': 'text/css',
19+ '.json': 'application/json',
20+ '.wasm': 'application/wasm',
21+};
22+
23+const server = createServer(async (req, res) => {
24+ try {
25+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
26+ const data = await readFile(join(DIST, path));
27+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
28+ res.end(data);
29+ } catch {
30+ res.writeHead(404);
31+ res.end('not found');
32+ }
33+});
34+await new Promise((r) => server.listen(0, '127.0.0.1', r));
35+const port = server.address().port;
36+
37+const flagSets = [
38+ // hardware first, then SwiftShader (software) WebGPU
39+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu', '--enable-features=Vulkan'],
40+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu', '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
41+];
42+
43+let final = null;
44+for (const flags of flagSets) {
45+ const browser = await puppeteer.launch({ executablePath: CHROME, args: flags });
46+ try {
47+ const page = await browser.newPage();
48+ page.on('console', (msg) => console.log(` [page] ${msg.text()}`));
49+ page.on('pageerror', (err) => console.log(` [pageerror] ${err.message}`));
50+ await page.goto(`http://127.0.0.1:${port}/test.html`, { waitUntil: 'load' });
51+ const results = await page.waitForFunction(() => window.__RESULTS__, { timeout: 600_000 });
52+ final = await results.jsonValue();
53+ } catch (e) {
54+ console.error(`run with flags [${flags.join(' ')}] failed: ${e.message}`);
55+ } finally {
56+ await browser.close();
57+ }
58+ if (final && !final.fatal) break;
59+ console.log('retrying with next flag set…');
60+}
61+server.close();
62+
63+if (!final || final.fatal) {
64+ console.error(`GPU tests could not run: ${final?.fatal ?? 'no results'}`);
65+ process.exit(2);
66+}
67+console.log(final.ok ? 'GPU SUITE: PASS' : 'GPU SUITE: FAIL');
68+process.exit(final.ok ? 0 : 1);
scripts/test-node.tsadded+151−0View file
@@ -0,0 +1,151 @@
1+/**
2+ * Solver correctness tests against the f64 CPU transform backend.
3+ *
4+ * A. Linear reaction + diffusion, single mode: every (l,m) mode of
5+ * f = c*u with implicit diffusion follows the exact scalar recurrence
6+ * g = (1 + dt*c) / (1 + dt*D*l(l+1)).
7+ * B. Uniform state, nonlinear reaction: the l=0 mode follows the explicit
8+ * Euler map of the reaction ODE exactly.
9+ * C. Turing linear stability: a small single-mode perturbation of the
10+ * Schnakenberg fixed point follows the 2x2 linearized IMEX recurrence,
11+ * and the (24, 7) mode lies in the unstable band.
12+ *
13+ * Run: node scripts/test-node.ts
14+ */
15+import { CpuBackend } from '../src/solver/backend.ts';
16+import { Simulation, gridForLmax } from '../src/solver/simulation.ts';
17+import { models, defaultParams } from '../src/solver/models.ts';
18+import type { ModelSpec } from '../src/solver/models.ts';
19+import { lmIndex } from '../src/sht/layout.ts';
20+
21+let failures = 0;
22+function check(name: string, ok: boolean, detail: string): void {
23+ console.log(`${ok ? 'PASS' : 'FAIL'} ${name} ${detail}`);
24+ if (!ok) failures++;
25+}
26+
27+// ---------------------------------------------------------------- test A
28+{
29+ const lmax = 15;
30+ const { nlat, nphi } = gridForLmax(lmax, 1);
31+ const backend = new CpuBackend({ lmax, mmax: lmax, nlat, nphi });
32+ const c = -0.3;
33+ const D = 0.01;
34+ const model: ModelSpec = {
35+ key: 'linear', label: 'linear', blurb: '', species: ['u'],
36+ params: [], pdeg: 1, seedAmp: 0,
37+ diffusivities: () => [D],
38+ reaction(_p, _t, _x, _y, _z, V, out) {
39+ for (let i = 0; i < out[0].length; i++) out[0][i] = c * V[0][i];
40+ },
41+ init() {},
42+ };
43+ const sim = new Simulation(backend, model, { dt: 0.1 });
44+ const l = 5, m = 2;
45+ const idx = lmIndex(lmax, l, m);
46+ sim.U[0][2 * idx] = 0.8;
47+ sim.U[0][2 * idx + 1] = -0.35;
48+
49+ const nsteps = 20;
50+ for (let s = 0; s < nsteps; s++) await sim.step();
51+
52+ const g = (1 + 0.1 * c) / (1 + 0.1 * D * l * (l + 1));
53+ const gn = Math.pow(g, nsteps);
54+ const errRe = Math.abs(sim.U[0][2 * idx] - 0.8 * gn);
55+ const errIm = Math.abs(sim.U[0][2 * idx + 1] - -0.35 * gn);
56+ let leak = 0;
57+ for (let i = 0; i < backend.nlm; i++) {
58+ if (i === idx) continue;
59+ leak = Math.max(leak, Math.abs(sim.U[0][2 * i]), Math.abs(sim.U[0][2 * i + 1]));
60+ }
61+ check('A: single-mode linear recurrence', errRe < 1e-12 && errIm < 1e-12,
62+ `err=(${errRe.toExponential(2)}, ${errIm.toExponential(2)})`);
63+ check('A: no leakage into other modes', leak < 1e-12, `leak=${leak.toExponential(2)}`);
64+}
65+
66+// ---------------------------------------------------------------- test B
67+{
68+ const schnak = models[0];
69+ const p = defaultParams(schnak);
70+ const lmax = 15;
71+ const { nlat, nphi } = gridForLmax(lmax, schnak.pdeg);
72+ const backend = new CpuBackend({ lmax, mmax: lmax, nlat, nphi });
73+ const uniform: ModelSpec = {
74+ ...schnak,
75+ seedAmp: 0,
76+ init(pp, x, _y, _z, _randn, out) {
77+ out[0].fill(1.2);
78+ out[1].fill(0.8);
79+ void pp; void x;
80+ },
81+ };
82+ const sim = new Simulation(backend, uniform, p);
83+ await sim.init(1);
84+ const nsteps = 50;
85+ for (let s = 0; s < nsteps; s++) await sim.step();
86+
87+ // reference: explicit Euler on the 2-species ODE (l=0 is untouched by diffusion)
88+ let u = 1.2, v = 0.8;
89+ for (let s = 0; s < nsteps; s++) {
90+ const fu = p.a - u + u * u * v;
91+ const fv = p.b - u * u * v;
92+ u += p.dt * fu;
93+ v += p.dt * fv;
94+ }
95+ // the area mean is the l=0 coefficient of U (V lags U by one step)
96+ const sqrt4pi = Math.sqrt(4 * Math.PI);
97+ const i00 = 2 * lmIndex(lmax, 0, 0);
98+ const errU = Math.abs(sim.U[0][i00] / sqrt4pi - u);
99+ const errV = Math.abs(sim.U[1][i00] / sqrt4pi - v);
100+ check('B: uniform nonlinear reaction ODE', errU < 1e-10 && errV < 1e-10,
101+ `err=(${errU.toExponential(2)}, ${errV.toExponential(2)}) u=${u.toFixed(6)} v=${v.toFixed(6)}`);
102+}
103+
104+// ---------------------------------------------------------------- test C
105+{
106+ const schnak = models[0];
107+ const p = defaultParams(schnak);
108+ const lmax = 31;
109+ const { nlat, nphi } = gridForLmax(lmax, schnak.pdeg);
110+ const backend = new CpuBackend({ lmax, mmax: lmax, nlat, nphi });
111+ const sim = new Simulation(backend, schnak, p);
112+
113+ const us = p.a + p.b; // 1.0
114+ const vs = p.b / (us * us); // 0.9
115+ const sqrt4pi = Math.sqrt(4 * Math.PI);
116+ const l = 24, m = 7;
117+ const idx = lmIndex(lmax, l, m);
118+ const eps = 1e-6;
119+ const c0 = [eps, 0.5 * eps];
120+ // fixed point + single-mode perturbation, set directly in spectral space
121+ sim.U[0][2 * lmIndex(lmax, 0, 0)] = us * sqrt4pi;
122+ sim.U[1][2 * lmIndex(lmax, 0, 0)] = vs * sqrt4pi;
123+ sim.U[0][2 * idx] = c0[0];
124+ sim.U[1][2 * idx] = c0[1];
125+
126+ const nsteps = 20;
127+ for (let s = 0; s < nsteps; s++) await sim.step();
128+
129+ // linearized IMEX recurrence: c' = diag(1/(1+dt*Dk*lam)) * (I + dt*J) * c
130+ const lam = l * (l + 1);
131+ const J = [
132+ [-1 + 2 * us * vs, us * us],
133+ [-2 * us * vs, -us * us],
134+ ];
135+ let c = [...c0];
136+ for (let s = 0; s < nsteps; s++) {
137+ const r0 = c[0] + p.dt * (J[0][0] * c[0] + J[0][1] * c[1]);
138+ const r1 = c[1] + p.dt * (J[1][0] * c[0] + J[1][1] * c[1]);
139+ c = [r0 / (1 + p.dt * p.D1 * lam), r1 / (1 + p.dt * p.D2 * lam)];
140+ }
141+ const got = [sim.U[0][2 * idx], sim.U[1][2 * idx]];
142+ const errU = Math.abs(got[0] - c[0]) / Math.abs(c[0]);
143+ const errV = Math.abs(got[1] - c[1]) / Math.abs(c[1]);
144+ check('C: linearized Turing-mode recurrence', errU < 1e-4 && errV < 1e-4,
145+ `rel err=(${errU.toExponential(2)}, ${errV.toExponential(2)})`);
146+ check('C: (l=24, m=7) is growing', Math.abs(got[0]) > Math.abs(c0[0]),
147+ `|c|: ${Math.abs(c0[0]).toExponential(2)} -> ${Math.abs(got[0]).toExponential(2)}`);
148+}
149+
150+console.log(failures === 0 ? '\nAll tests passed.' : `\n${failures} test(s) FAILED.`);
151+process.exit(failures === 0 ? 0 : 1);
src/main.tsadded+363−0View file
@@ -0,0 +1,363 @@
1+import {
2+ GpuBackend,
3+ CpuBackend,
4+ requestShtDevice,
5+ type ShtBackend,
6+} from './solver/backend.ts';
7+import { Simulation, gridForLmax } from './solver/simulation.ts';
8+import {
9+ models,
10+ presets,
11+ defaultParams,
12+ type ModelSpec,
13+ type Params,
14+} from './solver/models.ts';
15+import {
16+ buildTopology,
17+ fillFieldValues,
18+ fillColors,
19+ type SphereMeshTopology,
20+} from './render/sphereMesh.ts';
21+import { SphereScene } from './render/SphereScene.ts';
22+import { Colorbar } from './render/colorbar.ts';
23+import { colormaps, colormapNames } from './render/colormaps.ts';
24+
25+const $ = <T extends HTMLElement>(id: string): T =>
26+ document.getElementById(id) as T;
27+
28+const elModel = $<HTMLSelectElement>('model');
29+const elLmax = $<HTMLSelectElement>('lmax');
30+const elColormap = $<HTMLSelectElement>('colormap');
31+const elBackend = $<HTMLSelectElement>('backend');
32+const elRunPause = $<HTMLButtonElement>('runpause');
33+const elReseed = $<HTMLButtonElement>('reseed');
34+const elResetView = $<HTMLButtonElement>('resetview');
35+const elParams = $('params');
36+const elPanels = $('panels');
37+const elStats = $('stats');
38+const elBlurb = $('blurb');
39+const elErr = $('err');
40+
41+for (const p of presets) {
42+ const o = document.createElement('option');
43+ o.value = p.key;
44+ o.textContent = p.label;
45+ elModel.append(o);
46+}
47+for (const name of colormapNames) {
48+ const o = document.createElement('option');
49+ o.value = name;
50+ o.textContent = name;
51+ elColormap.append(o);
52+}
53+elColormap.value = 'jet';
54+
55+// ---------------------------------------------------------------- state
56+let device: GPUDevice | null = null;
57+let backend: ShtBackend | null = null;
58+let sim: Simulation | null = null;
59+let topo: SphereMeshTopology | null = null;
60+let scenes: SphereScene[] = [];
61+let colorbars: Colorbar[] = [];
62+let valueBufs: Float32Array[] = [];
63+let colorBufs: Float32Array[] = [];
64+let ranges: { lo: number; hi: number }[] = [];
65+let resizeObs: ResizeObserver | null = null;
66+
67+let model: ModelSpec = models[0];
68+let params: Params = defaultParams(model);
69+let seed = 1;
70+let running = false;
71+let adapterName = '';
72+let pumping = false;
73+let stepMs = 0;
74+let generation = 0; // bumped on every rebuild to cancel stale pumps
75+
76+// ---------------------------------------------------------------- UI wiring
77+function buildParamInputs(): void {
78+ elParams.replaceChildren();
79+ for (const spec of model.params) {
80+ const label = document.createElement('label');
81+ label.textContent = `${spec.label} `;
82+ const input = document.createElement('input');
83+ input.type = 'number';
84+ input.min = String(spec.min);
85+ input.max = String(spec.max);
86+ input.step = String(spec.step);
87+ input.value = String(params[spec.key]);
88+ input.addEventListener('change', () => {
89+ const v = Number(input.value);
90+ if (Number.isFinite(v)) params[spec.key] = v;
91+ });
92+ label.append(input);
93+ elParams.append(label);
94+ }
95+}
96+
97+function applyPreset(presetKey: string): void {
98+ const preset = presets.find((p) => p.key === presetKey) ?? presets[0];
99+ model = models.find((m) => m.key === preset.modelKey) ?? models[0];
100+ params = { ...defaultParams(model), ...preset.params };
101+ buildParamInputs();
102+ elBlurb.textContent = model.blurb;
103+}
104+
105+elModel.addEventListener('change', () => {
106+ applyPreset(elModel.value);
107+ void rebuild();
108+});
109+elLmax.addEventListener('change', () => void rebuild());
110+elBackend.addEventListener('change', () => void rebuild());
111+elColormap.addEventListener('change', () => draw());
112+function setRunning(next: boolean): void {
113+ running = next;
114+ elRunPause.textContent = running ? 'Pause' : 'Run';
115+ if (running) void pump();
116+}
117+
118+elRunPause.addEventListener('click', () => setRunning(!running));
119+elReseed.addEventListener('click', () => {
120+ seed = (Math.random() * 2 ** 31) >>> 0;
121+ setRunning(false);
122+ void reseed();
123+});
124+elResetView.addEventListener('click', () => {
125+ for (const s of scenes) s.resetCamera();
126+});
127+
128+// ---------------------------------------------------------------- setup
129+function disposeView(): void {
130+ for (const s of scenes) s.dispose();
131+ scenes = [];
132+ colorbars = [];
133+ resizeObs?.disconnect();
134+ resizeObs = null;
135+ elPanels.replaceChildren();
136+}
137+
138+async function rebuild(): Promise<void> {
139+ generation++;
140+ const gen = generation;
141+ // a rebuild restarts from a fresh initial state, so pause like Re-seed does
142+ setRunning(false);
143+ disposeView();
144+ backend?.destroy();
145+ backend = null;
146+ sim = null;
147+ stepMs = 0;
148+
149+ const lmax = Number(elLmax.value);
150+ const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
151+ const cfg = { lmax, mmax: lmax, nlat, nphi };
152+ const wantGpu = elBackend.value === 'webgpu' && device !== null;
153+ try {
154+ backend = wantGpu
155+ ? await GpuBackend.create(device!, cfg)
156+ : new CpuBackend(cfg);
157+ } catch (e) {
158+ elErr.textContent = `Failed to create transform plan: ${e}`;
159+ return;
160+ }
161+ if (gen !== generation) return;
162+ elErr.textContent =
163+ !wantGpu && lmax > 31
164+ ? 'Heads up: the CPU backend is a direct-summation f64 reference — expect well under 10 steps/s at this lmax.'
165+ : '';
166+
167+ sim = new Simulation(backend, model, params);
168+ await sim.init(seed);
169+ if (gen !== generation) return;
170+
171+ // mesh + scenes
172+ const phi = new Float64Array(nphi);
173+ for (let j = 0; j < nphi; j++) phi[j] = (2 * Math.PI * j) / nphi;
174+ topo = buildTopology(backend.cosTheta, phi);
175+
176+ const sphereBg = getComputedStyle(document.documentElement)
177+ .getPropertyValue('--sphere-bg')
178+ .trim();
179+ for (let k = 0; k < sim.nspecies; k++) {
180+ const panel = document.createElement('div');
181+ panel.className = 'panel';
182+ const box = document.createElement('div');
183+ box.className = 'sphere-box';
184+ const tag = document.createElement('div');
185+ tag.className = 'species-tag';
186+ tag.textContent = model.species[k];
187+ box.append(tag);
188+ const side = document.createElement('div');
189+ panel.append(box, side);
190+ elPanels.append(panel);
191+
192+ const scene = new SphereScene(
193+ box,
194+ topo.numVertices,
195+ topo.indices,
196+ topo.sphereRef,
197+ sphereBg || undefined,
198+ );
199+ scene.fitCamera();
200+ scenes.push(scene);
201+ colorbars.push(new Colorbar(side));
202+ valueBufs[k] = new Float32Array(topo.numVertices);
203+ colorBufs[k] = new Float32Array(topo.numVertices * 3);
204+ ranges[k] = { lo: NaN, hi: NaN };
205+ }
206+ for (let k = 1; k < scenes.length; k++) scenes[0].syncCamerasWith(scenes[k]);
207+
208+ resizeObs = new ResizeObserver(() => {
209+ const boxes = elPanels.querySelectorAll<HTMLElement>('.sphere-box');
210+ boxes.forEach((box, i) => {
211+ scenes[i]?.resize(box.clientWidth, box.clientHeight);
212+ });
213+ });
214+ elPanels
215+ .querySelectorAll<HTMLElement>('.sphere-box')
216+ .forEach((box) => resizeObs!.observe(box));
217+
218+ draw();
219+ updateStats();
220+ void pump();
221+}
222+
223+async function reseed(): Promise<void> {
224+ if (!sim) return;
225+ const gen = generation;
226+ await sim.init(seed);
227+ if (gen !== generation) return;
228+ for (const r of ranges) {
229+ r.lo = NaN;
230+ r.hi = NaN;
231+ }
232+ draw();
233+}
234+
235+// ---------------------------------------------------------------- drawing
236+function draw(): void {
237+ if (!sim || !topo) return;
238+ const cmap = colormaps[elColormap.value] ?? colormaps.viridis;
239+ for (let k = 0; k < sim.nspecies; k++) {
240+ fillFieldValues(valueBufs[k], sim.V[k], topo);
241+ let lo = Infinity;
242+ let hi = -Infinity;
243+ for (const v of valueBufs[k]) {
244+ if (v < lo) lo = v;
245+ if (v > hi) hi = v;
246+ }
247+ // smooth the color range in both directions so the shading evolves
248+ // gently as the pattern grows (out-of-range values clamp meanwhile)
249+ const r = ranges[k];
250+ if (!Number.isFinite(r.lo)) {
251+ r.lo = lo;
252+ r.hi = hi;
253+ } else {
254+ const a = 0.15;
255+ r.lo += a * (lo - r.lo);
256+ r.hi += a * (hi - r.hi);
257+ }
258+ if (r.hi - r.lo < 1e-9) {
259+ const mid = (r.hi + r.lo) / 2;
260+ r.lo = mid - 5e-10;
261+ r.hi = mid + 5e-10;
262+ }
263+ fillColors(colorBufs[k], valueBufs[k], r.lo, r.hi, cmap);
264+ scenes[k]?.updateColors(colorBufs[k]);
265+ colorbars[k]?.update(cmap, r.lo, r.hi);
266+ }
267+}
268+
269+function updateStats(): void {
270+ if (!sim || !backend) return;
271+ const { nlat, nphi } = backend.cfg;
272+ const kind =
273+ backend.kind === 'webgpu'
274+ ? `WebGPU fp32${adapterName ? ` — ${adapterName}` : ''}`
275+ : 'CPU f64 (direct summation)';
276+ const rate = stepMs > 0 ? `${(1000 / stepMs).toFixed(1)} steps/s` : '—';
277+ elStats.innerHTML =
278+ `<b>${kind}</b> · grid ${nlat}×${nphi} · nlm ${backend.nlm.toLocaleString()} · ` +
279+ `${stepMs > 0 ? stepMs.toFixed(1) : '—'} ms/step · ${rate} · ` +
280+ `t = <b>${sim.t.toFixed(2)}</b> (${sim.stepCount} steps)`;
281+}
282+
283+// ---------------------------------------------------------------- sim loop
284+const nextFrame = () => new Promise<number>(requestAnimationFrame);
285+
286+async function pump(): Promise<void> {
287+ if (pumping) return;
288+ pumping = true;
289+ const gen = generation;
290+ let lastYield = performance.now();
291+ try {
292+ while (running && sim && gen === generation) {
293+ const t0 = performance.now();
294+ await sim.step();
295+ const dtMs = performance.now() - t0;
296+ stepMs = stepMs === 0 ? dtMs : stepMs + 0.05 * (dtMs - stepMs);
297+ const now = performance.now();
298+ if (now - lastYield > 25 || backend?.kind === 'cpu') {
299+ draw();
300+ updateStats();
301+ await nextFrame();
302+ lastYield = performance.now();
303+ }
304+ }
305+ // final frame after pausing
306+ if (gen === generation) {
307+ draw();
308+ updateStats();
309+ }
310+ } finally {
311+ pumping = false;
312+ }
313+}
314+
315+// ---------------------------------------------------------------- boot
316+/** Best-effort human-readable adapter name, so it is clear which GPU (or
317+ * software rasterizer) is actually running the transforms. */
318+async function describeAdapter(dev: GPUDevice): Promise<string> {
319+ const fmt = (info: GPUAdapterInfo | undefined): string => {
320+ if (!info) return '';
321+ const parts = [info.description, info.device, info.vendor].filter(
322+ (s): s is string => !!s && s.length > 0,
323+ );
324+ const name = parts[0] ?? '';
325+ return info.architecture && !name.includes(info.architecture)
326+ ? `${name} (${info.architecture})`.trim()
327+ : name;
328+ };
329+ const own = fmt((dev as GPUDevice & { adapterInfo?: GPUAdapterInfo }).adapterInfo);
330+ if (own) return own;
331+ try {
332+ const adapter = await navigator.gpu.requestAdapter();
333+ return fmt(adapter?.info);
334+ } catch {
335+ return '';
336+ }
337+}
338+
339+async function boot(): Promise<void> {
340+ elModel.value = presets[0].key;
341+ applyPreset(presets[0].key);
342+ try {
343+ device = await requestShtDevice();
344+ adapterName = await describeAdapter(device);
345+ } catch (e) {
346+ device = null;
347+ elLmax.value = '31';
348+ elBackend.value = 'cpu';
349+ elBackend.options[0].disabled = true;
350+ elErr.textContent =
351+ `WebGPU is not available (${e instanceof Error ? e.message : e}); ` +
352+ `falling back to the slow CPU transform at low resolution. ` +
353+ `Use a WebGPU-capable browser (Chrome/Edge 113+) for the full experience.`;
354+ }
355+ device?.lost.then((info) => {
356+ if (info.reason !== 'destroyed') {
357+ elErr.textContent = `WebGPU device lost: ${info.message}`;
358+ }
359+ });
360+ await rebuild();
361+}
362+
363+void boot();
src/render/SphereScene.tsadded+175−0View file
@@ -0,0 +1,175 @@
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 static
6+ * per-vertex positions and dynamic per-vertex colors, orbit controls, and
7+ * optional camera synchronization with sibling scenes.
8+ *
9+ * Adapted from figpack's SphereEmbedding view (figpack_experimental).
10+ */
11+export class SphereScene {
12+ #scene: THREE.Scene;
13+ #camera: THREE.PerspectiveCamera;
14+ #renderer: THREE.WebGLRenderer;
15+ #controls: OrbitControls;
16+ #geometry: THREE.BufferGeometry;
17+ #mesh: THREE.Mesh;
18+ #animationId: number | null = null;
19+ #defaultCameraState: {
20+ position: THREE.Vector3;
21+ target: THREE.Vector3;
22+ } | null = null;
23+ #syncing = false;
24+ #lastW = -1;
25+ #lastH = -1;
26+
27+ constructor(
28+ container: HTMLElement,
29+ numVertices: number,
30+ indices: Uint32Array,
31+ positions: Float32Array,
32+ background = '#14161c',
33+ ) {
34+ this.#scene = new THREE.Scene();
35+ this.#scene.background = new THREE.Color(background);
36+
37+ this.#camera = new THREE.PerspectiveCamera(50, 1, 0.01, 1000);
38+
39+ this.#renderer = new THREE.WebGLRenderer({ antialias: true });
40+ this.#renderer.setPixelRatio(window.devicePixelRatio || 1);
41+ // The canvas always fills its container via CSS; resize() then only
42+ // updates the drawing buffer
43+ this.#renderer.domElement.style.width = '100%';
44+ this.#renderer.domElement.style.height = '100%';
45+ this.#renderer.domElement.style.display = 'block';
46+ container.appendChild(this.#renderer.domElement);
47+
48+ // Lighting: ambient plus a headlight attached to the camera so the
49+ // surface stays lit from the viewing direction as it is rotated
50+ this.#scene.add(new THREE.AmbientLight(0xffffff, 0.65));
51+ const headlight = new THREE.DirectionalLight(0xffffff, 1.6);
52+ headlight.position.set(0.5, 0.8, 1);
53+ this.#camera.add(headlight);
54+ this.#scene.add(this.#camera);
55+
56+ this.#geometry = new THREE.BufferGeometry();
57+ const positionAttr = new THREE.BufferAttribute(positions, 3);
58+ const colorAttr = new THREE.BufferAttribute(
59+ new Float32Array(numVertices * 3),
60+ 3,
61+ );
62+ colorAttr.setUsage(THREE.DynamicDrawUsage);
63+ this.#geometry.setAttribute('position', positionAttr);
64+ this.#geometry.setAttribute('color', colorAttr);
65+ this.#geometry.setIndex(new THREE.BufferAttribute(indices, 1));
66+ this.#geometry.computeVertexNormals();
67+ this.#geometry.computeBoundingSphere();
68+
69+ const material = new THREE.MeshPhongMaterial({
70+ vertexColors: true,
71+ side: THREE.DoubleSide,
72+ shininess: 25,
73+ specular: new THREE.Color(0x222222),
74+ });
75+ this.#mesh = new THREE.Mesh(this.#geometry, material);
76+ this.#scene.add(this.#mesh);
77+
78+ this.#controls = new OrbitControls(this.#camera, this.#renderer.domElement);
79+ this.#controls.enableDamping = true;
80+ this.#controls.dampingFactor = 0.1;
81+
82+ this.#animate();
83+ }
84+
85+ #animate = () => {
86+ this.#animationId = requestAnimationFrame(this.#animate);
87+ this.#controls.update();
88+ this.#renderer.render(this.#scene, this.#camera);
89+ };
90+
91+ updateColors(colors: Float32Array): void {
92+ const attr = this.#geometry.getAttribute('color') as THREE.BufferAttribute;
93+ (attr.array as Float32Array).set(colors);
94+ attr.needsUpdate = true;
95+ }
96+
97+ /** Mirror this scene's camera whenever the other scene's controls move. */
98+ syncCamerasWith(other: SphereScene): void {
99+ const follow = (src: SphereScene, dst: SphereScene) => {
100+ src.#controls.addEventListener('change', () => {
101+ if (dst.#syncing) return;
102+ src.#syncing = true;
103+ dst.#camera.position.copy(src.#camera.position);
104+ dst.#camera.zoom = src.#camera.zoom;
105+ dst.#camera.updateProjectionMatrix();
106+ dst.#controls.target.copy(src.#controls.target);
107+ dst.#controls.update();
108+ src.#syncing = false;
109+ });
110+ };
111+ follow(this, other);
112+ follow(other, this);
113+ }
114+
115+ /** Position the camera to comfortably frame the geometry. */
116+ fitCamera(): void {
117+ this.#geometry.computeBoundingSphere();
118+ const bs = this.#geometry.boundingSphere;
119+ if (!bs) return;
120+ const radius = Math.max(bs.radius, 1e-6);
121+ const distance = radius * 2.6;
122+ this.#controls.target.copy(bs.center);
123+ this.#camera.position.set(
124+ bs.center.x + distance * 0.55,
125+ bs.center.y + distance * 0.35,
126+ bs.center.z + distance * 0.75,
127+ );
128+ this.#camera.near = radius * 0.01;
129+ this.#camera.far = radius * 100;
130+ this.#camera.updateProjectionMatrix();
131+ this.#controls.update();
132+ this.#defaultCameraState = {
133+ position: this.#camera.position.clone(),
134+ target: this.#controls.target.clone(),
135+ };
136+ }
137+
138+ resetCamera(): void {
139+ if (this.#defaultCameraState) {
140+ this.#camera.position.copy(this.#defaultCameraState.position);
141+ this.#controls.target.copy(this.#defaultCameraState.target);
142+ this.#controls.update();
143+ } else {
144+ this.fitCamera();
145+ }
146+ }
147+
148+ resize(width: number, height: number): void {
149+ // Setting canvas.width clears the canvas even at the same value, which
150+ // shows as a blank flash until the next render — skip no-op resizes.
151+ if (width === this.#lastW && height === this.#lastH) return;
152+ this.#lastW = width;
153+ this.#lastH = height;
154+ this.#camera.aspect = width / Math.max(1, height);
155+ this.#camera.updateProjectionMatrix();
156+ // updateStyle=false: the canvas keeps its 100%/100% CSS sizing
157+ this.#renderer.setSize(width, height, false);
158+ }
159+
160+ dispose(): void {
161+ if (this.#animationId !== null) {
162+ cancelAnimationFrame(this.#animationId);
163+ this.#animationId = null;
164+ }
165+ this.#controls.dispose();
166+ this.#geometry.dispose();
167+ (this.#mesh.material as THREE.Material).dispose();
168+ if (this.#renderer.domElement.parentNode) {
169+ this.#renderer.domElement.parentNode.removeChild(
170+ this.#renderer.domElement,
171+ );
172+ }
173+ this.#renderer.dispose();
174+ }
175+}
src/render/colorbar.tsadded+36−0View file
@@ -0,0 +1,36 @@
1+import type { ColormapFunc } from './colormaps.ts';
2+
3+/** Vertical colorbar drawn on a small canvas, with min/max labels. */
4+export class Colorbar {
5+ #canvas: HTMLCanvasElement;
6+ #minLabel: HTMLElement;
7+ #maxLabel: HTMLElement;
8+
9+ constructor(container: HTMLElement) {
10+ container.classList.add('colorbar');
11+ this.#maxLabel = document.createElement('div');
12+ this.#maxLabel.className = 'colorbar-label';
13+ this.#canvas = document.createElement('canvas');
14+ this.#canvas.width = 12;
15+ this.#canvas.height = 160;
16+ this.#minLabel = document.createElement('div');
17+ this.#minLabel.className = 'colorbar-label';
18+ container.append(this.#maxLabel, this.#canvas, this.#minLabel);
19+ }
20+
21+ update(cmap: ColormapFunc, vmin: number, vmax: number): void {
22+ const ctx = this.#canvas.getContext('2d');
23+ if (!ctx) return;
24+ const h = this.#canvas.height;
25+ for (let y = 0; y < h; y++) {
26+ const t = 1 - y / (h - 1);
27+ const [r, g, b] = cmap(t);
28+ ctx.fillStyle = `rgb(${r},${g},${b})`;
29+ ctx.fillRect(0, y, this.#canvas.width, 1);
30+ }
31+ const fmt = (v: number) =>
32+ Number.isFinite(v) ? v.toPrecision(3).replace(/\.?0+$/, '') : '—';
33+ this.#maxLabel.textContent = fmt(vmax);
34+ this.#minLabel.textContent = fmt(vmin);
35+ }
36+}
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+186−0View file
@@ -0,0 +1,186 @@
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+ * Expand a field frame (nlat * nphi) to per-vertex values (numVertices),
130+ * with cap values averaged from the adjacent ring.
131+ */
132+export const fillFieldValues = (
133+ out: Float32Array,
134+ fieldFrame: Float32Array | Float64Array,
135+ topo: SphereMeshTopology,
136+): void => {
137+ const { nlat, nphi } = topo;
138+ const n = nlat * nphi;
139+ for (let p = 0; p < n; p++) {
140+ out[p] = fieldFrame[p];
141+ }
142+ const ringMean = (rowIndex: number) => {
143+ let sum = 0;
144+ let count = 0;
145+ for (let j = 0; j < nphi; j++) {
146+ const v = fieldFrame[rowIndex * nphi + j];
147+ if (!Number.isNaN(v)) {
148+ sum += v;
149+ count++;
150+ }
151+ }
152+ return count > 0 ? sum / count : NaN;
153+ };
154+ if (topo.startCapIndex >= 0) out[topo.startCapIndex] = ringMean(0);
155+ if (topo.endCapIndex >= 0) out[topo.endCapIndex] = ringMean(nlat - 1);
156+};
157+
158+/**
159+ * Fill the color buffer (numVertices * 3, floats in [0, 1]) from per-vertex
160+ * field values using the given colormap and range. NaN values render gray.
161+ */
162+export const fillColors = (
163+ out: Float32Array,
164+ values: Float32Array,
165+ valueMin: number,
166+ valueMax: number,
167+ cmap: ColormapFunc,
168+): void => {
169+ const span = valueMax - valueMin;
170+ const invSpan = span !== 0 ? 1 / span : 0;
171+ for (let i = 0; i < values.length; i++) {
172+ const v = values[i];
173+ const p = i * 3;
174+ if (Number.isNaN(v)) {
175+ out[p] = 0.35;
176+ out[p + 1] = 0.35;
177+ out[p + 2] = 0.35;
178+ } else {
179+ const t = span !== 0 ? (v - valueMin) * invSpan : 0.5;
180+ const [r, g, b] = cmap(t);
181+ out[p] = r / 255;
182+ out[p + 1] = g / 255;
183+ out[p + 2] = b / 255;
184+ }
185+ }
186+};
src/sht/coeffs.tsadded+85−0View file
@@ -0,0 +1,85 @@
1+/**
2+ * Recurrence coefficients for orthonormal associated Legendre functions
3+ * ytilde_l^m(theta) (spherical-harmonic normalized, Condon-Shortley phase
4+ * included), matching SHTNS legendre_precomp() with norm=sht_orthonormal:
5+ *
6+ * ytilde_m^m(theta) = amm * sin(theta)^m
7+ * ytilde_{m+1}^m = a_{m+1}^m * cos(theta) * ytilde_m^m
8+ * ytilde_l^m = a_l^m * cos(theta) * ytilde_{l-1}^m + b_l^m * ytilde_{l-2}^m
9+ *
10+ * with (cf. sht_legendre.c lines 442-447):
11+ * a_{m+1}^m = sqrt(2m+3)
12+ * a_l^m = sqrt( (2l+1)(2l-1) / ((l+m)(l-m)) )
13+ * b_l^m = -sqrt( (2l+1)/(2l-3) * ((l-1+m)(l-1-m)) / ((l+m)(l-m)) )
14+ * amm = cs^m * sqrt( 1/(4pi) * prod_{k=1..m} (2k+1)/(2k) )
15+ *
16+ * With this normalization, Y_lm(theta,phi) = ytilde_l^m(theta) e^{i m phi}
17+ * and integral |Y_lm|^2 dOmega = 1.
18+ */
19+import { lmIndex, nlmCalc } from './layout.ts';
20+
21+export interface LegendreCoeffs {
22+ /** amm[m]: seed value (includes Condon-Shortley phase (-1)^m). */
23+ amm: Float64Array;
24+ /** ab[2*lm], ab[2*lm+1] = (a_l^m, b_l^m); entries at l=m unused (0), b at l=m+1 unused (0). */
25+ ab: Float64Array;
26+}
27+
28+export function legendreCoeffs(lmax: number, mmax: number): LegendreCoeffs {
29+ const nlm = nlmCalc(lmax, mmax);
30+ const amm = new Float64Array(mmax + 1);
31+ const ab = new Float64Array(2 * nlm);
32+
33+ let t = 1.0 / (4.0 * Math.PI);
34+ amm[0] = Math.sqrt(t);
35+ for (let m = 1; m <= mmax; m++) {
36+ t *= (2 * m + 1) / (2 * m);
37+ amm[m] = -Math.sqrt(t); // (-1)^m accumulates: Condon-Shortley phase
38+ if (m % 2 === 0) amm[m] = -amm[m];
39+ }
40+
41+ for (let m = 0; m <= mmax; m++) {
42+ if (m + 1 <= lmax) {
43+ const lm = lmIndex(lmax, m + 1, m);
44+ ab[2 * lm] = Math.sqrt(2 * m + 3); // a_{m+1}^m
45+ ab[2 * lm + 1] = 0;
46+ }
47+ for (let l = m + 2; l <= lmax; l++) {
48+ const lm = lmIndex(lmax, l, m);
49+ const t1 = (l + m) * (l - m);
50+ const t2 = (l - 1 + m) * (l - 1 - m);
51+ ab[2 * lm] = Math.sqrt(((2 * l + 1) * (2 * l - 1)) / t1);
52+ ab[2 * lm + 1] = -Math.sqrt(((2 * l + 1) / (2 * l - 3)) * (t2 / t1));
53+ }
54+ }
55+ return { amm, ab };
56+}
57+
58+/**
59+ * Evaluate ytilde_l^m(theta) for l = m..lmax at one point, in f64.
60+ * ct = cos(theta), st = sin(theta). Plain (unscaled) recurrence: fine in
61+ * f64 for the moderate lmax this library targets (underflow of st^m only
62+ * matters for m of several hundred very close to the poles).
63+ */
64+export function legendreRow(
65+ coeffs: LegendreCoeffs,
66+ lmax: number,
67+ m: number,
68+ ct: number,
69+ st: number,
70+ out: Float64Array, // length lmax - m + 1
71+): void {
72+ let y0 = coeffs.amm[m] * Math.pow(st, m);
73+ out[0] = y0;
74+ if (m === lmax) return;
75+ const base = lmIndex(lmax, m, m);
76+ let y1 = coeffs.ab[2 * (base + 1)] * ct * y0;
77+ out[1] = y1;
78+ for (let l = m + 2; l <= lmax; l++) {
79+ const lm = base + (l - m);
80+ const y2 = coeffs.ab[2 * lm] * ct * y1 + coeffs.ab[2 * lm + 1] * y0;
81+ y0 = y1;
82+ y1 = y2;
83+ out[l - m] = y2;
84+ }
85+}
src/sht/gauss.tsadded+51−0View file
@@ -0,0 +1,51 @@
1+/**
2+ * Gauss-Legendre quadrature nodes and weights, computed in double
3+ * precision by Newton iteration on P_n (cf. gauss_nodes() in SHTNS
4+ * sht_legendre.c).
5+ *
6+ * Returns nodes x_i = cos(theta_i) in DECREASING order (theta increasing,
7+ * north pole first), and weights w_i for integration over x in [-1, 1]:
8+ * integral f(x) dx ~= sum_i w_i f(x_i), exact for polynomials of
9+ * degree <= 2n - 1.
10+ */
11+export function gaussNodesWeights(n: number): { x: Float64Array; w: Float64Array } {
12+ const x = new Float64Array(n);
13+ const w = new Float64Array(n);
14+ const m = (n + 1) >> 1;
15+ for (let i = 0; i < m; i++) {
16+ // initial guess (Tricomi-like), then Newton
17+ let z = Math.cos((Math.PI * (i + 0.75)) / (n + 0.5));
18+ let pp = 0;
19+ for (let iter = 0; iter < 100; iter++) {
20+ // evaluate P_n(z) and P_{n-1}(z) by recurrence
21+ let p1 = 1.0;
22+ let p2 = 0.0;
23+ for (let j = 1; j <= n; j++) {
24+ const p3 = p2;
25+ p2 = p1;
26+ p1 = ((2 * j - 1) * z * p2 - (j - 1) * p3) / j;
27+ }
28+ pp = (n * (z * p1 - p2)) / (z * z - 1.0);
29+ const dz = p1 / pp;
30+ z -= dz;
31+ if (Math.abs(dz) < 1e-15 * Math.abs(z) + 1e-300) {
32+ // one extra iteration for full convergence
33+ let q1 = 1.0, q2 = 0.0;
34+ for (let j = 1; j <= n; j++) {
35+ const q3 = q2; q2 = q1;
36+ q1 = ((2 * j - 1) * z * q2 - (j - 1) * q3) / j;
37+ }
38+ pp = (n * (z * q1 - q2)) / (z * z - 1.0);
39+ z -= q1 / pp;
40+ break;
41+ }
42+ }
43+ x[i] = z; // largest roots first => theta increasing
44+ x[n - 1 - i] = -z;
45+ const wi = 2.0 / ((1.0 - z * z) * pp * pp);
46+ w[i] = wi;
47+ w[n - 1 - i] = wi;
48+ }
49+ if (n & 1) x[m - 1] = 0.0; // exact for odd n
50+ return { x, w };
51+}
src/sht/layout.tsadded+47−0View file
@@ -0,0 +1,47 @@
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+}
src/sht/reference.tsadded+136−0View file
@@ -0,0 +1,136 @@
1+/**
2+ * Double-precision reference implementation of the scalar spherical
3+ * harmonic transform, by direct summation. Slow (O(nlat*nlm) Legendre +
4+ * O(nlat*nphi*mmax) Fourier) but simple, and serves as ground truth for
5+ * validating the fp32 WebGPU implementation.
6+ *
7+ * Conventions are identical to the GPU path (see layout.ts).
8+ */
9+import { gaussNodesWeights } from './gauss.ts';
10+import { legendreCoeffs, legendreRow, type LegendreCoeffs } from './coeffs.ts';
11+import { lmIndex, nlmCalc, validateConfig, type ShtConfig } from './layout.ts';
12+
13+export class ShtReference {
14+ readonly cfg: ShtConfig;
15+ readonly nlm: number;
16+ readonly ct: Float64Array;
17+ readonly st: Float64Array;
18+ readonly wg: Float64Array; // Gauss weights (for integral over cos(theta))
19+ readonly coeffs: LegendreCoeffs;
20+
21+ constructor(cfg: ShtConfig) {
22+ validateConfig(cfg);
23+ this.cfg = cfg;
24+ this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
25+ const { x, w } = gaussNodesWeights(cfg.nlat);
26+ this.ct = x;
27+ this.wg = w;
28+ this.st = new Float64Array(cfg.nlat);
29+ for (let i = 0; i < cfg.nlat; i++) this.st[i] = Math.sqrt(1 - x[i] * x[i]);
30+ this.coeffs = legendreCoeffs(cfg.lmax, cfg.mmax);
31+ }
32+
33+ /**
34+ * Legendre stage of the synthesis: F_m(theta_i) = sum_l Q_lm ytilde_l^m(theta_i).
35+ * Returns complex array indexed [m * nlat + ilat], interleaved re/im.
36+ */
37+ legendreSynth(qlm: ArrayLike<number>): Float64Array {
38+ const { lmax, mmax, nlat } = this.cfg;
39+ const fm = new Float64Array(2 * (mmax + 1) * nlat);
40+ const row = new Float64Array(lmax + 1);
41+ for (let i = 0; i < nlat; i++) {
42+ for (let m = 0; m <= mmax; m++) {
43+ legendreRow(this.coeffs, lmax, m, this.ct[i], this.st[i], row);
44+ let re = 0, im = 0;
45+ const base = lmIndex(lmax, m, m);
46+ for (let l = m; l <= lmax; l++) {
47+ const y = row[l - m];
48+ re += y * qlm[2 * (base + l - m)];
49+ im += y * qlm[2 * (base + l - m) + 1];
50+ }
51+ const o = 2 * (m * nlat + i);
52+ fm[o] = re;
53+ fm[o + 1] = im;
54+ }
55+ }
56+ return fm;
57+ }
58+
59+ /** Full synthesis: spectral -> spatial grid [ilat * nphi + iphi]. */
60+ synth(qlm: ArrayLike<number>): Float64Array {
61+ const { mmax, nlat, nphi } = this.cfg;
62+ const fm = this.legendreSynth(qlm);
63+ const spat = new Float64Array(nlat * nphi);
64+ for (let i = 0; i < nlat; i++) {
65+ for (let j = 0; j < nphi; j++) {
66+ const phi = (2 * Math.PI * j) / nphi;
67+ let v = fm[2 * (0 * nlat + i)]; // m=0: real part (imag must be 0)
68+ for (let m = 1; m <= mmax; m++) {
69+ const o = 2 * (m * nlat + i);
70+ const c = Math.cos(m * phi);
71+ const s = Math.sin(m * phi);
72+ v += 2 * (fm[o] * c - fm[o + 1] * s);
73+ }
74+ spat[i * nphi + j] = v;
75+ }
76+ }
77+ return spat;
78+ }
79+
80+ /** Full analysis: spatial grid -> spectral coefficients (interleaved re/im). */
81+ analys(spat: ArrayLike<number>): Float64Array {
82+ const { lmax, mmax, nlat, nphi } = this.cfg;
83+ const qlm = new Float64Array(2 * this.nlm);
84+ const row = new Float64Array(lmax + 1);
85+ // forward Fourier: G_m(theta_i) = (2*pi/nphi) * sum_j f_ij e^{-i m phi_j}
86+ const gm = new Float64Array(2 * (mmax + 1) * nlat);
87+ for (let i = 0; i < nlat; i++) {
88+ for (let m = 0; m <= mmax; m++) {
89+ let re = 0, im = 0;
90+ for (let j = 0; j < nphi; j++) {
91+ const phi = (2 * Math.PI * j) / nphi;
92+ const f = spat[i * nphi + j];
93+ re += f * Math.cos(m * phi);
94+ im -= f * Math.sin(m * phi);
95+ }
96+ const o = 2 * (m * nlat + i);
97+ const norm = (2 * Math.PI) / nphi;
98+ gm[o] = re * norm;
99+ gm[o + 1] = im * norm;
100+ }
101+ }
102+ // Legendre stage with Gauss quadrature: Q_lm = sum_i w_i ytilde_l^m(theta_i) G_m(theta_i)
103+ for (let m = 0; m <= mmax; m++) {
104+ const base = lmIndex(lmax, m, m);
105+ for (let i = 0; i < nlat; i++) {
106+ legendreRow(this.coeffs, lmax, m, this.ct[i], this.st[i], row);
107+ const o = 2 * (m * nlat + i);
108+ const wr = this.wg[i] * gm[o];
109+ const wi = this.wg[i] * gm[o + 1];
110+ for (let l = m; l <= lmax; l++) {
111+ const y = row[l - m];
112+ qlm[2 * (base + l - m)] += y * wr;
113+ qlm[2 * (base + l - m) + 1] += y * wi;
114+ }
115+ }
116+ }
117+ return qlm;
118+ }
119+}
120+
121+/** Random band-limited spectrum for testing (m=0 imaginary parts zeroed). */
122+export function randomSpectrum(cfg: ShtConfig, seed = 12345): Float32Array {
123+ const nlm = nlmCalc(cfg.lmax, cfg.mmax);
124+ const q = new Float32Array(2 * nlm);
125+ let s = seed >>> 0;
126+ const rnd = () => {
127+ // xorshift32
128+ s ^= s << 13; s >>>= 0;
129+ s ^= s >> 17;
130+ s ^= s << 5; s >>>= 0;
131+ return (s / 4294967296) * 2 - 1;
132+ };
133+ for (let k = 0; k < 2 * nlm; k++) q[k] = rnd();
134+ for (let l = 0; l <= cfg.lmax; l++) q[2 * lmIndex(cfg.lmax, l, 0) + 1] = 0; // m=0 real
135+ return q;
136+}
src/sht/sht.tsadded+288−0View file
@@ -0,0 +1,288 @@
1+/**
2+ * WebGPU spherical harmonic transform plan (scalar transforms, fp32).
3+ *
4+ * Mirrors the structure of the SHTNS CUDA backend (sht_gpu.cu):
5+ * host-side f64 precomputation of grid + recurrence coefficients, shader
6+ * source generated with sizes baked in (SHTNS uses NVRTC; WGSL is always
7+ * runtime-compiled), then per-transform: Legendre stage + Fourier stage.
8+ */
9+import { gaussNodesWeights } from './gauss.ts';
10+import { legendreCoeffs } from './coeffs.ts';
11+import { nlmCalc, validateConfig, isPowerOfTwo, type ShtConfig } from './layout.ts';
12+import { legSynthWGSL, legAnalysWGSL } from './wgsl/leg.ts';
13+import {
14+ fftSynthWGSL,
15+ fftAnalysWGSL,
16+ dftSynthWGSL,
17+ dftAnalysWGSL,
18+ fftThreads,
19+} from './wgsl/fourier.ts';
20+
21+export type FourierMode = 'auto' | 'fft' | 'dft';
22+
23+export interface ShtOptions {
24+ /** Fourier stage implementation. 'auto' picks fft when nphi is a power of two that fits in workgroup memory. */
25+ fourier?: FourierMode;
26+}
27+
28+const WG_SYNTH = 64;
29+const WG_ANALYS = 256;
30+
31+async function makePipeline(
32+ device: GPUDevice,
33+ code: string,
34+ entryPoint: string,
35+): Promise<GPUComputePipeline> {
36+ device.pushErrorScope('validation');
37+ const module = device.createShaderModule({ code, label: entryPoint });
38+ const info = await module.getCompilationInfo();
39+ const errors = info.messages.filter((m) => m.type === 'error');
40+ if (errors.length) {
41+ throw new Error(
42+ `WGSL compile error in ${entryPoint}:\n` +
43+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
44+ );
45+ }
46+ const pipeline = await device.createComputePipelineAsync({
47+ layout: 'auto',
48+ compute: { module, entryPoint },
49+ label: entryPoint,
50+ });
51+ const err = await device.popErrorScope();
52+ if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
53+ return pipeline;
54+}
55+
56+export class ShtPlan {
57+ readonly cfg: ShtConfig;
58+ readonly nlm: number;
59+ readonly fourierMode: 'fft' | 'dft';
60+ /** Colatitudes theta_i (f64, increasing: north to south). */
61+ readonly theta: Float64Array;
62+ readonly cosTheta: Float64Array;
63+ readonly gaussWeights: Float64Array;
64+
65+ private device: GPUDevice;
66+ private bufAb!: GPUBuffer;
67+ private bufAmm!: GPUBuffer;
68+ private bufCtstw!: GPUBuffer;
69+ private bufTrig!: GPUBuffer;
70+ /** Spectral input (synthesis) — write with queue.writeBuffer or use synth(). */
71+ readonly qlmIn!: GPUBuffer;
72+ /** Spectral output (analysis). */
73+ readonly qlmOut!: GPUBuffer;
74+ /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. */
75+ readonly fmBuf!: GPUBuffer;
76+ /** Spatial field [ilat*nphi + iphi], f32. */
77+ readonly spatBuf!: GPUBuffer;
78+ private stageSpat!: GPUBuffer;
79+ private stageQ!: GPUBuffer;
80+
81+ private pipeLegSynth!: GPUComputePipeline;
82+ private pipeLegAnalys!: GPUComputePipeline;
83+ private pipeFourSynth!: GPUComputePipeline;
84+ private pipeFourAnalys!: GPUComputePipeline;
85+ private bgLegSynth!: GPUBindGroup;
86+ private bgLegAnalys!: GPUBindGroup;
87+ private bgFourSynth!: GPUBindGroup;
88+ private bgFourAnalys!: GPUBindGroup;
89+
90+ private constructor(device: GPUDevice, cfg: ShtConfig, fourierMode: 'fft' | 'dft') {
91+ this.device = device;
92+ this.cfg = cfg;
93+ this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
94+ this.fourierMode = fourierMode;
95+ const { x, w } = gaussNodesWeights(cfg.nlat);
96+ this.cosTheta = x;
97+ this.gaussWeights = w;
98+ this.theta = new Float64Array(cfg.nlat);
99+ for (let i = 0; i < cfg.nlat; i++) this.theta[i] = Math.acos(x[i]);
100+ }
101+
102+ static async create(device: GPUDevice, cfg: ShtConfig, opts: ShtOptions = {}): Promise<ShtPlan> {
103+ validateConfig(cfg);
104+ const want = opts.fourier ?? 'auto';
105+ const fftFits =
106+ isPowerOfTwo(cfg.nphi) &&
107+ 16 * cfg.nphi <= device.limits.maxComputeWorkgroupStorageSize &&
108+ fftThreads(cfg.nphi) <= device.limits.maxComputeInvocationsPerWorkgroup;
109+ if (want === 'fft' && !fftFits) {
110+ throw new Error(
111+ `fourier:'fft' requires power-of-two nphi with 16*nphi <= maxComputeWorkgroupStorageSize ` +
112+ `(nphi=${cfg.nphi}, limit=${device.limits.maxComputeWorkgroupStorageSize})`,
113+ );
114+ }
115+ const mode: 'fft' | 'dft' = want === 'dft' ? 'dft' : fftFits ? 'fft' : 'dft';
116+ const plan = new ShtPlan(device, cfg, mode);
117+ await plan.init();
118+ return plan;
119+ }
120+
121+ private async init(): Promise<void> {
122+ const { lmax, mmax, nlat, nphi } = this.cfg;
123+ const dev = this.device;
124+ const self = this as {
125+ -readonly [k in keyof ShtPlan]: ShtPlan[k];
126+ };
127+
128+ // --- host precomputation (f64), then downcast to f32 for upload ---
129+ const { amm, ab } = legendreCoeffs(lmax, mmax);
130+ const ctstw = new Float32Array(3 * nlat);
131+ for (let i = 0; i < nlat; i++) {
132+ ctstw[i] = this.cosTheta[i];
133+ ctstw[nlat + i] = Math.sqrt(1 - this.cosTheta[i] * this.cosTheta[i]);
134+ ctstw[2 * nlat + i] = this.gaussWeights[i] * ((2 * Math.PI) / nphi);
135+ }
136+ // twiddle/phase table in f64 (device sin/cos is too inaccurate: ~2^-11 under Vulkan)
137+ const trig = new Float32Array(2 * nphi);
138+ for (let k = 0; k < nphi; k++) {
139+ trig[2 * k] = Math.cos((2 * Math.PI * k) / nphi);
140+ trig[2 * k + 1] = Math.sin((2 * Math.PI * k) / nphi);
141+ }
142+
143+ const mkBuf = (label: string, size: number, usage: GPUBufferUsageFlags) =>
144+ dev.createBuffer({ label, size, usage });
145+ this.bufAb = mkBuf('sht-ab', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
146+ this.bufAmm = mkBuf('sht-amm', 4 * (mmax + 1), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
147+ this.bufCtstw = mkBuf('sht-ctstw', 4 * 3 * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
148+ this.bufTrig = mkBuf('sht-trig', 8 * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
149+ self.qlmIn = mkBuf('sht-qlm-in', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
150+ self.qlmOut = mkBuf('sht-qlm-out', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
151+ self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE);
152+ self.spatBuf = mkBuf('sht-spat', 4 * nlat * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
153+ this.stageSpat = mkBuf('sht-stage-spat', 4 * nlat * nphi, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
154+ this.stageQ = mkBuf('sht-stage-q', 8 * this.nlm, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
155+
156+ dev.queue.writeBuffer(this.bufAb, 0, new Float32Array(ab));
157+ dev.queue.writeBuffer(this.bufAmm, 0, new Float32Array(amm));
158+ dev.queue.writeBuffer(this.bufCtstw, 0, ctstw);
159+ dev.queue.writeBuffer(this.bufTrig, 0, trig);
160+
161+ // --- shaders / pipelines ---
162+ const legP = { lmax, mmax, nlat, wgSynth: WG_SYNTH, wgAnalys: WG_ANALYS };
163+ const fourP = { mmax, nlat, nphi };
164+ const [pLegS, pLegA, pFourS, pFourA] = await Promise.all([
165+ makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),
166+ makePipeline(dev, legAnalysWGSL(legP), 'leg_analys'),
167+ makePipeline(
168+ dev,
169+ this.fourierMode === 'fft' ? fftSynthWGSL(fourP) : dftSynthWGSL(fourP),
170+ this.fourierMode === 'fft' ? 'fft_synth' : 'dft_synth',
171+ ),
172+ makePipeline(
173+ dev,
174+ this.fourierMode === 'fft' ? fftAnalysWGSL(fourP) : dftAnalysWGSL(fourP),
175+ this.fourierMode === 'fft' ? 'fft_analys' : 'dft_analys',
176+ ),
177+ ]);
178+ this.pipeLegSynth = pLegS;
179+ this.pipeLegAnalys = pLegA;
180+ this.pipeFourSynth = pFourS;
181+ this.pipeFourAnalys = pFourA;
182+
183+ const entries = (bufs: GPUBuffer[]) =>
184+ bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
185+ this.bgLegSynth = dev.createBindGroup({
186+ layout: pLegS.getBindGroupLayout(0),
187+ entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.qlmIn, this.fmBuf]),
188+ });
189+ this.bgLegAnalys = dev.createBindGroup({
190+ layout: pLegA.getBindGroupLayout(0),
191+ entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, this.qlmOut]),
192+ });
193+ this.bgFourSynth = dev.createBindGroup({
194+ layout: pFourS.getBindGroupLayout(0),
195+ entries: entries([this.fmBuf, this.spatBuf, this.bufTrig]),
196+ });
197+ this.bgFourAnalys = dev.createBindGroup({
198+ layout: pFourA.getBindGroupLayout(0),
199+ entries: entries([this.spatBuf, this.fmBuf, this.bufTrig]),
200+ });
201+ }
202+
203+ /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
204+ encodeSynth(encoder: GPUCommandEncoder): void {
205+ const { mmax, nlat, nphi } = this.cfg;
206+ const pass = encoder.beginComputePass({ label: 'sht-synth' });
207+ pass.setPipeline(this.pipeLegSynth);
208+ pass.setBindGroup(0, this.bgLegSynth);
209+ pass.dispatchWorkgroups(Math.ceil(nlat / WG_SYNTH), mmax + 1);
210+ pass.setPipeline(this.pipeFourSynth);
211+ pass.setBindGroup(0, this.bgFourSynth);
212+ if (this.fourierMode === 'fft') {
213+ pass.dispatchWorkgroups(nlat);
214+ } else {
215+ pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
216+ }
217+ pass.end();
218+ }
219+
220+ /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
221+ encodeAnalys(encoder: GPUCommandEncoder): void {
222+ const { mmax, nlat } = this.cfg;
223+ const pass = encoder.beginComputePass({ label: 'sht-analys' });
224+ pass.setPipeline(this.pipeFourAnalys);
225+ pass.setBindGroup(0, this.bgFourAnalys);
226+ if (this.fourierMode === 'fft') {
227+ pass.dispatchWorkgroups(nlat);
228+ } else {
229+ pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
230+ }
231+ pass.setPipeline(this.pipeLegAnalys);
232+ pass.setBindGroup(0, this.bgLegAnalys);
233+ pass.dispatchWorkgroups(mmax + 1);
234+ pass.end();
235+ }
236+
237+ /**
238+ * Spectral -> spatial. qlm: interleaved [re, im], SHTNS LM ordering,
239+ * length 2*nlm. Returns the spatial field, length nlat*nphi.
240+ */
241+ async synth(qlm: Float32Array): Promise<Float32Array> {
242+ const { nlat, nphi } = this.cfg;
243+ if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
244+ this.device.queue.writeBuffer(this.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
245+ const enc = this.device.createCommandEncoder();
246+ this.encodeSynth(enc);
247+ enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
248+ this.device.queue.submit([enc.finish()]);
249+ await this.stageSpat.mapAsync(GPUMapMode.READ);
250+ const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
251+ this.stageSpat.unmap();
252+ return out;
253+ }
254+
255+ /** Spatial -> spectral. spat: length nlat*nphi. Returns interleaved qlm, length 2*nlm. */
256+ async analys(spat: Float32Array): Promise<Float32Array> {
257+ const { nlat, nphi } = this.cfg;
258+ if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
259+ this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
260+ const enc = this.device.createCommandEncoder();
261+ this.encodeAnalys(enc);
262+ enc.copyBufferToBuffer(this.qlmOut, 0, this.stageQ, 0, 8 * this.nlm);
263+ this.device.queue.submit([enc.finish()]);
264+ await this.stageQ.mapAsync(GPUMapMode.READ);
265+ const out = new Float32Array(this.stageQ.getMappedRange().slice(0));
266+ this.stageQ.unmap();
267+ return out;
268+ }
269+
270+ destroy(): void {
271+ for (const b of [
272+ this.bufAb, this.bufAmm, this.bufCtstw, this.bufTrig, this.qlmIn, this.qlmOut,
273+ this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ,
274+ ]) b?.destroy();
275+ }
276+}
277+
278+/** Request an adapter/device suitable for the transforms. */
279+export async function requestShtDevice(): Promise<GPUDevice> {
280+ if (!navigator.gpu) throw new Error('WebGPU is not available in this browser');
281+ const adapter = await navigator.gpu.requestAdapter();
282+ if (!adapter) throw new Error('No WebGPU adapter available');
283+ // ask for a larger workgroup storage if the adapter offers it (bigger FFTs)
284+ const wgStorage = Math.min(adapter.limits.maxComputeWorkgroupStorageSize, 32768);
285+ return adapter.requestDevice({
286+ requiredLimits: { maxComputeWorkgroupStorageSize: wgStorage },
287+ });
288+}
src/sht/wgsl/common.tsadded+56−0View file
@@ -0,0 +1,56 @@
1+/**
2+ * Shared WGSL fragments. Shaders are generated as strings with all sizes
3+ * baked in as compile-time constants (the WGSL analog of what SHTNS does
4+ * with NVRTC on CUDA: cf. init_cuda_program() in sht_gpu.cu).
5+ *
6+ * fp32 extended-range constants: same values SHTNS injects for a
7+ * single-precision recurrence (sht_gpu.cu):
8+ * SHT_ACCURACY = 1e-15
9+ * SHT_SCALE_FACTOR = 2^56 = 7.2057594037927936e16
10+ * A per-thread integer exponent `ny` counts how many times the running
11+ * Legendre value has been multiplied by SCALE to stay in fp32 range;
12+ * contributions are only accumulated once ny == 0 (value back in normal
13+ * range and significant).
14+ */
15+
16+export const RESCALE_WGSL = /* wgsl */ `
17+const SCALE: f32 = 7.2057594e16; // rounds to exactly 2^56 in f32
18+const INV_SCALE: f32 = 1.0 / 7.2057594e16;
19+const ACCURACY: f32 = 1e-15;
20+const RESCALE_THR: f32 = ACCURACY * SCALE + 1.0; // ~73: value became significant again
21+
22+struct Seed { y0: f32, ny: i32 }
23+
24+// Seed of the recurrence: y0 ~ sin(theta)^m by binary exponentiation with
25+// rescaling (ports the HI_LLIM path of SHT/cuda_legendre.gen.cu, ~651-691).
26+// The caller multiplies by amm afterwards (|amm| is O(1)).
27+fn sinpow_rescaled(st: f32, m: u32) -> Seed {
28+ var y0: f32 = 1.0;
29+ var ny: i32 = 0;
30+ if (m > 0u) {
31+ var s: f32 = st;
32+ var lb: u32 = m;
33+ if ((lb & 1u) != 0u) { y0 = s; }
34+ var nsint: i32 = 0;
35+ lb = lb >> 1u;
36+ while (lb > 0u) {
37+ s = s * s;
38+ nsint = nsint + nsint;
39+ if (s < INV_SCALE) {
40+ nsint = nsint - 1;
41+ s = s * SCALE;
42+ }
43+ if ((lb & 1u) != 0u) {
44+ y0 = y0 * s;
45+ ny = ny + nsint;
46+ if (y0 < (ACCURACY + INV_SCALE)) {
47+ y0 = y0 * SCALE;
48+ ny = ny - 1;
49+ }
50+ }
51+ lb = lb >> 1u;
52+ }
53+ }
54+ return Seed(y0, ny);
55+}
56+`;
src/sht/wgsl/fourier.tsadded+194−0View file
@@ -0,0 +1,194 @@
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+}
27+
28+const TRIG_BINDING = /* wgsl */ `
29+@group(0) @binding(2) var<storage, read> trig: array<vec2f>; // (cos,sin)(2*pi*k/NPHI), k < NPHI
30+`;
31+
32+function stockham(nphi: number, threads: number, sign: number): string {
33+ const log2n = Math.log2(nphi);
34+ if (!Number.isInteger(log2n)) throw new Error('fft requires power-of-two nphi');
35+ // twiddle for pass with half-block ns: w = e^{sign*i*pi*j/ns} = T[j * (N/(2*ns))]^sign
36+ return /* wgsl */ `
37+var<workgroup> bufA: array<vec2f, ${nphi}>;
38+var<workgroup> bufB: array<vec2f, ${nphi}>;
39+
40+fn cmul(a: vec2f, b: vec2f) -> vec2f {
41+ return vec2f(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
42+}
43+
44+fn ld(sel: u32, i: u32) -> vec2f {
45+ if (sel == 0u) { return bufA[i]; }
46+ return bufB[i];
47+}
48+fn st_(sel: u32, i: u32, v: vec2f) {
49+ if (sel == 0u) { bufA[i] = v; } else { bufB[i] = v; }
50+}
51+
52+// radix-2 Stockham, natural order in and out; data starts in bufA (sel 0)
53+// and ends in sel = LOG2N % 2. Unnormalized: X_k = sum_j x_j e^{s*2*pi*i*jk/N}.
54+fn fft_inplace(lid: u32) {
55+ for (var p = 0u; p < ${log2n}u; p++) {
56+ workgroupBarrier();
57+ let ns = 1u << p;
58+ let sel = p & 1u;
59+ let stride = ${nphi / 2}u >> p; // N/(2*ns)
60+ for (var t = lid; t < ${nphi / 2}u; t += ${threads}u) {
61+ let j = t & (ns - 1u);
62+ let tw = trig[j * stride];
63+ let w = vec2f(tw.x, ${sign > 0 ? '' : '-'}tw.y);
64+ let u = ld(sel, t);
65+ let v = cmul(ld(sel, t + ${nphi / 2}u), w);
66+ let idst = 2u * (t - j) + j;
67+ st_(1u - sel, idst, u + v);
68+ st_(1u - sel, idst + ns, u - v);
69+ }
70+ }
71+ workgroupBarrier();
72+}
73+const FFT_OUT_SEL: u32 = ${log2n % 2}u;
74+`;
75+}
76+
77+/** Choose FFT workgroup size: enough threads for the butterflies, capped at 256. */
78+export function fftThreads(nphi: number): number {
79+ return Math.max(32, Math.min(256, nphi / 2));
80+}
81+
82+export function fftSynthWGSL(p: FourierParams): string {
83+ const T = fftThreads(p.nphi);
84+ return /* wgsl */ `
85+const MMAX: u32 = ${p.mmax}u;
86+const NLAT: u32 = ${p.nlat}u;
87+const NPHI: u32 = ${p.nphi}u;
88+@group(0) @binding(0) var<storage, read> fm: array<vec2f>;
89+@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
90+${TRIG_BINDING}
91+${stockham(p.nphi, T, +1)}
92+
93+@compute @workgroup_size(${T})
94+fn fft_synth(@builtin(local_invocation_id) lid3: vec3u,
95+ @builtin(workgroup_id) wid: vec3u) {
96+ let lid = lid3.x;
97+ let ilat = wid.x;
98+ // assemble Hermitian spectrum: X[0] = Re F_0, X[m] = F_m, X[N-m] = conj(F_m)
99+ for (var k = lid; k < NPHI; k += ${T}u) {
100+ var v = vec2f(0.0);
101+ if (k == 0u) {
102+ v = vec2f(fm[ilat].x, 0.0);
103+ } else if (k <= MMAX) {
104+ v = fm[k * NLAT + ilat];
105+ } else if (k >= NPHI - MMAX) {
106+ let c = fm[(NPHI - k) * NLAT + ilat];
107+ v = vec2f(c.x, -c.y);
108+ }
109+ bufA[k] = v;
110+ }
111+ fft_inplace(lid);
112+ for (var k = lid; k < NPHI; k += ${T}u) {
113+ spat[ilat * NPHI + k] = ld(FFT_OUT_SEL, k).x;
114+ }
115+}
116+`;
117+}
118+
119+export function fftAnalysWGSL(p: FourierParams): string {
120+ const T = fftThreads(p.nphi);
121+ return /* wgsl */ `
122+const MMAX: u32 = ${p.mmax}u;
123+const NLAT: u32 = ${p.nlat}u;
124+const NPHI: u32 = ${p.nphi}u;
125+@group(0) @binding(0) var<storage, read> spat: array<f32>;
126+@group(0) @binding(1) var<storage, read_write> fm: array<vec2f>;
127+${TRIG_BINDING}
128+${stockham(p.nphi, T, -1)}
129+
130+@compute @workgroup_size(${T})
131+fn fft_analys(@builtin(local_invocation_id) lid3: vec3u,
132+ @builtin(workgroup_id) wid: vec3u) {
133+ let lid = lid3.x;
134+ let ilat = wid.x;
135+ for (var k = lid; k < NPHI; k += ${T}u) {
136+ bufA[k] = vec2f(spat[ilat * NPHI + k], 0.0);
137+ }
138+ fft_inplace(lid);
139+ for (var m = lid; m <= MMAX; m += ${T}u) {
140+ fm[m * NLAT + ilat] = ld(FFT_OUT_SEL, m);
141+ }
142+}
143+`;
144+}
145+
146+export function dftSynthWGSL(p: FourierParams): string {
147+ return /* wgsl */ `
148+const MMAX: u32 = ${p.mmax}u;
149+const NLAT: u32 = ${p.nlat}u;
150+const NPHI: u32 = ${p.nphi}u;
151+@group(0) @binding(0) var<storage, read> fm: array<vec2f>;
152+@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
153+${TRIG_BINDING}
154+
155+@compute @workgroup_size(64)
156+fn dft_synth(@builtin(global_invocation_id) gid: vec3u) {
157+ let iphi = gid.x;
158+ let ilat = gid.y;
159+ if (iphi >= NPHI) { return; }
160+ var v: f32 = fm[ilat].x; // m = 0: real part
161+ for (var m = 1u; m <= MMAX; m++) {
162+ let w = trig[(m * iphi) % NPHI]; // e^{+i m phi}
163+ let c = fm[m * NLAT + ilat];
164+ v += 2.0 * (c.x * w.x - c.y * w.y);
165+ }
166+ spat[ilat * NPHI + iphi] = v;
167+}
168+`;
169+}
170+
171+export function dftAnalysWGSL(p: FourierParams): string {
172+ return /* wgsl */ `
173+const MMAX: u32 = ${p.mmax}u;
174+const NLAT: u32 = ${p.nlat}u;
175+const NPHI: u32 = ${p.nphi}u;
176+@group(0) @binding(0) var<storage, read> spat: array<f32>;
177+@group(0) @binding(1) var<storage, read_write> fm: array<vec2f>;
178+${TRIG_BINDING}
179+
180+@compute @workgroup_size(64)
181+fn dft_analys(@builtin(global_invocation_id) gid: vec3u) {
182+ let m = gid.x;
183+ let ilat = gid.y;
184+ if (m > MMAX) { return; }
185+ var acc = vec2f(0.0);
186+ for (var j = 0u; j < NPHI; j++) {
187+ let w = trig[(m * j) % NPHI]; // conj => e^{-i m phi}
188+ let f = spat[ilat * NPHI + j];
189+ acc += f * vec2f(w.x, -w.y);
190+ }
191+ fm[m * NLAT + ilat] = acc;
192+}
193+`;
194+}
src/sht/wgsl/leg.tsadded+182−0View file
@@ -0,0 +1,182 @@
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+}
27+
28+const BINDINGS = /* wgsl */ `
29+@group(0) @binding(0) var<storage, read> ab: array<vec2f>; // (a_l^m, b_l^m) per lm
30+@group(0) @binding(1) var<storage, read> amm: array<f32>; // seed per m
31+@group(0) @binding(2) var<storage, read> ctstw: array<f32>; // [ct | st | w], each NLAT
32+`;
33+
34+export function legSynthWGSL(p: LegParams): string {
35+ return /* wgsl */ `
36+${RESCALE_WGSL}
37+const LMAX: u32 = ${p.lmax}u;
38+const NLAT: u32 = ${p.nlat}u;
39+${BINDINGS}
40+@group(0) @binding(3) var<storage, read> qlm: array<vec2f>;
41+@group(0) @binding(4) var<storage, read_write> fm: array<vec2f>; // [(m)*NLAT + ilat]
42+
43+@compute @workgroup_size(${p.wgSynth})
44+fn leg_synth(@builtin(global_invocation_id) gid: vec3u,
45+ @builtin(workgroup_id) wid: vec3u) {
46+ let ilat = gid.x;
47+ let m = wid.y;
48+ if (ilat >= NLAT) { return; }
49+
50+ let ct = ctstw[ilat];
51+ let st = ctstw[NLAT + ilat];
52+ let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u; // lm index of (l=m, m)
53+
54+ var seed = sinpow_rescaled(st, m);
55+ var y0 = seed.y0 * amm[m];
56+ var ny = seed.ny;
57+ var y1: f32 = 0.0;
58+ if (m < LMAX) {
59+ y1 = ab[base + 1u].x * ct * y0;
60+ }
61+
62+ var acc = vec2f(0.0);
63+ var l = m;
64+ loop {
65+ if (ny == 0) {
66+ acc += y0 * qlm[base + (l - m)];
67+ if (l + 1u <= LMAX) {
68+ acc += y1 * qlm[base + (l + 1u - m)];
69+ }
70+ } else if (abs(y0) > RESCALE_THR) {
71+ ny += 1;
72+ y0 *= INV_SCALE;
73+ y1 *= INV_SCALE;
74+ }
75+ if (l + 2u > LMAX) { break; }
76+ let c0 = ab[base + (l + 2u - m)];
77+ y0 = c0.x * ct * y1 + c0.y * y0;
78+ if (l + 3u <= LMAX) {
79+ let c1 = ab[base + (l + 3u - m)];
80+ y1 = c1.x * ct * y0 + c1.y * y1;
81+ }
82+ l += 2u;
83+ }
84+ fm[m * NLAT + ilat] = acc;
85+}
86+`;
87+}
88+
89+export function legAnalysWGSL(p: LegParams): string {
90+ const K = Math.ceil(p.nlat / p.wgAnalys); // latitudes per thread
91+ return /* wgsl */ `
92+${RESCALE_WGSL}
93+const LMAX: u32 = ${p.lmax}u;
94+const NLAT: u32 = ${p.nlat}u;
95+const WG: u32 = ${p.wgAnalys}u;
96+const K: u32 = ${K}u;
97+${BINDINGS}
98+@group(0) @binding(3) var<storage, read> fm: array<vec2f>; // [(m)*NLAT + ilat]
99+@group(0) @binding(4) var<storage, read_write> qout: array<vec2f>;
100+
101+var<workgroup> red: array<vec4f, ${p.wgAnalys}>;
102+
103+@compute @workgroup_size(${p.wgAnalys})
104+fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,
105+ @builtin(workgroup_id) wid: vec3u) {
106+ let lid = lid3.x;
107+ let m = wid.x;
108+ let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
109+
110+ // per-thread recurrence state for K latitudes
111+ var y0v: array<f32, ${K}>;
112+ var y1v: array<f32, ${K}>;
113+ var nyv: array<i32, ${K}>;
114+ var ctv: array<f32, ${K}>;
115+ var wfv: array<vec2f, ${K}>;
116+
117+ for (var k = 0u; k < K; k++) {
118+ let lat = lid + k * WG;
119+ var ct: f32 = 0.0;
120+ var st: f32 = 0.0;
121+ var wf = vec2f(0.0);
122+ if (lat < NLAT) {
123+ ct = ctstw[lat];
124+ st = ctstw[NLAT + lat];
125+ wf = fm[m * NLAT + lat] * ctstw[2u * NLAT + lat]; // Gauss weight (incl. 2*pi/nphi)
126+ }
127+ ctv[k] = ct;
128+ let seed = sinpow_rescaled(st, m);
129+ y0v[k] = seed.y0 * amm[m];
130+ nyv[k] = seed.ny;
131+ y1v[k] = 0.0;
132+ if (m < LMAX) {
133+ y1v[k] = ab[base + 1u].x * ct * y0v[k];
134+ }
135+ wfv[k] = wf;
136+ }
137+
138+ var l = m;
139+ loop {
140+ var c0 = vec2f(0.0);
141+ var c1 = vec2f(0.0);
142+ for (var k = 0u; k < K; k++) {
143+ if (nyv[k] == 0) {
144+ c0 += wfv[k] * y0v[k];
145+ c1 += wfv[k] * y1v[k];
146+ } else if (abs(y0v[k]) > RESCALE_THR) {
147+ nyv[k] += 1;
148+ y0v[k] *= INV_SCALE;
149+ y1v[k] *= INV_SCALE;
150+ }
151+ }
152+ // workgroup tree reduction of (c0, c1)
153+ red[lid] = vec4f(c0, c1);
154+ workgroupBarrier();
155+ var s = WG / 2u;
156+ while (s > 0u) {
157+ if (lid < s) { red[lid] += red[lid + s]; }
158+ workgroupBarrier();
159+ s = s >> 1u;
160+ }
161+ if (lid == 0u) {
162+ qout[base + (l - m)] = red[0].xy;
163+ if (l + 1u <= LMAX) {
164+ qout[base + (l + 1u - m)] = red[0].zw;
165+ }
166+ }
167+ if (l + 2u > LMAX) { break; }
168+ let a0 = ab[base + (l + 2u - m)];
169+ var a1 = vec2f(0.0);
170+ if (l + 3u <= LMAX) {
171+ a1 = ab[base + (l + 3u - m)];
172+ }
173+ for (var k = 0u; k < K; k++) {
174+ let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
175+ y0v[k] = t0;
176+ y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
177+ }
178+ l += 2u;
179+ }
180+}
181+`;
182+}
src/solver/backend.tsadded+87−0View file
@@ -0,0 +1,87 @@
1+/**
2+ * Backend abstraction over the spherical harmonic transform, mirroring the
3+ * websph "porting boundary": the solver only ever needs coeffs->vals,
4+ * vals->coeffs, and the grid. Spectral layout is the SHTNS convention used
5+ * by shtns-webgpu (see src/sht/layout.ts): complex interleaved [re, im],
6+ * m >= 0 only, m-major ordering, orthonormal + Condon-Shortley.
7+ */
8+import { ShtPlan, requestShtDevice } from '../sht/sht.ts';
9+import { ShtReference } from '../sht/reference.ts';
10+import type { ShtConfig } from '../sht/layout.ts';
11+
12+export interface ShtBackend {
13+ readonly cfg: ShtConfig;
14+ readonly nlm: number;
15+ /** cos(colatitude), length nlat, decreasing (north to south). */
16+ readonly cosTheta: Float64Array;
17+ readonly kind: 'webgpu' | 'cpu';
18+ synth(qlm: Float64Array): Promise<Float32Array | Float64Array>;
19+ analys(spat: Float64Array): Promise<Float32Array | Float64Array>;
20+ destroy(): void;
21+}
22+
23+/** fp32 WebGPU backend (fast path). */
24+export class GpuBackend implements ShtBackend {
25+ readonly kind = 'webgpu';
26+ readonly cfg: ShtConfig;
27+ readonly nlm: number;
28+ readonly cosTheta: Float64Array;
29+ #plan: ShtPlan;
30+ #qlm32: Float32Array;
31+ #spat32: Float32Array;
32+
33+ private constructor(plan: ShtPlan) {
34+ this.#plan = plan;
35+ this.cfg = plan.cfg;
36+ this.nlm = plan.nlm;
37+ this.cosTheta = plan.cosTheta;
38+ this.#qlm32 = new Float32Array(2 * plan.nlm);
39+ this.#spat32 = new Float32Array(plan.cfg.nlat * plan.cfg.nphi);
40+ }
41+
42+ static async create(device: GPUDevice, cfg: ShtConfig): Promise<GpuBackend> {
43+ return new GpuBackend(await ShtPlan.create(device, cfg));
44+ }
45+
46+ synth(qlm: Float64Array): Promise<Float32Array> {
47+ this.#qlm32.set(qlm);
48+ return this.#plan.synth(this.#qlm32);
49+ }
50+
51+ analys(spat: Float64Array): Promise<Float32Array> {
52+ this.#spat32.set(spat);
53+ return this.#plan.analys(this.#spat32);
54+ }
55+
56+ destroy(): void {
57+ this.#plan.destroy();
58+ }
59+}
60+
61+/** f64 CPU backend by direct summation (slow; tests and no-WebGPU fallback). */
62+export class CpuBackend implements ShtBackend {
63+ readonly kind = 'cpu';
64+ readonly cfg: ShtConfig;
65+ readonly nlm: number;
66+ readonly cosTheta: Float64Array;
67+ #ref: ShtReference;
68+
69+ constructor(cfg: ShtConfig) {
70+ this.#ref = new ShtReference(cfg);
71+ this.cfg = cfg;
72+ this.nlm = this.#ref.nlm;
73+ this.cosTheta = this.#ref.ct;
74+ }
75+
76+ synth(qlm: Float64Array): Promise<Float64Array> {
77+ return Promise.resolve(this.#ref.synth(qlm));
78+ }
79+
80+ analys(spat: Float64Array): Promise<Float64Array> {
81+ return Promise.resolve(this.#ref.analys(spat));
82+ }
83+
84+ destroy(): void {}
85+}
86+
87+export { requestShtDevice };
src/solver/models.tsadded+198−0View file
@@ -0,0 +1,198 @@
1+/**
2+ * Reaction-diffusion model presets, ported from websph's
3+ * SphericalReactionDiffusionDriver.m. Each species k solves
4+ *
5+ * d(u_k)/dt = D_k * lap_s(u_k) + f_k(t, x, y, z, u_1, ..., u_N)
6+ *
7+ * on the unit sphere. Reactions are vectorized over the grid.
8+ */
9+
10+export type Params = Record<string, number>;
11+
12+export interface ParamSpec {
13+ key: string;
14+ label: string;
15+ value: number;
16+ min: number;
17+ max: number;
18+ step: number;
19+}
20+
21+export interface ModelSpec {
22+ key: string;
23+ label: string;
24+ blurb: string;
25+ species: string[];
26+ params: ParamSpec[];
27+ /** Polynomial degree of the reaction in the fields (for dealiasing). */
28+ pdeg: number;
29+ /** Amplitude of the random perturbation seeded into the initial state. */
30+ seedAmp: number;
31+ diffusivities(p: Params): number[];
32+ /** Fill out[k][i] with f_k evaluated at every grid point. */
33+ reaction(
34+ p: Params,
35+ t: number,
36+ x: Float64Array,
37+ y: Float64Array,
38+ z: Float64Array,
39+ V: ArrayLike<number>[],
40+ out: Float64Array[],
41+ ): void;
42+ /** Fill out[k][i] with the initial condition (noise added via randn). */
43+ init(
44+ p: Params,
45+ x: Float64Array,
46+ y: Float64Array,
47+ z: Float64Array,
48+ randn: () => number,
49+ out: Float64Array[],
50+ ): void;
51+}
52+
53+const schnakenberg: ModelSpec = {
54+ key: 'schnakenberg',
55+ label: 'Schnakenberg',
56+ blurb:
57+ 'Turing spots. The homogeneous state is stable to uniform perturbations ' +
58+ 'but unstable to degrees 14 ≤ l ≤ 40, most strongly at l = 24.',
59+ species: ['u', 'v'],
60+ params: [
61+ { key: 'a', label: 'a', value: 0.1, min: 0.01, max: 0.5, step: 0.01 },
62+ { key: 'b', label: 'b', value: 0.9, min: 0.1, max: 2, step: 0.05 },
63+ { key: 'D1', label: 'D₁', value: 4e-4, min: 1e-5, max: 5e-3, step: 1e-5 },
64+ { key: 'D2', label: 'D₂', value: 8e-3, min: 1e-4, max: 5e-2, step: 1e-4 },
65+ { key: 'dt', label: 'dt', value: 0.05, min: 0.005, max: 0.5, step: 0.005 },
66+ ],
67+ pdeg: 3,
68+ seedAmp: 1e-2,
69+ diffusivities: (p) => [p.D1, p.D2],
70+ reaction(p, _t, _x, _y, _z, V, out) {
71+ const [u, v] = V;
72+ const [fu, fv] = out;
73+ const a = p.a, b = p.b;
74+ const n = fu.length;
75+ for (let i = 0; i < n; i++) {
76+ const ui = u[i], vi = v[i];
77+ const uuv = ui * ui * vi;
78+ fu[i] = a - ui + uuv;
79+ fv[i] = b - uuv;
80+ }
81+ },
82+ init(p, x, _y, _z, randn, out) {
83+ const [u, v] = out;
84+ const a = p.a, b = p.b;
85+ const us = a + b;
86+ const vs = b / (us * us);
87+ const n = x.length;
88+ for (let i = 0; i < n; i++) {
89+ u[i] = us + this.seedAmp * randn();
90+ v[i] = vs;
91+ }
92+ },
93+};
94+
95+const brusselator: ModelSpec = {
96+ key: 'brusselator',
97+ label: 'Brusselator',
98+ blurb:
99+ 'Turing stripes and spots, from a smaller diffusivity contrast than ' +
100+ 'Schnakenberg but with a stiffer reaction.',
101+ species: ['u', 'v'],
102+ params: [
103+ { key: 'A', label: 'A', value: 3, min: 0.5, max: 6, step: 0.1 },
104+ { key: 'B', label: 'B', value: 9, min: 1, max: 15, step: 0.25 },
105+ { key: 'D1', label: 'D₁', value: 3.33e-3, min: 1e-4, max: 2e-2, step: 1e-4 },
106+ { key: 'D2', label: 'D₂', value: 1.67e-2, min: 1e-3, max: 1e-1, step: 1e-3 },
107+ { key: 'dt', label: 'dt', value: 0.02, min: 0.002, max: 0.1, step: 0.002 },
108+ ],
109+ pdeg: 3,
110+ seedAmp: 1e-2,
111+ diffusivities: (p) => [p.D1, p.D2],
112+ reaction(p, _t, _x, _y, _z, V, out) {
113+ const [u, v] = V;
114+ const [fu, fv] = out;
115+ const A = p.A, B = p.B;
116+ const n = fu.length;
117+ for (let i = 0; i < n; i++) {
118+ const ui = u[i], vi = v[i];
119+ const uuv = ui * ui * vi;
120+ fu[i] = A - (B + 1) * ui + uuv;
121+ fv[i] = B * ui - uuv;
122+ }
123+ },
124+ init(p, x, _y, _z, randn, out) {
125+ const [u, v] = out;
126+ const n = x.length;
127+ for (let i = 0; i < n; i++) {
128+ u[i] = p.A + this.seedAmp * randn();
129+ v[i] = p.B / p.A;
130+ }
131+ },
132+};
133+
134+const allenCahn: ModelSpec = {
135+ key: 'allencahn',
136+ label: 'Allen–Cahn',
137+ blurb:
138+ 'A single species: interfaces form and then coarsen until one domain ' +
139+ 'swallows the sphere.',
140+ species: ['u'],
141+ params: [
142+ { key: 'eps2', label: 'ε²', value: 1e-3, min: 1e-4, max: 1e-2, step: 1e-4 },
143+ { key: 'dt', label: 'dt', value: 0.02, min: 0.002, max: 0.2, step: 0.002 },
144+ ],
145+ pdeg: 3,
146+ seedAmp: 1e-2,
147+ diffusivities: (p) => [p.eps2],
148+ reaction(_p, _t, _x, _y, _z, V, out) {
149+ const [u] = V;
150+ const [fu] = out;
151+ const n = fu.length;
152+ for (let i = 0; i < n; i++) {
153+ const ui = u[i];
154+ fu[i] = ui - ui * ui * ui;
155+ }
156+ },
157+ init(_p, x, _y, _z, randn, out) {
158+ const [u] = out;
159+ const n = x.length;
160+ for (let i = 0; i < n; i++) {
161+ u[i] = this.seedAmp * randn();
162+ }
163+ },
164+};
165+
166+export const models: ModelSpec[] = [schnakenberg, brusselator, allenCahn];
167+
168+export const defaultParams = (m: ModelSpec): Params =>
169+ Object.fromEntries(m.params.map((p) => [p.key, p.value]));
170+
171+/** Named parameter presets shown in the UI dropdown. The pattern length
172+ * scale goes as 1/sqrt(D), so scaling both diffusivities moves the spot
173+ * size without changing the dynamics. */
174+export interface Preset {
175+ key: string;
176+ label: string;
177+ modelKey: string;
178+ /** Overrides applied on top of the model's default parameters. */
179+ params?: Params;
180+}
181+
182+export const presets: Preset[] = [
183+ { key: 'schnak-spots', label: 'Schnakenberg — spots', modelKey: 'schnakenberg' },
184+ {
185+ key: 'schnak-coarse',
186+ label: 'Schnakenberg — coarse spots',
187+ modelKey: 'schnakenberg',
188+ params: { D1: 1e-3, D2: 2e-2 },
189+ },
190+ {
191+ key: 'schnak-fine',
192+ label: 'Schnakenberg — fine spots',
193+ modelKey: 'schnakenberg',
194+ params: { D1: 1.6e-4, D2: 3.2e-3 },
195+ },
196+ { key: 'brussel', label: 'Brusselator — stripes & spots', modelKey: 'brusselator' },
197+ { key: 'allencahn', label: 'Allen–Cahn — coarsening', modelKey: 'allencahn' },
198+];
src/solver/simulation.tsadded+164−0View file
@@ -0,0 +1,164 @@
1+/**
2+ * IMEX Euler reaction-diffusion timestepper on the sphere, ported from
3+ * websph's SphericalReactionDiffusion.m. Diffusion is implicit and diagonal
4+ * in spherical-harmonic space (Laplace-Beltrami eigenvalues -l(l+1));
5+ * reaction is explicit on the grid:
6+ *
7+ * (I - dt*D_k*lap_s) u_k^{n+1} = u_k^n + dt*f_k(u^n)
8+ */
9+import type { ShtBackend } from './backend.ts';
10+import type { ModelSpec, Params } from './models.ts';
11+import { lmIndex } from '../sht/layout.ts';
12+
13+/** Grid sizes for a given lmax, dealiased for a reaction of degree pdeg
14+ * (see websph README): nlat >= ((pdeg+1)*lmax+1)/2, nlon >= (pdeg+1)*lmax+1.
15+ * nphi is rounded up to a power of two to keep the GPU FFT path. */
16+export function gridForLmax(lmax: number, pdeg: number): { nlat: number; nphi: number } {
17+ const minLat = Math.max(lmax + 1, ((pdeg + 1) * lmax + 1) / 2);
18+ const nlat = 2 * Math.ceil(minLat / 2);
19+ let nphi = 1;
20+ while (nphi < (pdeg + 1) * lmax + 1) nphi *= 2;
21+ return { nlat, nphi };
22+}
23+
24+/** Seeded normal deviates: mulberry32 + Box-Muller. */
25+export function makeRandn(seed: number): () => number {
26+ let s = seed >>> 0;
27+ const rand = () => {
28+ s = (s + 0x6d2b79f5) >>> 0;
29+ let t = s;
30+ t = Math.imul(t ^ (t >>> 15), t | 1);
31+ t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
32+ return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
33+ };
34+ let spare: number | null = null;
35+ return () => {
36+ if (spare !== null) {
37+ const v = spare;
38+ spare = null;
39+ return v;
40+ }
41+ let u = 0;
42+ while (u === 0) u = rand();
43+ const r = Math.sqrt(-2 * Math.log(u));
44+ const th = 2 * Math.PI * rand();
45+ spare = r * Math.sin(th);
46+ return r * Math.cos(th);
47+ };
48+}
49+
50+export class Simulation {
51+ readonly backend: ShtBackend;
52+ readonly model: ModelSpec;
53+ readonly params: Params;
54+ readonly nspecies: number;
55+
56+ /** Spectral state, one interleaved-complex Float64Array (2*nlm) per species. */
57+ U: Float64Array[];
58+ /** Grid values per species as of the START of the last step (one step
59+ * behind U; recomputed as the first stage of the next step). */
60+ V: (Float32Array | Float64Array)[];
61+ t = 0;
62+ stepCount = 0;
63+
64+ /** Cartesian coordinates of the grid points (nlat*nphi each). */
65+ readonly x: Float64Array;
66+ readonly y: Float64Array;
67+ readonly z: Float64Array;
68+ /** Laplace-Beltrami eigenvalues l*(l+1) per spectral index (length nlm). */
69+ readonly lam: Float64Array;
70+
71+ #R: Float64Array[]; // reaction scratch, one grid array per species
72+ #m0Imag: number[]; // interleaved-array positions of m=0 imaginary parts
73+
74+ constructor(backend: ShtBackend, model: ModelSpec, params: Params) {
75+ this.backend = backend;
76+ this.model = model;
77+ this.params = params;
78+ this.nspecies = model.species.length;
79+
80+ const { lmax, mmax, nlat, nphi } = backend.cfg;
81+ const npts = nlat * nphi;
82+ this.x = new Float64Array(npts);
83+ this.y = new Float64Array(npts);
84+ this.z = new Float64Array(npts);
85+ for (let i = 0; i < nlat; i++) {
86+ const ct = backend.cosTheta[i];
87+ const st = Math.sqrt(Math.max(0, 1 - ct * ct));
88+ for (let j = 0; j < nphi; j++) {
89+ const phi = (2 * Math.PI * j) / nphi;
90+ const idx = i * nphi + j;
91+ this.x[idx] = st * Math.cos(phi);
92+ this.y[idx] = st * Math.sin(phi);
93+ this.z[idx] = ct;
94+ }
95+ }
96+
97+ this.lam = new Float64Array(backend.nlm);
98+ for (let m = 0; m <= mmax; m++) {
99+ for (let l = m; l <= lmax; l++) {
100+ this.lam[lmIndex(lmax, l, m)] = l * (l + 1);
101+ }
102+ }
103+ this.#m0Imag = [];
104+ for (let l = 0; l <= lmax; l++) {
105+ this.#m0Imag.push(2 * lmIndex(lmax, l, 0) + 1);
106+ }
107+
108+ this.U = [];
109+ this.V = [];
110+ this.#R = [];
111+ for (let k = 0; k < this.nspecies; k++) {
112+ this.U.push(new Float64Array(2 * backend.nlm));
113+ this.V.push(new Float64Array(npts));
114+ this.#R.push(new Float64Array(npts));
115+ }
116+ }
117+
118+ /** Project the initial conditions (band-limiting the seed noise). */
119+ async init(seed: number): Promise<void> {
120+ const grids = this.#R;
121+ this.model.init(this.params, this.x, this.y, this.z, makeRandn(seed), grids);
122+ for (let k = 0; k < this.nspecies; k++) {
123+ const q = await this.backend.analys(grids[k]);
124+ this.U[k].set(q);
125+ this.#cleanM0(this.U[k]);
126+ this.V[k] = await this.backend.synth(this.U[k]);
127+ }
128+ this.t = 0;
129+ this.stepCount = 0;
130+ }
131+
132+ /** One IMEX Euler step. */
133+ async step(): Promise<void> {
134+ const dt = this.params.dt;
135+ const D = this.model.diffusivities(this.params);
136+
137+ // Evaluate every species on the grid before reacting any of them
138+ for (let k = 0; k < this.nspecies; k++) {
139+ this.V[k] = await this.backend.synth(this.U[k]);
140+ }
141+
142+ this.model.reaction(this.params, this.t, this.x, this.y, this.z, this.V, this.#R);
143+
144+ for (let k = 0; k < this.nspecies; k++) {
145+ const Rlm = await this.backend.analys(this.#R[k]);
146+ const U = this.U[k];
147+ const dD = dt * D[k];
148+ for (let i = 0; i < this.backend.nlm; i++) {
149+ const fac = 1 / (1 + dD * this.lam[i]);
150+ U[2 * i] = (U[2 * i] + dt * Rlm[2 * i]) * fac;
151+ U[2 * i + 1] = (U[2 * i + 1] + dt * Rlm[2 * i + 1]) * fac;
152+ }
153+ this.#cleanM0(U);
154+ }
155+
156+ this.t += dt;
157+ this.stepCount++;
158+ }
159+
160+ /** m=0 coefficients of a real field are purely real; drop numerical junk. */
161+ #cleanM0(U: Float64Array): void {
162+ for (const p of this.#m0Imag) U[p] = 0;
163+ }
164+}
test.htmladded+12−0View file
@@ -0,0 +1,12 @@
1+<!doctype html>
2+<html lang="en">
3+ <head>
4+ <meta charset="utf-8" />
5+ <title>turing-sphere validation</title>
6+ </head>
7+ <body>
8+ <h1>turing-sphere validation suite</h1>
9+ <pre id="log">starting…</pre>
10+ <script type="module" src="/test/test-page.ts"></script>
11+ </body>
12+</html>
test/test-page.tsadded+130−0View file
@@ -0,0 +1,130 @@
1+/**
2+ * Browser validation: the fp32 WebGPU solver against the f64 CPU solver.
3+ * Runs identical seeded simulations on both backends and compares fields.
4+ * Results are posted to window.__RESULTS__ for the headless runner.
5+ */
6+import { GpuBackend, CpuBackend, requestShtDevice } from '../src/solver/backend.ts';
7+import { Simulation, gridForLmax } from '../src/solver/simulation.ts';
8+import { models, defaultParams } from '../src/solver/models.ts';
9+import { randomSpectrum } from '../src/sht/reference.ts';
10+
11+declare global {
12+ interface Window {
13+ __RESULTS__?: { ok: boolean; fatal?: string; lines: string[] };
14+ }
15+}
16+
17+const logEl = document.getElementById('log')!;
18+const lines: string[] = [];
19+let failures = 0;
20+
21+function log(s: string): void {
22+ lines.push(s);
23+ logEl.textContent = lines.join('\n');
24+ console.log(s);
25+}
26+
27+function check(name: string, ok: boolean, detail: string): void {
28+ log(`${ok ? 'PASS' : 'FAIL'} ${name} ${detail}`);
29+ if (!ok) failures++;
30+}
31+
32+function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
33+ let num = 0;
34+ let den = 0;
35+ for (let i = 0; i < a.length; i++) {
36+ const d = a[i] - b[i];
37+ num += d * d;
38+ den += b[i] * b[i];
39+ }
40+ return Math.sqrt(num / Math.max(den, 1e-300));
41+}
42+
43+async function main(): Promise<void> {
44+ const device = await requestShtDevice();
45+
46+ // --- transform cross-check: GPU vs CPU on a random spectrum ---
47+ {
48+ const lmax = 31;
49+ const { nlat, nphi } = gridForLmax(lmax, 1);
50+ const cfg = { lmax, mmax: lmax, nlat, nphi };
51+ const gpu = await GpuBackend.create(device, cfg);
52+ const cpu = new CpuBackend(cfg);
53+ const q = randomSpectrum(cfg, 42);
54+ const q64 = new Float64Array(q);
55+ const sGpu = await gpu.synth(q64);
56+ const sCpu = await cpu.synth(q64);
57+ const errSynth = relL2(sGpu, sCpu);
58+ const aGpu = await gpu.analys(new Float64Array(sCpu));
59+ const aCpu = await cpu.analys(new Float64Array(sCpu));
60+ const errAnalys = relL2(aGpu, aCpu);
61+ check('transforms: GPU vs CPU', errSynth < 1e-4 && errAnalys < 1e-4,
62+ `synth ${errSynth.toExponential(2)}, analys ${errAnalys.toExponential(2)}`);
63+ gpu.destroy();
64+ }
65+
66+ // --- solver cross-check: identical seeded runs on both backends ---
67+ {
68+ const schnak = models[0];
69+ const params = defaultParams(schnak);
70+ const lmax = 31;
71+ const { nlat, nphi } = gridForLmax(lmax, schnak.pdeg);
72+ const cfg = { lmax, mmax: lmax, nlat, nphi };
73+ const gpu = await GpuBackend.create(device, cfg);
74+ const cpu = new CpuBackend(cfg);
75+ const simGpu = new Simulation(gpu, schnak, { ...params });
76+ const simCpu = new Simulation(cpu, schnak, { ...params });
77+ await simGpu.init(7);
78+ await simCpu.init(7);
79+ const nsteps = 10;
80+ const t0 = performance.now();
81+ for (let s = 0; s < nsteps; s++) await simGpu.step();
82+ const gpuMs = (performance.now() - t0) / nsteps;
83+ for (let s = 0; s < nsteps; s++) await simCpu.step();
84+ let worst = 0;
85+ for (let k = 0; k < simGpu.nspecies; k++) {
86+ worst = Math.max(worst, relL2(simGpu.V[k], simCpu.V[k]));
87+ }
88+ let nan = false;
89+ for (let k = 0; k < simGpu.nspecies; k++) {
90+ for (const v of simGpu.V[k]) if (!Number.isFinite(v)) nan = true;
91+ }
92+ check('solver: GPU vs CPU after 10 steps', worst < 2e-3 && !nan,
93+ `worst rel L2 ${worst.toExponential(2)}${nan ? ', NaN!' : ''} (${gpuMs.toFixed(1)} ms/step GPU)`);
94+ gpu.destroy();
95+ }
96+
97+ // --- longer GPU-only run stays finite and patterned ---
98+ {
99+ const schnak = models[0];
100+ const params = defaultParams(schnak);
101+ const lmax = 63;
102+ const { nlat, nphi } = gridForLmax(lmax, schnak.pdeg);
103+ const gpu = await GpuBackend.create(device, { lmax, mmax: lmax, nlat, nphi });
104+ const sim = new Simulation(gpu, schnak, params);
105+ await sim.init(3);
106+ const nsteps = 100;
107+ const t0 = performance.now();
108+ for (let s = 0; s < nsteps; s++) await sim.step();
109+ const ms = (performance.now() - t0) / nsteps;
110+ let lo = Infinity;
111+ let hi = -Infinity;
112+ for (const v of sim.V[0]) {
113+ if (v < lo) lo = v;
114+ if (v > hi) hi = v;
115+ }
116+ const finite = Number.isFinite(lo) && Number.isFinite(hi);
117+ check('solver: 100 steps at lmax 63 stay finite', finite && lo > -10 && hi < 10,
118+ `u range [${lo.toFixed(4)}, ${hi.toFixed(4)}], ${ms.toFixed(1)} ms/step`);
119+ gpu.destroy();
120+ }
121+
122+ window.__RESULTS__ = { ok: failures === 0, lines };
123+ log(failures === 0 ? 'ALL PASS' : `${failures} FAILURE(S)`);
124+}
125+
126+main().catch((e) => {
127+ const msg = e instanceof Error ? `${e.message}\n${e.stack ?? ''}` : String(e);
128+ log(`fatal: ${msg}`);
129+ window.__RESULTS__ = { ok: false, fatal: msg, lines };
130+});
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+15−0View file
@@ -0,0 +1,15 @@
1+import { defineConfig } from 'vite';
2+import { resolve } from 'node:path';
3+
4+export default defineConfig({
5+ base: './',
6+ build: {
7+ target: 'es2022',
8+ rollupOptions: {
9+ input: {
10+ main: resolve(import.meta.dirname, 'index.html'),
11+ test: resolve(import.meta.dirname, 'test.html'),
12+ },
13+ },
14+ },
15+});