When placing ode_solver out of the loop takes long to run

Modeling
1

Hi,

When I run the following ODEs with random effects, it takes too long, even if it says that it takes 0.013 seconds, which I do not see long.
Previously I ran this model without random effects, in which I called the ode function in the transformed parameters (without any loop), where I also specified the rest of the parameter transformations and was working well.
Now that I have the random effects and I have to define it in a loop (to specify them for each individual), I need to define this in the model statement, and therefore I had to move everything in the model statement.
What could improve the memory of this code?

functions {
  real[] sldecay(real t, real[] y, real[] theta,  real[] x_r, int[] x_i) {
    real dydt[3];
    dydt[1] = -theta[1]*y[1];
    dydt[2] = -theta[2]*y[2];
    dydt[3] = theta[1]*y[1] + theta[2]*y[2] - theta[3]*y[3];
    
    return dydt;
  }
}
data {
  int<lower=1> N;
  int<lower=1> nsubjects;
  int<lower=1> subj_ids[N];
  real <lower=0> y[N];
  int <lower=1> ntimes;
  real ts[ntimes];
  real t0;
}
transformed data {
    real x_r[0];
    int x_i[0]; 
}
parameters {
  real logps;
  real logpl;
  real logmuab;
  real <lower = 0> sigmalogmuab;
  
  real logAB0;
  real <lower = 0> sigmalogAB0;
  real logps0;
  real logpl0;
  real <lower= 0> sigma;
  real rho_rhobit;
  
  vector[2] r[nsubjects];
}
model {
  real inits[3];
  real mu[N, 3];
  real theta[3];
  real new_mu[N,3];
  vector[nsubjects] rAB0;
  vector[nsubjects] rmuab;
  // transformed parameters 
  real ps = exp(logps);
  real pl = exp(logpl);
  real muab = exp(logmuab);
  real sigma_muab = exp(sigmalogmuab);
  real AB0 = exp(logAB0);
  real sigma_AB0 = exp(sigmalogAB0);
  real ps0 = exp(logps0);
  real pl0 = exp(logpl0);
  
  // defining the correlation 
  vector[2] mean_vec;
  matrix[2,2] Sigma; 
  real rho = (exp(rho_rhobit) - 1)/(exp(rho_rhobit) + 1);
  mean_vec[1] = -pow(sigma_AB0, 2.0)/2;
  mean_vec[2] = -pow(sigma_muab, 2.0)/2;
  
  Sigma[1,1] = pow(sigma_AB0, 2.0);
  Sigma[2,2] = pow(sigma_muab, 2.0);
  Sigma[1,2] = rho*sigma_AB0*sigma_muab;
  Sigma[2,1] = rho*sigma_AB0*sigma_muab;
  
  //prior distributions
  logps ~ normal(0, 1);
  logpl ~ normal(0, 1);
  logmuab ~ normal(0, 10);
  sigmalogmuab ~ normal(0, 1);
  logAB0 ~ normal(5, 10);
  sigmalogAB0 ~ normal(0, 1);
  logps0 ~ normal(0, 1);
  logpl0 ~ normal(0, 1);
  sigma ~ normal(1, 10);
  rho_rhobit ~ uniform(-10,10);
  
  theta[1] = ps;
  theta[2] = pl;
  
  // defining the distribution for the random effects
   for (subj_id in 1:nsubjects) {
        r[subj_id] ~ multi_normal(mean_vec, Sigma);
        rAB0[subj_id] = r[subj_id,1];
        rmuab[subj_id] = r[subj_id,2];
  }
  
  for (obs_id in 1:N){
    inits[1] = AB0*exp(rAB0[subj_ids[obs_id]]);
    inits[2] = inits[1]/ps0;
    inits[3] = inits[1]/pl0;
    
    theta[3] = muab*exp(rmuab[subj_ids[obs_id]]);
  }
  
  //defining ode_solver
  mu = integrate_ode_rk45(sldecay, inits, t0, ts, theta, x_r, x_i);

    for (obs in 1:N){
      new_mu[obs,1] = mu[obs,1];
      new_mu[obs,2] = mu[obs,2];
      new_mu[obs,3] = log(mu[obs,3]) - pow(sigma, 2.0)/2;
    //likelihood function for AB levels
    y[obs] ~ lognormal(new_mu[obs,3], sigma); 
}
}
generated quantities {
  real z_pred[N];
  real sigma2 = sigma;
  for (t in 1:N){
   z_pred[t] = lognormal_rng(log(y[t] + 1e-5), sigma2); 
  }
}
2

