/ concept-collection / turing-surface
Sign in
concept-collection / turing-surface
23.1 KBPreviewCodeBlameHistoryRaw

turing-surface#

Reaction–diffusion systems (Turing patterns) on closed surfaces given by spherical-harmonic embeddings, solved live in the browser with a spectral method whose transforms run on the GPU via WebGPU.

This is the sibling of turing-sphere, which solves the same systems on the round sphere. Everything there is here; what is added is a surface.

The geometry is in the operator: the models evaluate the surface Laplace–Beltrami operator lap_g inside the implicit solve, in a flux form that costs 6 spherical-harmonic transforms per species per iteration where the textbook Cartesian-gradient form needs 12. See The geometry in the operator and docs/reduced-transforms.md.

What a surface is here#

A geometry is an embedding of the sphere into R³: three scalar fields x, y, z over the (θ, φ) parametrization, each carried as spherical-harmonic coefficients. The unit sphere is the case where all three are pure degree-1 harmonics.

You write one down as MATLAB, in geometries/:

function [gx, gy, gz] = shape(theta, phi, waist, stretch)
  st = sin(theta);
  r = 1 - waist * (st .^ 2);
  gx = r .* (st .* cos(phi));
  gy = r .* (st .* sin(phi));
  gz = (1 + stretch) * (r .* cos(theta));
end

That is ordinary element-wise MATLAB and goes through the same compiler and the same WGSL backend the models do. It is evaluated once on the solver's grid, and then analysed into coefficients, which is the form everything downstream uses. Two things follow from going through the coefficients rather than keeping the pointwise values:

Four geometries ship: sphere (the reference case), ellipsoid, peanut — a dumbbell whose waist is a saddle — and bumpy. Each is editable in the page, with its own parameters. Changing a shape does not recompile the solver and does not disturb the run: the geometry is data whose shape in the bindings depends only on the grid, so a swap is sixteen buffer writes and the pattern carries straight on.

A morph slider blends the drawn surface back to the unit sphere. The parametrization is the sphere's either way, so sweeping it shows which point went where.

The scheme, and where the geometry enters#

It solves the N-species system

d(u_k)/dt = D_k*lap_g(u_k) + f_k(t, u_1, ..., u_N),    k = 1, ..., N

where lap_g is the Laplace–Beltrami operator of the surface. On the round sphere lap_g is diagonal in spherical-harmonic space with eigenvalues -l(l+1), which is what makes turing-sphere's implicit diffusion a single divide. On a general surface it is not diagonal, and not even constant- coefficient, so that divide has to become a solve.

The models split the operator:

lap_g = lap_s + dlap

with lap_s the round-sphere one. (I - dt*D*lap_s) is still exactly invertible, so the implicit step

(I - dt*D*lap_g) Unew = B

rearranges into a fixed point that keeps the whole geometry on the right-hand side,

Unew = (B + dt*D*dlap(Unew)) ./ (1 + dt*D*lam)

and the loop iterates it from the round-sphere answer. That is preconditioned Richardson, with the operator we can invert exactly as the preconditioner; it converges while dt*D*dlap stays small against (I - dt*D*lap_s), which is what keeps the cost to a few transforms per step rather than a full elliptic solve (see docs/richardson-iteration.md). One species of models/schnakenberg.m's solve loop:

for k = 1:niter
  Fu = Un .* filt;              % zero the top 2 degrees before differentiating
  Ftu = synth(dthetac(Fu));     % sin(theta) * dtheta(u) -- smooth on the sphere
  Fpu = synth(dphic(Fu));       % dphi(u)                -- smooth on the sphere
  Pu = p1 .* Ftu + p2 .* Fpu;   % the two fluxes, also smooth: the precomputed
  Qu = p2 .* Ftu + q2 .* Fpu;   % weights carry every 1/sin(theta) there is
  Pcu = analys(Pu) .* filt;
  Qcu = analys(Qu) .* filt;
  scu = dthetac(Pcu) + dphic(Qcu);   % divergence, in coefficient space
  lapu = r .* synth(scu);            % = lap_g(u) on the grid
  dLu = analys(lapu) + lam .* Un;    % dlap = lap_g - lap_s
  Un = (Bu + (dt * D1) * dLu) ./ (1 + (dt * D1) * lam);
end

On the sphere dlap is mathematically zero — p1 = q2 = 1, p2 = 0, r = 1/sin²θ, and the composition collapses to lap_s — so the sphere case reproduces turing-sphere to fp32 round-off, and the tests assert the state stays put across 0, 1 and 4 iterations.

The geometry in the operator#

