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';
15export interface ModelSessionOptions {
16 device: GPUDevice;
17 model: MModel;
18 params: ModelParams;
19 lmax: number;
20 /** Override the model source — the editor's working copy. */
21 source?: string;
22}
24export class ModelSession {
25 readonly model: MModel;
26 readonly cfg: ShtConfig;
27 readonly sht: ShtPlan;
28 readonly gpu: GpuModel;
29 readonly npts: number;
31 /** Model time and step count since the last seeding. */
32 t = 0;
33 steps = 0;
35 #params: ModelParams;
37 private constructor(init: {
38 model: MModel;
39 cfg: ShtConfig;
40 sht: ShtPlan;
41 gpu: GpuModel;
42 params: ModelParams;
43 }) {
44 this.model = init.model;
45 this.cfg = init.cfg;
46 this.sht = init.sht;
47 this.gpu = init.gpu;
48 this.npts = init.cfg.nlat * init.cfg.nphi;
49 this.#params = init.params;
50 }
52 static async create(opts: ModelSessionOptions): Promise<ModelSession> {
53 const { device, model, params, lmax } = opts;
54 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
55 const cfg = { lmax, mmax: lmax, nlat, nphi };
56 const sht = await ShtPlan.create(device, cfg);
57 try {
58 const gpu = await GpuModel.create({
59 device,
60 sht,
61 cfg,
62 source: opts.source ?? model.source,
63 paramNames: model.params.map((p) => p.key),
64 state: model.state,
65 view: model.species,
66 });
67 gpu.setParams(params);
68 return new ModelSession({ model, cfg, sht, gpu, params });
69 } catch (e) {
70 // The transform plan owns GPU buffers; do not leak them on a compile error.
71 sht.destroy();
72 throw e;
73 }
74 }
76 /** Run `init` from a seeded perturbation, resetting model time. */
77 seed(seed: number): void {
78 this.gpu.init(seededNoise(this.npts, this.model.seedAmp, seed));
79 this.t = 0;
80 this.steps = 0;
81 }
83 setParams(params: ModelParams): void {
84 this.#params = params;
85 this.gpu.setParams(params);
86 }
88 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
89 step(n = 1): void {
90 this.gpu.step(n);
91 this.t += n * (this.#params.dt ?? 0);
92 this.steps += n;
93 }
95 /** Read a named value (a grid field or the spectral state). */
96 read(name: string): Promise<Float32Array> {
97 return this.gpu.read(name);
98 }
100 describe(): { init: string[]; step: string[] } {
101 return this.gpu.describe();
102 }
104 destroy(): void {
105 this.gpu.destroy();
106 this.sht.destroy();
107 }
108}