Does every subject have the same number of observations and the same t0/ts? You may want to organize the code per-subject

data {
  int<lower=1> N;
  int<lower=1> nsubjects;
  int <lower=1> ntimes;
  real <lower=0> y[nsubjects,ntimes];
  real ts[ntimes];
  real t0;
}
...
model {
  ...
  inits[2] = inits[1]/ps0;
  inits[3] = inits[1]/pl0;
  for (subj_id in 1:N){
    inits[1] = AB0*exp(rAB0[subj_id]);
    theta[3] = muab*exp(rmuab[subj_ids[obs_id]]);
    real mu[ntimes, 3];
    mu = integrate_ode_rk45(sldecay, inits, t0, ts, theta, x_r, x_i);
    real new_mu[ntimes];
    for (t in 1:ntimes) {
      new_mu[t] = mu[t,3] - pow(sigma, 2.0)/2;
    }
    y[subj_id] ~ lognormal(new_mu, sigma);
  }
}

Note that the first two components in the differential equation, y_1 and y_2, can be integrated analytically.

y[1] = init[1]*exp(-theta[1]*t);
y[2] = init[2]*exp(-theta[2]*t);

Substituting these simplifies the integration

real[] sldecay(real t, real[] y, real[] theta,  real[] x_r, int[] x_i) {
  real dydt[1];
  dydt[1] = theta[4]*exp(-theta[1]*t) + theta[5]*exp(-theta[2]*t) - theta[3]*y[1];
  return dydt;
}
...
real I = AB0*exp(rAB0[subj_ids[obs_id]]);
real inits[1];
inits[1] = I/pl0;
real theta[5];
theta[1] = ps;
theta[2] = pl;
theta[3] = muab*exp(rmuab[subj_ids[obs_id]]);
theta[4] = theta[1]*I;
theta[5] = theta[2]*I/ps0;
real mu[ntimes,1];
mu = integrate_ode_rk45(sldecay, inits, t0, ts, theta, x_r, x_i);

And WolframAlpha says even that ODE can be solved in closed-form.

2 Likes
3

Actually, ntimes refers to a vector that goes from 1 to 96 (which is the range of the time when the blood sample was taken), to specify that the time to solve the equation has these values. Did you refer to times as the number of observations that each individual has? The model does not run right now, it says: Initialization between (-2, 2) failed, error in the line where I call the ode_solver

4

Err, you mean “ts refers to a vector…”? The model you posted has

  int <lower=1> ntimes;
  real ts[ntimes];

But does the error say anything else? Like, times is not sorted or some array is the wrong size?

Regardless, if the closed-form solution from WolframAlpha is correct then you don’t need to call integrate_ode_rk45 at all.

5

