shtns-webgpu#

Spherical harmonic transforms on WebGPU (browser, fp32), modeled on SHTNS. This is a from-scratch TypeScript + WGSL implementation of the scalar transforms, structured after the SHTNS CUDA backend (sht_gpu.cu / SHT/cuda_legendre.gen.cu).

Live demo: https://concept-collection.github.io/shtns-webgpu/ (validation suite: test.html)

Scope (v0.1)#

• Scalar transforms of real fields, both directions:
  • synth() — spectral → spatial (SHTNS SH_to_spat)
  • analys() — spatial → spectral (SHTNS spat_to_SH)
• Gauss–Legendre grid in latitude, uniform longitude grid.
• fp32 on the GPU end to end; all precomputation (Gauss nodes/weights, recurrence coefficients, twiddle factors) is done host-side in f64.
• Not (yet) implemented: vector (spheroidal/toroidal) transforms, complex fields, regular grids, mres > 1, Schmidt/4π normalizations, on-the-fly truncation (llim < lmax).

Conventions (= SHTNS defaults)#

• Orthonormal spherical harmonics with Condon–Shortley phase.
• Spectral coefficients Q_lm are complex, stored for m >= 0 with the SHTNS LM ordering: for m = 0..mmax: for l = m..lmax, interleaved [re, im] (Float32Array of length 2*nlm). Real fields imply Q_{l,-m} = (-1)^m conj(Q_lm); m = 0 coefficients must be real.
• Spatial fields are Float32Array[nlat * nphi], phi-contiguous, latitudes ordered north → south (colatitude increasing).

Usage#

import { ShtPlan, requestShtDevice, lmIndex } from 'shtns-webgpu';

const device = await requestShtDevice();           // or your own GPUDevice
const plan = await ShtPlan.create(device, { lmax: 127, mmax: 127, nlat: 128, nphi: 256 });

const qlm = new Float32Array(2 * plan.nlm);
qlm[2 * lmIndex(127, 8, 5)] = 1.0;                 // Y_8^5

const spat = await plan.synth(qlm);                // nlat*nphi field
const qBack = await plan.analys(spat);             // back to spectral
plan.destroy();

For GPU-resident pipelines (no readback), use plan.encodeSynth(encoder) / plan.encodeAnalys(encoder) with the exposed qlmIn / qlmOut / spatBuf buffers.

Constraints checked at plan creation: nlat > lmax (Gauss quadrature exactness), nphi >= 2*mmax + 1 (no aliasing).

How it works#

Same two-stage split as SHTNS:

1. Legendre stage (src/wgsl/leg.ts): associated Legendre functions are generated on the fly inside the shader by the standard 3-term recurrence over l (coefficients from legendre_precomp()-equivalent host code, src/coeffs.ts). Underflow of sin^m(theta) — fatal in fp32 beyond m ≈ 75 — is handled with the SHTNS extended-range scheme (SHT_SCALE_FACTOR = 2^56, SHT_ACCURACY = 1e-15, per-thread integer exponent), ported from the HI_LLIM path of SHT/cuda_legendre.gen.cu. Synthesis runs one thread per latitude and one workgroup row per m; analysis runs one workgroup per m with a shared-memory tree reduction over latitudes (the portable equivalent of SHTNS's warp shuffles).
2. Fourier stage (src/wgsl/fourier.ts): batched radix-2 Stockham FFT in workgroup memory (one workgroup per latitude row) when nphi is a power of two that fits (16*nphi <= maxComputeWorkgroupStorageSize); otherwise a direct band-limited DFT. Twiddles come from a host-computed f64 table — device sin/cos is only guaranteed to ~2^-11 under Vulkan, which would otherwise dominate the error budget.

All problem sizes are baked into the WGSL at plan creation (the WGSL equivalent of SHTNS's NVRTC runtime compilation).

Accuracy (fp32)#

Relative L2 errors vs the double-precision reference (src/reference.ts), random spectra, measured on SwiftShader (results on hardware GPUs are the same to within noise since the arithmetic is IEEE fp32):

SHTNS itself switches its fp32 GPU recurrence to f64 above lmax = 128 (SHT_L_RESCALE_FLY_FLOAT); WGSL has no f64, so past that point accuracy degrades gracefully as above. Fine for visualization; for scientific use keep lmax ≲ 128 or wait for the float-float recurrence (planned).

Develop / test#

npm install
npm run dev        # demo at http://localhost:5173
npm run test:node  # f64 math tests (no GPU needed)
npm run test:gpu   # builds, then runs the browser suite in headless Chrome
                   # (falls back to SwiftShader software WebGPU; CHROME_PATH to override)

Roadmap#

• Vector transforms (spheroidal/toroidal), gradients — port of leg_m_kernel<1>.
• Float-float (double-single) recurrence option for full accuracy at high lmax.
• Latitude parity folding (2× Legendre work reduction, as in SHTNS).
• Subgroup (warp) reductions where available, replacing the shared-memory tree.
• mres > 1, truncated transforms, regular grids.