Simulation

In this section, we provide a tutorial of the simulation methods in EpiLearn. In general, we focus on the simulation of static and dynamic properties including graph structure and node features.

Note

Coming from 0.0.x? Time_geo was removed with no replacement, and Gravity_model changed signature. Both changes, and everything else, are listed in MIGRATION.md.

Static Properties

Random Static Graph

epilearn.utils.simulation.get_random_graph(num_nodes=None, connect_prob=None, block_sizes=None, num_edges=None, graph_type='erdos_renyi')

Generates a random static graph using one of the supported graph types: Erdos-Renyi, Stochastic Blockmodel, or Barabasi-Albert.

Parameters:

num_nodes (int) – Number of nodes in the graph.
connect_prob (float, optional) – Probability of edge creation (for Erdos-Renyi and Stochastic Blockmodel graphs).
block_sizes (list of int, optional) – Sizes of blocks (for Stochastic Blockmodel graph).
num_edges (int, optional) – Number of edges (for Barabasi-Albert graph).
graph_type (str) – Type of graph to generate. Options are ‘erdos_renyi’, ‘stochastic_blockmodel’, ‘barabasi_albert’. Default is ‘erdos_renyi’.

Returns:

Adjacency matrix of the generated graph.

Return type:

torch.Tensor

graph_type selects the generator: 'erdos_renyi' needs num_nodes and connect_prob, 'stochastic_blockmodel' takes block_sizes plus a block-wise probability matrix, and 'barabasi_albert' takes num_edges (edges attached per new node). All three return a dense torch.float32 adjacency matrix, ready to pass as Dataset(graph=...).

from epilearn.utils.simulation import get_random_graph

adj = get_random_graph(num_nodes=25, connect_prob=0.2, graph_type='erdos_renyi')
sbm = get_random_graph(block_sizes=[10, 15],
                       connect_prob=[[0.4, 0.05], [0.05, 0.4]],
                       graph_type='stochastic_blockmodel')
ba = get_random_graph(num_nodes=25, num_edges=2, graph_type='barabasi_albert')
print(adj.shape, sbm.shape, ba.shape)
# torch.Size([25, 25]) torch.Size([25, 25]) torch.Size([25, 25])

Static Features

epilearn.utils.simulation.get_graph_from_features(features, adj=None, G=1)

Generate a graph from node features using cosine similarity.

This function generates a graph where each edge weight is computed based on the cosine similarity between the feature vectors of the connected nodes. If an adjacency matrix is provided, the cosine similarity is adjusted by the corresponding entry in the adjacency matrix.

Parameters:

features (torch.Tensor) – A tensor of shape (num_nodes, feat_dim) where num_nodes is the number of nodes and feat_dim is the dimensionality of the feature vectors.
adj (torch.Tensor, optional) – A tensor of shape (num_nodes, num_nodes) representing the adjacency matrix, where adj[i, j] denotes the distance or weight between node i and node j. If None, the cosine similarity is used directly as the edge weight. Default is None.

Returns:

A tensor of shape (num_nodes, num_nodes) representing the generated graph’s adjacency matrix, where each entry [i, j] contains the adjusted cosine similarity between nodes i and j.

Return type:

torch.Tensor

Given fixed node features, a graph can be built from the cosine similarity between nodes. The optional adj holds the distances between nodes: pass it and each similarity is divided by the corresponding distance, penalizing far-apart pairs.

import torch
from epilearn.utils.simulation import get_graph_from_features

feature = torch.rand(10, 20)              # 10 nodes, 20 features each
adj = torch.randint(10, 100, (10, 10))    # pairwise distances

graph1 = get_graph_from_features(features=feature, adj=None)
graph2 = get_graph_from_features(features=feature, adj=adj)

Gravity Model

The gravity model turns populations and connection strengths into mobility flows — in epidemics, the regional contact and transmission driven by human movement. It is parameterized by two population exponents and a connectivity decay:

\[F_{ij} = N_i^{\rho} \cdot N_j^{\theta} \cdot \exp\!\big((w_{ij} - 1) / \delta\big)\]

