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.
33 */
34import { ShtPlan } from '../src/sht/sht.ts';
35import { DerivPlan } from '../src/sht/deriv.ts';
36import { ShtReference } from '../src/sht/reference.ts';
37import { gridForLmax, lmIndex, nlmCalc, type ShtConfig } from '../src/sht/layout.ts';
38import { ModelSession } from '../src/mgpu/session.ts';
39import { mModelByKey, defaultParams } from '../src/mgpu/registry.ts';
40import { Geometry } from '../src/geom/geometry.ts';
41import { mGeometryByKey, defaultGeometryParams } from '../src/geom/registry.ts';
42import { computeFluxMetric } from '../src/geom/metric.ts';
43import type { Check, Log } from './analyticChecks.ts';
45/** Band limit of the test surface and field. */
46const LMAX = 24;
47/** Band limit of the oversampled analysis grid the tails are measured on. */
48const LMAX_HI = 63;
49/** Degrees at and above this count as "beyond-band tail": LMAX+1 is the last
50 * degree with direct content, and the smooth-but-not-band-limited metric
51 * weights spread it upward with (their own) exponentially decaying spectra,
52 * so the window starts well above the band edge. */
53const TAIL_START = 44;
55/** The part of a transform layout the spectral helpers need. */
56interface Band {
57 lmax: number;
58 mmax: number;
59}
61/** Per-degree spectral amplitude: E(l) = sqrt(sum_m |q_l^m|^2). */
62function degreeEnergy(band: Band, qlm: ArrayLike<number>): Float64Array {
63 const E = new Float64Array(band.lmax + 1);
64 for (let m = 0; m <= band.mmax; m++) {
65 for (let l = m; l <= band.lmax; l++) {
66 const i = lmIndex(band.lmax, l, m);
67 E[l] += qlm[2 * i] ** 2 + qlm[2 * i + 1] ** 2;
68 }
69 }
70 for (let l = 0; l <= band.lmax; l++) E[l] = Math.sqrt(E[l]);
71 return E;
72}
74/** max E(l) over l >= TAIL_START, relative to max E(l) overall. */
75function tailRel(band: Band, qlm: ArrayLike<number>): number {
76 const E = degreeEnergy(band, qlm);
77 let bulk = 0;
78 let tail = 0;
79 for (let l = 0; l <= band.lmax; l++) {
80 if (E[l] > bulk) bulk = E[l];
81 if (l >= TAIL_START && E[l] > tail) tail = E[l];
82 }
83 return tail / Math.max(bulk, 1e-300);
84}
86/** Re-index coefficients from the lo layout into the hi layout (zero-padded). */
87function padSpectrum(qlo: ArrayLike<number>, lo: Band, hi: Band): Float64Array {
88 const out = new Float64Array(2 * nlmCalc(hi.lmax, hi.mmax));
89 for (let m = 0; m <= lo.mmax; m++) {
90 for (let l = m; l <= lo.lmax; l++) {
91 const src = lmIndex(lo.lmax, l, m);
92 const dst = lmIndex(hi.lmax, l, m);
93 out[2 * dst] = qlo[2 * src];
94 out[2 * dst + 1] = qlo[2 * src + 1];
95 }
96 }
97 return out;
98}
100/** Deterministic random band-limited spectrum with O(1) coefficients. */
101function flatSpectrum(band: Band, seed: number): Float64Array {
102 const nlm = nlmCalc(band.lmax, band.mmax);
103 const q = new Float64Array(2 * nlm);
104 let s = seed >>> 0;
105 const rnd = () => {
106 s ^= s << 13; s >>>= 0;
107 s ^= s >> 17;
108 s ^= s << 5; s >>>= 0;
109 return (s / 4294967296) * 2 - 1;
110 };
111 for (let k = 0; k < 2 * nlm; k++) q[k] = rnd();
112 for (let l = 0; l <= band.lmax; l++) q[2 * lmIndex(band.lmax, l, 0) + 1] = 0;
113 return q;
114}
116export interface FluxCheckOptions {
117 /**
118 * Run the live flux-vs-Algorithm-4 A/B (4 sessions at lmax 63). On by
119 * default, but — like geometryChecks' sweep, and for the same reason — a
120 * browser recompiles every session's unrolled step from scratch on software
121 * WebGPU, so the page leaves it out unless asked (?sweep=1) to keep CI
122 * short. The f64 smoothness checks always run; they are CPU work.
123 */
124 ab?: boolean;
125}
127export async function fluxChecks(
128 device: GPUDevice,
129 check: Check,
130 log: Log,
131 opts: FluxCheckOptions = {},
132): Promise<void> {
133 // ---- 1a. round sphere: closed-form weights, decisive discrimination -----
134 {
135 const hiGrid = gridForLmax(LMAX_HI, 1);
136 const hi = { lmax: LMAX_HI, mmax: LMAX_HI, nlat: hiGrid.nlat, nphi: hiGrid.nphi };
137 const ref = new ShtReference(hi);
138 const npts = hi.nlat * hi.nphi;
140 // The unit sphere needs no GPU build: analyse the closed-form embedding
141 // on the fine grid directly, in f64.
142 const xg = new Float64Array(npts);
143 const yg = new Float64Array(npts);
144 const zg = new Float64Array(npts);
145 for (let i = 0; i < hi.nlat; i++) {
146 const ct = ref.ct[i];
147 const st = ref.st[i];
148 for (let j = 0; j < hi.nphi; j++) {
149 const phi = (2 * Math.PI * j) / hi.nphi;
150 const k = i * hi.nphi + j;
151 xg[k] = st * Math.cos(phi);
152 yg[k] = st * Math.sin(phi);
153 zg[k] = ct;
154 }
155 }
156 const X = ref.analys(xg);
157 const Y = ref.analys(yg);
158 const Z = ref.analys(zg);
160 const sXt = [ref.sinDtheta(X), ref.sinDtheta(Y), ref.sinDtheta(Z)];
161 const Xp = [ref.dphi(X), ref.dphi(Y), ref.dphi(Z)];
162 const { p1, p2, q2, r } = computeFluxMetric(
163 npts, sXt[0], sXt[1], sXt[2], Xp[0], Xp[1], Xp[2],
164 );
166 // On the sphere the weights have a closed form: p1 = q2 = 1, p2 = 0,
167 // r = 1/sin^2(theta) — the flux-form counterpart of geometryChecks'
168 // closed-form V check, pinning computeFluxMetric before it is buried
169 // under the operator. f64 throughout, so the tolerance is conditioning
170 // at the polar rings, not fp32.
171 let worst = 0;
172 for (let i = 0; i < hi.nlat; i++) {
173 const st2 = ref.st[i] * ref.st[i];
174 for (let j = 0; j < hi.nphi; j++) {
175 const k = i * hi.nphi + j;
176 worst = Math.max(
177 worst,
178 Math.abs(p1[k] - 1),
179 Math.abs(p2[k]),
180 Math.abs(q2[k] - 1),
181 Math.abs(r[k] * st2 - 1),
182 );
183 }
184 }
185 check(
186 'flux: sphere weights match the closed form (p1 = q2 = 1, p2 = 0, r = 1/sin^2)',
187 worst < 1e-9,
188 `max deviation ${worst.toExponential(2)} in f64`,
189 );
191 // Flat random u, band-limited at LMAX. The properly weighted fluxes are
192 // then *exactly* band-limited (P = sin(theta) dtheta u, Qtilde = dphi u),
193 // so their beyond-band tails are pure round-off; the Sec 8 control
194 // Qtilde/sin(theta) is not a function on the sphere and keeps a fat tail.
195 const band = { lmax: LMAX, mmax: LMAX };
196 const u = padSpectrum(flatSpectrum(band, 777), band, hi);
197 const A = ref.sinDtheta(u);
198 const B = ref.dphi(u);
199 const P = new Float64Array(npts);
200 const Qt = new Float64Array(npts);
201 const control = new Float64Array(npts);
202 for (let i = 0; i < hi.nlat; i++) {
203 const st = ref.st[i];
204 for (let j = 0; j < hi.nphi; j++) {
205 const k = i * hi.nphi + j;
206 P[k] = p1[k] * A[k] + p2[k] * B[k];
207 Qt[k] = p2[k] * A[k] + q2[k] * B[k];
208 control[k] = Qt[k] / st;
209 }
210 }
211 const tails = {
212 P: tailRel(hi, ref.analys(P)),
213 Qt: tailRel(hi, ref.analys(Qt)),
214 control: tailRel(hi, ref.analys(control)),
215 };
216 log(
217 ` flux smoothness on the sphere (f64, band ${LMAX}, analysed to ${LMAX_HI}, ` +
218 `tail l >= ${TAIL_START}): P ${tails.P.toExponential(2)}, ` +
219 `Qt ${tails.Qt.toExponential(2)}, control ${tails.control.toExponential(2)}`,
220 );
221 check(
222 'flux: on the sphere the fluxes are band-limited and the non-smooth control is not',
223 tails.P < 1e-10 && tails.Qt < 1e-10 &&
224 tails.control > 1e3 * Math.max(tails.P, tails.Qt, 1e-14),
225 `P ${tails.P.toExponential(2)}, Qt ${tails.Qt.toExponential(2)}, ` +
226 `control ${tails.control.toExponential(2)}`,
227 );
228 }
230 // ---- 1b. bumpy: the fluxes sit on the Cartesian gradient's footing ------
231 {
232 // The surface: bumpy, the one shipped geometry that is genuinely
233 // non-axisymmetric (g_thetaphi != 0), so the off-diagonal weight p2 is
234 // exercised. Built by the real pipeline at LMAX, then everything below is
235 // CPU f64 from its band-limited coefficients.
236 const g = mGeometryByKey('bumpy')!;
237 const { nlat, nphi } = gridForLmax(LMAX, 3);
238 const cfg = { lmax: LMAX, mmax: LMAX, nlat, nphi };
239 const sht = await ShtPlan.create(device, cfg);
240 const deriv = await DerivPlan.create(device, sht);
241 const geometry = await Geometry.create({
242 device, sht, cfg,
243 source: g.source,
244 paramNames: g.params.map((p) => p.key),
245 params: defaultGeometryParams(g),
246 deriv,
247 });
248 deriv.destroy();
249 sht.destroy();
251 const hiGrid = gridForLmax(LMAX_HI, 1);
252 const hi = { lmax: LMAX_HI, mmax: LMAX_HI, nlat: hiGrid.nlat, nphi: hiGrid.nphi };
253 const ref = new ShtReference(hi);
254 const npts = hi.nlat * hi.nphi;
256 // Embedding and test field, zero-padded into the fine layout. Both are
257 // band-limited at LMAX, so on the fine grid every derived field's content
258 // beyond the band is genuinely the non-band-limited part of the weights —
259 // the thing being measured — and not aliasing.
260 const X = padSpectrum(geometry.X, cfg, hi);
261 const Y = padSpectrum(geometry.Y, cfg, hi);
262 const Z = padSpectrum(geometry.Z, cfg, hi);
263 const u = padSpectrum(flatSpectrum(cfg, 777), cfg, hi);
265 // Tangents, both weightings, all f64.
266 const sXt = [ref.sinDtheta(X), ref.sinDtheta(Y), ref.sinDtheta(Z)];
267 const Xp = [ref.dphi(X), ref.dphi(Y), ref.dphi(Z)];
268 const Xt = [ref.dtheta(X), ref.dtheta(Y), ref.dtheta(Z)];
269 const { p1, p2, q2 } = computeFluxMetric(
270 npts, sXt[0], sXt[1], sXt[2], Xp[0], Xp[1], Xp[2],
271 );
273 const A = ref.sinDtheta(u); // sin(theta) dtheta u
274 const B = ref.dphi(u); // dphi u
276 // The two fluxes. No non-smooth control here: on a deformed surface
277 // every smooth field's beyond-band tail is set by the weights' own
278 // (slowly decaying) spectra, which swamps a pole singularity at this
279 // resolution — the sphere block above is where the discrimination has
280 // teeth. This block asserts the doc's relative criterion instead.
281 const P = new Float64Array(npts);
282 const Qt = new Float64Array(npts);
283 for (let k = 0; k < npts; k++) {
284 P[k] = p1[k] * A[k] + p2[k] * B[k];
285 Qt[k] = p2[k] * A[k] + q2[k] * B[k];
286 }
288 // The known-smooth yardstick (doc Sec 2): the x component of the
289 // Cartesian surface gradient, built the Algorithm-4 way from the inverse
290 // metric quantities, in f64.
291 const gradx = new Float64Array(npts);
292 {
293 const ut = ref.dtheta(u);
294 const up = ref.dphi(u);
295 for (let k = 0; k < npts; k++) {
296 const gtt = Xt[0][k] ** 2 + Xt[1][k] ** 2 + Xt[2][k] ** 2;
297 const gtp = Xt[0][k] * Xp[0][k] + Xt[1][k] * Xp[1][k] + Xt[2][k] * Xp[2][k];
298 const gpp = Xp[0][k] ** 2 + Xp[1][k] ** 2 + Xp[2][k] ** 2;
299 const det = gtt * gpp - gtp * gtp;
300 const Vtx = (gpp * Xt[0][k] - gtp * Xp[0][k]) / det;
301 const Vpx = (gtt * Xp[0][k] - gtp * Xt[0][k]) / det;
302 gradx[k] = ut[k] * Vtx + up[k] * Vpx;
303 }
304 }
306 const tails = {
307 P: tailRel(hi, ref.analys(P)),
308 Qt: tailRel(hi, ref.analys(Qt)),
309 gradx: tailRel(hi, ref.analys(gradx)),
310 };
311 log(
312 ` flux smoothness on bumpy (f64, band ${LMAX}, analysed to ${LMAX_HI}, ` +
313 `tail l >= ${TAIL_START}): P ${tails.P.toExponential(2)}, ` +
314 `Qt ${tails.Qt.toExponential(2)}, gradx ${tails.gradx.toExponential(2)}`,
315 );
316 // "Matching tails" (Sec 7.1): same footing as the Cartesian component,
317 // with an order of magnitude of headroom on top of it. A wrong weighting
318 // (a missing sin factor, say) puts genuinely non-smooth content into P or
319 // Qtilde and the tail lands at O(bulk), far above this.
320 const ceiling = Math.max(30 * tails.gradx, 1e-10);
321 check(
322 'flux: on bumpy, P and Qtilde tails match the Cartesian gradient component',
323 tails.P < ceiling && tails.Qt < ceiling,
324 `P ${tails.P.toExponential(2)}, Qt ${tails.Qt.toExponential(2)} vs ` +
325 `ceiling ${ceiling.toExponential(2)}`,
326 );
327 }
329 // ---- 2. flux form vs Algorithm 4, live, on a curved surface -------------
330 if (!(opts.ab ?? true)) {
331 log(
332 ' flux A/B: skipped — run `npm run test:node` (desktop Dawn) or ' +
333 '`npm run test:gpu -- --sweep` for the flux-vs-Algorithm-4 comparison.',
334 );
335 } else {
336 const geometry = mGeometryByKey('bumpy')!;
337 const geometryParams = defaultGeometryParams(geometry);
338 const LMAX_AB = 63;
339 const STEPS = 20;
340 const states: Float32Array[] = [];
341 const xformsPerIter: number[] = [];
343 for (const key of ['schnakenberg', 'schnakenberg-alg4']) {
344 const model = mModelByKey(key)!;
345 const params = defaultParams(model);
346 // Real transforms added by one solve iteration: synth/analys ops plus
347 // dtheta/dphi (each of which contains a synthesis); the coefficient-
348 // space dthetac/dphic shuffles are O(nlm) index gathers, not transforms.
349 const counts: number[] = [];
350 for (const niter of [0, 1]) {
351 const session = await ModelSession.create({
352 device, model, params, lmax: LMAX_AB,
353 geometry, geometryParams, niter,
354 });
355 counts.push(
356 session.describe().step.filter((l) =>
357 l.startsWith('synth') || l.startsWith('analys') ||
358 l.startsWith('dtheta ') || l.startsWith('dphi '),
359 ).length,
360 );
361 if (niter === 1) {
362 session.seed(1);
363 session.step(STEPS);
364 states.push(await session.read('U'));
365 }
366 session.destroy();
367 }
368 xformsPerIter.push(counts[1] - counts[0]);
369 }
371 // The headline number, from the compiled op sequences: 5 Legendre
372 // transforms per species per iteration against Algorithm 4's 12
373 // (2 species here). The phi flux's derivative runs as dphig -- two
374 // Fourier stages, no Legendre work -- and is deliberately not counted.
375 check(
376 'flux: 5 Legendre transforms per species per iteration, versus 12',
377 xformsPerIter[0] === 10 && xformsPerIter[1] === 24,
378 `flux form adds ${xformsPerIter[0]} transforms/iteration, ` +
379 `Algorithm 4 adds ${xformsPerIter[1]}`,
380 );
382 // Same operator, same discretization, different arithmetic path: after
383 // STEPS steps the two states may differ only by fp32 accumulation. A
384 // formulation error (wrong weight, wrong shift, missing sin) would show
385 // up at O(1), not O(1e-3). Identical states would mean the A/B compared
386 // one path to itself.
387 let worst = 0;
388 let identical = true;
389 let finite = true;
390 for (let i = 0; i < states[0].length; i++) {
391 const d = Math.abs(states[0][i] - states[1][i]);
392 if (d > worst) worst = d;
393 if (states[0][i] !== states[1][i]) identical = false;
394 if (!Number.isFinite(states[0][i]) || !Number.isFinite(states[1][i])) finite = false;
395 }
396 check(
397 'flux: tracks the Algorithm-4 reference through a real simulation',
398 finite && !identical && worst < 5e-3,
399 `max |U_flux - U_alg4| = ${worst.toExponential(2)} after ${STEPS} steps ` +
400 `on bumpy at lmax ${LMAX_AB}`,
401 );
402 }
403}