9b71d5bexport standalone matlabJeremy Magland 1% ---------------------------------------------------------------- transforms
2%
3% Double-precision MATLAB ports of the operations the app provides to a .m
4% around its compiled GPU pipeline. Conventions follow src/sht/layout.ts:
5% orthonormal spherical harmonics with the Condon-Shortley phase, coefficients
6% stored for m >= 0 only in m-major order (m = 0..mmax, l = m..lmax within
7% each m) -- here as complex nlm x 1 column vectors where the GPU carries
8% interleaved [re, im] pairs. Grid fields are npts x 1 columns, phi-fastest:
9% point (itheta, iphi) sits at row (itheta-1)*nphi + iphi, north row first.
11% Holds the precomputed tables between calls: set once from the top of the
12% run, read back by every transform below.
13function S = sht_tables(S)
14 persistent stored
15 if nargin > 0
16 stored = S;
17 end
18 S = stored;
19end
21% Everything the transforms need for one grid: Gauss nodes and weights,
22% per-m Legendre tables, the coefficient layout, the derivative shuffles,
23% and the eigenvalue/filter vectors the models take as `lam` and `filt`.
24function S = sht_setup(lmax, mmax, nlat, nphi)
25 S.lmax = lmax;
26 S.mmax = mmax;
27 S.nlat = nlat;
28 S.nphi = nphi;
29 S.npts = nlat * nphi;
30 [ct, wg] = gauss_legendre(nlat);
31 S.ct = ct;
32 S.st = sqrt(1 - ct.^2);
33 S.wg = wg;
34 S.nlm = (mmax + 1) * (lmax + 1) - mmax * (mmax + 1) / 2;
36 % The grid angles as npts x 1 fields, phi-fastest like everything else.
37 S.theta = repelem(acos(min(1, max(-1, ct))), nphi);
38 S.phi = repmat(2*pi*(0:nphi-1)'/nphi, nlat, 1);
39 S.stpt = repelem(S.st, nphi);
41 % Degree and order of each coefficient, and each m block's start.
42 off = zeros(mmax + 1, 1);
43 lv = zeros(S.nlm, 1);
44 mv = zeros(S.nlm, 1);
45 pos = 1;
46 for m = 0:mmax
47 n = lmax - m + 1;
48 off(m + 1) = pos;
49 lv(pos:pos + n - 1) = (m:lmax)';
50 mv(pos:pos + n - 1) = m;
51 pos = pos + n;
52 end
53 S.off = off;
54 S.lv = lv;
55 S.mv = mv;
56 % Laplace-Beltrami eigenvalues l(l+1) and the top-mode filter: 1 below
57 % lmax-2, 0 at the top two degrees, where the derivative recurrences cannot
58 % exactly represent a derivative (src/mgpu/model.ts).
59 S.lam = lv .* (lv + 1);
60 S.filt = double(lv < lmax - 2);
62 % Orthonormal Legendre tables ytilde_l^m(theta_i), one nlat x (lmax-m+1)
63 % block per m, by the standard three-term recurrence (src/sht/coeffs.ts;
64 % SHTNS normalization, Condon-Shortley phase carried in the seed's sign).
65 S.Y = cell(mmax + 1, 1);
66 t = 1 / (4*pi);
67 amm = sqrt(t);
68 for m = 0:mmax
69 if m > 0
70 t = t * (2*m + 1) / (2*m);
71 amm = (-1)^m * sqrt(t);
72 end
73 n = lmax - m + 1;
74 Y = zeros(nlat, n);
75 y0 = amm * S.st.^m;
76 Y(:, 1) = y0;
77 if n > 1
78 y1 = sqrt(2*m + 3) * ct .* y0;
79 Y(:, 2) = y1;
80 for l = m + 2:lmax
81 t1 = (l + m) * (l - m);
82 a = sqrt((2*l + 1) * (2*l - 1) / t1);
83 b = -sqrt(((2*l + 1) / (2*l - 3)) * ((l - 1 + m) * (l - 1 - m) / t1));
84 y2 = a * ct .* y1 + b * y0;
85 Y(:, l - m + 1) = y2;
86 y0 = y1;
87 y1 = y2;
88 end
89 end
90 S.Y{m + 1} = Y;
91 end
93 % sin(theta)*dtheta in coefficient space: v_l^m = ap(lm) u_{l-1}^m +
94 % am(lm) u_{l+1}^m (src/sht/derivCoeffs.ts). Neighbors sit at +-1 within
95 % each m block; ap/am are zero at the block edges, so the clamped index
96 % vectors never read across a boundary.
97 l = lv;
98 m = mv;
99 ap = (l - 1) .* sqrt(max(0, (l - m) .* (l + m)) ./ ((2*l - 1) .* (2*l + 1)));
100 ap(l <= m) = 0;
101 am = -(l + 2) .* sqrt((l + 1 - m) .* (l + 1 + m) ./ ((2*l + 1) .* (2*l + 3)));
102 am(l >= lmax) = 0;
103 S.ap = ap;
104 S.am = am;
105 S.iprev = max((1:S.nlm)' - 1, 1);
106 S.inext = min((1:S.nlm)' + 1, S.nlm);
108 % dphig's Fourier multiplier: i*m on fft's frequency layout, masked past
109 % the filter's reach (mcut = lmax-3), mirroring src/sht/wgsl/deriv.ts.
110 freq = [(0:nphi/2)'; (-nphi/2 + 1:-1)'];
111 S.dmul = 1i * freq .* (abs(freq) <= max(0, lmax - 3));
112end
114% Gauss-Legendre nodes cos(theta), in decreasing order (north pole first),
115% and weights for integration over cos(theta) -- Newton iteration on P_n,
116% as src/sht/gauss.ts.
117function [x, w] = gauss_legendre(n)
118 x = zeros(n, 1);
119 w = zeros(n, 1);
120 half = floor((n + 1) / 2);
121 for i = 1:half
122 z = cos(pi * (i - 0.25) / (n + 0.5));
123 pp = 0;
124 for it = 1:100
125 p1 = 1;
126 p2 = 0;
127 for j = 1:n
128 p3 = p2;
129 p2 = p1;
130 p1 = ((2*j - 1) * z * p2 - (j - 1) * p3) / j;
131 end
132 pp = n * (z * p1 - p2) / (z^2 - 1);
133 dz = p1 / pp;
134 z = z - dz;
135 if abs(dz) < 1e-15 * abs(z) + 1e-300
136 p1 = 1;
137 p2 = 0;
138 for j = 1:n
139 p3 = p2;
140 p2 = p1;
141 p1 = ((2*j - 1) * z * p2 - (j - 1) * p3) / j;
142 end
143 pp = n * (z * p1 - p2) / (z^2 - 1);
144 z = z - p1 / pp;
145 break;
146 end
147 end
148 x(i) = z;
149 x(n + 1 - i) = -z;
150 wi = 2 / ((1 - z^2) * pp^2);
151 w(i) = wi;
152 w(n + 1 - i) = wi;
153 end
154 if mod(n, 2) == 1
155 x(half) = 0;
156 end
157end
159% Synthesis, spectral -> grid. Grouped calls -- [a, b] = synth(x, y) -- are
160% the app's batching hint; here each member simply runs in turn.
161function varargout = synth(varargin)
162 S = sht_tables();
163 varargout = cell(1, nargin);
164 for k = 1:nargin
165 varargout{k} = synth_one(S, varargin{k});
166 end
167end
169function f = synth_one(S, Q)
170 % Legendre stage per m, then one inverse FFT per latitude ring with the
171 % m < 0 modes filled in by conjugate symmetry (the field is real).
172 G = zeros(S.nphi, S.nlat);
173 for m = 0:S.mmax
174 Fm = (S.Y{m + 1} * Q(S.off(m + 1):S.off(m + 1) + S.lmax - m)).';
175 G(m + 1, :) = Fm;
176 if m > 0
177 G(S.nphi + 1 - m, :) = conj(Fm);
178 end
179 end
180 f = S.nphi * real(ifft(G, [], 1));
181 f = f(:);
182end
184% Analysis, grid -> spectral: forward FFT per ring, then Gauss quadrature
185% against the same Legendre tables.
186function varargout = analys(varargin)
187 S = sht_tables();
188 varargout = cell(1, nargin);
189 for k = 1:nargin
190 varargout{k} = analys_one(S, varargin{k});
191 end
192end
194function Q = analys_one(S, f)
195 F = fft(reshape(f, S.nphi, S.nlat), [], 1) * (2*pi/S.nphi);
196 Q = complex(zeros(S.nlm, 1));
197 for m = 0:S.mmax
198 Q(S.off(m + 1):S.off(m + 1) + S.lmax - m) = S.Y{m + 1}.' * (S.wg .* F(m + 1, :).');
199 end
200end
202% The coefficients of sin(theta)*dtheta(u): the alpha^+/alpha^- shift by one
203% degree within each m block.
204function V = dthetac(Q)
205 S = sht_tables();
206 V = S.ap .* Q(S.iprev) + S.am .* Q(S.inext);
207end
209% The coefficients of dphi(u): i*m, diagonal.
210function V = dphic(Q)
211 S = sht_tables();
212 V = 1i * (S.mv .* Q);
213end
215% Grid-space derivatives, coefficients in: compositions of the shuffles and
216% the synthesis (src/sht/deriv.ts). dtheta divides by sin(theta) afterwards.
217function f = dtheta(Q)
218 S = sht_tables();
219 f = synth(dthetac(Q)) ./ S.stpt;
220end
222function f = dphi(Q)
223 f = synth(dphic(Q));
224end
226% Grid-space phi derivative, grid in: two FFT stages and a pointwise i*m,
227% no Legendre work -- d/dphi is diagonal in the Fourier index.
228function g = dphig(f)
229 S = sht_tables();
230 F = fft(reshape(f, S.nphi, S.nlat), [], 1);
231 g = real(ifft(S.dmul .* F, [], 1));
232 g = g(:);
233end
235% ---------------------------------------------------------------- the surface
236%
237% What the app precomputes from a shape's raw grid values: the band-limited
238% embedding and both metric formulations built on it (src/geom/geometry.ts,
239% src/geom/metric.ts). The solver runs on the synthesis of the coefficients,
240% not on the raw values -- for a shape with sharp features the two differ.
241function G = surface_tables(gxr, gyr, gzr)
242 S = sht_tables();
243 [G.Gx, G.Gy, G.Gz] = analys(gxr, gyr, gzr);
244 [G.gx, G.gy, G.gz] = synth(G.Gx, G.Gy, G.Gz);
246 % Flux-form metric weights, from the sin-weighted theta tangent
247 % sin(theta)*X_theta and X_phi, both smooth on the sphere:
248 % gtt~ = sin^2 g_tt, gtp~ = sin g_tp, D = J sin^2(theta).
249 [sXtx, sXty, sXtz] = synth(dthetac(G.Gx), dthetac(G.Gy), dthetac(G.Gz));
250 [Xpx, Xpy, Xpz] = synth(dphic(G.Gx), dphic(G.Gy), dphic(G.Gz));
251 gtt = sXtx.^2 + sXty.^2 + sXtz.^2;
252 gtp = sXtx.*Xpx + sXty.*Xpy + sXtz.*Xpz;
253 gpp = Xpx.^2 + Xpy.^2 + Xpz.^2;
254 D = sqrt(gtt .* gpp - gtp.^2);
255 G.p1 = gpp ./ D;
256 G.p2 = -gtp ./ D;
257 G.q2 = gtt ./ D;
258 G.r = 1 ./ D;
260 % The sphere-subtracted weights and the bounded 1/J = r sin^2(theta) --
261 % what keeps the concentrated division off the round sphere's share of the
262 % flux divergence. Formed here in f64, as the app forms them.
263 G.jinv = G.r .* S.stpt.^2;
264 G.dp1 = G.p1 - 1;
265 G.dq2 = G.q2 - 1;
267 % Preconditioner scale Jhat = 2/(muMin + muMax) over the eigenvalues of
268 % the operator's symbol S = (1/J)[[p1, p2], [p2, q2]].
269 s11 = G.p1 .* G.jinv;
270 s12 = G.p2 .* G.jinv;
271 s22 = G.q2 .* G.jinv;
272 mn = (s11 + s22) / 2;
273 disc = sqrt(((s11 - s22) / 2).^2 + s12.^2);
274 G.Jhat = 2 / (min(mn - disc) + max(mn + disc));
276 % Inverse metric quantities V_theta/V_phi, for the Algorithm-4 models.
277 Xtx = sXtx ./ S.stpt;
278 Xty = sXty ./ S.stpt;
279 Xtz = sXtz ./ S.stpt;
280 g11 = Xtx.^2 + Xty.^2 + Xtz.^2;
281 g12 = Xtx.*Xpx + Xty.*Xpy + Xtz.*Xpz;
282 g22 = gpp;
283 det = g11 .* g22 - g12.^2;
284 G.Vtx = (g22 .* Xtx - g12 .* Xpx) ./ det;
285 G.Vty = (g22 .* Xty - g12 .* Xpy) ./ det;
286 G.Vtz = (g22 .* Xtz - g12 .* Xpz) ./ det;
287 G.Vpx = (g11 .* Xpx - g12 .* Xtx) ./ det;
288 G.Vpy = (g11 .* Xpy - g12 .* Xty) ./ det;
289 G.Vpz = (g11 .* Xpz - g12 .* Xtz) ./ det;
290end
292% ---------------------------------------------------------------- random field
293%
294% chebfun-style smooth random field in 3D, restricted to the surface by
295% evaluating it at the grid points -- the way surfacefun seeds a run. Two
296% signatures, as in the app:
297% [k, c] = randnfun3(lambda, dom) the Fourier-mode draw (tools/randnfun3.m)
298% f = randnfun3(lambda, gx, gy, gz) that draw, summed at the surface points
299% Seed with rng(...) before calling.
300function varargout = randnfun3(lambda, varargin)
301 if nargin == 2
302 [k, c] = randnfun3_modes(lambda, varargin{1});
303 varargout = {k, c};
304 return;
305 end
306 [gx, gy, gz] = deal(varargin{1:3});
307 dom = [min(gx) max(gx) min(gy) max(gy) min(gz) max(gz)];
308 [k, c] = randnfun3_modes(lambda, dom);
309 % Summed in blocks of modes: the full npts x nmodes phase matrix can reach
310 % hundreds of MB at a fine wavelength.
311 f = zeros(numel(gx), 1);
312 blk = 2048;
313 for j0 = 1:blk:size(k, 1)
314 j1 = min(j0 + blk - 1, size(k, 1));
315 t = gx * k(j0:j1, 1)' + gy * k(j0:j1, 2)' + gz * k(j0:j1, 3)';
316 f = f + cos(t) * c(j0:j1, 1) - sin(t) * c(j0:j1, 2);
317 end
318 varargout = {f};
319end
321% ---------------------------------------------------------------- display
322%
323% The pattern on the surface, one panel per species. The solver grid has no
324% pole rows and an open phi seam; wrap_grid closes both for display, capping
325% each pole with the mean of its nearest ring.
326function h = plot_setup(gx, gy, gz, fields, names)
327 fig = figure('Name', 'turing-surface', 'Color', 'w');
328 Xs = wrap_grid(gx);
329 Ys = wrap_grid(gy);
330 Zs = wrap_grid(gz);
331 n = numel(fields);
332 h.surf = gobjects(1, n);
333 h.ax = gobjects(1, n);
334 for k = 1:n
335 ax = subplot(1, n, k, 'Parent', fig);
336 h.surf(k) = surf(ax, Xs, Ys, Zs, wrap_grid(fields{k}), 'EdgeColor', 'none');
337 shading(ax, 'interp');
338 axis(ax, 'equal');
339 axis(ax, 'off');
340 colormap(ax, 'jet');
341 colorbar(ax);
342 h.ax(k) = ax;
343 end
344 h.names = names;
345end
347function plot_update(h, fields, t, k, nsteps)
348 for i = 1:numel(fields)
349 C = wrap_grid(fields{i});
350 set(h.surf(i), 'CData', C);
351 lo = min(C(:));
352 hi = max(C(:));
353 if ~(hi > lo)
354 hi = lo + 1;
355 end
356 caxis(h.ax(i), [lo hi]);
357 title(h.ax(i), sprintf('%s t = %.3f (step %d/%d)', h.names{i}, t, k, nsteps));
358 end
359 drawnow;
360end
362function M = wrap_grid(f)
363 S = sht_tables();
364 M = reshape(f, S.nphi, S.nlat).';
365 M = [M, M(:, 1)];
366 M = [mean(M(1, :)) * ones(1, S.nphi + 1); M; mean(M(end, :)) * ones(1, S.nphi + 1)];
367end
369% ---------------------------------------------------------------- results file
370%
371% Complex coefficients -> flat float32 [re, im] per (l, m), the layout the
372% app's reference-file reader expects (docs/ellipsoid-reference-spec.md).
373function write_coeffs(fname, path, Q)
374 flat = zeros(2 * numel(Q), 1);
375 flat(1:2:end) = real(Q);
376 flat(2:2:end) = imag(Q);
377 h5create(fname, path, numel(flat), 'Datatype', 'single');
378 h5write(fname, path, single(flat));
379end
381% h5writeatt cannot create a bare group, so the attribute-only groups of the
382% reference layout are made through the low-level API.
383function make_group(fname, path)
384 fid = H5F.open(fname, 'H5F_ACC_RDWR', 'H5P_DEFAULT');
385 gid = H5G.create(fid, path, 'H5P_DEFAULT', 'H5P_DEFAULT', 'H5P_DEFAULT');
386 H5G.close(gid);
387 H5F.close(fid);
388end