epiworldpy: Python bindings for epiworld

This is a python wrapper of the epiworld c++ library, an ABM simulation engine. This is possible using the pybind11 library (which rocks!).

The epiworld module is already implemented in R.

Installation

API

You can find API documentation on the API page.

Examples

Basic

Here we show how to create a SEIR object and add terms to it. We will use the following data:

# Loading the module
import epiworldpy as epiworld

# Create a SEIR model (susceptible, exposed, infectious, recovered), representing COVID-19.
covid19 = epiworld.ModelSEIRCONN(
  name              = 'covid-19',
  n                 = 10000,
  prevalence        = .01,
  contact_rate      = 2.0,
  transmission_rate = .1,
  incubation_days   = 7.0,
  recovery_rate     = 0.14
)

# Taking a look
covid19.print(False)

________________________________________________________________________________
________________________________________________________________________________
SIMULATION STUDY

Name of the model   : Susceptible-Exposed-Infected-Removed (SEIR) (connected)
Population size     : 10000
Agents' data        : (none)
Number of entities  : 0
Days (duration)     : 0 (of 0)
Number of viruses   : 1
Last run elapsed t  : -
Rewiring            : off

Global events:
 - Update infected individuals (runs daily)

Virus(es):
 - covid-19

Tool(s):
 (none)

Model parameters:
 - Avg. Incubation days : 7.0000
 - Contact rate         : 2.0000
 - Prob. Recovery       : 0.1400
 - Prob. Transmission   : 0.1000

<epiworldpy._core.ModelSEIRCONN at 0x7efd602994f0>

Let’s run it and to see what we get:

# Run for 100 days with a seed of 223.
covid19.run(100, 223)

# Print an overview.
covid19.print(False)

_________________________________________________________________________
Running the model...
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||| done.
 done.
________________________________________________________________________________
________________________________________________________________________________
SIMULATION STUDY

Name of the model   : Susceptible-Exposed-Infected-Removed (SEIR) (connected)
Population size     : 10000
Agents' data        : (none)
Number of entities  : 0
Days (duration)     : 100 (of 100)
Number of viruses   : 1
Last run elapsed t  : 41.00ms
Last run speed      : 24.09 million agents x day / second
Rewiring            : off

Global events:
 - Update infected individuals (runs daily)

Virus(es):
 - covid-19

Tool(s):
 (none)

Model parameters:
 - Avg. Incubation days : 7.0000
 - Contact rate         : 2.0000
 - Prob. Recovery       : 0.1400
 - Prob. Transmission   : 0.1000

Distribution of the population at time 100:
  - (0) Susceptible :  9900 -> 7540
  - (1) Exposed     :   100 -> 229
  - (2) Infected    :     0 -> 280
  - (3) Recovered   :     0 -> 1951

Transition Probabilities:
 - Susceptible  1.00  0.00  0.00  0.00
 - Exposed      0.00  0.86  0.14  0.00
 - Infected     0.00  0.00  0.86  0.14
 - Recovered    0.00  0.00  0.00  1.00

<epiworldpy._core.ModelSEIRCONN at 0x7efd602994f0>

We can now visualize the model’s compartments:

import numpy as np
import matplotlib.pyplot as plt

# Get the data from the database
history = covid19.get_db().get_hist_total()

# Extract unique states and dates
unique_states = np.unique(history['states'])
unique_dates = np.unique(history['dates'])

# Remove some data that will mess with scaling
unique_states = np.delete(unique_states, np.where(unique_states == 'Susceptible'))

# Initialize a dictionary to store time series data for each state
time_series_data = {state: [] for state in unique_states}

# Populate the time series data for each state
for state in unique_states:
  for date in unique_dates:
    # Get the count for the current state and date
    mask = (history['states'] == state) & (history['dates'] == date)
    count = history['counts'][mask][0]
    time_series_data[state].append(count)

# Start the plotting!
plt.figure(figsize=(10, 6))

for state in unique_states:
  plt.plot(unique_dates, time_series_data[state], label=state)

plt.xlabel('Day')
plt.ylabel('Count')
plt.title('COVID-19 SEIR Model Data')
plt.legend()
plt.grid(True)
plt.show()

We can get the effective reproductive number, over time, too:

reproductive_data = covid19.get_db().get_reproductive_number()

# Start the plotting!
plt.figure(figsize=(10, 6))

for virus_id, virus_data in enumerate(reproductive_data):
    average_rts = list()

    for date_data in virus_data:
        if not date_data:
            continue

        keys_array = np.array(list(date_data.values()), dtype=np.float64)
        average_rts.append(np.mean(keys_array))

    plt.plot(range(0, len(virus_data)-1), average_rts, label=f"Virus {virus_id}")

plt.xlabel('Date')
plt.ylabel('Effective Reproductive Rate')
plt.title('COVID-19 SEIR Model Effective Reproductive Rate')
plt.legend()
plt.grid(True)
plt.show()

Let’s do the same for generation time:

from collections import defaultdict

generation_time = covid19.get_db().get_generation_time()
agents = generation_time['agents']
viruses = generation_time['viruses']
times = generation_time['times']
gentimes = generation_time['gentimes']

# Data formatting
unique_viruses = np.unique(viruses)
data = defaultdict(lambda: defaultdict(list))

for agent, virus, time, gentime in zip(agents, viruses, times, gentimes):
    data[virus][time].append(gentime)

average_data = {virus: {} for virus in unique_viruses}

for virus, time_dict in data.items():
    for time, gentime_list in time_dict.items():
        average_data[virus][time] = np.mean(gentime_list)

# Plotting
plt.figure(figsize=(10, 6))
for virus, time_dict in average_data.items():
    times = sorted(time_dict.keys())
    gentimes = [time_dict[time] for time in times]
    plt.plot(times, gentimes, label=f'Virus {virus}')

plt.xlabel('Date')
plt.ylabel('Generation Time')
plt.title('COVID-19 SEIR Model Generation Time')
plt.legend()
plt.grid(True)
plt.show()

Epiworld records agent-agent interactions, and we can graph those too. In the below example, we only track all cases stemming from a specific index case, despite the model having a prevalence of 0.01.

import networkx as nx
from matplotlib.animation import FuncAnimation

transmissions = covid19.get_db().get_transmissions()
start = transmissions['source_exposure_dates']
end = transmissions['dates']
source = transmissions['sources']
target = transmissions['targets']
days = max(end)

graph = nx.Graph()
fig, ax = plt.subplots(figsize=(6,4))

# Animation function
to_track = { source[0] }
def update(frame):
    ax.clear()

    agents_involved_today = set()
    agents_relationships_we_care_about = []

    # Get only the agents involved in the current frame.
    for i in range(len(start)):
        if start[i] <= frame <= end[i]:
            agents_involved_today.add((source[i], target[i]))

    # Get only today's agents who have some connection to agents
    # we've seen before.
    for agent in agents_involved_today:
        if agent[0] in to_track or agent[1] in to_track:
            to_track.add(agent[0])
            to_track.add(agent[1])
            graph.add_edge(agent[0], agent[1])

    # Lay and space them out.
    pos = nx.kamada_kawai_layout(graph)

    options = {
        "with_labels": True,
        "node_size": 300,
        "font_size": 6,
        "node_color": "white",
        "edgecolors": "white",
        "linewidths": 1,
        "width": 1,
    }

    # Graph!
    nx.draw_networkx(graph, pos, **options)
    ax.set_title(f"COVID-19 SEIR Model Agent Contact (Day {frame})")

ani = FuncAnimation(fig, update, frames=int(days/3), interval=200, repeat=False)
plt.show()