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 — the same element-wise MATLAB the models
10 * are written in, compiled by the same backend into the same kind of WGSL
11 * kernel. It is evaluated once, on the CPU's behalf, and then *analysed*: the
12 * canonical geometry this project carries is the three sets of coefficients
13 * `X`, `Y`, `Z`, one per Cartesian component of the embedding.
14 *
15 * Going through the coefficients rather than keeping the pointwise values is
16 * what makes the geometry usable by a spectral method, for two reasons:
17 *
18 * - it is exactly band-limited at lmax afterwards, so the surface has as many
19 * derivatives as the scheme needs and no aliased content the solver cannot
20 * see. `x`, `y`, `z` below are the synthesis of the coefficients, not the
21 * raw output of the .m — the shape actually being solved on, which for a
22 * shape with sharp features is not quite the shape that was written down.
23 * - it can be evaluated on any grid. The renderer draws the surface on the
24 * (possibly finer) display grid by synthesizing the same coefficients
25 * there, which is exact interpolation rather than subdivision — the same
26 * argument that lets the species fields be oversampled.
27 *
28 * The unit sphere is the case where `x`, `y`, `z` are pure degree-1 harmonics
29 * and everything downstream reduces to turing-sphere.
30 */
31import { ShtPlan } from '../sht/sht.ts';
32import type { ShtConfig } from '../sht/layout.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 33import type { DerivPlan } from '../sht/deriv.ts';
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 34import { computeMetric, computeFluxMetric } from './metric.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 35import { HostBuffers, ModelPlan } from '../mgpu/plan.ts';
36import { CompiledModel, type Binding } from '../mgpu/compile.ts';
37import { inFunction, inFunctionAsync, inModel } from '../mgpu/errors.ts';
38import type { ModelParams } from '../mgpu/model.ts';
40/** The function a geometry file must define. */
41export const SHAPE_FN = 'shape';
43export interface GeometryOptions {
44 device: GPUDevice;
45 /** The solver's transform plan — the grid the shape is evaluated on. */
46 sht: ShtPlan;
47 cfg: ShtConfig;
48 /** Geometry source (.m text). */
49 source: string;
50 /** Parameter names the .m may take beyond `theta` and `phi`. */
51 paramNames: string[];
52 params: ModelParams;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 53 /** Computes the theta/phi derivatives the inverse metric quantities need. */
54 deriv: DerivPlan;
57export class Geometry {
58 /** Coordinates on the solver grid, npts each — synthesis of the coefficients. */
59 readonly x: Float32Array;
60 readonly y: Float32Array;
61 readonly z: Float32Array;
62 /** Their spherical-harmonic coefficients, 2 x nlm each. */
63 readonly X: Float32Array;
64 readonly Y: Float32Array;
65 readonly Z: Float32Array;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 66 /**
67 * Inverse metric quantities (src/geom/metric.ts), grid space, npts each.
68 * 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 69 * one-off computed here, not per-solve-step work. Used by the Algorithm-4
70 * (12-transform) Laplace-Beltrami path.
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 71 */
72 readonly Vtx: Float32Array;
73 readonly Vty: Float32Array;
74 readonly Vtz: Float32Array;
75 readonly Vpx: Float32Array;
76 readonly Vpy: Float32Array;
77 readonly Vpz: Float32Array;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 78 /**
79 * Flux-form metric weights (src/geom/metric.ts computeFluxMetric), grid
80 * space, npts each — the six-transform Laplace-Beltrami scheme's
81 * replacement for the six V arrays (docs/reduced-transforms.md
82 * Sec 3). Both sets are carried so either operator formulation can run.
83 */
84 readonly p1: Float32Array;
85 readonly p2: Float32Array;
86 readonly q2: Float32Array;
87 readonly r: Float32Array;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 88 /**
89 * Preconditioner scale for the implicit solve (docs/reduced-transforms.md
90 * Sec 10). At high degree the Richardson iteration's per-mode factor is
91 * governed by the operator's principal symbol: in the orthonormal frame
92 * the surface symbol matrix is S = (1/J)[[p1, p2], [p2, q2]], whose
93 * eigenvalues mu(x) are the inverse squared principal stretches of the
94 * embedding — the round sphere has mu = 1. Preconditioning with lam/Jhat
95 * contracts every mode and every direction iff Jhat*mu stays in (0, 2),
96 * so the minimax constant is the harmonic mean of the symbol extremes,
97 *
98 * Jhat = 2/(muMin + muMax), rate = (muMax - muMin)/(muMax + muMin) < 1.
99 *
100 * The direction dependence is the point: a det-based mean of the area
101 * factor J (mu's geometric mean, exact only for conformal surfaces)
102 * under-corrects anisotropic stretching — on the shipped ellipsoid it
103 * leaves a band of directional high-degree modes with amplification > 1,
104 * which inflates the pattern's spectrum at moderate niter/lmax and
105 * diverges at larger ones. The plain scheme (Jhat = 1) diverges wherever
106 * muMax > 2. The solve's fixed point never depends on Jhat; only the
107 * convergence rate does.
108 */
109 readonly Jhat: number;
110 /** Symbol-eigenvalue range over the grid (see Jhat), for diagnostics. */
111 readonly muMin: number;
112 readonly muMax: number;
113 /** Area-factor range over the grid, for diagnostics. */
114 readonly Jmin: number;
115 readonly Jmax: number;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 116
117 private constructor(init: {
118 x: Float32Array; y: Float32Array; z: Float32Array;
119 X: Float32Array; Y: Float32Array; Z: Float32Array;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 120 Vtx: Float32Array; Vty: Float32Array; Vtz: Float32Array;
121 Vpx: Float32Array; Vpy: Float32Array; Vpz: Float32Array;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 122 p1: Float32Array; p2: Float32Array; q2: Float32Array; r: Float32Array;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 123 Jhat: number; muMin: number; muMax: number; Jmin: number; Jmax: number;
125 this.x = init.x;
126 this.y = init.y;
127 this.z = init.z;
128 this.X = init.X;
129 this.Y = init.Y;
130 this.Z = init.Z;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 131 this.Vtx = init.Vtx;
132 this.Vty = init.Vty;
133 this.Vtz = init.Vtz;
134 this.Vpx = init.Vpx;
135 this.Vpy = init.Vpy;
136 this.Vpz = init.Vpz;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 137 this.p1 = init.p1;
138 this.p2 = init.p2;
139 this.q2 = init.q2;
140 this.r = init.r;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 141 this.Jhat = init.Jhat;
142 this.muMin = init.muMin;
143 this.muMax = init.muMax;
144 this.Jmin = init.Jmin;
145 this.Jmax = init.Jmax;
148 /**
149 * Compile the shape file, evaluate it once on the solver grid, and reduce it
150 * to coefficients. Everything here happens at build time — a geometry never
151 * takes part in the timestep — so it reads back through the CPU freely.
152 */
153 static async create(opts: GeometryOptions): Promise<Geometry> {
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 154 const { device, sht, cfg, source, paramNames, params, deriv } = opts;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 155 const npts = cfg.nlat * cfg.nphi;
156 const nlm = sht.nlm;
158 const bindings: Record<string, Binding> = {
159 theta: { kind: 'tensor', shape: [npts, 1] },
160 phi: { kind: 'tensor', shape: [npts, 1] },
161 npts: { kind: 'const', value: npts },
162 };
163 for (const p of paramNames) bindings[p] = { kind: 'param' };
165 const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
166 const fn = inFunction(SHAPE_FN, () => compiled.specialize(SHAPE_FN, 3));
167 compiled.finish();
169 const host = new HostBuffers(device);
170 host.ensure('theta', npts);
171 host.ensure('phi', npts);
173 const plan = await inFunctionAsync(SHAPE_FN, () =>
174 // Nothing feeds back: the three outputs are read once and the plan is
175 // thrown away.
176 ModelPlan.create(device, sht, { fn, feedback: [null, null, null] }, host),
177 );
179 try {
180 const { theta, phi } = gridAngles(sht, cfg);
181 host.upload('theta', theta);
182 host.upload('phi', phi);
183 plan.setParams(params);
185 const enc = device.createCommandEncoder({ label: 'geometry-shape' });
186 plan.encodeSteps(enc, 1);
187 device.queue.submit([enc.finish()]);
189 const raw = await Promise.all(
190 fn.outputs.map((out) => readBuffer(device, plan, out.name, npts)),
191 );
192 // Coefficients first, then back to the grid: what the solver and the
193 // renderer both see is the band-limited surface, not the raw .m output.
194 const [X, Y, Z] = [
195 await sht.analys(raw[0]),
196 await sht.analys(raw[1]),
197 await sht.analys(raw[2]),
198 ];
199 const [x, y, z] = [
200 await sht.synth(X),
201 await sht.synth(Y),
202 await sht.synth(Z),
203 ];
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 204
205 // Inverse metric quantities (algos.tex Algorithm 2): theta/phi
206 // derivatives of the embedding's coefficients, contracted through the
207 // inverse first fundamental form. Depends only on the geometry, so
208 // this is a one-off alongside x,y,z above, not per-step work.
209 const Xt = await deriv.dtheta(X);
210 const Xp = await deriv.dphi(X);
211 const Yt = await deriv.dtheta(Y);
212 const Yp = await deriv.dphi(Y);
213 const Zt = await deriv.dtheta(Z);
214 const Zp = await deriv.dphi(Z);
215 const { Vtx, Vty, Vtz, Vpx, Vpy, Vpz } = computeMetric(npts, Xt, Xp, Yt, Yp, Zt, Zp);
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 217 // Flux-form metric weights for the six-transform scheme, built from the
218 // *undivided* theta tangents sin(theta)*X_theta (smooth on the sphere,
219 // unlike X_theta itself) and the same X_phi as above. Also a one-off;
220 // the f64 combination happens on the CPU, rounded to f32 for upload.
221 const sXtx = await deriv.sinDtheta(X);
222 const sXty = await deriv.sinDtheta(Y);
223 const sXtz = await deriv.sinDtheta(Z);
224 const flux = computeFluxMetric(npts, sXtx, sXty, sXtz, Xp, Yp, Zp);
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 226 // The preconditioner scale — see the Jhat field comment. The symbol
227 // matrix in the orthonormal frame is S = (1/J)[[p1,p2],[p2,q2]] with
228 // 1/J = r sin^2(theta); its entries are the bounded quantities
229 // g^tt, sin g^tp, sin^2 g^pp, so the eigenvalue extremes are clean to
230 // take over the grid. det S = 1/J^2, so the area factor comes along
231 // for free. f64 throughout.
232 let muMin = Infinity;
233 let muMax = 0;
234 let Jmin = Infinity;
235 let Jmax = 0;
236 for (let i = 0; i < cfg.nlat; i++) {
237 const ct = sht.cosTheta[i];
238 const st2 = Math.max(0, 1 - ct * ct);
239 for (let j = 0; j < cfg.nphi; j++) {
240 const k = i * cfg.nphi + j;
241 const invJ = flux.r[k] * st2;
242 const s11 = flux.p1[k] * invJ;
243 const s12 = flux.p2[k] * invJ;
244 const s22 = flux.q2[k] * invJ;
245 const mean = (s11 + s22) / 2;
246 const disc = Math.sqrt(((s11 - s22) / 2) ** 2 + s12 * s12);
247 if (mean - disc < muMin) muMin = mean - disc;
248 if (mean + disc > muMax) muMax = mean + disc;
249 const J = 1 / invJ;
250 if (J < Jmin) Jmin = J;
251 if (J > Jmax) Jmax = J;
252 }
253 }
254 const Jhat = 2 / (muMin + muMax);
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 256 return new Geometry({
257 x, y, z, X, Y, Z, Vtx, Vty, Vtz, Vpx, Vpy, Vpz,
258 p1: new Float32Array(flux.p1),
259 p2: new Float32Array(flux.p2),
260 q2: new Float32Array(flux.q2),
261 r: new Float32Array(flux.r),
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 262 Jhat, muMin, muMax, Jmin, Jmax,
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 263 });
265 plan.destroy();
266 host.destroy();
267 }
268 }
270 /**
271 * The surface evaluated on another plan's grid, as interleaved xyz vertex
272 * positions (nlat * nphi * 3) — for rendering at display resolution. Exact
273 * interpolation: the same coefficients, more evaluation points.
274 */
275 async positionsOn(view: ShtPlan): Promise<Float32Array> {
276 const [x, y, z] = [
277 await view.synth(this.X),
278 await view.synth(this.Y),
279 await view.synth(this.Z),
280 ];
281 const out = new Float32Array(x.length * 3);
282 for (let i = 0; i < x.length; i++) {
283 out[3 * i] = x[i];
284 out[3 * i + 1] = y[i];
285 out[3 * i + 2] = z[i];
286 }
287 return out;
288 }
290 /** How far the surface departs from the unit sphere, as min/max radius. */
291 radiusRange(): { lo: number; hi: number } {
292 let lo = Infinity;
293 let hi = -Infinity;
294 for (let i = 0; i < this.x.length; i++) {
295 const r = Math.hypot(this.x[i], this.y[i], this.z[i]);
296 if (r < lo) lo = r;
297 if (r > hi) hi = r;
298 }
299 return { lo, hi };
300 }
301}
303/** The (theta, phi) of every grid point, flattened phi-fastest as the fields are. */
304function gridAngles(
305 sht: ShtPlan,
306 cfg: ShtConfig,
307): { theta: Float32Array; phi: Float32Array } {
308 const { nlat, nphi } = cfg;
309 const theta = new Float32Array(nlat * nphi);
310 const phi = new Float32Array(nlat * nphi);
311 for (let i = 0; i < nlat; i++) {
312 const th = Math.acos(Math.max(-1, Math.min(1, sht.cosTheta[i])));
313 for (let j = 0; j < nphi; j++) {
314 theta[i * nphi + j] = th;
315 phi[i * nphi + j] = (2 * Math.PI * j) / nphi;
316 }
317 }
318 return { theta, phi };
319}
321async function readBuffer(
322 device: GPUDevice,
323 plan: ModelPlan,
324 name: string,
325 count: number,
326): Promise<Float32Array> {
327 const buffer = plan.buffer(name);
328 if (!buffer) {
329 throw new Error(`the geometry never assigns '${name}'`);
330 }
331 const staging = device.createBuffer({
332 label: `geometry-read-${name}`,
333 size: 4 * count,
334 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
335 });
336 try {
337 const enc = device.createCommandEncoder({ label: `geometry-read-${name}` });
338 enc.copyBufferToBuffer(buffer, 0, staging, 0, 4 * count);
339 device.queue.submit([enc.finish()]);
340 await staging.mapAsync(GPUMapMode.READ);
341 const out = new Float32Array(staging.getMappedRange().slice(0));
342 staging.unmap();
343 return out;
344 } finally {
345 staging.destroy();
346 }
347}