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