In [1]:
from pygom import SquareLoss, common_models
import numpy
import scipy.integrate
import matplotlib.pyplot as plt
from scipy.optimize import minimize
In [2]:
# Generate some data for fitting
# Standard SIR model with 2 parameters
# Define the parameters
paramEval = [('beta',0.5), ('gamma',1.0/3.0)]
# initialize the model
ode = common_models.SIR(paramEval)
# Initial conditions (normalized to 1) and time scale
x0 = [1-1.27e-6, 1.27e-6, 0]
t = numpy.linspace(0, 70, 36)
# Find the solution.
solution = scipy.integrate.odeint(ode.ode, x0, t)
y = solution[:,1:3].copy()
Data over 70 days, with observations from every other day for infected and removed populations.
In [3]:
plt.plot(t,y[:,0], 'bo') # infected observations
plt.plot(t,y[:,1], 'go') # removed observations (recoverd/died)
plt.show()
In [4]:
# Add noise
numpy.random.seed(seed=6)
noise = numpy.random.normal(0,1.e-2,[len(t),2])
ynoise = y + noise
ynoise[ynoise<0] = 0
plt.plot(t,ynoise[:,0], 'bo') # infected observations
plt.plot(t,ynoise[:,1], 'go') # removed observations (recoverd/died)
plt.show()
We provide a guess for $\beta$ and $\gamma$.
This example assumes we have information about the infected and removed population, up to 50 days of the epidemic.
In [5]:
# Initial guess of parameters, and bounding constraints
theta = [0.2, 0.2]
boxBounds = [(0.0,5.0),(0.0,5.0)]
objSIR = SquareLoss(theta, ode, x0, t[0], t[1::], ynoise[1::,:], ['I','R'])
In [ ]:
numpy.shape(y[1::,:])
In [14]:
objSIR.residual()
Out[14]:
In [ ]:
In [ ]:
In [ ]:
In [11]:
# perform optimization
res = minimize(fun=objSIR.cost,
jac=objSIR.sensitivity,
x0=theta,
bounds=boxBounds,
method='SLSQP')
print(res)
Our estimated parameters, with the provided data, are:
In [7]:
res.x
Out[7]:
We can now simulate the epidemic with the fitted parameters.
In [8]:
# model with fitted parameters
ode_fit = common_models.SIR({'beta':res.x[0], 'gamma':res.x[1]})
x0_fit = [1-1.27e-6, 1.27e-6, 0]
t_fit = numpy.linspace(0, 150, 1000)
ode_fit.initial_values = (x0_fit, t_fit[0])
sol_fit = ode_fit.integrate(t_fit[1::])
ode_fit.plot()
peak_i = numpy.argmax(sol_fit[:,1])
print('Peak infection (days)', t_fit[peak_i])
print('R0 (beta/gamma) = ', res.x[0]/res.x[1])
In [10]:
plt.plot(t,ynoise[:,0], 'bo') # infected observations
plt.plot(t,ynoise[:,1], 'go') # removed observations (recoverd/died)
plt.plot(t_fit, sol_fit)
plt.show(())
In [ ]: