Problem 2

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

# Define starting conditions
S0 = 0.99
I0 = 0.01
R0 = 1-S0-I0
days = 30
dt = 0.01
steps = round(days / dt)

# Define values for beta and gamma
beta = 1
gamma = 0.4

# Define differential equations
def SIR_equations(y, t, beta, gamma):
    S, I, R = y
    dS = -beta * S * I
    dI = (beta * S - gamma) * I
    dR = gamma * I
    return dS, dI, dR

# Time vector, for the plot
t = np.linspace(0,days,steps)

# Find solution
final = odeint(SIR_equations, [S0, I0, R0], t, args=(beta, gamma))
S, I, R = np.transpose(final)

# Plot the distribution
plt.title('For β= '+str(beta)+' and γ = '+str(gamma))
plt.plot(t, S, label='Susceptible')
plt.plot(t, I, label='Infected')
plt.plot(t, R, label='Recovered')
plt.ylim(0, 1)
plt.xlim(0, days)
plt.ylabel('Fraction of people')
plt.xlabel('Amount of days')
plt.legend()
plt.show()