rho and theta are the source/target population exponents (typically 0.5-1.0), delta is the connectivity decay (0.2-0.5 local, 0.5-1.0 regional, 1.0-2.0 long range), and normalize divides the flow by the two regions’ total population. Setting rho=0, theta=0, normalize=False gives the diffusive special case — flow is just the edge weight times the population difference — which is what simulate_spatiotemporal_regions uses when no model is given.

Warning

\(w_{ij}\) is connection strength: higher means more flow, and 0 means no edge. That is the opposite convention from a distance matrix, so convert distances to strengths and normalize them to [0, 1] before use.

class epilearn.utils.simulation.Gravity_model(rho: float = 0.0, theta: float = 0.0, delta: float = 1.0, normalize: bool = False)

Elegant gravity model for human mobility in epidemic simulations.

Computes flow between regions based on population attraction and connectivity strength:

\[F_{ij} = N_i^{\rho} \cdot N_j^{\theta} \cdot \exp((w_{ij} - 1) / \delta)\]

where higher w_{ij} indicates stronger connection between regions.

Parameters:

rho (float) – Source population exponent (typically 0.5-1.0).
theta (float) – Target population exponent (typically 0.5-1.0).
delta (float) – Connectivity decay parameter. Controls sensitivity to connectivity variations. Recommended: 0.2-2.0 for normalized connectivity [0, 1].
normalize (bool) – If True, normalize flows by total population (default True).

Notes

Unified Connectivity Semantics:

Connectivity values represent connection strength where: - Higher values = Stronger connection = More flow - Lower values = Weaker connection = Less flow - 0 = No connection

This is consistent across both diffusive and gravity models.

It is recommended to normalize your connectivity matrix to [0, 1] range:

>>> # Normalize connectivity matrix (edge weights)
>>> max_conn = connectivity_matrix.max()
>>> normalized_connectivity = connectivity_matrix / max_conn
>>>
>>> # Choose delta based on desired spatial spread
>>> # Sharp decay (local): delta = 0.2-0.5
>>> # Moderate (regional): delta = 0.5-1.0
>>> # Gradual (long-range): delta = 1.0-2.0

Examples

>>> # Diffusive flow (special case: rho=0, theta=0)
>>> diffusive = Gravity_model(rho=0, theta=0, delta=1.0, normalize=False)
>>>
>>> # Gravity model with normalized connectivity
>>> gravity = Gravity_model(rho=1.0, theta=1.0, delta=0.5, normalize=True)
>>>
>>> # Using with normalized edge weights
>>> weights_normalized = adjacency / adjacency.max()
>>> result = simulate_spatiotemporal_regions(
...     model, states, weights_normalized, steps=100, gravity_model=gravity)

__init__(rho: float = 0.0, theta: float = 0.0, delta: float = 1.0, normalize: bool = False)

compute_flow(pop_i: float, pop_j: float, connectivity: float) → float

Compute flow between two regions.

Parameters:

pop_i (float) – Population sizes of regions i and j.
pop_j (float) – Population sizes of regions i and j.
connectivity (float) –
Connection strength between regions (recommended range: 0-1). Higher values = stronger connection = more flow. - 0: No connection - 1: Maximum connection strength

Interpretation is unified across all models: - Diffusive: Direct multiplier on population flow - Gravity: Exponential enhancement of attraction

Returns:

Flow magnitude or edge weight for further computation.

Return type:

float

Notes

Unified Connectivity Semantics:

Connectivity always represents connection strength where higher = stronger:

Diffusive model (rho=0, theta=0):

Formula: F = connectivity × (N_i - N_j)
Range: [0.0, 1.0]
Interpretation: Fraction of population difference that flows per timestep
Example: 0.1 = 10% of population difference travels per day

Gravity model (rho>0, theta>0):

