Simple Spatial SIR Model with Synthetic Data¶
This example simulates a spatial Susceptible-Infected-Recovered (SIR) model with modular components. The population is distributed across 20 nodes in a 1-D chain, with infection spreading spatially from node 0 and agents migrating based on a migration matrix.
Spatial Arrangement¶
The model uses a 1-D spatial structure: - Node connections: Each node is connected to its next neighbor in the chain, allowing migration to occur sequentially (e.g., 0 → 1 → 2 …). - Infection propagation: Infection starts in node 0 and spreads to neighboring nodes as agents migrate.
Two migration matrix options are available: 1. Sequential Migration Matrix: Agents can only move to their next node in the chain. 2. Gravity Model Migration Matrix: This provides a more realistic 2-D spatial dynamic, where migration probabilities depend on node distances and population sizes.
Population Initialization¶
Skewed Distribution: The population is distributed across nodes using a rural-to-urban skew.
Timers: Migration timers are assigned to control agent migration frequency.
Model Components¶
The simulation uses modular components for migration, transmission, and recovery dynamics. Each component encapsulates the logic for its specific task.
Full Code Implementation¶
Imports¶
1import numpy as np
2import matplotlib.pyplot as plt
3import csv
4from laser_core.laserframe import LaserFrame
5from laser_core.demographics.spatialpops import distribute_population_skewed as dps
6from laser_core.migration import gravity
Model Class¶
1# Define the model
2class MultiNodeSIRModel:
3 """
4 A multi-node SIR (Susceptible-Infected-Recovered) disease transmission model.
5
6 Attributes:
7 params (dict): Configuration parameters for the model.
8 nodes (int): Number of nodes in the simulation.
9 population (LaserFrame): Represents the population with agent-level properties.
10 results (np.ndarray): Stores simulation results for S, I, and R per node.
11 components (list): List of components (e.g., Migration, Transmission) added to the model.
12 """
13
14 def __init__(self, params):
15 """
16 Initializes the SIR model with the given parameters.
17
18 Args:
19 params (dict): Dictionary containing parameters such as population size,
20 number of nodes, timesteps, and rates for transmission/migration.
21 """
22 self.params = params
23 self.nodes = params["nodes"]
24 self.population = LaserFrame(capacity=params["population_size"], initial_count=params["population_size"])
25
26 # Define properties
27 self.population.add_scalar_property("node_id", dtype=np.int32)
28 self.population.add_scalar_property("disease_state", dtype=np.int32, default=0) # 0=S, 1=I, 2=R
29 self.population.add_scalar_property("recovery_timer", dtype=np.int32, default=0)
30 self.population.add_scalar_property("migration_timer", dtype=np.int32, default=0)
31
32 # Initialize population distribution
33 node_pops = dps(params["population_size"], self.nodes, frac_rural=0.3)
34 node_ids = np.concatenate([np.full(count, i) for i, count in enumerate(node_pops)])
35 np.random.shuffle(node_ids)
36 self.population.node_id[:params["population_size"]] = node_ids
37
38 # Reporting structure
39 self.results = np.zeros((params["timesteps"], self.nodes, 3)) # S, I, R per node
40
41 # Components
42 self.components = []
43
44 def add_component(self, component):
45 """
46 Adds a component (e.g., Migration, Transmission, Recovery) to the model.
47
48 Args:
49 component: An instance of a component to be added.
50 """
51 self.components.append(component)
52
53 def step(self):
54 """
55 Advances the simulation by one timestep, updating all components and recording results.
56 """
57 for component in self.components:
58 component.step()
59
60 # Record results
61 for i in range(self.nodes):
62 in_node = self.population.node_id == i
63 self.results[self.current_timestep, i, 0] = (self.population.disease_state[in_node] == 0).sum()
64 self.results[self.current_timestep, i, 1] = (self.population.disease_state[in_node] == 1).sum()
65 self.results[self.current_timestep, i, 2] = (self.population.disease_state[in_node] == 2).sum()
66
67 def run(self):
68 """
69 Runs the simulation for the configured number of timesteps.
70 """
71 from tqdm import tqdm
72 for self.current_timestep in tqdm(range(self.params["timesteps"])):
73 self.step()
74
75 def save_results(self, filename):
76 """
77 Saves the simulation results to a CSV file.
78
79 Args:
80 filename (str): Path to the output file.
81 """
82 with open(filename, mode='w', newline='') as file:
83 writer = csv.writer(file)
84 writer.writerow(["Timestep", "Node", "Susceptible", "Infected", "Recovered"])
85 for t in range(self.params["timesteps"]):
86 for node in range(self.nodes):
87 writer.writerow([t, node, *self.results[t, node]])
88
89 def plot_results(self):
90 """
91 Plots the prevalence of infected agents over time for all nodes.
92 """
93 plt.figure(figsize=(10, 6))
94 for i in range(self.nodes):
95 prevalence = self.results[:, i, 1] / self.results[:, i, :].sum(axis=1)
96 plt.plot(prevalence, label=f"Node {i}")
97 plt.title("Prevalence Across All Nodes")
98 plt.xlabel("Timesteps")
99 plt.ylabel("Prevalence of Infected Agents")
100 plt.legend()
101 plt.show()
Migration Component Class¶
1class MigrationComponent:
2 """
3 Handles migration behavior of agents between nodes in the model.
4
5 Attributes:
6 model (MultiNodeSIRModel): The simulation model instance.
7 migration_matrix (ndarray): A matrix representing migration probabilities between nodes.
8 """
9
10 def __init__(self, model):
11 """
12 Initializes the MigrationComponent.
13
14 Args:
15 model (MultiNodeSIRModel): The simulation model instance.
16 """
17 self.model = model
18
19 # Set initial migration timers
20 max_timer = int(1 / model.params["migration_rate"])
21 model.population.migration_timer[:] = np.random.randint(1, max_timer + 1, size=model.params["population_size"])
22
23 self.migration_matrix = self.get_sequential_migration_matrix(model.nodes)
24
25 # Example customization: Disable migration from node 13 to 14
26 def break_matrix_node(matrix, from_node, to_node):
27 matrix[from_node][to_node] = 0
28 break_matrix_node(self.migration_matrix, 13, 14)
29
30 def get_gravity_migration_matrix(self, nodes):
31 """
32 Generates a gravity-based migration matrix based on population and distances between nodes.
33
34 Args:
35 nodes (int): Number of nodes in the simulation.
36
37 Returns:
38 ndarray: A migration matrix where each row represents probabilities of migration to other nodes.
39 """
40 pops = np.array([np.sum(self.model.population.node_id == i) for i in range(nodes)])
41 distances = np.ones((nodes, nodes)) - np.eye(nodes)
42 migration_matrix = gravity(pops, distances, k=1.0, a=0.5, b=0.5, c=2.0)
43 migration_matrix = migration_matrix / migration_matrix.sum(axis=1, keepdims=True)
44 return migration_matrix
45
46 def get_sequential_migration_matrix(self, nodes):
47 """
48 Creates a migration matrix where agents can only migrate to the next sequential node.
49
50 Args:
51 nodes (int): Number of nodes in the simulation.
52
53 Returns:
54 ndarray: A migration matrix where migration is allowed only to the next node.
55 """
56 migration_matrix = np.zeros((nodes, nodes))
57 for i in range(nodes - 1):
58 migration_matrix[i, i + 1] = 1.0
59 return migration_matrix
60
61 def step(self):
62 """
63 Updates the migration state of the population by determining which agents migrate
64 and their destinations based on the migration matrix.
65 """
66 node_ids = self.model.population.node_id
67
68 # Decrement migration timers
69 self.model.population.migration_timer -= 1
70
71 # Identify agents ready to migrate
72 migrating_indices = np.where(self.model.population.migration_timer <= 0)[0]
73 if migrating_indices.size == 0:
74 return
75
76 # Shuffle migrants and assign destinations based on migration matrix
77 np.random.shuffle(migrating_indices)
78 destinations = np.empty(len(migrating_indices), dtype=int)
79 for origin in range(self.model.nodes):
80 origin_mask = node_ids[migrating_indices] == origin
81 num_origin_migrants = origin_mask.sum()
82
83 if num_origin_migrants > 0:
84 # Assign destinations proportionally to migration matrix
85 destination_counts = np.round(self.migration_matrix[origin] * num_origin_migrants).astype(int)
86 destination_counts = np.maximum(destination_counts, 0) # Clip negative values
87 if destination_counts.sum() == 0: # No valid destinations
88 destinations[origin_mask] = origin # Stay in the same node
89 continue
90 destination_counts[origin] += num_origin_migrants - destination_counts.sum() # Adjust rounding errors
91
92 # Create ordered destination assignments
93 destination_indices = np.repeat(np.arange(self.model.nodes), destination_counts)
94 destinations[origin_mask] = destination_indices[:num_origin_migrants]
95
96 # Update node IDs of migrants
97 node_ids[migrating_indices] = destinations
98
99 # Reset migration timers for migrated agents
100 self.model.population.migration_timer[migrating_indices] = np.random.randint(
101 1, int(1 / self.model.params["migration_rate"]) + 1, size=migrating_indices.size
102 )
Transmission Component Class¶
1class TransmissionComponent:
2 """
3 Handles the disease transmission dynamics within the population.
4
5 Attributes:
6 model (MultiNodeSIRModel): The simulation model instance.
7 """
8
9 def __init__(self, model):
10 """
11 Initializes the TransmissionComponent and infects initial agents.
12
13 Args:
14 model (MultiNodeSIRModel): The simulation model instance.
15 """
16 self.model = model
17
18 def step(self):
19 """
20 Simulates disease transmission for each node in the current timestep.
21 """
22 for i in range(self.model.nodes):
23 in_node = self.model.population.node_id == i
24 susceptible = in_node & (self.model.population.disease_state == 0)
25 infected = in_node & (self.model.population.disease_state == 1)
26
27 num_susceptible = susceptible.sum()
28 num_infected = infected.sum()
29 total_in_node = in_node.sum()
30
31 if total_in_node > 0 and num_infected > 0 and num_susceptible > 0:
32 infectious_fraction = num_infected / total_in_node
33 susceptible_fraction = num_susceptible / total_in_node
34
35 new_infections = int(
36 self.model.params["transmission_rate"] * infectious_fraction * susceptible_fraction * total_in_node
37 )
38
39 susceptible_indices = np.where(susceptible)[0]
40 newly_infected_indices = np.random.choice(susceptible_indices, size=new_infections, replace=False)
41
42 self.model.population.disease_state[newly_infected_indices] = 1
43 self.model.population.recovery_timer[newly_infected_indices] = np.random.randint(5, 15, size=new_infections)
Recovery Component Class¶
1class RecoveryComponent:
2 """
3 Handles the recovery dynamics of infected individuals in the population.
4
5 Attributes:
6 model (MultiNodeSIRModel): The simulation model instance.
7 """
8
9 def __init__(self, model):
10 """
11 Initializes the RecoveryComponent.
12
13 Args:
14 model (MultiNodeSIRModel): The simulation model instance.
15 """
16 self.model = model
17
18 def step(self):
19 """
20 Updates the recovery state of infected individuals, moving them to the recovered state
21 if their recovery timer has elapsed.
22 """
23 infected = self.model.population.disease_state == 1
24 self.model.population.recovery_timer[infected] -= 1
25 self.model.population.disease_state[(infected) & (self.model.population.recovery_timer <= 0)] = 2
Run Everything¶
1# Parameters
2params = {
3 "population_size": 1_000_000,
4 "nodes": 20,
5 "timesteps": 600,
6 "initial_infected_fraction": 0.01,
7 "transmission_rate": 0.25,
8 "migration_rate": 0.001
9}
10
11# Run simulation
12model = MultiNodeSIRModel(params)
13model.add_component(MigrationComponent(model))
14model.add_component(TransmissionComponent(model))
15model.add_component(RecoveryComponent(model))
16model.run()
17model.save_results("simulation_results.csv")
18model.plot_results()