1// WebGPU time-stepping driver of the full (geometrically nonlinear)
2// simulation. The run advances in batches of one movie frame (1/fps s). A
3// batch encodes one compute dispatch per time step (with a per-step uniform
4// addressed by dynamic offset), one retarded-time listener sum over the
5// batch's audio samples, and one extraction of the transverse displacement
6// at the end. Audio, frame, and stick state are read back after every batch;
7// one batch is kept encoded ahead of the one being read back so that the
8// GPU never waits for JavaScript. The recording starts at the impact, with
9// the plate at rest in its stress-free shape.
11import {
12 type Plate,
13 type SimParams,
14 type StandRates,
15 CFL,
16 HIST_MARGIN,
17 massProperties,
18 standRates,
19 strikePatch,
20 strikePoint,
21 thicknessRange,
22 nodeWeights,
23 listenerTerms,
24} from "../params.ts";
25import { type FullLoss, maxAngularFrequency } from "../model/reference.ts";
26import commonWgsl from "./shaders/common.wgsl?raw";
27import stepWgsl from "./shaders/step.wgsl?raw";
28import audioWgsl from "./shaders/audio.wgsl?raw";
29import extractWgsl from "./shaders/extract.wgsl?raw";
31export type Phase = "recording" | "done";
33export interface StickState {
34 x: number;
35 v: number;
36 force: number;
37 maxForce: number;
38 withdrawn: boolean;
39}
41export interface SimStatus {
42 phase: Phase;
43 /** Simulation time since the impact of the state read back so far [s]. */
44 t: number;
45 totalSteps: number;
46 samples: number;
47 frames: number;
48 peakPressure: number;
49 /** Largest transverse displacement seen in the frames [m]. */
50 wmax: number;
51 delayClamped: boolean;
52 stick: StickState;
53}
55const PARAMS_LEN = 40;
56const STEP_LEN = 12;
57const STICK_LEN = 16;
58/** Nodes per workgroup along each axis in step.wgsl. */
59export const TILE = 12;
61interface Pending {
62 staging: GPUBuffer;
63 steps: number;
64 nAudio: number;
65 stickParity: number;
66}
68export class Simulation {
69 readonly plate: Plate;
70 readonly dt: number;
71 readonly stepsPerSample: number;
72 readonly samplesPerBatch: number;
73 readonly stepsPerBatch: number;
74 readonly omegaMax: number;
75 readonly histLen: number;
76 readonly mass: number;
77 readonly stand: StandRates | null;
79 audio: Float32Array = new Float32Array(0);
80 samples = 0;
81 /** Movie frames of w [m]; frame k is at t = k / fps. */
82 frames: Float32Array[] = [];
83 phase: Phase = "recording";
84 totalSteps = 0;
85 peakPressure = 0;
86 wmax = 0;
87 delayClamped = false;
89 private stick: StickState;
90 private inFlight: Pending | null = null;
91 private device: GPUDevice;
92 private params: SimParams;
93 private loss: FullLoss;
94 private paramsBuf: GPUBuffer;
95 private stepBuf: GPUBuffer;
96 private stepStride: number;
97 private stepData: ArrayBuffer;
98 private pos: GPUBuffer[];
99 private vel: GPUBuffer[];
100 private hist: GPUBuffer;
101 private stickBuf: GPUBuffer;
102 private nd: GPUBuffer;
103 private lw: GPUBuffer;
104 private audioOut: GPUBuffer;
105 private frameBuf: GPUBuffer;
106 private staging: GPUBuffer[];
107 private stagingIdx = 0;
108 private cur = 0;
109 private stepPipe: GPUComputePipeline;
110 private audioPipe: GPUComputePipeline;
111 private extractPipe: GPUComputePipeline;
112 private stepBG: GPUBindGroup[];
113 private audioBG: GPUBindGroup;
114 private extractBG: GPUBindGroup[];
116 constructor(device: GPUDevice, plate: Plate, params: SimParams, loss: FullLoss) {
117 this.device = device;
118 this.plate = plate;
119 this.params = { ...params };
120 this.loss = loss;
121 const g = plate.grid, m = plate.mat;
123 // Stable time step, rounded so that an integer number of steps make one audio sample.
124 this.omegaMax = maxAngularFrequency(g, m, thicknessRange(plate).max);
125 const dtMax = (CFL * 2) / this.omegaMax;
126 this.stepsPerSample = Math.ceil(1 / (params.fs * dtMax));
127 this.dt = 1 / (params.fs * this.stepsPerSample);
128 this.samplesPerBatch = Math.max(1, Math.round(params.fs / params.fps));
129 this.stepsPerBatch = this.samplesPerBatch * this.stepsPerSample;
130 this.histLen = this.samplesPerBatch + HIST_MARGIN;
132 const n = g.n;
133 const f32 = (size: number, usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST) =>
134 device.createBuffer({ size: Math.max(16, size * 4), usage });
135 this.pos = [f32(4 * n), f32(4 * n)];
136 this.vel = [f32(3 * n), f32(3 * n)];
137 this.hist = f32(this.histLen * n);
138 this.stickBuf = f32(STICK_LEN);
139 this.stand = standRates(plate);
140 const weights = nodeWeights(g);
141 const packed = new Float32Array(4 * n);
142 for (let k = 0; k < n; k++) {
143 packed[4 * k] = plate.w0[k]; packed[4 * k + 1] = weights[k]; packed[4 * k + 2] = plate.h[k];
144 packed[4 * k + 3] = this.stand ? this.stand.kwBottom[k] : 0;
145 }
146 this.nd = f32(4 * n);
147 device.queue.writeBuffer(this.nd, 0, packed);
148 this.lw = f32(2 * n);
149 this.audioOut = f32(this.samplesPerBatch);
150 this.frameBuf = f32(n);
151 const stagingSize = (this.samplesPerBatch + n + STICK_LEN) * 4;
152 this.staging = [0, 1].map(() => device.createBuffer({ size: stagingSize, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }));
154 const { mass, Ix, Iy } = massProperties(plate);
155 this.mass = mass;
156 const cp = strikePatch(plate, params);
157 let i0 = g.nx, i1 = 0, j0 = g.ny, j1 = 0;
158 for (const k of cp.nodes) {
159 const i = k % g.nx, j = Math.floor(k / g.nx);
160 i0 = Math.min(i0, i); i1 = Math.max(i1, i); j0 = Math.min(j0, j); j1 = Math.max(j1, j);
161 }
162 const stick = new Float32Array(STICK_LEN);
163 stick[1] = -params.strikeVelocity; stick[9] = -params.strikeVelocity;
164 device.queue.writeBuffer(this.stickBuf, 0, stick);
165 this.stick = { x: 0, v: -params.strikeVelocity, force: 0, maxForce: 0, withdrawn: false };
167 const sp = strikePoint(plate, params);
168 const pf = new Float32Array(PARAMS_LEN);
169 const pu = new Uint32Array(pf.buffer);
170 pu[0] = g.nx; pu[1] = g.ny; pu[2] = n; pu[3] = this.histLen;
171 pf[4] = g.hx; pf[5] = g.hy; pf[6] = this.dt; pf[7] = m.nu;
172 pf[8] = m.E / (12 * (1 - m.nu * m.nu)); pf[9] = m.E / (1 - m.nu * m.nu); pf[10] = m.E / (2 * (1 + m.nu)); pf[11] = m.rho;
173 pf[12] = mass; pf[13] = Ix; pf[14] = Iy; pf[15] = params.fs;
174 pf[16] = sp.x; pf[17] = sp.y; pf[18] = params.contactRadius ** 2; pf[19] = cp.mass;
175 pu[20] = cp.nearest; pf[21] = params.contactStiffness; pf[22] = params.contactExponent; pf[23] = params.stickMass;
176 pu[24] = i0; pu[25] = i1; pu[26] = j0; pu[27] = j1;
177 pf[28] = cp.cx; pf[29] = cp.cy;
178 const st = this.stand;
179 pu[30] = st && st.nodes.length > 0 ? 1 : 0;
180 pu[31] = params.padLiftoff ? 1 : 0;
181 pf[32] = st ? st.preloadTop : 0; pf[33] = st ? st.preloadBottom : 0;
182 pf[34] = loss.padDamping;
183 pu[35] = plate.stand?.topPad ? 1 : 0;
184 this.paramsBuf = device.createBuffer({ size: pf.byteLength, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
185 device.queue.writeBuffer(this.paramsBuf, 0, pf);
187 this.stepStride = Math.max(STEP_LEN * 4, device.limits.minUniformBufferOffsetAlignment);
188 this.stepData = new ArrayBuffer(this.stepStride * this.stepsPerBatch);
189 this.stepBuf = device.createBuffer({ size: this.stepData.byteLength, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
191 const shader = (code: string) => device.createShaderModule({ code: commonWgsl + code });
192 const entry = (binding: number, type: GPUBufferBindingType, hasDynamicOffset = false): GPUBindGroupLayoutEntry => ({
193 binding, visibility: GPUShaderStage.COMPUTE, buffer: { type, hasDynamicOffset },
194 });
195 const pipeline = (module: GPUShaderModule, layout: GPUBindGroupLayout) =>
196 device.createComputePipeline({ layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }), compute: { module, entryPoint: "main" } });
197 const stepLayout = device.createBindGroupLayout({
198 entries: [entry(0, "uniform"), entry(1, "uniform", true), entry(2, "read-only-storage"), entry(3, "read-only-storage"), entry(4, "storage"), entry(5, "storage"), entry(6, "storage"), entry(7, "storage"), entry(8, "read-only-storage")],
199 });
200 this.stepPipe = pipeline(shader(stepWgsl), stepLayout);
201 const audioLayout = device.createBindGroupLayout({
202 entries: [entry(0, "uniform"), entry(1, "uniform", true), entry(2, "read-only-storage"), entry(3, "read-only-storage"), entry(4, "storage"), entry(8, "read-only-storage")],
203 });
204 this.audioPipe = pipeline(shader(audioWgsl), audioLayout);
205 const extractLayout = device.createBindGroupLayout({ entries: [entry(0, "uniform"), entry(1, "read-only-storage"), entry(2, "storage"), entry(8, "read-only-storage")] });
206 this.extractPipe = pipeline(shader(extractWgsl), extractLayout);
208 const spb = { buffer: this.stepBuf, size: STEP_LEN * 4 };
209 this.stepBG = [0, 1].map((p) =>
210 device.createBindGroup({
211 layout: stepLayout,
212 entries: [
213 { binding: 0, resource: { buffer: this.paramsBuf } },
214 { binding: 1, resource: spb },
215 { binding: 2, resource: { buffer: this.pos[p] } },
216 { binding: 3, resource: { buffer: this.vel[p] } },
217 { binding: 4, resource: { buffer: this.pos[1 - p] } },
218 { binding: 5, resource: { buffer: this.vel[1 - p] } },
219 { binding: 6, resource: { buffer: this.hist } },
220 { binding: 7, resource: { buffer: this.stickBuf } },
221 { binding: 8, resource: { buffer: this.nd } },
222 ],
223 }),
224 );
225 this.audioBG = device.createBindGroup({
226 layout: audioLayout,
227 entries: [
228 { binding: 0, resource: { buffer: this.paramsBuf } },
229 { binding: 1, resource: spb },
230 { binding: 2, resource: { buffer: this.hist } },
231 { binding: 3, resource: { buffer: this.lw } },
232 { binding: 4, resource: { buffer: this.audioOut } },
233 { binding: 8, resource: { buffer: this.nd } },
234 ],
235 });
236 this.extractBG = [0, 1].map((p) =>
237 device.createBindGroup({
238 layout: extractLayout,
239 entries: [
240 { binding: 0, resource: { buffer: this.paramsBuf } },
241 { binding: 1, resource: { buffer: this.pos[p] } },
242 { binding: 2, resource: { buffer: this.frameBuf } },
243 { binding: 8, resource: { buffer: this.nd } },
244 ],
245 }),
246 );
247 this.setListener(params);
248 this.frames.push(new Float32Array(n));
249 }
251 /** Recomputes the per-node weights and delays of the listener sum; takes effect on later samples. */
252 setListener(p: SimParams): void {
253 const g = this.plate.grid;
254 this.params.listenerDist = p.listenerDist; this.params.listenerAngle = p.listenerAngle; this.params.listenerAzimuth = p.listenerAzimuth;
255 const { weight, delay } = listenerTerms(this.plate, this.params);
256 const data = new Float32Array(2 * g.n);
257 let dmin = Infinity;
258 for (let k = 0; k < g.n; k++) if (weight[k] > 0) dmin = Math.min(dmin, delay[k]);
259 this.delayClamped = false;
260 for (let k = 0; k < g.n; k++) {
261 data[2 * k] = weight[k];
262 let d = weight[k] > 0 ? (delay[k] - dmin) * this.params.fs : 0;
263 if (d > HIST_MARGIN - 3) { d = HIST_MARGIN - 3; this.delayClamped = true; }
264 data[2 * k + 1] = d;
265 }
266 this.device.queue.writeBuffer(this.lw, 0, data);
267 }
269 setDuration(duration: number): void {
270 this.params.duration = duration;
271 if (this.canContinue()) this.phase = "recording";
272 }
273 canContinue(): boolean {
274 return this.phase === "done" && this.samples < this.targetSamples();
275 }
276 stop(): void {
277 this.phase = "done";
278 }
279 targetSamples(): number {
280 return Math.round(this.params.duration * this.params.fs);
281 }
282 status(): SimStatus {
283 const ahead = this.inFlight ? this.inFlight.steps * this.dt : 0;
284 return { phase: this.phase, t: this.totalSteps * this.dt - ahead, totalSteps: this.totalSteps, samples: this.samples, frames: this.frames.length, peakPressure: this.peakPressure, wmax: this.wmax, delayClamped: this.delayClamped, stick: this.stick };
285 }
286 recordedAudio(): Float32Array {
287 return this.audio.subarray(0, this.samples);
288 }
289 private stepWorkgroups(): [number, number] {
290 return [Math.ceil(this.plate.grid.nx / TILE), Math.ceil(this.plate.grid.ny / TILE)];
291 }
292 private appendAudio(raw: Float32Array): void {
293 const need = this.samples + raw.length;
294 if (this.audio.length < need) {
295 const a = new Float32Array(Math.max(need, 2 * this.audio.length, this.params.fs));
296 a.set(this.audio.subarray(0, this.samples));
297 this.audio = a;
298 }
299 this.audio.set(raw, this.samples);
300 for (let k = 0; k < raw.length; k++) this.peakPressure = Math.max(this.peakPressure, Math.abs(raw[k]));
301 this.samples = need;
302 }
303 private writeStepParams(s: number, fields: { flags: number; slot: number; outIdx: number }): void {
304 const o = (s * this.stepStride) / 4;
305 const u32 = new Uint32Array(this.stepData), i32 = new Int32Array(this.stepData), f32 = new Float32Array(this.stepData);
306 u32[o] = this.totalSteps + s;
307 u32[o + 1] = fields.flags;
308 i32[o + 2] = fields.slot;
309 i32[o + 3] = fields.outIdx;
310 f32[o + 4] = this.loss.sigma0;
311 f32[o + 5] = this.loss.sigma1;
312 f32[o + 6] = this.loss.sigmaIn;
313 u32[o + 7] = (this.totalSteps + s) & 1;
314 }
316 /** Runs one batch (one movie frame of time); resolves when its results have been read back. */
317 async runBatch(): Promise<void> {
318 if (this.phase === "done") return;
319 const cur = this.inFlight ?? this.encodeBatch();
320 this.inFlight = null;
321 if (!cur) { this.phase = "done"; return; }
322 if (this.stepsLeft() > 0) this.inFlight = this.encodeBatch();
323 await this.collect(cur);
324 }
325 private stepsLeft(): number {
326 return this.targetSamples() * this.stepsPerSample - this.totalSteps;
327 }
328 private encodeBatch(): Pending | null {
329 const steps = Math.min(this.stepsPerBatch, Math.max(0, this.stepsLeft()));
330 if (steps === 0) return null;
331 const audioSteps: number[] = [];
332 for (let s = 0; s < steps; s++) {
333 const local = this.totalSteps + s;
334 let slot = -1, outIdx = -1;
335 if (local % this.stepsPerSample === 0) {
336 slot = (local / this.stepsPerSample) % this.histLen;
337 outIdx = audioSteps.length;
338 audioSteps.push(s);
339 }
340 this.writeStepParams(s, { flags: 1, slot, outIdx });
341 }
342 this.device.queue.writeBuffer(this.stepBuf, 0, this.stepData, 0, steps * this.stepStride);
343 const CHUNK = 1024;
344 let enc = this.device.createCommandEncoder();
345 let pass = enc.beginComputePass();
346 const [wgX, wgY] = this.stepWorkgroups();
347 const n = this.plate.grid.n;
348 let cur = this.cur;
349 pass.setPipeline(this.stepPipe);
350 for (let s = 0; s < steps; s++) {
351 if (s > 0 && s % CHUNK === 0) {
352 pass.end();
353 this.device.queue.submit([enc.finish()]);
354 enc = this.device.createCommandEncoder();
355 pass = enc.beginComputePass();
356 pass.setPipeline(this.stepPipe);
357 }
358 pass.setBindGroup(0, this.stepBG[cur], [s * this.stepStride]);
359 pass.dispatchWorkgroups(wgX, wgY);
360 cur = 1 - cur;
361 }
362 const nAudio = audioSteps.length;
363 if (nAudio > 0) {
364 pass.setPipeline(this.audioPipe);
365 pass.setBindGroup(0, this.audioBG, [audioSteps[0] * this.stepStride]);
366 pass.dispatchWorkgroups(nAudio);
367 }
368 pass.setPipeline(this.extractPipe);
369 pass.setBindGroup(0, this.extractBG[cur]);
370 pass.dispatchWorkgroups(Math.ceil(n / 256));
371 pass.end();
372 const staging = this.staging[this.stagingIdx];
373 this.stagingIdx = 1 - this.stagingIdx;
374 enc.copyBufferToBuffer(this.audioOut, 0, staging, 0, Math.max(4, nAudio * 4));
375 enc.copyBufferToBuffer(this.frameBuf, 0, staging, this.samplesPerBatch * 4, n * 4);
376 enc.copyBufferToBuffer(this.stickBuf, 0, staging, (this.samplesPerBatch + n) * 4, STICK_LEN * 4);
377 this.device.queue.submit([enc.finish()]);
378 this.cur = cur;
379 this.totalSteps += steps;
380 return { staging, steps, nAudio, stickParity: this.totalSteps & 1 };
381 }
382 private async collect(b: Pending): Promise<void> {
383 const n = this.plate.grid.n;
384 await b.staging.mapAsync(GPUMapMode.READ);
385 const view = new Float32Array(b.staging.getMappedRange());
386 const base = this.samplesPerBatch + n + 8 * b.stickParity;
387 this.stick = { x: view[base], v: view[base + 1], force: view[base + 2], maxForce: view[base + 3], withdrawn: view[base + 5] !== 0 };
388 this.appendAudio(view.slice(0, b.nAudio));
389 const frame = view.slice(this.samplesPerBatch, this.samplesPerBatch + n);
390 for (let k = 0; k < n; k++) this.wmax = Math.max(this.wmax, Math.abs(frame[k]));
391 this.frames.push(frame);
392 b.staging.unmap();
393 if (this.samples >= this.targetSamples()) this.phase = "done";
394 }
396 // ---------- Testing support ----------
398 /** Advances the state by exactly one free step (no drumstick, no audio) with the losses and the stand. */
399 async stepOnce(): Promise<void> {
400 this.writeStepParams(0, { flags: 0, slot: -1, outIdx: -1 });
401 this.device.queue.writeBuffer(this.stepBuf, 0, this.stepData, 0, this.stepStride);
402 const enc = this.device.createCommandEncoder();
403 const pass = enc.beginComputePass();
404 pass.setPipeline(this.stepPipe);
405 pass.setBindGroup(0, this.stepBG[this.cur], [0]);
406 pass.dispatchWorkgroups(...this.stepWorkgroups());
407 pass.end();
408 this.device.queue.submit([enc.finish()]);
409 this.cur = 1 - this.cur;
410 this.totalSteps++;
411 await this.device.queue.onSubmittedWorkDone();
412 }
413 /** Overwrites the state: pos as [u | v | w | wLo], vel as [du | dv | dw]. */
414 writeState(pos: Float32Array<ArrayBuffer>, vel: Float32Array<ArrayBuffer>): void {
415 const n = this.plate.grid.n;
416 if (pos.length !== 4 * n || vel.length !== 3 * n) throw new Error("state arrays have the wrong length");
417 this.device.queue.writeBuffer(this.pos[this.cur], 0, pos);
418 this.device.queue.writeBuffer(this.vel[this.cur], 0, vel);
419 }
420 /** Reads back the full state as [u | v | w | wLo] and [du | dv | dw]. */
421 async readState(): Promise<{ pos: Float32Array; vel: Float32Array }> {
422 const n = this.plate.grid.n;
423 const stg = this.device.createBuffer({ size: 7 * n * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
424 const enc = this.device.createCommandEncoder();
425 enc.copyBufferToBuffer(this.pos[this.cur], 0, stg, 0, 4 * n * 4);
426 enc.copyBufferToBuffer(this.vel[this.cur], 0, stg, 4 * n * 4, 3 * n * 4);
427 this.device.queue.submit([enc.finish()]);
428 await stg.mapAsync(GPUMapMode.READ);
429 const all = new Float32Array(stg.getMappedRange()).slice();
430 stg.unmap();
431 stg.destroy();
432 return { pos: all.subarray(0, 4 * n), vel: all.subarray(4 * n) };
433 }
435 destroy(): void {
436 this.inFlight = null;
437 for (const b of [...this.pos, ...this.vel, this.hist, this.stickBuf, this.nd, this.lw, this.audioOut, this.frameBuf, this.paramsBuf, this.stepBuf, ...this.staging]) b.destroy();
438 }
439 get fs(): number { return this.params.fs; }
440 get fps(): number { return this.params.fps; }
441 get simParams(): SimParams { return this.params; }
442}