Replacing a constant parameter with a time-dependent function

[edit: escaped code]

Hi, I am a just learning how to use Rstan. I want to replace the constant parameter “psi” in my stan code with a time-dependent function defined as

psi_t <- ifelse(t >= 1 & t < 31, 0,
                         ifelse(t >= 31 & t < 32, 0.94,
                                ifelse(t >= 32 & t < 33, 1.72,
                                       ifelse(t >= 33 & t < 34, 2.12,
                                              ifelse(t >= 34 & t <35, 2.04, 
                                                     ifelse(t>=35 & t <= 48, 0,0)
                                              )
                                       )
                                )
                         )
         )  

How do I go about this? This is my code below. The code runs perfectly well with the constant parameter “psi”.

  functions {
  real seasonal_forcing(real t, real amp, real phi){
  return(1 + amp * cos(2 * pi() * (t - phi) / 12));
  }
  real human_pop(real Sh, real Mc, real Eh, real Ah, real Ih, real Rh){
  return(Sh + Mc + Eh + Ah + Ih + Rh);
  }
  real mosquito_pop(real Sm, real Em, real Im){
  return(Sm + Em + Im);
  }
  real[] malaria(real t, real[] y, real[] params, 
             real[] x_r, int[] x_i) {
    
  real Sh = y[1];
  real Mc = y[2];
  real Eh = y[3];
  real Ah = y[4];
  real Ih = y[5];
  real Rh = y[6];
  real Mh = y[7];
  real Sm = y[8];
  real Em = y[9];
  real Im = y[10];
  
  real Lambda_h = x_r[1];
    real Lambda_v = x_r[2];
    real mu_h = x_r[3];
    real mu_v = x_r[4];
    real rho_h = x_r[5];
    real rho_v = x_r[6];
    real kappa = x_r[7];
    real delta_h    = x_r[8];
    real alpha_a = x_r[9];
    real alpha_i = x_r[10];
    real psi = x_r[11]; 
  
  real beta_h = params[1];            
  real beta_v = params[2];
  real sigma = params[3];
  real p = params[4];
  real amp = params[5];
  real phi = params[6];
  real eta = params[7];
  real nu = params[8];
  real gamma = params[9];
  real tau = params[10];


  real N_h = human_pop(Sh, Mc, Eh, Ah, Ih, Rh);
  real N_m = mosquito_pop(Sm, Em, Im);
    
    real forcing_function = seasonal_forcing(t,amp,phi); 
    real beta_1 =  beta_h * forcing_function;
    real beta_2 =  beta_v * forcing_function;
   
    real dSh_dt = Lambda_h-(beta_1*Im*Sh)/N_h-(psi+mu_h)*Sh+eta*Mc+kappa*nu*Rh;
    real dMc_dt = psi*Sh-(1-tau)*(beta_1*Im*Mc)/N_h-(eta+mu_h)*Mc + (1-kappa)*nu*Rh;
    real dEh_dt = (beta_1*Im*Sh)/N_h+(1-tau)*(beta_1*Im*Mc)/N_h-(rho_h+mu_h)*Eh;
    real dAh_dt = (1-p)*rho_h*Eh - (gamma*beta_1*Im*Ah)/N_h-(mu_h+alpha_a)*Ah;
    real dIh_dt = p*rho_h*Eh+(gamma*beta_1*Im*Ah)/N_h-(alpha_i+mu_h+delta_h)*Ih;
    real dRh_dt = alpha_i*Ih+alpha_a*Ah-(mu_h+nu)*Rh;
    real dMh_dt = p*rho_h*Eh+(gamma*beta_1*Im*Ah)/N_h;
    real dSm_dt = Lambda_v-mu_v*Sm-(beta_2*(Ih+sigma*Ah)*Sm)/N_h;
    real dEm_dt = (beta_2*(Ih+sigma*Ah)*Sm)/N_h - (rho_v+mu_v)*Em;
    real dIm_dt = rho_v*Em - mu_v*Im;
     
  return{dSh_dt,dMc_dt,dEh_dt,dAh_dt,dIh_dt,dRh_dt,dMh_dt,dSm_dt,dEm_dt,dIm_dt};
  }
}
data {
  int <lower=1> n_months;
  real y0[10];/
  real t0;
  real ts[n_months];
  real Lambda_h;
  real Lambda_v;
  real mu_h;
  real mu_v;
  real rho_h;
  real rho_v;
  real kappa;
  real delta_h;
  real alpha_a;
  real alpha_i;
  real psi;


  int total_cases[n_months];
}
transformed data {
  real x_r[11] = {Lambda_h, Lambda_v, mu_h, mu_v, rho_h, rho_v, 
                    kappa, delta_h, alpha_a, alpha_i, psi};
  int x_i[0];
}
parameters {
  real<lower=0, upper=1> beta_h;
  real<lower=0, upper=1> beta_v;
  real<lower=0, upper=1> sigma;
  real<lower=0, upper=1> p;
  real<lower=0, upper=1> amp;
  real phi;
  real<lower=0, upper=1> eta;
  real<lower=0, upper=1> nu;
  real<lower=0, upper=1> gamma;
  real<lower=0, upper=1> tau;
  real<lower=0> zeta_inv;
} 
transformed parameters{
  real y[n_months, 10];
  real zeta = 1. / zeta_inv;
  real incidence[n_months - 1];
  real params[10];
  params[1] = beta_h;            
  params[2] = beta_v;
  params[3] = sigma;
  params[4] = p;
  params[5] = amp;
  params[6] = phi;
  params[7] = eta;
  params[8] = nu;
  params[9] = gamma;
  params[10] = tau;
  
    y = integrate_ode_rk45(diph, y0, t0, ts, params, x_r, x_i);
    
    for (i in 1:n_months-1){
    incidence[i] = y[i+1, 7] - y[i, 7];
    }
  }
model {
  //priors
  beta_h ~ normal(0.5,0.5);             
  beta_v ~ normal(0.5,0.5);
  sigma ~ normal(0.5,0.5);
  p ~ normal(0.5,0.5);
  amp ~ normal(0.5,0.5);
  phi ~ normal(10,10);
  eta ~ normal(0.5,0.5);
  nu ~ normal(0.5, 0.5);
  gamma ~ normal(0.5,0.5);
  tau ~ normal(0.5,0.5);
  zeta_inv ~ exponential(5);
   
  
  //sampling distribution
  //col(matrix x, int n) - The n-th column of matrix x. Here the number of infected people 
  total_cases[1:(n_months-1)] ~ neg_binomial_2(incidence, phi);
}
generated quantities {
  real pred_cases[n_months-1];
  pred_cases = neg_binomial_2_rng(incidence, phi);
}

The simplest thing to do is to define phi as a function and then, whenever you would have used phi, use phi(t) in its place.

The problem you’re likely to have with this is that a discrete function like this based on comparisons isn’t smoothly differentiable in time. If you look at what autodiff is going to do here, you will have a derivative of zero everywhere other than the decision points, where it’s undefined. This can be a problem with the sampler, which can run into acceptance problems around boundaries.

The alternative is to use some kind of smooth spline which is differentiable w.r.t. t rather than something piecewise constant.

Thanks a lot for your response