# S0, IWT0, R0: initial sizes of subpopulations
# a: removal rate of infected, rWT infection rate
# T: solve the equations until time T
# ts: array of time points
def solve_SIR(S0, IWT0, R0, rWT, a, ts):
#initialize the populations vectors
S = np.zeros(ts.size)
IWT = np.zeros(ts.size)
R = np.zeros(ts.size)
N = S0 + IWT0 + R0
S[0] = S0
IWT[0] = IWT0
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) * dt
IWT[tidx] = IWT[tidx-1] + (rWT*IWT[tidx-1]*S[tidx-1]/N - a*IWT[tidx-1]) * dt
R[tidx] = R[tidx-1] + (a*IWT[tidx-1]) * dt
return S, IWT, R