import numpy as np
import matplotlib.pyplot as plt

T = 30
dt = 1
Nt = int(np.round(T/dt))
t = np.linspace(0,T,Nt+1)

S = np.zeros(Nt+1)
I = np.zeros(Nt+1)
R = np.zeros(Nt+1)
N = 500
S[0] = 400
I[0] = 50
R[0] = N-S[0]-I[0]
S = S/N; I = I/N; R = R/N
print(S[0]+I[0]+R[0])

bet = 0.5; gam1 = 0.2; gam2 = 0.3

for n in range(Nt):
S[n+1] = S[n]+dt*(-bet*S[n]*I[n]-gam2*S[n]*R[n])
I[n+1] = I[n]+dt*(bet*S[n]*I[n]-gam1*I[n]*R[n])
R[n+1] = 1-S[n+1]-I[n+1]

fig = plt.figure()
plt.grid()
plt.plot(t,S,'k-',label='S')
plt.plot(t,I,'r-',label='I')
plt.plot(t,R,'b-',label='R')
plt.axis([0,T,0,1])

plt.legend()
plt.show()