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
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 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.
27import { ShtPlan } from '../src/sht/sht.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 28import { DerivPlan } from '../src/sht/deriv.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 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;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 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;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 50
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);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 57 const deriv = await DerivPlan.create(device, sht);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 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),
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 65 deriv,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 67 return { g, sht, deriv, cfg, geometry };
70export async function geometryChecks(
71 device: GPUDevice,
72 check: Check,
73 log: Log,
74): Promise<void> {
75 // ---- every geometry compiles and closes ---------------------------------
76 for (const spec of mGeometries) {
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 77 const { sht, deriv, geometry } = await buildGeometry(device, spec.key);
79 for (const a of [geometry.x, geometry.y, geometry.z]) {
80 for (const v of a) if (!Number.isFinite(v)) finite = false;
81 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 82 for (const a of [geometry.Vtx, geometry.Vty, geometry.Vtz, geometry.Vpx, geometry.Vpy, geometry.Vpz]) {
83 for (const v of a) if (!Number.isFinite(v)) finite = false;
84 }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 85 const { lo, hi } = geometry.radiusRange();
86 check(
87 `geometry: ${spec.key}.m evaluates to a finite surface`,
88 finite && lo > 1e-3,
89 `radius ${lo.toFixed(4)}–${hi.toFixed(4)}`,
90 );
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 91 deriv.destroy();
93 }
95 // ---- the sphere is the unit sphere, exactly, and is degree 1 ------------
96 {
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 97 const { sht, deriv, geometry } = await buildGeometry(device, SPHERE_KEY);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 98
99 let maxRadiusErr = 0;
100 for (let i = 0; i < geometry.x.length; i++) {
101 const r = Math.hypot(geometry.x[i], geometry.y[i], geometry.z[i]);
102 maxRadiusErr = Math.max(maxRadiusErr, Math.abs(r - 1));
103 }
104 // Tolerance is fp32 through a full analysis/synthesis round trip, not the
105 // geometry: the exact answer is representable, and what is measured here
106 // is the transforms' own round-off. It is set by the loosest stack this
107 // runs on — SwiftShader in CI is an order of magnitude worse than Dawn on
108 // real hardware (4e-4 against 2e-5). A geometry that was actually wrong
109 // would miss by O(1), so the slack costs nothing.
110 check(
111 'geometry: sphere.m has radius 1 everywhere',
112 maxRadiusErr < 2e-3,
113 `max |r - 1| = ${maxRadiusErr.toExponential(2)}`,
114 );
116 // x, y, z of the unit sphere are the three degree-1 harmonics and nothing
117 // else, so analysing them must leave every other coefficient at zero.
118 // This is what makes the sphere case exact rather than merely accurate:
119 // there is no content for the band limit to throw away.
120 const degreeOne = new Set([
121 lmIndex(LMAX, 1, 0),
122 lmIndex(LMAX, 1, 1),
123 ]);
124 let leak = 0;
125 for (const coeffs of [geometry.X, geometry.Y, geometry.Z]) {
126 for (let i = 0; i < coeffs.length / 2; i++) {
127 if (degreeOne.has(i)) continue;
128 leak = Math.max(leak, Math.abs(coeffs[2 * i]), Math.abs(coeffs[2 * i + 1]));
129 }
130 }
131 check(
132 'geometry: sphere.m is exactly degree 1 in the harmonics',
133 leak < 1e-3,
134 `max |coefficient| outside l = 1 is ${leak.toExponential(2)}`,
135 );
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 136
137 // The inverse metric quantities have a closed form on the unit sphere:
138 // V_theta = (cos(theta)cos(phi), cos(theta)sin(phi), -sin(theta)),
139 // V_phi = (-sin(phi)/sin(theta), cos(phi)/sin(theta), 0). Checking these
140 // pins the sign convention of computeMetric (src/geom/metric.ts) before
141 // it is buried under the Laplace-Beltrami operator built on top of it.
142 let maxMetricErr = 0;
143 for (let i = 0; i < sht.cfg.nlat; i++) {
144 const ct = sht.cosTheta[i];
145 const st = Math.sqrt(Math.max(0, 1 - ct * ct));
146 for (let j = 0; j < sht.cfg.nphi; j++) {
147 const phi = (2 * Math.PI * j) / sht.cfg.nphi;
148 const k = i * sht.cfg.nphi + j;
149 const cphi = Math.cos(phi);
150 const sphi = Math.sin(phi);
151 const wantVtx = ct * cphi;
152 const wantVty = ct * sphi;
153 const wantVtz = -st;
154 const wantVpx = -sphi / st;
155 const wantVpy = cphi / st;
156 const wantVpz = 0;
157 maxMetricErr = Math.max(
158 maxMetricErr,
159 Math.abs(geometry.Vtx[k] - wantVtx),
160 Math.abs(geometry.Vty[k] - wantVty),
161 Math.abs(geometry.Vtz[k] - wantVtz),
162 Math.abs(geometry.Vpx[k] - wantVpx),
163 Math.abs(geometry.Vpy[k] - wantVpy),
164 Math.abs(geometry.Vpz[k] - wantVpz),
165 );
166 }
167 }
168 check(
169 'geometry: sphere.m has the closed-form inverse metric quantities',
170 maxMetricErr < 2e-3,
171 `max |V - closed form| = ${maxMetricErr.toExponential(2)}`,
172 );
174 deriv.destroy();
176 }
178 // ---- a deformed surface matches its own formula, on any grid ------------
179 {
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 180 const { g, sht, deriv, cfg, geometry } = await buildGeometry(device, 'peanut');
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 181 const p = defaultGeometryParams(g);
183 // peanut.m written out: r = 1 - waist*sin(theta)^2 scales the unit sphere,
184 // and z is then stretched, so the distance from the origin depends on
185 // theta alone. Checking every point against this closed form checks the
186 // whole path at once — the compiled shape kernel, the analysis into
187 // coefficients, the synthesis back — and, because the formula has no phi
188 // in it, that the surface really is a surface of revolution.
189 const peanutRadius = (ct: number): number => {
190 const st2 = Math.max(0, 1 - ct * ct);
191 const r = 1 - p.waist * st2;
192 return r * Math.hypot(Math.sqrt(st2), (1 + p.stretch) * ct);
193 };
195 const onGrid = (
196 cosTheta: Float64Array,
197 nlat: number,
198 nphi: number,
199 at: (i: number) => number,
200 ): number => {
201 let worst = 0;
202 for (let i = 0; i < nlat; i++) {
203 const want = peanutRadius(cosTheta[i]);
204 for (let j = 0; j < nphi; j++) {
205 worst = Math.max(worst, Math.abs(at(i * nphi + j) - want));
206 }
207 }
208 return worst;
209 };
211 const coarse = onGrid(sht.cosTheta, cfg.nlat, cfg.nphi, (k) =>
212 Math.hypot(geometry.x[k], geometry.y[k], geometry.z[k]),
213 );
214 check(
215 'geometry: peanut.m matches its own radial formula on the solver grid',
216 coarse < 1e-3,
217 `max |dr| = ${coarse.toExponential(2)}`,
218 );
220 // And the same on a finer grid, from the same coefficients. This is what
221 // "the rendered surface is the surface being solved on" means: display
222 // oversampling evaluates the embedding at more points, it does not
223 // subdivide or smooth it. The 2x Gauss latitudes share no point with the
224 // 1x ones, so agreeing here is agreeing everywhere, not at samples.
225 const fine = await ShtPlan.create(device, {
226 lmax: cfg.lmax,
227 mmax: cfg.mmax,
228 nlat: 2 * cfg.nlat,
229 nphi: 2 * cfg.nphi,
230 });
231 const finePos = await geometry.positionsOn(fine);
232 const refined = onGrid(fine.cosTheta, 2 * cfg.nlat, 2 * cfg.nphi, (k) =>
233 Math.hypot(finePos[3 * k], finePos[3 * k + 1], finePos[3 * k + 2]),
234 );
235 check(
236 'geometry: the same coefficients give the same surface on a 2x grid',
237 refined < 1e-3,
238 `max |dr| = ${refined.toExponential(2)} at ${2 * cfg.nlat}×${2 * cfg.nphi} points`,
239 );
240 fine.destroy();
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 241 deriv.destroy();
243 }
245 // ---- the unrolled loop: more ops, identical answer ----------------------
246 {
247 const model = mModelByKey('schnakenberg')!;
248 const params = defaultParams(model);
249 const counts = [0, 1, 4];
250 const ops: number[] = [];
251 const states: Float32Array[] = [];
253 for (const niter of counts) {
254 const session = await ModelSession.create({
255 device, model, params, lmax: LMAX, niter,
256 });
257 ops.push(session.describe().step.length);
258 session.seed(1);
259 session.step(STEPS);
260 states.push(await session.read('U'));
261 session.destroy();
262 }
264 log(` schnakenberg.m ops/step by solve iterations: ${
265 counts.map((n, i) => `${n} -> ${ops[i]}`).join(', ')
266 }`);
267 check(
268 'loop: each solve iteration adds GPU operations',
269 ops[0] < ops[1] && ops[1] < ops[2],
270 `${ops.join(' < ')} ops for ${counts.join(', ')} iterations`,
271 );
272 // Unrolling has to be exactly linear in the trip count: the body planned
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 273 // once per iteration, no more and no less. Per species per iteration: 8
274 // dtheta/dphi + 4 analys transforms (Algorithm 3's cost, applied to the
275 // field and to each of its three Cartesian gradient components) plus 15
276 // generated kernels -- see test/modelChecks.ts's KERNELS_PER_ITERATION,
277 // which counts the kernels alone; this counts every op, transforms
278 // included.
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 279 const perIteration = ops[1] - ops[0];
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 280 const want = 54;
282 'loop: unrolling is exactly linear in the trip count',
283 perIteration === want && ops[2] - ops[0] === 4 * perIteration,
284 `${perIteration} ops per iteration (expected ${want}), ` +
285 `${ops[2] - ops[0]} for 4 iterations`,
286 );
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 288 // On the sphere lap_g = lap_s exactly, so the correction should compute
289 // (numerically) close to zero regardless of niter -- not bit-identical
290 // (it is a real computation now, through 8+ chained fp32 transforms per
291 // iteration, not the literal `0 * Un` placeholder that used to make this
292 // exact), but close. The tolerance is set by that chain's fp32 roundoff,
293 // not by the scheme: a real geometry-correction bug would miss by orders
294 // of magnitude more than this.
296 for (let k = 1; k < states.length; k++) {
297 for (let i = 0; i < states[0].length; i++) {
298 worst = Math.max(worst, Math.abs(states[k][i] - states[0][i]));
299 }
300 }
301 check(
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 302 'loop: on the sphere, the correction stays near zero across niter',
303 worst < 2e-3,
304 `states differ by up to ${worst.toExponential(2)} after ${STEPS} steps at ${counts.join('/')} iterations`,
305 );
306 }
308 // ---- on a curved surface, the correction actually changes the answer ----
309 {
310 const model = mModelByKey('schnakenberg')!;
311 const params = defaultParams(model);
312 const peanut = mGeometryByKey('peanut')!;
313 const peanutParams = defaultGeometryParams(peanut);
314 // niter 0 vs 1 only -- deliberately not the 4/8 the sweep below already
315 // documents as outside the Richardson iteration's convergence radius on
316 // this geometry. The point here is just that the correction is not a
317 // no-op, which a much smaller, still-converging niter already shows.
318 const states: Float32Array[] = [];
319 for (const niter of [0, 1]) {
320 const session = await ModelSession.create({
321 device, model, params, lmax: SWEEP_LMAX,
322 geometry: peanut, geometryParams: peanutParams, niter,
323 });
324 session.seed(1);
325 session.step(STEPS);
326 states.push(await session.read('U'));
327 session.destroy();
328 }
329 let worst = 0;
330 for (let i = 0; i < states[0].length; i++) {
331 worst = Math.max(worst, Math.abs(states[1][i] - states[0][i]));
332 }
333 check(
334 'loop: on peanut, the correction measurably changes the answer',
335 worst > 1e-4 && states[1].every((v) => Number.isFinite(v)),
336 `states differ by ${worst.toExponential(2)} after ${STEPS} steps at niter 0 vs 1`,
338 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 340 // ---- niter x geometry sweep: catch a "doesn't run" regression early -----
341 // This is what actually turned up the two real issues found while building
342 // the correction: peanut diverging at niter >= 4 with schnak-spots'
343 // shipped default dt (a genuine Richardson-convergence-radius limit, not a
344 // bug -- see docs/richardson-iteration.md), and a since-fixed compiler bug
345 // where a loop-body statement could silently reuse a *different*
346 // statement's compiled kernel (test/modelChecks.ts's pipeline-cache check
347 // guards that one directly). Every shipped geometry x every niter the
348 // app's <select> actually offers, so a regression anywhere in that grid is
349 // caught -- without asserting away the one combination already known to be
350 // outside the convergence radius.
351 {
352 const model = mModelByKey('schnakenberg')!;
353 const params = defaultParams(model);
354 const SWEEP_NITER = [0, 1, 2, 4, 8];
355 const KNOWN_DIVERGENT = new Set(['peanut/2', 'peanut/4', 'peanut/8']);
357 for (const geomSpec of mGeometries) {
358 for (const niter of SWEEP_NITER) {
359 const session = await ModelSession.create({
360 device, model, params, lmax: SWEEP_LMAX,
361 geometry: geomSpec, geometryParams: defaultGeometryParams(geomSpec),
362 niter,
363 });
364 session.seed(1);
365 session.step(STEPS);
366 const values = await session.read('u');
367 const finite = values.every((v) => Number.isFinite(v));
368 session.destroy();
370 const key = `${geomSpec.key}/${niter}`;
371 const expectDivergent = KNOWN_DIVERGENT.has(key);
372 check(
373 expectDivergent
374 ? `sweep: ${key} is known to diverge (outside the Richardson convergence radius)`
375 : `sweep: ${key} stays finite after ${STEPS} steps`,
376 expectDivergent ? !finite : finite,
377 expectDivergent
378 ? finite
379 ? 'now finite -- the convergence radius may have improved; update KNOWN_DIVERGENT'
380 : 'diverged as expected'
381 : finite
382 ? 'finite'
383 : 'NOT FINITE -- unexpected divergence, investigate before treating this as another known case',
384 );
385 }
386 }
387 }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 389 // ---- a loop whose length is not known at compile time is refused --------
390 {
391 const model = mModelByKey('allencahn')!;
392 // `dt` is a tunable parameter, so it reaches the compiler with no value:
393 // the plan cannot know how many iterations to emit.
394 const bad = model.source.replace('for k = 1:niter', 'for k = 1:dt');
395 let message = '';
396 try {
397 const session = await ModelSession.create({
398 device, model, params: defaultParams(model), lmax: LMAX, source: bad, niter: 1,
399 });
400 session.destroy();
401 } catch (e) {
402 message = e instanceof ModelCompileError ? e.message : `wrong error type: ${e}`;
403 }
404 check(
405 'loop: a runtime loop bound is refused at compile time',
406 message.includes('known when the model is compiled'),
407 message ? `refused: ${message.slice(0, 72)}…` : 'compiled anyway',
408 );
409 }
411 // ---- swapping the surface leaves the simulation alone ------------------
412 {
413 const model = mModelByKey('schnakenberg')!;
414 const session = await ModelSession.create({
415 device, model, params: defaultParams(model), lmax: LMAX,
416 });
417 session.seed(1);
418 session.step(STEPS);
419 const before = await session.read('U');
421 const peanut = mGeometryByKey('peanut')!;
422 await session.setGeometry(peanut, defaultGeometryParams(peanut));
423 const after = await session.read('U');
425 let survived = before.length === after.length;
426 for (let i = 0; survived && i < before.length; i++) {
427 if (before[i] !== after[i]) survived = false;
428 }
429 const { lo, hi } = session.geometry.radiusRange();
430 check(
431 'geometry: swapping the surface mid-run does not disturb the state',
432 survived && session.geometryModel.key === 'peanut' && hi - lo > 0.1,
433 survived
434 ? `state identical, now on ${session.geometryModel.key} (radius ${lo.toFixed(3)}–${hi.toFixed(3)})`
435 : 'state changed',
436 );
437 session.destroy();
438 }
439}
441/** Index of the entry minimizing `score`, over the first `n` entries. */
442function argMin(
443 xs: Float64Array | Float32Array,
444 n: number,
445 score: (v: number) => number,
446): number {
447 let best = 0;
448 let bestScore = Infinity;
449 for (let i = 0; i < n; i++) {
450 const s = score(xs[i]);
451 if (s < bestScore) {
452 bestScore = s;
453 best = i;
454 }
455 }
456 return best;
457}