2 * The surface: a .m shape file, compiled and evaluated into spherical-harmonic
3 * coefficients.
4 *
5 * A geometry file is ordinary MATLAB defining one function,
6 *
7 * function [gx, gy, gz] = shape(theta, phi, <parameters>)
8 *
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 9 * over the solver's (theta, phi) grid. Unlike the models it is *not* compiled
10 * to WGSL: a model's step runs every frame and must lower to a fixed sequence
11 * of GPU dispatches, but a shape is evaluated exactly once at build time and
12 * survives only as coefficients. So it runs through numbl's CPU interpreter
13 * instead, which buys the full MATLAB subset — loops, arrays, reductions,
14 * `legendre`, seeded randomness via `rng`/`randn` — and f64 evaluation, where
15 * the step dialect is element-wise f32. The result is then *analysed*: the
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 16 * canonical geometry this project carries is the three sets of coefficients
17 * `X`, `Y`, `Z`, one per Cartesian component of the embedding.
18 *
19 * Going through the coefficients rather than keeping the pointwise values is
20 * what makes the geometry usable by a spectral method, for two reasons:
21 *
22 * - it is exactly band-limited at lmax afterwards, so the surface has as many
23 * derivatives as the scheme needs and no aliased content the solver cannot
24 * see. `x`, `y`, `z` below are the synthesis of the coefficients, not the
25 * raw output of the .m — the shape actually being solved on, which for a
26 * shape with sharp features is not quite the shape that was written down.
27 * - it can be evaluated on any grid. The renderer draws the surface on the
28 * (possibly finer) display grid by synthesizing the same coefficients
29 * there, which is exact interpolation rather than subdivision — the same
30 * argument that lets the species fields be oversampled.
31 *
32 * The unit sphere is the case where `x`, `y`, `z` are pure degree-1 harmonics
33 * and everything downstream reduces to turing-sphere.
34 */
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 35import { parseMFile, type FunctionStmt } from 'numbl-src/numbl-core/parser/index.ts';
36import { executeCode } from 'numbl-src/numbl-core/executeCode.ts';
37import {
38 RuntimeTensor,
39 isRuntimeTensor,
40 type RuntimeValue,
41} from 'numbl-src/numbl-core/runtime/types.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 42import { ShtPlan } from '../sht/sht.ts';
43import type { ShtConfig } from '../sht/layout.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 44import type { DerivPlan } from '../sht/deriv.ts';
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 45import { computeMetric, computeFluxMetric } from './metric.ts';
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 46import { toolFiles } from '../tools.ts';
47import { inFunction, inModel, ModelCompileError } from '../mgpu/errors.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 48import type { ModelParams } from '../mgpu/model.ts';
50/** The function a geometry file must define. */
51export const SHAPE_FN = 'shape';
53export interface GeometryOptions {
54 /** The solver's transform plan — the grid the shape is evaluated on. */
55 sht: ShtPlan;
56 cfg: ShtConfig;
57 /** Geometry source (.m text). */
58 source: string;
59 /** Parameter names the .m may take beyond `theta` and `phi`. */
60 paramNames: string[];
61 params: ModelParams;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 62 /** Computes the theta/phi derivatives the inverse metric quantities need. */
63 deriv: DerivPlan;
66export class Geometry {
67 /** Coordinates on the solver grid, npts each — synthesis of the coefficients. */
68 readonly x: Float32Array;
69 readonly y: Float32Array;
70 readonly z: Float32Array;
71 /** Their spherical-harmonic coefficients, 2 x nlm each. */
72 readonly X: Float32Array;
73 readonly Y: Float32Array;
74 readonly Z: Float32Array;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 75 /**
76 * Inverse metric quantities (src/geom/metric.ts), grid space, npts each.
77 * Depend only on the geometry, so — like x,y,z,X,Y,Z above — these are a
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 78 * one-off computed here, not per-solve-step work. Used by the Algorithm-4
79 * (12-transform) Laplace-Beltrami path.
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 80 */
81 readonly Vtx: Float32Array;
82 readonly Vty: Float32Array;
83 readonly Vtz: Float32Array;
84 readonly Vpx: Float32Array;
85 readonly Vpy: Float32Array;
86 readonly Vpz: Float32Array;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 87 /**
88 * Flux-form metric weights (src/geom/metric.ts computeFluxMetric), grid
89 * space, npts each — the six-transform Laplace-Beltrami scheme's
90 * replacement for the six V arrays (docs/reduced-transforms.md
91 * Sec 3). Both sets are carried so either operator formulation can run.
92 */
93 readonly p1: Float32Array;
94 readonly p2: Float32Array;
95 readonly q2: Float32Array;
96 readonly r: Float32Array;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 97 /**
98 * Preconditioner scale for the implicit solve (docs/reduced-transforms.md
99 * Sec 10). At high degree the Richardson iteration's per-mode factor is
100 * governed by the operator's principal symbol: in the orthonormal frame
101 * the surface symbol matrix is S = (1/J)[[p1, p2], [p2, q2]], whose
102 * eigenvalues mu(x) are the inverse squared principal stretches of the
103 * embedding — the round sphere has mu = 1. Preconditioning with lam/Jhat
104 * contracts every mode and every direction iff Jhat*mu stays in (0, 2),
105 * so the minimax constant is the harmonic mean of the symbol extremes,
106 *
107 * Jhat = 2/(muMin + muMax), rate = (muMax - muMin)/(muMax + muMin) < 1.
108 *
109 * The direction dependence is the point: a det-based mean of the area
110 * factor J (mu's geometric mean, exact only for conformal surfaces)
111 * under-corrects anisotropic stretching — on the shipped ellipsoid it
112 * leaves a band of directional high-degree modes with amplification > 1,
113 * which inflates the pattern's spectrum at moderate niter/lmax and
114 * diverges at larger ones. The plain scheme (Jhat = 1) diverges wherever
115 * muMax > 2. The solve's fixed point never depends on Jhat; only the
116 * convergence rate does.
117 */
118 readonly Jhat: number;
119 /** Symbol-eigenvalue range over the grid (see Jhat), for diagnostics. */
120 readonly muMin: number;
121 readonly muMax: number;
122 /** Area-factor range over the grid, for diagnostics. */
123 readonly Jmin: number;
124 readonly Jmax: number;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 125
126 private constructor(init: {
127 x: Float32Array; y: Float32Array; z: Float32Array;
128 X: Float32Array; Y: Float32Array; Z: Float32Array;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 129 Vtx: Float32Array; Vty: Float32Array; Vtz: Float32Array;
130 Vpx: Float32Array; Vpy: Float32Array; Vpz: Float32Array;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 131 p1: Float32Array; p2: Float32Array; q2: Float32Array; r: Float32Array;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 132 Jhat: number; muMin: number; muMax: number; Jmin: number; Jmax: number;
134 this.x = init.x;
135 this.y = init.y;
136 this.z = init.z;
137 this.X = init.X;
138 this.Y = init.Y;
139 this.Z = init.Z;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 140 this.Vtx = init.Vtx;
141 this.Vty = init.Vty;
142 this.Vtz = init.Vtz;
143 this.Vpx = init.Vpx;
144 this.Vpy = init.Vpy;
145 this.Vpz = init.Vpz;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 146 this.p1 = init.p1;
147 this.p2 = init.p2;
148 this.q2 = init.q2;
149 this.r = init.r;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 150 this.Jhat = init.Jhat;
151 this.muMin = init.muMin;
152 this.muMax = init.muMax;
153 this.Jmin = init.Jmin;
154 this.Jmax = init.Jmax;
157 /**
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 158 * Evaluate the shape file once on the solver grid and reduce it to
159 * coefficients. Everything here happens at build time — a geometry never
160 * takes part in the timestep — so the .m runs on the CPU (see
161 * `evaluateShape`) and only the analysis onward touches the GPU.
163 static async create(opts: GeometryOptions): Promise<Geometry> {
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 164 const { sht, cfg, source, paramNames, params, deriv } = opts;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 165 const npts = cfg.nlat * cfg.nphi;
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 167 const { theta, phi } = gridAngles(sht, cfg);
168 const raw = evaluateShape(source, paramNames, params, theta, phi, npts);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 169
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 170 // Coefficients first, then back to the grid: what the solver and the
171 // renderer both see is the band-limited surface, not the raw .m output.
172 const [X, Y, Z] = [
173 await sht.analys(raw[0]),
174 await sht.analys(raw[1]),
175 await sht.analys(raw[2]),
176 ];
177 const [x, y, z] = [
178 await sht.synth(X),
179 await sht.synth(Y),
180 await sht.synth(Z),
181 ];
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 182
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 183 // Inverse metric quantities (algos.tex Algorithm 2): theta/phi
184 // derivatives of the embedding's coefficients, contracted through the
185 // inverse first fundamental form. Depends only on the geometry, so
186 // this is a one-off alongside x,y,z above, not per-step work.
187 const Xt = await deriv.dtheta(X);
188 const Xp = await deriv.dphi(X);
189 const Yt = await deriv.dtheta(Y);
190 const Yp = await deriv.dphi(Y);
191 const Zt = await deriv.dtheta(Z);
192 const Zp = await deriv.dphi(Z);
193 const { Vtx, Vty, Vtz, Vpx, Vpy, Vpz } = computeMetric(npts, Xt, Xp, Yt, Yp, Zt, Zp);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 194
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 195 // Flux-form metric weights for the six-transform scheme, built from the
196 // *undivided* theta tangents sin(theta)*X_theta (smooth on the sphere,
197 // unlike X_theta itself) and the same X_phi as above. Also a one-off;
198 // the f64 combination happens on the CPU, rounded to f32 for upload.
199 const sXtx = await deriv.sinDtheta(X);
200 const sXty = await deriv.sinDtheta(Y);
201 const sXtz = await deriv.sinDtheta(Z);
202 const flux = computeFluxMetric(npts, sXtx, sXty, sXtz, Xp, Yp, Zp);
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 203
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 204 // The preconditioner scale — see the Jhat field comment. The symbol
205 // matrix in the orthonormal frame is S = (1/J)[[p1,p2],[p2,q2]] with
206 // 1/J = r sin^2(theta); its entries are the bounded quantities
207 // g^tt, sin g^tp, sin^2 g^pp, so the eigenvalue extremes are clean to
208 // take over the grid. det S = 1/J^2, so the area factor comes along
209 // for free. f64 throughout.
210 let muMin = Infinity;
211 let muMax = 0;
212 let Jmin = Infinity;
213 let Jmax = 0;
214 for (let i = 0; i < cfg.nlat; i++) {
215 const ct = sht.cosTheta[i];
216 const st2 = Math.max(0, 1 - ct * ct);
217 for (let j = 0; j < cfg.nphi; j++) {
218 const k = i * cfg.nphi + j;
219 const invJ = flux.r[k] * st2;
220 const s11 = flux.p1[k] * invJ;
221 const s12 = flux.p2[k] * invJ;
222 const s22 = flux.q2[k] * invJ;
223 const mean = (s11 + s22) / 2;
224 const disc = Math.sqrt(((s11 - s22) / 2) ** 2 + s12 * s12);
225 if (mean - disc < muMin) muMin = mean - disc;
226 if (mean + disc > muMax) muMax = mean + disc;
227 const J = 1 / invJ;
228 if (J < Jmin) Jmin = J;
229 if (J > Jmax) Jmax = J;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 230 }
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 232 const Jhat = 2 / (muMin + muMax);
234 return new Geometry({
235 x, y, z, X, Y, Z, Vtx, Vty, Vtz, Vpx, Vpy, Vpz,
236 p1: new Float32Array(flux.p1),
237 p2: new Float32Array(flux.p2),
238 q2: new Float32Array(flux.q2),
239 r: new Float32Array(flux.r),
240 Jhat, muMin, muMax, Jmin, Jmax,
241 });
244 /**
245 * The surface evaluated on another plan's grid, as interleaved xyz vertex
246 * positions (nlat * nphi * 3) — for rendering at display resolution. Exact
247 * interpolation: the same coefficients, more evaluation points.
248 */
249 async positionsOn(view: ShtPlan): Promise<Float32Array> {
250 const [x, y, z] = [
251 await view.synth(this.X),
252 await view.synth(this.Y),
253 await view.synth(this.Z),
254 ];
255 const out = new Float32Array(x.length * 3);
256 for (let i = 0; i < x.length; i++) {
257 out[3 * i] = x[i];
258 out[3 * i + 1] = y[i];
259 out[3 * i + 2] = z[i];
260 }
261 return out;
262 }
264 /** How far the surface departs from the unit sphere, as min/max radius. */
265 radiusRange(): { lo: number; hi: number } {
266 let lo = Infinity;
267 let hi = -Infinity;
268 for (let i = 0; i < this.x.length; i++) {
269 const r = Math.hypot(this.x[i], this.y[i], this.z[i]);
270 if (r < lo) lo = r;
271 if (r > hi) hi = r;
272 }
273 return { lo, hi };
274 }
275}
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 277/** The (theta, phi) of every grid point, flattened phi-fastest as the fields
278 * are — in f64, the precision the shape is evaluated at. */
280 sht: ShtPlan,
281 cfg: ShtConfig,
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 282): { theta: Float64Array; phi: Float64Array } {
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 283 const { nlat, nphi } = cfg;
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 284 const theta = new Float64Array(nlat * nphi);
285 const phi = new Float64Array(nlat * nphi);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 286 for (let i = 0; i < nlat; i++) {
287 const th = Math.acos(Math.max(-1, Math.min(1, sht.cosTheta[i])));
288 for (let j = 0; j < nphi; j++) {
289 theta[i * nphi + j] = th;
290 phi[i * nphi + j] = (2 * Math.PI * j) / nphi;
291 }
292 }
293 return { theta, phi };
294}
297 * Evaluate the shape file on the grid, through numbl's CPU interpreter.
298 *
299 * The .m keeps the same contract it had as a compiled model: it names the
300 * arguments it wants — `theta`, `phi`, and any of the registry's parameters —
301 * and the host supplies them by name, so their order in the signature is the
302 * .m's own business. A one-line driver script calls `shape` with exactly the
303 * arguments its signature declares, with those names pre-bound in the
304 * driver's workspace.
305 */
306function evaluateShape(
307 source: string,
308 paramNames: string[],
309 params: ModelParams,
310 theta: Float64Array,
311 phi: Float64Array,
312 npts: number,
313): [Float32Array, Float32Array, Float32Array] {
314 const file = `${SHAPE_FN}.m`;
315 const ast = inModel(() => parseMFile(source, file));
316 const fn = ast.body.find(
317 (s): s is FunctionStmt =>
318 s.type === 'Function' && (s as FunctionStmt).name === SHAPE_FN,
319 );
320 if (!fn) {
321 throw new ModelCompileError(
322 `the geometry defines no function named '${SHAPE_FN}'`,
323 );
324 }
325 if (fn.outputs.length !== 3) {
326 throw new ModelCompileError(
327 `'${SHAPE_FN}' must return three outputs [gx, gy, gz], not ${fn.outputs.length}`,
328 { fn: SHAPE_FN, start: fn.span.start, end: fn.span.end },
329 );
330 }
331 const known = new Set(['theta', 'phi', ...paramNames]);
332 for (const p of fn.params) {
333 if (!known.has(p)) {
334 throw new ModelCompileError(
335 `'${SHAPE_FN}' takes an argument '${p}' that is neither the grid ` +
336 `(theta, phi) nor one of this geometry's parameters` +
337 (paramNames.length ? ` (${paramNames.join(', ')})` : ''),
338 { fn: SHAPE_FN, start: fn.span.start, end: fn.span.end },
339 );
340 }
341 }
343 const vars: Record<string, RuntimeValue> = {
344 theta: new RuntimeTensor(theta, [npts, 1]),
345 phi: new RuntimeTensor(phi, [npts, 1]),
346 };
347 for (const name of paramNames) {
348 const v = params[name];
349 // Missing parameters read as 0, as ModelPlan.setParams has it.
350 vars[name] = Number.isFinite(v) ? v : 0;
351 }
353 const driver = `[gx__, gy__, gz__] = ${SHAPE_FN}(${fn.params.join(', ')});`;
354 const result = inFunction(SHAPE_FN, () =>
355 executeCode(
356 driver,
357 { initialVariableValues: vars, displayResults: false, implicitCwdPath: null },
358 [...toolFiles, { name: file, source }],
359 'geometry-driver.m',
360 ),
361 );
363 return [
364 toGridField(result.variableValues['gx__'], fn.outputs[0], npts),
365 toGridField(result.variableValues['gy__'], fn.outputs[1], npts),
366 toGridField(result.variableValues['gz__'], fn.outputs[2], npts),
367 ];
368}
370/** One returned coordinate → npts values, rounded to the transforms' f32. */
371function toGridField(
372 value: RuntimeValue | undefined,
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 374 npts: number,
375): Float32Array {
376 // A constant coordinate stays scalar in MATLAB; spread it over the grid.
377 if (typeof value === 'number') return new Float32Array(npts).fill(value);
378 if (value !== undefined && isRuntimeTensor(value)) {
379 if (value.imag) {
380 throw new ModelCompileError(
381 `the geometry's '${name}' is complex; coordinates must be real`,
382 { fn: SHAPE_FN },
383 );
384 }
385 // A vector of npts values, either orientation. A 2-D reshape is refused
386 // rather than reordered: the tensor's column-major layout would not match
387 // the grid's phi-fastest rows.
388 if (value.data.length === npts && value.shape.every((d) => d === 1 || d === npts)) {
389 return new Float32Array(value.data);
390 }
391 throw new ModelCompileError(
392 `the geometry's '${name}' is ${value.shape.join(' x ')}, but the grid ` +
393 `wants one value per point (${npts} x 1)`,
394 { fn: SHAPE_FN },
395 );
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 397 throw new ModelCompileError(`the geometry's '${name}' is not numeric`, {
398 fn: SHAPE_FN,
400}