functions {
vector tb_ode(real t, vector y, vector theta, array real x_r_time_varying, array int x_i) {
// Extract parameters from theta
real alpha = theta[1];
real alpha1 = theta[2];
real theta_param = theta[3];
real theta1 = theta[4];
real q = theta[5];
real h = theta[6];
real psi = theta[7];
real psi1 = theta[8];
// Extract time-varying parameters from x_r_time_varying
real mu = x_r_time_varying[1];
real mu1 = x_r_time_varying[2];
real sigma = x_r_time_varying[3];
real ch = x_r_time_varying[4];
real ad = x_r_time_varying[5];
real w_fs = x_r_time_varying[6];
real omega = x_r_time_varying[7];
real chi = x_r_time_varying[8];
real p = x_r_time_varying[9];
real n = x_r_time_varying[10];
real z = x_r_time_varying[11];
real zr = x_r_time_varying[12];
real f_ar = x_r_time_varying[13];
real s_ar = x_r_time_varying[14];
real f_c = x_r_time_varying[15];
real s_c = x_r_time_varying[16];
real mu_c = x_r_time_varying[17];
real mu_a = x_r_time_varying[18];
real mu_r = x_r_time_varying[19];
real tau = x_r_time_varying[20];
real tau1 = x_r_time_varying[21];
real c_Ac = x_r_time_varying[22];
real c_M = x_r_time_varying[23];
real c_Ar = x_r_time_varying[24];
real c_Tr = x_r_time_varying[25];
real c_As = x_r_time_varying[26];
real pseek = x_r_time_varying[27];
real ptest = x_r_time_varying[28];
real psensi = x_r_time_varying[29];
real ptreat = x_r_time_varying[30];
real pseek_s = x_r_time_varying[31];
real ptest_s = x_r_time_varying[32];
real psensi_s = x_r_time_varying[33];
real ptreat_s = x_r_time_varying[34];
real pseek_c = x_r_time_varying[35];
real ptest_c = x_r_time_varying[36];
real psensi_c = x_r_time_varying[37];
real ptreat_c = x_r_time_varying[38];
real pseek_r = x_r_time_varying[39];
real ptest_r = x_r_time_varying[40];
real psensi_r = x_r_time_varying[41];
real ptreat_r = x_r_time_varying[42];
real PR = x_r_time_varying[43];
real PR_a = x_r_time_varying[44];
real PR_r = x_r_time_varying[45];
real PR_rr = x_r_time_varying[46];
real f_ac = x_r_time_varying[47];
real s_ac = x_r_time_varying[48];
real f_as = x_r_time_varying[49];
real s_as = x_r_time_varying[50];
real p_ch = x_r_time_varying[51];
real p_ad = x_r_time_varying[52];
real iota = x_r_time_varying[53];
real mu_sev = x_r_time_varying[54];
real rho = x_r_time_varying[55];
// Extract state variables
real S = y[1]; real Lf = y[2]; real Ls = y[3]; real Ac = y[4]; real M = y[5];
real Rc = y[6]; real R = y[7]; real S1 = y[8]; real Lf1 = y[9]; real Ls1 = y[10];
real As1 = y[11]; real Ras = y[12]; real Ms = y[13]; real Ac1 = y[14]; real Rac = y[15];
real Mc = y[16]; real R1 = y[17]; real Lr = y[18]; real Ar = y[19]; real Tr = y[20];
real Rr = y[21]; real R2 = y[22];
// Compute populations
real P_children = S + Lf + Ls + Ac + M + Rc + R;
real P_adults = S1 + Lf1 + Ls1 + As1 + Ras + Ms + Ac1 + Rac + Mc + R1 + Lr + Ar + Tr + Rr + R2;
real Total_pop = P_children + P_adults;
// Birth calculation
real Birth = mu_sev * (Ms + Mc + M) + mu_r * Tr + mu * P_children + mu1 * P_adults;
// Treatment probabilities
real ptrt = pseek * ptest * psensi * ptreat;
real ptrt_s = pseek_s * ptest_s * psensi_s * ptreat_s;
real ptrt_c = pseek_c * ptest_c * psensi_c * ptreat_c;
real ptrt_r = pseek_r * ptest_r * psensi_r * ptreat_r;
// Total infectious calculation
real Total_infectious = c_As * As1 + c_Ac * (Ac + Ac1) + c_M * (Ms + Mc + M) + Ar + c_Tr * Tr;
// Forces of infection
real lambda = p_ch * ch * Total_infectious / Total_pop;
real lambda1 = p_ad * ad * Total_infectious / Total_pop;
vector[22] dydt;
// ODE equations... (these are unchanged)
dydt[1] = Birth - (mu + sigma + lambda) * S;
dydt[2] = lambda * S + (1 - PR) * lambda * Ls + tau * (1 - PR_a) * R - (w_fs + theta_param + psi + sigma + mu) * Lf;
dydt[3] = w_fs * Lf - ((1 - PR) * lambda + alpha + sigma + psi + mu) * Ls;
dydt[4] = alpha * Ls + q * Rc + theta_param * Lf - (chi + sigma + ptrt + mu) * Ac;
dydt[5] = ptrt * Ac - ((f_c + s_c) / z + sigma + mu_sev + mu) * M;
dydt[6] = f_c / z * M - (q + sigma + mu) * Rc;
dydt[7] = s_c / z * M + chi * Ac - (tau * (1 - PR_a) + sigma + mu) * R;
dydt[8] = sigma * S - (mu1 + lambda1) * S1;
dydt[9] = p * lambda1 * S1 + sigma * Lf + tau1 * (1 - PR_a) * R1 + (1 - PR) * p * lambda1 * Ls1 - (w_fs + theta1 + psi + mu1) * Lf1;
dydt[10]= w_fs * Lf1 + sigma * Ls - ((1 - PR) * p * lambda1 + alpha1 + psi + mu1) * Ls1;
dydt[11]= alpha1 * n * Ls1 + q * Ras + theta1 * Lf1 + rho * Ac1 - (chi + ptrt + omega + psi1 + mu1) * As1;
dydt[12]= f_as / z * Ms - (q + mu1) * Ras;
dydt[13]= ptrt * As1 - ((f_as + s_as) / z + mu_sev + mu1) * Ms;
dydt[14]= (1 - n) * alpha1 * Ls1 + omega * As1 + sigma * Ac + q * Rac - (rho + ptrt + chi + mu1) * Ac1;
dydt[15]= f_ac / z * Mc + sigma * Rc - (q + mu1) * Rac;
dydt[16]= ptrt * Ac1 + sigma * M - ((f_ac + s_ac) / z + mu_sev + mu1) * Mc;
dydt[17]= sigma * R + chi * (Ac1 + As1) + s_as / z * Ms + s_ac / z * Mc - (tau1 * (1 - PR_a) + psi + mu1) * R1;
dydt[18] = (1 - p) * lambda1 * S1 + psi * (Ls + Lf + Lf1 + Ls1 + R1) + (1 - PR_r) * (1 - p) * lambda1 * Ar + tau1 * (1 - PR_rr) * R2 - (iota + mu1) * Lr;
dydt[19]= psi1 * As1 + iota * Lr + h * Rr - (ptrt_r + mu1 + (1 - PR_r) * (1 - p) * lambda1 + chi) * Ar;
dydt[20]= ptrt_r * Ar - ((f_ar + s_ar) / zr + mu1 + mu_r) * Tr;
dydt[21]= f_ar / zr * Tr - (h + mu1) * Rr;
dydt[22]= chi * Ar + s_ar / zr * Tr - (tau1 * (1 - PR_rr) + mu1) * R2;
return dydt;
}
}
data {
int<lower=1> n_times;
int<lower=1> n_obs;
real t0;
array[n_times] real ts;
array[n_times, n_obs] int<lower=0> y_observed;
vector[22] y0;
array[55] real x_r;
array[0] int x_i; // Required by ode_rk45
}
transformed data {
real x_r_t[n_times, 55];
for (t in 1:n_times) {
real year = ts[t];
// Create a temporary copy for modification
array[55] real x_r_temp = x_r;
// --- Apply interventions based on the year ---
if (year >= 2002 && year < 2006) {
x_r_temp[30] *= 1.05; x_r_temp[27] *= 1.10; x_r_temp[28] *= 1.03;
} else if (year >= 2006 && year < 2013) {
x_r_temp[30] *= 1.05; x_r_temp[27] *= 1.04; x_r_temp[28] *= 1.10;
} else if (year >= 2013 && year < 2015) {
x_r_temp[30] *= 1.35; x_r_temp[27] *= 1.20; x_r_temp[28] *= 1.08;
} else if (year >= 2015 && year < 2022) {
x_r_temp[30] *= 1.20; x_r_temp[27] *= 1.20; x_r_temp[28] *= 1.08;
} else if (year >= 2022) {
x_r_temp[30] *= 1.15; x_r_temp[27] *= 1.10; x_r_temp[28] *= 1.08;
}
if (year >= 2006 && year < 2019) {
x_r_temp[40] *= 1.40; x_r_temp[42] *= 1.25; x_r_temp[28] *= 1.40;
} else if (year >= 2019) {
x_r_temp[14] *= 1.265; x_r_temp[13] *= 0.53;
}
if (year >= 2005 && year < 2012) {
x_r_temp[28] *= 1.05; x_r_temp[27] *= 1.01;
} else if (year >= 2012 && year < 2016) {
x_r_temp[28] *= 1.10; x_r_temp[27] *= 1.15;
} else if (year >= 2016 && year < 2024) {
x_r_temp[28] *= 1.15; x_r_temp[27] *= 1.20;
}
// Assign the modified temp array to the row of the 2D array
for (j in 1:55) {
x_r_t[t, j] = x_r_temp[j];
}
}
}
parameters {
real<lower=0.00001, upper=0.005> alpha_base;
real<lower=0.00001, upper=0.005> alpha1_base;
real<lower=0.05, upper=0.2> theta_param_base;
real<lower=0.05, upper=0.2> theta1_base;
real<lower=0.009, upper=0.8> q;
real<lower=0.0001, upper=0.0004> psi;
real<lower=0.0002, upper=0.0006> psi1;
real<lower=0.009, upper=0.75> h;
}
transformed parameters {
vector[22] y_hat[n_times];
vector[22] y_curr = y0;
for (t in 1:n_times) {
real t_start = (t == 1) ? t0 : ts[t-1];
array[1] real t_end = {ts[t]};
vector[8] theta;
real year = ts[t];
theta[1] = alpha_base;
theta[2] = alpha1_base;
theta[3] = theta_param_base;
theta[4] = theta1_base;
if (year >= 2014 && year < 2020) {
theta[3] *= (1 - 0.15);
theta[1] *= (1 - 0.15);
} else if (year >= 2020) {
theta[3] *= (1 - 0.15);
theta[1] *= (1 - 0.15);
theta[4] *= (1 - 0.30);
theta[2] *= (1 - 0.30);
}
theta[5] = q;
theta[6] = h;
theta[7] = psi;
theta[8] = psi1;
// ** Extract the correct row from the 2D array x_r_t
array[55] real x_r_row = x_r_t[t];
array[1] vector[22] y_hat_segment = ode_rk45(tb_ode, y_curr, t_start, t_end, theta, x_r_row, x_i);
y_hat[t] = y_hat_segment[1];
y_curr = y_hat[t];
}
}
model {
// Priors
alpha_base ~ gamma(2, 1000);
alpha1_base ~ gamma(2, 1000);
theta_param_base ~ gamma(4, 33);
theta1_base ~ gamma(4, 33);
q ~ gamma(2, 5);
h ~ gamma(2, 5);
psi ~ gamma(5, 20000);
psi1 ~ gamma(6, 15000);
// Likelihood
for (t in 1:n_times) {
real mu1 = fmax(y_hat[t, 4] + y_hat[t, 5], 1e-6);
y_observed[t, 1] ~ poisson(mu1);
real mu2 = fmax(y_hat[t, 11] + y_hat[t, 14] + y_hat[t, 16] + y_hat[t, 13], 1e-6);
y_observed[t, 2] ~ poisson(mu2);
real mu3 = fmax(y_hat[t, 19] + y_hat[t, 20], 1e-6);
y_observed[t, 3] ~ poisson(mu3);
}
}
generated quantities {
real y_pred[n_times, n_obs];
real log_lik[n_times];
for (t in 1:n_times) {
log_lik[t] = 0;
real mu1 = fmax(y_hat[t, 4] + y_hat[t, 5], 1e-6);
y_pred[t, 1] = poisson_rng(mu1);
log_lik[t] += poisson_lpmf(y_observed[t, 1] | mu1);
real mu2 = fmax(y_hat[t, 11] + y_hat[t, 14] + y_hat[t, 16] + y_hat[t, 13], 1e-6);
y_pred[t, 2] = poisson_rng(mu2);
log_lik[t] += poisson_lpmf(y_observed[t, 2] | mu2);
real mu3 = fmax(y_hat[t, 19] + y_hat[t, 20], 1e-6);
y_pred[t, 3] = poisson_rng(mu3);
log_lik[t] += poisson_lpmf(y_observed[t, 3] | mu3);
}
}