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;
 
@@ -159,9 +161,7 @@ export default class Geant4Worker {
 
 	async loadDeps() {
 		if (!this.worker || !this.isInitialized) {
-			console.error('Worker is not initialized. Call init() first.');
-
-			return;
+			throw new Error('Worker is not initialized. Call init() first.');
 		}
 
 		if (this.depsLoaded) {
@@ -177,9 +177,7 @@ export default class Geant4Worker {
 
 	async loadDepsLazy() {
 		if (!this.worker || !this.isInitialized) {
-			console.error('Worker is not initialized. Call init() first.');
-
-			return;
+			throw new Error('Worker is not initialized. Call init() first.');
 		}
 
 		if (this.depsLoaded) {
@@ -195,9 +193,7 @@ export default class Geant4Worker {
 
 	async includeFile(name: string, data: string) {
 		if (!this.worker || !this.isInitialized) {
-			console.error('Worker is not initialized. Call init() first.');
-
-			return;
+			throw new Error('Worker is not initialized. Call init() first.');
 		}
 
 		return await this.makePromise({
@@ -208,9 +204,7 @@ export default class Geant4Worker {
 
 	async start() {
 		if (!this.worker || !this.isInitialized) {
-			console.error('Worker is not initialized. Call init() first.');
-
-			return;
+			throw new Error('Worker is not initialized. Call init() first.');
 		}
 
 		// Worker will acknowledge the message and wait for the deps to load
@@ -227,9 +221,7 @@ export default class Geant4Worker {
 
 	async fetchResultsFile(name: string) {
 		if (!this.worker || this.state !== StatusState.COMPLETED) {
-			console.error('Worker state is invalid.');
-
-			return;
+			throw new Error('Worker is not initialized. Call init() first.');
 		}
 
 		const result = await this.makePromise<Geant4WorkerMessageFile>({
@@ -243,4 +235,14 @@ export default class Geant4Worker {
 
 		return result;
 	}
+
+	async pollDatasetProgress() {
+		if (!this.worker || !this.isInitialized) {
+			throw new Error('Worker is not initialized. Call init() first.');
+		}
+
+		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,12 @@ export class Geant4WorkerDownloadProgressMonitor {
 				return;
 			}
 
-			datasetProgress.progress = progress;
+			datasetProgress.stage = 'downloading';
+			datasetProgress.progress =
+				total - this.downloadOffset === 0
+					? 1
+					: (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 => {