Add option to download full datasets prior to simulation

src/Geant4Worker/Geant4DatasetDownloadManager.ts

@@ -0,0 +1,123 @@
+// Additional credits:
+// - @kmichalik
+
+import { Dispatch, SetStateAction, useCallback, useEffect, useRef, useState } from 'react';
+
+import Geant4Worker from './Geant4Worker';
+
+export enum DownloadManagerStatus {
+	IDLE,
+	WORKING,
+	FINISHED,
+	ERROR
+}
+
+export enum DatasetDownloadStatus {
+	IDLE,
+	DOWNLOADING,
+	PROCESSING,
+	DONE
+}
+
+const statusTypeMap: Record<string, DatasetDownloadStatus> = {
+	downloading: DatasetDownloadStatus.DOWNLOADING,
+	preparing: DatasetDownloadStatus.PROCESSING,
+	done: DatasetDownloadStatus.DONE
+};
+
+export interface DatasetStatus {
+	name: string;
+	status: DatasetDownloadStatus;
+	done?: number;
+	total?: number;
+}
+
+async function fetchProgress(
+	worker: Geant4Worker,
+	setDatasetStates: Dispatch<SetStateAction<Record<string, DatasetStatus>>>
+) {
+	const progress = await worker.pollDatasetProgress();
+
+	if (progress) {
+		const newDatasetStates: Record<string, DatasetStatus> = {};
+
+		for (const [datasetName, datasetProgress] of Object.entries(progress)) {
+			let status = statusTypeMap[datasetProgress.stage] ?? DatasetDownloadStatus.IDLE;
+
+			newDatasetStates[datasetName] = {
+				name: datasetName,
+				status,
+				done: Math.floor(datasetProgress.progress * 100),
+				total: 100
+			};
+		}
+
+		setDatasetStates(prev => ({ ...prev, ...newDatasetStates }));
+	}
+}
+
+type StartDownloadArgs = {
+	worker: Geant4Worker;
+	setManagerState: Dispatch<SetStateAction<DownloadManagerStatus>>;
+	setDatasetStates: Dispatch<SetStateAction<Record<string, DatasetStatus>>>;
+	setIdle: Dispatch<SetStateAction<boolean>>;
+};
+
+function startDownload({ worker, setManagerState, setDatasetStates, setIdle }: StartDownloadArgs) {
+	const loadDepsPromise = worker.loadDeps();
+
+	const interval = setInterval(async () => {
+		await fetchProgress(worker, setDatasetStates);
+	}, 500);
+
+	loadDepsPromise
+		.then(async () => {
+			clearInterval(interval);
+
+			await fetchProgress(worker, setDatasetStates);
+
+			setManagerState(DownloadManagerStatus.FINISHED);
+			worker.destroy();
+		})
+		.catch(error => {
+			console.error('Dataset download error:', error);
+			setManagerState(DownloadManagerStatus.ERROR);
+			clearInterval(interval);
+		});
+	setManagerState(DownloadManagerStatus.WORKING);
+	setIdle(false);
+}
+
+export function useDatasetDownloadManager() {
+	const [managerState, setManagerState] = useState<DownloadManagerStatus>(
+		DownloadManagerStatus.IDLE
+	);
+	const [datasetStates, setDatasetStates] = useState<Record<string, DatasetStatus>>({});
+	const [idle, setIdle] = useState<boolean>(false);
+	const [worker] = useState<Geant4Worker>(new Geant4Worker());
+	const initCalledRef = useRef(false);
+
+	useEffect(() => {
+		if (initCalledRef.current) return;
+		worker
+			.init()
+			.then(() => setIdle(true))
+			.catch(error => {
+				setManagerState(DownloadManagerStatus.ERROR);
+				console.error('Failed to initialize Geant4 worker for dataset download:', error);
+			});
+		initCalledRef.current = true;
+	}, [worker]);
+
+	const startDownloadSimple = useCallback(() => {
+		if (!idle) return;
+
+		startDownload({ worker, setManagerState, setDatasetStates, setIdle });
+	}, [worker, idle]);
+
+	return {
+		managerState,
+		datasetStates: Object.values(datasetStates),
+		startDownload: startDownloadSimple
+	};
+}

src/Geant4Worker/Geant4Worker.ts

