concept-collection / turing-surface
117 lines · 3.9 KBBlameHistoryRaw
1/**
2 * The WGSL spherical-harmonic transforms against the f64 CPU reference.
3 *
4 * This is the one place a second implementation is still the right oracle: the
5 * transforms are vendored shtns-webgpu, and `src/sht/reference.ts` is its direct-
6 * summation f64 twin. Everything above them (the .m models) is checked against
7 * closed-form answers instead — see analyticChecks.ts.
8 */
9import { ShtPlan } from '../src/sht/sht.ts';
10import { ShtReference, randomSpectrum } from '../src/sht/reference.ts';
11import { DerivPlan } from '../src/sht/deriv.ts';
12import { gridForLmax } from '../src/sht/layout.ts';
13import type { Check, Log } from './analyticChecks.ts';
15function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
16 let num = 0;
17 let den = 0;
18 for (let i = 0; i < a.length; i++) {
19 const d = a[i] - b[i];
20 num += d * d;
21 den += b[i] * b[i];
22 }
23 return Math.sqrt(num / Math.max(den, 1e-300));
26export async function transformChecks(
27 device: GPUDevice,
28 check: Check,
29 _log: Log,
30): Promise<void> {
31 const lmax = 31;
32 const { nlat, nphi } = gridForLmax(lmax, 1);
33 const cfg = { lmax, mmax: lmax, nlat, nphi };
35 const plan = await ShtPlan.create(device, cfg);
36 const ref = new ShtReference(cfg);
38 const q = randomSpectrum(cfg, 42);
39 const q64 = new Float64Array(q);
41 const spatGpu = await plan.synth(new Float32Array(q64));
42 const spatCpu = ref.synth(q64);
43 const errSynth = relL2(spatGpu, spatCpu);
45 const qGpu = await plan.analys(new Float32Array(spatCpu));
46 const qCpu = ref.analys(new Float64Array(spatCpu));
47 const errAnalys = relL2(qGpu, qCpu);
49 check(
50 'transforms: WGSL fp32 vs f64 CPU reference',
51 errSynth < 1e-4 && errAnalys < 1e-4,
52 `synth ${errSynth.toExponential(2)}, analys ${errAnalys.toExponential(2)}`,
53 );
55 // ---- f64 reference dtheta/dphi vs an independent closed form -----------
56 // x(theta,phi) = sin(theta)*cos(phi) is exactly degree 1, so quadrature
57 // recovers it to f64 round-off; comparing its dtheta/dphi against the
58 // grid-space analytic derivatives (not derived from the same recurrence
59 // being tested) catches a sign or indexing error the random-spectrum check
60 // below, which compares two implementations of the same formula, would not.
61 {
62 const x = new Float64Array(nlat * nphi);
63 for (let i = 0; i < nlat; i++) {
64 const st = ref.st[i];
65 for (let j = 0; j < nphi; j++) {
66 const phi = (2 * Math.PI * j) / nphi;
67 x[i * nphi + j] = st * Math.cos(phi);
68 }
69 }
70 const X = ref.analys(x);
71 const dThetaX = ref.dtheta(X);
72 const dPhiX = ref.dphi(X);
74 let errNum = 0;
75 let norm = 0;
76 for (let i = 0; i < nlat; i++) {
77 const ct = ref.ct[i];
78 const st = ref.st[i];
79 for (let j = 0; j < nphi; j++) {
80 const phi = (2 * Math.PI * j) / nphi;
81 const k = i * nphi + j;
82 const wantTheta = ct * Math.cos(phi);
83 const wantPhi = -st * Math.sin(phi);
84 errNum += (dThetaX[k] - wantTheta) ** 2 + (dPhiX[k] - wantPhi) ** 2;
85 norm += wantTheta * wantTheta + wantPhi * wantPhi;
86 }
87 }
88 const relErr = Math.sqrt(errNum / Math.max(norm, 1e-300));
89 check(
90 'deriv: f64 reference dtheta/dphi match the closed form on x = sin(theta)cos(phi)',
91 relErr < 1e-6,
92 `rel L2 error ${relErr.toExponential(2)}`,
93 );
94 }
96 // ---- WGSL fp32 dtheta/dphi vs the (now closed-form-verified) f64 reference
97 {
98 const deriv = await DerivPlan.create(device, plan);
100 const dThetaGpu = await deriv.dtheta(new Float32Array(q64));
101 const dThetaCpu = ref.dtheta(q64);
102 const errDtheta = relL2(dThetaGpu, dThetaCpu);
104 const dPhiGpu = await deriv.dphi(new Float32Array(q64));
105 const dPhiCpu = ref.dphi(q64);
106 const errDphi = relL2(dPhiGpu, dPhiCpu);
108 check(
109 'deriv: WGSL fp32 dtheta/dphi vs f64 CPU reference',
110 errDtheta < 1e-4 && errDphi < 1e-4,
111 `dtheta ${errDtheta.toExponential(2)}, dphi ${errDphi.toExponential(2)}`,
112 );
113 deriv.destroy();
114 }
116 plan.destroy();