# Solving of SIR differential equations with Python
# Copyright (c) 2020 Jiri Kriz, www.nosco.ch
#
# Remark:
# The program integrates all 3 quantities:
# s(t) = susceptible
# i(t) = infectious
# r(t) = recovered
# It would suffice to integrate only s(t) and i(t)
# because r(t) = 1 - s(t) - i(t)

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
%matplotlib qt5

# function that returns dz/dt
def model(z, t):
    s = z[0]
    i = z[1]
    r = z[2]

    beta = 0.5
    gamma = 0.05
    
    ds_dt = - beta * i * s 
    di_dt =   beta * i * s  - gamma * i
    dr_dt =   gamma * i
    dz_dt = [ds_dt, di_dt, dr_dt]
    return dz_dt

def main():
    # initial condition
    i0 = 0.01
    s0 = 1 - i0
    r0 = 0
    z0 = [s0, i0, r0]
    
    # time points
    t = np.linspace(0, 100, 100)
    
    # solve ODE
    z = odeint(model, z0, t)
    
    # plot results
    plt.title("SIR-model  i0=0.01  beta=0.5  gamma=0.05")
    plt.plot(t, z[:, 0], 'b-', label='Susceptible s')
    plt.plot(t, z[:, 1], 'r-', label='Infectious i')
    plt.plot(t, z[:, 2], 'g-', label='Recovered r')
    plt.ylabel('Proportions')
    plt.xlabel('Days t')
    plt.legend(loc='best')
    plt.show()    
    
main()