diff --git a/covid/python/seir.py b/covid/python/seir.py index a9d852f2cf1ef0ae8945470190a57ff65b1eb463..9142342ba441a42c9dc255ac59ed750530ae68cd 100644 --- a/covid/python/seir.py +++ b/covid/python/seir.py @@ -1,12 +1,22 @@ import numpy as np import matplotlib.pyplot as plt import functools +import matplotlib.animation as animation swiss = np.array([2.100350058, 3.150525088, 4.900816803, 6.534422404, 10.501750292, 13.302217036, 24.970828471, 31.271878646, 39.323220537, 43.640606768, 57.292882147, 76.079346558, 100.116686114, 131.271878646, 158.576429405, 256.709451575, 256.709451575, 268.378063011, 309.218203034, 353.325554259, 453.675612602]) swiss = np.array([18, 27, 42, 56, 90, 114, 214, 268, 337, 374, 491, 652, 858, 1125, 1359, 2200, 2200, 2300, 2650, 3028, 3888]) days = np.array(range(1,len(swiss)+1)) +def update(i, lines, t, p): + n_pop, length = p.shape + for j in range(0, n_pop): + lines[j].set_data(t[0:i], p[j][0:i]) + return lines + +def policy_r0(R_i, t): + + def seir(y, t, R_0, Tinf, Tinc): N = np.sum(y) y1 = np.zeros(4) @@ -16,6 +26,18 @@ def seir(y, t, R_0, Tinf, Tinc): y1[3] = 1.0 / Tinf * y[2] return y1 + +def seihse(y, t, R_0, Tinf, Tinc, Thos, Tsev): + N = np.sum(y) + y1 = np.zeros(6) + y1[0] = - R_0 / Tinf * y[2] * y[0] / N # Sane + y1[1] = R_0 / Tinf * y[2] * y[0] / N - 1.0 / Tinc * y[1] # Exposed + y1[2] = 1.0 / Tinc * y[1] - 1.0 / Tinf * y[2] - 1.0 / Tinf * y[2] # Infectious + y1[3] = 1.0 / Tinf * y[2] + 1.0 / Thos * y[4] + 1.0 / Tsev * y[5] # Recovered + y1[4] = - 1.0 / Thos * y[4] + 1.0 / Tinf * y[2] - 1.0 / Tsev * y[4] # Hospitalized + y1[5] = - 1.0 / Tsev * y[5] + 1.0 / Thos * y[4] # Severe + + return y1 def rk4(F, y, t, dt): k1 = dt * F(y, t) @@ -30,33 +52,35 @@ R_02 = 1 - 1/R_01 Tinf = 7.0 Tinc = 5.1 +Thos = 14.0 +Tsev = 14.0 + N = 500000 I0 = 214 / 0.2 E0 = 2000 R0 = 0 S0 = N-E0-I0 +H0 = 0 +Sev0 = 0 t0 = 24 # max_t = 5*days[len(swiss)-1] max_t = 200.0 -n_steps = 10000 +n_steps = 1000 dt = max_t / n_steps -y0 = np.array([S0, E0, I0, R0]) +y0 = np.array([S0, E0, I0, R0, H0, Sev0]) y_list = [y0] t_list = [t0] for i in range(0, n_steps): t = t_list[i] + dt - foo = functools.partial(seir, R_0=R_0, Tinf=Tinf, Tinc=Tinc) + # foo = functools.partial(seir, R_0=R_0, Tinf=Tinf, Tinc=Tinc) + foo = functools.partial(seihse, R_0=R_0, Tinf=Tinf, Tinc=Tinc, Thos=Thos, Tsev=Tsev) y1 = rk4(foo, y_list[i], t, dt) y_list.append(y1) t_list.append(t) - if (t > t0 and t <= t0 + 30) or (t >= t0 + 60 and t <= t0 + 90) : - R_0 = R_02 - else: - R_0 = R_01 if (y1[1] + y1[2] < 1): break @@ -66,11 +90,29 @@ y = np.array(y_list) p = np.transpose(y) -plt.semilogy(t, p[0], 'b') -plt.semilogy(t, p[1], 'r') -plt.semilogy(t, p[2], 'k') -plt.semilogy(t, p[3], 'g') -# plt.semilogy(days, swiss, 'k*') -plt.legend(['S', 'E', 'I', 'R', 'swiss']) + +fig = plt.figure() +# ax = plt.axes(xlim=(-0.1*np.max(t),np.max(t)*1.1), ylim=(-0.1*np.max(p),np.max(p)*1.1)) +ax = plt.axes(xlim=(1,np.max(t)*1.1), ylim=(1,np.max(p)*1.1)) + +lines = [plt.semilogy([], [], 'r-')[0], + plt.semilogy([], [], 'b-')[0], + plt.semilogy([], [], 'g-')[0], + plt.semilogy([], [], 'k-')[0], + plt.semilogy([], [], 'c-')[0], + plt.semilogy([], [], 'y-')[0] ] + +# anim = animation.FuncAnimation(fig, functools.partial(update, lines=lines, t=t, p=p), +# frames=n_steps, interval=10, blit=True) + +anim = animation.FuncAnimation(fig, functools.partial(update, lines=lines, t=t, p=p), + frames=n_steps, interval=10, blit=True) + +# plt.semilogy(t, p[0], 'b') +# plt.semilogy(t, p[1], 'r') +# plt.semilogy(t, p[2], 'k') +# plt.semilogy(t, p[3], 'g') +# # plt.semilogy(days, swiss, 'k*') +# plt.legend(['S', 'E', 'I', 'R', 'swiss']) plt.show()