# shtns-webgpu

Spherical harmonic transforms on **WebGPU** (browser, fp32), modeled on
[SHTNS](https://nschaeff.bitbucket.io/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](https://concept-collection.github.io/shtns-webgpu/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

```ts
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):

| lmax | synthesis | analysis | round trip |
|-----:|----------:|---------:|-----------:|
|   15 |     6e-7 |     3e-7 |      6e-7 |
|   63 |     5e-6 |     2e-6 |      3e-6 |
|  127 |     7e-6 |     3e-6 |      5e-6 |
|  255 |     2e-5 |     6e-6 |      1e-5 |
|  399 |     9e-5 |     1e-5 |      2e-5 |

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

```sh
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.