2 Lattice mobility

Example for a setup of the model with non-trivial connectivity between the metapopulations.

To keep it simple, we use a 2D lattice.

import numpy as np
import networkx as nx

from EpiCommute import SIRModel

Model setup

Create lattice

We use networkx to create a simple 2D lattice

L = 6   # Grid length
M = L*L # Number of subpopulations

G = nx.lattice.grid_2d_graph(L,L,periodic=False)
nx.draw(G)

Mobility

To add mobility, we use the adjacency matrix of the lattice and add some random noise

# --- Create directed graph ---
G = nx.DiGraph()

# Add nodes with population attribute
for i in range(M):
    G.add_node(i, population=subpopulation_sizes[i])

# Add edges with mobility weights
for i in range(M):
    for j in range(M):
        if mobility[i, j] > 0:  # you can threshold here if needed
            G.add_edge(i, j, weight=mobility[i, j])

# --- Layout ---
pos = nx.spring_layout(G, seed=42)  # nice spatial layout

# --- Node sizes (scaled) ---
node_sizes = subpopulation_sizes * 20  # adjust scale if needed

# --- Edge widths (scaled) ---
weights = np.array([G[u][v]['weight'] for u, v in G.edges()])
edge_widths = 5 * (weights / weights.max())  # normalize + scale

# --- Draw ---
plt.figure(figsize=(10, 10))

# Nodes
nx.draw_networkx_nodes(
    G, pos,
    node_size=node_sizes,
    node_color='skyblue',
    alpha=0.8
)

# Edges
nx.draw_networkx_edges(
    G, pos,
    width=edge_widths,
    alpha=0.5,
    arrows=True,
    arrowsize=15,
    arrowstyle='->'
)

# Labels
nx.draw_networkx_labels(G, pos, font_size=8)

plt.title("Red de movilidad (metapoblación)")
plt.axis('off')
plt.show()

# Adjacency matrix
A = nx.adjacency_matrix(G).toarray() + np.identity(M)

# Mobility matrix is adjacency matrix with some noise
mobility = (np.random.rand(M, M) + 1) * A

# Choose random subpopulation sizes
subpopulation_sizes = np.random.randint(50,100,M)

len(subpopulation_sizes)

# outbreak source is the first compartment
outbreak_source=0

# Initialize the model
model = SIRModel(
            mobility,
            subpopulation_sizes,
            outbreak_source=outbreak_source,
            dt=0.1,                   # simulation time interval
            dt_save=1,                # time interval when to save observables
            I0=10,                    # number of initial infected
            VERBOSE=True              # print verbose output
        )

result = model.run_simulation()
result['S'][0].shape

Starting Simulation ...
Simulation completed
Time: 0min 5.53s

(36,)

Results

Epidemic curve

import matplotlib
import matplotlib.pyplot as plt
matplotlib.rc('figure', dpi=200)

figure = plt.figure(figsize=(3.5,2))
plt.plot(result['t'], result['S_total'], label='S', color='purple')
plt.plot(result['t'], result['I_total'], label='I', color='firebrick')
plt.plot(result['t'], result['R_total'], label='R', color='k')
plt.legend(frameon=False, loc='center right')
plt.xlabel("time")
plt.ylabel("compartment")
plt.show()

Spread of infected

fig, axes = plt.subplots(1,6,figsize=(10,2))
t_indices = [0,5,10,20,30,40]
for t_index, ax in zip(t_indices, axes):
    
    ax.imshow(result['I'][t_index].reshape(L,L), vmax=1, interpolation=None, cmap='Reds')
    ax.set_title("t = {:.0f}".format(result['t'][t_index]))
    ax.axis('off')

plt.show()

Susceptibles

fig, axes = plt.subplots(1,6,figsize=(10,2))
t_indices = [0,5,10,20,30,40]
for t_index, ax in zip(t_indices, axes):
    
    ax.imshow(result['S'][t_index].reshape(L,L), vmax=1, interpolation=None, cmap='Greens')
    ax.set_title("t = {:.0f}".format(result['t'][t_index]))
    ax.axis('off')

plt.show()

Arrival times

fig = plt.figure(figsize=(3,3))
im = plt.imshow(result['T_arrival'].reshape(L,L), cmap='viridis_r', interpolation=None)
cax = fig.colorbar(im)
cax.set_label("arrival time t")
# plt.gca().axis('off')
plt.show()