test rk4 seir

covid/python/seir.py

+import numpy as np
+import matplotlib.pyplot as plt
+
+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))
+
+R_0 = 4.26
+R_01 = 4.26
+R_02 = 1 - 1/R_01
+
+Tinf = 7.0
+Tinc = 5.1
+N = 8000000.0
+
+def seir(y, t):
+    y1 = np.zeros(4)
+    y1[0] = - R_0 / Tinf * y[2] * y[0] / N
+    y1[1] =  R_0 / Tinf * y[2] * y[0] / N - 1.0 / Tinc * y[1]
+    y1[2] = 1.0 / Tinc * y[1] - 1.0 / Tinf * y[2]
+    y1[3] = 1.0 / Tinf * y[2]
+
+    return y1
+    
+def rk4(F, y, t, dt):
+    k1 = dt * F(y, t)
+    k2 = dt * F(y + k1 / 2, t + dt/2)
+    k3 = dt * F(y + k2 / 2, t + dt/2)
+    k4 = dt * F(y + k3, t + dt)
+    return y + (k1 + 2 * (k2 + k3) + k4) / 6
+
+I0 = 0
+E0 = swiss[0]
+R0 = 0
+S0 = N-E0-I0
+t0 = -Tinc
+
+# max_t = 5*days[len(swiss)-1]
+max_t = 2000
+n_steps = 1000000
+dt = max_t / n_steps
+
+y0 = np.array([S0, E0, I0, R0])
+
+y_list = [y0]
+t_list = [t0]
+for i in range(0, n_steps):
+    t = t_list[i] + dt
+    y1 = rk4(seir, y_list[i], t, dt)
+    y_list.append(y1)
+    t_list.append(t)
+    if (t >= 27 and t <= 67) or (t >= 98 and t <= 122) :
+        R_0 = R_02
+    else:
+        R_0 = R_01
+
+    if (y1[1] + y1[2] < 1):
+        break
+
+t = np.array(t_list) 
+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'])
+plt.show()
+