/ proof-of-concept / shtns-webgpu
Sign in
proof-of-concept / shtns-webgpu
shtns-webgpu / src / demo / main.ts
191 lines · 6.3 KBCodeBlameHistory
2 * Demo: synthesize a spherical harmonic field on the GPU and render it on
3 * an equirectangular map, with a round-trip (analysis) error readout.
4 */
5import { ShtPlan, requestShtDevice } from '../sht.ts';
6import { randomSpectrum } from '../reference.ts';
7import { lmIndex, type ShtConfig } from '../layout.ts';
9const $ = <T extends HTMLElement>(id: string) => document.getElementById(id) as T;
10const mapCanvas = $<HTMLCanvasElement>('map');
11const statsEl = $('stats');
12const errEl = $('err');
13const hoverEl = $('hover');
15/** Diverging cool-warm colormap (blue — neutral — red), t in [0, 1]. */
16function coolwarm(t: number): [number, number, number] {
17 // endpoints and neutral midpoint (Moreland-style, linearly interpolated)
18 const lo = [59, 76, 192], mid = [221, 221, 221], hi = [180, 4, 38];
19 const a = t < 0.5 ? lo : hi;
20 const u = t < 0.5 ? t * 2 : 2 - t * 2; // 0 at endpoint, 1 at midpoint
21 return [
22 Math.round(a[0] + (mid[0] - a[0]) * u),
23 Math.round(a[1] + (mid[1] - a[1]) * u),
24 Math.round(a[2] + (mid[2] - a[2]) * u),
25 ];
28let device: GPUDevice;
29let plan: ShtPlan | null = null;
30let cfg: ShtConfig;
31let qlm: Float32Array;
32let spat: Float32Array | null = null;
33let vmax = 1;
34let seed = 1;
36function currentConfig(): ShtConfig {
37 const lmax = parseInt($<HTMLSelectElement>('preset').value, 10);
38 const nlat = lmax + 1;
39 let nphi = 4;
40 while (nphi < 2 * lmax + 2) nphi *= 2;
41 return { lmax, mmax: lmax, nlat, nphi };
44function buildSpectrum(): Float32Array {
45 const source = $<HTMLSelectElement>('source').value;
46 const n = 2 * plan!.nlm;
47 if (source === 'single') {
48 const q = new Float32Array(n);
49 const l = Math.min(Math.max(0, $<HTMLInputElement>('inpL').valueAsNumber || 0), cfg.lmax);
50 const m = Math.min(Math.max(0, $<HTMLInputElement>('inpM').valueAsNumber || 0), l);
51 $<HTMLInputElement>('inpL').value = String(l);
52 $<HTMLInputElement>('inpM').value = String(m);
53 q[2 * lmIndex(cfg.lmax, l, m)] = 1;
54 return q;
55 }
56 const q = randomSpectrum(cfg, seed);
57 if (source === 'bandpass') {
58 const lo = Math.floor(cfg.lmax / 3), hi = Math.ceil((2 * cfg.lmax) / 3);
59 for (let m = 0; m <= cfg.mmax; m++)
60 for (let l = m; l <= cfg.lmax; l++) {
61 if (l < lo || l > hi) {
62 const k = lmIndex(cfg.lmax, l, m);
63 q[2 * k] = 0;
64 q[2 * k + 1] = 0;
65 }
66 }
67 }
68 return q;
71function draw() {
72 if (!spat) return;
73 const { nlat, nphi } = cfg;
74 const off = new OffscreenCanvas(nphi, nlat);
75 const octx = off.getContext('2d')!;
76 const img = octx.createImageData(nphi, nlat);
77 vmax = 1e-30;
78 for (let i = 0; i < spat.length; i++) vmax = Math.max(vmax, Math.abs(spat[i]));
79 for (let i = 0; i < nlat; i++) {
80 for (let j = 0; j < nphi; j++) {
81 const v = spat[i * nphi + j];
82 const [r, g, b] = coolwarm(0.5 + (0.5 * v) / vmax);
83 const o = 4 * (i * nphi + j);
84 img.data[o] = r;
85 img.data[o + 1] = g;
86 img.data[o + 2] = b;
87 img.data[o + 3] = 255;
88 }
89 }
90 octx.putImageData(img, 0, 0);
91 const W = Math.min(908, mapCanvas.parentElement!.clientWidth - 2);
92 const H = Math.round(W / 2);
93 mapCanvas.width = W;
94 mapCanvas.height = H;
95 const ctx = mapCanvas.getContext('2d')!;
96 ctx.imageSmoothingEnabled = true;
97 ctx.imageSmoothingQuality = 'high';
98 ctx.drawImage(off, 0, 0, W, H);
100 // colorbar
101 const cb = $<HTMLCanvasElement>('colorbar');
102 const cctx = cb.getContext('2d')!;
103 for (let x = 0; x < cb.width; x++) {
104 const [r, g, b] = coolwarm(x / (cb.width - 1));
105 cctx.fillStyle = `rgb(${r},${g},${b})`;
106 cctx.fillRect(x, 0, 1, cb.height);
107 }
108 $('vmin').textContent = (-vmax).toPrecision(3);
109 $('vmax').textContent = '+' + vmax.toPrecision(3);
112async function recompute() {
113 if (!plan) return;
114 errEl.textContent = '';
115 try {
116 qlm = buildSpectrum();
117 const t0 = performance.now();
118 spat = await plan.synth(qlm);
119 const tSynth = performance.now() - t0;
120 const t1 = performance.now();
121 const qBack = await plan.analys(spat);
122 const tAnalys = performance.now() - t1;
123 let num = 0, den = 0;
124 for (let k = 0; k < qlm.length; k++) {
125 num += (qBack[k] - qlm[k]) ** 2;
126 den += qlm[k] ** 2;
127 }
128 const rel = Math.sqrt(num / (den || 1));
129 draw();
130 statsEl.innerHTML =
131 `grid <b>${cfg.nlat}×${cfg.nphi}</b>, lmax <b>${cfg.lmax}</b>, nlm <b>${plan.nlm}</b>, ` +
132 `Fourier <b>${plan.fourierMode}</b> · synthesis <b>${tSynth.toFixed(1)} ms</b>, ` +
133 `analysis <b>${tAnalys.toFixed(1)} ms</b> (incl. transfers) · ` +
134 `round-trip rel. error <b>${rel.toExponential(2)}</b>`;
135 } catch (e) {
136 errEl.textContent = String(e);
137 }
140async function rebuild() {
141 cfg = currentConfig();
142 plan?.destroy();
143 plan = null;
144 statsEl.textContent = 'building plan…';
145 try {
146 plan = await ShtPlan.create(device, cfg);
147 await recompute();
148 } catch (e) {
149 statsEl.textContent = '';
150 errEl.textContent = String(e);
151 }
154mapCanvas.addEventListener('mousemove', (ev) => {
155 if (!spat || !plan) return;
156 const r = mapCanvas.getBoundingClientRect();
157 const j = Math.min(cfg.nphi - 1, Math.max(0, Math.floor(((ev.clientX - r.left) / r.width) * cfg.nphi)));
158 const i = Math.min(cfg.nlat - 1, Math.max(0, Math.floor(((ev.clientY - r.top) / r.height) * cfg.nlat)));
159 const latDeg = 90 - (plan.theta[i] * 180) / Math.PI;
160 const lonDeg = (j * 360) / cfg.nphi;
161 hoverEl.textContent = `lat ${latDeg.toFixed(1)}°, lon ${lonDeg.toFixed(1)}°: ${spat[i * cfg.nphi + j].toPrecision(4)}`;
162});
163mapCanvas.addEventListener('mouseleave', () => (hoverEl.textContent = ''));
165$('preset').addEventListener('change', rebuild);
166$('source').addEventListener('change', () => {
167 $('lmWrap').style.display = $<HTMLSelectElement>('source').value === 'single' ? '' : 'none';
168 recompute();
169});
170$('inpL').addEventListener('change', recompute);
171$('inpM').addEventListener('change', recompute);
172$('reroll').addEventListener('click', () => {
173 seed = (seed * 16807) % 2147483647 || 1;
174 recompute();
175});
176window.addEventListener('resize', draw);
178async function start() {
179 $('lmWrap').style.display = '';
180 try {
181 device = await requestShtDevice();
182 } catch (e) {
183 statsEl.textContent = '';
184 errEl.textContent =
185 `WebGPU unavailable: ${e instanceof Error ? e.message : e}\n` +
186 'Use a browser with WebGPU support (Chrome/Edge 113+, recent Firefox/Safari).';
187 return;
188 }
189 await rebuild();
191start();
↑↓ move↵openescclose