2 * Correctness of the .m -> WGSL path, against closed-form answers.
3 *
4 * These replace what used to be a comparison against a second TypeScript
5 * implementation of the same scheme. Checking against arithmetic is stronger:
6 * two implementations agreeing only shows they share assumptions, whereas an
7 * exact recurrence pins the result. Each test picks a case whose evolution is
8 * known in closed form, runs it through the real pipeline — MATLAB source,
9 * numbl lowering, generated WGSL, GPU transforms — and compares.
10 *
11 * A: a linear reaction makes every spherical-harmonic mode independent, with a
12 * known growth factor per degree. Checks the transform round-trip, the
13 * eigenvalue mapping, the IMEX update and the state feedback.
14 * B: a nonlinear reaction on a uniform field follows the scalar ODE map
15 * exactly. Checks that a generated kernel evaluates a nonlinear reaction.
16 * C: a small perturbation of the Schnakenberg fixed point follows the
17 * linearized 2x2 IMEX recurrence, and the expected mode is unstable. Checks
18 * a real two-species model.
19 *
20 * Everything runs in fp32 on the GPU, so tolerances are set by fp32 round-off
21 * (~1e-7 relative) rather than by the scheme.
22 */
23import { ShtPlan } from '../src/sht/sht.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 24import { DerivPlan } from '../src/sht/deriv.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 25import { gridForLmax, lmIndex, nlmCalc, type ShtConfig } from '../src/sht/layout.ts';
26import { GpuModel } from '../src/mgpu/model.ts';
27import { mModelByKey, defaultParams, type MModel, type ParamSpec } from '../src/mgpu/registry.ts';
28import { Geometry } from '../src/geom/geometry.ts';
29import { mGeometryByKey, SPHERE_KEY } from '../src/geom/registry.ts';
30import linearSource from './models/linear.m?raw';
31import logisticSource from './models/logistic.m?raw';
33export type Check = (name: string, ok: boolean, detail: string) => void;
34export type Log = (s: string) => void;
36const param = (key: string, value: number): ParamSpec => ({
37 key, label: key, value, min: -1e9, max: 1e9, step: 1,
38});
40/** A one-species test model with an arbitrary parameter list. */
41const testModel = (key: string, source: string, params: string[]): MModel => ({
42 key,
43 label: key,
44 blurb: '',
45 species: ['u'],
46 state: ['U'],
47 params: params.map((p) => param(p, 0)),
48 pdeg: 1,
49 seedAmp: 1,
50 source,
51});
53/**
54 * Every closed-form case here is a statement about the *round sphere*, so
55 * every model here is built on the sphere geometry. The models that take a
56 * surface still get one — a geometry is always supplied, and for the sphere it
57 * is the degree-1 embedding, which is what makes these answers exact.
58 */
59async function makeModel(
60 device: GPUDevice,
61 model: MModel,
62 cfg: ShtConfig,
63 niter = 1,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 64): Promise<{ sht: ShtPlan; deriv: DerivPlan; gpu: GpuModel }> {
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 65 const sht = await ShtPlan.create(device, cfg);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 66 const deriv = await DerivPlan.create(device, sht);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 67 const geometry = await Geometry.create({
68 sht,
69 cfg,
70 source: mGeometryByKey(SPHERE_KEY)!.source,
71 paramNames: [],
72 params: {},
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 73 deriv,
75 const gpu = await GpuModel.create({
76 device,
77 sht,
78 cfg,
79 source: model.source,
80 paramNames: model.params.map((p) => p.key),
81 state: model.state,
82 view: model.species,
83 geometry,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 84 deriv,
86 });
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 87 return { sht, deriv, gpu };
90export async function analyticChecks(
91 device: GPUDevice,
92 check: Check,
93 log: Log,
94): Promise<void> {
95 // ---- A: linear reaction, exact per-mode growth factor -----------------
96 {
97 const lmax = 15;
98 const { nlat, nphi } = gridForLmax(lmax, 1);
99 const cfg = { lmax, mmax: lmax, nlat, nphi };
100 const nlm = nlmCalc(lmax, lmax);
101 const c = -0.3;
102 const D = 0.01;
103 const dt = 0.1;
104 const nsteps = 20;
106 const model = testModel('linear', linearSource, ['c', 'D', 'dt']);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 107 const { sht, deriv, gpu } = await makeModel(device, model, cfg);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 108 gpu.setParams({ c, D, dt });
110 // A single (l, m) mode, written straight into the spectral state.
111 const l = 5;
112 const m = 2;
113 const idx = lmIndex(lmax, l, m);
114 const U0 = new Float32Array(2 * nlm);
115 U0[2 * idx] = 0.8;
116 U0[2 * idx + 1] = -0.35;
117 gpu.upload('U', U0);
119 gpu.step(nsteps);
120 const U = await gpu.read('U');
122 const g = (1 + dt * c) / (1 + dt * D * l * (l + 1));
123 const factor = g ** nsteps;
124 const wantRe = 0.8 * factor;
125 const wantIm = -0.35 * factor;
126 const errRe = Math.abs(U[2 * idx] - wantRe);
127 const errIm = Math.abs(U[2 * idx + 1] - wantIm);
128 check(
129 'A: linear reaction follows the exact per-mode recurrence',
130 errRe < 2e-6 && errIm < 2e-6,
131 `err (${errRe.toExponential(2)}, ${errIm.toExponential(2)}) after ${nsteps} steps`,
132 );
134 // Nothing may leak into the other modes.
135 let leak = 0;
136 for (let i = 0; i < nlm; i++) {
137 if (i === idx) continue;
138 leak = Math.max(leak, Math.abs(U[2 * i]), Math.abs(U[2 * i + 1]));
139 }
140 check('A: no leakage into other modes', leak < 2e-6, `max |other| ${leak.toExponential(2)}`);
142 gpu.destroy();
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 143 deriv.destroy();
145 }
147 // ---- B: nonlinear reaction on a uniform field, exact ODE map ----------
148 {
149 const lmax = 15;
150 const { nlat, nphi } = gridForLmax(lmax, 3);
151 const cfg = { lmax, mmax: lmax, nlat, nphi };
152 const npts = nlat * nphi;
153 const r = 0.7;
154 const D = 0.01;
155 const dt = 0.05;
156 const nsteps = 25;
157 const u0 = 0.3;
159 const model = testModel('logistic', logisticSource, ['r', 'D', 'dt']);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 160 const { sht, deriv, gpu } = await makeModel(device, model, cfg);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 161 gpu.setParams({ r, D, dt });
163 // Uniform initial field: stays uniform, and diffusion cannot touch it.
164 const field = new Float32Array(npts).fill(u0);
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 165 await gpu.init(field, null);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 166 const Ustart = await gpu.read('U');
167 gpu.step(nsteps);
168 const Uend = await gpu.read('U');
170 // Read the *state*, not the `u` output: a model computes its grid fields
171 // from the state at the START of the step (`u = synth(U)` precedes the
172 // update), so the rendered field lags the state by one step. The l=0
173 // coefficient of a uniform field scales linearly with its value, so the
174 // ratio gives the value back without needing Y_00's normalization.
175 const got = u0 * (Uend[0] / Ustart[0]);
177 let want = u0;
178 for (let s = 0; s < nsteps; s++) want += dt * r * want * (1 - want);
180 const err = Math.abs(got - want);
181 check(
182 'B: uniform nonlinear reaction follows the scalar ODE map',
183 err < 5e-6,
184 `${got.toFixed(7)} vs ${want.toFixed(7)}, err ${err.toExponential(2)}`,
185 );
187 // And it must still be uniform: any structure would mean the kernel is
188 // reading the wrong elements.
189 const u = await gpu.read('u');
190 let lo = Infinity;
191 let hi = -Infinity;
192 for (const v of u) {
193 if (v < lo) lo = v;
194 if (v > hi) hi = v;
195 }
196 check(
197 'B: the field stays uniform',
198 hi - lo < 1e-6,
199 `spread ${(hi - lo).toExponential(2)}`,
200 );
202 gpu.destroy();
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 203 deriv.destroy();
205 }
207 // ---- C: linearized Turing recurrence on a real two-species model ------
208 {
209 const model = mModelByKey('schnakenberg')!;
210 const p = defaultParams(model);
211 const lmax = 31;
212 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
213 const cfg = { lmax, mmax: lmax, nlat, nphi };
214 const nlm = nlmCalc(lmax, lmax);
215 const npts = nlat * nphi;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 217 const { sht, deriv, gpu } = await makeModel(device, model, cfg);
220 // Seed the exact homogeneous fixed point by handing init a zero
221 // perturbation, then add a small single-mode bump to u only.
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 222 await gpu.init(new Float32Array(npts), null);
224 const m = 7;
225 const idx = lmIndex(lmax, l, m);
226 const eps = 1e-6;
227 const U0 = await gpu.read('U');
228 const V0 = await gpu.read('V');
229 const Upert = Float32Array.from(U0);
230 Upert[2 * idx] += eps;
231 gpu.upload('U', Upert);
232 gpu.upload('V', V0);
234 const nsteps = 40;
235 gpu.step(nsteps);
236 const U = await gpu.read('U');
237 const V = await gpu.read('V');
239 // Jacobian of (a - u + u^2 v, b - u^2 v) at the fixed point us = a+b,
240 // vs = b/us^2, with diffusion applied implicitly per species.
241 const us = p.a + p.b;
242 const vs = p.b / (us * us);
243 const J = [
244 [-1 + 2 * us * vs, us * us],
245 [-2 * us * vs, -us * us],
246 ];
247 const lam = l * (l + 1);
248 const du = 1 / (1 + p.dt * p.D1 * lam);
249 const dv = 1 / (1 + p.dt * p.D2 * lam);
250 let cu = eps;
251 let cv = 0;
252 for (let s = 0; s < nsteps; s++) {
253 const nu = (cu + p.dt * (J[0][0] * cu + J[0][1] * cv)) * du;
254 const nv = (cv + p.dt * (J[1][0] * cu + J[1][1] * cv)) * dv;
255 cu = nu;
256 cv = nv;
257 }
259 const gotU = U[2 * idx] - U0[2 * idx];
260 const gotV = V[2 * idx] - V0[2 * idx];
261 const relU = Math.abs(gotU - cu) / Math.max(Math.abs(cu), 1e-30);
262 const relV = Math.abs(gotV - cv) / Math.max(Math.abs(cv), 1e-30);
263 // Looser than A and B by design: a 1e-6 perturbation sits on a state of
264 // order 1, so fp32 keeps only ~4 significant digits of it.
265 check(
266 'C: perturbation follows the linearized 2x2 IMEX recurrence',
267 relU < 5e-3 && relV < 5e-3,
268 `rel err (${relU.toExponential(2)}, ${relV.toExponential(2)})`,
269 );
270 check(
271 `C: the (l=${l}, m=${m}) mode is unstable`,
272 Math.abs(cu) > eps && Math.abs(gotU) > eps,
273 `|c_u| ${eps.toExponential(2)} -> ${Math.abs(gotU).toExponential(2)}`,
274 );
276 log(` C: growth over ${nsteps} steps = ${(Math.abs(cu) / eps).toFixed(3)}x (predicted)`);
278 gpu.destroy();
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 279 deriv.destroy();
281 }
282}