#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Created on Wed Apr 8 15:46:15 2020 @author: shubha """ # Use scipy.integrate.odeint to solve the SIR model import numpy as np import matplotlib.pyplot as plt from scipy.integrate import odeint def fd(yvec,t): """ Computes dsdt,di/dt, dr/dt""" s,i,r = yvec # unpack vector into the three dependent variables dsdt = -trate * s * i # susceptible didt = trate * s * i - rrate * i #infected drdt = rrate * i #recovered return dsdt,didt,drdt # pack derivatives into vector and return t = np.linspace(0,100) # times for simulation #i0 = float(input("What percent of the population is infected at the start? ")) #trate = float(input("What is the transmission rate? ")) #rrate = float(input("What is the recovery rate? ")) i0 = 0.1 trate = 0.1 #transmission rate rrate = 0.05 #recovery rate R0 = trate/rrate #basic reproduction rate s0 = 1.0 - i0 # initial susceptible population r0 = 0 # initial recovered population sir = odeint(fd,(s0,i0,r0),t) # integrate using odeint, with initial conditions provided as a tuple plt.plot(t,sir[:,0],label='Susceptible') # plot susceptible vs time plt.plot(t,sir[:,1],label='Infected') # plot infected vs time plt.plot(t,sir[:,2],label='Recovered') # plot recovered vs time plt.xlabel('Time') plt.title('R0 = %f' %R0) plt.legend() plt.show()