1/**
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';
10import { DerivPlan } from '../sht/deriv.ts';
11import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
12import { GpuModel, type ModelParams } from './model.ts';
13import { seededNoise } from './noise.ts';
14import { boundingBox, drawModesAsync, DEFAULT_LAMBDA } from './randnfun3.ts';
15import { resolveLambda } from './plan.ts';
16import type { MModel } from './registry.ts';
17import { Geometry } from '../geom/geometry.ts';
18import { mGeometryByKey, defaultGeometryParams, SPHERE_KEY, type MGeometry } from '../geom/registry.ts';
20export interface ModelSessionOptions {
21 device: GPUDevice;
22 model: MModel;
23 params: ModelParams;
24 lmax: number;
25 /** Override the model source — the editor's working copy. */
26 source?: string;
27 /** Linear render oversampling: read the species fields on a grid this many
28 * times finer than the solver's in each direction (default 1). The state is
29 * band-limited at lmax, so the finer evaluation is exact interpolation. */
30 oversample?: number;
31 /** The surface to solve on. Defaults to the unit sphere. */
32 geometry?: MGeometry;
33 geometryParams?: ModelParams;
34 /** Override the geometry source — the editor's working copy. */
35 geometrySource?: string;
36 /**
37 * Iterations of the .m's implicit solve. Structural, not tunable: the loop
38 * is unrolled into the op sequence, so a change recompiles.
39 */
40 niter?: number;
41 /** Wavelength of the seeded random field a model's `init` draws
42 * (src/mgpu/randnfun3.ts). Redrawn on the next seed, never recompiled. */
43 lam3?: number;
44}
46export class ModelSession {
47 readonly device: GPUDevice;
48 readonly model: MModel;
49 readonly cfg: ShtConfig;
50 readonly sht: ShtPlan;
51 readonly gpu: GpuModel;
52 readonly npts: number;
53 /** Iterations of the implicit solve compiled into the step. */
54 readonly niter: number;
56 /** The surface being solved on, as spherical-harmonic coefficients. */
57 #geometry: Geometry;
58 #geometryModel: MGeometry;
59 /** Computes the theta/phi derivatives a geometry's metric quantities need. */
60 #deriv: DerivPlan;
62 /** Model time and step count since the last seeding. */
63 t = 0;
64 steps = 0;
66 #params: ModelParams;
67 /** Display-only transforms on the oversampled grid; null at 1x. */
68 #displaySht: ShtPlan | null;
69 #oversample: number;
70 /** Wavelength of the seeded random field, and the seed it was drawn from —
71 * kept so changing one can redraw with the other unchanged. */
72 #lam3: number;
73 #seed = 1;
75 private constructor(init: {
76 device: GPUDevice;
77 model: MModel;
78 cfg: ShtConfig;
79 sht: ShtPlan;
80 displaySht: ShtPlan | null;
81 gpu: GpuModel;
82 params: ModelParams;
83 oversample: number;
84 geometry: Geometry;
85 geometryModel: MGeometry;
86 deriv: DerivPlan;
87 niter: number;
88 lam3: number;
89 }) {
90 this.device = init.device;
91 this.model = init.model;
92 this.cfg = init.cfg;
93 this.sht = init.sht;
94 this.gpu = init.gpu;
95 this.npts = init.cfg.nlat * init.cfg.nphi;
96 this.#oversample = init.oversample;
97 this.#params = init.params;
98 this.#displaySht = init.displaySht;
99 this.#geometry = init.geometry;
100 this.#geometryModel = init.geometryModel;
101 this.#deriv = init.deriv;
102 this.niter = init.niter;
103 this.#lam3 = init.lam3;
104 }
106 get geometry(): Geometry {
107 return this.#geometry;
108 }
110 get geometryModel(): MGeometry {
111 return this.#geometryModel;
112 }
114 /** Linear render oversampling factor (1 = read on the solver grid). */
115 get oversample(): number {
116 return this.#oversample;
117 }
119 static async create(opts: ModelSessionOptions): Promise<ModelSession> {
120 const { device, model, params, lmax } = opts;
121 const oversample = Math.max(1, Math.round(opts.oversample ?? 1));
122 const niter = Math.max(0, Math.round(opts.niter ?? 1));
123 const geometryModel = opts.geometry ?? mGeometryByKey(SPHERE_KEY)!;
124 const geometryParams = opts.geometryParams ?? defaultGeometryParams(geometryModel);
125 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
126 const cfg = { lmax, mmax: lmax, nlat, nphi };
127 const sht = await ShtPlan.create(device, cfg);
128 let displaySht: ShtPlan | null = null;
129 let deriv: DerivPlan | null = null;
130 try {
131 // The display plan shares nothing with the solver's beyond the
132 // coefficients copied into it per readback; its grid is the solver's
133 // scaled by the oversampling factor, so nphi stays a power of two (the
134 // FFT path) for power-of-two factors.
135 if (oversample > 1) {
136 displaySht = await ShtPlan.create(device, {
137 lmax,
138 mmax: lmax,
139 nlat: oversample * nlat,
140 nphi: oversample * nphi,
141 });
142 }
143 // Computes the theta/phi derivatives the geometry's metric quantities
144 // (and, per step, the surface Laplace-Beltrami correction) need.
145 deriv = await DerivPlan.create(device, sht);
146 // The surface is built before the model, because the model takes it as
147 // an argument. It is a one-off: compiled, evaluated, read back, and its
148 // plan discarded — nothing of it survives into the timestep but sixteen
149 // buffers of numbers (the embedding, and both metric formulations built
150 // on it: the inverse metric quantities and the flux-form weights).
151 const geometry = await Geometry.create({
152 sht,
153 cfg,
154 source: opts.geometrySource ?? geometryModel.source,
155 paramNames: geometryModel.params.map((p) => p.key),
156 params: geometryParams,
157 deriv,
158 });
159 const gpu = await GpuModel.create({
160 device,
161 sht,
162 cfg,
163 source: opts.source ?? model.source,
164 paramNames: model.params.map((p) => p.key),
165 state: model.state,
166 view: model.species,
167 geometry,
168 deriv,
169 niter,
170 });
171 const lam3 = opts.lam3 ?? DEFAULT_LAMBDA;
172 gpu.setParams({ lam3, ...params });
173 return new ModelSession({
174 device, model, cfg, sht, displaySht, gpu, params, oversample,
175 geometry, geometryModel, deriv, niter, lam3,
176 });
177 } catch (e) {
178 // The transform plans own GPU buffers; do not leak them on a compile error.
179 deriv?.destroy();
180 displaySht?.destroy();
181 sht.destroy();
182 throw e;
183 }
184 }
186 /**
187 * Vertex positions for the current render grid: the surface synthesized on
188 * `viewSht`, interleaved xyz. Exact interpolation of the same coefficients
189 * the solver sees, so the drawn surface is the one being solved on however
190 * finely it is sampled.
191 */
192 renderPositions(): Promise<Float32Array> {
193 return this.#geometry.positionsOn(this.viewSht);
194 }
196 /**
197 * Change the surface in place, without recompiling or disturbing the run.
198 * The geometry's shape in the bindings depends only on the grid, so the
199 * compiled step does not change — only the numbers it reads. The caller
200 * still has to rebuild the mesh from `renderPositions()`.
201 */
202 async setGeometry(
203 geometryModel: MGeometry,
204 params: ModelParams,
205 source?: string,
206 ): Promise<void> {
207 const next = await Geometry.create({
208 sht: this.sht,
209 cfg: this.cfg,
210 source: source ?? geometryModel.source,
211 paramNames: geometryModel.params.map((p) => p.key),
212 params,
213 deriv: this.#deriv,
214 });
215 this.#geometry = next;
216 this.#geometryModel = geometryModel;
217 this.gpu.uploadGeometry(next);
218 // The new surface brings a new preconditioner scale (GpuModel folds its
219 // current geometry's jhat into every params upload).
220 this.gpu.setParams(this.#params);
221 }
223 /** The plan whose grid `readSpecies` samples on — the display plan when
224 * oversampling, otherwise the solver's. Its cosTheta/nphi define the mesh. */
225 get viewSht(): ShtPlan {
226 return this.#displaySht ?? this.sht;
227 }
229 /**
230 * Change the display oversampling in place. Display-only: the simulation
231 * state, time and parameters are untouched, so the run continues seamlessly
232 * on the new render grid. The caller must not have a readSpecies in flight —
233 * its readback maps a buffer of the plan being destroyed.
234 */
235 async setOversample(oversample: number): Promise<void> {
236 const os = Math.max(1, Math.round(oversample));
237 await this.setDisplayGrid(os * this.cfg.nlat, os * this.cfg.nphi);
238 }
240 /**
241 * Point the display plan at an arbitrary grid, rather than an integer
242 * multiple of the solver's. Same contract as setOversample — display-only,
243 * no readback may be in flight — and the same exactness argument, which does
244 * not care about the ratio: the state is band-limited at lmax, so
245 * synthesizing it anywhere is evaluation, not resampling. What this adds is a
246 * grid that need not be *finer*: several sessions at different lmax can be
247 * put on one common grid, which is what makes their fields directly
248 * comparable point by point and lets one mesh serve all of them.
249 */
250 async setDisplayGrid(nlat: number, nphi: number): Promise<void> {
251 const view = this.viewSht.cfg;
252 if (nlat === view.nlat && nphi === view.nphi) return;
253 const onSolverGrid = nlat === this.cfg.nlat && nphi === this.cfg.nphi;
254 const next = onSolverGrid
255 ? null
256 : await ShtPlan.create(this.device, {
257 lmax: this.cfg.lmax,
258 mmax: this.cfg.mmax,
259 nlat,
260 nphi,
261 });
262 const old = this.#displaySht;
263 this.#displaySht = next;
264 this.#oversample = nlat / this.cfg.nlat;
265 old?.destroy();
266 }
268 /**
269 * The coefficient table a model that calls `randnfun3` seeds from — drawn
270 * over the current surface's bounding box, at the wavelength its own .m asked
271 * for — or null for a model that seeds some other way. The draw is host-side
272 * MATLAB (a few ms); the evaluation at every grid point is the GPU kernel
273 * inside `init`.
274 *
275 * Separate from `seed` because the table is a function of space, not of a
276 * grid: drawn once, it is the *same field* wherever it is evaluated, which is
277 * how every variant of a comparison across lmax seeds from one random field
278 * (see src/compare/sharedStart.ts).
279 */
280 drawSeedModes(seed: number): Promise<Float32Array | null> {
281 const lambda = this.gpu.randnfun3Lambda;
282 if (!lambda) return Promise.resolve(null);
283 // Drawn on a worker: at a fine wavelength this is seconds of interpreter
284 // time, and it must not be seconds of frozen page.
285 return drawModesAsync(
286 resolveLambda(lambda, this.#mergedParams()),
287 boundingBox(this.#geometry.x, this.#geometry.y, this.#geometry.z),
288 seed,
289 this.npts,
290 );
291 }
293 /** Run `init` from a seeded perturbation, resetting model time. */
294 async seed(seed: number): Promise<void> {
295 const modes = await this.drawSeedModes(seed);
296 await this.seedWith(seededNoise(this.npts, this.model.seedAmp, seed), modes);
297 this.#seed = seed;
298 }
300 /**
301 * Run `init` from a caller-supplied perturbation, resetting model time.
302 * `seed()` is this with what this session would draw for itself: the host
303 * RNG's field on its own grid, and its own random-field table. Supplying them
304 * instead is how several sessions on *different* grids can be started from
305 * the same initial condition, which is the only way a comparison across lmax
306 * compares one problem rather than two (see src/compare/sharedStart.ts).
307 */
308 async seedWith(noise: Float32Array, modes: Float32Array | null = null): Promise<void> {
309 if (noise.length !== this.npts) {
310 throw new Error(`seedWith: noise must have length ${this.npts} (got ${noise.length})`);
311 }
312 await this.gpu.init(noise, modes);
313 this.t = 0;
314 this.steps = 0;
315 }
317 /**
318 * Push an exact spectral state into the running model, bypassing seeded
319 * init, and reset model time like seed() does. `gpu.step(0)` flips which
320 * of the init/step buffer aliases a read resolves to, without otherwise
321 * touching the state — see GpuModel's `#lastRan`.
322 */
323 loadState(coeffs: Record<string, Float32Array>): void {
324 for (const name of this.model.state) {
325 const data = coeffs[name];
326 if (!data) throw new Error(`loadState: missing state '${name}'`);
327 this.gpu.upload(name, data);
328 }
329 this.gpu.step(0);
330 this.t = 0;
331 this.steps = 0;
332 }
334 /** The model's parameters plus the ones the host owns. */
335 #mergedParams(): ModelParams {
336 return { lam3: this.#lam3, ...this.#params };
337 }
339 setParams(params: ModelParams): void {
340 this.#params = params;
341 this.gpu.setParams(this.#mergedParams());
342 }
344 /** Wavelength of the seeded random field. Redraws on the next seed. */
345 get lam3(): number {
346 return this.#lam3;
347 }
349 /**
350 * Change the random field's wavelength. Nothing recompiles — lam3 is a
351 * uniform and the coefficient table is host-drawn — but the field itself
352 * only changes on the next `seed`, which is where it is drawn.
353 */
354 setLam3(lambda: number): void {
355 if (lambda === this.#lam3) return;
356 this.#lam3 = lambda;
357 this.gpu.setParams(this.#mergedParams());
358 }
360 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
361 step(n = 1): void {
362 this.gpu.step(n);
363 this.t += n * (this.#params.dt ?? 0);
364 this.steps += n;
365 }
367 /**
368 * Wait for the submitted steps to finish, without reading anything back.
369 * This is the honest way to time the solver: a readback would add a GPU->CPU
370 * round trip, which in a browser also crosses a process boundary and can cost
371 * more than the steps themselves.
372 */
373 sync(): Promise<undefined> {
374 return this.device.queue.onSubmittedWorkDone();
375 }
377 /**
378 * Time a batch of `n` steps and return ms/step, leaving the simulation
379 * exactly where it was: the spectral state is snapshotted before the batch
380 * and restored after, and `t`/`steps` do not advance. One sync amortized
381 * over the batch — the same measurement the desktop benchmark makes. The
382 * grid view fields hold the batch's output until the next real step, so
383 * step before reading them.
384 */
385 async measure(n: number): Promise<number> {
386 this.gpu.snapshotState();
387 const t0 = performance.now();
388 this.gpu.step(n);
389 await this.sync();
390 const ms = (performance.now() - t0) / n;
391 this.gpu.restoreState();
392 return ms;
393 }
395 /** Read a named value (a grid field or the spectral state). */
396 read(name: string): Promise<Float32Array> {
397 return this.gpu.read(name);
398 }
400 /**
401 * Read species `k` at render resolution (`viewSht`'s grid): the spectral
402 * state synthesized there. The models define each species as synth of its
403 * state, so this is the field the .m returned — evaluated exactly, whatever
404 * the grid — and it is current however the state last changed, including a
405 * `loadState`, which runs no kernel that would write the grid-space fields.
406 */
407 readSpecies(k: number): Promise<Float32Array> {
408 const state = this.model.state[k];
409 const buf = this.gpu.valueBuffer(state);
410 if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
411 return this.viewSht.synthFrom(buf);
412 }
414 describe(): { init: string[]; step: string[] } {
415 return this.gpu.describe();
416 }
418 destroy(): void {
419 this.gpu.destroy();
420 this.#deriv.destroy();
421 this.#displaySht?.destroy();
422 this.sht.destroy();
423 }
424}