Formula: F = [N_i^ρ × N_j^θ × exp((connectivity-1)/δ)] × (N_i - N_j)/(N_i + N_j)
Range: [0.0, 1.0] normalized
connectivity=1.0 gives baseline gravity attraction
connectivity>1.0 enhances flow (if using unnormalized weights)
connectivity<1.0 reduces flow
Delta parameter controls sensitivity: * Small delta (0.2-0.5): Sharp response to connectivity differences * Large delta (1.0-2.0): Gradual response to connectivity differences

Normalization Strategy:

For any connectivity matrix (edge weights, similarity scores, etc.): 1. Find max value: max_conn = connectivity_matrix.max() 2. Normalize: connectivity_matrix = connectivity_matrix / max_conn 3. Choose delta based on desired sensitivity:

Sharp (local spread): delta = 0.2-0.5

Moderate: delta = 0.5-1.0

Gradual (long-range): delta = 1.0-2.0

compute_mobility_matrix(populations: Tensor, connectivity_matrix: Tensor) → Tensor

Vectorized computation of full mobility matrix.

Parameters:

populations (torch.Tensor) – Population of each region, shape (num_regions,).
connectivity_matrix (torch.Tensor) – Connectivity matrix, shape (num_regions, num_regions). Higher values = stronger connection = more flow (unified semantics).

Returns:

Mobility flow matrix, shape (num_regions, num_regions).

Return type:

torch.Tensor

compute_net_flow(pop_i: float, pop_j: float, connectivity: float) → float

Compute net directional flow from region i to region j.

Parameters:

pop_i (float) – Population sizes of regions i and j.
pop_j (float) – Population sizes of regions i and j.
connectivity (float) – Connection strength between regions. Higher values = stronger connection = more flow (unified across all models).

Returns:

Net flow (positive = i→j, negative = j→i).

Return type:

float

Given population numbers in each node (or say region) and a connectivity matrix, edge weights can be obtained for a pair of nodes or for the whole graph at once:

import torch
from epilearn.utils.simulation import Gravity_model

node_populations = torch.tensor([1000., 2000., 1500.])
connectivity = torch.tensor([[0.0, 0.8, 0.2],      # connection strength in [0, 1]
                             [0.8, 0.0, 0.5],
                             [0.2, 0.5, 0.0]])

gravity = Gravity_model(rho=1.0, theta=1.0, delta=0.5, normalize=True)
print(round(gravity.compute_flow(1000.0, 2000.0, 0.8), 2))       # 446.88
print(round(gravity.compute_net_flow(1000.0, 2000.0, 0.8), 2))   # -148.96, i.e. j -> i
print(gravity.compute_mobility_matrix(node_populations, connectivity))
# tensor([[  0.0000, 446.8800, 121.1379],
#         [446.8800,   0.0000, 315.3253],
#         [121.1379, 315.3253,   0.0000]])

# diffusive special case: the mobility matrix is the connectivity itself
diffusive = Gravity_model(rho=0.0, theta=0.0, delta=1.0, normalize=False)
print(diffusive.is_diffusive)                                    # True

Dynamic Properties

The three simulators below all take a compartmental model and return a dict of tensors. They differ only in what a “unit” is:

Function	Level	`trajectory` shape
`simulate_temporal_epidemic`	one population	`(steps+1, n_compartments)`
`simulate_spatiotemporal_individual`	individuals on a contact graph	`(steps+1, n_nodes)` of compartment indices
`simulate_spatiotemporal_regions`	regions with mobility	`(steps+1, n_regions, n_compartments)`

The models live in epilearn.utils.compartmental_models (SIRModel, SEIRModel, SIRSModel, SEIRVIModel); see Utilities for their parameters and for the parameter_schedule /input_schedule mechanism used below.

Temporal Simulation

simulate_temporal_epidemic integrates the ODE at the population level. State vectors are counts, and their length must match model.compartments. Interventions enter through parameter_schedule: a callable f(step_idx, t, state) -> dict | None, a {step_idx: dict} mapping, or a sequence indexed by step, where None leaves the base parameters alone.

import epilearn
from epilearn.utils.compartmental_models import SIRModel
from epilearn.utils.simulation import simulate_temporal_epidemic

