Mechanistic model with stochastic forcing

Modeling
1

I am working on an ODE-based mechanistic model of malaria transmission dynamics. It is necessary for me to introduce a forcing function into the force of infection, which is the product of a sinusoidal function and a non-centred random walk. However, the sampling could only achieve 1% after more than 12 hours of sampling. How can I possibly speed things up. The code is pasted for your consideration

  real human_pop(real Sh, real Eh, real Ih, real Ah){
    return (Sh + Eh + Ih + Ah);
  }
  
  real mosquito_pop(real Svi, real Svo, real Evi, real Evo, real Ivi, real Ivo){
    return (Svi + Svo + Evi + Evo + Ivi + Ivo);
  }
  
  real llin_prot(real t, real u_1, real u_2, real omegap,
                 real tau_1, real tau_2, real T_1, real T_last) {
    if (0 <= t && t < tau_2)
      return u_1 * omegap * exp(-log(2) * square((t - tau_1) / T_1));
    else if (tau_2 <= t && t <= T_last)
      return u_2 * omegap * exp(-log(2) * square((t - tau_2) / T_1));
    else
      return 0;
  }
  
  real llin_kill(real t, real u_1, real u_2, real omegak,
                 real tau_1, real tau_2, real T_1, real T_last) {
    if (0 <= t && t < tau_2)
      return u_1 * omegak * exp(-log(2) * square((t - tau_1) / T_1));
    else if (tau_2 <= t && t <= T_last)
      return u_2 * omegak * exp(-log(2) * square((t - tau_2) / T_1));
    else
      return 0;
  }
  
  vector akwa_ibom_malaria(real t, vector y, vector parms,
                           data vector x_r, data array[] int x_i) {
    
    int n_months = x_i[1];  
    
    real Lamda_h = x_r[1];
    real sigma_h = x_r[2];
    real gamma = x_r[3];
    real mu_h = x_r[4];
    real delta_h = x_r[5];
    real Lamda_v1= x_r[6];
    real mu_v = x_r[7];
    real sigma_v = x_r[8];
    real a = x_r[9];
    real u_1 = x_r[10];
    real u_2 = x_r[11];
    real q = x_r[12];
    real omegap = x_r[13];
    real omegak = x_r[14];
    real tau_1 = x_r[15];
    real tau_2 = x_r[16];
    real T_last = x_r[17];
    real T_1 = x_r[18];
    real psi = x_r[19];
    real p_1 = x_r[20];
    real varepsilon_vi = x_r[21];
    real varepsilon_vo = x_r[22];
    
    real beta_h = parms[1];
    real p_2   = parms[2];
    real beta_v = parms[3];
    real alpha_io = parms[4];
    real alpha_oi = parms[5];
    real amp = parms[6];
    real b   = parms[7];
    
    vector[n_months] RW;
    for (k in 1:n_months)
      RW[k] = parms[7 + k];
    
    // unpack states
    real Sh = y[1];
    real Eh = y[2];
    real Ih = y[3];
    real Ah = y[4];
    real Qh = y[5];
    real Svi = y[6];
    real Svo = y[7];
    real Evi = y[8];
    real Evo = y[9];
    real Ivi = y[10];
    real Ivo = y[11];
    
    real N_h = human_pop(Sh, Eh, Ih, Ah);
    real N_v = mosquito_pop(Svi, Svo, Evi, Evo, Ivi, Ivo);
    
    real u_p = llin_prot(t, u_1, u_2, omegap, tau_1, tau_2, T_1, T_last);
    real u_k = llin_kill(t, u_1, u_2, omegak, tau_1, tau_2, T_1, T_last);
    
    // ✔ MATCHED TO CODE 1 — correct linear RW interpolation with fabs()
    real RW_t = 0;
    for (k in 1:n_months) {
    real w = 1 - abs(t - k);
    if (w > 0)
      RW_t += RW[k] * w;
  }
    
    real seasonal = 1 + amp * cos(2 * pi() * (t - b) / 12);
    
    // ✔ MATCHED TO CODE 1
  real forcing_t = seasonal * exp(RW_t);
    
    // Infection terms
    real lamda_h1 = a * beta_h * varepsilon_vi * Ivi / N_h;
    real lamda_h2 = (1 - a) * beta_h * varepsilon_vi * Ivi / N_h;
    real lamda_h3 = beta_h * varepsilon_vo * Ivo / N_h;
    real lamda_h = forcing_t * ((1 - u_p)*lamda_h1 + lamda_h2 + lamda_h3);
    
    real lamda_vi_1 = a * beta_v * varepsilon_vi * (psi*Ih + Ah) / N_h;
    real lamda_vi_2 = (1 - a) * beta_v * varepsilon_vi * (psi*Ih + Ah) / N_h;
    real lamda_vi =  forcing_t *((1 - u_p)*lamda_vi_1 + lamda_vi_2);
    
    real lamda_vo = forcing_t *(beta_v * varepsilon_vo * (psi*Ih + Ah) / N_h);
    
    real mu_k = a * varepsilon_vi * u_k;
    
    vector[11] dydt;
    
    dydt[1] = Lamda_h + gamma*Ih - lamda_h*Sh - mu_h*Sh;
    dydt[2] = lamda_h*Sh - sigma_h*Eh - mu_h*Eh;
    dydt[3] = p_1*sigma_h*Eh + p_2*lamda_h*Ah - gamma*Ih - (mu_h+delta_h)*Ih;
    dydt[4] = (1-p_1)*sigma_h*Eh - p_2*lamda_h*Ah - mu_h*Ah;
    dydt[5] = p_1*sigma_h*Eh + p_2*lamda_h*Ah;
    dydt[6] = q*Lamda_v1 - lamda_vi*Svi - (mu_v+mu_k)*Svi - alpha_io*Svi + alpha_oi*Svo;
    dydt[7] = (1-q)*Lamda_v1 - lamda_vo*Svo - mu_v*Svo - alpha_oi*Svo + alpha_io*Svi;
    dydt[8] = lamda_vi*Svi - sigma_v*Evi - (mu_v+mu_k)*Evi - alpha_io*Evi + alpha_oi*Evo;
    dydt[9] = lamda_vo*Svo - sigma_v*Evo - mu_v*Evo - alpha_oi*Evo + alpha_io*Evi;
    dydt[10] = sigma_v*Evi - (mu_v+mu_k)*Ivi - alpha_io*Ivi + alpha_oi*Ivo;
    dydt[11] = sigma_v*Evo - mu_v*Ivo - alpha_oi*Ivo + alpha_io*Ivi;
    
    return dydt;
  }
}
data {
  int <lower=1> n_months;
  //int <lower=1> n_years;
  int <lower=1> n_survey_years;
  int <lower=1> n_DHS_MIS_years;
  vector[11] y0;
  real t0;
  array[n_months] real ts;
  real Lamda_h;
  real sigma_h;
  real gamma;
  //real gamma_2;
  //real gamma_3;
  real mu_h;
  real delta_h;
  real Lamda_v1;
  real mu_v;
  real sigma_v;
  real a;
  real u_1;
  real u_2;
  real q;
  //real rho;
  real omegap;
  real omegak;
  real tau_1;
  real tau_2;
  real T_last;
  real T_1;
  real psi;
  real p_1;
  real varepsilon_vi;
  real varepsilon_vo;
  
  
  
  array[n_months] int akwa_ibom_cases;
 // array[n_years] int hum_pop;
  array[n_survey_years] int outdoor_mosq_pop;
  array[n_survey_years] int indoor_mosq_pop;
  int sporo_tested_indoor;
  int sporo_tested_outdoor;
  int sporo_pos_indoor;
  int sporo_pos_outdoor;
  array[n_DHS_MIS_years] int prev_n_tested;
  array[n_DHS_MIS_years] int prev_n_pos;

  }
