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';
24import { gridForLmax, lmIndex, nlmCalc, type ShtConfig } from '../src/sht/layout.ts';
25import { GpuModel } from '../src/mgpu/model.ts';
26import { mModelByKey, defaultParams, type MModel, type ParamSpec } from '../src/mgpu/registry.ts';
27import linearSource from './models/linear.m?raw';
28import logisticSource from './models/logistic.m?raw';
30export type Check = (name: string, ok: boolean, detail: string) => void;
31export type Log = (s: string) => void;
33const param = (key: string, value: number): ParamSpec => ({
34 key, label: key, value, min: -1e9, max: 1e9, step: 1,
35});
37/** A one-species test model with an arbitrary parameter list. */
38const testModel = (key: string, source: string, params: string[]): MModel => ({
39 key,
40 label: key,
41 blurb: '',
42 species: ['u'],
43 state: ['U'],
44 params: params.map((p) => param(p, 0)),
45 pdeg: 1,
46 seedAmp: 1,
47 source,
48});
50async function makeModel(
51 device: GPUDevice,
52 model: MModel,
53 cfg: ShtConfig,
54): Promise<{ sht: ShtPlan; gpu: GpuModel }> {
55 const sht = await ShtPlan.create(device, cfg);
56 const gpu = await GpuModel.create({
57 device,
58 sht,
59 cfg,
60 source: model.source,
61 paramNames: model.params.map((p) => p.key),
62 state: model.state,
63 view: model.species,
64 });
65 return { sht, gpu };
66}
68export async function analyticChecks(
69 device: GPUDevice,
70 check: Check,
71 log: Log,
72): Promise<void> {
73 // ---- A: linear reaction, exact per-mode growth factor -----------------
74 {
75 const lmax = 15;
76 const { nlat, nphi } = gridForLmax(lmax, 1);
77 const cfg = { lmax, mmax: lmax, nlat, nphi };
78 const nlm = nlmCalc(lmax, lmax);
79 const c = -0.3;
80 const D = 0.01;
81 const dt = 0.1;
82 const nsteps = 20;
84 const model = testModel('linear', linearSource, ['c', 'D', 'dt']);
85 const { sht, gpu } = await makeModel(device, model, cfg);
86 gpu.setParams({ c, D, dt });
88 // A single (l, m) mode, written straight into the spectral state.
89 const l = 5;
90 const m = 2;
91 const idx = lmIndex(lmax, l, m);
92 const U0 = new Float32Array(2 * nlm);
93 U0[2 * idx] = 0.8;
94 U0[2 * idx + 1] = -0.35;
95 gpu.upload('U', U0);
97 gpu.step(nsteps);
98 const U = await gpu.read('U');
100 const g = (1 + dt * c) / (1 + dt * D * l * (l + 1));
101 const factor = g ** nsteps;
102 const wantRe = 0.8 * factor;
103 const wantIm = -0.35 * factor;
104 const errRe = Math.abs(U[2 * idx] - wantRe);
105 const errIm = Math.abs(U[2 * idx + 1] - wantIm);
106 check(
107 'A: linear reaction follows the exact per-mode recurrence',
108 errRe < 2e-6 && errIm < 2e-6,
109 `err (${errRe.toExponential(2)}, ${errIm.toExponential(2)}) after ${nsteps} steps`,
110 );
112 // Nothing may leak into the other modes.
113 let leak = 0;
114 for (let i = 0; i < nlm; i++) {
115 if (i === idx) continue;
116 leak = Math.max(leak, Math.abs(U[2 * i]), Math.abs(U[2 * i + 1]));
117 }
118 check('A: no leakage into other modes', leak < 2e-6, `max |other| ${leak.toExponential(2)}`);
120 gpu.destroy();
121 sht.destroy();
122 }
124 // ---- B: nonlinear reaction on a uniform field, exact ODE map ----------
125 {
126 const lmax = 15;
127 const { nlat, nphi } = gridForLmax(lmax, 3);
128 const cfg = { lmax, mmax: lmax, nlat, nphi };
129 const npts = nlat * nphi;
130 const r = 0.7;
131 const D = 0.01;
132 const dt = 0.05;
133 const nsteps = 25;
134 const u0 = 0.3;
136 const model = testModel('logistic', logisticSource, ['r', 'D', 'dt']);
137 const { sht, gpu } = await makeModel(device, model, cfg);
138 gpu.setParams({ r, D, dt });
140 // Uniform initial field: stays uniform, and diffusion cannot touch it.
141 const field = new Float32Array(npts).fill(u0);
142 gpu.init(field);
143 const Ustart = await gpu.read('U');
144 gpu.step(nsteps);
145 const Uend = await gpu.read('U');
147 // Read the *state*, not the `u` output: a model computes its grid fields
148 // from the state at the START of the step (`u = synth(U)` precedes the
149 // update), so the rendered field lags the state by one step. The l=0
150 // coefficient of a uniform field scales linearly with its value, so the
151 // ratio gives the value back without needing Y_00's normalization.
152 const got = u0 * (Uend[0] / Ustart[0]);
154 let want = u0;
155 for (let s = 0; s < nsteps; s++) want += dt * r * want * (1 - want);
157 const err = Math.abs(got - want);
158 check(
159 'B: uniform nonlinear reaction follows the scalar ODE map',
160 err < 5e-6,
161 `${got.toFixed(7)} vs ${want.toFixed(7)}, err ${err.toExponential(2)}`,
162 );
164 // And it must still be uniform: any structure would mean the kernel is
165 // reading the wrong elements.
166 const u = await gpu.read('u');
167 let lo = Infinity;
168 let hi = -Infinity;
169 for (const v of u) {
170 if (v < lo) lo = v;
171 if (v > hi) hi = v;
172 }
173 check(
174 'B: the field stays uniform',
175 hi - lo < 1e-6,
176 `spread ${(hi - lo).toExponential(2)}`,
177 );
179 gpu.destroy();
180 sht.destroy();
181 }
183 // ---- C: linearized Turing recurrence on a real two-species model ------
184 {
185 const model = mModelByKey('schnakenberg')!;
186 const p = defaultParams(model);
187 const lmax = 31;
188 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
189 const cfg = { lmax, mmax: lmax, nlat, nphi };
190 const nlm = nlmCalc(lmax, lmax);
191 const npts = nlat * nphi;
193 const { sht, gpu } = await makeModel(device, model, cfg);
194 gpu.setParams(p);
196 // Seed the exact homogeneous fixed point by handing init a zero
197 // perturbation, then add a small single-mode bump to u only.
198 gpu.init(new Float32Array(npts));
199 const l = 24;
200 const m = 7;
201 const idx = lmIndex(lmax, l, m);
202 const eps = 1e-6;
203 const U0 = await gpu.read('U');
204 const V0 = await gpu.read('V');
205 const Upert = Float32Array.from(U0);
206 Upert[2 * idx] += eps;
207 gpu.upload('U', Upert);
208 gpu.upload('V', V0);
210 const nsteps = 40;
211 gpu.step(nsteps);
212 const U = await gpu.read('U');
213 const V = await gpu.read('V');
215 // Jacobian of (a - u + u^2 v, b - u^2 v) at the fixed point us = a+b,
216 // vs = b/us^2, with diffusion applied implicitly per species.
217 const us = p.a + p.b;
218 const vs = p.b / (us * us);
219 const J = [
220 [-1 + 2 * us * vs, us * us],
221 [-2 * us * vs, -us * us],
222 ];
223 const lam = l * (l + 1);
224 const du = 1 / (1 + p.dt * p.D1 * lam);
225 const dv = 1 / (1 + p.dt * p.D2 * lam);
226 let cu = eps;
227 let cv = 0;
228 for (let s = 0; s < nsteps; s++) {
229 const nu = (cu + p.dt * (J[0][0] * cu + J[0][1] * cv)) * du;
230 const nv = (cv + p.dt * (J[1][0] * cu + J[1][1] * cv)) * dv;
231 cu = nu;
232 cv = nv;
233 }
235 const gotU = U[2 * idx] - U0[2 * idx];
236 const gotV = V[2 * idx] - V0[2 * idx];
237 const relU = Math.abs(gotU - cu) / Math.max(Math.abs(cu), 1e-30);
238 const relV = Math.abs(gotV - cv) / Math.max(Math.abs(cv), 1e-30);
239 // Looser than A and B by design: a 1e-6 perturbation sits on a state of
240 // order 1, so fp32 keeps only ~4 significant digits of it.
241 check(
242 'C: perturbation follows the linearized 2x2 IMEX recurrence',
243 relU < 5e-3 && relV < 5e-3,
244 `rel err (${relU.toExponential(2)}, ${relV.toExponential(2)})`,
245 );
246 check(
247 `C: the (l=${l}, m=${m}) mode is unstable`,
248 Math.abs(cu) > eps && Math.abs(gotU) > eps,
249 `|c_u| ${eps.toExponential(2)} -> ${Math.abs(gotU).toExponential(2)}`,
250 );
252 log(` C: growth over ${nsteps} steps = ${(Math.abs(cu) / eps).toFixed(3)}x (predicted)`);
254 gpu.destroy();
255 sht.destroy();
256 }
257}