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