shtns-webgpu: fp32 spherical harmonic transforms on WebGPU
Scalar synthesis and analysis of real fields on a Gauss grid, modeled on
the SHTNS CUDA backend: on-the-fly Legendre recurrence with the SHTNS
extended-range fp32 rescaling, Stockham FFT / band-limited DFT Fourier
stage with f64 host-computed twiddle tables, f64 reference implementation
and headless-Chrome validation suite.
24 changed files+4052−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+3−0View file
@@ -0,0 +1,3 @@
1+node_modules/
2+dist/
3+*.log
README.mdadded+115−0View file
@@ -0,0 +1,115 @@
1+# shtns-webgpu
2+
3+Spherical harmonic transforms on **WebGPU** (browser, fp32), modeled on
4+[SHTNS](https://nschaeff.bitbucket.io/shtns/). This is a from-scratch
5+TypeScript + WGSL implementation of the scalar transforms, structured after
6+the SHTNS CUDA backend (`sht_gpu.cu` / `SHT/cuda_legendre.gen.cu`).
7+
8+**Live demo:** https://concept-collection.github.io/shtns-webgpu/
9+(validation suite: [test.html](https://concept-collection.github.io/shtns-webgpu/test.html))
10+
11+## Scope (v0.1)
12+
13+- **Scalar transforms of real fields**, both directions:
14+ - `synth()` — spectral → spatial (SHTNS `SH_to_spat`)
15+ - `analys()` — spatial → spectral (SHTNS `spat_to_SH`)
16+- **Gauss–Legendre grid** in latitude, uniform longitude grid.
17+- **fp32 on the GPU** end to end; all precomputation (Gauss nodes/weights,
18+ recurrence coefficients, twiddle factors) is done host-side in f64.
19+- Not (yet) implemented: vector (spheroidal/toroidal) transforms, complex
20+ fields, regular grids, `mres > 1`, Schmidt/4π normalizations, on-the-fly
21+ truncation (`llim < lmax`).
22+
23+## Conventions (= SHTNS defaults)
24+
25+- **Orthonormal** spherical harmonics **with Condon–Shortley phase**.
26+- Spectral coefficients `Q_lm` are complex, stored for `m >= 0` with the
27+ SHTNS LM ordering: `for m = 0..mmax: for l = m..lmax`, interleaved
28+ `[re, im]` (`Float32Array` of length `2*nlm`). Real fields imply
29+ `Q_{l,-m} = (-1)^m conj(Q_lm)`; `m = 0` coefficients must be real.
30+- Spatial fields are `Float32Array[nlat * nphi]`, **phi-contiguous**,
31+ latitudes ordered north → south (colatitude increasing).
32+
33+## Usage
34+
35+```ts
36+import { ShtPlan, requestShtDevice, lmIndex } from 'shtns-webgpu';
37+
38+const device = await requestShtDevice(); // or your own GPUDevice
39+const plan = await ShtPlan.create(device, { lmax: 127, mmax: 127, nlat: 128, nphi: 256 });
40+
41+const qlm = new Float32Array(2 * plan.nlm);
42+qlm[2 * lmIndex(127, 8, 5)] = 1.0; // Y_8^5
43+
44+const spat = await plan.synth(qlm); // nlat*nphi field
45+const qBack = await plan.analys(spat); // back to spectral
46+plan.destroy();
47+```
48+
49+For GPU-resident pipelines (no readback), use `plan.encodeSynth(encoder)` /
50+`plan.encodeAnalys(encoder)` with the exposed `qlmIn` / `qlmOut` / `spatBuf`
51+buffers.
52+
53+Constraints checked at plan creation: `nlat > lmax` (Gauss quadrature
54+exactness), `nphi >= 2*mmax + 1` (no aliasing).
55+
56+## How it works
57+
58+Same two-stage split as SHTNS:
59+
60+1. **Legendre stage** (`src/wgsl/leg.ts`): associated Legendre functions are
61+ generated *on the fly* inside the shader by the standard 3-term
62+ recurrence over `l` (coefficients from `legendre_precomp()`-equivalent
63+ host code, `src/coeffs.ts`). Underflow of `sin^m(theta)` — fatal in fp32
64+ beyond `m ≈ 75` — is handled with the SHTNS extended-range scheme
65+ (`SHT_SCALE_FACTOR = 2^56`, `SHT_ACCURACY = 1e-15`, per-thread integer
66+ exponent), ported from the `HI_LLIM` path of `SHT/cuda_legendre.gen.cu`.
67+ Synthesis runs one thread per latitude and one workgroup row per `m`;
68+ analysis runs one workgroup per `m` with a shared-memory tree reduction
69+ over latitudes (the portable equivalent of SHTNS's warp shuffles).
70+2. **Fourier stage** (`src/wgsl/fourier.ts`): batched radix-2 Stockham FFT
71+ in workgroup memory (one workgroup per latitude row) when `nphi` is a
72+ power of two that fits (`16*nphi <= maxComputeWorkgroupStorageSize`);
73+ otherwise a direct band-limited DFT. Twiddles come from a host-computed
74+ f64 table — device `sin`/`cos` is only guaranteed to ~2^-11 under
75+ Vulkan, which would otherwise dominate the error budget.
76+
77+All problem sizes are baked into the WGSL at plan creation (the WGSL
78+equivalent of SHTNS's NVRTC runtime compilation).
79+
80+## Accuracy (fp32)
81+
82+Relative L2 errors vs the double-precision reference (`src/reference.ts`),
83+random spectra, measured on SwiftShader (results on hardware GPUs are the
84+same to within noise since the arithmetic is IEEE fp32):
85+
86+| lmax | synthesis | analysis | round trip |
87+|-----:|----------:|---------:|-----------:|
88+| 15 | 6e-7 | 3e-7 | 6e-7 |
89+| 63 | 5e-6 | 2e-6 | 3e-6 |
90+| 127 | 7e-6 | 3e-6 | 5e-6 |
91+| 255 | 2e-5 | 6e-6 | 1e-5 |
92+| 399 | 9e-5 | 1e-5 | 2e-5 |
93+
94+SHTNS itself switches its fp32 GPU recurrence to f64 above `lmax = 128`
95+(`SHT_L_RESCALE_FLY_FLOAT`); WGSL has no f64, so past that point accuracy
96+degrades gracefully as above. Fine for visualization; for scientific use
97+keep `lmax ≲ 128` or wait for the float-float recurrence (planned).
98+
99+## Develop / test
100+
101+```sh
102+npm install
103+npm run dev # demo at http://localhost:5173
104+npm run test:node # f64 math tests (no GPU needed)
105+npm run test:gpu # builds, then runs the browser suite in headless Chrome
106+ # (falls back to SwiftShader software WebGPU; CHROME_PATH to override)
107+```
108+
109+## Roadmap
110+
111+- Vector transforms (spheroidal/toroidal), gradients — port of `leg_m_kernel<1>`.
112+- Float-float (double-single) recurrence option for full accuracy at high lmax.
113+- Latitude parity folding (2× Legendre work reduction, as in SHTNS).
114+- Subgroup (warp) reductions where available, replacing the shared-memory tree.
115+- `mres > 1`, truncated transforms, regular grids.
index.htmladded+102−0View file
@@ -0,0 +1,102 @@
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+ <title>shtns-webgpu — spherical harmonic transforms in the browser</title>
7+ <style>
8+ :root {
9+ --bg: #ffffff;
10+ --ink: #1f2328;
11+ --ink-2: #57606a;
12+ --line: #d0d7de;
13+ --accent: #0969da;
14+ color-scheme: light dark;
15+ }
16+ @media (prefers-color-scheme: dark) {
17+ :root {
18+ --bg: #14171a;
19+ --ink: #e6e9ec;
20+ --ink-2: #9aa4af;
21+ --line: #333b44;
22+ --accent: #58a6ff;
23+ }
24+ }
25+ body {
26+ margin: 0;
27+ background: var(--bg);
28+ color: var(--ink);
29+ font: 15px/1.5 system-ui, -apple-system, sans-serif;
30+ }
31+ main { max-width: 940px; margin: 0 auto; padding: 20px 16px 48px; }
32+ h1 { font-size: 20px; margin: 0 0 2px; }
33+ .sub { color: var(--ink-2); margin: 0 0 16px; font-size: 13px; }
34+ .controls {
35+ display: flex; flex-wrap: wrap; gap: 10px 14px; align-items: center;
36+ padding: 10px 0 14px;
37+ }
38+ .controls label { color: var(--ink-2); font-size: 13px; }
39+ select, input[type="number"], button {
40+ font: inherit; font-size: 13px;
41+ color: var(--ink); background: var(--bg);
42+ border: 1px solid var(--line); border-radius: 6px;
43+ padding: 4px 8px;
44+ }
45+ input[type="number"] { width: 5em; }
46+ button { cursor: pointer; }
47+ button:hover { border-color: var(--accent); }
48+ #map {
49+ width: 100%; display: block; border: 1px solid var(--line); border-radius: 6px;
50+ background: var(--bg);
51+ }
52+ .bar-row { display: flex; align-items: center; gap: 10px; margin-top: 10px; }
53+ #colorbar { border: 1px solid var(--line); border-radius: 3px; }
54+ .bar-row span { font-size: 12px; color: var(--ink-2); font-variant-numeric: tabular-nums; }
55+ .stats { margin-top: 10px; font-size: 13px; color: var(--ink-2); }
56+ .stats b { color: var(--ink); font-weight: 600; }
57+ #hover { min-height: 1.2em; font-variant-numeric: tabular-nums; }
58+ #err { color: #b35900; white-space: pre-wrap; }
59+ </style>
60+ </head>
61+ <body>
62+ <main>
63+ <h1>shtns-webgpu</h1>
64+ <p class="sub">
65+ fp32 spherical harmonic transforms on WebGPU, modeled on
66+ <a href="https://nschaeff.bitbucket.io/shtns/">SHTNS</a> — scalar synthesis
67+ and analysis on a Gauss grid.
68+ </p>
69+ <div class="controls">
70+ <label>size
71+ <select id="preset">
72+ <option value="31">lmax 31 (32×64)</option>
73+ <option value="63">lmax 63 (64×128)</option>
74+ <option value="127" selected>lmax 127 (128×256)</option>
75+ <option value="255">lmax 255 (256×512)</option>
76+ </select>
77+ </label>
78+ <label>source
79+ <select id="source">
80+ <option value="single" selected>single mode Y<sub>l</sub><sup>m</sup></option>
81+ <option value="random">random spectrum</option>
82+ <option value="bandpass">random, l ∈ [lmax/3, 2·lmax/3]</option>
83+ </select>
84+ </label>
85+ <label id="lmWrap">l <input id="inpL" type="number" min="0" value="8" />
86+ m <input id="inpM" type="number" min="0" value="5" /></label>
87+ <button id="reroll">new random seed</button>
88+ </div>
89+ <canvas id="map"></canvas>
90+ <div class="bar-row">
91+ <span id="vmin"></span>
92+ <canvas id="colorbar" width="256" height="12"></canvas>
93+ <span id="vmax"></span>
94+ <span style="flex:1"></span>
95+ <span id="hover"></span>
96+ </div>
97+ <p class="stats" id="stats"></p>
98+ <p class="stats" id="err"></p>
99+ </main>
100+ <script type="module" src="/src/demo/main.ts"></script>
101+ </body>
102+</html>
package-lock.jsonadded+2118−0View file
This diff is 2,123 lines long and is not shown.
package.jsonadded+21−0View file
@@ -0,0 +1,21 @@
1+{
2+ "name": "shtns-webgpu",
3+ "version": "0.1.0",
4+ "description": "Spherical harmonic transforms on WebGPU (fp32), modeled on SHTNS",
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+ "devDependencies": {
15+ "@types/node": "^26.1.1",
16+ "@webgpu/types": "^0.1.44",
17+ "puppeteer-core": "^23.0.0",
18+ "typescript": "^5.5.0",
19+ "vite": "^5.4.0"
20+ }
21+}
scripts/screenshot.mjsadded+51−0View file
@@ -0,0 +1,51 @@
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 DIST = new URL('../dist/', import.meta.url).pathname;
10+const MIME = { '.html': 'text/html', '.js': 'text/javascript', '.css': 'text/css' };
11+
12+const server = createServer(async (req, res) => {
13+ try {
14+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
15+ const data = await readFile(join(DIST, path));
16+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
17+ res.end(data);
18+ } catch {
19+ res.writeHead(404);
20+ res.end();
21+ }
22+});
23+await new Promise((r) => server.listen(0, '127.0.0.1', r));
24+const port = server.address().port;
25+
26+const browser = await puppeteer.launch({
27+ executablePath: process.env.CHROME_PATH ?? '/usr/bin/google-chrome',
28+ args: ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu',
29+ '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
30+});
31+const page = await browser.newPage();
32+await page.setViewport({ width: 1000, height: 900 });
33+await page.emulateMediaFeatures([{ name: 'prefers-color-scheme', value: scheme }]);
34+page.on('console', (m) => console.log(' [page]', m.text()));
35+await page.goto(`http://127.0.0.1:${port}/index.html`, { waitUntil: 'load' });
36+await page.waitForFunction(
37+ () => {
38+ const s = document.getElementById('stats');
39+ const e = document.getElementById('err');
40+ return (s && /round-trip/.test(s.textContent)) || (e && e.textContent.length > 4);
41+ },
42+ { timeout: 120_000 },
43+);
44+await new Promise((r) => setTimeout(r, 300));
45+await page.screenshot({ path: out });
46+console.log('screenshot:', out);
47+console.log('stats:', await page.$eval('#stats', (el) => el.textContent));
48+const err = await page.$eval('#err', (el) => el.textContent);
49+if (err) console.log('err:', err);
50+await browser.close();
51+server.close();
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+111−0View file
@@ -0,0 +1,111 @@
1+/**
2+ * Node-side tests for the double-precision planner and reference transform.
3+ * Run: node scripts/test-node.ts
4+ */
5+import { gaussNodesWeights } from '../src/gauss.ts';
6+import { legendreCoeffs, legendreRow } from '../src/coeffs.ts';
7+import { lmIndex } from '../src/layout.ts';
8+import { ShtReference, randomSpectrum } from '../src/reference.ts';
9+
10+let failures = 0;
11+function check(name: string, ok: boolean, detail = '') {
12+ console.log(`${ok ? 'PASS' : 'FAIL'} ${name}${detail ? ' (' + detail + ')' : ''}`);
13+ if (!ok) failures++;
14+}
15+
16+// --- Gauss quadrature ---
17+{
18+ const { x, w } = gaussNodesWeights(64);
19+ let sw = 0, sx2 = 0;
20+ for (let i = 0; i < 64; i++) {
21+ sw += w[i];
22+ sx2 += w[i] * x[i] * x[i];
23+ }
24+ check('gauss: sum(w) == 2', Math.abs(sw - 2) < 1e-13, `err=${Math.abs(sw - 2).toExponential(2)}`);
25+ check('gauss: int x^2 == 2/3', Math.abs(sx2 - 2 / 3) < 1e-13, `err=${Math.abs(sx2 - 2 / 3).toExponential(2)}`);
26+ check('gauss: nodes decreasing', x[0] > x[1] && x[1] > x[63]);
27+}
28+
29+// --- Analytic Legendre values (orthonormal + Condon-Shortley) ---
30+{
31+ const lmax = 8;
32+ const coeffs = legendreCoeffs(lmax, lmax);
33+ const theta = 0.7;
34+ const ct = Math.cos(theta), st = Math.sin(theta);
35+ const row = new Float64Array(lmax + 1);
36+
37+ legendreRow(coeffs, lmax, 0, ct, st, row);
38+ const y00 = Math.sqrt(1 / (4 * Math.PI));
39+ const y10 = Math.sqrt(3 / (4 * Math.PI)) * ct;
40+ const y20 = Math.sqrt(5 / (16 * Math.PI)) * (3 * ct * ct - 1);
41+ check('Y_0^0', Math.abs(row[0] - y00) < 1e-14);
42+ check('Y_1^0', Math.abs(row[1] - y10) < 1e-14);
43+ check('Y_2^0', Math.abs(row[2] - y20) < 1e-14);
44+
45+ legendreRow(coeffs, lmax, 1, ct, st, row);
46+ const y11 = -Math.sqrt(3 / (8 * Math.PI)) * st; // CS phase => negative
47+ const y21 = -Math.sqrt(15 / (8 * Math.PI)) * st * ct;
48+ check('Y_1^1 (CS phase)', Math.abs(row[0] - y11) < 1e-14, `got ${row[0]}, want ${y11}`);
49+ check('Y_2^1', Math.abs(row[1] - y21) < 1e-14);
50+
51+ legendreRow(coeffs, lmax, 2, ct, st, row);
52+ const y22 = Math.sqrt(15 / (32 * Math.PI)) * st * st;
53+ check('Y_2^2', Math.abs(row[0] - y22) < 1e-14);
54+}
55+
56+// --- Orthonormality under Gauss quadrature ---
57+{
58+ const lmax = 42, nlat = 48;
59+ const coeffs = legendreCoeffs(lmax, lmax);
60+ const { x, w } = gaussNodesWeights(nlat);
61+ let worst = 0;
62+ for (const m of [0, 1, 7, 25]) {
63+ const rowsI = new Float64Array(lmax + 1);
64+ const rowsJ = new Float64Array(lmax + 1);
65+ for (const [la, lb] of [[m, m], [m, m + 3], [lmax, lmax], [m + 1, lmax]] as const) {
66+ if (la > lmax || lb > lmax) continue;
67+ let s = 0;
68+ for (let i = 0; i < nlat; i++) {
69+ const st = Math.sqrt(1 - x[i] * x[i]);
70+ legendreRow(coeffs, lmax, m, x[i], st, rowsI);
71+ legendreRow(coeffs, lmax, m, x[i], st, rowsJ);
72+ s += w[i] * rowsI[la - m] * rowsJ[lb - m];
73+ }
74+ const want = la === lb ? 1 / (2 * Math.PI) : 0;
75+ worst = Math.max(worst, Math.abs(s - want));
76+ }
77+ }
78+ check('orthonormality: max err < 1e-12', worst < 1e-12, `worst=${worst.toExponential(2)}`);
79+}
80+
81+// --- Reference round trip ---
82+{
83+ const cfg = { lmax: 31, mmax: 31, nlat: 34, nphi: 64 };
84+ const ref = new ShtReference(cfg);
85+ const q0 = randomSpectrum(cfg, 999);
86+ const spat = ref.synth(q0);
87+ const q1 = ref.analys(spat);
88+ let num = 0, den = 0;
89+ for (let k = 0; k < q0.length; k++) {
90+ num += (q1[k] - q0[k]) ** 2;
91+ den += q0[k] ** 2;
92+ }
93+ const rel = Math.sqrt(num / den);
94+ check('reference round trip rel L2 < 1e-12', rel < 1e-12, `rel=${rel.toExponential(2)}`);
95+}
96+
97+// --- Mean value: Y_00 coefficient of a constant field ---
98+{
99+ const cfg = { lmax: 15, mmax: 15, nlat: 16, nphi: 32 };
100+ const ref = new ShtReference(cfg);
101+ const spat = new Float64Array(cfg.nlat * cfg.nphi).fill(1.0);
102+ const q = ref.analys(spat);
103+ const want = Math.sqrt(4 * Math.PI); // <1, Y00> = sqrt(4pi)
104+ check('constant field -> Q_00 = sqrt(4pi)', Math.abs(q[0] - want) < 1e-12, `got ${q[0]}`);
105+ let rest = 0;
106+ for (let k = 2; k < q.length; k++) rest = Math.max(rest, Math.abs(q[k]));
107+ check('constant field -> other coeffs ~ 0', rest < 1e-12, `max=${rest.toExponential(2)}`);
108+}
109+
110+console.log(failures === 0 ? '\nALL NODE TESTS PASSED' : `\n${failures} TEST(S) FAILED`);
111+process.exit(failures === 0 ? 0 : 1);
src/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/demo/main.tsadded+191−0View file
@@ -0,0 +1,191 @@
1+/**
2+ * Demo: synthesize a spherical harmonic field on the GPU and render it on
3+ * an equirectangular map, with a round-trip (analysis) error readout.
4+ */
5+import { ShtPlan, requestShtDevice } from '../sht.ts';
6+import { randomSpectrum } from '../reference.ts';
7+import { lmIndex, type ShtConfig } from '../layout.ts';
8+
9+const $ = <T extends HTMLElement>(id: string) => document.getElementById(id) as T;
10+const mapCanvas = $<HTMLCanvasElement>('map');
11+const statsEl = $('stats');
12+const errEl = $('err');
13+const hoverEl = $('hover');
14+
15+/** Diverging cool-warm colormap (blue — neutral — red), t in [0, 1]. */
16+function coolwarm(t: number): [number, number, number] {
17+ // endpoints and neutral midpoint (Moreland-style, linearly interpolated)
18+ const lo = [59, 76, 192], mid = [221, 221, 221], hi = [180, 4, 38];
19+ const a = t < 0.5 ? lo : hi;
20+ const u = t < 0.5 ? t * 2 : 2 - t * 2; // 0 at endpoint, 1 at midpoint
21+ return [
22+ Math.round(a[0] + (mid[0] - a[0]) * u),
23+ Math.round(a[1] + (mid[1] - a[1]) * u),
24+ Math.round(a[2] + (mid[2] - a[2]) * u),
25+ ];
26+}
27+
28+let device: GPUDevice;
29+let plan: ShtPlan | null = null;
30+let cfg: ShtConfig;
31+let qlm: Float32Array;
32+let spat: Float32Array | null = null;
33+let vmax = 1;
34+let seed = 1;
35+
36+function currentConfig(): ShtConfig {
37+ const lmax = parseInt($<HTMLSelectElement>('preset').value, 10);
38+ const nlat = lmax + 1;
39+ let nphi = 4;
40+ while (nphi < 2 * lmax + 2) nphi *= 2;
41+ return { lmax, mmax: lmax, nlat, nphi };
42+}
43+
44+function buildSpectrum(): Float32Array {
45+ const source = $<HTMLSelectElement>('source').value;
46+ const n = 2 * plan!.nlm;
47+ if (source === 'single') {
48+ const q = new Float32Array(n);
49+ const l = Math.min(Math.max(0, $<HTMLInputElement>('inpL').valueAsNumber || 0), cfg.lmax);
50+ const m = Math.min(Math.max(0, $<HTMLInputElement>('inpM').valueAsNumber || 0), l);
51+ $<HTMLInputElement>('inpL').value = String(l);
52+ $<HTMLInputElement>('inpM').value = String(m);
53+ q[2 * lmIndex(cfg.lmax, l, m)] = 1;
54+ return q;
55+ }
56+ const q = randomSpectrum(cfg, seed);
57+ if (source === 'bandpass') {
58+ const lo = Math.floor(cfg.lmax / 3), hi = Math.ceil((2 * cfg.lmax) / 3);
59+ for (let m = 0; m <= cfg.mmax; m++)
60+ for (let l = m; l <= cfg.lmax; l++) {
61+ if (l < lo || l > hi) {
62+ const k = lmIndex(cfg.lmax, l, m);
63+ q[2 * k] = 0;
64+ q[2 * k + 1] = 0;
65+ }
66+ }
67+ }
68+ return q;
69+}
70+
71+function draw() {
72+ if (!spat) return;
73+ const { nlat, nphi } = cfg;
74+ const off = new OffscreenCanvas(nphi, nlat);
75+ const octx = off.getContext('2d')!;
76+ const img = octx.createImageData(nphi, nlat);
77+ vmax = 1e-30;
78+ for (let i = 0; i < spat.length; i++) vmax = Math.max(vmax, Math.abs(spat[i]));
79+ for (let i = 0; i < nlat; i++) {
80+ for (let j = 0; j < nphi; j++) {
81+ const v = spat[i * nphi + j];
82+ const [r, g, b] = coolwarm(0.5 + (0.5 * v) / vmax);
83+ const o = 4 * (i * nphi + j);
84+ img.data[o] = r;
85+ img.data[o + 1] = g;
86+ img.data[o + 2] = b;
87+ img.data[o + 3] = 255;
88+ }
89+ }
90+ octx.putImageData(img, 0, 0);
91+ const W = Math.min(908, mapCanvas.parentElement!.clientWidth - 2);
92+ const H = Math.round(W / 2);
93+ mapCanvas.width = W;
94+ mapCanvas.height = H;
95+ const ctx = mapCanvas.getContext('2d')!;
96+ ctx.imageSmoothingEnabled = true;
97+ ctx.imageSmoothingQuality = 'high';
98+ ctx.drawImage(off, 0, 0, W, H);
99+
100+ // colorbar
101+ const cb = $<HTMLCanvasElement>('colorbar');
102+ const cctx = cb.getContext('2d')!;
103+ for (let x = 0; x < cb.width; x++) {
104+ const [r, g, b] = coolwarm(x / (cb.width - 1));
105+ cctx.fillStyle = `rgb(${r},${g},${b})`;
106+ cctx.fillRect(x, 0, 1, cb.height);
107+ }
108+ $('vmin').textContent = (-vmax).toPrecision(3);
109+ $('vmax').textContent = '+' + vmax.toPrecision(3);
110+}
111+
112+async function recompute() {
113+ if (!plan) return;
114+ errEl.textContent = '';
115+ try {
116+ qlm = buildSpectrum();
117+ const t0 = performance.now();
118+ spat = await plan.synth(qlm);
119+ const tSynth = performance.now() - t0;
120+ const t1 = performance.now();
121+ const qBack = await plan.analys(spat);
122+ const tAnalys = performance.now() - t1;
123+ let num = 0, den = 0;
124+ for (let k = 0; k < qlm.length; k++) {
125+ num += (qBack[k] - qlm[k]) ** 2;
126+ den += qlm[k] ** 2;
127+ }
128+ const rel = Math.sqrt(num / (den || 1));
129+ draw();
130+ statsEl.innerHTML =
131+ `grid <b>${cfg.nlat}×${cfg.nphi}</b>, lmax <b>${cfg.lmax}</b>, nlm <b>${plan.nlm}</b>, ` +
132+ `Fourier <b>${plan.fourierMode}</b> · synthesis <b>${tSynth.toFixed(1)} ms</b>, ` +
133+ `analysis <b>${tAnalys.toFixed(1)} ms</b> (incl. transfers) · ` +
134+ `round-trip rel. error <b>${rel.toExponential(2)}</b>`;
135+ } catch (e) {
136+ errEl.textContent = String(e);
137+ }
138+}
139+
140+async function rebuild() {
141+ cfg = currentConfig();
142+ plan?.destroy();
143+ plan = null;
144+ statsEl.textContent = 'building plan…';
145+ try {
146+ plan = await ShtPlan.create(device, cfg);
147+ await recompute();
148+ } catch (e) {
149+ statsEl.textContent = '';
150+ errEl.textContent = String(e);
151+ }
152+}
153+
154+mapCanvas.addEventListener('mousemove', (ev) => {
155+ if (!spat || !plan) return;
156+ const r = mapCanvas.getBoundingClientRect();
157+ const j = Math.min(cfg.nphi - 1, Math.max(0, Math.floor(((ev.clientX - r.left) / r.width) * cfg.nphi)));
158+ const i = Math.min(cfg.nlat - 1, Math.max(0, Math.floor(((ev.clientY - r.top) / r.height) * cfg.nlat)));
159+ const latDeg = 90 - (plan.theta[i] * 180) / Math.PI;
160+ const lonDeg = (j * 360) / cfg.nphi;
161+ hoverEl.textContent = `lat ${latDeg.toFixed(1)}°, lon ${lonDeg.toFixed(1)}°: ${spat[i * cfg.nphi + j].toPrecision(4)}`;
162+});
163+mapCanvas.addEventListener('mouseleave', () => (hoverEl.textContent = ''));
164+
165+$('preset').addEventListener('change', rebuild);
166+$('source').addEventListener('change', () => {
167+ $('lmWrap').style.display = $<HTMLSelectElement>('source').value === 'single' ? '' : 'none';
168+ recompute();
169+});
170+$('inpL').addEventListener('change', recompute);
171+$('inpM').addEventListener('change', recompute);
172+$('reroll').addEventListener('click', () => {
173+ seed = (seed * 16807) % 2147483647 || 1;
174+ recompute();
175+});
176+window.addEventListener('resize', draw);
177+
178+async function start() {
179+ $('lmWrap').style.display = '';
180+ try {
181+ device = await requestShtDevice();
182+ } catch (e) {
183+ statsEl.textContent = '';
184+ errEl.textContent =
185+ `WebGPU unavailable: ${e instanceof Error ? e.message : e}\n` +
186+ 'Use a browser with WebGPU support (Chrome/Edge 113+, recent Firefox/Safari).';
187+ return;
188+ }
189+ await rebuild();
190+}
191+start();
src/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/index.tsadded+5−0View file
@@ -0,0 +1,5 @@
1+export { ShtPlan, requestShtDevice, type ShtOptions, type FourierMode } from './sht.ts';
2+export { nlmCalc, lmIndex, validateConfig, type ShtConfig } from './layout.ts';
3+export { ShtReference, randomSpectrum } from './reference.ts';
4+export { gaussNodesWeights } from './gauss.ts';
5+export { legendreCoeffs, legendreRow } from './coeffs.ts';
src/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/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.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/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/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/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+}
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>shtns-webgpu validation</title>
6+ </head>
7+ <body>
8+ <h1>shtns-webgpu 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+129−0View file
@@ -0,0 +1,129 @@
1+/**
2+ * Browser validation suite: fp32 WebGPU transforms vs the f64 reference.
3+ * Results are written to #log, console, and window.__RESULTS__ (read by
4+ * scripts/test-gpu.mjs).
5+ */
6+import { ShtPlan, requestShtDevice } from '../src/sht.ts';
7+import { ShtReference, randomSpectrum } from '../src/reference.ts';
8+import { lmIndex, type ShtConfig } from '../src/layout.ts';
9+import type { FourierMode } from '../src/sht.ts';
10+
11+interface CaseResult {
12+ name: string;
13+ pass: boolean;
14+ detail: string;
15+}
16+const results: CaseResult[] = [];
17+const logEl = document.getElementById('log')!;
18+function log(line: string) {
19+ console.log(line);
20+ logEl.textContent += '\n' + line;
21+}
22+
23+function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
24+ let num = 0,
25+ den = 0;
26+ for (let i = 0; i < a.length; i++) {
27+ num += (a[i] - b[i]) ** 2;
28+ den += b[i] ** 2;
29+ }
30+ return Math.sqrt(num / (den || 1));
31+}
32+
33+async function runCase(
34+ device: GPUDevice,
35+ name: string,
36+ cfg: ShtConfig,
37+ fourier: FourierMode,
38+ tolSynth: number,
39+ tolAnalys: number,
40+ tolRound: number,
41+) {
42+ const t0 = performance.now();
43+ const plan = await ShtPlan.create(device, cfg, { fourier });
44+ const ref = new ShtReference(cfg);
45+ const q0 = randomSpectrum(cfg, 42 + cfg.lmax);
46+
47+ // synthesis vs f64 reference
48+ const spatGpu = await plan.synth(q0);
49+ const spatRef = ref.synth(q0);
50+ const eSynth = relL2(spatGpu, spatRef);
51+
52+ // analysis of the reference field vs f64 reference
53+ const spatRef32 = Float32Array.from(spatRef);
54+ const qGpu = await plan.analys(spatRef32);
55+ const qRef = ref.analys(spatRef32);
56+ const eAnalys = relL2(qGpu, qRef);
57+
58+ // GPU round trip: synth -> analys, compare to original spectrum
59+ const qRound = await plan.analys(spatGpu);
60+ const eRound = relL2(qRound, q0);
61+
62+ const dt = (performance.now() - t0).toFixed(0);
63+ const pass = eSynth < tolSynth && eAnalys < tolAnalys && eRound < tolRound;
64+ const detail =
65+ `mode=${plan.fourierMode} synth=${eSynth.toExponential(2)}/${tolSynth} ` +
66+ `analys=${eAnalys.toExponential(2)}/${tolAnalys} round=${eRound.toExponential(2)}/${tolRound} (${dt}ms)`;
67+ results.push({ name, pass, detail });
68+ log(`${pass ? 'PASS' : 'FAIL'} ${name} ${detail}`);
69+ plan.destroy();
70+}
71+
72+async function runSingleModeCase(device: GPUDevice, name: string, cfg: ShtConfig, l: number, m: number) {
73+ // synthesize a single (l, m) mode and analyze it back: spectrum should
74+ // come back as the unit vector, and the field should match Y_lm exactly.
75+ const plan = await ShtPlan.create(device, cfg, {});
76+ const q = new Float32Array(2 * plan.nlm);
77+ q[2 * lmIndex(cfg.lmax, l, m)] = 1.0;
78+ const spat = await plan.synth(q);
79+ const qBack = await plan.analys(spat);
80+ const e = relL2(qBack, q);
81+ const pass = e < 2e-5;
82+ results.push({ name, pass, detail: `round=${e.toExponential(2)}` });
83+ log(`${pass ? 'PASS' : 'FAIL'} ${name} round=${e.toExponential(2)}`);
84+ plan.destroy();
85+}
86+
87+async function main() {
88+ if (!navigator.gpu) {
89+ (window as any).__RESULTS__ = { fatal: 'navigator.gpu undefined (WebGPU unavailable)' };
90+ log('FATAL: WebGPU unavailable');
91+ return;
92+ }
93+ const adapter = await navigator.gpu.requestAdapter();
94+ const info = adapter ? `${adapter.info?.vendor ?? '?'} / ${adapter.info?.architecture ?? '?'}` : 'none';
95+ log(`adapter: ${info}`);
96+ const device = await requestShtDevice();
97+ device.addEventListener('uncapturederror', (ev) => {
98+ log(`UNCAPTURED GPU ERROR: ${(ev as GPUUncapturedErrorEvent).error.message}`);
99+ });
100+
101+ try {
102+ // small, FFT path
103+ await runCase(device, 'lmax=15 fft', { lmax: 15, mmax: 15, nlat: 32, nphi: 32 }, 'fft', 2e-6, 2e-6, 2e-6);
104+ // small, DFT path (cross-check of the Fourier stages)
105+ await runCase(device, 'lmax=15 dft', { lmax: 15, mmax: 15, nlat: 32, nphi: 36 }, 'dft', 2e-6, 2e-6, 2e-6);
106+ // moderate
107+ await runCase(device, 'lmax=63 fft', { lmax: 63, mmax: 63, nlat: 64, nphi: 128 }, 'auto', 5e-6, 5e-6, 5e-6);
108+ // the SHTNS fp32 comfort zone boundary (SHT_L_RESCALE_FLY_FLOAT = 128)
109+ await runCase(device, 'lmax=127 fft', { lmax: 127, mmax: 127, nlat: 128, nphi: 256 }, 'auto', 1e-5, 1e-5, 1e-5);
110+ // reduced mmax
111+ await runCase(device, 'lmax=127 mmax=40', { lmax: 127, mmax: 40, nlat: 144, nphi: 128 }, 'auto', 1e-5, 1e-5, 1e-5);
112+ // beyond the comfort zone: rescaling must kick in (sin^m underflows f32
113+ // around m ~ 90 at mid-latitudes); accuracy degrades gracefully
114+ await runCase(device, 'lmax=255', { lmax: 255, mmax: 255, nlat: 256, nphi: 512 }, 'auto', 5e-5, 5e-5, 5e-5);
115+ await runCase(device, 'lmax=399', { lmax: 399, mmax: 399, nlat: 400, nphi: 1024 }, 'auto', 2e-4, 2e-4, 2e-4);
116+ // single-mode checks incl. a high-m sectoral mode (pure rescale territory)
117+ await runSingleModeCase(device, 'mode (l=3,m=2)', { lmax: 15, mmax: 15, nlat: 32, nphi: 32 }, 3, 2);
118+ await runSingleModeCase(device, 'mode (l=200,m=200)', { lmax: 200, mmax: 200, nlat: 224, nphi: 512 }, 200, 200);
119+ } catch (e) {
120+ results.push({ name: 'exception', pass: false, detail: String(e) });
121+ log(`EXCEPTION: ${e instanceof Error ? e.stack ?? e.message : e}`);
122+ }
123+
124+ const failed = results.filter((r) => !r.pass);
125+ log(failed.length === 0 ? 'ALL GPU TESTS PASSED' : `${failed.length} GPU TEST(S) FAILED`);
126+ (window as any).__RESULTS__ = { results, ok: failed.length === 0 };
127+}
128+
129+main();
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+});