Epidemic Models

Epidemic models originate from mathematical epidemiology and describe how diseases, information, or behaviors spread through a population represented as a network. In these models, each node occupies a discrete state (e.g., Susceptible, Infected, Recovered) and transitions between states according to stochastic rules governed by local interactions with neighbors.

Epidemic models have broad applications beyond public health, including viral marketing, rumor propagation, cybersecurity threat analysis, and influence maximization in social networks.

All epidemic models in FS_GPlib operate on static networks — directed or undirected graphs with a fixed topology throughout the simulation. Both unweighted and weighted edges are supported; when edge weights are enabled, transmission probabilities are modulated by the corresponding weights. The models support execution on CPU and GPU, and can be combined with the batch-parallel and distributed acceleration features described in the Tutorial.

Base Class & Common API

Every epidemic model inherits from DiffusionModel, which provides the shared initialisation pipeline and a uniform simulation interface. Understanding the base class helps you leverage the full API of any concrete model.

class fs_gplib.Epidemics.base.DiffusionModel(data, seeds: float | List[int] | None, rand_seed=None, device='cpu', use_weight=False, **kwargs)[source]

Base class for all epidemic / diffusion models on static graphs.

DiffusionModel provides the shared initialisation logic and a uniform simulation interface that every concrete model (SI, SIR, SEIR, …) inherits. Users typically do not instantiate this class directly; instead they create a subclass such as SIModel.

Initialisation pipeline (executed in __init__):

Validate the input graph via _validate_graph().
Initialise seed nodes via _initialize_seeds().
Validate and store model-specific parameters via _validate_parameters().
Transfer the model to the target device via _set_device().

Simulation interface – four progressively higher-level entry points:

Method	Description
`run_iteration()`	Advance one time step from the current state.
`run_iterations(times)`	Advance times steps from the current state.
`run_epoch(times)`	Reset to initial state, then run times steps (one independent realisation).
`run_epochs(epochs, times, batch_size)`	Run epochs independent realisations in batches (Monte-Carlo simulation).

Parameters:

data (torch_geometric.data.Data) -- PyTorch Geometric Data object representing the graph. Must contain edge_index; if use_weight=True, edge_attr must also be present.
seeds (float | list[int] | None) -- Initial infected (or activated) node set. A list of node IDs or a float in (0, 1) representing the fraction of nodes to infect uniformly at random. None is accepted for models that do not require seeds.
rand_seed (int | None) -- Random seed for NumPy, used when seeds is a float to make the random selection reproducible. Defaults to None.
device (str | int) -- 'cpu' or a GPU device index (e.g. 0). When a GPU index is given the model automatically selects CUDA or MPS depending on hardware availability. Defaults to 'cpu'.
use_weight (bool) -- If True, edge weights from data.edge_attr are used to modulate transmission probabilities. Defaults to False.
kwargs -- Model-specific parameters (e.g. infection_beta, recovery_lambda). Each key-value pair is validated to lie in [0, 1] and stored as an instance attribute.

Utility Methods

The following utility methods can be called on any epidemic model instance to re-configure the model after construction:

DiffusionModel._init_node_status()[source]

Reset all node states to the initial configuration.

This method rebuilds the internal node_status tensor(s) so that only the seed nodes are in the Infected (or Activated) state and all other nodes are Susceptible.

Subclasses must override this method to create the status tensors appropriate for their compartmental structure (e.g. a single boolean tensor for SI, or a dict of boolean tensors for SIR).

Calling this method is useful when you want to rerun a simulation from scratch without reconstructing the model object.

Example

result1 = model.run_iterations(100)
model._init_node_status()   # reset
result2 = model.run_iterations(100)

DiffusionModel._set_device(device)[source]

Set (or switch) the computation device for the model.

When called during initialisation the graph data and internal tensors are moved to the target device. It can also be called after construction to migrate an existing model to a different device (subclasses should override this method to move additional tensors).

Parameters:: device (str | int) -- 'cpu' or a GPU device index (e.g. 0). If a GPU index is given, CUDA is preferred; on Apple Silicon the MPS backend is used as a fallback.
Raises:: Exception -- If a GPU is requested but neither CUDA nor MPS is available.

Example

model = SIModel(data, seeds, infection_beta=0.05, device='cpu')
# switch to GPU 0
model._set_device(0)

DiffusionModel._set_seed(seeds)[source]

Replace the current seed set and refresh node states.

This is the recommended way to change seed nodes after construction. It calls _initialize_seeds() to parse and store the new seeds, then (if the model is fully initialised) _init_node_status() to rebuild the node-state tensors on the current device.

Parameters:: seeds (float | list[int] | None) -- New seed specification – same format as the constructor parameter seeds (a float fraction, a list of node IDs, or None).

Example

model = SIRModel(data, seeds=[0, 1, 2], ...)
# later, switch to a random 10 % seed set
model._set_seed(0.1)

Simulation Interface

All epidemic models expose four progressively higher-level simulation methods. The table below summarises their behaviour; each concrete model page documents the return type for that specific model.

Method	Description	Resets state?
`Model.run_iteration()`	Advance one time step from the current state.	No
`Model.run_iterations(times)`	Advance times steps from the current state.	No
`Model.run_epoch(times)`	Reset to initial state, then run times steps (one independent realisation).	Yes
`Model.run_epochs(epochs, times, batch_size)`	Run epochs independent realisations in batches (Monte-Carlo simulation).	Yes

Tip

run_iteration / run_iterations do not reset node states, so they can be chained to inspect intermediate states or implement custom stopping criteria. run_epoch / run_epochs always start from the initial seed configuration and are the recommended entry points for Monte-Carlo studies.

Available Models

The following models are available: