Skip to content

Instantly share code, notes, and snippets.

@a-voel
Created April 12, 2022 07:44
Show Gist options
  • Save a-voel/3dfb7acd5e65606bd51a1c148650ae36 to your computer and use it in GitHub Desktop.
Save a-voel/3dfb7acd5e65606bd51a1c148650ae36 to your computer and use it in GitHub Desktop.
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('ggplot')
Sbar = 0.15
Ibar = 0.001
beta = 2
gamma = 1/2
exponent= beta*Sbar-gamma
print("exponent:", exponent)
# numerics params
N = 1000
T = 15
dt = T/N
def rk4(f, s, dt):
k1 = f(s)
k2 = f(s+dt*k1/2)
k3 = f(s+dt*k2/2)
k4 = f(s+dt*k3)
return s+dt*1/6*(k1+2*k2+2*k3+k4)
def f(x):
I, S = x
return np.array([
beta*S*I-gamma*I,
-beta*S*I
])
# numerics
ts = []
nss = []
nis = []
x = np.array([Ibar, Sbar])
for i in range(N):
t = dt*i
x = rk4(f, x, dt)
nis.append(x[0])
nss.append(x[1])
ts.append(t)
ts = np.array(ts)
nss = np.array(nss)
nis = np.array(nis)
#perturbation theory (eps=1)
pss = Sbar-beta*Sbar*Ibar/(beta*Sbar-gamma)*(np.exp((beta*Sbar-gamma)*ts)-1)
pis = Ibar*np.exp((beta*Sbar-gamma)*ts)
#plotting
fig, (ax1, ax2) = plt.subplots(2, 1)
ax1.plot(ts, nis, label="Numerics")
ax1.plot(ts, pis, "--", label="Perturbation theory")
ax1.set_title("Infections $I(t)$")
ax1.autoscale(enable=True, axis='x', tight=True)
ax2.plot(ts, nss)
ax2.plot(ts, pss, "--")
ax2.set_title("Susceptible $S(t)$")
ax2.autoscale(enable=True, axis='x', tight=True)
plt.figlegend()
plt.tight_layout()
plt.savefig("peturbation.svg")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment