591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 1/**
2 * The flux-form (six-transform) Laplace-Beltrami scheme of
3 * docs/reduced-transforms.md, against the two things that can
4 * silently go wrong with it:
5 *
6 * 1. The smoothness claim (doc Sec 2, validation Sec 7.1). The whole scheme
7 * rests on the analysed fluxes P and Qtilde being smooth functions on the
8 * sphere — that is a mathematical property of the p1/p2/q2 weighting, so it
9 * is checked in f64 on the CPU, where a failure is a wrong formula and not
10 * round-off. The fields are synthesized and re-analysed on a grid with
11 * twice the band limit: content beyond the band is exactly the non-smooth
12 * residue the weighting is supposed to remove.
13 *
14 * Two surfaces split the claim's two halves. On the *round sphere* the
15 * correctly weighted fluxes are exactly band-limited, so their beyond-band
16 * tail is f64 round-off, while the doc's Sec 8 counterexample
17 * Qtilde/sin(theta) — bounded but with a phi-dependent polar limit — keeps
18 * an algebraically decaying tail orders of magnitude above it: the
19 * decisive smooth-vs-non-smooth discrimination, plus the closed-form check
20 * p1 = q2 = 1, p2 = 0, r = 1/sin^2(theta). On *bumpy* (non-axisymmetric,
21 * so the off-diagonal p2 does real work) nothing is band-limited and every
22 * smooth field's tail is set by the weights' own spectral decay, so the
23 * check there is the doc's relative one: P and Qtilde must sit on the same
24 * footing as the Cartesian gradient component Algorithm 4 analyses.
25 *
26 * 2. The operator identity (validation Sec 7.2/7.4). The flux form and the
27 * Cartesian-gradient form (models/schnakenberg.m vs
28 * models/schnakenberg_alg4.m) are the same operator, so a real simulation
29 * driven by one must track the other to fp32 accumulation — checked on a
30 * non-axisymmetric surface, where the off-diagonal weight p2 actually does
31 * something. The headline transform count (6 vs 12 per species per
32 * iteration) is asserted from the compiled op sequences, not the doc.
34 * 3. The polar conditioning of the divergence (doc Sec 5). r ~ 1/sin^2(theta)
35 * multiplies a bracket that must cancel to O(sin^2(theta)) at the poles,
36 * so it amplifies the polar round-off of whatever it is handed. Splitting
37 * the round sphere out of the divergence (models/schnakenberg.m, and
38 * dp1/dq2/jinv in src/geom/geometry.ts) keeps r off all but the geometry
39 * deviation; without the split, the amplified round-off is a static polar
40 * forcing that a Turing instability grows into a spot at the pole,
41 * regardless of the seed. That is the failure this checks for: it is
42 * invisible to 1 and 2, which compare operators rather than watch what a
43 * run nucleates from.
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 44 */
45import { ShtPlan } from '../src/sht/sht.ts';
46import { DerivPlan } from '../src/sht/deriv.ts';
47import { ShtReference } from '../src/sht/reference.ts';
48import { gridForLmax, lmIndex, nlmCalc, type ShtConfig } from '../src/sht/layout.ts';
49import { ModelSession } from '../src/mgpu/session.ts';
50import { mModelByKey, defaultParams } from '../src/mgpu/registry.ts';
51import { Geometry } from '../src/geom/geometry.ts';
52import { mGeometryByKey, defaultGeometryParams } from '../src/geom/registry.ts';
53import { computeFluxMetric } from '../src/geom/metric.ts';
54import type { Check, Log } from './analyticChecks.ts';
56/** Band limit of the test surface and field. */
57const LMAX = 24;
58/** Band limit of the oversampled analysis grid the tails are measured on. */
59const LMAX_HI = 63;
60/** Degrees at and above this count as "beyond-band tail": LMAX+1 is the last
61 * degree with direct content, and the smooth-but-not-band-limited metric
62 * weights spread it upward with (their own) exponentially decaying spectra,
63 * so the window starts well above the band edge. */
64const TAIL_START = 44;
66/** The part of a transform layout the spectral helpers need. */
67interface Band {
68 lmax: number;
69 mmax: number;
70}
72/** Per-degree spectral amplitude: E(l) = sqrt(sum_m |q_l^m|^2). */
73function degreeEnergy(band: Band, qlm: ArrayLike<number>): Float64Array {
74 const E = new Float64Array(band.lmax + 1);
75 for (let m = 0; m <= band.mmax; m++) {
76 for (let l = m; l <= band.lmax; l++) {
77 const i = lmIndex(band.lmax, l, m);
78 E[l] += qlm[2 * i] ** 2 + qlm[2 * i + 1] ** 2;
79 }
80 }
81 for (let l = 0; l <= band.lmax; l++) E[l] = Math.sqrt(E[l]);
82 return E;
83}
85/** max E(l) over l >= TAIL_START, relative to max E(l) overall. */
86function tailRel(band: Band, qlm: ArrayLike<number>): number {
87 const E = degreeEnergy(band, qlm);
88 let bulk = 0;
89 let tail = 0;
90 for (let l = 0; l <= band.lmax; l++) {
91 if (E[l] > bulk) bulk = E[l];
92 if (l >= TAIL_START && E[l] > tail) tail = E[l];
93 }
94 return tail / Math.max(bulk, 1e-300);
95}
97/** Re-index coefficients from the lo layout into the hi layout (zero-padded). */
98function padSpectrum(qlo: ArrayLike<number>, lo: Band, hi: Band): Float64Array {
99 const out = new Float64Array(2 * nlmCalc(hi.lmax, hi.mmax));
100 for (let m = 0; m <= lo.mmax; m++) {
101 for (let l = m; l <= lo.lmax; l++) {
102 const src = lmIndex(lo.lmax, l, m);
103 const dst = lmIndex(hi.lmax, l, m);
104 out[2 * dst] = qlo[2 * src];
105 out[2 * dst + 1] = qlo[2 * src + 1];
106 }
107 }
108 return out;
109}
111/** Deterministic random band-limited spectrum with O(1) coefficients. */
112function flatSpectrum(band: Band, seed: number): Float64Array {
113 const nlm = nlmCalc(band.lmax, band.mmax);
114 const q = new Float64Array(2 * nlm);
115 let s = seed >>> 0;
116 const rnd = () => {
117 s ^= s << 13; s >>>= 0;
118 s ^= s >> 17;
119 s ^= s << 5; s >>>= 0;
120 return (s / 4294967296) * 2 - 1;
121 };
122 for (let k = 0; k < 2 * nlm; k++) q[k] = rnd();
123 for (let l = 0; l <= band.lmax; l++) q[2 * lmIndex(band.lmax, l, 0) + 1] = 0;
124 return q;
125}
127export interface FluxCheckOptions {
128 /**
129 * Run the live flux-vs-Algorithm-4 A/B (4 sessions at lmax 63). On by
130 * default, but — like geometryChecks' sweep, and for the same reason — a
131 * browser recompiles every session's unrolled step from scratch on software
132 * WebGPU, so the page leaves it out unless asked (?sweep=1) to keep CI
133 * short. The f64 smoothness checks always run; they are CPU work.
134 */
135 ab?: boolean;
136}
138export async function fluxChecks(
139 device: GPUDevice,
140 check: Check,
141 log: Log,
142 opts: FluxCheckOptions = {},
143): Promise<void> {
144 // ---- 1a. round sphere: closed-form weights, decisive discrimination -----
145 {
146 const hiGrid = gridForLmax(LMAX_HI, 1);
147 const hi = { lmax: LMAX_HI, mmax: LMAX_HI, nlat: hiGrid.nlat, nphi: hiGrid.nphi };
148 const ref = new ShtReference(hi);
149 const npts = hi.nlat * hi.nphi;
151 // The unit sphere needs no GPU build: analyse the closed-form embedding
152 // on the fine grid directly, in f64.
153 const xg = new Float64Array(npts);
154 const yg = new Float64Array(npts);
155 const zg = new Float64Array(npts);
156 for (let i = 0; i < hi.nlat; i++) {
157 const ct = ref.ct[i];
158 const st = ref.st[i];
159 for (let j = 0; j < hi.nphi; j++) {
160 const phi = (2 * Math.PI * j) / hi.nphi;
161 const k = i * hi.nphi + j;
162 xg[k] = st * Math.cos(phi);
163 yg[k] = st * Math.sin(phi);
164 zg[k] = ct;
165 }
166 }
167 const X = ref.analys(xg);
168 const Y = ref.analys(yg);
169 const Z = ref.analys(zg);
171 const sXt = [ref.sinDtheta(X), ref.sinDtheta(Y), ref.sinDtheta(Z)];
172 const Xp = [ref.dphi(X), ref.dphi(Y), ref.dphi(Z)];
173 const { p1, p2, q2, r } = computeFluxMetric(
174 npts, sXt[0], sXt[1], sXt[2], Xp[0], Xp[1], Xp[2],
175 );
177 // On the sphere the weights have a closed form: p1 = q2 = 1, p2 = 0,
178 // r = 1/sin^2(theta) — the flux-form counterpart of geometryChecks'
179 // closed-form V check, pinning computeFluxMetric before it is buried
180 // under the operator. f64 throughout, so the tolerance is conditioning
181 // at the polar rings, not fp32.
182 let worst = 0;
183 for (let i = 0; i < hi.nlat; i++) {
184 const st2 = ref.st[i] * ref.st[i];
185 for (let j = 0; j < hi.nphi; j++) {
186 const k = i * hi.nphi + j;
187 worst = Math.max(
188 worst,
189 Math.abs(p1[k] - 1),
190 Math.abs(p2[k]),
191 Math.abs(q2[k] - 1),
192 Math.abs(r[k] * st2 - 1),
193 );
194 }
195 }
196 check(
197 'flux: sphere weights match the closed form (p1 = q2 = 1, p2 = 0, r = 1/sin^2)',
198 worst < 1e-9,
199 `max deviation ${worst.toExponential(2)} in f64`,
200 );
202 // Flat random u, band-limited at LMAX. The properly weighted fluxes are
203 // then *exactly* band-limited (P = sin(theta) dtheta u, Qtilde = dphi u),
204 // so their beyond-band tails are pure round-off; the Sec 8 control
205 // Qtilde/sin(theta) is not a function on the sphere and keeps a fat tail.
206 const band = { lmax: LMAX, mmax: LMAX };
207 const u = padSpectrum(flatSpectrum(band, 777), band, hi);
208 const A = ref.sinDtheta(u);
209 const B = ref.dphi(u);
210 const P = new Float64Array(npts);
211 const Qt = new Float64Array(npts);
212 const control = new Float64Array(npts);
213 for (let i = 0; i < hi.nlat; i++) {
214 const st = ref.st[i];
215 for (let j = 0; j < hi.nphi; j++) {
216 const k = i * hi.nphi + j;
217 P[k] = p1[k] * A[k] + p2[k] * B[k];
218 Qt[k] = p2[k] * A[k] + q2[k] * B[k];
219 control[k] = Qt[k] / st;
220 }
221 }
222 const tails = {
223 P: tailRel(hi, ref.analys(P)),
224 Qt: tailRel(hi, ref.analys(Qt)),
225 control: tailRel(hi, ref.analys(control)),
226 };
227 log(
228 ` flux smoothness on the sphere (f64, band ${LMAX}, analysed to ${LMAX_HI}, ` +
229 `tail l >= ${TAIL_START}): P ${tails.P.toExponential(2)}, ` +
230 `Qt ${tails.Qt.toExponential(2)}, control ${tails.control.toExponential(2)}`,
231 );
232 check(
233 'flux: on the sphere the fluxes are band-limited and the non-smooth control is not',
234 tails.P < 1e-10 && tails.Qt < 1e-10 &&
235 tails.control > 1e3 * Math.max(tails.P, tails.Qt, 1e-14),
236 `P ${tails.P.toExponential(2)}, Qt ${tails.Qt.toExponential(2)}, ` +
237 `control ${tails.control.toExponential(2)}`,
238 );
239 }
241 // ---- 1b. bumpy: the fluxes sit on the Cartesian gradient's footing ------
242 {
243 // The surface: bumpy, the one shipped geometry that is genuinely
244 // non-axisymmetric (g_thetaphi != 0), so the off-diagonal weight p2 is
245 // exercised. Built by the real pipeline at LMAX, then everything below is
246 // CPU f64 from its band-limited coefficients.
247 const g = mGeometryByKey('bumpy')!;
248 const { nlat, nphi } = gridForLmax(LMAX, 3);
249 const cfg = { lmax: LMAX, mmax: LMAX, nlat, nphi };
250 const sht = await ShtPlan.create(device, cfg);
251 const deriv = await DerivPlan.create(device, sht);
252 const geometry = await Geometry.create({
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 254 source: g.source,
255 paramNames: g.params.map((p) => p.key),
256 params: defaultGeometryParams(g),
257 deriv,
258 });
259 deriv.destroy();
260 sht.destroy();
262 const hiGrid = gridForLmax(LMAX_HI, 1);
263 const hi = { lmax: LMAX_HI, mmax: LMAX_HI, nlat: hiGrid.nlat, nphi: hiGrid.nphi };
264 const ref = new ShtReference(hi);
265 const npts = hi.nlat * hi.nphi;
267 // Embedding and test field, zero-padded into the fine layout. Both are
268 // band-limited at LMAX, so on the fine grid every derived field's content
269 // beyond the band is genuinely the non-band-limited part of the weights —
270 // the thing being measured — and not aliasing.
271 const X = padSpectrum(geometry.X, cfg, hi);
272 const Y = padSpectrum(geometry.Y, cfg, hi);
273 const Z = padSpectrum(geometry.Z, cfg, hi);
274 const u = padSpectrum(flatSpectrum(cfg, 777), cfg, hi);
276 // Tangents, both weightings, all f64.
277 const sXt = [ref.sinDtheta(X), ref.sinDtheta(Y), ref.sinDtheta(Z)];
278 const Xp = [ref.dphi(X), ref.dphi(Y), ref.dphi(Z)];
279 const Xt = [ref.dtheta(X), ref.dtheta(Y), ref.dtheta(Z)];
280 const { p1, p2, q2 } = computeFluxMetric(
281 npts, sXt[0], sXt[1], sXt[2], Xp[0], Xp[1], Xp[2],
282 );
284 const A = ref.sinDtheta(u); // sin(theta) dtheta u
285 const B = ref.dphi(u); // dphi u
287 // The two fluxes. No non-smooth control here: on a deformed surface
288 // every smooth field's beyond-band tail is set by the weights' own
289 // (slowly decaying) spectra, which swamps a pole singularity at this
290 // resolution — the sphere block above is where the discrimination has
291 // teeth. This block asserts the doc's relative criterion instead.
292 const P = new Float64Array(npts);
293 const Qt = new Float64Array(npts);
294 for (let k = 0; k < npts; k++) {
295 P[k] = p1[k] * A[k] + p2[k] * B[k];
296 Qt[k] = p2[k] * A[k] + q2[k] * B[k];
297 }
299 // The known-smooth yardstick (doc Sec 2): the x component of the
300 // Cartesian surface gradient, built the Algorithm-4 way from the inverse
301 // metric quantities, in f64.
302 const gradx = new Float64Array(npts);
303 {
304 const ut = ref.dtheta(u);
305 const up = ref.dphi(u);
306 for (let k = 0; k < npts; k++) {
307 const gtt = Xt[0][k] ** 2 + Xt[1][k] ** 2 + Xt[2][k] ** 2;
308 const gtp = Xt[0][k] * Xp[0][k] + Xt[1][k] * Xp[1][k] + Xt[2][k] * Xp[2][k];
309 const gpp = Xp[0][k] ** 2 + Xp[1][k] ** 2 + Xp[2][k] ** 2;
310 const det = gtt * gpp - gtp * gtp;
311 const Vtx = (gpp * Xt[0][k] - gtp * Xp[0][k]) / det;
312 const Vpx = (gtt * Xp[0][k] - gtp * Xt[0][k]) / det;
313 gradx[k] = ut[k] * Vtx + up[k] * Vpx;
314 }
315 }
317 const tails = {
318 P: tailRel(hi, ref.analys(P)),
319 Qt: tailRel(hi, ref.analys(Qt)),
320 gradx: tailRel(hi, ref.analys(gradx)),
321 };
322 log(
323 ` flux smoothness on bumpy (f64, band ${LMAX}, analysed to ${LMAX_HI}, ` +
324 `tail l >= ${TAIL_START}): P ${tails.P.toExponential(2)}, ` +
325 `Qt ${tails.Qt.toExponential(2)}, gradx ${tails.gradx.toExponential(2)}`,
326 );
327 // "Matching tails" (Sec 7.1): same footing as the Cartesian component,
328 // with an order of magnitude of headroom on top of it. A wrong weighting
329 // (a missing sin factor, say) puts genuinely non-smooth content into P or
330 // Qtilde and the tail lands at O(bulk), far above this.
331 const ceiling = Math.max(30 * tails.gradx, 1e-10);
332 check(
333 'flux: on bumpy, P and Qtilde tails match the Cartesian gradient component',
334 tails.P < ceiling && tails.Qt < ceiling,
335 `P ${tails.P.toExponential(2)}, Qt ${tails.Qt.toExponential(2)} vs ` +
336 `ceiling ${ceiling.toExponential(2)}`,
337 );
338 }
340 // ---- 2. flux form vs Algorithm 4, live, on a curved surface -------------
341 if (!(opts.ab ?? true)) {
342 log(
343 ' flux A/B: skipped — run `npm run test:node` (desktop Dawn) or ' +
344 '`npm run test:gpu -- --sweep` for the flux-vs-Algorithm-4 comparison.',
345 );
346 } else {
347 const geometry = mGeometryByKey('bumpy')!;
348 const geometryParams = defaultGeometryParams(geometry);
349 const LMAX_AB = 63;
350 const STEPS = 20;
351 const states: Float32Array[] = [];
352 const xformsPerIter: number[] = [];
354 for (const key of ['schnakenberg', 'schnakenberg-alg4']) {
355 const model = mModelByKey(key)!;
356 const params = defaultParams(model);
357 // Real transforms added by one solve iteration: synth/analys ops plus
358 // dtheta/dphi (each of which contains a synthesis); the coefficient-
359 // space dthetac/dphic shuffles are O(nlm) index gathers, not transforms.
360 const counts: number[] = [];
361 for (const niter of [0, 1]) {
362 const session = await ModelSession.create({
363 device, model, params, lmax: LMAX_AB,
364 geometry, geometryParams, niter,
365 });
366 counts.push(
367 session.describe().step.filter((l) =>
368 l.startsWith('synth') || l.startsWith('analys') ||
369 l.startsWith('dtheta ') || l.startsWith('dphi '),
370 ).length,
371 );
372 if (niter === 1) {
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 373 await session.seed(1);
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 374 session.step(STEPS);
375 states.push(await session.read('U'));
376 }
377 session.destroy();
378 }
379 xformsPerIter.push(counts[1] - counts[0]);
380 }
3d078cfSplit the flux-form divergence against the round sphereDan Fortunato 382 // The headline number, from the compiled op sequences: 6 Legendre
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 383 // transforms per species per iteration against Algorithm 4's 12
3d078cfSplit the flux-form divergence against the round sphereDan Fortunato 384 // (2 species here). Five of the six are the flux matvec; the sixth is
385 // the round-sphere synthesis the divergence split buys its polar
386 // conditioning with. The phi flux's derivative runs as dphig -- two
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 387 // Fourier stages, no Legendre work -- and is deliberately not counted.
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 388 check(
3d078cfSplit the flux-form divergence against the round sphereDan Fortunato 389 'flux: 6 Legendre transforms per species per iteration, versus 12',
390 xformsPerIter[0] === 12 && xformsPerIter[1] === 24,
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 391 `flux form adds ${xformsPerIter[0]} transforms/iteration, ` +
392 `Algorithm 4 adds ${xformsPerIter[1]}`,
393 );
395 // Same operator, same discretization, different arithmetic path: after
396 // STEPS steps the two states may differ only by fp32 accumulation. A
397 // formulation error (wrong weight, wrong shift, missing sin) would show
398 // up at O(1), not O(1e-3). Identical states would mean the A/B compared
399 // one path to itself.
400 let worst = 0;
401 let identical = true;
402 let finite = true;
403 for (let i = 0; i < states[0].length; i++) {
404 const d = Math.abs(states[0][i] - states[1][i]);
405 if (d > worst) worst = d;
406 if (states[0][i] !== states[1][i]) identical = false;
407 if (!Number.isFinite(states[0][i]) || !Number.isFinite(states[1][i])) finite = false;
408 }
409 check(
410 'flux: tracks the Algorithm-4 reference through a real simulation',
411 finite && !identical && worst < 5e-3,
412 `max |U_flux - U_alg4| = ${worst.toExponential(2)} after ${STEPS} steps ` +
413 `on bumpy at lmax ${LMAX_AB}`,
414 );
416 // ---- 3. the correction must not manufacture its own perturbation -----
417 //
418 // From the exact uniform steady state, with the Turing band switched off
419 // (D1 = D2) so nothing can grow on its own, the only thing driving the
420 // state away from uniform is round-off. niter = 0 never touches the flux
421 // machinery and sets the floor; niter = 3 runs it three times per step.
422 // The ratio is the correction's noise gain. Sphere-split it is O(1); with
423 // r multiplying the whole divergence it was ~50 at lmax 63, and that
424 // margin is what decides where a pattern nucleates. The ellipsoid is the
425 // case to run it on: axisymmetric grid, strongly non-spherical geometry.
426 const model = mModelByKey('schnakenberg')!;
427 const quiet = model.source.replace(
428 /function \[U, V, u, v\] = init\([\s\S]*?\nend/,
429 `function [U, V, u, v] = init(lam3, gx, gy, gz, a, b)
430 us = a + b;
431 vs = b / (us * us);
432 [U, V] = analys(us * ones(numel(gx), 1), vs * ones(numel(gx), 1));
433 [u, v] = synth(U, V);
434end`,
435 );
436 if (quiet === model.source) throw new Error('quiet-start fixture no longer matches schnakenberg.m');
437 const ell = mGeometryByKey('ellipsoid')!;
438 const noise: number[] = [];
439 for (const niter of [0, 3]) {
440 const session = await ModelSession.create({
441 device,
442 model,
443 params: { ...defaultParams(model), D2: defaultParams(model).D1 },
444 lmax: LMAX_AB,
445 source: quiet,
446 niter,
447 geometry: ell,
448 geometryParams: defaultGeometryParams(ell),
449 });
450 await session.seed(1);
451 session.step(400);
452 const U = await session.read('U');
453 // Everything above the mean: l = 0, m = 0 is the uniform state itself.
454 let sum = 0;
455 for (let m = 0; m <= LMAX_AB; m++) {
456 for (let l = Math.max(m, 1); l <= LMAX_AB; l++) {
457 const i = lmIndex(LMAX_AB, l, m);
458 sum += (U[2 * i] ** 2 + U[2 * i + 1] ** 2) * (m === 0 ? 1 : 2);
459 }
460 }
461 noise.push(Math.sqrt(sum));
462 session.destroy();
463 }
464 const gain = noise[1] / noise[0];
465 log(
466 ` flux polar noise gain on ellipsoid: ||U'|| ${noise[0].toExponential(2)} ` +
467 `at niter 0, ${noise[1].toExponential(2)} at niter 3`,
468 );
469 check(
470 'flux: the geometric correction does not amplify polar round-off',
471 Number.isFinite(gain) && gain < 5,
472 `niter-3 round-off is ${gain.toFixed(1)}x the niter-0 floor ` +
473 `(sphere-split: ~1; r on the whole divergence: ~50)`,
474 );
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 475 }
476}