model = SIRModel(beta=0.3, gamma=0.1)
run = simulate_temporal_epidemic(model,
                                 initial_state=[9990.0, 10.0, 0.0],   # [S, I, R]
                                 steps=160, dt=1.0)
print(run.keys())                                  # time, trajectory, compartments
print(run['trajectory'].shape, run['compartments'])
# torch.Size([161, 3]) ('S', 'I', 'R')
epilearn.visualize.plot_series(run['trajectory'].numpy(),
                               columns=['Susceptible', 'Infected', 'Recovered'])

def lockdown(step_idx, t, state):
    """Halve transmission between day 20 and day 70."""
    return {'beta': 0.15} if 20 <= step_idx < 70 else None

npi = simulate_temporal_epidemic(model, [9990.0, 10.0, 0.0], steps=200,
                                 parameter_schedule=lockdown)
for r in (run, npi):
    print(f"peak {r['trajectory'][:, 1].max():.0f} on day {r['trajectory'][:, 1].argmax()}")
# peak 3006 on day 38     <- unmitigated
# peak 1078 on day 86     <- with the lockdown

noisy = simulate_temporal_epidemic(model, [9990.0, 10.0, 0.0], steps=160,
                                   process_noise=5.0, seed=42)

process_noise (a scalar or a per-compartment vector, with seed for reproducibility) adds Gaussian noise after each step and re-projects to non-negative values, which is the easiest way to get a realistic-looking series to train a model on. Without it, simulate_temporal_epidemic is just model.simulate(initial_state, steps).

epilearn.utils.simulation.simulate_temporal_epidemic(model: CompartmentalModel, initial_state: Sequence[float] | Tensor | ndarray, steps: int, dt: float = 1.0, parameter_schedule: Mapping[int, Dict[str, float]] | Sequence | None = None, input_schedule: Mapping[int, Dict[str, float]] | Sequence | None = None, process_noise: float | Sequence[float] | Tensor | None = None, method: str = 'rk4', seed: int | None = None): Run a temporal (population-level) simulation for the provided compartmental model.

Individual-Level Simulation

simulate_spatiotemporal_individual runs a stochastic node-level process on a contact graph: susceptible nodes are infected at rate \(1 - e^{-\beta\,\Delta t\,\sum_j A_{ij}\mathbb{1}[j \text{ infectious}]}\), then progress and recover with the model’s sigma / gamma / omega. create_initial_conditions_individual builds a random Erdos-Renyi contact network plus matching initial states.

from epilearn.utils.compartmental_models import SEIRModel
from epilearn.utils.simulation import (create_initial_conditions_individual,
                                      simulate_spatiotemporal_individual)

model = SEIRModel(beta=0.35, gamma=0.1, sigma=0.2)
node_states, contact_graph = create_initial_conditions_individual(
    model, num_individuals=200, p_edge=0.03,
    initial_compartment_fractions={'S': 0.98, 'I': 0.02}, seed=0)
print(node_states.shape, contact_graph.shape)
# torch.Size([200]) torch.Size([200, 200])

res = simulate_spatiotemporal_individual(model, contact_graph, node_states,
                                         steps=60, stochastic=True, seed=0)
print(res.keys())
# time, trajectory, counts, compartments, contact_graph, dynamic_graph, node_features
print(res['trajectory'].shape, res['counts'].shape)
# torch.Size([61, 200]) compartment index, torch.Size([61, 4]) S/E/I/R totals
print(res['node_features'].shape, res['dynamic_graph'].shape)
# torch.Size([61, 200, 4]) one-hot, model-ready, torch.Size([61, 200, 200])

# node_features is already shaped [T, N, F], so this is a training set:
from epilearn.data import Dataset
dataset = Dataset(x=res['node_features'],
                  y=res['node_features'][:, :, model.compartments.index('I')],
                  graph=contact_graph,
                  dynamic_graph=res['dynamic_graph'])

contact_graph may also be time-varying: pass a (steps+1, N, N) tensor, or a callable f(t, step_idx) -> adjacency to rewire the network as the epidemic runs (for example, to model contact reduction).

