Started adding Experiment 3 (SIR Simulation).

SIR_simulation.py

+import numpy as np
+from scipy.stats import multinomial
+import matplotlib.pyplot as plt
+
+def sim(N=1000, ts=np.linspace(0, 150, 150), S0=0, I0=10, R0=0, D0=0, beta=0.2, gamma=0.1, delta=0.01):
+  S, I, R, D = S0, I0, R0, D0
+  Ss = [S0]
+  Is = [I0]
+  Rs = [R0]
+  Ds = [D0]
+
+  for t in ts[1:]:
+    
+    dSdt = -np.random.binomial(S, beta*I/N)
+    S += dSdt
+    R, D, I = multinomial.rvs(I, [gamma, delta, 1 - (gamma + delta)])
+    I -= dSdt
+    Ss.append(S)
+    Is.append(I)
+    Rs.append(Rs[-1] + R)
+    Ds.append(Ds[-1] + D)
+
+  return np.array(Ss), np.array(Is), np.array(Rs), np.array(Ds)
+
+
+if __name__ == "__main__":
+  # Total population, N.
+  N = 1000
+  # Initial number of infected, recovered, and dead individuals, I0, R0, and D0.
+  I, R, D = 10, 0, 0
+  # Everyone else, S0, is susceptible to infection initially.
+  S = N - I - R - D
+  # Contact rate, beta, and mean recovery rate, gamma, (in 1/days).
+  beta, gamma, delta = 0.2, 0.1, 0.01
+  # A grid of time points (in days)
+  ts = np.linspace(0, 150, 150)
+  Ss, Is, Rs, Ds = sim(N, ts, S, I, R, D)
+  
+  # Plot the data on three separate curves for S(t), I(t) and R(t)
+  fig = plt.figure(facecolor='w')
+  ax = fig.add_subplot(111, facecolor='#dddddd', axisbelow=True)
+  ax.plot(ts, Ss/N, alpha=0.5, lw=2, label='Susceptible')
+  ax.plot(ts, Is/N, alpha=0.5, lw=2, label='Infected')
+  ax.plot(ts, Rs/N, alpha=0.5, lw=2, label='Recovered with immunity')
+  ax.plot(ts, Ds/N, alpha=0.5, lw=2, label='Dead')
+  ax.set_xlabel('Time /days')
+  ax.set_ylabel('Number (1000s)')
+  ax.set_ylim(0,1.2)
+  ax.yaxis.set_tick_params(length=0)
+  ax.xaxis.set_tick_params(length=0)
+  ax.grid(visible=True, which='major', c='w', lw=2, ls='-')
+  legend = ax.legend()
+  legend.get_frame().set_alpha(0.5)
+  for spine in ('top', 'right', 'bottom', 'left'):
+      ax.spines[spine].set_visible(False)
+  print(f'Total deaths: {Ds[-1]}')
+  plt.show()

experiment3.py

+from SIR_simulation import sim
+
+Ss, Is, Rs, Ds = sim()