Merge pull request #123 from flatironinstitute/use-stan-sampler-tests

Tests for useStanSampler hooks

gui/src/app/StanSampler/StanModelWorker.ts

@@ -30,12 +30,13 @@ const parseProgress = (msg: string): Progress => {
   if (msg.startsWith("Iteration:")) {
     msg = "Chain [1] " + msg;
   }
+  msg = msg.replace(/\[|\]/g, "");
   const parts = msg.split(/\s+/);
-  const chain = parseInt(parts[1].slice(1, -1));
+  const chain = parseInt(parts[1]);
   const iteration = parseInt(parts[3]);
   const totalIterations = parseInt(parts[5]);
-  const percent = parseInt(parts[7].slice(0, -2));
-  const warmup = parts[8] === "(Warmup)";
+  const percent = parseInt(parts[6].slice(0, -1));
+  const warmup = parts[7] === "(Warmup)";
   return { chain, iteration, totalIterations, percent, warmup };
 };
 
@@ -64,7 +65,7 @@ self.onmessage = (e) => {
               m.stanVersion(),
           );
           self.postMessage({ purpose: Replies.ModelLoaded });
-        });
+        }, console.error);
       break;
     }
     case Requests.Sample: {

gui/src/app/StanSampler/StanSampler.ts

@@ -33,6 +33,7 @@ class StanSampler {
   } {
     const sampler = new StanSampler(compiledUrl);
     const cleanup = () => {
+      console.log("terminating model worker");
       sampler.#worker && sampler.#worker.terminate();
       sampler.#worker = undefined;
     };
@@ -81,38 +82,23 @@ class StanSampler {
   sample(data: any, samplingOpts: SamplingOpts) {
     const refresh = calculateReasonableRefreshRate(samplingOpts);
     const sampleConfig: Partial<SamplerParams> = {
+      ...samplingOpts,
       data,
-      num_chains: samplingOpts.num_chains,
-      num_warmup: samplingOpts.num_warmup,
-      num_samples: samplingOpts.num_samples,
-      init_radius: samplingOpts.init_radius,
       seed: samplingOpts.seed !== undefined ? samplingOpts.seed : null,
       refresh,
     };
     if (!this.#worker) return;
-    if (this.#status === "") {
-      console.warn("Model not loaded yet");
-      return;
-    }
-    if (sampleConfig.num_chains === undefined) {
-      console.warn("Number of chains not specified");
-      return;
-    }
     if (this.#status === "sampling") {
       console.warn("Already sampling");
       return;
     }
-    if (this.#status === "loading") {
-      console.warn("Model not loaded yet");
-      return;
-    }
     this.#samplingOpts = samplingOpts;
     this.#draws = [];
     this.#paramNames = [];
-    this.#worker.postMessage({ purpose: Requests.Sample, sampleConfig });
     this.#samplingStartTimeSec = Date.now() / 1000;
     this.#status = "sampling";
     this.#onStatusChangedCallbacks.forEach((cb) => cb());
+    this.#worker.postMessage({ purpose: Requests.Sample, sampleConfig });
   }
   onProgress(callback: (progress: Progress) => void) {
     this.#onProgressCallbacks.push(callback);

gui/test/app/StanSampler/MockStanModel.ts

@@ -0,0 +1,70 @@
+import { vi } from "vitest";
+
+import type StanModel from "tinystan";
+import type { PrintCallback } from "tinystan";
+import { defaultSamplingOpts } from "../../../src/app/Project/ProjectDataModel";
+
+import fakeURL from "./empty.ts?url";
+import erroringURL from "./fail.ts?url";
+import failSentinel from "./fail.ts";
+
+export const mockCompiledMainJsUrl = fakeURL;
+export const erroringCompiledMainJsUrl = erroringURL;
+
+const erroring_num_chains = 999;
+
+export const erroringSamplingOpts = {
+  ...defaultSamplingOpts,
+  num_chains: erroring_num_chains,
+};
+
+export const mockedParamNames = ["a", "b"];
+export const mockedDraws = [
+  [1, 2],
+  [3, 4],
+];
+
+export const mockedProgress = {
+  chain: 1,
+  iteration: 123,
+  totalIterations: 1000,
+  percent: 45,
+  warmup: false,
+};
+
+const mockedProgressString = `Chain [${mockedProgress.chain}] \
+Iteration: ${mockedProgress.iteration} \
+/ ${mockedProgress.totalIterations} \
+[${mockedProgress.percent.toString().padStart(3)}%]  \
+(${mockedProgress.warmup ? "Warmup" : "Sampling"})`;
+
+const mockedLoad = async (
+  _create: any,
+  printCallback: PrintCallback | null,
+) => {
+  await new Promise((resolve) => setTimeout(resolve, 50));
+
+  if (_create === failSentinel) {
+    return Promise.reject("error for testing in load!");
+  }
+
+  const model = {
+    stanVersion: vi.fn(() => "1.2.3"),
+    sample: vi.fn(({ num_chains }) => {
+      if (num_chains === erroring_num_chains) {
+        throw new Error("error for testing in sample!");
+      }
+
+      printCallback && printCallback(mockedProgressString);
+
+      return {
+        paramNames: mockedParamNames,
+        draws: mockedDraws,
+      };
+    }),
+  } as unknown as StanModel;
+
+  return model;
+};
+
+export default mockedLoad;

gui/test/app/StanSampler/empty.ts

@@ -0,0 +1 @@
+// intentionally empty, used to create a URL vitest can resolve

gui/test/app/StanSampler/fail.ts

@@ -0,0 +1,3 @@
+// used to create a URL vitest can resolve
+const failSentinel = "fail";
+export default failSentinel;