forked from github.com/concept-collection/turing-surface
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 { boundingBox, drawModes, DEFAULT_LAMBDA } from '../src/mgpu/randnfun3.ts';
41import type { Check, Log } from './analyticChecks.ts';
43const LMAX = 31;
44const STEPS = 20;
45/** The app's actual default lmax (README: "at the default lmax 63 that is a
46 * 128x256 grid"), used for the niter/geometry sweep below and the peanut
47 * check next to it -- the divergence they're both about is a real, lmax-
48 * dependent numerical property of the Richardson iteration, not one this
49 * file's other, smaller LMAX happens to reproduce. */
50const SWEEP_LMAX = 63;
52/** Build one geometry on its own transform plan, for inspection. */
53async function buildGeometry(device: GPUDevice, key: string) {
54 const g = mGeometryByKey(key)!;
55 const { nlat, nphi } = gridForLmax(LMAX, 3);
56 const cfg = { lmax: LMAX, mmax: LMAX, nlat, nphi };
57 const sht = await ShtPlan.create(device, cfg);
58 const deriv = await DerivPlan.create(device, sht);
59 const geometry = await Geometry.create({
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 // ---- a surface handed over as coefficients is the same surface -----------
286 // The import path (reharm's "Export to turing-surface"): Geometry.create
287 // with `coeffs` set evaluates no .m — the coefficients are re-indexed onto
288 // the session's band and everything from synthesis onward runs unchanged.
289 // Checked in all three band relations: same (must reproduce the original
290 // geometry, arrays and all), wider (zero-padding; the surface must still
291 // match its closed form on the finer plan's own grid), and narrower
292 // (truncation; exact when the content fits the smaller band, as the
293 // degree-1 sphere does).
294 {
295 const importGeometry = async (
296 coeffs: { lmax: number; mmax: number; X: Float32Array; Y: Float32Array; Z: Float32Array },
297 lmax: number,
298 ) => {
299 const { nlat, nphi } = gridForLmax(lmax, 3);
300 const cfg = { lmax, mmax: lmax, nlat, nphi };
301 const sht = await ShtPlan.create(device, cfg);
302 const deriv = await DerivPlan.create(device, sht);
303 const geometry = await Geometry.create({
304 sht, cfg, source: '', paramNames: [], params: {}, deriv, coeffs,
305 });
306 return { sht, deriv, cfg, geometry };
307 };
309 const peanut = await buildGeometry(device, 'peanut');
310 const coeffs = {
311 lmax: peanut.cfg.lmax,
312 mmax: peanut.cfg.mmax,
313 X: peanut.geometry.X,
314 Y: peanut.geometry.Y,
315 Z: peanut.geometry.Z,
316 };
318 // Same band: the identical pipeline on identical inputs, so the surface
319 // and both metric formulations must come back numerically unchanged.
320 const same = await importGeometry(coeffs, peanut.cfg.lmax);
321 let worst = 0;
322 const pairs: [Float32Array, Float32Array][] = [
323 [same.geometry.x, peanut.geometry.x],
324 [same.geometry.y, peanut.geometry.y],
325 [same.geometry.z, peanut.geometry.z],
326 [same.geometry.p1, peanut.geometry.p1],
327 [same.geometry.p2, peanut.geometry.p2],
328 [same.geometry.q2, peanut.geometry.q2],
329 [same.geometry.r, peanut.geometry.r],
330 [same.geometry.Vtx, peanut.geometry.Vtx],
331 [same.geometry.Vpx, peanut.geometry.Vpx],
332 ];
333 for (const [a, b] of pairs) {
334 for (let i = 0; i < a.length; i++) {
335 worst = Math.max(worst, Math.abs(a[i] - b[i]));
336 }
337 }
338 check(
339 'import: coefficients rebuild the identical geometry on the same band',
340 worst <= 1e-7,
341 `max |difference| = ${worst.toExponential(2)} over surface and metric`,
342 );
343 same.deriv.destroy();
344 same.sht.destroy();
346 // Wider band: zero-padded onto a finer plan, then checked against the
347 // .m's own closed form on that plan's grid — points the original build
348 // never evaluated (Gauss latitudes do not nest).
349 const p = defaultGeometryParams(peanut.g);
350 const peanutRadius = (ct: number): number => {
351 const st2 = Math.max(0, 1 - ct * ct);
352 const r = 1 - p.waist * st2;
353 return r * Math.hypot(Math.sqrt(st2), (1 + p.stretch) * ct);
354 };
355 const wide = await importGeometry(coeffs, 2 * peanut.cfg.lmax + 1);
356 let dr = 0;
357 for (let i = 0; i < wide.cfg.nlat; i++) {
358 const want = peanutRadius(wide.sht.cosTheta[i]);
359 for (let j = 0; j < wide.cfg.nphi; j++) {
360 const k = i * wide.cfg.nphi + j;
361 const got = Math.hypot(wide.geometry.x[k], wide.geometry.y[k], wide.geometry.z[k]);
362 dr = Math.max(dr, Math.abs(got - want));
363 }
364 }
365 check(
366 'import: zero-padding onto a wider band keeps the surface',
367 dr < 1e-3,
368 `max |dr| = ${dr.toExponential(2)} at lmax ${wide.cfg.lmax}`,
369 );
370 wide.deriv.destroy();
371 wide.sht.destroy();
372 peanut.deriv.destroy();
373 peanut.sht.destroy();
375 // Narrower band: truncation is orthogonal projection, so a surface whose
376 // content fits the smaller band comes through exactly — the sphere is
377 // degree 1 and must still have radius 1 after import at any lmax.
378 const sphere = await buildGeometry(device, SPHERE_KEY);
379 const narrow = await importGeometry(
380 {
381 lmax: sphere.cfg.lmax,
382 mmax: sphere.cfg.mmax,
383 X: sphere.geometry.X,
384 Y: sphere.geometry.Y,
385 Z: sphere.geometry.Z,
386 },
387 Math.floor(sphere.cfg.lmax / 2),
388 );
389 let radiusErr = 0;
390 for (let k = 0; k < narrow.geometry.x.length; k++) {
391 const r = Math.hypot(narrow.geometry.x[k], narrow.geometry.y[k], narrow.geometry.z[k]);
392 radiusErr = Math.max(radiusErr, Math.abs(r - 1));
393 }
394 check(
395 'import: truncation onto a narrower band is exact when the content fits',
396 radiusErr < 2e-3,
397 `sphere from lmax ${sphere.cfg.lmax} coefficients at lmax ${narrow.cfg.lmax}: ` +
398 `max |r - 1| = ${radiusErr.toExponential(2)}`,
399 );
400 narrow.deriv.destroy();
401 narrow.sht.destroy();
402 sphere.deriv.destroy();
403 sphere.sht.destroy();
404 }
406 // ---- the unrolled loop: more ops, identical answer ----------------------
407 {
408 const model = mModelByKey('schnakenberg')!;
409 const params = defaultParams(model);
410 const counts = [0, 1, 4];
411 const ops: number[] = [];
412 const states: Float32Array[] = [];
414 for (const niter of counts) {
415 const session = await ModelSession.create({
416 device, model, params, lmax: LMAX, niter,
417 });
418 ops.push(session.describe().step.length);
419 await session.seed(1);
420 session.step(STEPS);
421 states.push(await session.read('U'));
422 session.destroy();
423 }
425 log(` schnakenberg.m ops/step by solve iterations: ${
426 counts.map((n, i) => `${n} -> ${ops[i]}`).join(', ')
427 }`);
428 check(
429 'loop: each solve iteration adds GPU operations',
430 ops[0] < ops[1] && ops[1] < ops[2],
431 `${ops.join(' < ')} ops for ${counts.join(', ')} iterations`,
432 );
433 // Unrolling has to be exactly linear in the trip count: the body planned
434 // once per iteration, no more and no less. Per species per iteration: 4
435 // synths + 2 analyses (the flux-form matvec's five Legendre transforms,
436 // docs/reduced-transforms.md Sec 4 with the dphig variation, plus the
437 // round-sphere synthesis of the divergence split) + the grid-space
438 // phi-derivative + 3 coefficient-space shuffles plus 8 generated kernels
439 // -- see test/modelChecks.ts's KERNELS_PER_ITERATION, which counts the
440 // kernels alone; this counts every op.
441 const perIteration = ops[1] - ops[0];
442 const want = 36;
443 check(
444 'loop: unrolling is exactly linear in the trip count',
445 perIteration === want && ops[2] - ops[0] === 4 * perIteration,
446 `${perIteration} ops per iteration (expected ${want}), ` +
447 `${ops[2] - ops[0]} for 4 iterations`,
448 );
450 // On the sphere lap_g = lap_s exactly, so the correction should compute
451 // (numerically) close to zero regardless of niter -- not bit-identical
452 // (it is a real computation now, through 8+ chained fp32 transforms per
453 // iteration, not the literal `0 * Un` placeholder that used to make this
454 // exact), but close. The tolerance is set by that chain's fp32 roundoff,
455 // not by the scheme: a real geometry-correction bug would miss by orders
456 // of magnitude more than this.
457 let worst = 0;
458 for (let k = 1; k < states.length; k++) {
459 for (let i = 0; i < states[0].length; i++) {
460 worst = Math.max(worst, Math.abs(states[k][i] - states[0][i]));
461 }
462 }
463 check(
464 'loop: on the sphere, the correction stays near zero across niter',
465 worst < 2e-3,
466 `states differ by up to ${worst.toExponential(2)} after ${STEPS} steps at ${counts.join('/')} iterations`,
467 );
468 }
470 // ---- on a curved surface, the correction actually changes the answer ----
471 {
472 const model = mModelByKey('schnakenberg')!;
473 const params = defaultParams(model);
474 const peanut = mGeometryByKey('peanut')!;
475 const peanutParams = defaultGeometryParams(peanut);
476 // niter 0 vs 1 only -- deliberately not the 4/8 the sweep below already
477 // documents as outside the Richardson iteration's convergence radius on
478 // this geometry. The point here is just that the correction is not a
479 // no-op, which a much smaller, still-converging niter already shows.
480 const states: Float32Array[] = [];
481 for (const niter of [0, 1]) {
482 const session = await ModelSession.create({
483 device, model, params, lmax: SWEEP_LMAX,
484 geometry: peanut, geometryParams: peanutParams, niter,
485 });
486 await session.seed(1);
487 session.step(STEPS);
488 states.push(await session.read('U'));
489 session.destroy();
490 }
491 let worst = 0;
492 for (let i = 0; i < states[0].length; i++) {
493 worst = Math.max(worst, Math.abs(states[1][i] - states[0][i]));
494 }
495 check(
496 'loop: on peanut, the correction measurably changes the answer',
497 worst > 1e-4 && states[1].every((v) => Number.isFinite(v)),
498 `states differ by ${worst.toExponential(2)} after ${STEPS} steps at niter 0 vs 1`,
499 );
500 }
502 // ---- niter x geometry sweep: catch a "doesn't run" regression early -----
503 // This is what actually turned up the two real issues found while building
504 // the correction: peanut diverging at niter >= 2 with schnak-spots'
505 // shipped default dt (the plain round-sphere preconditioner's convergence
506 // radius -- the mean-J preconditioner has since lifted it; see the control
507 // check after the sweep), and a since-fixed compiler bug where a loop-body
508 // statement could silently reuse a *different* statement's compiled kernel
509 // (test/modelChecks.ts's pipeline-cache check guards that one directly).
510 // Every shipped geometry x every niter the app's <select> actually offers,
511 // so a regression anywhere in that grid is caught.
512 //
513 // SWEEP_LMAX stays at the app's default: the divergence the control check
514 // pins down is lmax-dependent (at lmax 31 or 15 even the plain
515 // preconditioner stays finite on peanut), so a smaller grid would stop
516 // testing the thing the control exists to demonstrate.
517 if (!runSweep) {
518 log(
519 ' sweep: skipped — run `npm run test:node` (desktop Dawn, ~3 s) or ' +
520 '`npm run test:gpu -- --sweep` for the niter x geometry sweep.',
521 );
522 } else {
523 const model = mModelByKey('schnakenberg')!;
524 const params = defaultParams(model);
525 const SWEEP_NITER = [0, 1, 2, 4, 8];
526 // Empty since the symbol-based preconditioner: its high-degree
527 // contraction rate (muMax - muMin)/(muMax + muMin) < 1 on any surface
528 // (mu = the symbol eigenvalues, i.e. inverse squared principal
529 // stretches), where the plain preconditioner diverges wherever mu > 2
530 // -- which is exactly what used to make peanut/2, /4 and /8 diverge.
531 // The mechanism stays: a regression lands here with its evidence.
532 const KNOWN_DIVERGENT = new Set<string>([]);
534 for (const geomSpec of mGeometries) {
535 for (const niter of SWEEP_NITER) {
536 const session = await ModelSession.create({
537 device, model, params, lmax: SWEEP_LMAX,
538 geometry: geomSpec, geometryParams: defaultGeometryParams(geomSpec),
539 niter,
540 });
541 await session.seed(1);
542 session.step(STEPS);
543 const values = await session.read('u');
544 const finite = values.every((v) => Number.isFinite(v));
545 session.destroy();
547 const key = `${geomSpec.key}/${niter}`;
548 const expectDivergent = KNOWN_DIVERGENT.has(key);
549 check(
550 expectDivergent
551 ? `sweep: ${key} is known to diverge (outside the Richardson convergence radius)`
552 : `sweep: ${key} stays finite after ${STEPS} steps`,
553 expectDivergent ? !finite : finite,
554 expectDivergent
555 ? finite
556 ? 'now finite -- the convergence radius may have improved; update KNOWN_DIVERGENT'
557 : 'diverged as expected'
558 : finite
559 ? 'finite'
560 : 'NOT FINITE -- unexpected divergence, investigate before treating this as another known case',
561 );
562 }
563 }
565 // ---- the mean-J control: what the sweep's health is owed to ----------
566 // peanut at niter 4 was the canonical divergent case before the mean-J
567 // preconditioner. Pinning jhat to 1 reproduces the plain round-sphere
568 // preconditioner on today's code, so this asserts both directions at
569 // once: mean-J converges where plain diverges, on the same operator,
570 // same surface, same dt. If this check ever finds jhat = 1 finite, the
571 // sweep above has stopped exercising the regime the preconditioner
572 // exists for (e.g. someone lowered SWEEP_LMAX or dt).
573 {
574 const peanut = mGeometryByKey('peanut')!;
575 const outcomes: boolean[] = [];
576 let jstats = '';
577 for (const jhat of [undefined, 1]) {
578 const session = await ModelSession.create({
579 device, model,
580 params: jhat === undefined ? params : { ...params, jhat },
581 lmax: SWEEP_LMAX,
582 geometry: peanut, geometryParams: defaultGeometryParams(peanut),
583 niter: 4,
584 });
585 if (jhat === undefined) {
586 const g = session.geometry;
587 jstats =
588 `mu in [${g.muMin.toFixed(3)}, ${g.muMax.toFixed(3)}] ` +
589 `(J in [${g.Jmin.toFixed(3)}, ${g.Jmax.toFixed(3)}]), ` +
590 `Jhat ${g.Jhat.toFixed(3)}, ` +
591 `rate ${((g.muMax - g.muMin) / (g.muMax + g.muMin)).toFixed(3)} ` +
592 `vs plain ${(g.muMax - 1).toFixed(2)}`;
593 }
594 await session.seed(1);
595 session.step(STEPS);
596 const values = await session.read('u');
597 outcomes.push(values.every((v) => Number.isFinite(v)));
598 session.destroy();
599 }
600 log(` mean-J on peanut: ${jstats}`);
601 check(
602 'mean-J: converges on peanut/4 where the plain preconditioner diverges',
603 outcomes[0] && !outcomes[1],
604 `mean-J finite: ${outcomes[0]}, jhat=1 finite: ${outcomes[1]}`,
605 );
606 }
607 }
609 // ---- a loop whose length is not known at compile time is refused --------
610 {
611 const model = mModelByKey('allencahn')!;
612 // `dt` is a tunable parameter, so it reaches the compiler with no value:
613 // the plan cannot know how many iterations to emit.
614 const bad = model.source.replace('for k = 1:niter', 'for k = 1:dt');
615 let message = '';
616 try {
617 const session = await ModelSession.create({
618 device, model, params: defaultParams(model), lmax: LMAX, source: bad, niter: 1,
619 });
620 session.destroy();
621 } catch (e) {
622 message = e instanceof ModelCompileError ? e.message : `wrong error type: ${e}`;
623 }
624 check(
625 'loop: a runtime loop bound is refused at compile time',
626 message.includes('known when the model is compiled'),
627 message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
628 );
629 }
631 // ---- swapping the surface leaves the simulation alone ------------------
632 {
633 const model = mModelByKey('schnakenberg')!;
634 const session = await ModelSession.create({
635 device, model, params: defaultParams(model), lmax: LMAX,
636 });
637 await session.seed(1);
638 session.step(STEPS);
639 const before = await session.read('U');
641 const peanut = mGeometryByKey('peanut')!;
642 await session.setGeometry(peanut, defaultGeometryParams(peanut));
643 const after = await session.read('U');
645 let survived = before.length === after.length;
646 for (let i = 0; survived && i < before.length; i++) {
647 if (before[i] !== after[i]) survived = false;
648 }
649 const { lo, hi } = session.geometry.radiusRange();
650 check(
651 'geometry: swapping the surface mid-run does not disturb the state',
652 survived && session.geometryModel.key === 'peanut' && hi - lo > 0.1,
653 survived
654 ? `state identical, now on ${session.geometryModel.key} (radius ${lo.toFixed(3)}–${hi.toFixed(3)})`
655 : 'state changed',
656 );
657 session.destroy();
658 }
660 await randnfun3Checks(device, check, log);
661}
663/**
664 * The seeded initial condition: chebfun's randnfun3, drawn on the host and
665 * summed on the GPU (src/mgpu/randnfun3.ts).
666 *
667 * The split is the thing worth testing. The draw is MATLAB whose distribution
668 * is checked directly, and the sum is a WGSL kernel checked against the same
669 * modes evaluated in f64 on the CPU — if the kernel's indexing into the packed
670 * mode table were wrong it would still produce a smooth random-looking field,
671 * which is exactly the kind of wrong no "looks patterned" check would catch.
672 */
673async function randnfun3Checks(
674 device: GPUDevice,
675 check: Check,
676 log: Log,
677): Promise<void> {
678 const model = mModelByKey('schnakenberg')!;
679 const params = defaultParams(model);
680 const make = (lam3: number): Promise<ModelSession> =>
681 ModelSession.create({ device, model, params, lmax: LMAX, lam3 });
683 // ---- the GPU sum matches the same modes evaluated on the CPU -----------
684 {
685 const session = await make(DEFAULT_LAMBDA);
686 await session.seed(3);
687 // `u` after init is the steady state plus 0.01*f, so the field is
688 // recovered by removing the model's own uniform offset.
689 const u = await session.read('u');
690 const g = session.geometry;
691 const modes = drawModes(
692 DEFAULT_LAMBDA,
693 boundingBox(g.x, g.y, g.z),
694 3,
695 g.x.length,
696 );
697 const nmodes = modes[0];
699 // The same sum in f64, straight from the packed table the GPU read.
700 let maxErr = 0;
701 let amp = 0;
702 const us = params.a + params.b;
703 for (let i = 0; i < g.x.length; i++) {
704 let f = 0;
705 for (let j = 0; j < nmodes; j++) {
706 const b = 4 + 5 * j;
707 const t = modes[b] * g.x[i] + modes[b + 1] * g.y[i] + modes[b + 2] * g.z[i];
708 f += modes[b + 3] * Math.cos(t) - modes[b + 4] * Math.sin(t);
709 }
710 const want = us + 0.01 * f;
711 maxErr = Math.max(maxErr, Math.abs(u[i] - want));
712 amp = Math.max(amp, Math.abs(0.01 * f));
713 }
714 log(` randnfun3: ${nmodes} modes at lambda ${DEFAULT_LAMBDA}, |perturbation| up to ${amp.toExponential(2)}`);
715 check(
716 'randnfun3: the GPU sum matches the same modes summed on the CPU',
717 // fp32 over ~1400 terms against f64. What is being bounded is the
718 // summation floor, and its size is the backend's accumulation order:
719 // Metal lands at 2.0e-6, SwiftShader at 3.8e-6, so an absolute constant
720 // tuned on one is a coin flip on the other. Scale it to the field
721 // instead. The bug this exists to catch -- a mis-indexed read into the
722 // packed table, which would still look like a smooth random field -- is
723 // wrong by O(amp), a thousand times over the bound.
724 maxErr < 1e-3 * amp && amp > 1e-3,
725 `max |GPU - CPU| = ${maxErr.toExponential(2)}, perturbation amplitude ${amp.toExponential(2)}`,
726 );
727 session.destroy();
728 }
730 // ---- a seed reproduces, a different seed does not ----------------------
731 {
732 const a = await make(DEFAULT_LAMBDA);
733 await a.seed(11);
734 const first = await a.read('u');
735 await a.seed(11);
736 const again = await a.read('u');
737 await a.seed(12);
738 const other = await a.read('u');
739 let same = true;
740 let differs = false;
741 for (let i = 0; i < first.length; i++) {
742 if (first[i] !== again[i]) same = false;
743 if (first[i] !== other[i]) differs = true;
744 }
745 check(
746 'randnfun3: the same seed redraws the same field, a different one does not',
747 same && differs,
748 same ? (differs ? 'reproducible and seed-dependent' : 'seed 12 gave seed 11 back') : 'not reproducible',
749 );
750 a.destroy();
751 }
753 // ---- the field is smooth, and lambda sets how smooth -------------------
754 //
755 // This is what randnfun3 buys over the white noise it replaced: the seed is
756 // band-limited, so it is fully resolved by the grid instead of being
757 // whatever the grid happened to alias. Measured as the share of spectral
758 // energy above degree 20 — near zero for a smooth field, and larger for a
759 // shorter wavelength, which is the direction lambda is supposed to move it.
760 {
761 const tail = async (lam3: number): Promise<number> => {
762 const session = await make(lam3);
763 await session.seed(5);
764 const U = await session.read('U');
765 let lo = 0;
766 let hi = 0;
767 for (let m = 0; m <= LMAX; m++) {
768 for (let l = m; l <= LMAX; l++) {
769 const i = lmIndex(LMAX, l, m);
770 const e = U[2 * i] ** 2 + U[2 * i + 1] ** 2;
771 if (l > 20) hi += e;
772 else lo += e;
773 }
774 }
775 session.destroy();
776 return hi / (lo + hi);
777 };
778 const coarse = await tail(1);
779 const fine = await tail(0.4);
780 log(` randnfun3: energy above l=20 is ${coarse.toExponential(2)} at lambda 1, ${fine.toExponential(2)} at lambda 0.4`);
781 check(
782 'randnfun3: the seed is band-limited, and lambda sets its scale',
783 coarse < 1e-3 && fine > coarse,
784 `tail ${coarse.toExponential(2)} (lambda 1) < ${fine.toExponential(2)} (lambda 0.4)`,
785 );
786 }
788 // ---- a finer wavelength grows the table rather than being capped -------
789 //
790 // The mode table is sized to the wavelength asked for, so going finer
791 // reallocates it and rebinds the dispatch. Getting that wrong would leave
792 // the kernel reading a destroyed buffer or a stale one, so check that a
793 // fine field is actually there and actually different.
794 {
795 const session = await make(DEFAULT_LAMBDA);
796 await session.seed(21);
797 const coarse = await session.read('u');
798 session.setLam3(0.12);
799 await session.seed(21);
800 const fine = await session.read('u');
801 let differs = false;
802 let finite = true;
803 for (let i = 0; i < fine.length; i++) {
804 if (!Number.isFinite(fine[i])) finite = false;
805 if (fine[i] !== coarse[i]) differs = true;
806 }
807 check(
808 'randnfun3: a finer wavelength grows the mode table and rebinds',
809 finite && differs,
810 finite ? 'redrew finer, buffer rebound' : 'field went non-finite after resize',
811 );
812 session.destroy();
813 }
815 // ---- a wavelength past the cost budget is refused, not truncated -------
816 {
817 const session = await make(DEFAULT_LAMBDA);
818 let message = '';
819 try {
820 session.setLam3(1e-4);
821 await session.seed(1);
822 } catch (e) {
823 message = e instanceof Error ? e.message : String(e);
824 }
825 check(
826 'randnfun3: a wavelength whose table could not be built is refused',
827 message.includes('Fourier modes on this surface'),
828 message ? `refused: ${message.slice(0, 62)}…` : 'drew it anyway',
829 );
830 session.destroy();
831 }
832}
834/** Index of the entry minimizing `score`, over the first `n` entries. */
835function argMin(
836 xs: Float64Array | Float32Array,
837 n: number,
838 score: (v: number) => number,
839): number {
840 let best = 0;
841 let bestScore = Infinity;
842 for (let i = 0; i < n; i++) {
843 const s = score(xs[i]);
844 if (s < bestScore) {
845 bestScore = s;
846 best = i;
847 }
848 }
849 return best;
850}