%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import math
N = 1000
I = 1
S = N - I
R = 0
beta = 0.0002
gamma = 0.05
sim_length = 120
dt = 0.01
xs = []
y1s = []
y2s = []
y3s = []
num_iter = int(sim_length / dt) + 1
for i in range(num_iter):
xs.append(i * dt)
y1s.append(S)
y2s.append(I)
y3s.append(R)
Sold = S
Iold = I
Rold = R
S += (-beta * Iold * Sold) * dt
I += (beta * Iold * Sold - gamma * Iold) * dt
R += (gamma * Iold) * dt
plt.plot(xs, y1s)
plt.plot(xs, y2s)
plt.plot(xs, y3s)
plt.xlabel("Time")
plt.ylabel("Population")
plt.title("SIR Epidemiology Model")
plt.legend(['Susceptible', 'Infected', 'Recovered'])
<matplotlib.legend.Legend at 0x108766b90>