@@ -9,6 +9,7 @@
 import { JobUnknownStatus, StatusState } from '../types/ResponseTypes';
 import {
 	Geant4WorkerMessage,
+	Geant4WorkerMessageDatasetProgress,
 	Geant4WorkerMessageFile,
 	Geant4WorkerMessageType,
 	Geant4WorkerPromise
@@ -99,6 +100,10 @@ export default class Geant4Worker {
 		return this.simulatedPrimaries;
 	}
 
+	getIsInitialized() {
+		return this.isInitialized;
+	}
+
 	init() {
 		if (this.worker) {
 			throw new Error('init() already called on this object');
@@ -118,9 +123,6 @@ export default class Geant4Worker {
 					console.error('From worker: ', message.data);
 
 					break;
-
-				case Geant4WorkerMessageType.DOWNLOAD_STATUS:
-					break;
 				case Geant4WorkerMessageType.PROGRESS:
 					this.simulatedPrimaries = message.data as number;
 
@@ -243,4 +245,16 @@ export default class Geant4Worker {
 
 		return result;
 	}
+
+	async pollDatasetProgress() {
+		if (!this.worker || !this.isInitialized) {
+			console.error('Worker is not initialized. Call init() first.');
+
+			return;
+		}
+
+		return await this.makePromise<Geant4WorkerMessageDatasetProgress>({
+			type: Geant4WorkerMessageType.POLL_DATASET_PROGRESS
+		});
+	}
 }

src/Geant4Worker/Geant4WorkerCallbacks.ts

@@ -46,6 +46,7 @@ const initDatasets: Geant4WorkerCallbacksType = async args => {
 	// and other workers should be able to access the data without downloading again.
 	// This, in theory, could enable the app to run offline.
 	try {
+		downloadTracker.reset();
 		downloadTracker.setCurrentDataset('G4ENSDFSTATE');
 		await initG4ENSDFSTATE(wasmModule);
 		downloadTracker.setCurrentDataset('G4EMLOW');

src/Geant4Worker/Geant4WorkerDownloadProgressMonitor.ts

@@ -5,6 +5,15 @@ type DatasetProgress = Geant4WorkerDatasetProgress & { totalDependencies: number
 export class Geant4WorkerDownloadProgressMonitor {
 	private datasetsProgressTracker: { [key: string]: DatasetProgress } = {};
 	private currentDataset?: string;
+	private downloadOffset: number = 0;
+	private downloadCurrentTotal: number = 0;
+
+	reset() {
+		this.datasetsProgressTracker = {};
+		this.currentDataset = undefined;
+		this.downloadOffset = 0;
+		this.downloadCurrentTotal = 0;
+	}
 
 	setCurrentDataset(dataset: string) {
 		this.datasetsProgressTracker[dataset] = {
@@ -13,10 +22,11 @@ export class Geant4WorkerDownloadProgressMonitor {
 			totalDependencies: 0
 		};
 
+		this.downloadOffset = this.downloadCurrentTotal;
 		this.currentDataset = dataset;
 	}
 
-	setDownloadProgress(progress: number) {
+	setDownloadProgress(loaded: number, total: number) {
 		if (!this.currentDataset) {
 			return;
 		}
@@ -28,7 +38,10 @@ export class Geant4WorkerDownloadProgressMonitor {
 				return;
 			}
 
-			datasetProgress.progress = progress;
+			datasetProgress.stage = 'downloading';
+			datasetProgress.progress =
+				(loaded - this.downloadOffset) / (total - this.downloadOffset);
+			this.downloadCurrentTotal = Math.max(this.downloadCurrentTotal, total);
 		}
 	}
 

src/Geant4Worker/Geant4WorkerPreModule.ts

@@ -14,6 +14,11 @@ export class Geant4WorkerPreModule {
 
 	printErr(...args: any[]) {
 		const data = args.join('');
+
+		if (data.includes('dependency')) return;
+		if (data.includes('still waiting on run dependencies:')) return;
+		if (data.includes('(end of list)')) return;
+
 		workerPostMessage({ type: Geant4WorkerMessageType.PRINT_ERROR, data });
 	}
 
@@ -24,16 +29,13 @@ export class Geant4WorkerPreModule {
 
 	setStatus(text: string) {
 		// Try parsing "text" to find "(x/y)" pattern
-		console.log('Status:', text);
 		const match = text.match(/\((\d+)\/(\d+)\)/);
 
 		if (match) {
 			const current = parseInt(match[1], 10);
 			const total = parseInt(match[2], 10);
 
-			const progress = current / total;
-
-			this.progressMonitor.setDownloadProgress(progress);
+			this.progressMonitor.setDownloadProgress(current, total);
 		}
 	}
 

src/Geant4Worker/Geant4WorkerTypes.ts

@@ -10,8 +10,7 @@ export enum Geant4WorkerMessageType {
 	ERROR,
 	PRINT,
 	PRINT_ERROR,
-	PROGRESS,
-	DOWNLOAD_STATUS
+	PROGRESS
 }
 
 export type Geant4WorkerMessageFile = {

src/WrapperApp/UseRunGeant4LocalWorkerSimulation.tsx

@@ -6,6 +6,7 @@ import { useStore } from '../services/StoreService';
 import { EditorJson } from '../ThreeEditor/js/EditorJson';
 import { SimulatorType } from '../types/RequestTypes';
 import { SimulationInputFiles } from '../types/ResponseTypes';
+import { Geant4DatasetsType } from './components/Simulation/Geant4DatasetDownload';
 import {
 	BatchOptionsType,
 	RunSimulationFunctionType,
@@ -25,7 +26,8 @@ export function useRunGeant4LocalWorkerSimulation(): RunSimulationFunctionType {
 		simName: string,
 		nTasks: number,
 		simulator: SimulatorType,
-		batchOptions: BatchOptionsType
+		batchOptions: BatchOptionsType,
+		geant4DatasetType: Geant4DatasetsType
 	) => {
 		const simData = sourceType === 'editor' ? editorJson : inputFiles;
 
@@ -43,6 +45,7 @@ export function useRunGeant4LocalWorkerSimulation(): RunSimulationFunctionType {
 			simulator,
 			simName,
 			options,
+			geant4DatasetType,
 			controller.signal
 		)
 			.then(res => {

src/WrapperApp/UseRunRemoteWorkerSimulation.tsx

@@ -6,6 +6,7 @@ import { useStore } from '../services/StoreService';
 import { EditorJson } from '../ThreeEditor/js/EditorJson';
 import { SimulatorType } from '../types/RequestTypes';
 import { SimulationInputFiles } from '../types/ResponseTypes';
+import { Geant4DatasetsType } from './components/Simulation/Geant4DatasetDownload';
 import {
 	BatchOptionsType,
 	RunSimulationFunctionType,
@@ -26,7 +27,8 @@ export function useRunRemoteWorkerSimulation(): RunSimulationFunctionType {
 		simName: string,
 		nTasks: number,
 		simulator: SimulatorType,
-		batchOptions: BatchOptionsType
+		batchOptions: BatchOptionsType,
+		geant4DatasetType: Geant4DatasetsType
 	) => {
 		const simData = sourceType === 'editor' ? editorJson : inputFiles;
 
@@ -62,6 +64,7 @@ export function useRunRemoteWorkerSimulation(): RunSimulationFunctionType {
 			simulator,
 			simName,
 			options,
+			geant4DatasetType,
 			controller.signal
 		)
 			.then(res => {

src/WrapperApp/WrapperApp.tsx

@@ -3,6 +3,7 @@ import { styled } from '@mui/material/styles';
 import { SyntheticEvent, useEffect, useState } from 'react';
 
 import { useConfig } from '../config/ConfigService';
+import { useDatasetDownloadManager } from '../Geant4Worker/Geant4DatasetDownloadManager';
 import { useAuth } from '../services/AuthService';
 import { useStore } from '../services/StoreService';
 import { EditorToolbar } from '../ThreeEditor/components/Editor/EditorToolbar';
@@ -21,6 +22,7 @@ import { AboutPanel } from './components/Panels/AboutPanel';
 import { ExamplePanel } from './components/Panels/ExamplePanel';
 import { TabPanel } from './components/Panels/TabPanel';
 import ResultsPanel from './components/Results/ResultsPanel';
+import { Geant4DatasetsType } from './components/Simulation/Geant4DatasetDownload';
 import RunSimulationPanel from './components/Simulation/RunSimulationPanel';
 import SimulationPanel from './components/Simulation/SimulationPanel';
 import { useRunGeant4LocalWorkerSimulation } from './UseRunGeant4LocalWorkerSimulation';
@@ -64,6 +66,12 @@ function WrapperApp() {
 	const [providedInputFiles, setProvidedInputFiles] = useState<SimulationInputFiles>();
 	const [highlightRunForm, setHighLightRunForm] = useState(false);
 
+	const {
+		managerState: geant4DownloadManagerState,
+		datasetStates: geant4DatasetStates,
+		startDownload: geant4DatasetDownload
+	} = useDatasetDownloadManager();
+
 	useEffect(() => {
 		if (Object.keys(providedInputFiles ?? {}).length > 0) {
 			setHighLightRunForm(true);
@@ -119,6 +127,10 @@ function WrapperApp() {
 			? runGeant4LocalWorkerSimulation
 			: runRemoteWorkerSimulation;
 
+	const [geant4DatasetType, setGeant4DatasetType] = useState<Geant4DatasetsType>(
+		Geant4DatasetsType.DOWNLOADED
+	);
+
 	return (
 		<NavDrawerContext value={tabsValue}>
 			<StyledAppGrid>
@@ -228,6 +240,11 @@ function WrapperApp() {
 						clearInputFiles={() => setProvidedInputFiles(undefined)}
 						runSimulation={runSimulation}
 						setSource={setSimulationPanelPresentedSimulationsSource}
+						geant4DownloadManagerState={geant4DownloadManagerState}
+						geant4DatasetDownloadStart={geant4DatasetDownload}
+						geant4DatasetStates={geant4DatasetStates}
+						geant4DatasetType={geant4DatasetType}
+						setGeant4DatasetType={setGeant4DatasetType}
 					/>
 				</TabPanel>
 				{/* end Simulations screen */}