Please see the code below. I suspect that I make a lot of mistakes in the covid_at_t()… If you see some bad coding practice in Stan, please let me know! Thanks!
functions {
vector switch_h_ex(int t,
int nag,
int ncg,
real[,] p_h_ex_o
) {
vector[nag * ncg] out;
if ((t >= 40) && (t < 101)){
out = to_vector(p_h_ex_o[2, :]);
} else if ((t >= 101) && (t < 163)){
out = to_vector(p_h_ex_o[3, :]);
} else if (t >= 163){
out = to_vector(p_h_ex_o[4, :]);
} else {
out = to_vector(p_h_ex_o[1, :]);
}
return out;
}
vector switch_icu_ex(int t,
real[,] p_icu_ex_o
) {
if ((t >= 40) && (t < 101)){
return(to_vector(p_icu_ex_o[2, :]));
} else if ((t >= 101) && (t < 163)){
return(to_vector(p_icu_ex_o[3, :]));
} else if (t >= 163){
return(to_vector(p_icu_ex_o[4, :]));
} else {
return(to_vector(p_icu_ex_o[1, :]));
}
}
vector calc_fraction_tested(real[,] y,
real p_I_test,
real p_nonI_test,
real n_test_pd,
real n_test_pd_hosp,
real n_test_pd_icu
) {
real n_s_pop;
real n_exp_pop;
real n_inf_pop;
real n_asym_pop;
real n_rec_pop;
real n_icu_pop;
real n_hosp_pop;
real n_nonI_pop;
real n_test_per_day;
real n_test_demand;
real p_I_tested;
real p_nonI_tested;
real n_pd[2];
vector[2] out;
n_s_pop = sum(y[1, :]);
n_exp_pop = sum(y[2, :]) + sum(y[3, :]);
n_inf_pop = sum(y[7, :]) + sum(y[8, :]) + sum(y[9, :]);
n_asym_pop = sum(y[4, :]) + sum(y[5, :]) + sum(y[6, :]);
n_rec_pop = sum(y[21, :]);
n_icu_pop = sum(y[20, :]);
n_hosp_pop = sum(y[19, :]);
n_nonI_pop = n_s_pop + n_exp_pop + n_asym_pop + n_rec_pop;
n_pd[1] = 0;
n_pd[2] = n_test_pd - n_test_pd_hosp * n_hosp_pop - n_test_pd_icu * n_icu_pop;
n_test_per_day = max(n_pd);
n_test_demand = p_I_test * n_inf_pop + p_nonI_test * n_nonI_pop;
if (n_test_demand > n_test_pd) {
p_I_tested = p_I_test * (n_test_per_day / n_test_demand);
p_nonI_tested = p_nonI_test * (n_test_per_day / n_test_demand);
} else {
p_I_tested = p_I_test;
p_nonI_tested = p_nonI_test;
}
out[1] = p_nonI_tested;
out[2] = p_I_tested;
return(out);
}
matrix[,] switch_mixing_mat(int t,
real[,] change_time,
real[,] mx_mat_o,
real[,,] mx_array
) {
if (t >= change_time[1, 1] && t < change_time[2, 1]) {
return(mx_array[:, :, 1]);
} else if (t >= change_time[1, 2] && t < change_time[2, 2]) {
return(mx_array[:, :, 2]);
} else if (t >= change_time[1, 3] && t < change_time[2, 3]) {
return(mx_array[:, :, 3]);
} else {
return(mx_mat_o);
}
}
real[] calculate_lambda(real quar_cr,
int ncg, // number of comorbidity groups
int nag, // number of age groups
real[] N_by_age,
matrix[,] mixing_matrix,
real[,] y,
int[,] co_ix
) {
real inf_by_age[nag * ncg];
real q_inf_by_age[nag * ncg];
real d_by_age[nag * ncg];
real q_noninf_by_age[nag * ncg];
real n_inf_by_age[nag * ncg];
real n_alive_by_age[nag * ncg];
real n_inf_contacts[nag * ncg];
int ix1[nag];
int ix2[nag];
real tmp1[nag];
real tmp2[nag];
real tmp3[nag * ncg];
ix1 = co_ix[1, :];
ix2 = co_ix[2, :];
for (k in 1:nag) {
tmp1[k] = y[4, ix1[k]] + y[5, ix1[k]] + y[6, ix1[k]] + y[7, ix1[k]] + y[8, ix1[k]] + y[9, ix1[k]];
tmp2[k] = y[4, ix2[k]] + y[5, ix2[k]] + y[6, ix2[k]] + y[7, ix2[k]] + y[8, ix2[k]] + y[9, ix2[k]];
tmp1[k] = tmp1[k] + tmp2[k];
inf_by_age[ix1[k]] = tmp1[k];
inf_by_age[ix2[k]] = tmp1[k];
tmp1[k] = y[13, ix1[k]] + y[14, ix1[k]] + y[15, ix1[k]] + y[16, ix1[k]] + y[17, ix1[k]] + y[18, ix1[k]];
tmp2[k] = y[13, ix2[k]] + y[14, ix2[k]] + y[15, ix2[k]] + y[16, ix2[k]] + y[17, ix2[k]] + y[18, ix2[k]];
tmp1[k] = tmp1[k] + tmp2[k];
q_inf_by_age[ix1[k]] = tmp1[k];
q_inf_by_age[ix2[k]] = tmp1[k];
tmp1[k] = y[23, ix1[k]];
tmp2[k] = y[23, ix2[k]];
tmp1[k] = tmp1[k] + tmp2[k];
d_by_age[ix1[k]] = tmp1[k];
d_by_age[ix2[k]] = tmp1[k];
tmp1[k] = y[10, ix1[k]] + y[11, ix1[k]] + y[12, ix1[k]];
tmp2[k] = y[10, ix2[k]] + y[11, ix2[k]] + y[12, ix2[k]];
tmp1[k] = tmp1[k] + tmp2[k];
q_noninf_by_age[ix1[k]] = tmp1[k];
q_noninf_by_age[ix2[k]] = tmp1[k];
}
for (k in 1:(nag*ncg)) {
n_inf_by_age[k] = inf_by_age[k] + (1 - quar_cr) * q_inf_by_age[k];
n_alive_by_age[k] = N_by_age[k] - d_by_age[k] - quar_cr * q_noninf_by_age[k];
n_inf_by_age[k] = n_inf_by_age[k] / n_alive_by_age[k];
}
for (k in 1:(nag * ncg)) {
for (l in 1:(nag * ncg)) {
tmp3[l] = mixing_matrix[k, l] * n_inf_by_age[l];
}
n_inf_contacts[k] = sum(tmp3);
}
return n_inf_contacts;
}
real[] covid_at_t(int t,
real[,] y,
real[] x_s, // vector of scalars
real[,] x_v, // matrix of vectors
real[,] mx_mat_o, // an 18x18 matrix; the original mixing matrix
real[,,] mx_array, // an 18x18x3 array; the 3rd dimention is corresponding to the NPI changes below
real[,] change_time, // a 2x3 matrix, four NPI changes, row1 is start times; row2 is end times: 1. social distancing; 2. SIP; 3 behavior change
int nag,
int ncg,
int ns
) {
// scalars
real beta;
real spec;
real sens;
real p_I_test;
real p_nonI_test;
real p_quar_ex;
real p_trans_exp;
real p_trans_inf;
real quar_cr;
real p_icu_ex_nb;
real n_icu_beds;
real n_test_pd; // number of tests per day
real n_test_pd_hosp; // number of tests per day in hosplital
real n_test_pd_icu; // number of test per day in ICU
// vectors or matrices
vector[nag * ncg] N_by_age;
int co_ix[ncg, nag]; // row1 is c0, row2 is c1; the column represents the age groups
vector[nag * ncg] prop_asym;
vector[nag * ncg] prop_hosp;
vector[nag * ncg] prop_inf_d;
vector[nag * ncg] prop_icu;
vector[nag * ncg] p_icu_d_nb;
vector[nag * ncg] prop_hosp_d; // no medical advance built in
vector[nag * ncg] prop_icu_d; // no medical advance built in
matrix[4, nag * ncg] p_h_ex_o; // 3 cutoff dates
matrix[4, nag * ncg] p_icu_ex_o; // 3 cutoff dates
// created variables to store results
matrix[ns, nag * ncg] dy; // ns: number of states; nag: number of age groups; ncg: number of comorbidity groups
vector[nag * ncg] lambda;
vector[nag * ncg] adj_beta[nag * ncg]; // adjusted beta by lambda
matrix[nag * ncg, nag * ncg] mixing_matrix;
real fr_vec[2];
real fr_nonI_test;
real fr_I_test;
real pr_f_pos; // tested false positive
real pr_t_pos_a; // tested true positive among asymptomatic
real pr_t_pos_s; // tested true positive among symptomatic
real n_icu_t;
real p_icu_ovfl;
real p_h_ex[nag * ncg];
real p_icu_ex[nag * ncg];
p_h_ex = switch_h_ex(t, p_h_ex_o);
p_icu_ex = switch_icu_ex(t, p_icu_ex_o);
fr_vec = calc_fraction_tested(y, p_I_test, p_nonI_test, n_test_pd, n_test_pd_hosp, n_test_pd_icu);
fr_nonI_test = fr_vec[1];
fr_I_test = fr_vec[2];
n_icu_t = sum(y[20, :]);
if (n_icu_t > n_icu_beds) {
p_icu_ovfl = (n_icu_t - n_icu_beds) / n_icu_t;
} else {
p_icu_ovfl = 0;
}
// get a new mixing matrix
mixing_matrix = switch_mixing_mat(t, change_time, mx_mat_o, mx_array);
// calculate lambda
lambda = calculate_lambda(quar_cr, ncg, nag, N_by_age, mixing_matrix, y, co_ix);
adj_beta = beta * lambda;
pr_f_pos = fr_nonI_test * (1 - spec);
pr_t_pos_a = fr_nonI_test * sens;
pr_t_pos_s = fr_I_test * sens;
// separate comorbidity group first (need to delete!!)
// 1: changes in S
dy[1, :] = -(adj_beta * (1 - pr_f_pos) + pr_f_pos) * y[1, :] + p_quar_ex * y[10, :];
// 2: changes in E1
dy[2, :] = adj_beta * (1 - pr_f_pos) * y[1, :] - ((1 - pr_f_pos) * p_trans_exp + pr_f_pos) * y[2, :];
// 3: changes in E2
dy[3, :] = p_trans_exp * (1 - pr_f_pos) * (y[2, :] - y[3, :]) - pr_f_pos * y[3, :];
// 4: changes in AI1
dy[4, :] = prop_asym * (1 - pr_f_pos) * p_trans_exp * y[3, :] - (p_trans_inf * (1 - pr_t_pos_a) + pr_t_pos_a) * y[4, :];
// 5: changes in AI2
dy[5, :] = p_trans_inf * (1 - pr_t_pos_a) * (y[4, :] - y[5, :]) - pr_t_pos_a * y[5, :];
// 6: changes in AI3
dy[6, :] = p_trans_inf * (1 - pr_t_pos_a) * (y[5, :] - y[6, :]) - pr_t_pos_a * y[6, :];
// 7: changes in I1
dy[7, :] = (1 - prop_asym) * (1 - pr_f_pos) * p_trans_exp * y[3, :] - (p_trans_inf * (1 - pr_t_pos_s) + pr_t_pos_s) * y[7, :];
// 8: changes in I2
dy[8, :] = p_trans_inf * (1 - pr_t_pos_s) * (y[7, :] - y[8, :]) - pr_t_pos_s * y[8, :];
// 9: changes in I3
dy[9, :] = p_trans_inf * (1 - pr_t_pos_s) * (y[8, :] - y[9, :]) - pr_t_pos_s * y[9, :];
// 10: changes in QS
dy[10, :] = pr_f_pos * y[1, :] - (p_quar_ex + adj_beta * (1 - quar_cr)) * y[10, :];
// 11: changes in QE1
dy[11, :] = adj_beta * (1 - quar_cr) * y[10, :] + pr_f_pos * y[2, :] - p_trans_exp * y[11, :];
// 12: changes in QE2
dy[12, :] = p_trans_exp * (y[11, :] - y[12, :]) + p_f_pos * y[12, :];
// 13: changes in QAI1
dy[13, :] = prop_asym * p_trans_exp * y[12, :] - p_trans_inf * y[13, :] + p_t_pos_a * y[4, :];
// 14: changes in QAI2
dy[14, :] = p_trans_inf * (y[13, :] - y[14, :]) + p_t_pos_a * y[5, :];
// 15: changes in QAI3
dy[15, :] = p_trans_inf * (y[14, :] - y[15, :]) + p_t_pos_a * y[6, :];
// 16: changes in QI1
dy[16, :] = (1 - prop_asym) * p_trans_exp * y[12, :] - p_trans_inf * y[16, :] + pr_t_pos_s * y[7, :];
// 17: changes in QI2
dy[17, :] = p_trans_inf * (y[16, :] - y[17, :]) + pr_t_pos_s * y[8, :];
// 18: changes in QI3
dy[18, :] = p_trans_inf * (y[17, :] - y[18, :]) + pr_t_pos_s * y[9, :];
// 19: changes in H
dy[19, :] = p_trans_inf * prop_hosp * (1 - prop_ICU) * (y[9, :] + y[18, :]) - p_h_ex * y[19, :];
// 20: changes in ICU
dy[20, :] = p_trans_inf * prop_hosp * prop_ICU * (y[9, :] + y[18, :]) - p_icu_ex_ov * y[20, :];
// 21: changes in R
dy[21, :] = p_trans_inf * ((1 - prop_inf_d) * (1 - prop_hosp) * y[9, :] + y[6, :]);
// 22: changes in RD
dy[22, :] = p_icu2rec_ov * y[20, :] + (1 - prop_hosp_d) * p_h_ex * y[19, :] + p_trans_inf * ((1 - prop_inf_d) * (1 - prop_hosp) * y[18, :] + y[15, :]);
// 23: changes in D
dy[23, :] = (p_icu_ex * (1 - p_icu_ovfl) * prop_icu_d + p_icu_ex_nb * p_icu_ovfl * p_icu_d_nb) * y[20, :] + prop_inf_d * p_trans_inf * (1 - prop_hosp) * (y[9, :] + y[18, :]) + prop_hosp_d * p_h_ex * y[19, :];
// 24: changes in HD
dy[24, :] = prop_inf_d * p_trans_inf * (1 - prop_hosp) * (y[9, :] + y[18, :]);
// newly hospitalized at time t
new_hosp = p_trans_inf * prop_hosp * (1 - prop_ICU) * (y[9, :] + y[18, :]);
}
}
