from mitepid.utils import load_mat
from mitepid.epid_sim import epid_sim
#%% general variables
country = 'Germany'
t_end = 541
x00 = 1e-5
x0_vec=[x00, x00, x00, x00, x00, x00, x00, x00, x00,]
dir_source_mat = Path(Path(__file__).resolve().parent)
file_data_opt_SIR = Path(dir_source_mat, 'Optimised_B', 'SIR_B_opt_normalised.mat')
dir_save_plots_main = Path(Path(__file__).resolve().parent, 'sample_output')
file_data_opt_SIS = Path(dir_source_mat, 'Optimised_B', 'SIS_B_opt_normalised.mat')
file_data_opt_SIR = Path(dir_source_mat, 'Optimised_B', 'SIR_B_opt_normalised.mat')
B_opt_SIS = get_B_policy(file_data_opt_SIS, country, 'Uncontained', )
B_opt_SIR = get_B_policy(file_data_opt_SIR, country, 'Uncontained', )
# age groups, used as legends in plots
age_groups = ['[0-10]', '[10-20]', '[20-30]', '[30-40]',
'[40-50]', '[50-60]', '[60-70]', '[70-80]', '80+',]
Ng = len(age_groups)
alpha = load_mat(file_data_opt_SIR, 'alpha')
alpha_vec = np.ones((Ng, 1)) * alpha
Gamma = np.zeros((Ng,Ng))
np.fill_diagonal(Gamma, alpha_vec)
#
Sigma = Gamma *10
model_type = 'SIR'
x00,
x00,
x00,
x00,
x00,
x00,
]
model_type = 'SEIR'
file_Bopt = Path('B_opt_SEIR_Ti_5_R0_2.95.mat')
file_B_opt = Path(dir_B_opt, file_Bopt)
Sigma = load_mat(file_B_opt, 'Sigma')
Gamma = load_mat(file_B_opt, 'Gamma')
B_opt_country = get_B_policy(
file_B_opt,
country,
'Uncontained',
)
Ng = B_opt_country.shape[0]
assert len(
age_groups
) == 9, 'Check the defined age-groups. Something is wrong there.'
ref_obj = epid_sim(
model_type=model_type,
B=B_opt_country,
Gamma=Gamma,
Sigma=Sigma,
dir_save_plots_main=dir_save_plots,
country='Iran',
policy_list=['Uncontained'],
policy_switch_times=[0],
def calc_sol(self, ):
"""
epid_calss method.
calculates the solution for the model given the parameters, switchin times, etc.
results saved in self.sol_dict and self.sol_agg_dict
Parameters
----------
None.
Returns
-------
self.sol_dict: dict, keys are 'I', 'R', 'E', etc, values are numpy array (Nt x Ng)
a dictionary inlcuding solutions for each group and each compartment
self.sol_agg_dict: dict, keys are 'I', 'R', 'E', etc, values are numpy array (Nt x 1)
a dictionary inlcuding aggregate solution for each compartment
"""
from pathlib import Path
import numpy as np
from numpy.linalg import inv, eigvals
from scipy.integrate import odeint
from mitepid.utils import sort_out_t_policy_x0
from mitepid.models import SIS, SIR, SEIR
from mitepid.utils import sol_aggregate
from mitepid.policies import get_B_policy, get_pop_distr
list_t1, list_t2, list_policies, list_x0, list_t_switch, list_all_policies = \
sort_out_t_policy_x0(self.policy_definition,
self.xternal_inputs,
self.t_end,
self.Ng)
print('*************************************************************')
print('time (days) ---> policy (R0 of policy)')
print('---------------------------------------------')
x0_step = self.x0
for idx in np.arange(len(list_t1)):
t_switch1 = list_t1[idx]
t_switch2 = list_t2[idx]
policy = list_policies[idx]
x0_xtrnal = list_x0[idx]
B_step = get_B_policy(
B=self.B,
country=self.country,
policy=policy,
)
rho = np.max(np.abs(eigvals(np.matmul(-inv(self.D),
B_step)))) # spectral radius
self.list_policy_info.append((t_switch1, policy, rho))
#%% R_0 policy
print('%10.1f ---> %s (R0 = %2.2f)' % (t_switch1, policy, rho))
t_step = np.arange(t_switch1, t_switch2, step=self.t_step)
x0_xtrnal = self.correct_x0(x0_xtrnal)
x0_step = np.array(x0_step) + x0_xtrnal
# solve the ODE
if self.model_type == 'SIS':
sol_step = odeint(SIS,
x0_step,
t_step,
args=(B_step, self.Gamma, self.Mu))
elif self.model_type == 'SIR':
sol_step = odeint(SIR,
x0_step,
t_step,
args=(B_step, self.Gamma, self.Mu))
elif self.model_type == 'SEIR':
sol_step = odeint(SEIR,
x0_step,
t_step,
args=(B_step, self.Gamma, self.Mu,
self.Sigma))
# update x0
x0_step = sol_step[-1]
if idx == 0:
sol_all = sol_step
else:
sol_all = np.concatenate((sol_all, sol_step))
country = self.country
sol_dict = {}
sol_agg_dict = {}
Ng = self.Ng
if self.model_type == 'SIS':
sol_dict['I'] = sol_all[:, :Ng]
sol_agg_dict['I'] = sol_aggregate(sol_dict['I'],
get_pop_distr(country))
elif self.model_type == 'SIR':
sol_dict['I'] = sol_all[:, :Ng]
sol_agg_dict['I'] = sol_aggregate(sol_dict['I'],
get_pop_distr(country))
sol_dict['R'] = sol_all[:, Ng:2 * Ng]
sol_agg_dict['R'] = sol_aggregate(sol_dict['R'],
get_pop_distr(country))
elif self.model_type == 'SEIR':
sol_dict['I'] = sol_all[:, :Ng]
sol_agg_dict['I'] = sol_aggregate(sol_dict['I'],
get_pop_distr(country))
sol_dict['R'] = sol_all[:, Ng:2 * Ng]
sol_agg_dict['R'] = sol_aggregate(sol_dict['R'],
get_pop_distr(country))
sol_dict['E'] = sol_all[:, 2 * Ng:3 * Ng]
sol_agg_dict['E'] = sol_aggregate(sol_dict['E'],
get_pop_distr(country))
self.sol_dict = sol_dict
self.sol_agg_dict = sol_agg_dict