Warning

create_initial_conditions_individual renormalizes initial_compartment_fractions so they sum to 1. Passing only the infected share — including the default {'I': 0.01} — therefore seeds the whole population as infectious. Always give the full distribution, e.g. {'S': 0.98, 'I': 0.02}. The model must also have a compartment literally named 'S', otherwise the simulator raises ValueError.

Simulate an individual-level compartmental process on a contact graph. Nodes follow the provided model (SIR/SEIR/SIRS) with infection pressure driven by neighbors. Returns per-step node features (one-hot compartment indicators) and a dynamic graph tensor of shape (time, N, N).

Parameters:

model (CompartmentalModel) – The compartmental model (e.g., SIRModel, SEIRModel, SIRSModel).
contact_graph (Union[torch.Tensor, np.ndarray, nx.Graph, Callable]) –
Contact graph specification. Can be: - Static: Binary adjacency matrix (num_nodes, num_nodes) or NetworkX graph - Time-varying: Tensor of shape (steps+1, num_nodes, num_nodes) - Dynamic: Callable function f(t, step_idx) -> adjacency_matrix that returns

the contact graph at each time step
initial_states (Union[Sequence[int], Sequence[str], torch.Tensor, np.ndarray, Dict[str, Iterable[int]]]) – Initial compartment states for each individual.
steps (int) – Number of simulation steps.
dt (float) – Time step size (default 1.0).
stochastic (bool) – Whether to use stochastic transitions (default True).
seed (int, optional) – Random seed for reproducibility.

Returns:

Simulation results with keys: - ‘time’: Time axis (steps+1,) - ‘trajectory’: Individual state history (steps+1, num_nodes) with compartment indices - ‘counts’: Compartment counts over time (steps+1, num_compartments) - ‘compartments’: Compartment names - ‘contact_graph’: Static or initial contact graph - ‘dynamic_graph’: Time-varying contact graphs (steps+1, num_nodes, num_nodes) - ‘node_features’: One-hot encoded compartment states (steps+1, num_nodes, num_compartments)

Return type:

dict

epilearn.utils.simulation.create_initial_conditions_individual(model, num_individuals=100, p_edge=0.01, initial_compartment_fractions: dict = {'I': 0.01}, seed=42)

Create initial conditions for individual-level simulations.

Parameters:

model (CompartmentalModel) – The compartmental model (e.g., SIRModel, SIRSModel, SEIRModel) that defines the disease dynamics and compartment structure
num_individuals (int) – Number of individuals in the network
p_edge (float) – Edge probability for Erdos-Renyi random graph (contact network)
initial_compartment_fractions (dict) – Dictionary mapping compartment names to initial fractions. Example: {‘I’: 0.01} means 1% start infected, rest susceptible Example: {‘S’: 0.9, ‘E’: 0.05, ‘I’: 0.05} for SEIR model
seed (int) – Random seed for reproducibility

Returns:

node_states: torch.Tensor of shape (num_individuals,) containing compartment indices
adjacency: torch.Tensor of shape (num_individuals, num_individuals) contact graph

Return type:

tuple

Region-Level (Metapopulation) Simulation

simulate_spatiotemporal_regions alternates two phases per step: integrate each region’s internal dynamics, then move people along the graph using a Gravity_model. Flows are antisymmetric, so the total population is conserved. create_initial_conditions_region sets up the regions, and create_regional_forcing_params gives each one its own multi-frequency seasonal forcing, which desynchronizes the regional waves instead of leaving them in lock-step.

from epilearn.utils.compartmental_models import SIRSModel
from epilearn.utils.simulation import (create_initial_conditions_region,
                                      create_regional_forcing_params,
                                      simulate_spatiotemporal_regions,
                                      Gravity_model)

model = SIRSModel(beta=0.35, gamma=0.1, omega=0.02)   # SIRS never burns out
init = create_initial_conditions_region(model, n_regions=30, p_edge=0.12,
                                        pop_range=(5000, 20000),
                                        n_initial_infected=3,
                                        initial_infected_size=50, seed=1)
