def solve(self, t, init_cond, t_params, params):
"""4th order Runge-Kutta solver.
"""
t_solve = np.arange(np.min(t), np.max(t) + self.dt, self.dt / 2)
y_solve = np.zeros((init_cond.size, t_solve.size),
dtype=init_cond.dtype)
y_solve[:, 0] = init_cond
# linear interpolate the parameters
params = utils.linear_interpolate(t_solve, t_params, params)
for i in range(2, t_solve.size, 2):
k1 = self.system(t_solve[i - 2], y_solve[:, i - 2], params[:,
i - 2])
k2 = self.system(t_solve[i - 1],
y_solve[:, i - 2] + self.dt / 2 * k1,
params[:, i - 1])
k3 = self.system(t_solve[i - 1],
y_solve[:, i - 2] + self.dt / 2 * k2,
params[:, i - 1])
k4 = self.system(t_solve[i], y_solve[:, i - 2] + self.dt * k3,
params[:, i])
y_solve[:,
i] = y_solve[:, i -
2] + self.dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
# linear interpolate the solutions.
y_solve = utils.linear_interpolate(t, t_solve, y_solve)
return y_solve
def predict(self, t):
params = linear_interpolate(t, self.t_params, self.params)
components = {
c: linear_interpolate(t, self.t_params, self.components[c])
for c in self.components
}
components.update(
{'newE': linear_interpolate(t, self.t_params, self.rhs_newE)})
components.update(
{'newE_obs': linear_interpolate(t, self.t, self.obs)})
return params, components
def solve(self, t, init_cond, t_params, params):
"""Forward Euler solver.
"""
t_solve = np.arange(np.min(t), np.max(t) + self.dt, self.dt)
y_solve = np.zeros((init_cond.size, t_solve.size),
dtype=init_cond.dtype)
y_solve[:, 0] = init_cond
# linear interpolate the parameters
params = utils.linear_interpolate(t_solve, t_params, params)
for i in range(1, t_solve.size):
y_solve[:, i] = y_solve[:, i - 1] + self.dt * self.system(
t_solve[i - 1], y_solve[:, i - 1], params[:, i - 1])
# linear interpolate the solutions.
y_solve = utils.linear_interpolate(t, t_solve, y_solve)
return y_solve
def process(self):
"""Process the data.
"""
self.rhs_newE = linear_interpolate(self.t_params, self.t, self.obs)
# fit the E
self.step_ode_sys.update_given_params(c=self.sigma)
E = self.step_ode_sys.simulate(self.t_params,
np.array([self.init_cond['E']]),
self.t_params,
self.rhs_newE[None, :])[0]
# fit I1
self.step_ode_sys.update_given_params(c=self.gamma1)
# modify initial condition of I1
self.init_cond.update(
{'I1': (self.rhs_newE[0] / 5.0)**(1.0 / self.alpha)})
I1 = self.step_ode_sys.simulate(self.t_params,
np.array([self.init_cond['I1']]),
self.t_params,
self.sigma * E[None, :])[0]
# fit I2
self.step_ode_sys.update_given_params(c=self.gamma2)
I2 = self.step_ode_sys.simulate(self.t_params,
np.array([self.init_cond['I2']]),
self.t_params,
self.gamma1 * I1[None, :])[0]
# fit S
self.init_cond.update(
{'S': self.N - self.init_cond['E'] - self.init_cond['I1']})
self.step_ode_sys.update_given_params(c=0.0)
S = self.step_ode_sys.simulate(self.t_params,
np.array([self.init_cond['S']]),
self.t_params,
-self.rhs_newE[None, :])[0]
neg_S_idx = S < 0.0
# fit R
self.step_ode_sys.update_given_params(c=0.0)
R = self.step_ode_sys.simulate(self.t_params,
np.array([self.init_cond['R']]),
self.t_params,
self.gamma2 * I2[None, :])[0]
if np.any(neg_S_idx):
id_min = np.min(np.arange(S.size)[neg_S_idx])
S[id_min:] = S[id_min - 1]
E[id_min:] = E[id_min - 1]
I1[id_min:] = I1[id_min - 1]
I2[id_min:] = I2[id_min - 1]
R[id_min:] = R[id_min - 1]
self.components = {'S': S, 'E': E, 'I1': I1, 'I2': I2, 'R': R}
# get beta
self.params = (self.rhs_newE / ((S / self.N) *
((I1 + I2)**self.alpha)))[None, :]
def test_linear_interpolate(t, t_org, x_org, result):
my_result = utils.linear_interpolate(t, t_org, x_org)
assert np.allclose(result, my_result.ravel())
assert my_result.ndim == x_org.ndim