network-sir-gnn

Mathematical modeling of epidemic spreading on complex networks, combining deterministic network-coupled ODEs, spectral graph theory, and a from-scratch Graph Neural Network that learns the epidemic dynamics directly from simulated data.

Overview

This project is a portfolio evolution of my master's thesis in Biomedical Engineering, Epidemic Spreading in Complex Networks: In Silico Simulation and Machine Learning Approaches, which used stochastic agent-based simulations and a GNN trained for static graph-level regression (predicting epidemic peak size from a single network snapshot).

This project takes a different, complementary angle:

• Deterministic instead of stochastic: epidemic spread is modeled with coupled ordinary differential equations rather than agent-based Monte Carlo simulation.
• Node-level temporal dynamics instead of static graph-level regression: the GNN here learns to predict how each node's state evolves from one timestep to the next, rather than predicting a single scalar outcome from a fixed snapshot.
• Stronger emphasis on the mathematical structure connecting network topology to epidemic behavior, via spectral graph theory.

The project is organized in three parts, each building on the previous one.

Mathematical Background

Part 1 — Network-coupled SIR model

Instead of the classical mean-field SIR model (three scalar quantities S(t), I(t), R(t) assuming homogeneous mixing), each node $i$ in the graph has its own continuous state $(S_i(t), I_i(t), R_i(t))$, coupled through the network's adjacency matrix $A$:

$$ \frac{dS_i}{dt} = -\beta , S_i \sum_j A_{ij} I_j, \qquad \frac{dI_i}{dt} = \beta , S_i \sum_j A_{ij} I_j - \gamma I_i, \qquad \frac{dR_i}{dt} = \gamma I_i $$

This is sometimes called the N-intertwined SIR model (Van Mieghem et al., 2009). On a complete graph, dividing by $N$, it reduces exactly to the classical mean-field SIR equations — verified numerically in the notebook as a sanity check.

Part 2 — Spectral epidemic threshold

Linearizing the system around the disease-free equilibrium shows that the epidemic threshold depends on the principal eigenvalue $\lambda_{\max}$ of the adjacency matrix:

$$ \beta_c = \frac{\gamma}{\lambda_{\max}(A)} $$

The epidemic spreads if $\beta > \beta_c$, and dies out otherwise. This connects network structure directly to epidemic risk: scale-free networks (Barabási–Albert) have hubs that drive $\lambda_{\max}$ up and $\beta_c$ down, making them structurally more vulnerable to outbreaks than equally-sized random networks — even when their average degree is lower.

Part 3 — Learning the dynamics with a GNN

A message-passing Graph Neural Network, implemented from scratch in pure PyTorch (no PyTorch Geometric), is trained to predict each node's state at $t+1$ given its state at $t$ and the graph structure:

$$ H^{(l+1)} = \sigma\left(D^{-1}A , H^{(l)} W^{(l)}_{\text{neigh}} + H^{(l)} W^{(l)}_{\text{self}}\right) $$

This mean-aggregation layer (GraphSAGE-style) is a learned generalization of the fixed physical coupling term $\sum_j A_{ij} I_j$ from Part 1 — the network already told us that structure matters through a fixed weighted neighbor sum; the GNN replaces that fixed rule with one learned from data.

Project Structure

network-sir-gnn/
├── network_sir/
│   ├── graphs.py          # topology generators (Erdős–Rényi, Barabási–Albert, Watts–Strogatz)
│   ├── sir_model.py       # network-coupled SIR ODE + classical mean-field baseline
│   ├── spectral.py        # eigenvalue analysis, epidemic threshold, threshold sweeps
│   └── gnn/
│       ├── dataset.py     # simulation-to-training-pairs pipeline
│       ├── model.py       # from-scratch message-passing GNN
│       └── train.py       # training loop, one-step & rollout evaluation
├── demo.ipynb             # full narrative walkthrough with all results and plots
├── main.py                # CLI demo: Parts 1 & 2 end-to-end
├── train_gnn.py           # CLI demo: Part 3 end-to-end (data gen, training, rollout eval)
├── tests/                 # pytest suite (21 tests covering all modules)
└── checkpoints/           # trained model weights (generated, not tracked in git)

Installation & Usage

Requires Python 3.12 and uv.

git clone https://github.com/FraViss/network-sir-gnn.git
cd network-sir-gnn
uv sync

Run the test suite:

uv run pytest

Run the Parts 1 & 2 CLI demo (network generation, SIR simulation, spectral analysis):

uv run python main.py

Run the full Part 3 pipeline (data generation, GNN training, rollout evaluation — takes a few minutes):

uv run python train_gnn.py

Explore the full narrative with all visualizations:

uv run jupyter notebook demo.ipynb

Key Results

Network topologies and spectral properties

Topology	Nodes	Edges	Avg. Degree	Clustering	$\lambda_{\max}$	$\beta_c$
Erdős–Rényi	200	942	9.42	0.046	10.176	0.0098
Barabási–Albert	200	591	5.91	0.102	10.995	0.0091
Watts–Strogatz	200	600	6.00	0.438	6.154	0.0162

Despite having the lowest average degree, Barabási–Albert has the highest $\lambda_{\max}$ (driven by hub nodes) and consequently the lowest epidemic threshold — the clearest illustration of why scale-free networks are structurally more vulnerable to outbreaks.

Phase transition: epidemic threshold across topologies

Final epidemic size as a function of $\beta$, for each topology, with the theoretical threshold $\beta_c$ marked. Watts–Strogatz's threshold is clearly separated from the other two; Erdős–Rényi and Barabási–Albert have close thresholds (0.0098 vs. 0.0091, reflecting their similar $\lambda_{\max}$), but differ in the sharpness of their transition and in the network's spectral gap once above threshold.

GNN performance

A 3-layer, 2,566-parameter message-passing GNN trained for 60 epochs reaches:

• One-step prediction: overall weighted test MSE of $1.45 \times 10^{-4}$ (ER: $2.1\times 10^{-4}$, BA: $1.4\times 10^{-4}$, WS: $8.3\times 10^{-5}$)
• Autoregressive rollout (Barabási–Albert, $\beta=0.3$): MSE peaks near the epidemic peak ($t \approx 4$) before decaying as the dynamics settle, while conservation of $S_i+I_i+R_i \approx 1$ — never explicitly enforced during training — stays bounded throughout the rollout (max violation $\approx 1.0\times 10^{-2}$).

Known Limitations

• Train/test split: samples are split randomly across all (state_t, state_t+1) pairs. Temporally adjacent pairs from the same simulation can end up on both sides of the split, so the reported test MSE may be slightly optimistic for truly out-of-distribution generalization.
• Rollout stability vs. one-step accuracy: these are distinct properties. In repeated runs with the same random seed but different execution environments, one-step test MSE stayed essentially unchanged, while rollout error near the epidemic peak varied noticeably — small weight differences get amplified autoregressively in the most nonlinear part of the dynamics. Conservation of $S+I+R \approx 1$ was consistently more robust across reruns than the precise trajectory shape.
• No physical constraints during training: the model is trained with plain MSE, with no explicit penalty for violating $S_i + I_i + R_i = 1$. It nonetheless tends to respect this constraint reasonably well — an emergent property worth further study rather than a designed one.

License

MIT — see LICENSE.