perturbation.py

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")