Yes, I get something else the error comes from the line where I define the new_mu (log(mu[t, 3]) - pow(sigma, 2.0)/2. I still do not see what is wrong with that line.
When I get this error my R session aborts and unfortunately I am not able to see anything from the error message.

6

Are you using RStan or CmdStanR? The latter might be less likely to crash your entire session.

7

To answer to your previous question, not all the individuals have the same number of longitudinal observations.
The code that I have created now which is still giving me the same issues is the following, where nobs[subjects] is a vector that records the number of observations per subject.

functions {
  real[] sldecay(real t, real[] y, real[] theta,  real[] x_r, int[] x_i) {
    real dydt[3];
    dydt[1] = -theta[1]*y[1];
    dydt[2] = -theta[2]*y[2];
    dydt[3] = theta[1]*y[1] + theta[2]*y[2] - theta[3]*y[3];
    
    return dydt;
  }
}
data {
  int<lower=1> N;
  int<lower=1> nsubjects;
  int<lower=1> subj_ids[N];
  real <lower=0> y[N];
  int <lower=1> ntimes;
  int <lower=1> nobs[nsubjects];
  real ts[ntimes];
  real t0;
}
transformed data {
    real x_r[0];
    int x_i[0]; 
}
parameters {
  real logps;
  real logpl;
  real logmuab;
  real sigmalogmuab;
  
  real logAB0;
  real sigmalogAB0;
  real logps0;
  real logpl0;
  real <lower = 0> sigma;
  
  vector[2] r[nsubjects];
}
model{
  real mu[ntimes, 3];
  real new_mu[ntimes];
  vector[nsubjects] rAB0;
  vector[nsubjects] rmuab;
  // transformed parameters 
  real ps = exp(logps);
  real pl = exp(logpl);
  real muab = exp(logmuab);
  real sigma_muab = exp(sigmalogmuab);
  real AB0 = exp(logAB0);
  real sigma_AB0 = exp(sigmalogAB0);
  real ps0 = exp(logps0);
  real pl0 = exp(logpl0);
  
  // defining the correlation 
  vector[2] mean_vec;
  matrix[2,2] Sigma; 
  mean_vec[1] = -pow(sigma_AB0, 2.0)/2;
  mean_vec[2] = -pow(sigma_muab, 2.0)/2;
  
  Sigma[1,1] = pow(sigma_AB0, 2.0);
  Sigma[2,2] = pow(sigma_muab, 2.0);
  Sigma[1,2] = 0;
  Sigma[2,1] = 0;
  
  //prior distributions
  logps ~ normal(0, 1);
  logpl ~ normal(0, 1);
  logmuab ~ normal(0, 10);
  sigmalogmuab ~ normal(0, 1);
  logAB0 ~ uniform(5, 10);
  sigmalogAB0 ~ normal(0, 1);
  logps0 ~ normal(0, 1);
  logpl0 ~ normal(0, 1);
  sigma ~ normal(1, 10);
  
  for (j in 1:nsubjects) {
    real inits[3];
    real theta[3];
    //defining the initial conditions
    inits[1] = inits[3]/ps0;
    inits[2] = inits[3]/pl0;
    inits[3] = AB0*exp(rAB0[j]);
    //defining the parameters
    theta[1] = ps;
    theta[2] = pl;
    theta[3] = muab*exp(rmuab[j]);
    // ode solver
    mu = integrate_ode_rk45(sldecay, inits, t0, ts, theta, x_r, x_i);
    for (t in 1:nobs[nsubjects]){
      new_mu[t] = log(mu[t, 3]) - pow(sigma, 2.0)/2;
      y[t] ~ lognormal(new_mu[t], sigma);  
      r[j] ~ multi_normal(mean_vec, Sigma);
      rAB0[j] = r[j,1];
      rmuab[j] = r[j,2];
    }
  }
}
8

I am using Rstan. I will try out cmstanr

9

I see a couple of problems with that model.

First, unlike BUGS or JAGS, Stan executes the code strictly in source order and therefore variable definitions must come before use. You can’t leave rAB and rmuab at the end of the loop. Assign them at the beginning

  for (j in 1:nsubjects) {
    rAB0[j] = r[j,1];
    rmuab[j] = r[j,2];
    ...

Second, nobs[nsubjects] is the number of observations for the last subject, you need to use nobs[j] for the jth subject. Looks like you only need the first j times so I’d write this as

    mu[1:nobs[j]] = integrate_ode_rk45(sldecay, inits, t0, ts[1:nobs[j]], theta, x_r, x_i);
    for (t in 1:nobs[j]){
1 Like
10

Thanks! I will try that.
I have run the model with cmdstanr, and it turns out that the error is the following:

cmdstanr
cmdstanr1059Ă—249 19.6 KB

So, it comes when I define the initial conditions which are strange that the value is Nan, because I did transform it with the exponential. Do you know what can be causing this issue?

11

NaN is unassigned value. You define inits[1] in term of inits[3] but at that point inits[3] is not defined yet so it comes out as NaN. Switch the order around.

    rAB0[j] = r[j,1];
    inits[3] = AB0*exp(rAB0[j]);
    inits[1] = inits[3]/ps0;
    inits[2] = inits[3]/pl0;
2 Likes