added legend

covid/python/covid.py

@@ -6,6 +6,22 @@ swiss = np.array([18, 27, 42, 56, 90, 114, 214, 268, 337, 374, 491, 652, 858, 11
 
 days = np.array(range(1,len(swiss)+1))
 
+def F(yx, yy, yz, t):
+    return sigma * (yy - yz), rho * yx, (mu * yx + yz)
+
+def rk2(yx, yy, yz, t, dt):
+    yx_tmp, yy_tmp, yz_tmp = F(yx, yy, yz, t)
+    yx0 = yx + dt / 2 * yx_tmp
+    yy0 = yy + dt / 2 * yy_tmp
+    yz0 = yz + dt / 2 * yz_tmp
+
+    y1x, y1y, y1z = F(yx0, yy0, yz0, t+dt/2)
+
+    return yx + dt * y1x, yy + dt * y1y, yz + dt * y1z
+
+def yx(dt, yx0, yy0, yz0, sig):
+    return yx0 + dt * sigma * (yy0-yz0)
+
 def s(dt, beta, S0, I0, N):
     return S0 - dt * (beta * S0 * I0 / N)
 

covid/python/seir.py

@@ -14,7 +14,9 @@ def update(i, lines, t, p):
         lines[j].set_data(t[0:i], p[j][0:i])
     return lines
 
-def policy_r0(R_i, t):
+# def policy_r0(R_0, R_target, t, delta_t):
+#     if (t <= )
+
 
 
 def seir(y, t, R_0, Tinf, Tinc):
@@ -104,6 +106,7 @@ lines = [plt.semilogy([], [], 'r-')[0],
 
 # anim = animation.FuncAnimation(fig, functools.partial(update, lines=lines, t=t, p=p), 
 #            frames=n_steps, interval=10, blit=True)
+ax.legend(['Sain', 'Exposé', 'Infectieux', 'Rétablis', 'Hospitalisés', 'Soins intensifs'])
 
 anim = animation.FuncAnimation(fig, functools.partial(update, lines=lines, t=t, p=p), 
            frames=n_steps, interval=10, blit=True)