print(init['n_regions'], init['n_edges'])          # 30 59
print(init['region_states'].shape)                 # torch.Size([30, 3])

forcing = create_regional_forcing_params(init['n_regions'], seed=1)
print(round(forcing[0]['amp1'], 4), forcing[0]['period1'])   # 0.1967 30.0

res = simulate_spatiotemporal_regions(
    model,
    region_states=init['region_states'],           # counts, [n_regions, n_compartments]
    adjacency_graph=init['adjacency'], steps=120, dt=1.0,
    forcing_params=forcing, forcing_parameter='beta',
    method='euler',                                # 'euler' is the stable choice here
    gravity_model=Gravity_model(rho=1.0, theta=1.0, delta=0.5, normalize=True),
    travel_rate=0.01)                              # global mobility scaling

print(res['trajectory'].shape)                     # torch.Size([121, 30, 3])
print(res['dynamic_graph'].shape)                  # torch.Size([121, 30, 30])
print(res['effective_reproduction_number'].shape)  # torch.Size([121, 30])

# straight into a spatiotemporal task
from epilearn.data import Dataset
dataset = Dataset(x=res['trajectory'],                    # [T, N, n_compartments]
                  y=res['trajectory'][:, :, 1],           # infectious counts
                  graph=init['adjacency'],
                  dynamic_graph=res['directed_flow'])

Three of the returned tensors are worth calling out. dynamic_graph is the signed net-flow graph per step (F[i, j] > 0 means i → j, and F == -F.T), so feed a model directed_flow instead — the same tensor clamped at 0. parameter_history (steps+1, n_regions) records the value of forcing_parameter actually used, and effective_reproduction_number (steps+1, n_regions) is the per-region \(R_t\) derived from it (None for models without gamma or without an 'S' compartment). Also returned: counts, regional_totals, adjacency, adjacency_history and node_features.

Warning

create_initial_conditions_region samples the seeded regions without replacement, so n_initial_infected >= n_regions raises ValueError: Cannot take a larger sample than population when 'replace=False'. The default is n_initial_infected=20, so any run with fewer than 21 regions must pass a smaller value. ensure_connected=True (the default) also keeps only the largest connected component, so read the surviving region count back from init['n_regions'].

epilearn.utils.simulation.simulate_spatiotemporal_regions(model: CompartmentalModel, region_states: Tensor | ndarray, adjacency_graph: Tensor | ndarray | Graph | Callable, steps: int, dt: float = 1.0, forcing_params: List[Dict[str, float]] | None = None, forcing_parameter: str = 'beta', method: str = 'euler', gravity_model: Gravity_model | None = None, travel_rate: float = 1.0) → Dict[str, Tensor]

Simulate region-level epidemics with mobility-driven population flows.

This function implements a two-step process: 1. Internal dynamics update (with optional multi-frequency forcing) 2. Population flow based on mobility model (diffusive or gravity-based)

The flows are applied AFTER internal dynamics. Flow models are configured via the gravity_model parameter, which supports both diffusive and gravity-based flows.

Parameters:

model (CompartmentalModel) – Compartmental epidemic model (must be SIRS-like for infinite stability).
region_states (Union[torch.Tensor, np.ndarray]) – Initial states for each region as counts (not fractions), shape (num_regions, num_compartments). For SIRS: [S, I, R] counts for each region.
adjacency_graph (Union[torch.Tensor, np.ndarray, nx.Graph, Callable]) – Connectivity/adjacency graph specification where higher values = stronger connections. Edge weights represent connection strength (travel rates, similarity, interaction frequency). Can be: - Static: Weighted matrix (num_regions, num_regions) or NetworkX graph - Time-varying: Tensor of shape (steps+1, num_regions, num_regions) - Dynamic: Callable f(t, step_idx) -> adjacency_matrix
steps (int) – Number of simulation steps.
dt (float) – Time step size (default 1.0 for daily updates).
forcing_params (List[Dict[str, float]], optional) – Region-specific forcing parameters with keys: ‘amp1’, ‘phase1’, ‘period1’, ‘amp2’, ‘phase2’, ‘period2’, ‘amp3’, ‘phase3’, ‘period3’.
forcing_parameter (str) – Name of the parameter to apply forcing to (default ‘beta’).
method (str) – Integration method (‘euler’ recommended for stability, default ‘euler’).
gravity_model (Gravity_model, optional) – Mobility model instance. If None, defaults to diffusive model. - Diffusive: Gravity_model(rho=0, theta=0, delta=1.0, normalize=False) - Gravity: Gravity_model(rho=1.0, theta=1.0, delta=100.0, normalize=True)
travel_rate (float) – Global mobility scaling factor applied to every edge-wise flow (default 1.0). Use this to tune the overall level of inter-regional travel without modifying the adjacency weights.

