Hi all,
I am working on building an age-structured epidemiological model. Most examples I’ve found online write out each age group explicitly in the ODE function. However, my model is more complex, so that approach isn’t feasible for me.
I tried to implement the model using a for loop in the ODE function, but it’s not working as expected, and I can’t figure out why.
Does anyone have an example of how to implement an age-structured model in Stan or how to properly use a for loop in the ODE function? Any guidance or suggestions would be greatly appreciated!
Thanks in advance!
HI, @hassanhachem and welcome to the Stan forums.
There’s nothing special about for loops and how they interact with the ODE function. I assume you mean you use it in the function that calculates the derivatives? You can do that however you like.
If you cold share what’s going wrong, we might be able to offer more help. We often use age as a covariate in an ODE model to pick out a subject’s individual metabolic parameters in pharmacometric models, if that’s what you mean by an age-structured model.
Hello @Bob_Carpenter, thanks! :)
Yes, that’s what I’m trying to do. I’m using this in an epidemiological model. Below is an example of the code I’ve written, but I’m running into issues when trying to fit it. The error message seems to be related to the ODE system, not the RStan code itself. Here’s the error message I get:
SAMPLING FOR MODEL 'anon_model' NOW (CHAIN 1).
Chain 1: Rejecting initial value:
Chain 1: Error evaluating the log probability at the initial value.
Chain 1: Exception: integrate_ode_rk45: initial time is 0, but must be less than 0.000000 (in 'string', line 44, column 2 to column 98)
Here’s the code:
library(rstan)
stan_model_code <- '
functions {
real[] ode_system(real t, real[] y, real[] theta, real[] x_r, int[] x_i) {
int n = num_elements(y) / 3;
real dydt[num_elements(y)];
real beta = theta[1];
real mu = theta[2];
real gamma_oral = theta[3];
real pi_oral = theta[4];
real S[n];
real I[n];
real A[n];
// Extract S, I, A for each age group
for (i in 1:n) {
S[i] = y[3 * (i - 1) + 1];
I[i] = y[3 * (i - 1) + 2];
A[i] = y[3 * (i - 1) + 3];
}
// Calculate dydt for each age group
for (a in 1:n) {
dydt[3 * (a - 1) + 1] = -beta * S[a] * I[a] - mu * S[a];
dydt[3 * (a - 1) + 2] = beta * S[a] * I[a] - (gamma_oral + mu) * I[a] + pi_oral * A[a];
dydt[3 * (a - 1) + 3] = gamma_oral * I[a] - (pi_oral + mu) * A[a];
}
return dydt;
}
}
data {
int<lower=1> n; // Number of age groups
int<lower=1> T; // Number of time points
real y0[3 * n]; // Initial state
real times[T]; // Time points
real theta[4]; // Parameters: [beta, mu, gamma_oral, pi_oral]
}
parameters {
real<lower=0> sigma; // Observation noise
}
transformed parameters {
real y_hat[T, 3 * n]; // Model predictions
y_hat = integrate_ode_rk45(ode_system, y0, 0, times, theta, rep_array(0.0, 0), rep_array(0, 0));
}
model {
sigma ~ normal(0, 10);
for (t in 1:T) {
for (i in 1:(3 * n)) {
y_hat[t, i] ~ normal(0, sigma); // To test only
}
}
}
'
model <- stan_model(model_code = stan_model_code)
n <- 3
T <- 100
times <- seq(0, 10, length.out = T)
y0 <- c(rep(1000, n), rep(1, n), rep(0, n)) # S = 1000, I = 1, A = 0 for each age group
theta <- c(0.1, 0.01, 0.05, 0.02) # beta, mu, gamma_oral, pi_oral
data <- list(
n = n,
T = T,
y0 = y0,
times = times,
theta = theta
)
fit <- sampling(model, data = data, iter = 1000, chains = 4)
[edit: escaped code]
Line numbers in errors is one of the main motivations for putting Stan code in standalone files.
That error message doesn’t make sense. I’d try substituting 0.0 for 0 and see what happens. Usually Stan doesn’t care and autocasts integers to floats.
P.S. That looped sampling statement can be replaced with the much more efficient:
to_matrix(y_hat) ~ normal(0, sigma);
But with an ode solver also going per eval, you may not notice the saving. It will help with correctness, though, beaus you don’t have to figure out the indexing manually.