Graph neural network for epidemic source detection

Our code is a mix of fast C code to simulate epidemic outcomes on static networks that was originally written by Petter Holme (https://github.com/pholme/sir) as well as various Python routines. For the graph neural networks we rely on PyTorch and PyTorch Geometric. Moreover, we provide various batch scripts.

You need to set up a virtual environment. Here are the most important commands to create, use, and delete the virtual env. in which the code is run:

• With pip list you can see what is already installed.
• python3 -m venv .venv to create the virtual environment.
• source .venv/bin/activate to activate the virtual environment.
• deactivate to stop running the virtual env.
• which python to see which python installation is used.
• pip freeze --local > requirements.txt to create a text file with all the requirements.
• rm -rf .venv/ to delete the virtual environment.
• pip install -r requirements.txt to recreate the virual environment according to the specification in the text file.

The file requirements.txt is part of the repository and setting up the virtual env. should thus be fairly easy.

To run the code, you first need to activate the virtual environment with source .venv/bin/activate.

There are different options of how to run the training code. For every option, you can specify all crucial parameters in the corresponding bash script:

• ./train_one_model.sh <MODEL> <N_RUNS> if you just want to run training run for a given network and a given model. Example: ./train_one_model.sh SHAH 3
• ./train_all_models.sh if you want to train all GNN-based models for a given network.
• ./run_ablation.sh if you want to run a full ablation study for a given network.
• ./run_scaling.sh if you want to run training for different numbers of simulations $n$ (per node), with all other parameters fixed.
• ./run_sigma.sh if you want to run training for different levels of uncertainty $\sigma$ (around the true infection rate), with all other parameters fixed.
• ./run_T_range.sh if you want to run training for different durations T, with all other parameters fixed.

To run inference, you currently have two options:

• You can call the inference python routine directly providing it with the path to a trained model (basically the output directory of the training run that is stored in data/). This is done by python3 -u -m sourcedet.run_inference [name of data directory].
• If you want to run inference for several trained models within a directory, you can use ./make_inference.sh.