/ concept-collection / turing-surface
Sign in
concept-collection / turing-surface
208 lines · 7.1 KBCodeBlameHistory
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';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 11import { alphaPlus, alphaMinus } from './derivCoeffs.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 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 /**
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 122 * Coefficient-space theta shift, f64: v_l^m = alpha^+(l-1,m) u_{l-1}^m +
123 * alpha^-(l+1,m) u_{l+1}^m (algos.tex eq. v_coeffs) — the coefficients of
124 * sin(theta) * dtheta(u). The same shift serves the divergence side of the
125 * six-transform Laplace-Beltrami scheme (step 5 of
126 * docs/reduced-transforms.md), which is why it is exposed
127 * apart from the synthesis.
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 129 dthetac(qlm: ArrayLike<number>): Float64Array {
130 const { lmax, mmax } = this.cfg;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 131 const v = new Float64Array(2 * this.nlm);
132 for (let m = 0; m <= mmax; m++) {
133 for (let l = m; l <= lmax; l++) {
134 const lm = lmIndex(lmax, l, m);
135 let re = 0;
136 let im = 0;
137 if (l - 1 >= m) {
138 const lm1 = lmIndex(lmax, l - 1, m);
139 const a = alphaPlus(l - 1, m);
140 re += a * qlm[2 * lm1];
141 im += a * qlm[2 * lm1 + 1];
142 }
143 if (l + 1 <= lmax) {
144 const lm2 = lmIndex(lmax, l + 1, m);
145 const a = alphaMinus(l + 1, m);
146 re += a * qlm[2 * lm2];
147 im += a * qlm[2 * lm2 + 1];
148 }
149 v[2 * lm] = re;
150 v[2 * lm + 1] = im;
151 }
152 }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 156 /** Coefficient-space phi derivative, f64: (dphi u)_l^m = i*m*u_l^m. */
157 dphic(qlm: ArrayLike<number>): Float64Array {
159 const v = new Float64Array(2 * this.nlm);
160 for (let m = 0; m <= mmax; m++) {
161 for (let l = m; l <= lmax; l++) {
162 const lm = lmIndex(lmax, l, m);
163 v[2 * lm] = -m * qlm[2 * lm + 1];
164 v[2 * lm + 1] = m * qlm[2 * lm];
165 }
166 }
168 }
170 /** sin(theta) * dtheta(u) on the grid, f64: the undivided synthesis of the
171 * theta shift. Smooth on the sphere, unlike dtheta(u) itself. */
172 sinDtheta(qlm: ArrayLike<number>): Float64Array {
173 return this.synth(this.dthetac(qlm));
174 }
176 /** Theta-derivative, f64: synth(dthetac(u)) / sin(theta). */
177 dtheta(qlm: ArrayLike<number>): Float64Array {
178 const { nlat, nphi } = this.cfg;
179 const grid = this.sinDtheta(qlm);
180 for (let i = 0; i < nlat; i++) {
181 const st = this.st[i];
182 for (let j = 0; j < nphi; j++) grid[i * nphi + j] /= st;
183 }
184 return grid;
185 }
187 /** Phi-derivative, f64: (dphi u)_l^m = i*m*u_l^m, then synthesize. */
188 dphi(qlm: ArrayLike<number>): Float64Array {
189 return this.synth(this.dphic(qlm));
193/** Random band-limited spectrum for testing (m=0 imaginary parts zeroed). */
194export function randomSpectrum(cfg: ShtConfig, seed = 12345): Float32Array {
195 const nlm = nlmCalc(cfg.lmax, cfg.mmax);
196 const q = new Float32Array(2 * nlm);
197 let s = seed >>> 0;
198 const rnd = () => {
199 // xorshift32
200 s ^= s << 13; s >>>= 0;
201 s ^= s >> 17;
202 s ^= s << 5; s >>>= 0;
203 return (s / 4294967296) * 2 - 1;
204 };
205 for (let k = 0; k < 2 * nlm; k++) q[k] = rnd();
206 for (let l = 0; l <= cfg.lmax; l++) q[2 * lmIndex(cfg.lmax, l, 0) + 1] = 0; // m=0 real
207 return q;
moveopenescclose