1/**
2 * The two things this project adds to turing-sphere: a surface, and a `for`
3 * loop in the compiled step.
4 *
5 * The surface is checked against what it is supposed to be — the sphere really
6 * is the unit sphere and really is degree 1, a deformed shape really has the
7 * radius profile its .m says, and the coefficients really do evaluate to the
8 * same surface on a finer grid.
9 *
10 * The loop is checked for the property the whole design rests on: it is
11 * unrolled into the fixed op sequence, so more iterations means more GPU ops.
12 * On the sphere, where the surface Laplace-Beltrami correction is
13 * mathematically zero (lap_g = lap_s exactly), the answer must stay close
14 * across niter to fp32 tolerance — not bit-identical, since the correction is
15 * now a real (if numerically near-zero) computation rather than the literal
16 * `0 * Un` placeholder, so the op sequence differs even though the answer
17 * shouldn't move much. On a genuinely curved surface the correction must
18 * actually change the answer, and — since the Richardson iteration only
19 * converges while the correction stays small relative to what the
20 * round-sphere solve inverts (docs/richardson-iteration.md) — a niter/dt/
21 * geometry combination outside that radius is expected to diverge. The
22 * niter x geometry sweep below documents which shipped combinations that
23 * currently affects, so a regression that makes a *currently-healthy*
24 * combination diverge is caught without this file silently asserting away a
25 * real, known numerical limit.
26 */
27import { ShtPlan } from '../src/sht/sht.ts';
28import { DerivPlan } from '../src/sht/deriv.ts';
29import { gridForLmax, lmIndex } from '../src/sht/layout.ts';
30import { ModelSession } from '../src/mgpu/session.ts';
31import { mModelByKey, defaultParams } from '../src/mgpu/registry.ts';
32import { Geometry } from '../src/geom/geometry.ts';
33import {
34 mGeometries,
35 mGeometryByKey,
36 defaultGeometryParams,
37 SPHERE_KEY,
38} from '../src/geom/registry.ts';
39import { ModelCompileError } from '../src/mgpu/errors.ts';
40import type { Check, Log } from './analyticChecks.ts';
42const LMAX = 31;
43const STEPS = 20;
44/** The app's actual default lmax (README: "at the default lmax 63 that is a
45 * 128x256 grid"), used for the niter/geometry sweep below and the peanut
46 * check next to it -- the divergence they're both about is a real, lmax-
47 * dependent numerical property of the Richardson iteration, not one this
48 * file's other, smaller LMAX happens to reproduce. */
49const SWEEP_LMAX = 63;
51/** Build one geometry on its own transform plan, for inspection. */
52async function buildGeometry(device: GPUDevice, key: string) {
53 const g = mGeometryByKey(key)!;
54 const { nlat, nphi } = gridForLmax(LMAX, 3);
55 const cfg = { lmax: LMAX, mmax: LMAX, nlat, nphi };
56 const sht = await ShtPlan.create(device, cfg);
57 const deriv = await DerivPlan.create(device, sht);
58 const geometry = await Geometry.create({
59 device,
60 sht,
61 cfg,
62 source: g.source,
63 paramNames: g.params.map((p) => p.key),
64 params: defaultGeometryParams(g),
65 deriv,
66 });
67 return { g, sht, deriv, cfg, geometry };
68}
70export interface GeometryCheckOptions {
71 /**
72 * Run the niter x geometry sweep at the end. On by default, and nearly free
73 * on desktop Dawn (~3 s for all 20 combinations), but in a browser it
74 * dominates the whole suite: every session recompiles its unrolled step from
75 * scratch — there is no pipeline cache across sessions — so the sweep costs
76 * ~6 minutes on software WebGPU against ~30 s for every other check here put
77 * together. The browser page therefore leaves it out unless asked (?sweep=1),
78 * which is what keeps CI short.
79 */
80 sweep?: boolean;
81}
83export async function geometryChecks(
84 device: GPUDevice,
85 check: Check,
86 log: Log,
87 opts: GeometryCheckOptions = {},
88): Promise<void> {
89 const runSweep = opts.sweep ?? true;
90 // ---- every geometry compiles and closes ---------------------------------
91 for (const spec of mGeometries) {
92 const { sht, deriv, geometry } = await buildGeometry(device, spec.key);
93 let finite = true;
94 for (const a of [geometry.x, geometry.y, geometry.z]) {
95 for (const v of a) if (!Number.isFinite(v)) finite = false;
96 }
97 for (const a of [geometry.Vtx, geometry.Vty, geometry.Vtz, geometry.Vpx, geometry.Vpy, geometry.Vpz]) {
98 for (const v of a) if (!Number.isFinite(v)) finite = false;
99 }
100 const { lo, hi } = geometry.radiusRange();
101 check(
102 `geometry: ${spec.key}.m evaluates to a finite surface`,
103 finite && lo > 1e-3,
104 `radius ${lo.toFixed(4)}–${hi.toFixed(4)}`,
105 );
106 deriv.destroy();
107 sht.destroy();
108 }
110 // ---- the sphere is the unit sphere, exactly, and is degree 1 ------------
111 {
112 const { sht, deriv, geometry } = await buildGeometry(device, SPHERE_KEY);
114 let maxRadiusErr = 0;
115 for (let i = 0; i < geometry.x.length; i++) {
116 const r = Math.hypot(geometry.x[i], geometry.y[i], geometry.z[i]);
117 maxRadiusErr = Math.max(maxRadiusErr, Math.abs(r - 1));
118 }
119 // Tolerance is fp32 through a full analysis/synthesis round trip, not the
120 // geometry: the exact answer is representable, and what is measured here
121 // is the transforms' own round-off. It is set by the loosest stack this
122 // runs on — SwiftShader in CI is an order of magnitude worse than Dawn on
123 // real hardware (4e-4 against 2e-5). A geometry that was actually wrong
124 // would miss by O(1), so the slack costs nothing.
125 check(
126 'geometry: sphere.m has radius 1 everywhere',
127 maxRadiusErr < 2e-3,
128 `max |r - 1| = ${maxRadiusErr.toExponential(2)}`,
129 );
131 // x, y, z of the unit sphere are the three degree-1 harmonics and nothing
132 // else, so analysing them must leave every other coefficient at zero.
133 // This is what makes the sphere case exact rather than merely accurate:
134 // there is no content for the band limit to throw away.
135 const degreeOne = new Set([
136 lmIndex(LMAX, 1, 0),
137 lmIndex(LMAX, 1, 1),
138 ]);
139 let leak = 0;
140 for (const coeffs of [geometry.X, geometry.Y, geometry.Z]) {
141 for (let i = 0; i < coeffs.length / 2; i++) {
142 if (degreeOne.has(i)) continue;
143 leak = Math.max(leak, Math.abs(coeffs[2 * i]), Math.abs(coeffs[2 * i + 1]));
144 }
145 }
146 check(
147 'geometry: sphere.m is exactly degree 1 in the harmonics',
148 leak < 1e-3,
149 `max |coefficient| outside l = 1 is ${leak.toExponential(2)}`,
150 );
152 // The inverse metric quantities have a closed form on the unit sphere:
153 // V_theta = (cos(theta)cos(phi), cos(theta)sin(phi), -sin(theta)),
154 // V_phi = (-sin(phi)/sin(theta), cos(phi)/sin(theta), 0). Checking these
155 // pins the sign convention of computeMetric (src/geom/metric.ts) before
156 // it is buried under the Laplace-Beltrami operator built on top of it.
157 //
158 // Split by latitude, because the accuracy available here is not uniform.
159 // computeMetric divides by det = g_tt*g_pp - g_tp^2, and on the sphere
160 // g_pp and det are both O(sin^2 theta) -- 1.4e-3 at the outermost Gauss
161 // latitude of this grid. The transforms deliver X_theta/X_phi with an
162 // *absolute* fp32 error of ~1e-6, which is a large *relative* error once
163 // it is squared into quantities that small, so the error in both V's grows
164 // like 1/sin^2 theta toward the poles. That is conditioning, not a wrong
165 // formula: measured worst cases are
166 //
167 // sin(theta) >= 0.2 sin(theta) < 0.2 (8 of 64 latitudes)
168 // Dawn 3.2e-4 1.7e-3
169 // SwiftShader 1.5e-3 8.8e-3
170 //
171 // and a sign or formula error would be O(1) in either band, so tolerances
172 // a few times the looser stack still catch one.
173 const POLE_SIN = 0.2;
174 let maxMetricErr = 0;
175 let maxPoleErr = 0;
176 for (let i = 0; i < sht.cfg.nlat; i++) {
177 const ct = sht.cosTheta[i];
178 const st = Math.sqrt(Math.max(0, 1 - ct * ct));
179 for (let j = 0; j < sht.cfg.nphi; j++) {
180 const phi = (2 * Math.PI * j) / sht.cfg.nphi;
181 const k = i * sht.cfg.nphi + j;
182 const cphi = Math.cos(phi);
183 const sphi = Math.sin(phi);
184 const wantVtx = ct * cphi;
185 const wantVty = ct * sphi;
186 const wantVtz = -st;
187 const wantVpx = -sphi / st;
188 const wantVpy = cphi / st;
189 const wantVpz = 0;
190 const worst = Math.max(
191 Math.abs(geometry.Vtx[k] - wantVtx),
192 Math.abs(geometry.Vty[k] - wantVty),
193 Math.abs(geometry.Vtz[k] - wantVtz),
194 Math.abs(geometry.Vpx[k] - wantVpx),
195 Math.abs(geometry.Vpy[k] - wantVpy),
196 Math.abs(geometry.Vpz[k] - wantVpz),
197 );
198 if (st < POLE_SIN) maxPoleErr = Math.max(maxPoleErr, worst);
199 else maxMetricErr = Math.max(maxMetricErr, worst);
200 }
201 }
202 check(
203 'geometry: sphere.m has the closed-form inverse metric quantities',
204 maxMetricErr < 4e-3,
205 `max |V - closed form| = ${maxMetricErr.toExponential(2)} ` +
206 `away from the poles (sin theta >= ${POLE_SIN})`,
207 );
208 check(
209 'geometry: the polar caps stay within their conditioning',
210 maxPoleErr < 2e-2,
211 `max |V - closed form| = ${maxPoleErr.toExponential(2)} at sin theta < ${POLE_SIN}`,
212 );
214 deriv.destroy();
215 sht.destroy();
216 }
218 // ---- a deformed surface matches its own formula, on any grid ------------
219 {
220 const { g, sht, deriv, cfg, geometry } = await buildGeometry(device, 'peanut');
221 const p = defaultGeometryParams(g);
223 // peanut.m written out: r = 1 - waist*sin(theta)^2 scales the unit sphere,
224 // and z is then stretched, so the distance from the origin depends on
225 // theta alone. Checking every point against this closed form checks the
226 // whole path at once — the compiled shape kernel, the analysis into
227 // coefficients, the synthesis back — and, because the formula has no phi
228 // in it, that the surface really is a surface of revolution.
229 const peanutRadius = (ct: number): number => {
230 const st2 = Math.max(0, 1 - ct * ct);
231 const r = 1 - p.waist * st2;
232 return r * Math.hypot(Math.sqrt(st2), (1 + p.stretch) * ct);
233 };
235 const onGrid = (
236 cosTheta: Float64Array,
237 nlat: number,
238 nphi: number,
239 at: (i: number) => number,
240 ): number => {
241 let worst = 0;
242 for (let i = 0; i < nlat; i++) {
243 const want = peanutRadius(cosTheta[i]);
244 for (let j = 0; j < nphi; j++) {
245 worst = Math.max(worst, Math.abs(at(i * nphi + j) - want));
246 }
247 }
248 return worst;
249 };
251 const coarse = onGrid(sht.cosTheta, cfg.nlat, cfg.nphi, (k) =>
252 Math.hypot(geometry.x[k], geometry.y[k], geometry.z[k]),
253 );
254 check(
255 'geometry: peanut.m matches its own radial formula on the solver grid',
256 coarse < 1e-3,
257 `max |dr| = ${coarse.toExponential(2)}`,
258 );
260 // And the same on a finer grid, from the same coefficients. This is what
261 // "the rendered surface is the surface being solved on" means: display
262 // oversampling evaluates the embedding at more points, it does not
263 // subdivide or smooth it. The 2x Gauss latitudes share no point with the
264 // 1x ones, so agreeing here is agreeing everywhere, not at samples.
265 const fine = await ShtPlan.create(device, {
266 lmax: cfg.lmax,
267 mmax: cfg.mmax,
268 nlat: 2 * cfg.nlat,
269 nphi: 2 * cfg.nphi,
270 });
271 const finePos = await geometry.positionsOn(fine);
272 const refined = onGrid(fine.cosTheta, 2 * cfg.nlat, 2 * cfg.nphi, (k) =>
273 Math.hypot(finePos[3 * k], finePos[3 * k + 1], finePos[3 * k + 2]),
274 );
275 check(
276 'geometry: the same coefficients give the same surface on a 2x grid',
277 refined < 1e-3,
278 `max |dr| = ${refined.toExponential(2)} at ${2 * cfg.nlat}×${2 * cfg.nphi} points`,
279 );
280 fine.destroy();
281 deriv.destroy();
282 sht.destroy();
283 }
285 // ---- the unrolled loop: more ops, identical answer ----------------------
286 {
287 const model = mModelByKey('schnakenberg')!;
288 const params = defaultParams(model);
289 const counts = [0, 1, 4];
290 const ops: number[] = [];
291 const states: Float32Array[] = [];
293 for (const niter of counts) {
294 const session = await ModelSession.create({
295 device, model, params, lmax: LMAX, niter,
296 });
297 ops.push(session.describe().step.length);
298 session.seed(1);
299 session.step(STEPS);
300 states.push(await session.read('U'));
301 session.destroy();
302 }
304 log(` schnakenberg.m ops/step by solve iterations: ${
305 counts.map((n, i) => `${n} -> ${ops[i]}`).join(', ')
306 }`);
307 check(
308 'loop: each solve iteration adds GPU operations',
309 ops[0] < ops[1] && ops[1] < ops[2],
310 `${ops.join(' < ')} ops for ${counts.join(', ')} iterations`,
311 );
312 // Unrolling has to be exactly linear in the trip count: the body planned
313 // once per iteration, no more and no less. Per species per iteration: 8
314 // dtheta/dphi + 4 analys transforms (Algorithm 3's cost, applied to the
315 // field and to each of its three Cartesian gradient components) plus 15
316 // generated kernels -- see test/modelChecks.ts's KERNELS_PER_ITERATION,
317 // which counts the kernels alone; this counts every op, transforms
318 // included.
319 const perIteration = ops[1] - ops[0];
320 const want = 54;
321 check(
322 'loop: unrolling is exactly linear in the trip count',
323 perIteration === want && ops[2] - ops[0] === 4 * perIteration,
324 `${perIteration} ops per iteration (expected ${want}), ` +
325 `${ops[2] - ops[0]} for 4 iterations`,
326 );
328 // On the sphere lap_g = lap_s exactly, so the correction should compute
329 // (numerically) close to zero regardless of niter -- not bit-identical
330 // (it is a real computation now, through 8+ chained fp32 transforms per
331 // iteration, not the literal `0 * Un` placeholder that used to make this
332 // exact), but close. The tolerance is set by that chain's fp32 roundoff,
333 // not by the scheme: a real geometry-correction bug would miss by orders
334 // of magnitude more than this.
335 let worst = 0;
336 for (let k = 1; k < states.length; k++) {
337 for (let i = 0; i < states[0].length; i++) {
338 worst = Math.max(worst, Math.abs(states[k][i] - states[0][i]));
339 }
340 }
341 check(
342 'loop: on the sphere, the correction stays near zero across niter',
343 worst < 2e-3,
344 `states differ by up to ${worst.toExponential(2)} after ${STEPS} steps at ${counts.join('/')} iterations`,
345 );
346 }
348 // ---- on a curved surface, the correction actually changes the answer ----
349 {
350 const model = mModelByKey('schnakenberg')!;
351 const params = defaultParams(model);
352 const peanut = mGeometryByKey('peanut')!;
353 const peanutParams = defaultGeometryParams(peanut);
354 // niter 0 vs 1 only -- deliberately not the 4/8 the sweep below already
355 // documents as outside the Richardson iteration's convergence radius on
356 // this geometry. The point here is just that the correction is not a
357 // no-op, which a much smaller, still-converging niter already shows.
358 const states: Float32Array[] = [];
359 for (const niter of [0, 1]) {
360 const session = await ModelSession.create({
361 device, model, params, lmax: SWEEP_LMAX,
362 geometry: peanut, geometryParams: peanutParams, niter,
363 });
364 session.seed(1);
365 session.step(STEPS);
366 states.push(await session.read('U'));
367 session.destroy();
368 }
369 let worst = 0;
370 for (let i = 0; i < states[0].length; i++) {
371 worst = Math.max(worst, Math.abs(states[1][i] - states[0][i]));
372 }
373 check(
374 'loop: on peanut, the correction measurably changes the answer',
375 worst > 1e-4 && states[1].every((v) => Number.isFinite(v)),
376 `states differ by ${worst.toExponential(2)} after ${STEPS} steps at niter 0 vs 1`,
377 );
378 }
380 // ---- niter x geometry sweep: catch a "doesn't run" regression early -----
381 // This is what actually turned up the two real issues found while building
382 // the correction: peanut diverging at niter >= 4 with schnak-spots'
383 // shipped default dt (a genuine Richardson-convergence-radius limit, not a
384 // bug -- see docs/richardson-iteration.md), and a since-fixed compiler bug
385 // where a loop-body statement could silently reuse a *different*
386 // statement's compiled kernel (test/modelChecks.ts's pipeline-cache check
387 // guards that one directly). Every shipped geometry x every niter the
388 // app's <select> actually offers, so a regression anywhere in that grid is
389 // caught -- without asserting away the one combination already known to be
390 // outside the convergence radius.
391 //
392 // SWEEP_LMAX cannot be lowered to make this cheaper: at lmax 31 or 15 the
393 // peanut/2, /4 and /8 combinations below all stay finite, so a smaller grid
394 // would quietly turn KNOWN_DIVERGENT into three failures and stop testing
395 // the thing this sweep exists to pin down.
396 if (!runSweep) {
397 log(
398 ' sweep: skipped — run `npm run test:node` (desktop Dawn, ~3 s) or ' +
399 '`npm run test:gpu -- --sweep` for the niter x geometry sweep.',
400 );
401 } else {
402 const model = mModelByKey('schnakenberg')!;
403 const params = defaultParams(model);
404 const SWEEP_NITER = [0, 1, 2, 4, 8];
405 const KNOWN_DIVERGENT = new Set(['peanut/2', 'peanut/4', 'peanut/8']);
407 for (const geomSpec of mGeometries) {
408 for (const niter of SWEEP_NITER) {
409 const session = await ModelSession.create({
410 device, model, params, lmax: SWEEP_LMAX,
411 geometry: geomSpec, geometryParams: defaultGeometryParams(geomSpec),
412 niter,
413 });
414 session.seed(1);
415 session.step(STEPS);
416 const values = await session.read('u');
417 const finite = values.every((v) => Number.isFinite(v));
418 session.destroy();
420 const key = `${geomSpec.key}/${niter}`;
421 const expectDivergent = KNOWN_DIVERGENT.has(key);
422 check(
423 expectDivergent
424 ? `sweep: ${key} is known to diverge (outside the Richardson convergence radius)`
425 : `sweep: ${key} stays finite after ${STEPS} steps`,
426 expectDivergent ? !finite : finite,
427 expectDivergent
428 ? finite
429 ? 'now finite -- the convergence radius may have improved; update KNOWN_DIVERGENT'
430 : 'diverged as expected'
431 : finite
432 ? 'finite'
433 : 'NOT FINITE -- unexpected divergence, investigate before treating this as another known case',
434 );
435 }
436 }
437 }
439 // ---- a loop whose length is not known at compile time is refused --------
440 {
441 const model = mModelByKey('allencahn')!;
442 // `dt` is a tunable parameter, so it reaches the compiler with no value:
443 // the plan cannot know how many iterations to emit.
444 const bad = model.source.replace('for k = 1:niter', 'for k = 1:dt');
445 let message = '';
446 try {
447 const session = await ModelSession.create({
448 device, model, params: defaultParams(model), lmax: LMAX, source: bad, niter: 1,
449 });
450 session.destroy();
451 } catch (e) {
452 message = e instanceof ModelCompileError ? e.message : `wrong error type: ${e}`;
453 }
454 check(
455 'loop: a runtime loop bound is refused at compile time',
456 message.includes('known when the model is compiled'),
457 message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
458 );
459 }
461 // ---- swapping the surface leaves the simulation alone ------------------
462 {
463 const model = mModelByKey('schnakenberg')!;
464 const session = await ModelSession.create({
465 device, model, params: defaultParams(model), lmax: LMAX,
466 });
467 session.seed(1);
468 session.step(STEPS);
469 const before = await session.read('U');
471 const peanut = mGeometryByKey('peanut')!;
472 await session.setGeometry(peanut, defaultGeometryParams(peanut));
473 const after = await session.read('U');
475 let survived = before.length === after.length;
476 for (let i = 0; survived && i < before.length; i++) {
477 if (before[i] !== after[i]) survived = false;
478 }
479 const { lo, hi } = session.geometry.radiusRange();
480 check(
481 'geometry: swapping the surface mid-run does not disturb the state',
482 survived && session.geometryModel.key === 'peanut' && hi - lo > 0.1,
483 survived
484 ? `state identical, now on ${session.geometryModel.key} (radius ${lo.toFixed(3)}–${hi.toFixed(3)})`
485 : 'state changed',
486 );
487 session.destroy();
488 }
489}
491/** Index of the entry minimizing `score`, over the first `n` entries. */
492function argMin(
493 xs: Float64Array | Float32Array,
494 n: number,
495 score: (v: number) => number,
496): number {
497 let best = 0;
498 let bestScore = Infinity;
499 for (let i = 0; i < n; i++) {
500 const s = score(xs[i]);
501 if (s < bestScore) {
502 bestScore = s;
503 best = i;
504 }
505 }
506 return best;
507}