transformed data {
  vector[22] x_r;
  
  x_r = to_vector({Lamda_h, sigma_h, gamma, mu_h, delta_h,
      Lamda_v1, mu_v, sigma_v, a, u_1, u_2, q, omegap, omegak,
      tau_1, tau_2, T_last, T_1, psi, p_1, varepsilon_vi, varepsilon_vo});
  array[1] int x_i;
  x_i[1] = n_months;
}
parameters {
  real<lower=0, upper =1> beta_h; 
  real<lower=0, upper =1> p_2;
  real<lower=0, upper =1> beta_v;
  real <lower =0> alpha_io;
  real <lower =0> alpha_oi;
  real <lower = 0, upper=1> amp;
  real<lower=0, upper=12> b;
//  real<lower=0, upper=1> p_r;
  real<lower=0, upper=1> r_p;
  real<lower=0> phi_1_inv;
  //real<lower=0> phi_2_inv;
  real<lower=0> phi_3_inv;
  real<lower=0> phi_4_inv;
  vector[n_months] RW_raw;     //latent RW increments
  real<lower=0> sigma_rw;       // rw step std
        
  
}
transformed parameters{
  array[n_months] vector[11] y;
  vector[n_months] RW;
  RW[1] = RW_raw[1] * sigma_rw;
  for (k in 2:n_months)
  RW[k] = RW[k-1] + RW_raw[k] * sigma_rw;
  real phi_1 = 1. / phi_1_inv;
 // real phi_2 = 1. / phi_2_inv;
  real phi_3 = 1. / phi_3_inv;
  real phi_4 = 1. / phi_4_inv;
  vector[n_months - 1] incidence;
  //vector[n_years] akwa_hum_pop;
  vector[n_survey_years] akwa_outdoor_mosq_pop;
  vector[n_survey_years] akwa_indoor_mosq_pop;
  real akwa_sporo_pos_indoor_prob;
  real akwa_sporo_pos_outdoor_prob;
  vector[n_DHS_MIS_years] akwa_prevalence;

  // pack parms
  vector[7+n_months] parms;
  parms[1] = beta_h;
  parms[2] = p_2;
  parms[3] = beta_v;
  parms[4] = alpha_io;
  parms[5] = alpha_oi;
  parms[6] = amp;
  parms[7] = b;
  // attach RW
  for (k in 1:n_months)
   parms[7 + k] = RW[k];
 
  // Solve ODE
   y = ode_rk45(akwa_ibom_malaria, to_vector(y0), t0, ts, parms, x_r, x_i);
   
    for (i in 1:(n_months - 1)){
    incidence[i] = (y[i+1, 5] - y[i, 5]);
    }
    
//  for (k in 1:n_years){
// akwa_hum_pop[k] = sum(y[12*k-11, 1:4]);
//  }
 
    akwa_indoor_mosq_pop = to_vector({sum({mean(to_vector(y[14:24,6])),mean(to_vector(y[14:24,8])),mean(to_vector(y[14:24,10]))}),
                            sum({mean(to_vector(y[25:36,6])),mean(to_vector(y[25:36,8])),mean(to_vector(y[25:36,10]))}),
                            sum({mean(to_vector(y[40:45,6])),mean(to_vector(y[40:45,8])),mean(to_vector(y[40:45,10]))}),
                            sum({mean(to_vector(y[47:57,6])),mean(to_vector(y[47:57,8])),mean(to_vector(y[47:57,10]))}),
                            sum({mean(to_vector(y[58:69,6])),mean(to_vector(y[58:69,8])),mean(to_vector(y[58:69,10]))}),
                            sum({mean(to_vector(y[70:81,6])),mean(to_vector(y[70:81,8])),mean(to_vector(y[70:81,10]))}),
                            sum({mean(to_vector(y[82:93,6])),mean(to_vector(y[82:93,8])),mean(to_vector(y[82:93,10]))})
                            });
    akwa_outdoor_mosq_pop = to_vector({sum({mean(to_vector(y[14:24,7])),mean(to_vector(y[14:24,9])),mean(to_vector(y[14:24,11]))}),
                            sum({mean(to_vector(y[25:36,7])),mean(to_vector(y[25:36,9])),mean(to_vector(y[25:36,11]))}),
                            sum({mean(to_vector(y[40:45,7])),mean(to_vector(y[40:45,9])),mean(to_vector(y[40:45,11]))}),
                            sum({mean(to_vector(y[47:57,7])),mean(to_vector(y[47:57,9])),mean(to_vector(y[47:57,11]))}),
                            sum({mean(to_vector(y[58:69,7])),mean(to_vector(y[58:69,9])),mean(to_vector(y[58:69,11]))}),
                            sum({mean(to_vector(y[70:81,7])),mean(to_vector(y[70:81,9])),mean(to_vector(y[70:81,11]))}),
                            sum({mean(to_vector(y[82:93,7])),mean(to_vector(y[82:93,9])),mean(to_vector(y[82:93,11]))})
                            });
   akwa_sporo_pos_indoor_prob = sum(to_vector(y[25:93,10]))/sum({sum(to_vector(y[25:93,6])),sum(to_vector(y[25:93,8])), sum(to_vector(y[25:93,10]))});
   akwa_sporo_pos_outdoor_prob = sum(to_vector(y[25:93,11]))/sum({sum(to_vector(y[25:93,7])),sum(to_vector(y[25:93,9])), sum(to_vector(y[25:93,11]))});
   akwa_prevalence = to_vector({(sum(to_vector(y[1:12,3])) + sum(to_vector(y[1:12,4])))/(sum(to_vector(y[1:12,1]))+sum(to_vector(y[1:12,2]))+sum(to_vector(y[1:12,3]))+sum(to_vector(y[1:12,4]))),
                     (sum(to_vector(y[37:48,3])) + sum(to_vector(y[37:48,4])))/(sum(to_vector(y[37:48,1]))+sum(to_vector(y[37:48,2]))+sum(to_vector(y[37:48,3]))+sum(to_vector(y[37:48,4]))),
                     (sum(to_vector(y[73:84,3])) + sum(to_vector(y[73:84,4])))/(sum(to_vector(y[73:84,1]))+sum(to_vector(y[73:84,2]))+sum(to_vector(y[73:84,3]))+sum(to_vector(y[73:84,4])))
                     });
  }
