Pandemic modeling - SIR

First implementation of an epidemiology model, SIR.

About

In this notebook I implement a basic SIR model.

SIR model

From Wikipedia: SIR.

The SIR model without vital dynamics: The dynamics of an epidemic, for example the flu, are often much faster than the dynamics of birth and death, therefore, birth and death are often omitted in simple compartmental models. The SIR system without so-called vital dynamics (birth and death, sometimes called demography) described above can be expressed by the following set of ordinary differential equations:

\[ \begin{aligned} &{\frac {dS}{dt}}=-{\frac {\beta IS}{N}},\\ &{\frac {dI}{dt}}={\frac {\beta IS}{N}}-\gamma I,\\ &{\frac {dR}{dt}}=\gamma I, \end{aligned} \]

At time \(t\) the SIR model is characterised by \(S(t)\) for the number of susceptible in the population, \(I(t)\) for the number of infected, and \(R(t)\) for the number of recovered or deceased (or immune) individuals out of the population \(N = S(t)+I(t)+R(t)\).

People move through these stages as follows:

import attr
import pandas as pd
import numpy as np
from scipy.integrate import solve_ivp
import plotly.express as px

@attr.s
class SIR(object):
    N = attr.ib(converter=int)
    I = attr.ib(converter=float)
    beta = attr.ib(converter=float)
    gamma = attr.ib(converter=float)
    days = attr.ib(converter=int, default=200)
    
    S = attr.ib(init=False, converter=float)
    
    R = attr.ib(init=False, converter=float, default=0.0)
    
    
    def __attrs_post_init__(self):
        self.S = self.N - self.I - self.R
        self.beta = round(self.beta, 2)
        self.gamma = round(self.gamma, 2)
    
    
    @staticmethod
    def ode(t, y, beta, gamma, N):
        S, I, R = y
        
        new_cases = beta*S*I/N
        removed_cases = gamma*I
        S = -new_cases
        I = new_cases - removed_cases
        R = removed_cases

        y = (S,I,R)

        return y
    
    def solve(self):
        y = (self.S, self.I, self.R)
        days = self.days + 1
        sol = solve_ivp(SIR.ode, [0, days], y,t_eval=np.arange(0, days), args=(self.beta, self.gamma, self.N))
        df = pd.DataFrame(sol.y.T, index=pd.Index(sol.t,name="Time"), columns=["Susceptible","Infected","Removed"]).reset_index()
        df["beta"] = self.beta
        df["gamma"] = self.gamma
        df = df.melt(id_vars=["Time", "beta", "gamma"],
                     value_vars=["Susceptible", "Infected", "Removed"],
                     value_name="People",
                     var_name="Category")
        return (df, sol)
        
    def plot(self):
        df, sol = self.solve()
        fig = px.line(df,
                      x=df.Time,
                      y=df.People,
                      color=df.Category,
                      title="test",
                    #   title=dict(
                    #       text=self.__repr__(),
                    #       y=0.98,
                    #       x=0.5,
                    #       xanchor="center",
                    #       yanchor="top")
                     )
        fig.for_each_trace(lambda t: t.update(hovertemplate = "%{y:.0f}"))
        fig.update_layout(hovermode="x unified",
                          margin=dict(l=20, r=20, t=30, b=20),
                          legend=dict(orientation="h", title=""))
        return fig


m = SIR(5000,6,0.3,0.1)
m.plot().show()

Figure 1: A sir plot

Grid of \(\beta\) and \(\gamma\)

from itertools import product
N = 5000
I = 6

cols = np.arange(0.2,0.8, 0.1)
rows = np.arange(0.1, 0.7, 0.1)

res = []
for b, g in product(cols, rows):
    s = SIR(N, I, b, g)
    df,_ = s.solve()
    res.append(df)

dfs = pd.concat(res)

def text_to_symbol(t):
    if "beta" in t or "gamma" in t:
        s,v = t.split("=")
        return "$\{s}={v}$".format(s=s,v=v)

fig = px.line(dfs, x="Time", y="People", color="Category", facet_row="beta", facet_col="gamma")
# fig.for_each_annotation(lambda a: a.update(text=text_to_symbol(a.text)))
# fig.for_each_trace(lambda t: t.update(hovertemplate = "%{y:.0f}"))
# fig.for_each_xaxis(lambda x: x.update(title=""))
# fig.for_each_yaxis(lambda x: x.update(title=""))
# fig.update_layout(title=dict(text=r"$\text{SIR for various }\beta\text{ and }\gamma$",
#                              y=0.98,
#                              x=0.5,
#                              xanchor="center",
#                              yanchor="top"),
#                   hovermode="x unified", 
#                   legend=dict(orientation="h", title=""),
#                   margin=dict(l=20, r=20, t=40, b=10),
#                   annotations = list(fig.layout.annotations) + 
#                                 [dict(x=-0.07,
#                                       y=0.5,
#                                       showarrow=False,
#                                       text="People",
#                                       textangle=-90,
#                                       xref="paper",
#                                       yref="paper"
#                                     ),
#                                  dict(x=0.5,
#                                       y=-0.07,
#                                       showarrow=False,
#                                       text="Time",
#                                       xref="paper",
#                                       yref="paper"
#                                     )
#                                 ]
#                  )

fig.show()

Figure 2: A sir plot

test