1/**
2 * Double-precision reference implementation of the scalar spherical
3 * harmonic transform, by direct summation. Slow (O(nlat*nlm) Legendre +
4 * O(nlat*nphi*mmax) Fourier) but simple, and serves as ground truth for
5 * validating the fp32 WebGPU implementation.
6 *
7 * Conventions are identical to the GPU path (see layout.ts).
8 */
9import { gaussNodesWeights } from './gauss.ts';
10import { legendreCoeffs, legendreRow, type LegendreCoeffs } from './coeffs.ts';
11import { alphaPlus, alphaMinus } from './derivCoeffs.ts';
12import { lmIndex, nlmCalc, validateConfig, type ShtConfig } from './layout.ts';
14export class ShtReference {
15 readonly cfg: ShtConfig;
16 readonly nlm: number;
17 readonly ct: Float64Array;
18 readonly st: Float64Array;
19 readonly wg: Float64Array; // Gauss weights (for integral over cos(theta))
20 readonly coeffs: LegendreCoeffs;
22 constructor(cfg: ShtConfig) {
23 validateConfig(cfg);
24 this.cfg = cfg;
25 this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
26 const { x, w } = gaussNodesWeights(cfg.nlat);
27 this.ct = x;
28 this.wg = w;
29 this.st = new Float64Array(cfg.nlat);
30 for (let i = 0; i < cfg.nlat; i++) this.st[i] = Math.sqrt(1 - x[i] * x[i]);
31 this.coeffs = legendreCoeffs(cfg.lmax, cfg.mmax);
32 }
34 /**
35 * Legendre stage of the synthesis: F_m(theta_i) = sum_l Q_lm ytilde_l^m(theta_i).
36 * Returns complex array indexed [m * nlat + ilat], interleaved re/im.
37 */
38 legendreSynth(qlm: ArrayLike<number>): Float64Array {
39 const { lmax, mmax, nlat } = this.cfg;
40 const fm = new Float64Array(2 * (mmax + 1) * nlat);
41 const row = new Float64Array(lmax + 1);
42 for (let i = 0; i < nlat; i++) {
43 for (let m = 0; m <= mmax; m++) {
44 legendreRow(this.coeffs, lmax, m, this.ct[i], this.st[i], row);
45 let re = 0, im = 0;
46 const base = lmIndex(lmax, m, m);
47 for (let l = m; l <= lmax; l++) {
48 const y = row[l - m];
49 re += y * qlm[2 * (base + l - m)];
50 im += y * qlm[2 * (base + l - m) + 1];
51 }
52 const o = 2 * (m * nlat + i);
53 fm[o] = re;
54 fm[o + 1] = im;
55 }
56 }
57 return fm;
58 }
60 /** Full synthesis: spectral -> spatial grid [ilat * nphi + iphi]. */
61 synth(qlm: ArrayLike<number>): Float64Array {
62 const { mmax, nlat, nphi } = this.cfg;
63 const fm = this.legendreSynth(qlm);
64 const spat = new Float64Array(nlat * nphi);
65 for (let i = 0; i < nlat; i++) {
66 for (let j = 0; j < nphi; j++) {
67 const phi = (2 * Math.PI * j) / nphi;
68 let v = fm[2 * (0 * nlat + i)]; // m=0: real part (imag must be 0)
69 for (let m = 1; m <= mmax; m++) {
70 const o = 2 * (m * nlat + i);
71 const c = Math.cos(m * phi);
72 const s = Math.sin(m * phi);
73 v += 2 * (fm[o] * c - fm[o + 1] * s);
74 }
75 spat[i * nphi + j] = v;
76 }
77 }
78 return spat;
79 }
81 /** Full analysis: spatial grid -> spectral coefficients (interleaved re/im). */
82 analys(spat: ArrayLike<number>): Float64Array {
83 const { lmax, mmax, nlat, nphi } = this.cfg;
84 const qlm = new Float64Array(2 * this.nlm);
85 const row = new Float64Array(lmax + 1);
86 // forward Fourier: G_m(theta_i) = (2*pi/nphi) * sum_j f_ij e^{-i m phi_j}
87 const gm = new Float64Array(2 * (mmax + 1) * nlat);
88 for (let i = 0; i < nlat; i++) {
89 for (let m = 0; m <= mmax; m++) {
90 let re = 0, im = 0;
91 for (let j = 0; j < nphi; j++) {
92 const phi = (2 * Math.PI * j) / nphi;
93 const f = spat[i * nphi + j];
94 re += f * Math.cos(m * phi);
95 im -= f * Math.sin(m * phi);
96 }
97 const o = 2 * (m * nlat + i);
98 const norm = (2 * Math.PI) / nphi;
99 gm[o] = re * norm;
100 gm[o + 1] = im * norm;
101 }
102 }
103 // Legendre stage with Gauss quadrature: Q_lm = sum_i w_i ytilde_l^m(theta_i) G_m(theta_i)
104 for (let m = 0; m <= mmax; m++) {
105 const base = lmIndex(lmax, m, m);
106 for (let i = 0; i < nlat; i++) {
107 legendreRow(this.coeffs, lmax, m, this.ct[i], this.st[i], row);
108 const o = 2 * (m * nlat + i);
109 const wr = this.wg[i] * gm[o];
110 const wi = this.wg[i] * gm[o + 1];
111 for (let l = m; l <= lmax; l++) {
112 const y = row[l - m];
113 qlm[2 * (base + l - m)] += y * wr;
114 qlm[2 * (base + l - m) + 1] += y * wi;
115 }
116 }
117 }
118 return qlm;
119 }
121 /**
122 * Theta-derivative, f64: v_l^m = alpha^+(l-1,m) u_{l-1}^m + alpha^-(l+1,m)
123 * u_{l+1}^m (algos.tex eq. v_coeffs), then synth(v_l^m) / sin(theta).
124 */
125 dtheta(qlm: ArrayLike<number>): Float64Array {
126 const { lmax, mmax, nlat, nphi } = this.cfg;
127 const v = new Float64Array(2 * this.nlm);
128 for (let m = 0; m <= mmax; m++) {
129 for (let l = m; l <= lmax; l++) {
130 const lm = lmIndex(lmax, l, m);
131 let re = 0;
132 let im = 0;
133 if (l - 1 >= m) {
134 const lm1 = lmIndex(lmax, l - 1, m);
135 const a = alphaPlus(l - 1, m);
136 re += a * qlm[2 * lm1];
137 im += a * qlm[2 * lm1 + 1];
138 }
139 if (l + 1 <= lmax) {
140 const lm2 = lmIndex(lmax, l + 1, m);
141 const a = alphaMinus(l + 1, m);
142 re += a * qlm[2 * lm2];
143 im += a * qlm[2 * lm2 + 1];
144 }
145 v[2 * lm] = re;
146 v[2 * lm + 1] = im;
147 }
148 }
149 const grid = this.synth(v);
150 for (let i = 0; i < nlat; i++) {
151 const st = this.st[i];
152 for (let j = 0; j < nphi; j++) grid[i * nphi + j] /= st;
153 }
154 return grid;
155 }
157 /** Phi-derivative, f64: (dphi u)_l^m = i*m*u_l^m, then synthesize. */
158 dphi(qlm: ArrayLike<number>): Float64Array {
159 const { lmax, mmax } = this.cfg;
160 const v = new Float64Array(2 * this.nlm);
161 for (let m = 0; m <= mmax; m++) {
162 for (let l = m; l <= lmax; l++) {
163 const lm = lmIndex(lmax, l, m);
164 v[2 * lm] = -m * qlm[2 * lm + 1];
165 v[2 * lm + 1] = m * qlm[2 * lm];
166 }
167 }
168 return this.synth(v);
169 }
170}
172/** Random band-limited spectrum for testing (m=0 imaginary parts zeroed). */
173export function randomSpectrum(cfg: ShtConfig, seed = 12345): Float32Array {
174 const nlm = nlmCalc(cfg.lmax, cfg.mmax);
175 const q = new Float32Array(2 * nlm);
176 let s = seed >>> 0;
177 const rnd = () => {
178 // xorshift32
179 s ^= s << 13; s >>>= 0;
180 s ^= s >> 17;
181 s ^= s << 5; s >>>= 0;
182 return (s / 4294967296) * 2 - 1;
183 };
184 for (let k = 0; k < 2 * nlm; k++) q[k] = rnd();
185 for (let l = 0; l <= cfg.lmax; l++) q[2 * lmIndex(cfg.lmax, l, 0) + 1] = 0; // m=0 real
186 return q;
187}