model {
  //priors
  beta_h ~ normal(0,0.5);
  p_2 ~ normal(0,0.5);
  beta_v ~ normal(0,0.5);
  alpha_io ~ normal(0,10);
  alpha_oi ~ normal(0,10);
 // p_r ~ normal(0,0.5);
  r_p ~ normal(0,0.5);
  phi_1_inv ~ exponential(5);
//  phi_2_inv ~ exponential(5);
  phi_3_inv ~ exponential(5);
  phi_4_inv ~ exponential(5);
  amp ~ normal(0, 0.5);
  b ~ normal(0, 10); 
  RW_raw ~ normal(0, 1);
  sigma_rw ~ exponential(1);
  

  
  //sampling distribution
  //col(matrix x, int n) - The n-th column of matrix x. Here the number of infected people 
  akwa_ibom_cases[1:(n_months-1)] ~ neg_binomial_2(r_p*incidence, phi_1);
  //hum_pop[1:n_years] ~ neg_binomial_2(akwa_hum_pop, phi_2);
  sporo_pos_indoor ~ binomial(sporo_tested_indoor, akwa_sporo_pos_indoor_prob);
  sporo_pos_outdoor ~ binomial(sporo_tested_outdoor, akwa_sporo_pos_outdoor_prob);
  outdoor_mosq_pop[1:n_survey_years]~ neg_binomial_2(akwa_outdoor_mosq_pop, phi_3);
  indoor_mosq_pop[1:n_survey_years] ~ neg_binomial_2(akwa_indoor_mosq_pop, phi_4);
  for (j in 1:n_DHS_MIS_years) { prev_n_pos[j] ~ binomial(prev_n_tested[j], akwa_prevalence[j]);}
  }
