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 — 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';
33import type { DerivPlan } from '../sht/deriv.ts';
34import { computeMetric, computeFluxMetric } from './metric.ts';
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;
53 /** Computes the theta/phi derivatives the inverse metric quantities need. */
54 deriv: DerivPlan;
55}
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;
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
69 * one-off computed here, not per-solve-step work. Used by the Algorithm-4
70 * (12-transform) Laplace-Beltrami path.
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;
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;
89 private constructor(init: {
90 x: Float32Array; y: Float32Array; z: Float32Array;
91 X: Float32Array; Y: Float32Array; Z: Float32Array;
92 Vtx: Float32Array; Vty: Float32Array; Vtz: Float32Array;
93 Vpx: Float32Array; Vpy: Float32Array; Vpz: Float32Array;
94 p1: Float32Array; p2: Float32Array; q2: Float32Array; r: Float32Array;
95 }) {
96 this.x = init.x;
97 this.y = init.y;
98 this.z = init.z;
99 this.X = init.X;
100 this.Y = init.Y;
101 this.Z = init.Z;
102 this.Vtx = init.Vtx;
103 this.Vty = init.Vty;
104 this.Vtz = init.Vtz;
105 this.Vpx = init.Vpx;
106 this.Vpy = init.Vpy;
107 this.Vpz = init.Vpz;
108 this.p1 = init.p1;
109 this.p2 = init.p2;
110 this.q2 = init.q2;
111 this.r = init.r;
112 }
114 /**
115 * Compile the shape file, evaluate it once on the solver grid, and reduce it
116 * to coefficients. Everything here happens at build time — a geometry never
117 * takes part in the timestep — so it reads back through the CPU freely.
118 */
119 static async create(opts: GeometryOptions): Promise<Geometry> {
120 const { device, sht, cfg, source, paramNames, params, deriv } = opts;
121 const npts = cfg.nlat * cfg.nphi;
122 const nlm = sht.nlm;
124 const bindings: Record<string, Binding> = {
125 theta: { kind: 'tensor', shape: [npts, 1] },
126 phi: { kind: 'tensor', shape: [npts, 1] },
127 npts: { kind: 'const', value: npts },
128 };
129 for (const p of paramNames) bindings[p] = { kind: 'param' };
131 const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
132 const fn = inFunction(SHAPE_FN, () => compiled.specialize(SHAPE_FN, 3));
133 compiled.finish();
135 const host = new HostBuffers(device);
136 host.ensure('theta', npts);
137 host.ensure('phi', npts);
139 const plan = await inFunctionAsync(SHAPE_FN, () =>
140 // Nothing feeds back: the three outputs are read once and the plan is
141 // thrown away.
142 ModelPlan.create(device, sht, { fn, feedback: [null, null, null] }, host),
143 );
145 try {
146 const { theta, phi } = gridAngles(sht, cfg);
147 host.upload('theta', theta);
148 host.upload('phi', phi);
149 plan.setParams(params);
151 const enc = device.createCommandEncoder({ label: 'geometry-shape' });
152 plan.encodeSteps(enc, 1);
153 device.queue.submit([enc.finish()]);
155 const raw = await Promise.all(
156 fn.outputs.map((out) => readBuffer(device, plan, out.name, npts)),
157 );
158 // Coefficients first, then back to the grid: what the solver and the
159 // renderer both see is the band-limited surface, not the raw .m output.
160 const [X, Y, Z] = [
161 await sht.analys(raw[0]),
162 await sht.analys(raw[1]),
163 await sht.analys(raw[2]),
164 ];
165 const [x, y, z] = [
166 await sht.synth(X),
167 await sht.synth(Y),
168 await sht.synth(Z),
169 ];
171 // Inverse metric quantities (algos.tex Algorithm 2): theta/phi
172 // derivatives of the embedding's coefficients, contracted through the
173 // inverse first fundamental form. Depends only on the geometry, so
174 // this is a one-off alongside x,y,z above, not per-step work.
175 const Xt = await deriv.dtheta(X);
176 const Xp = await deriv.dphi(X);
177 const Yt = await deriv.dtheta(Y);
178 const Yp = await deriv.dphi(Y);
179 const Zt = await deriv.dtheta(Z);
180 const Zp = await deriv.dphi(Z);
181 const { Vtx, Vty, Vtz, Vpx, Vpy, Vpz } = computeMetric(npts, Xt, Xp, Yt, Yp, Zt, Zp);
183 // Flux-form metric weights for the six-transform scheme, built from the
184 // *undivided* theta tangents sin(theta)*X_theta (smooth on the sphere,
185 // unlike X_theta itself) and the same X_phi as above. Also a one-off;
186 // the f64 combination happens on the CPU, rounded to f32 for upload.
187 const sXtx = await deriv.sinDtheta(X);
188 const sXty = await deriv.sinDtheta(Y);
189 const sXtz = await deriv.sinDtheta(Z);
190 const flux = computeFluxMetric(npts, sXtx, sXty, sXtz, Xp, Yp, Zp);
192 return new Geometry({
193 x, y, z, X, Y, Z, Vtx, Vty, Vtz, Vpx, Vpy, Vpz,
194 p1: new Float32Array(flux.p1),
195 p2: new Float32Array(flux.p2),
196 q2: new Float32Array(flux.q2),
197 r: new Float32Array(flux.r),
198 });
199 } finally {
200 plan.destroy();
201 host.destroy();
202 }
203 }
205 /**
206 * The surface evaluated on another plan's grid, as interleaved xyz vertex
207 * positions (nlat * nphi * 3) — for rendering at display resolution. Exact
208 * interpolation: the same coefficients, more evaluation points.
209 */
210 async positionsOn(view: ShtPlan): Promise<Float32Array> {
211 const [x, y, z] = [
212 await view.synth(this.X),
213 await view.synth(this.Y),
214 await view.synth(this.Z),
215 ];
216 const out = new Float32Array(x.length * 3);
217 for (let i = 0; i < x.length; i++) {
218 out[3 * i] = x[i];
219 out[3 * i + 1] = y[i];
220 out[3 * i + 2] = z[i];
221 }
222 return out;
223 }
225 /** How far the surface departs from the unit sphere, as min/max radius. */
226 radiusRange(): { lo: number; hi: number } {
227 let lo = Infinity;
228 let hi = -Infinity;
229 for (let i = 0; i < this.x.length; i++) {
230 const r = Math.hypot(this.x[i], this.y[i], this.z[i]);
231 if (r < lo) lo = r;
232 if (r > hi) hi = r;
233 }
234 return { lo, hi };
235 }
236}
238/** The (theta, phi) of every grid point, flattened phi-fastest as the fields are. */
239function gridAngles(
240 sht: ShtPlan,
241 cfg: ShtConfig,
242): { theta: Float32Array; phi: Float32Array } {
243 const { nlat, nphi } = cfg;
244 const theta = new Float32Array(nlat * nphi);
245 const phi = new Float32Array(nlat * nphi);
246 for (let i = 0; i < nlat; i++) {
247 const th = Math.acos(Math.max(-1, Math.min(1, sht.cosTheta[i])));
248 for (let j = 0; j < nphi; j++) {
249 theta[i * nphi + j] = th;
250 phi[i * nphi + j] = (2 * Math.PI * j) / nphi;
251 }
252 }
253 return { theta, phi };
254}
256async function readBuffer(
257 device: GPUDevice,
258 plan: ModelPlan,
259 name: string,
260 count: number,
261): Promise<Float32Array> {
262 const buffer = plan.buffer(name);
263 if (!buffer) {
264 throw new Error(`the geometry never assigns '${name}'`);
265 }
266 const staging = device.createBuffer({
267 label: `geometry-read-${name}`,
268 size: 4 * count,
269 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
270 });
271 try {
272 const enc = device.createCommandEncoder({ label: `geometry-read-${name}` });
273 enc.copyBufferToBuffer(buffer, 0, staging, 0, 4 * count);
274 device.queue.submit([enc.finish()]);
275 await staging.mapAsync(GPUMapMode.READ);
276 const out = new Float32Array(staging.getMappedRange().slice(0));
277 staging.unmap();
278 return out;
279 } finally {
280 staging.destroy();
281 }
282}