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 if (os === this.#oversample) return;
238 const next =
239 os > 1
240 ? await ShtPlan.create(this.device, {
241 lmax: this.cfg.lmax,
242 mmax: this.cfg.mmax,
243 nlat: os * this.cfg.nlat,
244 nphi: os * this.cfg.nphi,
245 })
246 : null;
247 const old = this.#displaySht;
248 this.#displaySht = next;
249 this.#oversample = os;
250 old?.destroy();
251 }
253 /**
254 * Run `init` from a seeded perturbation, resetting model time.
255 *
256 * A model that calls `randnfun3` gets its coefficient table drawn here:
257 * over the current surface's bounding box, at the wavelength its own .m
258 * asked for. The draw is host-side MATLAB (a few ms); the evaluation at
259 * every grid point is the GPU kernel inside `init`.
260 */
261 async seed(seed: number): Promise<void> {
262 const lambda = this.gpu.randnfun3Lambda;
263 // Drawn on a worker: at a fine wavelength this is seconds of interpreter
264 // time, and it must not be seconds of frozen page.
265 const modes = lambda
266 ? await drawModesAsync(
267 resolveLambda(lambda, this.#mergedParams()),
268 boundingBox(this.#geometry.x, this.#geometry.y, this.#geometry.z),
269 seed,
270 this.npts,
271 )
272 : null;
273 await this.gpu.init(seededNoise(this.npts, this.model.seedAmp, seed), modes);
274 this.#seed = seed;
275 this.t = 0;
276 this.steps = 0;
277 }
279 /** The model's parameters plus the ones the host owns. */
280 #mergedParams(): ModelParams {
281 return { lam3: this.#lam3, ...this.#params };
282 }
284 setParams(params: ModelParams): void {
285 this.#params = params;
286 this.gpu.setParams(this.#mergedParams());
287 }
289 /** Wavelength of the seeded random field. Redraws on the next seed. */
290 get lam3(): number {
291 return this.#lam3;
292 }
294 /**
295 * Change the random field's wavelength. Nothing recompiles — lam3 is a
296 * uniform and the coefficient table is host-drawn — but the field itself
297 * only changes on the next `seed`, which is where it is drawn.
298 */
299 setLam3(lambda: number): void {
300 if (lambda === this.#lam3) return;
301 this.#lam3 = lambda;
302 this.gpu.setParams(this.#mergedParams());
303 }
305 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
306 step(n = 1): void {
307 this.gpu.step(n);
308 this.t += n * (this.#params.dt ?? 0);
309 this.steps += n;
310 }
312 /**
313 * Wait for the submitted steps to finish, without reading anything back.
314 * This is the honest way to time the solver: a readback would add a GPU->CPU
315 * round trip, which in a browser also crosses a process boundary and can cost
316 * more than the steps themselves.
317 */
318 sync(): Promise<undefined> {
319 return this.device.queue.onSubmittedWorkDone();
320 }
322 /**
323 * Time a batch of `n` steps and return ms/step, leaving the simulation
324 * exactly where it was: the spectral state is snapshotted before the batch
325 * and restored after, and `t`/`steps` do not advance. One sync amortized
326 * over the batch — the same measurement the desktop benchmark makes. The
327 * grid view fields hold the batch's output until the next real step, so
328 * step before reading them.
329 */
330 async measure(n: number): Promise<number> {
331 this.gpu.snapshotState();
332 const t0 = performance.now();
333 this.gpu.step(n);
334 await this.sync();
335 const ms = (performance.now() - t0) / n;
336 this.gpu.restoreState();
337 return ms;
338 }
340 /** Read a named value (a grid field or the spectral state). */
341 read(name: string): Promise<Float32Array> {
342 return this.gpu.read(name);
343 }
345 /**
346 * Read species `k` at render resolution (`viewSht`'s grid). Without
347 * oversampling this is the grid field the .m returned. With oversampling the
348 * spectral state is synthesized on the finer grid instead — the same field,
349 * since the models define each species as synth of its state, evaluated
350 * exactly on more points.
351 */
352 readSpecies(k: number): Promise<Float32Array> {
353 if (!this.#displaySht) return this.read(this.model.species[k]);
354 const state = this.model.state[k];
355 const buf = this.gpu.valueBuffer(state);
356 if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
357 return this.#displaySht.synthFrom(buf);
358 }
360 describe(): { init: string[]; step: string[] } {
361 return this.gpu.describe();
362 }
364 destroy(): void {
365 this.gpu.destroy();
366 this.#deriv.destroy();
367 this.#displaySht?.destroy();
368 this.sht.destroy();
369 }
370}