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