Returns:

Simulation results with keys: - ‘time’: Time axis (steps+1,) - ‘trajectory’: Regional state history (steps+1, num_regions, num_compartments) - ‘compartments’: Compartment names - ‘adjacency’: Static adjacency or provided tensor (for backward compatibility) - ‘adjacency_history’: Adjacency tensor for every timestep (steps+1, num_regions, num_regions) - ‘dynamic_graph’: Signed flow graphs (steps+1, num_regions, num_regions) - ‘directed_flow’: Non-negative directed flows derived from dynamic_graph - ‘parameter_history’: Value of forcing_parameter used per region and timestep - ‘effective_reproduction_number’: R_t estimates when applicable - ‘counts’: Aggregate compartment counts summed across regions - ‘regional_totals’: Population of each region at every timestep - ‘node_features’: Alias for trajectory (kept for convenience)

Return type:

dict

Notes

Flow Models:

Both models use unified connectivity semantics (higher = stronger connection):

Diffusive (default):: F(i,j) = connectivity[i,j] × (N_i - N_j)
Gravity:: F(i,j) = [N_i^ρ × N_j^θ × exp((connectivity[i,j]-1)/δ)] × (N_i - N_j)/(N_i + N_j)

Both ensure population conservation through antisymmetric flows.

Connectivity Normalization:

Normalize edge weights to [0, 1] for best results: >>> adjacency_normalized = adjacency / adjacency.max()

Examples

Diffusive flow (default):

>>> result = simulate_spatiotemporal_regions(
...     model, states, adjacency, steps=100)

Diffusive flow (explicit):

>>> diffusive = Gravity_model(rho=0, theta=0, delta=1.0, normalize=False)
>>> result = simulate_spatiotemporal_regions(
...     model, states, adjacency, steps=100, gravity_model=diffusive)

Gravity model:

>>> # Normalize connectivity (edge weights)
>>> adjacency_norm = adjacency / adjacency.max()
>>> gravity = Gravity_model(rho=1.0, theta=1.0, delta=0.5, normalize=True)
>>> result = simulate_spatiotemporal_regions(
...     model, states, adjacency_norm, steps=100, gravity_model=gravity)

epilearn.utils.simulation.create_initial_conditions_region(model, n_regions=100, p_edge=0.01, pop_range=(800, 20000), n_initial_infected=20, initial_infected_size=100, initial_compartment_fractions=None, ensure_connected=True, seed=42)

Create initial conditions for spatiotemporal epidemic simulation based on a compartmental model.

Parameters:

model (CompartmentalModel) – The compartmental model (e.g., SIRModel, SIRSModel, SEIRModel) that defines the disease dynamics and compartment structure
n_regions (int) – Number of regions in the network
p_edge (float) – Edge probability for Erdos-Renyi random graph
pop_range (tuple of (int, int)) – Range for random population initialization (min_pop, max_pop)
n_initial_infected (int) – Number of regions to seed with initial infections
initial_infected_size (int) – Number of individuals initially infected in each seeded region
initial_compartment_fractions (dict or None) – Optional dictionary mapping compartment names to initial fractions. If None, all individuals start in ‘S’ (susceptible). Example: {‘S’: 0.9, ‘E’: 0.05, ‘I’: 0.05} for SEIR model
ensure_connected (bool) – If True, extract largest connected component
seed (int) – Random seed for reproducibility

