1/**
2 * WGSL kernels for the coefficient-space step of the theta/phi first-
3 * derivative algorithm (evolving_surface/notes/algos.tex, Algorithm 1, theta
4 * and phi branches only -- the Laplace-Beltrami operator built on these never
5 * needs the second-derivative/curvature branches).
6 *
7 * Both derivatives are a shuffle across nearby spectral coefficients --
8 * independent of latitude/longitude, so neither touches the Legendre
9 * recurrence or Fourier stages in leg.ts/fourier.ts -- followed by the
10 * *existing*, unchanged scalar synthesis pipeline. dtheta additionally
11 * divides the synthesized grid field by sin(theta) afterwards.
12 */
14const WG = 64;
16export interface DerivCoeffParams {
17 nlm: number;
18}
20/**
21 * v_l^m = alpha^+(l-1,m) * u_{l-1}^m + alpha^-(l+1,m) * u_{l+1}^m, the
22 * coefficients of sin(theta) * dtheta(u) (algos.tex eq. v_coeffs). aPlus/
23 * aMinus are precomputed zero at each m-block's boundary
24 * (src/sht/derivCoeffs.ts), so the multiply is always mathematically
25 * correct; the bounds checks below exist only to avoid reading past the ends
26 * of the qlm array (an m-block-internal +-1 step never leaves the array, so
27 * this is the only place it could).
28 */
29export function dthetaShuffleWGSL(p: DerivCoeffParams): string {
30 return /* wgsl */ `
31const NLM: u32 = ${p.nlm}u;
33@group(0) @binding(0) var<storage, read> aPlus: array<f32>;
34@group(0) @binding(1) var<storage, read> aMinus: array<f32>;
35@group(0) @binding(2) var<storage, read> qlmIn: array<vec2f>;
36@group(0) @binding(3) var<storage, read_write> vOut: array<vec2f>;
38@compute @workgroup_size(${WG})
39fn dtheta_shuffle(@builtin(global_invocation_id) gid: vec3u) {
40 let lm = gid.x;
41 if (lm >= NLM) { return; }
42 var v = vec2f(0.0);
43 if (lm > 0u) { v += aPlus[lm] * qlmIn[lm - 1u]; }
44 if (lm + 1u < NLM) { v += aMinus[lm] * qlmIn[lm + 1u]; }
45 vOut[lm] = v;
46}
47`;
48}
50/**
51 * (dphi u)_l^m = i*m*u_l^m: in the [re, im] row layout this swaps and
52 * negates, re' = -m*im, im' = m*re (algos.tex eq. dYdphi).
53 */
54export function dphiShuffleWGSL(p: DerivCoeffParams): string {
55 return /* wgsl */ `
56const NLM: u32 = ${p.nlm}u;
58@group(0) @binding(0) var<storage, read> mOf: array<u32>;
59@group(0) @binding(1) var<storage, read> qlmIn: array<vec2f>;
60@group(0) @binding(2) var<storage, read_write> vOut: array<vec2f>;
62@compute @workgroup_size(${WG})
63fn dphi_shuffle(@builtin(global_invocation_id) gid: vec3u) {
64 let lm = gid.x;
65 if (lm >= NLM) { return; }
66 let m = f32(mOf[lm]);
67 let c = qlmIn[lm];
68 vOut[lm] = vec2f(-m * c.y, m * c.x);
69}
70`;
71}
73export interface DivideParams {
74 nlat: number;
75 nphi: number;
76}
78/**
79 * Elementwise divide by sin(theta): the grid-space finish of Algorithm 1's
80 * dtheta branch (dtheta(u) = synth(v_l^m) / sin(theta)). Gauss nodes never
81 * sit at the poles, so this never divides by zero.
82 */
83export function divideSinThetaWGSL(p: DivideParams): string {
84 const npts = p.nlat * p.nphi;
85 return /* wgsl */ `
86const NLAT: u32 = ${p.nlat}u;
87const NPHI: u32 = ${p.nphi}u;
88const NPTS: u32 = ${npts}u;
90@group(0) @binding(0) var<storage, read> sinTheta: array<f32>;
91@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
93@compute @workgroup_size(${WG})
94fn divide_sin_theta(@builtin(global_invocation_id) gid: vec3u) {
95 let i = gid.x;
96 if (i >= NPTS) { return; }
97 let ilat = i / NPHI;
98 spat[i] = spat[i] / sinTheta[ilat];
99}
100`;
101}
103export interface FmDphiParams {
104 mmax: number;
105 nlat: number;
106 nphi: number;
107 /** Highest m kept; modes above are zeroed (mirrors the l-space filt). */
108 mcut: number;
109}
111/**
112 * The Fourier-space middle of the grid-space phi-derivative `dphig`:
113 * fm holds the unnormalized DFT modes of each latitude row (what the
114 * Fourier analysis stage produces), so d/dphi is fm[m] *= i*m/NPHI --
115 * the 1/NPHI undoes the unnormalized analysis+synthesis round trip.
116 * Modes above MCUT are zeroed: the Fourier analysis stage already
117 * truncated m > mmax for free, and MCUT additionally mirrors the
118 * top-degree filt so the differentiated field carries no content the
119 * l-space route would not have kept.
120 */
121export function fmDphiWGSL(p: FmDphiParams): string {
122 const count = (p.mmax + 1) * p.nlat;
123 return /* wgsl */ `
124const NLAT: u32 = ${p.nlat}u;
125const COUNT: u32 = ${count}u;
126const MCUT: u32 = ${Math.max(0, p.mcut)}u;
127const INV_NPHI: f32 = ${1 / p.nphi};
129@group(0) @binding(0) var<storage, read_write> fm: array<vec2f>;
131@compute @workgroup_size(${WG})
132fn fm_dphi(@builtin(global_invocation_id) gid: vec3u) {
133 let i = gid.x;
134 if (i >= COUNT) { return; }
135 let m = i / NLAT;
136 var k: f32 = 0.0;
137 if (m <= MCUT) { k = f32(m) * INV_NPHI; }
138 let c = fm[i];
139 fm[i] = vec2f(-k * c.y, k * c.x);
140}
141`;
142}