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
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 137 // plan discarded — nothing of it survives into the timestep but sixteen
138 // buffers of numbers (the embedding, and both metric formulations built
139 // on it: the inverse metric quantities and the flux-form weights).
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 140 const geometry = await Geometry.create({
141 device,
142 sht,
143 cfg,
144 source: opts.geometrySource ?? geometryModel.source,
145 paramNames: geometryModel.params.map((p) => p.key),
146 params: geometryParams,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 147 deriv,
149 const gpu = await GpuModel.create({
150 device,
151 sht,
152 cfg,
153 source: opts.source ?? model.source,
154 paramNames: model.params.map((p) => p.key),
155 state: model.state,
156 view: model.species,
157 geometry,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 158 deriv,
160 });
161 gpu.setParams(params);
162 return new ModelSession({
163 device, model, cfg, sht, displaySht, gpu, params, oversample,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 164 geometry, geometryModel, deriv, niter,
166 } catch (e) {
167 // 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 168 deriv?.destroy();
170 sht.destroy();
171 throw e;
172 }
173 }
175 /**
176 * Vertex positions for the current render grid: the surface synthesized on
177 * `viewSht`, interleaved xyz. Exact interpolation of the same coefficients
178 * the solver sees, so the drawn surface is the one being solved on however
179 * finely it is sampled.
180 */
181 renderPositions(): Promise<Float32Array> {
182 return this.#geometry.positionsOn(this.viewSht);
183 }
185 /**
186 * Change the surface in place, without recompiling or disturbing the run.
187 * The geometry's shape in the bindings depends only on the grid, so the
188 * compiled step does not change — only the numbers it reads. The caller
189 * still has to rebuild the mesh from `renderPositions()`.
190 */
191 async setGeometry(
192 geometryModel: MGeometry,
193 params: ModelParams,
194 source?: string,
195 ): Promise<void> {
196 const next = await Geometry.create({
197 device: this.device,
198 sht: this.sht,
199 cfg: this.cfg,
200 source: source ?? geometryModel.source,
201 paramNames: geometryModel.params.map((p) => p.key),
202 params,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 203 deriv: this.#deriv,
205 this.#geometry = next;
206 this.#geometryModel = geometryModel;
207 this.gpu.uploadGeometry(next);
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 208 // The new surface brings a new preconditioner scale (GpuModel folds its
209 // current geometry's jhat into every params upload).
210 this.gpu.setParams(this.#params);
213 /** The plan whose grid `readSpecies` samples on — the display plan when
214 * oversampling, otherwise the solver's. Its cosTheta/nphi define the mesh. */
215 get viewSht(): ShtPlan {
216 return this.#displaySht ?? this.sht;
217 }
219 /**
220 * Change the display oversampling in place. Display-only: the simulation
221 * state, time and parameters are untouched, so the run continues seamlessly
222 * on the new render grid. The caller must not have a readSpecies in flight —
223 * its readback maps a buffer of the plan being destroyed.
224 */
225 async setOversample(oversample: number): Promise<void> {
226 const os = Math.max(1, Math.round(oversample));
f467ffaCompare several solver settings side by side, on one clockJeremy Magland 227 await this.setDisplayGrid(os * this.cfg.nlat, os * this.cfg.nphi);
228 }
230 /**
231 * Point the display plan at an arbitrary grid, rather than an integer
232 * multiple of the solver's. Same contract as setOversample — display-only,
233 * no readback may be in flight — and the same exactness argument, which does
234 * not care about the ratio: the state is band-limited at lmax, so
235 * synthesizing it anywhere is evaluation, not resampling. What this adds is a
236 * grid that need not be *finer*: several sessions at different lmax can be
237 * put on one common grid, which is what makes their fields directly
238 * comparable point by point and lets one mesh serve all of them.
239 */
240 async setDisplayGrid(nlat: number, nphi: number): Promise<void> {
241 const view = this.viewSht.cfg;
242 if (nlat === view.nlat && nphi === view.nphi) return;
243 const onSolverGrid = nlat === this.cfg.nlat && nphi === this.cfg.nphi;
244 const next = onSolverGrid
245 ? null
246 : await ShtPlan.create(this.device, {
247 lmax: this.cfg.lmax,
248 mmax: this.cfg.mmax,
249 nlat,
250 nphi,
251 });
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 252 const old = this.#displaySht;
253 this.#displaySht = next;
f467ffaCompare several solver settings side by side, on one clockJeremy Magland 254 this.#oversample = nlat / this.cfg.nlat;
256 }
258 /** Run `init` from a seeded perturbation, resetting model time. */
259 seed(seed: number): void {
f467ffaCompare several solver settings side by side, on one clockJeremy Magland 260 this.seedWith(seededNoise(this.npts, this.model.seedAmp, seed));
261 }
263 /**
264 * Run `init` from a caller-supplied perturbation field, resetting model time.
265 * `seed()` is this with the field the host's RNG produces on this session's
266 * grid; supplying the field instead is how several sessions on *different*
267 * grids can be started from the same band-limited initial condition, which is
268 * the only way a comparison across lmax compares one problem rather than two
269 * (see src/compare/sharedStart.ts).
270 */
271 seedWith(noise: Float32Array): void {
272 if (noise.length !== this.npts) {
273 throw new Error(`seedWith: noise must have length ${this.npts} (got ${noise.length})`);
274 }
275 this.gpu.init(noise);
277 this.steps = 0;
278 }
280 setParams(params: ModelParams): void {
281 this.#params = params;
282 this.gpu.setParams(params);
283 }
285 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
286 step(n = 1): void {
287 this.gpu.step(n);
288 this.t += n * (this.#params.dt ?? 0);
289 this.steps += n;
290 }
292 /**
293 * Wait for the submitted steps to finish, without reading anything back.
294 * This is the honest way to time the solver: a readback would add a GPU->CPU
295 * round trip, which in a browser also crosses a process boundary and can cost
296 * more than the steps themselves.
297 */
298 sync(): Promise<undefined> {
299 return this.device.queue.onSubmittedWorkDone();
300 }
302 /**
303 * Time a batch of `n` steps and return ms/step, leaving the simulation
304 * exactly where it was: the spectral state is snapshotted before the batch
305 * and restored after, and `t`/`steps` do not advance. One sync amortized
306 * over the batch — the same measurement the desktop benchmark makes. The
307 * grid view fields hold the batch's output until the next real step, so
308 * step before reading them.
309 */
310 async measure(n: number): Promise<number> {
311 this.gpu.snapshotState();
312 const t0 = performance.now();
313 this.gpu.step(n);
314 await this.sync();
315 const ms = (performance.now() - t0) / n;
316 this.gpu.restoreState();
317 return ms;
318 }
320 /** Read a named value (a grid field or the spectral state). */
321 read(name: string): Promise<Float32Array> {
322 return this.gpu.read(name);
323 }
325 /**
326 * Read species `k` at render resolution (`viewSht`'s grid). Without
327 * oversampling this is the grid field the .m returned. With oversampling the
328 * spectral state is synthesized on the finer grid instead — the same field,
329 * since the models define each species as synth of its state, evaluated
330 * exactly on more points.
331 */
332 readSpecies(k: number): Promise<Float32Array> {
333 if (!this.#displaySht) return this.read(this.model.species[k]);
334 const state = this.model.state[k];
335 const buf = this.gpu.valueBuffer(state);
336 if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
337 return this.#displaySht.synthFrom(buf);
338 }
340 describe(): { init: string[]; step: string[] } {
341 return this.gpu.describe();
342 }
344 destroy(): void {
345 this.gpu.destroy();
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 346 this.#deriv.destroy();
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 347 this.#displaySht?.destroy();
348 this.sht.destroy();
349 }
350}