1/**
2 * Which half of a transform is wrong on this device?
3 *
4 * npx vite-node scripts/diagnose-sht.ts [--lmax 63] [--seed 12345]
5 *
6 * A transform is two stages, and both directions share code, so a single
7 * pass/fail says very little:
8 *
9 * synthesis: qlm --[leg_synth]--> fm --[fft_synth | dft_synth]--> spat
10 * analysis: spat --[fft_analys | dft_analys]--> fm --[leg_analys]--> qlm
11 *
12 * This reads the intermediate `fm` back out and compares each stage against
13 * src/sht/reference.ts (f64, direct summation) on its own:
14 *
15 * - fm wrong -> the Legendre stage
16 * - fm right but spat wrong -> the Fourier stage
17 * - both right in DFT, wrong in FFT -> the WGSL FFT specifically
18 *
19 * and breaks the error down by m and by latitude, because "only high m" or "only
20 * near the poles" points straight at the rescaled recurrence, while "every m
21 * equally" points at indexing.
22 *
23 * Written for a report of `npm run test:node` failing on a GPU the transforms
24 * have not run on before. Nothing here is a benchmark.
25 */
26import { ShtPlan, requestShtDevice, describeAdapter } from '../src/sht/sht.ts';
27import { ShtReference, randomSpectrum } from '../src/sht/reference.ts';
28import { gridForLmax, isPowerOfTwo, type ShtConfig } from '../src/sht/layout.ts';
29import { fftThreads } from '../src/sht/wgsl/fourier.ts';
30import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
32const argv = process.argv.slice(2);
33if (argv.includes('--help') || argv.includes('-h')) {
34 console.log(`usage: npx vite-node scripts/diagnose-sht.ts [options]
36 --lmax <n> spherical harmonic truncation (default 63, the app's)
37 --seed <n> seed of the test spectrum (default 12345)
38 --fourier <m> only test this stage: fft | dft (default: both)
39 --help
41Compares each stage of each direction against the f64 CPU reference and says
42which one is wrong. See the header of this file.`);
43 process.exit(0);
44}
45const flag = (name: string, dflt: string): string => {
46 const i = argv.indexOf(`--${name}`);
47 if (i >= 0 && argv[i + 1] !== undefined) return argv[i + 1];
48 const eq = argv.find((a) => a.startsWith(`--${name}=`));
49 return eq ? eq.slice(name.length + 3) : dflt;
50};
51const lmax = Number(flag('lmax', '63'));
52const seed = Number(flag('seed', '12345'));
53const only = flag('fourier', '');
55/** Relative L2 of a against b, both flat. */
56function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
57 let num = 0;
58 let den = 0;
59 for (let i = 0; i < b.length; i++) {
60 const d = (a[i] ?? NaN) - b[i];
61 num += d * d;
62 den += b[i] * b[i];
63 }
64 return Math.sqrt(num / Math.max(den, 1e-300));
65}
67function anyNonFinite(a: ArrayLike<number>): boolean {
68 for (let i = 0; i < a.length; i++) if (!Number.isFinite(a[i])) return true;
69 return false;
70}
72/** The Fourier half of a synthesis, on the host, from whatever fm it is given.
73 * Mirrors ShtReference.synth's inner loop — so feeding it the GPU's own fm says
74 * what the Fourier stage should have produced from the input it actually had. */
75function fourierSynth(cfg: ShtConfig, fm: ArrayLike<number>): Float64Array {
76 const { mmax, nlat, nphi } = cfg;
77 const spat = new Float64Array(nlat * nphi);
78 for (let i = 0; i < nlat; i++) {
79 for (let j = 0; j < nphi; j++) {
80 const phi = (2 * Math.PI * j) / nphi;
81 let v = fm[2 * i];
82 for (let m = 1; m <= mmax; m++) {
83 const o = 2 * (m * nlat + i);
84 v += 2 * (fm[o] * Math.cos(m * phi) - fm[o + 1] * Math.sin(m * phi));
85 }
86 spat[i * nphi + j] = v;
87 }
88 }
89 return spat;
90}
92/** Worst offender along one axis of the [m][ilat] complex fm array. */
93function fmBreakdown(
94 cfg: ShtConfig,
95 got: ArrayLike<number>,
96 want: ArrayLike<number>,
97): { byM: { m: number; rel: number }[]; worstLat: { ilat: number; rel: number } } {
98 const { mmax, nlat } = cfg;
99 const byM: { m: number; rel: number }[] = [];
100 const latNum = new Float64Array(nlat);
101 const latDen = new Float64Array(nlat);
102 for (let m = 0; m <= mmax; m++) {
103 let num = 0;
104 let den = 0;
105 for (let i = 0; i < nlat; i++) {
106 for (let c = 0; c < 2; c++) {
107 const k = 2 * (m * nlat + i) + c;
108 const d = (got[k] ?? NaN) - want[k];
109 num += d * d;
110 den += want[k] * want[k];
111 latNum[i] += d * d;
112 latDen[i] += want[k] * want[k];
113 }
114 }
115 byM.push({ m, rel: Math.sqrt(num / Math.max(den, 1e-300)) });
116 }
117 let worstLat = { ilat: 0, rel: 0 };
118 for (let i = 0; i < nlat; i++) {
119 const rel = Math.sqrt(latNum[i] / Math.max(latDen[i], 1e-300));
120 if (rel > worstLat.rel) worstLat = { ilat: i, rel };
121 }
122 return { byM, worstLat };
123}
125const OK = 1e-4; // fp32 through these transforms lands near 1e-6; 1e-4 is generous
126const verdict = (rel: number): string => (rel < OK ? 'ok ' : 'WRONG');
128let device: GPUDevice | null = null;
129try {
130 const runtime = await installWebGpu();
131 device = await requestShtDevice().catch((e: unknown) => {
132 throw new Error(`${errMsg(e)}\n${NO_ADAPTER_HINT}`);
133 });
134 const adapter = await describeAdapter(device);
136 const { nlat, nphi } = gridForLmax(lmax, 3);
137 const cfg: ShtConfig = { lmax, mmax: lmax, nlat, nphi };
138 const ref = new ShtReference(cfg);
139 const qlm = randomSpectrum(cfg, seed);
141 console.log('turing-sphere — which stage of the transform is wrong?\n');
142 console.log(` device ${adapter || '(unknown)'}\n ${runtime}`);
143 console.log(` grid lmax ${cfg.lmax} · ${nlat}×${nphi} · nlm ${ref.nlm}`);
144 console.log(
145 ` limits maxComputeWorkgroupStorageSize ${device.limits.maxComputeWorkgroupStorageSize}` +
146 `, maxComputeInvocationsPerWorkgroup ${device.limits.maxComputeInvocationsPerWorkgroup}`,
147 );
148 console.log(
149 ` the FFT stage needs a power-of-two nphi (${isPowerOfTwo(nphi)}), ` +
150 `16*nphi = ${16 * nphi} bytes of\n workgroup storage and ` +
151 `${fftThreads(nphi)} invocations per workgroup\n`,
152 );
154 // reference values, computed once
155 const fmRef = ref.legendreSynth(qlm);
156 const spatRef = ref.synth(qlm);
157 const qlmRef = ref.analys(spatRef);
159 const modes: ('fft' | 'dft')[] =
160 only === 'fft' || only === 'dft' ? [only] : ['fft', 'dft'];
161 const summary: string[] = [];
163 for (const mode of modes) {
164 let plan: ShtPlan | null = null;
165 try {
166 plan = await ShtPlan.create(device, cfg, { fourier: mode });
167 } catch (e) {
168 console.log(`${mode.toUpperCase()} stage: unavailable — ${errMsg(e)}\n`);
169 continue;
170 }
172 const stageFm = device.createBuffer({
173 size: 8 * (cfg.mmax + 1) * nlat,
174 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
175 });
176 const stageSpat = device.createBuffer({
177 size: 4 * nlat * nphi,
178 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
179 });
180 const read = async (buf: GPUBuffer): Promise<Float32Array> => {
181 await buf.mapAsync(GPUMapMode.READ);
182 const out = new Float32Array(buf.getMappedRange().slice(0));
183 buf.unmap();
184 return out;
185 };
187 // --- synthesis, stopping to look at fm on the way through ---
188 device.queue.writeBuffer(plan.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
189 const enc = device.createCommandEncoder();
190 plan.encodeSynth(enc);
191 enc.copyBufferToBuffer(plan.fmBuf, 0, stageFm, 0, 8 * (cfg.mmax + 1) * nlat);
192 enc.copyBufferToBuffer(plan.spatBuf, 0, stageSpat, 0, 4 * nlat * nphi);
193 device.queue.submit([enc.finish()]);
194 const fmGpu = await read(stageFm);
195 const spatGpu = await read(stageSpat);
197 const legSynthRel = relL2(fmGpu, fmRef);
198 const synthRel = relL2(spatGpu, spatRef);
199 // The Fourier stage judged on its own input, not on the reference's: if the
200 // Legendre stage is already wrong, comparing spat with spatRef only repeats
201 // that. This asks whether the Fourier stage did the right thing with the fm
202 // it was actually handed.
203 const fourierRel = relL2(spatGpu, fourierSynth(cfg, fmGpu));
205 // --- analysis, for contrast: same two stages, opposite order ---
206 const spatIn = Float32Array.from(spatRef);
207 device.queue.writeBuffer(plan.spatBuf, 0, spatIn as Float32Array<ArrayBuffer>);
208 const enc2 = device.createCommandEncoder();
209 plan.encodeAnalys(enc2);
210 enc2.copyBufferToBuffer(plan.fmBuf, 0, stageFm, 0, 8 * (cfg.mmax + 1) * nlat);
211 device.queue.submit([enc2.finish()]);
212 const fmAnalysGpu = await read(stageFm);
213 const qlmGpu = await plan.analys(spatIn);
214 // forward Fourier of the reference field, in the reference's own normalization
215 const gmRef = new Float64Array(2 * (cfg.mmax + 1) * nlat);
216 for (let i = 0; i < nlat; i++) {
217 for (let m = 0; m <= cfg.mmax; m++) {
218 let re = 0;
219 let im = 0;
220 for (let j = 0; j < nphi; j++) {
221 const phi = (2 * Math.PI * j) / nphi;
222 re += spatRef[i * nphi + j] * Math.cos(m * phi);
223 im -= spatRef[i * nphi + j] * Math.sin(m * phi);
224 }
225 gmRef[2 * (m * nlat + i)] = re;
226 gmRef[2 * (m * nlat + i) + 1] = im;
227 }
228 }
229 const analysFourierRel = relL2(fmAnalysGpu, gmRef);
230 const analysRel = relL2(qlmGpu, qlmRef);
232 console.log(`${mode.toUpperCase()} stage — plan chose ${plan.fourierMode.toUpperCase()}\n`);
233 console.log(` synthesis qlm -> fm -> spat`);
234 console.log(
235 ` ${verdict(legSynthRel)} leg_synth fm vs f64 reference ` +
236 `${legSynthRel.toExponential(2)}${anyNonFinite(fmGpu) ? ' (has NaN/Inf)' : ''}`,
237 );
238 console.log(
239 ` ${verdict(fourierRel)} ${mode}_synth spat vs host Fourier of that fm ` +
240 `${fourierRel.toExponential(2)}${anyNonFinite(spatGpu) ? ' (has NaN/Inf)' : ''}`,
241 );
242 console.log(
243 ` ${verdict(synthRel)} end to end spat vs f64 reference ` +
244 `${synthRel.toExponential(2)}`,
245 );
246 console.log(`\n analysis spat -> fm -> qlm`);
247 console.log(
248 ` ${verdict(analysFourierRel)} ${mode}_analys fm vs f64 reference ` +
249 `${analysFourierRel.toExponential(2)}`,
250 );
251 console.log(
252 ` ${verdict(analysRel)} end to end qlm vs f64 reference ` +
253 `${analysRel.toExponential(2)}`,
254 );
256 if (legSynthRel >= OK) {
257 const { byM, worstLat } = fmBreakdown(cfg, fmGpu, fmRef);
258 const bad = byM.filter((e) => e.rel >= OK);
259 const good = byM.filter((e) => e.rel < OK);
260 console.log(`\n leg_synth is wrong. Where:`);
261 console.log(
262 ` ${bad.length} of ${byM.length} orders m are wrong` +
263 (good.length
264 ? `; the ones that are right are m = ${good.slice(0, 12).map((e) => e.m).join(', ')}` +
265 (good.length > 12 ? ', ...' : '')
266 : ' (all of them)'),
267 );
268 if (bad.length) {
269 const first = bad[0];
270 const worst = bad.reduce((a, b) => (b.rel > a.rel ? b : a));
271 console.log(
272 ` lowest wrong m = ${first.m} (${first.rel.toExponential(2)}), ` +
273 `worst m = ${worst.m} (${worst.rel.toExponential(2)})`,
274 );
275 }
276 const theta = (Math.acos(ref.ct[worstLat.ilat]) * 180) / Math.PI;
277 console.log(
278 ` worst latitude ilat = ${worstLat.ilat} of ${nlat} ` +
279 `(theta = ${theta.toFixed(1)}°, sin(theta) = ` +
280 `${ref.st[worstLat.ilat].toExponential(2)}), rel ${worstLat.rel.toExponential(2)}`,
281 );
282 console.log(
283 ` If only high m are wrong, or only latitudes near the poles where\n` +
284 ` sin(theta) is small, the rescaled seed (sinpow_rescaled in\n` +
285 ` src/sht/wgsl/common.ts) is the place to look. If every m is wrong by\n` +
286 ` a similar amount, it is indexing or the dispatch, not the recurrence.\n` +
287 ` Either way, follow it term by term from here:\n` +
288 ` npx vite-node scripts/diagnose-leg.ts --lmax ${lmax} --m 0`,
289 );
290 }
291 console.log();
293 summary.push(
294 `${mode}: leg_synth ${verdict(legSynthRel).trim()}, ${mode}_synth ` +
295 `${verdict(fourierRel).trim()}, ${mode}_analys ${verdict(analysFourierRel).trim()}, ` +
296 `leg_analys ${verdict(analysRel).trim()}`,
297 );
299 stageFm.destroy();
300 stageSpat.destroy();
301 plan.destroy();
302 }
304 console.log('summary');
305 for (const s of summary) console.log(` ${s}`);
306 device.destroy();
307} catch (e) {
308 device?.destroy();
309 console.error(`diagnose-sht: ${errMsg(e)}`);
310 process.exit(1);
311}