generated quantities {
  array[n_months-1] real pred_cases;
  pred_cases = neg_binomial_2_rng(r_p*incidence, phi_1);
  vector[n_months] forcing_t;
  
  for (i in 1:n_months) {
    real seasonal = 1 + amp * cos(2 * pi() * (i - b) / 12);
   forcing_t[i] = seasonal * exp(RW[i]);
}
} ```
1 Like
2

Hi, @Idowu, and sorry nobody’s replied until now. I’m afraid I don’t even know what a forcing function is—@charlesm93 should be able to help you with that if he’s hanging out on the forums.

The first question with something this complex is always whether you’ve implemented it correctly. It looks like you are using tests.

Suggestion 0. Get your indentation in order to make things easier to read. Pro tip: Don’t ever use tabs in computer code.

Suggestion 1. If your ODE winds up being stiff, which can happen with random initialization even if it’s not stiff where you eventually want to sample, then using a stiff solver in place of RK45 can make a huge difference. Even better is to supply good initial values.

Suggestion 2. Declare variables in the forms you will eventually us them. For example, you don’t need to apply to_vector before taking the mean of a what would otherwise have been a row vector.

Suggestion 3. Somehow don’t cut-and-paste all these indexes. Turn it into a function since it looks like indoor and outdoor only vary in using columns 6 and 8 vs. 7 and 9. Pull this into a function to cut down on potential errors. You could probably pull all this stuff out of being in a big matrix if that cutting is always going the same way. Or at least break the columns out into vectors. All this reshaping is really hard on memory, which is the main bottleneck in a big Stan program. (That’s why the new Mac ARM architecture is so much better for these programs.)

Suggesiton 4. Start simpler. Where you. have negative binomial, use a Poisson, for example. Sampling really gets frustrated when there are multiple explanations for an observation, as you always get with negative binomial (large values can indicate the mean’s too small or the over dispersion’s too small, which adds a lot of non-identifiability, which reduces step size and slows sampling).

Suggestion 5. Compute generated quantities after sampling using standalone generated quantities.

Usually the problem is a lot of duplicated computation, but your code looks good as far as that goes outside of the reshaping.