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));
24}
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 );
114 // The undivided theta derivative sin(theta)*dtheta(u) — the flux-form
115 // Laplace-Beltrami scheme's step 1 and the flux-metric precompute's
116 // input — is the same shuffle+synthesis with the divide skipped, so it
117 // gets the same oracle.
118 const sinDthetaGpu = await deriv.sinDtheta(new Float32Array(q64));
119 const sinDthetaCpu = ref.sinDtheta(q64);
120 const errSinDtheta = relL2(sinDthetaGpu, sinDthetaCpu);
121 check(
122 'deriv: WGSL fp32 sinDtheta (undivided) vs f64 CPU reference',
123 errSinDtheta < 1e-4,
124 `sinDtheta ${errSinDtheta.toExponential(2)}`,
125 );
126 deriv.destroy();
127 }
129 // ---- grid-space phi-derivative (dphig) vs the f64 reference -------------
130 // dphig differentiates in phi with two Fourier stages and an i*m multiply,
131 // no Legendre work. On a band-limited field whose m >= lmax-2 modes are
132 // zero (dphig masks those, mirroring filt), it must agree with the
133 // coefficient-space route dphi = synth(i*m*coeffs) to fp32.
134 {
135 const q64 = new Float64Array(randomSpectrum(cfg, 4242));
136 for (let m = Math.max(0, lmax - 2); m <= lmax; m++) {
137 for (let l = m; l <= lmax; l++) {
138 const i = 2 * (m * (lmax + 1) - (m * (m - 1)) / 2 + (l - m));
139 q64[i] = 0;
140 q64[i + 1] = 0;
141 }
142 }
143 const grid = ref.synth(q64);
144 const dPhiGpu = await plan.dphig(new Float32Array(grid));
145 const dPhiCpu = ref.dphi(q64);
146 const err = relL2(dPhiGpu, dPhiCpu);
147 check(
148 'dphig: grid-space FFT phi-derivative vs f64 CPU reference',
149 err < 1e-4,
150 `rel L2 ${err.toExponential(2)}`,
151 );
152 }
154 // ---- batched transforms reproduce the scalar transforms ------------------
155 // A batch walks the Legendre recurrence once for K fields with per-lane
156 // arithmetic textually identical to the scalar kernel's, so each lane must
157 // agree with the scalar path to shader-compiler latitude (FMA contraction
158 // may differ between the two modules; nothing else may).
159 {
160 const { nlat: gl, nphi: gp } = plan.cfg;
161 const npts = gl * gp;
162 const sizes = [];
163 for (let k = 2; k <= plan.batchK; k += 2) sizes.push(k);
164 check(
165 'batch: plan compiled batched pipelines',
166 plan.batchK >= 2,
167 `batchK = ${plan.batchK} (${sizes.map((s) => `x${s}`).join(', ') || 'none'})`,
168 );
169 for (const K of sizes) {
170 const qs = Array.from({ length: K }, (_, k) => randomSpectrum(cfg, 1000 + k));
171 const qBufs = qs.map((q, k) => {
172 const b = device.createBuffer({
173 label: `batch-test-q${k}`,
174 size: 8 * plan.nlm,
175 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
176 });
177 device.queue.writeBuffer(b, 0, q as Float32Array<ArrayBuffer>);
178 return b;
179 });
180 const spatBufs = qs.map((_, k) =>
181 device.createBuffer({
182 label: `batch-test-spat${k}`,
183 size: 4 * npts,
184 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
185 }),
186 );
187 const qOutBufs = qs.map((_, k) =>
188 device.createBuffer({
189 label: `batch-test-qout${k}`,
190 size: 8 * plan.nlm,
191 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
192 }),
193 );
194 const stage = device.createBuffer({
195 label: 'batch-test-stage',
196 size: K * (4 * npts + 8 * plan.nlm),
197 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
198 });
200 // One pass: batched synthesis of all K, then batched analysis back.
201 const synthB = plan.createSynthBatchBinding(
202 qs.map((_, k) => ({ qlmIn: qBufs[k], spatOut: spatBufs[k] })),
203 );
204 const analysB = plan.createAnalysBatchBinding(
205 qs.map((_, k) => ({ spatIn: spatBufs[k], qlmOut: qOutBufs[k] })),
206 );
207 const enc = device.createCommandEncoder({ label: 'batch-test' });
208 const pass = enc.beginComputePass();
209 plan.encodeSynthBatchInto(pass, synthB);
210 plan.encodeAnalysBatchInto(pass, analysB);
211 pass.end();
212 for (let k = 0; k < K; k++) {
213 enc.copyBufferToBuffer(spatBufs[k], 0, stage, k * 4 * npts, 4 * npts);
214 enc.copyBufferToBuffer(qOutBufs[k], 0, stage, K * 4 * npts + k * 8 * plan.nlm, 8 * plan.nlm);
215 }
216 device.queue.submit([enc.finish()]);
217 await stage.mapAsync(GPUMapMode.READ);
218 const raw = new Float32Array(stage.getMappedRange().slice(0));
219 stage.unmap();
221 let worstSynth = 0;
222 let worstAnalys = 0;
223 for (let k = 0; k < K; k++) {
224 const spatLane = raw.subarray(k * npts, (k + 1) * npts);
225 const qLane = raw.subarray(K * npts + k * 2 * plan.nlm, K * npts + (k + 1) * 2 * plan.nlm);
226 const spatScalar = await plan.synth(qs[k]);
227 const qScalar = await plan.analys(spatScalar);
228 worstSynth = Math.max(worstSynth, relL2(spatLane, spatScalar));
229 worstAnalys = Math.max(worstAnalys, relL2(qLane, qScalar));
230 }
231 check(
232 `batch: x${K} lanes match the scalar transforms`,
233 worstSynth < 1e-6 && worstAnalys < 1e-6,
234 `synth ${worstSynth.toExponential(2)}, analys ${worstAnalys.toExponential(2)} ` +
235 `across ${K} lanes`,
236 );
237 for (const b of [...qBufs, ...spatBufs, ...qOutBufs, stage]) b.destroy();
238 }
239 }
241 plan.destroy();
242}