2 * Drawing the pressure field, straight out of the buffer the solver wrote.
3 *
4 * There is no readback in the display path: the fragment shader reads the
5 * solver's storage buffer directly, so a frame costs one draw call and no
6 * GPU-to-CPU round trip. (The host does read a reduction back occasionally, to
7 * decide the colour scale.)
8 *
9 * The picture is a ray march. Each pixel casts one ray, intersects it with the
10 * cube, and steps along it accumulating emission front to back: the pressure
11 * goes through a diverging colormap about zero, and the opacity goes as a
12 * power of |p|, so quiet regions are transparent and the wavefronts are what
13 * you see. The medium is added as a grey emission proportional to how far the
14 * local sound speed departs from the background, which shows a hard scatterer
15 * as a distinct shape and a smooth one as a soft cloud without either needing
16 * its own kind of drawing.
17 *
18 * Two honest limitations. Sampling is nearest-neighbour, not trilinear: the
19 * field lives in a storage buffer rather than a filterable 3D texture, so
20 * trilinear would be eight fetches per sample and the march takes tens of
21 * millions of samples a frame. With the ray step set near the cell size the
22 * difference is visible mainly as a faint stippling on strong wavefronts. And
23 * the compositing is emission only, with no lighting and no shadowing, so what
24 * is behind a strong feature is dimmed but never occluded correctly.
25 *
26 * A clip plane on x is provided because a volume render of a wavefield is
27 * mostly the outside of a wavefield. Pulling the clip in is how you see the
28 * interior, and it is the closest thing here to the flat sibling's picture.
29 */
30import type { ColormapFunc } from './colormaps.ts';
32const SHADER = `
33struct View {
34 eye: vec4f, // xyz: eye position, w: tan(fov/2)
35 right: vec4f, // xyz: camera right, w: aspect ratio
36 up: vec4f, // xyz: camera up, w: domain side L
37 fwd: vec4f, // xyz: camera forward, w: pressure the colormap saturates at
38 bg: vec4f, // rgb: background
39 a: vec4f, // opacity, medium strength, reference speed, speed spread
40 b: vec4f, // ray steps, clip plane x, grid side n, contrast exponent
41 m: vec4f, // xyz: microphone position, w: whether to draw it
42};
44@group(0) @binding(0) var<uniform> V: View;
45@group(0) @binding(1) var<storage, read> p: array<f32>;
46@group(0) @binding(2) var<storage, read> cs: array<f32>;
47@group(0) @binding(3) var cmap: texture_2d<f32>;
48@group(0) @binding(4) var samp: sampler;
50struct VSOut {
51 @builtin(position) pos: vec4f,
52 @location(0) ndc: vec2f,
53};
55@vertex
56fn vs(@builtin(vertex_index) vi: u32) -> VSOut {
57 // One oversized triangle covering the viewport.
58 var xy = array<vec2f, 3>(vec2f(-1.0, -3.0), vec2f(-1.0, 1.0), vec2f(3.0, 1.0));
59 var out: VSOut;
60 let q = xy[vi];
61 out.pos = vec4f(q, 0.0, 1.0);
62 out.ndc = q;
63 return out;
64}
66fn voxel(q: vec3f) -> u32 {
67 let n = i32(V.b.z);
68 let h = V.up.w / f32(n);
69 let g = clamp(vec3i(floor((q + vec3f(0.5 * V.up.w)) / h)), vec3i(0), vec3i(n - 1));
70 return u32(g.x + n * g.y + n * n * g.z);
71}
73// A point on the cube's surface is on an edge when two of its three distances
74// to the bounding planes vanish, so the test is on the median of the three.
75fn edge(q: vec3f, half: f32, L: f32) -> f32 {
76 let d = abs(abs(q) - vec3f(half));
77 let lo = min(d.x, min(d.y, d.z));
78 let hi = max(d.x, max(d.y, d.z));
79 let mid = d.x + d.y + d.z - lo - hi;
80 let w = 0.004 * L;
81 return 1.0 - smoothstep(w, 2.0 * w, mid);
82}
84@fragment
85fn fs(in: VSOut) -> @location(0) vec4f {
86 let L = V.up.w;
87 let half = 0.5 * L;
88 let eye = V.eye.xyz;
89 let dir = normalize(V.fwd.xyz
90 + in.ndc.x * V.right.w * V.eye.w * V.right.xyz
91 + in.ndc.y * V.eye.w * V.up.xyz);
93 let inv = 1.0 / dir;
94 let ta = (vec3f(-half) - eye) * inv;
95 let tb = (vec3f(half) - eye) * inv;
96 let lo = min(ta, tb);
97 let hi = max(ta, tb);
98 let t1 = min(min(hi.x, hi.y), hi.z);
99 var t0 = max(max(lo.x, lo.y), lo.z);
101 var col = V.bg.rgb;
102 if (t1 <= max(t0, 0.0)) { return vec4f(col, 1.0); }
103 t0 = max(t0, 0.0);
105 let lineCol = vec3f(0.42, 0.47, 0.55);
106 col = mix(col, lineCol, 0.5 * edge(eye + dir * t1, half, L));
108 let steps = i32(V.b.x);
109 let dl = (t1 - t0) / f32(steps);
110 // Opacity is quoted per cell, so a longer ray step is proportionally more
111 // opaque and the picture does not change brightness with the quality knob.
112 let unit = dl / (L / V.b.z);
113 let scale = max(V.fwd.w, 1e-20);
114 let grey = vec3f(0.55, 0.58, 0.63);
116 var acc = vec3f(0.0);
117 var alpha = 0.0;
118 for (var k = 0; k < steps; k = k + 1) {
119 if (alpha > 0.995) { break; }
120 let q = eye + dir * (t0 + (f32(k) + 0.5) * dl);
121 if (q.x > V.b.y) { continue; }
122 let i = voxel(q);
123 let v = clamp(p[i] / scale, -1.0, 1.0);
124 var a = pow(abs(v), V.b.w) * V.a.x * unit;
125 var rgb = textureSampleLevel(cmap, samp, vec2f(0.5 + 0.5 * v, 0.5), 0.0).rgb;
126 if (V.a.y > 0.0) {
127 let m = clamp(abs(cs[i] - V.a.z) / max(V.a.w, 1e-20), 0.0, 1.0);
128 let am = m * V.a.y * unit;
129 let tot = a + am;
130 if (tot > 1e-12) { rgb = (a * rgb + am * grey) / tot; }
131 a = tot;
132 }
133 a = clamp(a, 0.0, 1.0);
134 acc = acc + (1.0 - alpha) * a * rgb;
135 alpha = alpha + (1.0 - alpha) * a;
136 }
137 col = acc + (1.0 - alpha) * col;
138 col = mix(col, lineCol, 0.75 * edge(eye + dir * t0, half, L));
140 // The microphone, a dot in a ring at the ray's closest approach to it.
141 // Drawn on top rather than composited into the march, like the 2D app's
142 // white ring: a marker, not a thing in the scene.
143 if (V.m.w > 0.0) {
144 let toM = V.m.xyz - eye;
145 let tm = dot(toM, dir);
146 if (tm > 0.0) {
147 let d = length(toM - tm * dir) / L;
148 let dotm = 1.0 - smoothstep(0.004, 0.007, d);
149 let ring = 1.0 - smoothstep(0.0015, 0.004, abs(d - 0.016));
150 col = mix(col, vec3f(1.0), 0.9 * max(dotm, ring));
151 }
152 }
153 return vec4f(col, 1.0);
154}
155`;
157/** Where the camera is looking from. Angles in radians, distance in metres. */
158export interface Camera {
159 az: number;
160 el: number;
161 dist: number;
162}
164export interface DrawOptions {
165 camera: Camera;
166 /** Pressure the colormap saturates at, in both directions. */
167 scale: number;
168 opacity: number;
169 contrast: number;
170 /** Samples along each ray. */
171 steps: number;
172 /** Everything with x above this is not drawn, in metres. */
173 clipX: number;
174 /** Strength of the medium wash, 0 to turn it off. */
175 medium: number;
176 cref: number;
177 cdev: number;
178 /** Microphone position in metres, or null to not draw it. */
179 mic?: { x: number; y: number; z: number } | null;
180}
182const TAN_HALF_FOV = Math.tan((32 * Math.PI) / 360);
184const cross = (a: number[], b: number[]) => [
185 a[1] * b[2] - a[2] * b[1],
186 a[2] * b[0] - a[0] * b[2],
187 a[0] * b[1] - a[1] * b[0],
188];
189const norm = (a: number[]) => {
190 const m = Math.hypot(a[0], a[1], a[2]) || 1;
191 return [a[0] / m, a[1] / m, a[2] / m];
192};
194export class VolumeView {
195 readonly canvas: HTMLCanvasElement;
197 #device: GPUDevice;
198 #context: GPUCanvasContext;
199 #pipeline: GPURenderPipeline;
200 #layout: GPUBindGroupLayout;
201 #uniform: GPUBuffer;
202 #host = new ArrayBuffer(8 * 16);
203 #sampler: GPUSampler;
204 #cmapTexture: GPUTexture;
205 #bindGroups: GPUBindGroup[] = [];
206 #n = 0;
207 #L = 1;
208 /** Background, matched to the CSS so the cube sits on the page. */
209 bg: [number, number, number] = [0.043, 0.055, 0.071];
211 constructor(device: GPUDevice, canvas: HTMLCanvasElement) {
212 this.#device = device;
213 this.canvas = canvas;
215 const context = canvas.getContext('webgpu');
216 if (!context) throw new Error('this canvas has no WebGPU context');
217 this.#context = context;
218 const format = navigator.gpu.getPreferredCanvasFormat();
219 context.configure({ device, format, alphaMode: 'opaque' });
221 this.#layout = device.createBindGroupLayout({
222 label: 'volume-view',
223 entries: [
224 { binding: 0, visibility: GPUShaderStage.FRAGMENT, buffer: { type: 'uniform' } },
225 { binding: 1, visibility: GPUShaderStage.FRAGMENT, buffer: { type: 'read-only-storage' } },
226 { binding: 2, visibility: GPUShaderStage.FRAGMENT, buffer: { type: 'read-only-storage' } },
227 { binding: 3, visibility: GPUShaderStage.FRAGMENT, texture: { sampleType: 'float' } },
228 { binding: 4, visibility: GPUShaderStage.FRAGMENT, sampler: { type: 'filtering' } },
229 ],
230 });
232 const module = device.createShaderModule({ code: SHADER, label: 'volume-view' });
233 this.#pipeline = device.createRenderPipeline({
234 label: 'volume-view',
235 layout: device.createPipelineLayout({ bindGroupLayouts: [this.#layout] }),
236 vertex: { module, entryPoint: 'vs' },
237 fragment: { module, entryPoint: 'fs', targets: [{ format }] },
238 primitive: { topology: 'triangle-list' },
239 });
241 this.#uniform = device.createBuffer({
242 label: 'volume-view-uniform',
243 size: this.#host.byteLength,
244 usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
245 });
246 this.#sampler = device.createSampler({ magFilter: 'linear', minFilter: 'linear' });
247 this.#cmapTexture = device.createTexture({
248 label: 'colormap',
249 size: [256, 1],
250 format: 'rgba8unorm',
251 usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST,
252 });
253 }
255 /**
256 * Point the view at the buffers of a (new) simulation. Two pressure buffers
257 * are given because the solver alternates between them; `draw` is told which
258 * one is current.
259 */
260 setSource(pressures: [GPUBuffer, GPUBuffer], speed: GPUBuffer, n: number, L: number): void {
261 this.#n = n;
262 this.#L = L;
263 this.#bindGroups = pressures.map((p) =>
264 this.#device.createBindGroup({
265 layout: this.#layout,
266 entries: [
267 { binding: 0, resource: { buffer: this.#uniform } },
268 { binding: 1, resource: { buffer: p } },
269 { binding: 2, resource: { buffer: speed } },
270 { binding: 3, resource: this.#cmapTexture.createView() },
271 { binding: 4, resource: this.#sampler },
272 ],
273 }),
274 );
275 }
277 setColormap(cmap: ColormapFunc): void {
278 const data = new Uint8Array(256 * 4);
279 for (let i = 0; i < 256; i++) {
280 const [r, g, b] = cmap(i / 255);
281 data[4 * i] = r;
282 data[4 * i + 1] = g;
283 data[4 * i + 2] = b;
284 data[4 * i + 3] = 255;
285 }
286 this.#device.queue.writeTexture(
287 { texture: this.#cmapTexture },
288 data,
289 { bytesPerRow: 256 * 4 },
290 { width: 256, height: 1 },
291 );
292 }
294 /** Draw one frame from the pressure buffer with the given index. */
295 draw(which: number, o: DrawOptions): void {
296 const bg = this.#bindGroups[which];
297 if (!bg) return;
299 const { az, el, dist } = o.camera;
300 const ce = Math.cos(el);
301 const toEye = [ce * Math.cos(az), ce * Math.sin(az), Math.sin(el)];
302 const eye = toEye.map((v) => v * dist);
303 const fwd = toEye.map((v) => -v);
304 const right = norm(cross([0, 0, 1], fwd));
305 const up = cross(fwd, right);
307 const rect = this.canvas.width / Math.max(this.canvas.height, 1);
308 const f = new Float32Array(this.#host);
309 f.set([eye[0], eye[1], eye[2], TAN_HALF_FOV], 0);
310 f.set([right[0], right[1], right[2], rect], 4);
311 f.set([up[0], up[1], up[2], this.#L], 8);
312 f.set([fwd[0], fwd[1], fwd[2], o.scale], 12);
313 f.set([this.bg[0], this.bg[1], this.bg[2], 1], 16);
314 f.set([o.opacity, o.medium, o.cref, o.cdev], 20);
315 f.set([o.steps, o.clipX, this.#n, o.contrast], 24);
316 f.set([o.mic?.x ?? 0, o.mic?.y ?? 0, o.mic?.z ?? 0, o.mic ? 1 : 0], 28);
317 this.#device.queue.writeBuffer(this.#uniform, 0, this.#host);
319 const enc = this.#device.createCommandEncoder({ label: 'volume-view' });
320 const pass = enc.beginRenderPass({
321 colorAttachments: [
322 {
323 view: this.#context.getCurrentTexture().createView(),
324 clearValue: { r: this.bg[0], g: this.bg[1], b: this.bg[2], a: 1 },
325 loadOp: 'clear',
326 storeOp: 'store',
327 },
328 ],
329 });
330 pass.setPipeline(this.#pipeline);
331 pass.setBindGroup(0, bg);
332 pass.draw(3);
333 pass.end();
334 this.#device.queue.submit([enc.finish()]);
335 }
337 /**
338 * Match the canvas's backing store to its CSS size. Device pixel ratio is
339 * ignored: the march is the whole cost of a frame and it scales with pixels,
340 * so a retina display would pay four times over for a picture that is
341 * already smooth.
342 */
343 resize(): void {
344 const rect = this.canvas.getBoundingClientRect();
345 const w = Math.max(1, Math.round(rect.width));
346 const h = Math.max(1, Math.round(rect.height));
347 if (this.canvas.width !== w || this.canvas.height !== h) {
348 this.canvas.width = w;
349 this.canvas.height = h;
350 }
351 }
353 destroy(): void {
354 this.#uniform.destroy();
355 this.#cmapTexture.destroy();
356 }
357}