TL:DR
How can one implement a bifurcation diagram of a seasonally forced epidemiological model such as SEIR (susceptible, exposed, infected, recovered) in Python? I already know how to implement the model itself and display a sampled time series (see this stackoverflow question), but I am struggling with reproducing figure from a textbook.
Context and My Attempt
I am trying to reproduce figures from the book "Modeling Infectious Diseases in Humans and Animals" (Keeling 2007) to both validate my implementations of models and to learn/visualize how different model parameters affect the evolution of a dynamical system.
I have found implementations of bifurcation diagrams for examples using the logistic map (see this ipython cookbook this pythonalgos bifurcation, and this stackoverflow question).
For the logistic map ($x_{n+1}=rx_n(1-x_n)$), the bifurcation diagram is plotted by computing a matrix $M \in \mathcal{R}^{T \times P}$ for $T$ iterations and $P$ parameters such that $\{r_j\}_{j=1}^{P}$ and $M_{ij}=LogisticMap(x=M_{(i-1)j}, r=r_j)$. The bifurcation diagram plots $M_{ij}\ \text{vs}\ r_j$. Based on this, I calculate the matrix $M$ for the SEIR model, which is a system of differential equations, by numerically integrating over $T$ timesteps for $P$ parameters such that $\{\beta_{1(j)}\}_{j=1}^{P}$ and the resultant matrix $M \in \mathcal{R}^{P \times T}$. I simply transpose this matrix to $M \in \mathcal{R}^{T \times P}$ and then plot this as $M_{ij}\ \text{vs}\ \beta_{1(j)}$, but with strange results as shown below.
Questions
Why do my panel outputs not match those in the figure? I assume I can't try to reproduce the bifurcation diagram from the textbook without my panels matching the output in the textbook.
I have tried to implement a function to produce a bifurcation diagram, but the output looks really strange. Am I misunderstanding something about the bifurcation diagram?
NOTE: I receive no warnings/errors during code execution.
Code to Reproduce my Figures
from typing import Callable, Dict, List, Optional, Any
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
def seasonal_seir(y: List, t: List, params: Dict[str, Any]):
"""Seasonally forced SEIR model.
Function parameters much match with those required
by `scipy.integrate.odeint`
Args:
y: Initial conditions.
t: Timesteps over which numerical solution will be computed.
params: Dict with the following key-value pairs:
beta_zero -- Average transmission rate.
beta_one -- Amplitude of seasonal forcing.
omega -- Period of forcing.
mu -- Natural mortality rate.
sigma -- Latent period for infection.
gamma -- Recovery from infection term.
Returns:
Tuple whose components are the derivatives of the
susceptible, exposed, and infected state variables
w.r.t to time.
References:
[SEIR Python Program from Textbook](http://homepages.warwick.ac.uk/~masfz/ModelingInfectiousDiseases/Chapter2/Program_2.6/Program_2_6.py)
[Seasonally Forced SIR Program from Textbook](http://homepages.warwick.ac.uk/~masfz/ModelingInfectiousDiseases/Chapter5/Program_5.1/Program_5_1.py)
"""
beta_zero = params['beta_zero']
beta_one = params['beta_one']
omega = params['omega']
mu = params['mu']
sigma = params['sigma']
gamma = params['gamma']
s, e, i = y
beta = beta_zero*(1 + beta_one*np.cos(omega*t))
sdot = mu - (beta * i + mu)*s
edot = beta*s*i - (mu + sigma)*e
idot = sigma*e - (mu + gamma)*i
return sdot, edot, idot
def plot_panels(
model: Callable,
model_params: Dict,
panel_param_space: List,
panel_param_name: str,
initial_conditions: List,
timesteps: List,
odeint_kwargs: Optional[Dict] = dict(),
x_ticks: Optional[List] = None,
time_slice: Optional[slice] = None,
state_var_ix: Optional[int] = None,
log_scale: bool = False):
"""Plot panels that are samples of the parameter space for bifurcation.
Args:
model: Function that models dynamical system. Returns dydt.
model_params: Dict whose key-value pairs are the names
of parameters in a given model and the values of those parameters.
bifurcation_parameter_space: List of varied bifurcation parameters.
bifuraction_parameter_name: The name o the bifurcation parameter.
initial_conditions: Initial conditions for numerical integration.
timesteps: Timesteps for numerical integration.
odeint_kwargs: Key word args for numerical integration.
state_var_ix: State variable in solutions to use for plot.
time_slice: Restrict the bifurcation plot to a subset
of the all solutions for numerical integration timestep space.
Returns:
Figure and axes tuple.
"""
# Set default ticks
if x_ticks is None:
x_ticks = timesteps
# Create figure
fig, axs = plt.subplots(ncols=len(panel_param_space))
# For each parameter that is varied for a given panel
# compute numerical solutions and plot
for ix, panel_param in enumerate(panel_param_space):
# update model parameters with the varied parameter
model_params[panel_param_name] = panel_param
# Compute solutions
solutions = odeint(
model,
initial_conditions,
timesteps,
args=(model_params,),
**odeint_kwargs)
# If there is a particular solution of interst, index it
# otherwise squeeze last dimension so that [T, 1] --> [T]
# where T is the max number of timesteps
if state_var_ix is not None:
solutions = solutions[:, state_var_ix]
elif state_var_ix is None and solutions.shape[-1] == 1:
solutions = np.squeeze(solutions)
else:
raise ValueError(
f'solutions to model are rank-2 tensor of shape {solutions.shape}'
' with the second dimension greater than 1. You must pass'
' a value to :param state_var_ix:')
# Slice the solutions based on the desired time range
if time_slice is not None:
solutions = solutions[time_slice]
# Natural log scale the results
if log_scale:
solutions = np.log(solutions)
# Plot the results
axs[ix].plot(x_ticks, solutions)
return fig, axs
def plot_bifurcation(
model: Callable,
model_params: Dict,
bifurcation_parameter_space: List,
bifurcation_param_name: str,
initial_conditions: List,
timesteps: List,
odeint_kwargs: Optional[Dict] = dict(),
state_var_ix: Optional[int] = None,
time_slice: Optional[slice] = None,
log_scale: bool = False):
"""Plot a bifurcation diagram of state variable from dynamical system.
Args:
model: Function that models system. Returns dydt.
model_params: Dict whose key-value pairs are the names
of parameters in a given model and the values of those parameters.
bifurcation_parameter_space: List of varied bifurcation parameters.
bifuraction_parameter_name: The name o the bifurcation parameter.
initial_conditions: Initial conditions for numerical integration.
timesteps: Timesteps for numerical integration.
odeint_kwargs: Key word args for numerical integration.
state_var_ix: State variable in solutions to use for plot.
time_slice: Restrict the bifurcation plot to a subset
of the all solutions for numerical integration timestep space.
log_scale: Flag to natural log scale solutions.
Returns:
Figure and axes tuple.
"""
# Track the solutions for each parameter
parameter_x_time_matrix = []
# Iterate through parameters
for param in bifurcation_parameter_space:
# Update the parameter dictionary for the model
model_params[bifurcation_param_name] = param
# Compute the solutions to the model using
# dictionary of parameters (including the bifurcation parameter)
solutions = odeint(
model,
initial_conditions,
timesteps,
args=(model_params, ),
**odeint_kwargs)
# If there is a particular solution of interst, index it
# otherwise squeeze last dimension so that [T, 1] --> [T]
# where T is the max number of timesteps
if state_var_ix is not None:
solutions = solutions[:, state_var_ix]
elif state_var_ix is None and solutions.shape[-1] == 1:
solutions = np.squeeze(solutions)
else:
raise ValueError(
f'solutions to model are rank-2 tensor of shape {solutions.shape}'
' with the second dimension greater than 1. You must pass'
' a value to :param state_var_ix:')
# Update the parent list of solutions for this particular
# bifurcation parameter
parameter_x_time_matrix.append(solutions)
# Cast to numpy array
parameter_x_time_matrix = np.array(parameter_x_time_matrix)
# Transpose: Bifurcation plots Function Output vs. Parameter
# This line ensures that each row in the matrix is the solution
# to a particular state variable in the system of ODEs
# a timestep t
# and each column is that solution for a particular value of
# the (varied) bifurcation parameter of interest
time_x_parameter_matrix = np.transpose(parameter_x_time_matrix)
# Slice the iterations to display to a smaller range
if time_slice is not None:
time_x_parameter_matrix = time_x_parameter_matrix[time_slice]
# Make bifurcation plot
fig, ax = plt.subplots()
# For the solutions vector at timestep plot the bifurcation
# NOTE: The elements of the solutions vector represent the
# numerical solutions at timestep t for all varied parameters
# in the parameter space
# e.g.,
# t beta1=0.025 beta1=0.030 .... beta1=0.30
# 0 solution00 solution01 .... solution0P
for sol_at_time_t_for_all_params in time_x_parameter_matrix:
if log_scale:
sol_at_time_t_for_all_params = np.log(sol_at_time_t_for_all_params)
ax.plot(
bifurcation_parameter_space,
sol_at_time_t_for_all_params,
',k',
alpha=0.25)
return fig, ax
# Define initial conditions based on figure
s0 = 6e-2
e0 = i0 = 1e-3
initial_conditions = [s0, e0, i0]
# Define model parameters based on figure
# NOTE: omega is not mentioned in the figure, but
# omega is defined elsewhere as 2pi/365
days_per_year = 365
mu = 0.02/days_per_year
beta_zero = 1250
sigma = 1/8
gamma = 1/5
omega = 2*np.pi / days_per_year
model_params = dict(
beta_zero=beta_zero,
omega=omega,
mu=mu,
sigma=sigma,
gamma=gamma)
# Define timesteps
nyears = 200
ndays = nyears * days_per_year
timesteps = np.arange(1, ndays + 1, 1)
# Define different levels of seasonality (from figure)
beta_ones = [0.025, 0.05, 0.25]
# Define the time range to actually show on the plot
min_year = 190
max_year = 200
# Create a slice of the iterations to display on the diagram
time_slice = slice(min_year*days_per_year, max_year*days_per_year)
# Get the xticks to display on the plot based on the time slice
x_ticks = timesteps[time_slice]/days_per_year
# Plot the panels using the infected state variable ix
infection_ix = 2
# Plot the panels
panel_fig, panel_ax = plot_panels(
model=seasonal_seir,
model_params=model_params,
panel_param_space=beta_ones,
panel_param_name='beta_one',
initial_conditions=initial_conditions,
timesteps=timesteps,
odeint_kwargs=dict(hmax=5),
x_ticks=x_ticks,
time_slice=time_slice,
state_var_ix=infection_ix,
log_scale=False)
# Label the panels
panel_fig.suptitle('Attempt to Reproduce Panels from Keeling 2007')
panel_fig.supxlabel('Time (years)')
panel_fig.supylabel('Fraction Infected')
panel_fig.set_size_inches(15, 8)
# Plot bifurcation
bi_fig, bi_ax = plot_bifurcation(
model=seasonal_seir,
model_params=model_params,
bifurcation_parameter_space=np.linspace(0.025, 0.3),
bifurcation_param_name='beta_one',
initial_conditions=initial_conditions,
timesteps=timesteps,
odeint_kwargs={'hmax':5},
state_var_ix=infection_ix,
time_slice=time_slice,
log_scale=False)
# Label the bifurcation
bi_fig.suptitle('Attempt to Reproduce Bifurcation Diagram from Keeling 2007')
bi_fig.supxlabel(r'$\beta_1$')
bi_fig.supylabel('Fraction Infected')
bi_fig.set_size_inches(15, 8)