Returns:

Dictionary containing: - ‘adjacency’: torch.Tensor of shape (n_regions, n_regions) - ‘region_states’: torch.Tensor of shape (n_regions, n_compartments) - ‘graph’: networkx.Graph object - ‘n_regions’: int (actual number of regions after connectivity check) - ‘n_edges’: int - ‘infected_nodes’: numpy.ndarray of initially infected region indices - ‘model’: the compartmental model used - ‘compartment_names’: list of compartment names - ‘total_population’: total population across all regions

Return type:

dict

epilearn.utils.simulation.create_regional_forcing_params(num_regions: int, base_amplitudes: Sequence[float] = (0.2, 0.15, 0.1), periods: Sequence[float] = (30.0, 47.0, 73.0), amplitude_noise: float = 0.1, seed: int | None = None) → List[Dict[str, List[float] | float]]

Create region-specific multi-frequency forcing parameters.

Each region gets slightly different amplitudes and random phases to desynchronize dynamics across regions.

Parameters:

num_regions (int) – Number of regions.
base_amplitudes (Sequence[float]) – Base amplitudes for each frequency component.
periods (Sequence[float]) – Periods for each frequency component (in time units).
amplitude_noise (float) – Relative noise level for amplitude perturbations (default 0.1 = 10%).
seed (int, optional) – Random seed for reproducibility.

Returns:

List of dictionaries, one per region, with keys: ‘amp1’, ‘phase1’, ‘period1’, ‘amp2’, ‘phase2’, ‘period2’, ‘amp3’, ‘phase3’, ‘period3’.

Return type:

List[Dict[str, Union[List[float], float]]]

Torch Compartmental Layers

Separately from the ODE utilities above, epilearn.models.Temporal.Compartmental ships nn.Module compartmental layers, useful when you want an SIR-like model that participates in autograd. SIR takes horizon (total simulation steps), infection_rate and recovery_rate; NetworkSIR spreads the disease over a graph and adds num_nodes. Both are called as model(states, [graph,] steps), where steps=None runs the full horizon. SIS, SEIR, NetworkSIS and NetworkSEIR live in the same module.

import epilearn
import torch
from epilearn.models.Temporal.Compartmental import SIR, NetworkSIR

# 25 nodes, all susceptible except nodes 3 and 10; columns are [S, I, R]
initial_states = torch.zeros(25, 3)
initial_states[:, 0] = 1
initial_states[[3, 10], 0] = 0
initial_states[[3, 10], 1] = 1

# population level: the layer takes the aggregated [S, I, R] counts
model = SIR(horizon=190, infection_rate=0.05, recovery_rate=0.05)
preds = model(initial_states.sum(0), steps=None)
print(preds.shape)   # torch.Size([190, 3])
epilearn.visualize.plot_series(preds.detach().numpy(),
                              columns=['Susceptible', 'Infected', 'Recovered'])

# node level: same states, plus a graph
initial_graph = epilearn.utils.simulation.get_random_graph(num_nodes=25, connect_prob=0.20)
net = NetworkSIR(num_nodes=initial_graph.shape[0], horizon=120,
                 infection_rate=0.05, recovery_rate=0.05)
preds = net(initial_states, initial_graph, steps=None)
print(preds.shape)   # torch.Size([120, 25, 3])
epilearn.visualize.plot_graph(preds.argmax(2)[15].detach().numpy(),
                              initial_graph.to_sparse().indices().detach().numpy(),
                              classes=['Susceptible', 'Infected', 'Recovered'])

Note

These classes moved from epilearn.models.Temporal.SIR to epilearn.models.Temporal.Compartmental in 0.1.0; the old path still works as an alias. epilearn.models.SpatialTemporal.NetworkSIR.NetSIR is the equivalent model in the SpatialTemporal family and takes the same (x, adj, steps) call.