Skip to content

Creating a custom compartmental model

Open In Colab

Note

You may want to create a custom compartmental model to finetune COVID19 analysis or realize models for new diseases. This tutorial will guide you through this creation.

Developer import
%matplotlib inline
%load_ext autoreload
%autoreload 2

# Developer import
import sys
sys.path.append("../../")
On Google Colab

Uncomment the following line to install the library locally

# !pip install pyepidemics
Verify the library is correctly installed
import pyepidemics

Creating a SIR model with newborns

Starting from a simple SIR model

To create a custom model, you need to override the CompartmentalModel class. There are plenty of options to customize the custom model, but here you can see a minimum viable example with :

  • Definition of the compartments
  • Definition of the transitions

Below you will learn how to create more complex models such as the COVID19 one.

# Verify you can import the base class
from pyepidemics.models import CompartmentalModel


class SIR(CompartmentalModel):
    def __init__(self,N,beta,gamma):

        # Define compartments name and number
        compartments = ["S","I","R"]
        super().__init__(compartments)

        # Parameters
        self.N = N # Total population

        # Add transition
        self.add_transition("S","I",lambda y,t: beta * y["S"] * y["I"] / self.N)
        self.add_transition("I","R",lambda y,t: gamma * y["I"])

Initialize the custom model

N = 67e6
beta = 3.3/4
gamma = 1/4

sir = SIR(N,beta,gamma)

Visualize the propagation

states = sir.solve(n_days = 100)
states.show(plotly = False)

Adding a newborn compartment

For some diseases like measles, we need to account for newborns with maternally-derived immunity, ie newborn are immune thanks to maternal antibodies. This can be modeled with a MSIR model we will learn to build in this section. To learn more about this model and other compartmental models the wikipedia page is pretty complete.

We want to add a newborn compartment with the following characteristics :

  • Children are born at a given birth_rate, ie the number of new children born per day
  • Children are not susceptible when they are born, they become susceptible after 1/susceptibility_rate days

For that we easily add the following :

  • Add a new compartment in the list
  • Specify the first state is the susceptible state and the starting state of the epidemics is the I compartment
  • Add a transition between N and S using the susceptibility_rate
  • Add a static derivative (ie a derivative only applied to a compartment and not a transition between compartments, it models when there is an individual is not changing from one state to another) representing the birth rate

NB: the number of people across compartments is not constant in this example compared to the classic SIR model, because of the birth rate we have constantly new individuals in the simulation over time. 

class MSIR(CompartmentalModel):
    def __init__(self,N,beta,gamma,birth_rate,susceptibility_rate):

        # Define compartments name and number
        compartments = ["M","S","I","R"]
        super().__init__(compartments,initial_state = "S",start_state = "I")

        # Parameters
        self.N = N # Total population

        # Add transition
        self.add_static_derivative("M",lambda y,t: birth_rate)
        self.add_transition("M","S",lambda y,t: susceptibility_rate * y["M"])
        self.add_transition("S","I",lambda y,t: beta * y["S"] * y["I"] / self.N)
        self.add_transition("I","R",lambda y,t: gamma * y["I"])

N = 1000
beta = 1/10 # 1 child is contaminated every 10 days
gamma = 1/30 # You recover in 30 days
birth_rate = 1 # 1 child is born per day
susceptibility_rate = 1/30 # a child become susceptible in 20 days

msir = MSIR(N,beta,gamma,birth_rate,susceptibility_rate)

states = msir.solve(n_days = 1000)
states.show(plotly = False)

As we could expect we have a first peak of infections in the first days and the current population reach herd immunity.
After a few days, sufficient children are born and a second wave starts.

Exercises

You can try changing this model to account for : - Vaccinations (for newborns and non-newborns) by adding a compartment V - Deceases by adding a compartment D

Modelling Fox & Rabbit populations

Such models can also be used to model the dynamics of any deterministic system. For example if you want to model the famous predator-prey model, aka Fox&Rabbit model, aka Lotka-Volterra competition problem you can also use the pyepidemics library as shown below :

class LotkaVolterra(CompartmentalModel):
    def __init__(self,alpha,beta,gamma,delta):

        # Define compartments name and number
        compartments = ["fox","rabbit"]
        super().__init__(compartments)

        # Add transition
        self.add_static_derivative("rabbit",lambda y,t: alpha * y["rabbit"] - beta * y["rabbit"] * y["fox"])
        self.add_static_derivative("fox",lambda y,t: beta * delta * y["rabbit"] * y["fox"] - gamma * y["fox"] )

alpha = 1.
beta = 0.1
gamma = 1.5
delta = 0.75

lotka = LotkaVolterra(alpha,beta,delta,gamma)
states = lotka.solve(n_days = 30,init_state = {"rabbit":10,"fox":5},d = 100)
states.show(plotly = False)

Advanced users

There are many other parameters you can configure in custom models:

  • I0, the number of infected at t = 0
  • start_date to deal with dates and not days
  • initial_state and start_state, start means at which state the epidemics start, for a SIR model initial_state="S" and start_state="I"
Super-advanced users
  • Overriding derivatives: Behind the CompartmentalModel a network is constructed and derivatives are computed using network transitions. Yet for any model you could override
  • Adding granularity: to build models by categories (age, region) using contact matrix, it's also easy to build this behavior with pyepidemics