Allow analysis scripts to read data.json

flatironinstitute · Nov 4, 2024 · 107bf5e · 107bf5e
1 parent 5c88cc6
commit 107bf5e
Show file tree

Hide file tree

Showing 10 changed files with 74 additions and 25 deletions.
diff --git a/gui/src/app/Scripting/Analysis/AnalysisPyWindow.tsx b/gui/src/app/Scripting/Analysis/AnalysisPyWindow.tsx
@@ -34,6 +34,7 @@ const AnalysisPyWindow: FunctionComponent<AnalysisWindowProps> = ({
     onStatus,
     runnable,
     notRunnableReason,
+    files,
   } = useAnalysisState(latestRun);
 
   const callbacks = useMemo(
@@ -51,13 +52,18 @@ const AnalysisPyWindow: FunctionComponent<AnalysisWindowProps> = ({
   const handleRun = useCallback(
     (code: string) => {
       clearOutputDivs(consoleRef, imagesRef);
-      run(code, spData, {
-        loadsDraws: true,
-        showsPlots: true,
-        producesData: false,
+      run({
+        code,
+        spData,
+        spRunSettings: {
+          loadsDraws: true,
+          showsPlots: true,
+          producesData: false,
+        },
+        files,
       });
     },
-    [consoleRef, imagesRef, run, spData],
+    [consoleRef, imagesRef, run, spData, files],
   );
 
   const contentOnEmpty = useTemplatedFillerText(

diff --git a/gui/src/app/Scripting/Analysis/AnalysisRWindow.tsx b/gui/src/app/Scripting/Analysis/AnalysisRWindow.tsx
@@ -26,16 +26,18 @@ const AnalysisRWindow: FunctionComponent<AnalysisWindowProps> = ({
     onStatus,
     runnable,
     notRunnableReason,
+    files,
   } = useAnalysisState(latestRun);
 
   const { run } = useWebR({ consoleRef, imagesRef, onStatus });
+
   const handleRun = useCallback(
     async (userCode: string) => {
       clearOutputDivs(consoleRef, imagesRef);
       const code = loadDrawsCode + userCode;
-      await run({ code, spData });
+      await run({ code, spData, files });
     },
-    [consoleRef, imagesRef, run, spData],
+    [consoleRef, imagesRef, run, spData, files],
   );
 
   const contentOnEmpty = useTemplatedFillerText(

diff --git a/gui/src/app/Scripting/Analysis/useAnalysisState.ts b/gui/src/app/Scripting/Analysis/useAnalysisState.ts
@@ -1,3 +1,4 @@
+import { FileNames } from "@SpCore/FileMapping";
 import {
   InterpreterStatus,
   isInterpreterBusy,
@@ -44,6 +45,16 @@ const useAnalysisState = (latestRun: StanRun) => {
 
   const isDataDefined = useMemo(() => spData !== undefined, [spData]);
 
+  const files = useMemo(() => {
+    if (latestRun.data === undefined) {
+      return undefined;
+    } else {
+      return {
+        [FileNames.DATAFILE]: latestRun.data,
+      };
+    }
+  }, [latestRun.data]);
+
   useEffect(() => {
     if (!isDataDefined) {
       setRunnable(false);
@@ -65,6 +76,7 @@ const useAnalysisState = (latestRun: StanRun) => {
     onStatus: setStatus,
     runnable,
     notRunnableReason,
+    files,
   };
 };
 

diff --git a/gui/src/app/Scripting/DataGeneration/DataPyWindow.tsx b/gui/src/app/Scripting/DataGeneration/DataPyWindow.tsx
@@ -39,15 +39,14 @@ const DataPyWindow: FunctionComponent<Props> = () => {
   const handleRun = useCallback(
     (code: string) => {
       clearOutputDivs(consoleRef);
-      run(
+      run({
         code,
-        {},
-        {
+        spRunSettings: {
           loadsDraws: false,
           showsPlots: false,
           producesData: true,
         },
-      );
+      });
     },
     [consoleRef, run],
   );

diff --git a/gui/src/app/Scripting/pyodide/pyodideWorker.ts b/gui/src/app/Scripting/pyodide/pyodideWorker.ts
@@ -67,13 +67,14 @@ self.onmessage = async (e) => {
     return;
   }
   const message = e.data;
-  await run(message.code, message.spData, message.spRunSettings);
+  await run(message.code, message.spData, message.spRunSettings, message.files);
 };
 
 const run = async (
   code: string,
   spData: Record<string, any> | undefined,
   spPySettings: PyodideRunSettings,
+  files: Record<string, string> | undefined,
 ) => {
   setStatus("loading");
   try {
@@ -113,12 +114,25 @@ const run = async (
         await f;
       }
 
+      if (files) {
+        const encoder = new TextEncoder();
+        for (const [filename, content] of Object.entries(files)) {
+          await pyodide.FS.writeFile(filename, encoder.encode(content + "\n"));
+        }
+      }
+
       setStatus("running");
       pyodide.runPython(script, { globals });
       succeeded = true;
     } catch (e: any) {
       console.error(e);
       sendStderr(e.toString());
+    } finally {
+      if (files) {
+        for (const filename of Object.keys(files)) {
+          await pyodide.FS.unlink(filename);
+        }
+      }
     }
 
     if (spPySettings.producesData) {

diff --git a/gui/src/app/Scripting/pyodide/pyodideWorkerTypes.ts b/gui/src/app/Scripting/pyodide/pyodideWorkerTypes.ts
@@ -15,6 +15,7 @@ export type MessageToPyodideWorker = {
   code: string;
   spData: Record<string, any> | undefined;
   spRunSettings: PyodideRunSettings;
+  files: Record<string, string> | undefined;
 };
 
 export const isMessageToPyodideWorker = (

diff --git a/gui/src/app/Scripting/pyodide/usePyodideWorker.ts b/gui/src/app/Scripting/pyodide/usePyodideWorker.ts
@@ -18,6 +18,13 @@ type PyodideWorkerCallbacks = {
   onImage?: (image: string) => void;
 };
 
+type RunPyProps = {
+  code: string;
+  spData?: Record<string, any>;
+  spRunSettings: PyodideRunSettings;
+  files?: Record<string, string>;
+};
+
 class PyodideWorkerInterface {
   #worker: Worker | undefined;
 
@@ -71,16 +78,13 @@ class PyodideWorkerInterface {
     return { worker, cleanup };
   }
 
-  run(
-    code: string,
-    spData: Record<string, any> | undefined,
-    spRunSettings: PyodideRunSettings,
-  ) {
+  run({ code, spData, spRunSettings, files }: RunPyProps) {
     const msg: MessageToPyodideWorker = {
       type: "run",
       code,
       spData,
       spRunSettings,
+      files,
     };
     if (this.#worker) {
       this.#worker.postMessage(msg);
@@ -114,13 +118,9 @@ const usePyodideWorker = (callbacks: {
   }, [callbacks]);
 
   const run = useCallback(
-    (
-      code: string,
-      spData: Record<string, any> | undefined,
-      spRunSettings: PyodideRunSettings,
-    ) => {
+    (p: RunPyProps) => {
       if (worker) {
-        worker.run(code, spData, spRunSettings);
+        worker.run(p);
       } else {
         throw new Error("pyodide worker is not defined");
       }

diff --git a/gui/src/app/Scripting/webR/useWebR.ts b/gui/src/app/Scripting/webR/useWebR.ts
@@ -23,6 +23,7 @@ type useWebRProps = {
 type RunRProps = {
   code: string;
   spData?: Record<string, any>;
+  files?: Record<string, string>;
 };
 
 const useWebR = ({ imagesRef, consoleRef, onStatus, onData }: useWebRProps) => {
@@ -63,7 +64,7 @@ const useWebR = ({ imagesRef, consoleRef, onStatus, onData }: useWebRProps) => {
   }, [consoleRef, imagesRef, onStatus, webR]);
 
   const run = useCallback(
-    async ({ code, spData }: RunRProps) => {
+    async ({ code, spData, files }: RunRProps) => {
       try {
         const webR = await loadWebRInstance();
         const shelter = await new webR.Shelter();
@@ -86,6 +87,12 @@ const useWebR = ({ imagesRef, consoleRef, onStatus, onData }: useWebRProps) => {
 
           const options = { ...captureOutputOptions, env };
 
+          if (files) {
+            const encoder = new TextEncoder();
+            for (const [name, content] of Object.entries(files)) {
+              await webR.FS.writeFile(name, encoder.encode(content + "\n"));
+            }
+          }
           // setup
           await webR.evalRVoid(webRPreamble, options);
           await webR.evalRVoid(code, options);
@@ -98,6 +105,11 @@ const useWebR = ({ imagesRef, consoleRef, onStatus, onData }: useWebRProps) => {
           }
         } finally {
           shelter.purge();
+          if (files) {
+            for (const [name] of Object.entries(files)) {
+              await webR.FS.unlink(name);
+            }
+          }
         }
         onStatus("completed");
       } catch (e: any) {

diff --git a/gui/src/app/StanSampler/StanSampler.ts b/gui/src/app/StanSampler/StanSampler.ts
@@ -100,7 +100,7 @@ class StanSampler {
     };
     if (!this.#stanWorker) throw new Error("model worker is undefined");
 
-    this.update({ type: "startSampling", samplingOpts });
+    this.update({ type: "startSampling", samplingOpts, data });
 
     this.#samplingStartTimeSec = Date.now() / 1000;
     this.postMessage({ purpose: Requests.Sample, sampleConfig });

diff --git a/gui/src/app/StanSampler/useStanSampler.ts b/gui/src/app/StanSampler/useStanSampler.ts
@@ -11,6 +11,7 @@ export type StanRun = {
   errorMessage: string;
   progress?: Progress;
   samplingOpts?: SamplingOpts;
+  data?: string;
   draws?: number[][];
   paramNames?: string[];
   computeTimeSec?: number;
@@ -35,6 +36,7 @@ export type StanRunAction =
   | {
       type: "startSampling";
       samplingOpts: SamplingOpts;
+      data: string;
     }
   | {
       type: "samplerReturn";
@@ -65,6 +67,7 @@ export const StanRunReducer = (
         status: "sampling",
         errorMessage: "",
         samplingOpts: action.samplingOpts,
+        data: action.data,
       };
     case "samplerReturn":
       return {