dlap = lap_g - lap_s is applied to the current iterate at every solve iteration, so its transform count is what the whole step's cost scales with. Two formulations ship:

  1. The flux form (above, all three models): lap_g u as the weighted divergence of two weighted fluxes of the sin-scaled derivatives. The weights p1, p2, q2, r are grid arrays precomputed once per surface from the embedding's θ/φ tangents (src/geom/metric.ts), chosen so that every field that gets analysed is a smooth function on the sphere — the property that makes spherical-harmonic analysis meaningful, and the entire difficulty near the poles. Cost: 6 transforms per species per iteration (3 syntheses + 3 analyses; dthetac/dphic are O(nlm) coefficient shuffles, not transforms). The derivation, the smoothness argument and the fp32 error analysis are in docs/reduced-transforms.md.
  2. The Cartesian-gradient form (Algorithm 4 of docs/algos.pdf), kept as a live reference in models/schnakenberg_alg4.m and selectable in the app: the surface gradient carried as three ambient components through the inverse metric quantities Vt*/Vp*. Cost: 12 transforms per species per iteration. The tests hold both forms to the same answer on a curved surface, and both metric formulations are precomputed and uploaded for every geometry, so either kind of model runs.

The θ-derivative machinery both forms need — the α± recurrence (sin θ ∂θ Y_l^m = α⁺Y_{l+1}^m + α⁻Y_{l-1}^m) as a coefficient-space shuffle feeding the existing scalar synthesis — lives in src/sht/deriv.ts; no Legendre-derivative tables are required.

for loops, unrolled#

A plan is a fixed list of GPU operations with no branching, which is what makes a timestep pure command recording — one submit, no CPU in the loop. A counted loop still fits: the planner (src/mgpu/plan.ts) unrolls it, planning the body once per iteration.

Nothing else had to change for that, because numbl gives a variable one cName for every assignment to it: the buffer an iteration writes is the buffer the next one reads, which is exactly a loop-carried value. The loop variable gets no buffer at all — it is bound as a derived scalar to that iteration's literal, so a kernel reading k folds the number in.

Two consequences worth stating:

Unrolling is exactly linear in the trip count: 19 GPU ops per species per iteration (6 transforms, 4 coefficient shuffles, 9 kernels), asserted in the tests.

MATLAB, compiled to WebGPU#

Unchanged from turing-sphere, and it now compiles the geometry files too. numbl parses and lowers each function for the concrete argument types of the current grid; its inline pass folds single-use temps back into their consumer, so one line of MATLAB becomes one expression tree; and this repo emits one WGSL compute kernel per element-wise statement (src/mgpu/wgsl.ts). synth / analys are external operations whose type rules numbl learns from a .mtoc2.js workspace file, and which the backend maps onto the spherical-harmonic pipelines. Anything it cannot express is refused at compile time with a source position.

The Schnakenberg step compiles to 51 GPU operations at one solve iteration: 16 transforms, 8 coefficient-space shuffles, 25 generated kernels, and 2 buffer copies feeding the new state back.

Transforms batch. The expensive part of every Legendre stage is generating the associated Legendre values on the fly by recurrence — work that depends only on the grid, not on the field. synth/analys therefore take multiple fields, and a grouped call runs as one batched dispatch: one walk of the recurrence, one accumulator lane per field —

[Ftu, Fpu, Ftv, Fpv] = synth(vtu, vpu, vtv, vpv);   % one Legendre dispatch

The grouping is a promise of independence, never of a lane width: the planner (src/mgpu/plan.ts, materializeTransforms) chunks each group into whatever the device supports — one ×4 batch under the default WebGPU limits, or scalar dispatches with SHT_BATCH=0 for A/B — so the same source runs anywhere. Ungrouped transforms that happen to sit on consecutive independent lines are batched the same way. Per-lane arithmetic is identical to the scalar kernels', so batched and scalar plans produce bit-identical states, asserted in the tests along with compile-time refusal of a group that drops one of its outputs. All 16 transforms of the step above land in batches, worth ~25% of the whole step (0.88 vs 1.14 ms/step at lmax 127, 2 iterations, on bumpy).

Two consequences carried over:

Provenance#

turing-sphere additionally carries a comparison against a native build of upstream SHTNS (bench/shtns/). That is not duplicated here: the transforms are the same code, and its C-side transcription of the model would have to be maintained against a step this project intends to change.

Because the algorithm is compiled to compute shaders, WebGPU is required — there is no CPU fallback (the f64 CPU transform remains, for tests).

Numerics#

Desktop vs browser#

scripts/bench.ts runs the same thing the app runs — same .m, same generated WGSL, same transforms — from Node on desktop WebGPU (Google Dawn), and the app prints the command line that reproduces whatever it is currently simulating:

npm run bench -- --preset schnak-spots --geometry ellipsoid --lmax 63 --niter 1 \
  --steps 2000 --seed 1 --a 0.1 --b 0.9 --D1 0.0004 --D2 0.008 --dt 0.05 \
  --gax 1.5 --gay 1 --gaz 0.6

Copy it from under the stats line and compare the ms/step it reports with the app's. Both sides go through the one shared src/bench/runSpec.ts — the app formats a run into that command, the benchmark parses it back — so there is no second copy of the defaults for the two runs to drift apart on. Geometry parameters take a g prefix (--gwaist) so a shape parameter can never collide with a model one.

The app reports two numbers and only the first is comparable to the benchmark: solver is the batch of steps alone, waited for but not read back; ms/frame additionally carries a GPU→CPU readback per species, the colormapping, and the vertex upload. Those per-frame costs are fixed and do not shrink when the GPU gets faster, so on a quick GPU a frame can easily cost ten times the steps inside it. That is expected and is not the solver being slower in the browser.

To attribute the gap rather than guess at it:

node scripts/compare-perf.mjs [--lmax 63] [--steps 300]

measures the same solver work in both — batched, nothing read back, no rendering on either side — and reports each with its CPU-encoding share, the Fourier stage, and the adapter. It stops you first if the two are not even the same device, which is a common cause of "the browser is much slower". Both sides resolve the geometry and the iteration count from the same constants, because the iteration count is unrolled into the step and a mismatch would compare two different amounts of work.

The app's Benchmark button runs the same measurement in the page, plus the ramp — the first third of the run against the last. GPUs downclock when idle and an animation-paced loop leaves them idle most of every frame, so a large ramp means the steady-state number is limited by clocks rather than work.

Is it really the same computation?#

node scripts/compare-env.mjs [--lmax 31] [--steps 200] [--preset schnak-spots]

runs one identical spec on the desktop and in a real browser and compares the final spectral state. The pipeline is deterministic given (model source, geometry, parameters, lmax, niter, seed, steps), so the two should agree to fp32 round-off — not bit for bit, since GPUs differ in fused-multiply-add and other latitude fp32 allows. It also reports which Fourier stage each side chose, since FFT and DFT are genuinely different algorithms that round differently.

Desktop WebGPU comes from the webgpu package (prebuilt Dawn, ~70 MB), an optional dependency so that an unsupported platform fails the install of that package alone. Its binaries need glibc 2.29+. Other flags: --steps, --warmup, --batch, --json, --help; DAWN_FLAGS='backend=vulkan' (;-separated) passes Dawn options through.

Tests#

There is no second implementation of the solver to diff against, so the .m path is checked against closed-form answers and against exact structural properties. Five modules, run in both environments:

test/analyticChecks.ts — cases whose evolution is known exactly, run through the whole real pipeline. All three are statements about the round sphere, so all three build on the sphere geometry:

test/geometryChecks.ts — the surface and the loop:

test/fluxChecks.ts — the six-transform flux-form Laplace-Beltrami scheme (docs/reduced-transforms.md):

test/modelChecks.ts compiles every model the app offers and asserts how many kernels it compiles to, split into the base step and what one solve iteration adds. That is a fusion guard: if numbl's inline pass stops folding, the results stay correct while every operator becomes its own dispatch, which is invisible in the numbers.

test/transformChecks.ts compares the WGSL transforms against shtns-webgpu's f64 CPU twin, and holds every compiled batch width to the scalar transforms lane by lane; a model run with SHT_BATCH=0 must reproduce the batched run's state exactly.

Other commands:

Development#

npm install
npm run dev       # local dev server
npm run build     # type-check + production build to dist/

The numbl dependency#

numbl is a local file:../../numbl dependency, so a sibling checkout of numbl is required. We use its compiler internals — parser, lowerer, IR, inline pass — which its package exports map does not publish, so they are reached through the numbl-src path alias in vite.config.ts.

The exact surface we depend on is written down in src/mgpu/numbl.d.ts and TypeScript checks against that, not against numbl's sources. This keeps this project's compiler settings independent of numbl's, and means a change to one of those shapes upstream breaks the build here with a clear diff rather than deep inside numbl's tree. The For IR node is spelled out there, since the planner now walks it.

CI clones numbl to the sibling path that the file: dependency expects, pinned to a commit, with --ignore-scripts (npm runs a linked package's prepare script, and numbl's is husky). numbl's own node_modules are not needed: the slice we import is self-contained TypeScript.

The scripts/*.ts entry points that touch the compiler go through vite-node, so they resolve imports exactly as the browser build does. Plain node cannot: numbl's sources import each other as ./foo.js while the files are .ts.

Deployed to GitHub Pages by .github/workflows/deploy.yml on push to main.

License#

CECILL-2.1 (inherited from SHTNS via shtns-webgpu, whose sources are vendored).

moveopenescclose