added swiss data for kind of fit

covid/python/covid.py

@@ -24,31 +24,35 @@ def timestep(S0, I0, R0, dt, beta, lamb, N):
     R1 = r(dt, lamb, R0, I0, N)
     return S1, I1, R1
 
-S0 = 400000
+S0 = 8000000
 I0 = swiss[0]
 R0 = 0
-max_t = 30
-n_steps = 1000
+max_t = 50*days[len(swiss)-1]
+n_steps = 10000
 dt = max_t / n_steps
 
 
 N = compute_n(S0, I0, R0) 
 
 lamb = 1.0 / 14.0
-beta = 0.33
+beta_1 = 0.34
+beta_2 = beta_1 / 100
+beta = beta_1
 
 s_list = [S0]
 r_list = [R0]
 i_list = [I0]
-t_list = [0]
+t_list = [1]
 for i in range(0, n_steps):
     S1, I1, R1 = timestep(s_list[i], i_list[i], r_list[i], dt, beta, lamb, N)
     s_list.append(S1)
     i_list.append(I1)
     r_list.append(R1)
-    t_list.append((i+1)*dt)
-    # if ((i+1)*dt >= 10 and (i+1)*dt <= 365):
-    #     beta = 0.1
+    t_list.append(t_list[i]+dt)
+    if (t_list[i+1] >= 27 and t_list[i+1] <= 127) :
+        beta = beta_2
+    else:
+        beta = beta_1
     # else:
     #     beta = 1
 
@@ -57,10 +61,12 @@ r = np.array(r_list)
 ii = np.array(i_list)
 t = np.array(t_list)
 
-# plt.plot(t, s, 'b')
-# plt.plot(t, r, 'r')
+print(t)
+
+plt.semilogy(t, s, 'b')
+plt.semilogy(t, r, 'r')
 plt.semilogy(t, ii, 'k')
-plt.loglog(days, swiss, 'k*')
-plt.legend(['S', 'I', 'R', swiss])
+plt.semilogy(days, swiss, 'k*')
+plt.legend(['S', 'I', 'R', 'swiss'])
 plt.show()