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