2 * Where does leg_synth's recurrence go wrong?
3 *
4 * npx vite-node scripts/diagnose-leg.ts [--lmax 63] [--m 0]
5 *
6 * Follow-up to scripts/diagnose-sht.ts, which narrows a bad transform down to
7 * one shader. This reads that shader's recurrence out term by term.
8 *
9 * The trick is to probe the production shader rather than a copy of it: with
10 * qlm set to a single 1 at coefficient (l0, m) and zero everywhere else,
11 *
12 * fm[m][ilat] = sum_l Q_lm ytilde_l^m(theta_i) = ytilde_l0^m(theta_i)
13 *
14 * so one synthesis per l0 hands back exactly the recurrence value at that l, for
15 * every latitude at once, computed by the same code the app runs. Sweeping l0
16 * from m to lmax gives the whole sequence, and comparing with legendreRow (f64)
17 * says which term first disagrees:
18 *
19 * - wrong at l = m -> the seed (amm, or sinpow_rescaled)
20 * - wrong at l = m+1 -> a_{m+1}^m, i.e. the ab buffer as the shader reads it
21 * - right until some l, then -> the two-at-a-time advance in the loop
22 * growing steadily
23 * - a constant wrong factor -> a scale error, not an instability
24 */
25import { ShtPlan, requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
26import { ShtReference } from '../src/sht/reference.ts';
27import { legendreRow } from '../src/sht/coeffs.ts';
28import { gridForLmax, lmIndex, type ShtConfig } from '../src/sht/layout.ts';
29import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
31const argv = process.argv.slice(2);
32if (argv.includes('--help') || argv.includes('-h')) {
33 console.log(`usage: npx vite-node scripts/diagnose-leg.ts [options]
35 --lmax <n> spherical harmonic truncation (default 63, the app's)
36 --m <n> the order to follow (default 0, which needs no rescaling at all
37 and so isolates the plain recurrence)
38 --lats <i,j> latitudes to sample (default 0,1,mid,last)
39 --all print every l, not just the interesting ones
40 --help`);
41 process.exit(0);
42}
43const flag = (name: string, dflt: string): string => {
44 const i = argv.indexOf(`--${name}`);
45 if (i >= 0 && argv[i + 1] !== undefined) return argv[i + 1];
46 const eq = argv.find((a) => a.startsWith(`--${name}=`));
47 return eq ? eq.slice(name.length + 3) : dflt;
48};
49const lmax = Number(flag('lmax', '63'));
50const m = Number(flag('m', '0'));
51const showAll = argv.includes('--all');
53if (!Number.isInteger(m) || m < 0 || m > lmax) {
54 console.error(`diagnose-leg: --m must be an integer in [0, lmax]`);
55 process.exit(2);
56}
58let device: GPUDevice | null = null;
59let plan: ShtPlan | null = null;
60try {
61 const runtime = await installWebGpu();
62 device = await requestShtDevice().catch((e: unknown) => {
63 throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
64 });
65 const adapter = await describeAdapter(device);
67 const { nlat, nphi } = gridForLmax(lmax, 3);
68 const cfg: ShtConfig = { lmax, mmax: lmax, nlat, nphi };
69 const ref = new ShtReference(cfg);
70 // The Fourier stage plays no part here — only fm is read.
71 plan = await ShtPlan.create(device, cfg, { fourier: 'dft' });
73 const lats = flag('lats', '')
74 ? flag('lats', '').split(',').map(Number)
75 : [0, 1, nlat >> 1, nlat - 1];
77 console.log('turing-surface — following leg_synth\'s recurrence term by term\n');
78 console.log(` device ${adapter || '(unknown)'}\n ${runtime}`);
79 console.log(` grid lmax ${lmax} · ${nlat}×${nphi}`);
80 console.log(` order m = ${m}${m === 0 ? ' (no rescaling: sinpow_rescaled returns 1, ny = 0)' : ''}`);
81 console.log(
82 ` latitudes ${lats
83 .map((i) => `${i} (theta ${((Math.acos(ref.ct[i]) * 180) / Math.PI).toFixed(1)}°)`)
84 .join(', ')}\n`,
85 );
87 const fmBytes = 8 * (cfg.mmax + 1) * nlat;
88 const stageFm = device.createBuffer({
89 size: fmBytes,
90 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
91 });
92 const qlm = new Float32Array(2 * ref.nlm);
93 const row = new Float64Array(lmax + 1);
95 /** fm[m][ilat] after a synthesis of the unit spectrum at (l0, m). */
96 const probe = async (l0: number): Promise<Float32Array> => {
97 qlm.fill(0);
98 qlm[2 * lmIndex(lmax, l0, m)] = 1;
99 device!.queue.writeBuffer(plan!.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
100 const enc = device!.createCommandEncoder();
101 plan!.encodeSynth(enc);
102 enc.copyBufferToBuffer(plan!.fmBuf, 0, stageFm, 0, fmBytes);
103 device!.queue.submit([enc.finish()]);
104 await stageFm.mapAsync(GPUMapMode.READ);
105 const out = new Float32Array(stageFm.getMappedRange().slice(0));
106 stageFm.unmap();
107 return out;
108 };
110 const rows: { l: number; rels: number[]; ratios: number[] }[] = [];
111 for (let l0 = m; l0 <= lmax; l0++) {
112 const fm = await probe(l0);
113 const rels: number[] = [];
114 const ratios: number[] = [];
115 for (const ilat of lats) {
116 legendreRow(ref.coeffs, lmax, m, ref.ct[ilat], ref.st[ilat], row);
117 const want = row[l0 - m];
118 const got = fm[2 * (m * nlat + ilat)];
119 rels.push(Math.abs(got - want) / Math.max(Math.abs(want), 1e-300));
120 ratios.push(want === 0 ? NaN : got / want);
121 }
122 rows.push({ l: l0, rels, ratios });
123 }
125 // Loose on purpose. A forward Legendre recurrence in fp32 loses relative
126 // accuracy as it goes — by l = 63 a few 1e-5 is normal, and worse at the
127 // latitudes where the terms nearly cancel. What we are hunting is a structural
128 // error, which shows up as a ratio far from 1, not as a slow drift.
129 const OK = 1e-2;
130 const bad = (r: { rels: number[] }): boolean => r.rels.some((x) => !(x < OK));
131 const firstBad = rows.find(bad);
133 const head = ` l ` + lats.map((i) => `ilat ${String(i).padStart(3)}`.padStart(14)).join('');
134 console.log(head);
135 console.log(` ${'-'.repeat(head.length)}`);
136 for (const r of rows) {
137 // every l when --all; otherwise the seed, the first step, the first failure
138 // and its neighbours, and a tail sample — enough to see the shape
139 const near = firstBad ? Math.abs(r.l - firstBad.l) <= 3 : false;
140 const interesting =
141 showAll || r.l <= m + 2 || near || r.l >= lmax - 1 || (r.l - m) % 8 === 0;
142 if (!interesting) continue;
143 const cells = r.rels
144 .map((rel, k) =>
145 (rel < OK
146 ? `ok ${rel.toExponential(1)}`
147 : `${r.ratios[k] > 1e3 || r.ratios[k] < -1e3 ? '' : 'x'}${r.ratios[k].toExponential(2)}`
148 ).padStart(14),
149 )
150 .join('');
151 console.log(` ${String(r.l).padStart(4)} ${cells}${bad(r) ? ' <-- wrong' : ''}`);
152 }
153 console.log(
154 `\n cells are "ok <relative error>" when the term is right, and the ratio got/want\n` +
155 ` when it is not. A few 1e-5 by l = ${lmax} is normal: an fp32 forward recurrence\n` +
156 ` loses relative accuracy as it goes, worst where consecutive terms nearly cancel.`,
157 );
159 if (!firstBad) {
160 console.log(`\n Every term of the m = ${m} recurrence is right on this device.`);
161 console.log(
162 ` So the problem is not the recurrence itself — try another --m, or look at\n` +
163 ` the accumulation into acc rather than the values going into it.`,
164 );
165 } else {
166 const steps = Math.floor((firstBad.l - m) / 2);
167 console.log(`\n First wrong term: l = ${firstBad.l}, which is`);
168 if (firstBad.l === m) {
169 console.log(
170 ` the seed itself — amm[m] or sinpow_rescaled, before any recurrence runs.`,
171 );
172 } else if (firstBad.l === m + 1) {
173 console.log(
174 ` y1's initializer, ab[base + 1].x * ct * y0 — so a_{m+1}^m as the shader\n` +
175 ` reads it, or the very first multiply. No loop iteration has run yet.`,
176 );
177 } else {
178 console.log(
179 ` ${steps} advance${steps === 1 ? '' : 's'} into the loop (l = m + ${firstBad.l - m}).`,
180 );
181 const ok = rows.filter((r) => !bad(r)).map((r) => r.l);
182 console.log(
183 ` Terms that are right: l = ${ok.slice(0, 10).join(', ')}` +
184 (ok.length > 10 ? ', ...' : ''),
185 );
186 const parity = new Set(ok.map((l) => (l - m) % 2));
187 if (parity.size === 1) {
188 console.log(
189 ` Every correct term has (l - m) % 2 == ${[...parity][0]}, and every wrong one\n` +
190 ` the other parity. The loop carries two values per iteration — y0 for even\n` +
191 ` offsets and y1 for odd — so one of the two is being updated wrongly while\n` +
192 ` the other is fine.`,
193 );
194 }
195 }
196 }
198 stageFm.destroy();
199 plan.destroy();
200 device.destroy();
201} catch (e) {
202 plan?.destroy();
203 device?.destroy();
204 console.error(`diagnose-leg: ${errMsg(e)}`);
205 process.exit(1);
206}