forked from github.com/concept-collection/turing-surface
1/**
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 *
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
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 */
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';
42import { ShtPlan } from '../sht/sht.ts';
43import { relayoutCoeffs, type ShtConfig } from '../sht/layout.ts';
44import type { DerivPlan } from '../sht/deriv.ts';
45import type { ImportedCoeffs } from './registry.ts';
46import { computeMetric, computeFluxMetric } from './metric.ts';
47import { toolFiles } from '../tools.ts';
48import { inFunction, inModel, ModelCompileError } from '../mgpu/errors.ts';
49import type { ModelParams } from '../mgpu/model.ts';
51/** The function a geometry file must define. */
52export const SHAPE_FN = 'shape';
54export interface GeometryOptions {
55 /** The solver's transform plan — the grid the shape is evaluated on. */
56 sht: ShtPlan;
57 cfg: ShtConfig;
58 /** Geometry source (.m text). */
59 source: string;
60 /** Parameter names the .m may take beyond `theta` and `phi`. */
61 paramNames: string[];
62 params: ModelParams;
63 /** Computes the theta/phi derivatives the inverse metric quantities need. */
64 deriv: DerivPlan;
65 /**
66 * The surface as coefficients instead of a .m. When set, `source` is never
67 * evaluated: the coefficients are re-indexed onto this plan's band
68 * (truncation is exact projection, padding is exact) and everything from
69 * synthesis onward runs unchanged.
70 */
71 coeffs?: ImportedCoeffs;
72}
74export class Geometry {
75 /** Coordinates on the solver grid, npts each — synthesis of the coefficients. */
76 readonly x: Float32Array;
77 readonly y: Float32Array;
78 readonly z: Float32Array;
79 /** Their spherical-harmonic coefficients, 2 x nlm each. */
80 readonly X: Float32Array;
81 readonly Y: Float32Array;
82 readonly Z: Float32Array;
83 /**
84 * Inverse metric quantities (src/geom/metric.ts), grid space, npts each.
85 * Depend only on the geometry, so — like x,y,z,X,Y,Z above — these are a
86 * one-off computed here, not per-solve-step work. Used by the Algorithm-4
87 * (12-transform) Laplace-Beltrami path.
88 */
89 readonly Vtx: Float32Array;
90 readonly Vty: Float32Array;
91 readonly Vtz: Float32Array;
92 readonly Vpx: Float32Array;
93 readonly Vpy: Float32Array;
94 readonly Vpz: Float32Array;
95 /**
96 * Flux-form metric weights (src/geom/metric.ts computeFluxMetric), grid
97 * space, npts each — the six-transform Laplace-Beltrami scheme's
98 * replacement for the six V arrays (docs/reduced-transforms.md
99 * Sec 3). Both sets are carried so either operator formulation can run.
100 */
101 readonly p1: Float32Array;
102 readonly p2: Float32Array;
103 readonly q2: Float32Array;
104 readonly r: Float32Array;
105 /**
106 * The same flux weights with the round sphere subtracted off, plus the
107 * bounded 1/J — what lets a model evaluate lap_g without ever multiplying
108 * the *whole* flux divergence by r ~ 1/sin^2(theta). Writing p1 = 1 + dp1,
109 * q2 = 1 + dq2 (p2 is already a pure deviation, zero on the sphere) splits
110 * the divergence into a round-sphere part, whose cancelling bracket
111 * sin(theta) dtheta(A) + dphi(B) = -sin^2(theta) lap_s u is known exactly in
112 * spectral space, and a remainder:
113 *
114 * lap_g u = -jinv * lap_s u + r * (sin(theta) dtheta(P') + dphi(Q'))
115 *
116 * with P' = dp1*A + p2*B, Q' = p2*A + dq2*B. Only the remainder meets the
117 * concentrated division, so the polar roundoff gain drops by |P'|/|P|
118 * instead of applying to the full flux. Subtracting 1 in f64 here is the
119 * point: on a near-sphere dp1 is the small quantity, and forming it as an
120 * f32 difference in the .m would lose it. See docs/reduced-transforms.md
121 * Sec 5 and models/schnakenberg.m.
122 */
123 readonly dp1: Float32Array;
124 readonly dq2: Float32Array;
125 readonly jinv: Float32Array;
126 /**
127 * Preconditioner scale for the implicit solve (docs/reduced-transforms.md
128 * Sec 10). At high degree the Richardson iteration's per-mode factor is
129 * governed by the operator's principal symbol: in the orthonormal frame
130 * the surface symbol matrix is S = (1/J)[[p1, p2], [p2, q2]], whose
131 * eigenvalues mu(x) are the inverse squared principal stretches of the
132 * embedding — the round sphere has mu = 1. Preconditioning with lam/Jhat
133 * contracts every mode and every direction iff Jhat*mu stays in (0, 2),
134 * so the minimax constant is the harmonic mean of the symbol extremes,
135 *
136 * Jhat = 2/(muMin + muMax), rate = (muMax - muMin)/(muMax + muMin) < 1.
137 *
138 * The direction dependence is the point: a det-based mean of the area
139 * factor J (mu's geometric mean, exact only for conformal surfaces)
140 * under-corrects anisotropic stretching — on the shipped ellipsoid it
141 * leaves a band of directional high-degree modes with amplification > 1,
142 * which inflates the pattern's spectrum at moderate niter/lmax and
143 * diverges at larger ones. The plain scheme (Jhat = 1) diverges wherever
144 * muMax > 2. The solve's fixed point never depends on Jhat; only the
145 * convergence rate does.
146 */
147 readonly Jhat: number;
148 /** Symbol-eigenvalue range over the grid (see Jhat), for diagnostics. */
149 readonly muMin: number;
150 readonly muMax: number;
151 /** Area-factor range over the grid, for diagnostics. */
152 readonly Jmin: number;
153 readonly Jmax: number;
155 private constructor(init: {
156 x: Float32Array; y: Float32Array; z: Float32Array;
157 X: Float32Array; Y: Float32Array; Z: Float32Array;
158 Vtx: Float32Array; Vty: Float32Array; Vtz: Float32Array;
159 Vpx: Float32Array; Vpy: Float32Array; Vpz: Float32Array;
160 p1: Float32Array; p2: Float32Array; q2: Float32Array; r: Float32Array;
161 dp1: Float32Array; dq2: Float32Array; jinv: Float32Array;
162 Jhat: number; muMin: number; muMax: number; Jmin: number; Jmax: number;
163 }) {
164 this.x = init.x;
165 this.y = init.y;
166 this.z = init.z;
167 this.X = init.X;
168 this.Y = init.Y;
169 this.Z = init.Z;
170 this.Vtx = init.Vtx;
171 this.Vty = init.Vty;
172 this.Vtz = init.Vtz;
173 this.Vpx = init.Vpx;
174 this.Vpy = init.Vpy;
175 this.Vpz = init.Vpz;
176 this.p1 = init.p1;
177 this.p2 = init.p2;
178 this.q2 = init.q2;
179 this.r = init.r;
180 this.dp1 = init.dp1;
181 this.dq2 = init.dq2;
182 this.jinv = init.jinv;
183 this.Jhat = init.Jhat;
184 this.muMin = init.muMin;
185 this.muMax = init.muMax;
186 this.Jmin = init.Jmin;
187 this.Jmax = init.Jmax;
188 }
190 /**
191 * Evaluate the shape file once on the solver grid and reduce it to
192 * coefficients. Everything here happens at build time — a geometry never
193 * takes part in the timestep — so the .m runs on the CPU (see
194 * `evaluateShape`) and only the analysis onward touches the GPU.
195 */
196 static async create(opts: GeometryOptions): Promise<Geometry> {
197 const { sht, cfg, source, paramNames, params, deriv } = opts;
198 const npts = cfg.nlat * cfg.nphi;
200 // Coefficients first, then back to the grid: what the solver and the
201 // renderer both see is the band-limited surface, not the raw .m output.
202 // An imported surface arrives already as coefficients, so its .m stage is
203 // a re-indexing onto this plan's band rather than an evaluation.
204 let X: Float32Array;
205 let Y: Float32Array;
206 let Z: Float32Array;
207 if (opts.coeffs) {
208 const from = { lmax: opts.coeffs.lmax, mmax: opts.coeffs.mmax };
209 X = relayoutCoeffs(opts.coeffs.X, from, cfg);
210 Y = relayoutCoeffs(opts.coeffs.Y, from, cfg);
211 Z = relayoutCoeffs(opts.coeffs.Z, from, cfg);
212 } else {
213 const { theta, phi } = gridAngles(sht, cfg);
214 const raw = evaluateShape(source, paramNames, params, theta, phi, npts);
215 [X, Y, Z] = [
216 await sht.analys(raw[0]),
217 await sht.analys(raw[1]),
218 await sht.analys(raw[2]),
219 ];
220 }
221 const [x, y, z] = [
222 await sht.synth(X),
223 await sht.synth(Y),
224 await sht.synth(Z),
225 ];
227 // Inverse metric quantities (algos.tex Algorithm 2): theta/phi
228 // derivatives of the embedding's coefficients, contracted through the
229 // inverse first fundamental form. Depends only on the geometry, so
230 // this is a one-off alongside x,y,z above, not per-step work.
231 const Xt = await deriv.dtheta(X);
232 const Xp = await deriv.dphi(X);
233 const Yt = await deriv.dtheta(Y);
234 const Yp = await deriv.dphi(Y);
235 const Zt = await deriv.dtheta(Z);
236 const Zp = await deriv.dphi(Z);
237 const { Vtx, Vty, Vtz, Vpx, Vpy, Vpz } = computeMetric(npts, Xt, Xp, Yt, Yp, Zt, Zp);
239 // Flux-form metric weights for the six-transform scheme, built from the
240 // *undivided* theta tangents sin(theta)*X_theta (smooth on the sphere,
241 // unlike X_theta itself) and the same X_phi as above. Also a one-off;
242 // the f64 combination happens on the CPU, rounded to f32 for upload.
243 const sXtx = await deriv.sinDtheta(X);
244 const sXty = await deriv.sinDtheta(Y);
245 const sXtz = await deriv.sinDtheta(Z);
246 const flux = computeFluxMetric(npts, sXtx, sXty, sXtz, Xp, Yp, Zp);
248 // The preconditioner scale — see the Jhat field comment. The symbol
249 // matrix in the orthonormal frame is S = (1/J)[[p1,p2],[p2,q2]] with
250 // 1/J = r sin^2(theta); its entries are the bounded quantities
251 // g^tt, sin g^tp, sin^2 g^pp, so the eigenvalue extremes are clean to
252 // take over the grid. det S = 1/J^2, so the area factor comes along
253 // for free. f64 throughout.
254 let muMin = Infinity;
255 let muMax = 0;
256 let Jmin = Infinity;
257 let Jmax = 0;
258 // The sphere-subtracted weights ride along on this loop: 1/J is already
259 // being formed here, and dp1/dq2 want the same f64 arithmetic.
260 const dp1 = new Float32Array(npts);
261 const dq2 = new Float32Array(npts);
262 const jinv = new Float32Array(npts);
263 for (let i = 0; i < cfg.nlat; i++) {
264 const ct = sht.cosTheta[i];
265 const st2 = Math.max(0, 1 - ct * ct);
266 for (let j = 0; j < cfg.nphi; j++) {
267 const k = i * cfg.nphi + j;
268 const invJ = flux.r[k] * st2;
269 dp1[k] = flux.p1[k] - 1;
270 dq2[k] = flux.q2[k] - 1;
271 jinv[k] = invJ;
272 const s11 = flux.p1[k] * invJ;
273 const s12 = flux.p2[k] * invJ;
274 const s22 = flux.q2[k] * invJ;
275 const mean = (s11 + s22) / 2;
276 const disc = Math.sqrt(((s11 - s22) / 2) ** 2 + s12 * s12);
277 if (mean - disc < muMin) muMin = mean - disc;
278 if (mean + disc > muMax) muMax = mean + disc;
279 const J = 1 / invJ;
280 if (J < Jmin) Jmin = J;
281 if (J > Jmax) Jmax = J;
282 }
283 }
284 const Jhat = 2 / (muMin + muMax);
286 return new Geometry({
287 x, y, z, X, Y, Z, Vtx, Vty, Vtz, Vpx, Vpy, Vpz,
288 p1: new Float32Array(flux.p1),
289 p2: new Float32Array(flux.p2),
290 q2: new Float32Array(flux.q2),
291 r: new Float32Array(flux.r),
292 dp1, dq2, jinv,
293 Jhat, muMin, muMax, Jmin, Jmax,
294 });
295 }
297 /**
298 * The surface evaluated on another plan's grid, as interleaved xyz vertex
299 * positions (nlat * nphi * 3) — for rendering at display resolution. Exact
300 * interpolation: the same coefficients, more evaluation points.
301 */
302 async positionsOn(view: ShtPlan): Promise<Float32Array> {
303 const [x, y, z] = [
304 await view.synth(this.X),
305 await view.synth(this.Y),
306 await view.synth(this.Z),
307 ];
308 const out = new Float32Array(x.length * 3);
309 for (let i = 0; i < x.length; i++) {
310 out[3 * i] = x[i];
311 out[3 * i + 1] = y[i];
312 out[3 * i + 2] = z[i];
313 }
314 return out;
315 }
317 /** How far the surface departs from the unit sphere, as min/max radius. */
318 radiusRange(): { lo: number; hi: number } {
319 let lo = Infinity;
320 let hi = -Infinity;
321 for (let i = 0; i < this.x.length; i++) {
322 const r = Math.hypot(this.x[i], this.y[i], this.z[i]);
323 if (r < lo) lo = r;
324 if (r > hi) hi = r;
325 }
326 return { lo, hi };
327 }
328}
330/** The (theta, phi) of every grid point, flattened phi-fastest as the fields
331 * are — in f64, the precision the shape is evaluated at. */
332function gridAngles(
333 sht: ShtPlan,
334 cfg: ShtConfig,
335): { theta: Float64Array; phi: Float64Array } {
336 const { nlat, nphi } = cfg;
337 const theta = new Float64Array(nlat * nphi);
338 const phi = new Float64Array(nlat * nphi);
339 for (let i = 0; i < nlat; i++) {
340 const th = Math.acos(Math.max(-1, Math.min(1, sht.cosTheta[i])));
341 for (let j = 0; j < nphi; j++) {
342 theta[i * nphi + j] = th;
343 phi[i * nphi + j] = (2 * Math.PI * j) / nphi;
344 }
345 }
346 return { theta, phi };
347}
349/**
350 * Evaluate the shape file on the grid, through numbl's CPU interpreter.
351 *
352 * The .m keeps the same contract it had as a compiled model: it names the
353 * arguments it wants — `theta`, `phi`, and any of the registry's parameters —
354 * and the host supplies them by name, so their order in the signature is the
355 * .m's own business. A one-line driver script calls `shape` with exactly the
356 * arguments its signature declares, with those names pre-bound in the
357 * driver's workspace.
358 */
359function evaluateShape(
360 source: string,
361 paramNames: string[],
362 params: ModelParams,
363 theta: Float64Array,
364 phi: Float64Array,
365 npts: number,
366): [Float32Array, Float32Array, Float32Array] {
367 const file = `${SHAPE_FN}.m`;
368 const ast = inModel(() => parseMFile(source, file));
369 const fn = ast.body.find(
370 (s): s is FunctionStmt =>
371 s.type === 'Function' && (s as FunctionStmt).name === SHAPE_FN,
372 );
373 if (!fn) {
374 throw new ModelCompileError(
375 `the geometry defines no function named '${SHAPE_FN}'`,
376 );
377 }
378 if (fn.outputs.length !== 3) {
379 throw new ModelCompileError(
380 `'${SHAPE_FN}' must return three outputs [gx, gy, gz], not ${fn.outputs.length}`,
381 { fn: SHAPE_FN, start: fn.span.start, end: fn.span.end },
382 );
383 }
384 const known = new Set(['theta', 'phi', ...paramNames]);
385 for (const p of fn.params) {
386 if (!known.has(p)) {
387 throw new ModelCompileError(
388 `'${SHAPE_FN}' takes an argument '${p}' that is neither the grid ` +
389 `(theta, phi) nor one of this geometry's parameters` +
390 (paramNames.length ? ` (${paramNames.join(', ')})` : ''),
391 { fn: SHAPE_FN, start: fn.span.start, end: fn.span.end },
392 );
393 }
394 }
396 const vars: Record<string, RuntimeValue> = {
397 theta: new RuntimeTensor(theta, [npts, 1]),
398 phi: new RuntimeTensor(phi, [npts, 1]),
399 };
400 for (const name of paramNames) {
401 const v = params[name];
402 // Missing parameters read as 0, as ModelPlan.setParams has it.
403 vars[name] = Number.isFinite(v) ? v : 0;
404 }
406 const driver = `[gx__, gy__, gz__] = ${SHAPE_FN}(${fn.params.join(', ')});`;
407 const result = inFunction(SHAPE_FN, () =>
408 executeCode(
409 driver,
410 { initialVariableValues: vars, displayResults: false, implicitCwdPath: null },
411 [...toolFiles, { name: file, source }],
412 'geometry-driver.m',
413 ),
414 );
416 return [
417 toGridField(result.variableValues['gx__'], fn.outputs[0], npts),
418 toGridField(result.variableValues['gy__'], fn.outputs[1], npts),
419 toGridField(result.variableValues['gz__'], fn.outputs[2], npts),
420 ];
421}
423/** One returned coordinate → npts values, rounded to the transforms' f32. */
424function toGridField(
425 value: RuntimeValue | undefined,
426 name: string,
427 npts: number,
428): Float32Array {
429 // A constant coordinate stays scalar in MATLAB; spread it over the grid.
430 if (typeof value === 'number') return new Float32Array(npts).fill(value);
431 if (value !== undefined && isRuntimeTensor(value)) {
432 if (value.imag) {
433 throw new ModelCompileError(
434 `the geometry's '${name}' is complex; coordinates must be real`,
435 { fn: SHAPE_FN },
436 );
437 }
438 // A vector of npts values, either orientation. A 2-D reshape is refused
439 // rather than reordered: the tensor's column-major layout would not match
440 // the grid's phi-fastest rows.
441 if (value.data.length === npts && value.shape.every((d) => d === 1 || d === npts)) {
442 return new Float32Array(value.data);
443 }
444 throw new ModelCompileError(
445 `the geometry's '${name}' is ${value.shape.join(' x ')}, but the grid ` +
446 `wants one value per point (${npts} x 1)`,
447 { fn: SHAPE_FN },
448 );
449 }
450 throw new ModelCompileError(`the geometry's '${name}' is not numeric`, {
451 fn: SHAPE_FN,
452 });
453}