import numpy as np
import matplotlib.pyplot as plt
# S0, IWT0, IMUT0, R0: initial sizes of subpopulations
# a: removal rate of infected, rWT and rMUT infection rates
# T: solve the equations until time T
# ts: array of time points
def solve_SIR(S0, IWT0, IMUT0, R0, rWT, rMUT, a, ts):
#initialize the populations vectors
S = np.zeros(ts.size)
IWT = np.zeros(ts.size)
IMUT = np.zeros(ts.size)
R = np.zeros(ts.size)
N = S0 + IWT0 + IMUT0 + R0
S[0] = S0
IWT[0] = IWT0
IMUT[0] = IMUT0
R[0] = R0
# Euler method loop
for tidx in range(1,len(ts)):
# update S I R using the values from the last timestep
S[tidx] = S[tidx-1] - (rWT*IWT[tidx-1]*S[tidx-1]/N + rMUT*IMUT[tidx-1]*S[tidx-1]/N) * dt
IWT[tidx] = IWT[tidx-1] + (rWT*IWT[tidx-1]*S[tidx-1]/N - a*IWT[tidx-1]) * dt
IMUT[tidx] = IMUT[tidx-1] + (rMUT*IMUT[tidx-1]*S[tidx-1]/N - a*IMUT[tidx-1]) * dt
R[tidx] = R[tidx-1] + a*(IWT[tidx-1]+IMUT[tidx-1]) * dt
return S, IWT, IMUT, R
a = 1
rWT = 3
rMUT = 6
T = 10
plt.figure()
dt = 1.0e-5
ts = np.arange(0, T, dt)
S, IWT, IMUT, R = solve_SIR(8000, 50, 10, 1940, rWT, rMUT, a, ts)
plt.plot(ts, S)
plt.plot(ts, IMUT)
plt.plot(ts, IWT)
plt.plot(ts, R)
print(ts[-1], IMUT[-1] / (IWT[-1] + IMUT[-1]))
dt = 1.0e-3
ts = np.arange(0, T, dt)
S, IWT, IMUT, R = solve_SIR(8000, 50, 10, 1940, rWT, rMUT, a, ts)
plt.plot(ts, S)
plt.plot(ts, IMUT)
plt.plot(ts, IWT)
plt.plot(ts, R)
print(ts[-1], IMUT[-1] / (IWT[-1] + IMUT[-1]))
plt.show()
plt.figure()
plt.plot(ts, IMUT/(IMUT+IWT))
plt.show()