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 } 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.
70 */
71 readonly Vtx: Float32Array;
72 readonly Vty: Float32Array;
73 readonly Vtz: Float32Array;
74 readonly Vpx: Float32Array;
75 readonly Vpy: Float32Array;
76 readonly Vpz: Float32Array;
78 private constructor(init: {
79 x: Float32Array; y: Float32Array; z: Float32Array;
80 X: Float32Array; Y: Float32Array; Z: Float32Array;
81 Vtx: Float32Array; Vty: Float32Array; Vtz: Float32Array;
82 Vpx: Float32Array; Vpy: Float32Array; Vpz: Float32Array;
83 }) {
84 this.x = init.x;
85 this.y = init.y;
86 this.z = init.z;
87 this.X = init.X;
88 this.Y = init.Y;
89 this.Z = init.Z;
90 this.Vtx = init.Vtx;
91 this.Vty = init.Vty;
92 this.Vtz = init.Vtz;
93 this.Vpx = init.Vpx;
94 this.Vpy = init.Vpy;
95 this.Vpz = init.Vpz;
96 }
98 /**
99 * Compile the shape file, evaluate it once on the solver grid, and reduce it
100 * to coefficients. Everything here happens at build time — a geometry never
101 * takes part in the timestep — so it reads back through the CPU freely.
102 */
103 static async create(opts: GeometryOptions): Promise<Geometry> {
104 const { device, sht, cfg, source, paramNames, params, deriv } = opts;
105 const npts = cfg.nlat * cfg.nphi;
106 const nlm = sht.nlm;
108 const bindings: Record<string, Binding> = {
109 theta: { kind: 'tensor', shape: [npts, 1] },
110 phi: { kind: 'tensor', shape: [npts, 1] },
111 npts: { kind: 'const', value: npts },
112 };
113 for (const p of paramNames) bindings[p] = { kind: 'param' };
115 const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
116 const fn = inFunction(SHAPE_FN, () => compiled.specialize(SHAPE_FN, 3));
117 compiled.finish();
119 const host = new HostBuffers(device);
120 host.ensure('theta', npts);
121 host.ensure('phi', npts);
123 const plan = await inFunctionAsync(SHAPE_FN, () =>
124 // Nothing feeds back: the three outputs are read once and the plan is
125 // thrown away.
126 ModelPlan.create(device, sht, { fn, feedback: [null, null, null] }, host),
127 );
129 try {
130 const { theta, phi } = gridAngles(sht, cfg);
131 host.upload('theta', theta);
132 host.upload('phi', phi);
133 plan.setParams(params);
135 const enc = device.createCommandEncoder({ label: 'geometry-shape' });
136 plan.encodeSteps(enc, 1);
137 device.queue.submit([enc.finish()]);
139 const raw = await Promise.all(
140 fn.outputs.map((out) => readBuffer(device, plan, out.name, npts)),
141 );
142 // Coefficients first, then back to the grid: what the solver and the
143 // renderer both see is the band-limited surface, not the raw .m output.
144 const [X, Y, Z] = [
145 await sht.analys(raw[0]),
146 await sht.analys(raw[1]),
147 await sht.analys(raw[2]),
148 ];
149 const [x, y, z] = [
150 await sht.synth(X),
151 await sht.synth(Y),
152 await sht.synth(Z),
153 ];
155 // Inverse metric quantities (algos.tex Algorithm 2): theta/phi
156 // derivatives of the embedding's coefficients, contracted through the
157 // inverse first fundamental form. Depends only on the geometry, so
158 // this is a one-off alongside x,y,z above, not per-step work.
159 const Xt = await deriv.dtheta(X);
160 const Xp = await deriv.dphi(X);
161 const Yt = await deriv.dtheta(Y);
162 const Yp = await deriv.dphi(Y);
163 const Zt = await deriv.dtheta(Z);
164 const Zp = await deriv.dphi(Z);
165 const { Vtx, Vty, Vtz, Vpx, Vpy, Vpz } = computeMetric(npts, Xt, Xp, Yt, Yp, Zt, Zp);
167 return new Geometry({ x, y, z, X, Y, Z, Vtx, Vty, Vtz, Vpx, Vpy, Vpz });
168 } finally {
169 plan.destroy();
170 host.destroy();
171 }
172 }
174 /**
175 * The surface evaluated on another plan's grid, as interleaved xyz vertex
176 * positions (nlat * nphi * 3) — for rendering at display resolution. Exact
177 * interpolation: the same coefficients, more evaluation points.
178 */
179 async positionsOn(view: ShtPlan): Promise<Float32Array> {
180 const [x, y, z] = [
181 await view.synth(this.X),
182 await view.synth(this.Y),
183 await view.synth(this.Z),
184 ];
185 const out = new Float32Array(x.length * 3);
186 for (let i = 0; i < x.length; i++) {
187 out[3 * i] = x[i];
188 out[3 * i + 1] = y[i];
189 out[3 * i + 2] = z[i];
190 }
191 return out;
192 }
194 /** How far the surface departs from the unit sphere, as min/max radius. */
195 radiusRange(): { lo: number; hi: number } {
196 let lo = Infinity;
197 let hi = -Infinity;
198 for (let i = 0; i < this.x.length; i++) {
199 const r = Math.hypot(this.x[i], this.y[i], this.z[i]);
200 if (r < lo) lo = r;
201 if (r > hi) hi = r;
202 }
203 return { lo, hi };
204 }
205}
207/** The (theta, phi) of every grid point, flattened phi-fastest as the fields are. */
208function gridAngles(
209 sht: ShtPlan,
210 cfg: ShtConfig,
211): { theta: Float32Array; phi: Float32Array } {
212 const { nlat, nphi } = cfg;
213 const theta = new Float32Array(nlat * nphi);
214 const phi = new Float32Array(nlat * nphi);
215 for (let i = 0; i < nlat; i++) {
216 const th = Math.acos(Math.max(-1, Math.min(1, sht.cosTheta[i])));
217 for (let j = 0; j < nphi; j++) {
218 theta[i * nphi + j] = th;
219 phi[i * nphi + j] = (2 * Math.PI * j) / nphi;
220 }
221 }
222 return { theta, phi };
223}
225async function readBuffer(
226 device: GPUDevice,
227 plan: ModelPlan,
228 name: string,
229 count: number,
230): Promise<Float32Array> {
231 const buffer = plan.buffer(name);
232 if (!buffer) {
233 throw new Error(`the geometry never assigns '${name}'`);
234 }
235 const staging = device.createBuffer({
236 label: `geometry-read-${name}`,
237 size: 4 * count,
238 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
239 });
240 try {
241 const enc = device.createCommandEncoder({ label: `geometry-read-${name}` });
242 enc.copyBufferToBuffer(buffer, 0, staging, 0, 4 * count);
243 device.queue.submit([enc.finish()]);
244 await staging.mapAsync(GPUMapMode.READ);
245 const out = new Float32Array(staging.getMappedRange().slice(0));
246 staging.unmap();
247 return out;
248 } finally {
249 staging.destroy();
250 }
251}