use main.ts

flatironinstitute · Jan 11, 2025 · 6726a9c · 6726a9c
1 parent 71655c4
commit 6726a9c
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 86 deletions.
diff --git a/package.json b/package.json
@@ -5,7 +5,7 @@
   "type": "module",
   "scripts": {
     "build": "tsc",
-    "start": "node dist/benchmark.js",
+    "start": "node dist/main.js",
     "benchmark": "npm run build && npm start"
   },
   "keywords": [],

diff --git a/src/benchmark.ts b/src/benchmark.ts
@@ -1,74 +1,49 @@
-import { init, wavedec, waverec, type Wavelet } from "wasmlets";
-import { loadPyodide } from "pyodide";
-import * as fs from "fs/promises";
-import * as path from "path";
-import * as os from "os";
+/* eslint-disable @typescript-eslint/no-explicit-any */
 import DiscreteWavelets from "discrete-wavelets";
+import { loadPyodide } from "pyodide";
+import { init, wavedec, Wavelet, waverec } from "wasmlets";
+import { BenchmarkResult } from "./benchmarkTypes.js";
+
+export async function runBenchmarks(o: {
+  systemInfo: any;
+  sizes?: number[];
+  wavelets?: string[];
+  targetDurationMs?: number;
+  setProgress?: (progress: string) => void;
+  setCurrentTest?: (test: string) => void;
+  setPercentComplete?: (percent: number) => void;
+  pyodideIndexUrl?: string;
+}) {
+  // Set defaults
+  const sizes = o.sizes ?? [1e5, 1e6];
+  const wavelets = o.wavelets ?? ["db2", "db4"];
+  const targetDurationMs = o.targetDurationMs ?? 1000;
+  const { setProgress, setCurrentTest, setPercentComplete } = o;
 
-async function runBenchmarks() {
   // Initialize both libraries
-  console.log("Initializing wasmlets...");
+  setProgress?.("Initializing wasmlets...");
   const preWasmlets = performance.now();
   await init();
   const postWasmlets = performance.now();
   const wasmletsInitTime = postWasmlets - preWasmlets;
   console.log(`wasmlets initialized in ${wasmletsInitTime}ms`);
 
-  console.log("Initializing Pyodide...");
-  // all this does is heat up the cache, so we can measure the time it takes to load the packages
-  // without internet speed affecting the results
-  await loadPyodide().then((pyodide) =>
-    pyodide.loadPackage(["numpy", "pywavelets"])
-  );
-
+  setProgress?.("Initializing Pyodide...");
   const prePyodide = performance.now();
-  const pyodide = await loadPyodide({ packages: ["numpy", "pywavelets"] });
+  const pyodide = await loadPyodide({
+    indexURL: o.pyodideIndexUrl,
+    packages: ["numpy", "pywavelets"]
+  });
   const postPyodide = performance.now();
   const pyodideInitTime = postPyodide - prePyodide;
   console.log(`Pyodide and packages initialized in ${pyodideInitTime}ms`);
 
-  // Benchmark parameters
-  const sizes = [1e5, 1e6];
-  const wavelets: Wavelet[] = ["db2", "db4"];
-  const targetDurationMs = 1000; // Run each benchmark for ~3 seconds
-
-  console.log("Running benchmarks...");
-  interface BenchmarkResult {
-    size: number;
-    wavelet: Wavelet;
-    wasmlets: {
-      timings: {
-        wavedec: number;
-        waverec: number;
-      };
-    };
-    pywavelets: {
-      timings: {
-        wavedec: number;
-        waverec: number;
-      };
-    };
-    pywaveletsWithMarshalling: {
-      timings: {
-        wavedec: number;
-        waverec: number;
-      };
-    };
-    discreteWavelets: {
-      timings: {
-        wavedec: number;
-        waverec?: number;
-      };
-    };
-  }
+  setProgress?.("Running benchmarks...");
 
   const results = {
     timestamp: new Date().toISOString(),
     systemInfo: {
-      hostname: os.hostname(),
-      os: process.platform,
-      arch: process.arch,
-      nodeVersion: process.version,
+      ...o.systemInfo,
       wasmlets: "0.0.3",
       pyodide: pyodide.version,
       discreteWavelets: "5.0.1",
@@ -77,32 +52,37 @@ async function runBenchmarks() {
       wavelets,
       targetDurationMs,
     },
-    info: `Timings are in milliseconds. 'wavedec' and 'waverec' are the average times taken to run the respective functions over ${targetDurationMs}ms.`,
     initializationTimings: {
       wasmlets: wasmletsInitTime,
       pyodide: pyodideInitTime,
     },
+    info: `Timings are in milliseconds. 'wavedec' and 'waverec' are the average times taken to run the respective functions over ${targetDurationMs}ms.`,
     benchmarks: [] as BenchmarkResult[],
   };
 
+  const totalTests = wavelets.length * sizes.length;
+  let completedTests = 0;
+
   for (const wavelet of wavelets) {
     for (const size of sizes) {
+      setCurrentTest?.(`Testing wavelet ${wavelet} with size ${size}`);
+      setPercentComplete?.(Math.round((completedTests / totalTests) * 100));
       // Generate test data
       const data = new Float64Array(size);
       for (let i = 0; i < size; i++) {
         data[i] = Math.sin(i / 10) + Math.random() * 0.1;
       }
 
       //////////////////////////////////////////////////////////////////////////
-      // Wasmlets benchmark
+      // Wavedec benchmark
       console.info(
         `Running wasmlets benchmarks for size ${size} and wavelet ${wavelet}`
       );
       const startDec = performance.now();
       let numDecTrials = 0;
       let coeffs: any;
       while (performance.now() - startDec < targetDurationMs) {
-        coeffs = wavedec(data, wavelet, "sym");
+        coeffs = wavedec(data, wavelet as Wavelet, "sym");
         numDecTrials++;
       }
       const endDec = performance.now();
@@ -113,12 +93,12 @@ async function runBenchmarks() {
       let numRecTrials = 0;
       let x: any;
       while (performance.now() - startRec < targetDurationMs) {
-        x = waverec(coeffs, wavelet, data.length);
+        x = waverec(coeffs, wavelet as Wavelet, data.length);
         numRecTrials++;
       }
       const endRec = performance.now();
       const wasmletsRecAvg = (endRec - startRec) / numRecTrials;
-      console.info(`wavedec: ${wasmletsDecAvg}ms (${numDecTrials} trials)`);
+      console.info(`waverec: ${wasmletsDecAvg}ms (${numDecTrials} trials)`);
       console.info("");
       if (!arraysAreClose(data, x)) {
         throw new Error("Round trip failed");
@@ -238,6 +218,7 @@ json.dumps({
       let numDiscreteDecTrials = 0;
       let discreteCoeffs: any;
       while (performance.now() - startDiscreteDec < targetDurationMs) {
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
         discreteCoeffs = (DiscreteWavelets as any).wavedec(
           regularArray,
           wavelet as any,
@@ -302,36 +283,18 @@ json.dumps({
       };
 
       results.benchmarks.push(wasmletResult);
+      completedTests++;
+      setPercentComplete?.(Math.round((completedTests / totalTests) * 100));
     }
   }
 
-  // Save results
-  const resultsJson = JSON.stringify(results, null, 2);
-  console.log("\nResults:");
-  console.log(resultsJson);
-
-  // Save results
-  const outputDir = "output";
-  await fs.mkdir(outputDir, { recursive: true });
-  const mainOutputPath = path.join(outputDir, "benchmark.json");
-  await fs.writeFile(mainOutputPath, resultsJson);
-  console.log(`\nResults saved to ${mainOutputPath}`);
-
-  // If running in GitHub Actions, save archived copy
-  if (process.env.GITHUB_ACTIONS) {
-    const archiveDir = path.join("benchmark-results", "archive");
-    await fs.mkdir(archiveDir, { recursive: true });
-    const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
-    const archivePath = path.join(archiveDir, `benchmark-${timestamp}.json`);
-    await fs.writeFile(archivePath, resultsJson);
-    console.log(`Archived copy saved to ${archivePath}`);
-  }
+  return results;
 }
 
 const pywaveletsRoundTripDec = async (
   pyodide: any,
   data: Float64Array,
-  wavelet: Wavelet
+  wavelet: string
 ) => {
   pyodide.globals.set("input_data", data);
 
@@ -352,7 +315,7 @@ coeffs
 const pywaveletsRoundTripRec = async (
   pyodide: any,
   coeffs: any,
-  wavelet: Wavelet
+  wavelet: string
 ) => {
   pyodide.globals.set("input_coeffs", coeffs);
 
@@ -383,7 +346,4 @@ const arraysAreClose = (a: Float64Array, b: Float64Array) => {
   }
 
   return true;
-};
-
-// Run benchmarks
-runBenchmarks().catch(console.error);
+};
diff --git a/src/benchmarkTypes.ts b/src/benchmarkTypes.ts
@@ -0,0 +1,28 @@
+export interface BenchmarkResult {
+    size: number;
+    wavelet: string;
+    wasmlets: {
+      timings: {
+        wavedec: number;
+        waverec: number;
+      };
+    };
+    pywavelets: {
+      timings: {
+        wavedec: number;
+        waverec: number;
+      };
+    };
+    pywaveletsWithMarshalling: {
+      timings: {
+        wavedec: number;
+        waverec: number;
+      };
+    };
+    discreteWavelets: {
+      timings: {
+        wavedec: number;
+        waverec?: number;
+      };
+    };
+  }
diff --git a/src/main.ts b/src/main.ts
@@ -0,0 +1,41 @@
+import * as fs from "fs/promises";
+import * as path from "path";
+import os from "os";
+import { runBenchmarks } from "./benchmark.js";
+
+async function main() {
+  const systemInfo = {
+    hostname: os.hostname(),
+    os: process.platform,
+    arch: process.arch,
+    nodeVersion: process.version
+  }
+  const results = await runBenchmarks({
+    systemInfo
+  });
+
+  // Save results
+  const resultsJson = JSON.stringify(results, null, 2);
+  console.log("\nResults:");
+  console.log(resultsJson);
+
+  // Save results
+  const outputDir = "output";
+  await fs.mkdir(outputDir, { recursive: true });
+  const mainOutputPath = path.join(outputDir, "benchmark.json");
+  await fs.writeFile(mainOutputPath, resultsJson);
+  console.log(`\nResults saved to ${mainOutputPath}`);
+
+  // If running in GitHub Actions, save archived copy
+  if (process.env.GITHUB_ACTIONS) {
+    const archiveDir = path.join("benchmark-results", "archive");
+    await fs.mkdir(archiveDir, { recursive: true });
+    const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
+    const archivePath = path.join(archiveDir, `benchmark-${timestamp}.json`);
+    await fs.writeFile(archivePath, resultsJson);
+    console.log(`Archived copy saved to ${archivePath}`);
+  }
+}
+
+// Run benchmarks
+main().catch(console.error);