Add scripts/diagnose-leg.ts: follow leg_synth's recurrence term by term
diagnose-sht.ts narrowed the RTX PRO 6000 failure to one shader:
leg_synth is wrong, both Fourier stages and leg_analys are right. The
breakdown said more than that. With lmax 63 the only correct orders are
m = 62 and 63 — exactly the two where the loop breaks before ever running
the recurrence advance — and m = 0 is wrong too, which needs no rescaling
at all and works entirely with O(1) values. So it is the advance, not the
rescaled seed.
Reading leg_synth against leg_analys does not settle it: the advance is
semantically the same in both, one via a temporary and one by relying on
y0 being assigned before y1 reads it.
So probe the production shader instead of a copy. qlm set to a single 1 at
(l0, m) makes fm[m][ilat] exactly ytilde_l0^m(theta_i), so one synthesis
per l0 returns the recurrence value at that l for every latitude, computed
by the code the app actually runs. Sweeping l0 gives the sequence, and the
first term that disagrees with legendreRow says which part is broken —
the seed, y1's initializer (and so the ab buffer as the shader reads it),
or the advance. If the correct terms all share one parity of (l - m), that
names which of the two carried values is being updated wrongly.
Tolerance is 1e-2, deliberately loose: an fp32 forward recurrence loses
relative accuracy with l, and on Intel it reaches ~1e-3 by l = 56 at the
latitudes where consecutive terms nearly cancel. That is normal and not
what we are looking for.
1 changed file+206−0
scripts/diagnose-leg.tsadded+206−0View file
@@ -0,0 +1,206 @@
1+/**
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+ */
25+import { ShtPlan, requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
26+import { ShtReference } from '../src/sht/reference.ts';
27+import { legendreRow } from '../src/sht/coeffs.ts';
28+import { gridForLmax, lmIndex, type ShtConfig } from '../src/sht/layout.ts';
29+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
30+
31+const argv = process.argv.slice(2);
32+if (argv.includes('--help') || argv.includes('-h')) {
33+ console.log(`usage: npx vite-node scripts/diagnose-leg.ts [options]
34+
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+}
43+const 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+};
49+const lmax = Number(flag('lmax', '63'));
50+const m = Number(flag('m', '0'));
51+const showAll = argv.includes('--all');
52+
53+if (!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+}
57+
58+let device: GPUDevice | null = null;
59+let plan: ShtPlan | null = null;
60+try {
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);
66+
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' });
72+
73+ const lats = flag('lats', '')
74+ ? flag('lats', '').split(',').map(Number)
75+ : [0, 1, nlat >> 1, nlat - 1];
76+
77+ console.log('turing-sphere — 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+ );
86+
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);
94+
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+ };
109+
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+ }
124+
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);
132+
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+ );
158+
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+ }
197+
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+}