SIR

The SIR Model [1]\(^,\) [2] assumes that infection spreads only through links between neighboring nodes in a graph \(G=(V,E)\). Each node is in one of three states: \(S\), \(I\), or \(R\). A susceptible node \(i\) is infected by its infected neighbors \(j \in N(i)\) with rate \(\beta\), while infected nodes recover with rate \(\gamma\):

\[\begin{split}\begin{aligned} \frac{dS_i}{dt} &= -\beta \sum_{j\in N(i)} S_i I_j, \\ \frac{dI_i}{dt} &= \beta \sum_{j\in N(i)} S_i I_j - \gamma I_i, \\ \frac{dR_i}{dt} &= \gamma I_i . \end{aligned}\end{split}\]

Implementation

Node transitions follow two rules: 1) if a \(S\) state node has \(I\) state neighbors, each \(I\) state neighbor transmits the infection to the \(S\) state node with probability \(\beta\); 2) the \(I\) state node is recovered to the \(R\) state with probability \(\gamma\).

Node states are represented by two Boolean indicator vectors \(h, r \in \{0,1\}^N\), where \(h_i=1\) denotes infected, \(r_i=1\) denotes recovered, and \((h_i,r_i)=(0,0)\) denotes susceptible. The update of the system at step \(k\) is decomposed into three stages:

1. Each infected neighbor \(j\) of node \(i\) transmits a log-probability contribution

\[m_{ji}^{(k)} = h_j^{(k-1)} \cdot \log(1-\beta)\]

2. Node \(i\) collects contributions from all neighbors \(N(i)\) to compute its infection probability

\[m_i^{(k)} = 1 - \exp\!\left( \sum_{j \in N(i)} m_{ji}^{(k)} \right)\]

3. The indicator variables are updated with independent uniform random variables \(U_i^{\mathrm{inf}}, U_i^{\mathrm{rec}} \sim \mathrm{Uniform}(0,1)\)

\[\begin{split}\begin{aligned} h_i^{(k)} &= \begin{cases} 1, & \text{if } (h_i^{(k-1)}=0 \land r_i^{(k-1)}=0) \land (U_i^{\mathrm{inf}} < m_i^{(k)}), \\[4pt] 0, & \text{if } h_i^{(k-1)}=1 \land (U_i^{\mathrm{rec}} \ge \gamma), \\[4pt] h_i^{(k-1)}, & \text{otherwise}, \end{cases} \\[6pt] r_i^{(k)} &= \begin{cases} 1, & \text{if } h_i^{(k-1)}=1 \land (U_i^{\mathrm{rec}} < \gamma), \\[4pt] r_i^{(k-1)}, & \text{otherwise}. \end{cases} \end{aligned}\end{split}\]

Status

During the simulation, a node can be in one of the following states:

Status	Code
Susceptible	0
Infected	1
Recovered	2

SIRModel

class fs_gplib.Epidemics.SIRModel(data, seeds, infection_beta: float, recovery_lambda: float, device='cpu', use_weight: bool = False, rand_seed=None)

Bases: DiffusionModel

SIR (Susceptible-Infected-Recovered) diffusion model on static graphs.

Each node starts as susceptible or infected (seed). At every step each infected neighbor independently transmits the disease with probability \(\beta\), while infected nodes recover with probability \(\gamma\) and become permanently immune.

Returned node states are encoded as: 0 = susceptible, 1 = infected, 2 = recovered.

Parameters:

• data (torch_geometric.data.Data) -- PyTorch Geometric Data object representing graph \(G=(V,E)\). Must contain edge_index (the edge set \(E\)) and num_nodes (\(|V|\)). When use_weight is True, edge_attr supplies per-edge weights \(w_{ji}\).

• seeds (list[int] | float) -- Nodes whose initial state is Infected. Pass a list of node IDs, or a float in (0, 1) to infect that fraction of nodes chosen uniformly at random.

• infection_beta (float) -- Per-contact infection probability \(\beta \in [0,1]\).

• recovery_lambda (float) -- Per-step recovery probability \(\gamma \in [0,1]\).

• device (str | int) -- (optional) 'cpu' or a CUDA device index. Defaults to 'cpu'.

• use_weight (bool) -- (optional) If True, each edge \((j,i)\) carries a weight \(w_{ji}\) from data.edge_attr and the infection probability becomes \(\beta w_{ji}\). If False all weights default to 1 (i.e. \(w_{ji}=1\)). Defaults to False.

• rand_seed (int | None) -- (optional) Random seed used when seeds is a float. Defaults to None.

run_iteration()

Execute a single simulation step.

The internal node_status is updated so that subsequent calls continue from the latest state.

Returns:

Node states after one step, shape (1, N).

Return type:

torch.Tensor

run_iterations(times)

Execute times simulation steps sequentially.

The internal node_status is updated in-place so that subsequent calls continue from the latest state.

Parameters:

times (int) -- Number of steps to run.

Returns:

Node states at final step, shape (1, N).

Return type:

torch.Tensor

run_epoch(iterations_times)

Run a single Monte-Carlo epoch (one independent realisation).

Node states are re-initialised before the epoch starts.

Parameters:

iterations_times (int) -- Number of simulation steps per epoch.

Returns:

Node states at final step of the epoch, shape (1, N).

Return type:

torch.Tensor

run_epochs(epochs, iterations_times, batch_size=200)

Run multiple independent Monte-Carlo epochs in batches.

Node states are re-initialised before the run.

Parameters:

• epochs (int) -- Total number of independent realisations.

• iterations_times (int) -- Number of simulation steps per epoch.

• batch_size (int) -- (optional) Number of epochs processed in parallel per batch. Defaults to 200.

Returns:

Node states at final step of all epochs, shape (epochs, N).

Return type:

torch.Tensor

References

[1]

Pastor-Satorras R, Castellano C, Van Mieghem P, et al. Epidemic processes in complex networks[J]. Reviews of modern physics, 2015, 87(3): 925-979.

[2]

Kermack W O, McKendrick A G. A contribution to the mathematical theory of epidemics[J]. Proceedings of the royal society of london. Series A, Containing papers of a mathematical and physical character, 1927, 115(772): 700-721.