21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 1// Web worker that loads a compiled Stan model (the emscripten module built
2// by the compile server) via tinystan and runs NUTS-HMC sampling. Mirrors
3// stan-playground's StanModelWorker: progress is parsed out of Stan's
4// stdout lines; everything else streams back as console messages.
6import StanModel from 'tinystan';
7import type { Progress, WorkerRequest, WorkerResponse } from './protocol';
9let model: StanModel | undefined;
11function post(message: WorkerResponse): void {
12 self.postMessage(message);
13}
15// The compiled models are threaded emscripten ES6 builds: they spawn their
16// pthread pool with new Worker(new URL('main.js', import.meta.url)), which
17// throws for a cross-origin script (the compile server). Workers cannot be
18// *constructed* from a cross-origin URL, but a module worker may *import*
19// one via CORS — so route cross-origin worker scripts through a same-origin
20// blob trampoline.
21const NativeWorker = Worker;
22(self as { Worker: unknown }).Worker = class extends NativeWorker {
23 constructor(scriptUrl: string | URL, options?: WorkerOptions) {
24 const resolved = new URL(scriptUrl, self.location.href);
25 if (resolved.origin !== self.location.origin) {
26 const blob = new Blob([`import ${JSON.stringify(resolved.href)};`], { type: 'text/javascript' });
27 super(URL.createObjectURL(blob), options);
28 } else {
29 super(scriptUrl, options);
30 }
31 }
32};
34// Stan progress lines look like (spacing varies):
35// Chain [1] Iteration: 2000 / 2000 [100%] (Sampling)
36// Chain [2] Iteration: 800 / 2000 [ 40%] (Warmup)
37// With a single chain the "Chain [x]" prefix is omitted.
38function parseProgress(line: string): Progress {
39 if (line.startsWith('Iteration:')) {
40 line = 'Chain [1] ' + line;
41 }
42 line = line.replace(/\[|\]/g, '');
43 const parts = line.split(/\s+/);
44 return {
45 chain: parseInt(parts[1], 10),
46 iteration: parseInt(parts[3], 10),
47 totalIterations: parseInt(parts[5], 10),
48 percent: parseInt(parts[6].slice(0, -1), 10),
49 warmup: parts[7] === '(Warmup)',
50 };
51}
53function onPrint(text: string): void {
54 if (!text) {
55 return;
56 }
57 if (text.startsWith('Chain') || text.startsWith('Iteration:')) {
58 const report = parseProgress(text);
59 if (Number.isFinite(report.chain) && Number.isFinite(report.iteration)) {
60 post({ type: 'progress', report });
61 return;
62 }
63 }
64 post({ type: 'console', text, level: 'log' });
65}
67function onPrintError(text: string): void {
68 if (text) {
69 post({ type: 'console', text, level: 'error' });
70 }
71}
73self.onmessage = (event: MessageEvent<WorkerRequest>) => {
74 const message = event.data;
75 switch (message.type) {
76 case 'load': {
77 if (!self.crossOriginIsolated) {
78 post({
79 type: 'console',
80 text: 'warning: not cross-origin isolated — SharedArrayBuffer is unavailable and the threaded Stan module may fail to load',
81 level: 'error',
82 });
83 }
84 (async () => {
85 const js = await import(/* @vite-ignore */ message.mainJsUrl);
86 model = await StanModel.load(js.default, onPrint, onPrintError);
87 post({ type: 'loaded', stanVersion: model.stanVersion() });
88 })().catch((error) => {
89 post({ type: 'error', message: `failed to load compiled model: ${error}` });
90 });
91 break;
92 }
93 case 'sample': {
94 if (!model) {
95 post({ type: 'error', message: 'model is not loaded' });
96 return;
97 }
98 try {
99 const { paramNames, draws } = model.sample(message.config);
100 post({ type: 'done', draws, paramNames });
101 } catch (error) {
102 post({ type: 'error', message: String(error) });
103 }
104 break;
105 }
106 }
107};