2 * One running model: grid, transforms, compiled .m, seeded state.
3 *
4 * Everything that is not rendering. The app, the desktop benchmark and the
5 * tests all go through this, so there is one place that decides how a model is
6 * turned into something running on the GPU — and nothing about it is
7 * browser-specific beyond needing a GPUDevice.
8 */
9import { ShtPlan } from '../sht/sht.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 10import { DerivPlan } from '../sht/deriv.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 11import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
12import { GpuModel, type ModelParams } from './model.ts';
13import { seededNoise } from './noise.ts';
14import type { MModel } from './registry.ts';
15import { Geometry } from '../geom/geometry.ts';
16import { mGeometryByKey, defaultGeometryParams, SPHERE_KEY, type MGeometry } from '../geom/registry.ts';
18export interface ModelSessionOptions {
19 device: GPUDevice;
20 model: MModel;
21 params: ModelParams;
22 lmax: number;
23 /** Override the model source — the editor's working copy. */
24 source?: string;
25 /** Linear render oversampling: read the species fields on a grid this many
26 * times finer than the solver's in each direction (default 1). The state is
27 * band-limited at lmax, so the finer evaluation is exact interpolation. */
28 oversample?: number;
29 /** The surface to solve on. Defaults to the unit sphere. */
30 geometry?: MGeometry;
31 geometryParams?: ModelParams;
32 /** Override the geometry source — the editor's working copy. */
33 geometrySource?: string;
34 /**
35 * Iterations of the .m's implicit solve. Structural, not tunable: the loop
36 * is unrolled into the op sequence, so a change recompiles.
37 */
38 niter?: number;
39}
41export class ModelSession {
42 readonly device: GPUDevice;
43 readonly model: MModel;
44 readonly cfg: ShtConfig;
45 readonly sht: ShtPlan;
46 readonly gpu: GpuModel;
47 readonly npts: number;
48 /** Iterations of the implicit solve compiled into the step. */
49 readonly niter: number;
51 /** The surface being solved on, as spherical-harmonic coefficients. */
52 #geometry: Geometry;
53 #geometryModel: MGeometry;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 54 /** Computes the theta/phi derivatives a geometry's metric quantities need. */
55 #deriv: DerivPlan;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 56
57 /** Model time and step count since the last seeding. */
58 t = 0;
59 steps = 0;
61 #params: ModelParams;
62 /** Display-only transforms on the oversampled grid; null at 1x. */
63 #displaySht: ShtPlan | null;
64 #oversample: number;
66 private constructor(init: {
67 device: GPUDevice;
68 model: MModel;
69 cfg: ShtConfig;
70 sht: ShtPlan;
71 displaySht: ShtPlan | null;
72 gpu: GpuModel;
73 params: ModelParams;
74 oversample: number;
75 geometry: Geometry;
76 geometryModel: MGeometry;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 77 deriv: DerivPlan;
79 }) {
80 this.device = init.device;
81 this.model = init.model;
82 this.cfg = init.cfg;
83 this.sht = init.sht;
84 this.gpu = init.gpu;
85 this.npts = init.cfg.nlat * init.cfg.nphi;
86 this.#oversample = init.oversample;
87 this.#params = init.params;
88 this.#displaySht = init.displaySht;
89 this.#geometry = init.geometry;
90 this.#geometryModel = init.geometryModel;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 91 this.#deriv = init.deriv;
93 }
95 get geometry(): Geometry {
96 return this.#geometry;
97 }
99 get geometryModel(): MGeometry {
100 return this.#geometryModel;
101 }
103 /** Linear render oversampling factor (1 = read on the solver grid). */
104 get oversample(): number {
105 return this.#oversample;
106 }
108 static async create(opts: ModelSessionOptions): Promise<ModelSession> {
109 const { device, model, params, lmax } = opts;
110 const oversample = Math.max(1, Math.round(opts.oversample ?? 1));
111 const niter = Math.max(0, Math.round(opts.niter ?? 1));
112 const geometryModel = opts.geometry ?? mGeometryByKey(SPHERE_KEY)!;
113 const geometryParams = opts.geometryParams ?? defaultGeometryParams(geometryModel);
114 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
115 const cfg = { lmax, mmax: lmax, nlat, nphi };
116 const sht = await ShtPlan.create(device, cfg);
117 let displaySht: ShtPlan | null = null;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 118 let deriv: DerivPlan | null = null;
120 // The display plan shares nothing with the solver's beyond the
121 // coefficients copied into it per readback; its grid is the solver's
122 // scaled by the oversampling factor, so nphi stays a power of two (the
123 // FFT path) for power-of-two factors.
124 if (oversample > 1) {
125 displaySht = await ShtPlan.create(device, {
126 lmax,
127 mmax: lmax,
128 nlat: oversample * nlat,
129 nphi: oversample * nphi,
130 });
131 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 132 // Computes the theta/phi derivatives the geometry's metric quantities
133 // (and, per step, the surface Laplace-Beltrami correction) need.
134 deriv = await DerivPlan.create(device, sht);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 135 // The surface is built before the model, because the model takes it as
136 // an argument. It is a one-off: compiled, evaluated, read back, and its
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 137 // plan discarded — nothing of it survives into the timestep but twelve
138 // buffers of numbers (the embedding and the metric quantities built on it).
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 139 const geometry = await Geometry.create({
140 device,
141 sht,
142 cfg,
143 source: opts.geometrySource ?? geometryModel.source,
144 paramNames: geometryModel.params.map((p) => p.key),
145 params: geometryParams,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 146 deriv,
148 const gpu = await GpuModel.create({
149 device,
150 sht,
151 cfg,
152 source: opts.source ?? model.source,
153 paramNames: model.params.map((p) => p.key),
154 state: model.state,
155 view: model.species,
156 geometry,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 157 deriv,
159 });
160 gpu.setParams(params);
161 return new ModelSession({
162 device, model, cfg, sht, displaySht, gpu, params, oversample,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 163 geometry, geometryModel, deriv, niter,
165 } catch (e) {
166 // The transform plans own GPU buffers; do not leak them on a compile error.
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 167 deriv?.destroy();
169 sht.destroy();
170 throw e;
171 }
172 }
174 /**
175 * Vertex positions for the current render grid: the surface synthesized on
176 * `viewSht`, interleaved xyz. Exact interpolation of the same coefficients
177 * the solver sees, so the drawn surface is the one being solved on however
178 * finely it is sampled.
179 */
180 renderPositions(): Promise<Float32Array> {
181 return this.#geometry.positionsOn(this.viewSht);
182 }
184 /**
185 * Change the surface in place, without recompiling or disturbing the run.
186 * The geometry's shape in the bindings depends only on the grid, so the
187 * compiled step does not change — only the numbers it reads. The caller
188 * still has to rebuild the mesh from `renderPositions()`.
189 */
190 async setGeometry(
191 geometryModel: MGeometry,
192 params: ModelParams,
193 source?: string,
194 ): Promise<void> {
195 const next = await Geometry.create({
196 device: this.device,
197 sht: this.sht,
198 cfg: this.cfg,
199 source: source ?? geometryModel.source,
200 paramNames: geometryModel.params.map((p) => p.key),
201 params,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 202 deriv: this.#deriv,
204 this.#geometry = next;
205 this.#geometryModel = geometryModel;
206 this.gpu.uploadGeometry(next);
207 }
209 /** The plan whose grid `readSpecies` samples on — the display plan when
210 * oversampling, otherwise the solver's. Its cosTheta/nphi define the mesh. */
211 get viewSht(): ShtPlan {
212 return this.#displaySht ?? this.sht;
213 }
215 /**
216 * Change the display oversampling in place. Display-only: the simulation
217 * state, time and parameters are untouched, so the run continues seamlessly
218 * on the new render grid. The caller must not have a readSpecies in flight —
219 * its readback maps a buffer of the plan being destroyed.
220 */
221 async setOversample(oversample: number): Promise<void> {
222 const os = Math.max(1, Math.round(oversample));
223 if (os === this.#oversample) return;
224 const next =
225 os > 1
226 ? await ShtPlan.create(this.device, {
227 lmax: this.cfg.lmax,
228 mmax: this.cfg.mmax,
229 nlat: os * this.cfg.nlat,
230 nphi: os * this.cfg.nphi,
231 })
232 : null;
233 const old = this.#displaySht;
234 this.#displaySht = next;
235 this.#oversample = os;
236 old?.destroy();
237 }
239 /** Run `init` from a seeded perturbation, resetting model time. */
240 seed(seed: number): void {
241 this.gpu.init(seededNoise(this.npts, this.model.seedAmp, seed));
242 this.t = 0;
243 this.steps = 0;
244 }
246 setParams(params: ModelParams): void {
247 this.#params = params;
248 this.gpu.setParams(params);
249 }
251 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
252 step(n = 1): void {
253 this.gpu.step(n);
254 this.t += n * (this.#params.dt ?? 0);
255 this.steps += n;
256 }
258 /**
259 * Wait for the submitted steps to finish, without reading anything back.
260 * This is the honest way to time the solver: a readback would add a GPU->CPU
261 * round trip, which in a browser also crosses a process boundary and can cost
262 * more than the steps themselves.
263 */
264 sync(): Promise<undefined> {
265 return this.device.queue.onSubmittedWorkDone();
266 }
268 /**
269 * Time a batch of `n` steps and return ms/step, leaving the simulation
270 * exactly where it was: the spectral state is snapshotted before the batch
271 * and restored after, and `t`/`steps` do not advance. One sync amortized
272 * over the batch — the same measurement the desktop benchmark makes. The
273 * grid view fields hold the batch's output until the next real step, so
274 * step before reading them.
275 */
276 async measure(n: number): Promise<number> {
277 this.gpu.snapshotState();
278 const t0 = performance.now();
279 this.gpu.step(n);
280 await this.sync();
281 const ms = (performance.now() - t0) / n;
282 this.gpu.restoreState();
283 return ms;
284 }
286 /** Read a named value (a grid field or the spectral state). */
287 read(name: string): Promise<Float32Array> {
288 return this.gpu.read(name);
289 }
291 /**
292 * Read species `k` at render resolution (`viewSht`'s grid). Without
293 * oversampling this is the grid field the .m returned. With oversampling the
294 * spectral state is synthesized on the finer grid instead — the same field,
295 * since the models define each species as synth of its state, evaluated
296 * exactly on more points.
297 */
298 readSpecies(k: number): Promise<Float32Array> {
299 if (!this.#displaySht) return this.read(this.model.species[k]);
300 const state = this.model.state[k];
301 const buf = this.gpu.valueBuffer(state);
302 if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
303 return this.#displaySht.synthFrom(buf);
304 }
306 describe(): { init: string[]; step: string[] } {
307 return this.gpu.describe();
308 }
310 destroy(): void {
311 this.gpu.destroy();
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 312 this.#deriv.destroy();
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 313 this.#displaySht?.destroy();
314 this.sht.destroy();
315 }
316}