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 { HostBuffers, ModelPlan } from '../mgpu/plan.ts';
34import { CompiledModel, type Binding } from '../mgpu/compile.ts';
35import { inFunction, inFunctionAsync, inModel } from '../mgpu/errors.ts';
36import type { ModelParams } from '../mgpu/model.ts';
38/** The function a geometry file must define. */
39export const SHAPE_FN = 'shape';
41export interface GeometryOptions {
42 device: GPUDevice;
43 /** The solver's transform plan — the grid the shape is evaluated on. */
44 sht: ShtPlan;
45 cfg: ShtConfig;
46 /** Geometry source (.m text). */
47 source: string;
48 /** Parameter names the .m may take beyond `theta` and `phi`. */
49 paramNames: string[];
50 params: ModelParams;
51}
53export class Geometry {
54 /** Coordinates on the solver grid, npts each — synthesis of the coefficients. */
55 readonly x: Float32Array;
56 readonly y: Float32Array;
57 readonly z: Float32Array;
58 /** Their spherical-harmonic coefficients, 2 x nlm each. */
59 readonly X: Float32Array;
60 readonly Y: Float32Array;
61 readonly Z: Float32Array;
63 private constructor(init: {
64 x: Float32Array; y: Float32Array; z: Float32Array;
65 X: Float32Array; Y: Float32Array; Z: Float32Array;
66 }) {
67 this.x = init.x;
68 this.y = init.y;
69 this.z = init.z;
70 this.X = init.X;
71 this.Y = init.Y;
72 this.Z = init.Z;
73 }
75 /**
76 * Compile the shape file, evaluate it once on the solver grid, and reduce it
77 * to coefficients. Everything here happens at build time — a geometry never
78 * takes part in the timestep — so it reads back through the CPU freely.
79 */
80 static async create(opts: GeometryOptions): Promise<Geometry> {
81 const { device, sht, cfg, source, paramNames, params } = opts;
82 const npts = cfg.nlat * cfg.nphi;
83 const nlm = sht.nlm;
85 const bindings: Record<string, Binding> = {
86 theta: { kind: 'tensor', shape: [npts, 1] },
87 phi: { kind: 'tensor', shape: [npts, 1] },
88 npts: { kind: 'const', value: npts },
89 };
90 for (const p of paramNames) bindings[p] = { kind: 'param' };
92 const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
93 const fn = inFunction(SHAPE_FN, () => compiled.specialize(SHAPE_FN, 3));
94 compiled.finish();
96 const host = new HostBuffers(device);
97 host.ensure('theta', npts);
98 host.ensure('phi', npts);
100 const plan = await inFunctionAsync(SHAPE_FN, () =>
101 // Nothing feeds back: the three outputs are read once and the plan is
102 // thrown away.
103 ModelPlan.create(device, sht, { fn, feedback: [null, null, null] }, host),
104 );
106 try {
107 const { theta, phi } = gridAngles(sht, cfg);
108 host.upload('theta', theta);
109 host.upload('phi', phi);
110 plan.setParams(params);
112 const enc = device.createCommandEncoder({ label: 'geometry-shape' });
113 plan.encodeSteps(enc, 1);
114 device.queue.submit([enc.finish()]);
116 const raw = await Promise.all(
117 fn.outputs.map((out) => readBuffer(device, plan, out.name, npts)),
118 );
119 // Coefficients first, then back to the grid: what the solver and the
120 // renderer both see is the band-limited surface, not the raw .m output.
121 const [X, Y, Z] = [
122 await sht.analys(raw[0]),
123 await sht.analys(raw[1]),
124 await sht.analys(raw[2]),
125 ];
126 const [x, y, z] = [
127 await sht.synth(X),
128 await sht.synth(Y),
129 await sht.synth(Z),
130 ];
131 return new Geometry({ x, y, z, X, Y, Z });
132 } finally {
133 plan.destroy();
134 host.destroy();
135 }
136 }
138 /**
139 * The surface evaluated on another plan's grid, as interleaved xyz vertex
140 * positions (nlat * nphi * 3) — for rendering at display resolution. Exact
141 * interpolation: the same coefficients, more evaluation points.
142 */
143 async positionsOn(view: ShtPlan): Promise<Float32Array> {
144 const [x, y, z] = [
145 await view.synth(this.X),
146 await view.synth(this.Y),
147 await view.synth(this.Z),
148 ];
149 const out = new Float32Array(x.length * 3);
150 for (let i = 0; i < x.length; i++) {
151 out[3 * i] = x[i];
152 out[3 * i + 1] = y[i];
153 out[3 * i + 2] = z[i];
154 }
155 return out;
156 }
158 /** How far the surface departs from the unit sphere, as min/max radius. */
159 radiusRange(): { lo: number; hi: number } {
160 let lo = Infinity;
161 let hi = -Infinity;
162 for (let i = 0; i < this.x.length; i++) {
163 const r = Math.hypot(this.x[i], this.y[i], this.z[i]);
164 if (r < lo) lo = r;
165 if (r > hi) hi = r;
166 }
167 return { lo, hi };
168 }
169}
171/** The (theta, phi) of every grid point, flattened phi-fastest as the fields are. */
172function gridAngles(
173 sht: ShtPlan,
174 cfg: ShtConfig,
175): { theta: Float32Array; phi: Float32Array } {
176 const { nlat, nphi } = cfg;
177 const theta = new Float32Array(nlat * nphi);
178 const phi = new Float32Array(nlat * nphi);
179 for (let i = 0; i < nlat; i++) {
180 const th = Math.acos(Math.max(-1, Math.min(1, sht.cosTheta[i])));
181 for (let j = 0; j < nphi; j++) {
182 theta[i * nphi + j] = th;
183 phi[i * nphi + j] = (2 * Math.PI * j) / nphi;
184 }
185 }
186 return { theta, phi };
187}
189async function readBuffer(
190 device: GPUDevice,
191 plan: ModelPlan,
192 name: string,
193 count: number,
194): Promise<Float32Array> {
195 const buffer = plan.buffer(name);
196 if (!buffer) {
197 throw new Error(`the geometry never assigns '${name}'`);
198 }
199 const staging = device.createBuffer({
200 label: `geometry-read-${name}`,
201 size: 4 * count,
202 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
203 });
204 try {
205 const enc = device.createCommandEncoder({ label: `geometry-read-${name}` });
206 enc.copyBufferToBuffer(buffer, 0, staging, 0, 4 * count);
207 device.queue.submit([enc.finish()]);
208 await staging.mapAsync(GPUMapMode.READ);
209 const out = new Float32Array(staging.getMappedRange().slice(0));
210 staging.unmap();
211 return out;
212 } finally {
213 staging.destroy();
214 }
215}