Issues getting non-linear ODE model inference going, starting off with `ode_ckrk` function error message

Hi,
Transitioning from this previous thread, where I asked about arranging data types and dimensions to facilitate inference conditional on multidimensional ODE states, as I’m now looking to resolve model issues. My ensuing questions correspond to the following Stan code:

functions {

  // Temperature function for ODE forcing.
  real temp_func(real t, real temp_ref, real temp_rise) {
    return temp_ref + (temp_rise * t) / (80 * 24 * 365) + 10 * sin((2 * pi() / 24) * t) + 10 * sin((2 * pi() / (24 * 365)) * t);
  }

  // Exogenous SOC input function.
  real i_s_func(real t) {
    return 0.001 + 0.0005 * sin((2 * pi() / (24 * 365)) * t);
  }

  // Exogenous DOC input function.
  real i_d_func(real t) {
    return 0.0001 + 0.00005 * sin((2 * pi() / (24 * 365)) * t);
  }

  // Function for enforcing Arrhenius temperature dependency of ODE parameter.
  real arrhenius_temp(real input, real temp, real Ea, real temp_ref) {
    return input * exp(-Ea / 0.008314 * (1 / temp - 1 / temp_ref));
  }

  // Function for enforcing linear temperature dependency of ODE parameter.
  real linear_temp(real input, real temp, real Q, real temp_ref) {
    return input - Q * (temp - temp_ref);
  }

  /*
  AWB-ECA (equilibrium chemistry approximation) variant of AWB model.
  C[1] is soil organic carbon (SOC) density.
  C[2] is dissolved organic carbon (DOC) density.
  C[3] is microbial biomass carbon (MBC) density.
  C[4] is extracellular enzyme carbon (EEC) density.
  */

  vector AWB_ECA_ODE(real t, vector C,
                     real u_Q_ref, // Reference carbon use efficiency.
                     real Q, // Carbon use efficiency linear dependence negative slope.
                     real a_MSA, // AWB MBC-to-SOC transfer fraction.
                     real K_DE, // SOC decomposition K_m.
                     real K_UE, // DOC uptake K_m.
                     real V_DE_ref, // Reference SOC decomposition V_max.
                     real V_UE_ref, // Reference DOC uptake V_max.
                     real Ea_V_DE, // SOC V_max activation energy.
                     real Ea_V_UE, // DOC V_max activation energy.
                     real r_M, // MBC turnover rate.
                     real r_E, // Enzyme production rate.
                     real r_L, // Enzyme loss rate.
                     real temp_ref,
                     real temp_rise) {

    // Initiate exogenous input and forcing variables for future assignment.
    real temp;
    real i_s;
    real i_d;
    real u_Q; // Forced u_Q_ref.
    real V_DE; // Forced V_DE.
    real V_UE; // Forced V_UE.

    // Initiate dependent expressions.
    real F_S; // SOC decomposition
    real F_D; // DOC uptake

    // Assign input and forcing variables to appropriate value at time t.
    temp = temp_func(t, temp_ref, temp_rise); // x_r[1] is temp_ref 283.
    i_s = i_s_func(t);
    i_d = i_d_func(t);

    // Force temperature dependent parameters.
    u_Q = linear_temp(u_Q_ref, temp, Q, temp_ref);
    V_DE = arrhenius_temp(V_DE_ref, temp, Ea_V_DE, temp_ref);
    V_UE = arrhenius_temp(V_UE_ref, temp, Ea_V_UE, temp_ref);

    // Assign dependent expressions.
    F_S = V_DE * C[4] * C[1] / (K_DE + C[4] + C[1]);
    F_D = V_UE * C[3] * C[2] / (K_UE + C[3] + C[2]);

    // Initiate vector for storing derivatives.
    vector[4] dCdt;

    // Compute derivatives.
    dCdt[1] = i_s + a_MSA * r_M * C[3] - F_S;
    dCdt[2] = i_d + (1 - a_MSA) * r_M * C[3] + F_S + r_L * C[4] - F_D;
    dCdt[3] = u_Q * F_D * - (r_M + r_E) * C[3];
    dCdt[4] = r_E * C[3] - r_L * C[4];

    return dCdt;
  }

  // From https://discourse.mc-stan.org/t/rng-for-truncated-distributions/3122/12.
  real normal_lb_ub_rng(real mu, real sigma, real lb, real ub) {
      real p1 = normal_cdf(lb, mu, sigma);  // cdf with lower bound
      real p2 = normal_cdf(ub, mu, sigma);  // cdf with upper bound
      real u = uniform_rng(p1, p2);
      return (sigma * inv_Phi(u)) + mu;  // inverse cdf 
  }
}

data {
  int<lower=1> state_dim; // Number of state dimensions (4 for AWB).
  int<lower=1> N_t; // Number of observations.
  array[N_t] real<lower=0> ts; // Univariate array of observation time steps.
  array[state_dim] vector<lower=0>[N_t] y; // Multidimensional array of state observations bounded at 0. y in [state_dim, N_t] shape to facilitate likelihood sampling.
  real<lower=0> temp_ref; // Reference temperature for temperature forcing.
  real<lower=0> temp_rise; // Assumed increase in temperature (°C/K) over next 80 years.
  real<lower=0> prior_scale_factor; // Factor multiplying parameter means to obtain prior standard deviations.
  real<lower=0> obs_error_scale; // Observation noise factor multiplying observations of model output x_hat.
  vector<lower=0>[state_dim] x_hat0; // Initial ODE conditions.
  real<lower=0> u_Q_ref_mean;
  real<lower=0> Q_mean;
  real<lower=0> a_MSA_mean;
  real<lower=0> K_DE_mean;
  real<lower=0> K_UE_mean;
  real<lower=0> V_DE_ref_mean;
  real<lower=0> V_UE_ref_mean;
  real<lower=0> Ea_V_DE_mean;
  real<lower=0> Ea_V_UE_mean;
  real<lower=0> r_M_mean;
  real<lower=0> r_E_mean;
  real<lower=0> r_L_mean;
}

transformed data {
  real t0 = 0; // Initial time.
}

parameters {
  real<lower=0> u_Q_ref; // Reference carbon use efficiency.
  real<lower=0> Q; // Carbon use efficiency linear dependence negative slope.
  real<lower=0> a_MSA; // AWB MBC-to-SOC transfer fraction.
  real<lower=0> K_DE; // SOC decomposition K_m.
  real<lower=0> K_UE; // DOC uptake K_m.
  real<lower=0> V_DE_ref; // Reference SOC decomposition V_max.
  real<lower=0> V_UE_ref; // Reference DOC uptake V_max.
  real<lower=0> Ea_V_DE; // SOC V_max activation energy.
  real<lower=0> Ea_V_UE; // DOC V_max activation energy.
  real<lower=0> r_M; // MBC turnover rate.
  real<lower=0> r_E; // Enzyme production rate.
  real<lower=0> r_L; // Enzyme loss rate.
}

transformed parameters {
  array[N_t] vector<lower=0>[state_dim] x_hat_intmd = ode_ckrk(AWB_ECA_ODE, x_hat0, t0, ts, u_Q_ref, Q, a_MSA, K_DE, K_UE, V_DE_ref, V_UE_ref, Ea_V_DE, Ea_V_UE, r_M, r_E, r_L, temp_ref, temp_rise);
  // Transform model output to match observations y in shape, [state_dim, N_t].
  array[state_dim] vector<lower=0>[N_t] x_hat;
  for (i in 1:N_t) {
    for (j in 1:state_dim) {
      x_hat[j, i] = x_hat_intmd[i, j];
    }
  }
}

model {
  u_Q_ref ~ normal(u_Q_ref_mean, u_Q_ref_mean * prior_scale_factor) T[0, 1];
  Q ~ normal(Q_mean, Q_mean * prior_scale_factor) T[0, 0.1];
  a_MSA ~ normal(a_MSA_mean, a_MSA_mean * prior_scale_factor) T[0, 1];
  K_DE ~ normal(K_DE_mean, K_DE_mean * prior_scale_factor) T[0, 5000];
  K_UE ~ normal(K_UE_mean, K_UE_mean * prior_scale_factor) T[0, 1];
  V_DE_ref ~ normal(V_DE_ref_mean, V_DE_ref_mean * prior_scale_factor) T[0, 1];
  V_UE_ref ~ normal(V_UE_ref_mean, V_UE_ref_mean * prior_scale_factor) T[0, 0.1];
  Ea_V_DE ~ normal(Ea_V_DE_mean, Ea_V_DE_mean * prior_scale_factor) T[5, 80];
  Ea_V_UE ~ normal(Ea_V_UE_mean, Ea_V_UE_mean * prior_scale_factor) T[5, 80];
  r_M ~ normal(r_M_mean, r_M_mean * prior_scale_factor) T[0, 0.1];
  r_E ~ normal(r_E_mean, r_E_mean * prior_scale_factor) T[0, 0.1];
  r_L ~ normal(r_L_mean, r_L_mean * prior_scale_factor) T[0, 0.1];

  // Likelihood evaluation.
  for (i in 1:state_dim) {
    y[i,] ~ normal(x_hat[i,], obs_error_scale * x_hat[i,]);
  }
}

generated quantities {
  array[N_t] vector<lower=0>[state_dim] x_hat_post_pred;
  array[state_dim] vector<lower=0>[N_t] x_hat_post_pred_intmd;  
  array[N_t, state_dim] real<lower=0> y_hat_post_pred;
  x_hat_post_pred_intmd = ode_ckrk(AWB_ECA_ODE, x_hat0, t0, ts, u_Q_ref, Q, a_MSA, K_DE, K_UE, V_DE_ref, V_UE_ref, Ea_V_DE, Ea_V_UE, r_M, r_E, r_L, temp_ref, temp_rise);
  // Transform posterior predictive model output to match observations y in dimensions, [state_dim, N_t].
  for (i in 1:N_t) {
    for (j in 1:state_dim) {
      x_hat_post_pred[j, i] = x_hat_post_pred_intmd[i, j];
    }
  }
  for (i in 1:state_dim) {
    y_hat_post_pred[i,] = normal_rng(x_hat_post_pred[i,], obs_error_scale * x_hat_post_pred[i,]);
  }
  real u_Q_ref_post = normal_lb_ub_rng(u_Q_ref_mean, u_Q_ref_mean * prior_scale_factor, 0, 1);
  real Q_post = normal_lb_ub_rng(Q_mean, Q_mean * prior_scale_factor, 0, 0.1);
  real a_MSA_post = normal_lb_ub_rng(a_MSA_mean, a_MSA_mean * prior_scale_factor, 0, 1);
  real K_DE_post = normal_lb_ub_rng(K_DE_mean, K_DE_mean * prior_scale_factor, 0, 5000);
  real K_UE_post = normal_lb_ub_rng(K_UE_mean, K_UE_mean * prior_scale_factor, 0, 1);
  real V_DE_ref_post = normal_lb_ub_rng(V_DE_ref_mean, V_DE_ref_mean * prior_scale_factor, 0, 1);
  real V_UE_ref_post = normal_lb_ub_rng(V_UE_ref_mean, V_UE_ref_mean * prior_scale_factor, 0, 0.1);
  real Ea_V_DE_post = normal_lb_ub_rng(Ea_V_DE_mean, Ea_V_DE_mean * prior_scale_factor, 5, 80);
  real Ea_V_UE_post = normal_lb_ub_rng(Ea_V_UE_mean, Ea_V_UE_mean * prior_scale_factor, 5, 80);
  real r_M_post = normal_lb_ub_rng(r_M_mean, r_M_mean * prior_scale_factor, 0, 0.1);
  real r_E_post = normal_lb_ub_rng(r_E_mean, r_E_mean * prior_scale_factor, 0, 0.1);
  real r_L_post = normal_lb_ub_rng(r_L_mean, r_L_mean * prior_scale_factor, 0, 0.1);
}

and R code for interfacing with CmdStanR:

library(cmdstanr)
library(posterior)
library(bayesplot)
library(tidyverse)

rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

color_scheme_set("viridisA")

#Data to be passed to Stan.
state_dim <- 4
temp_ref <- 283
temp_rise <- 5 #High estimate of 5 celsius temperature rise by 2100.
prior_scale_factor <- 0.25
obs_error_scale <- 0.1
obs_every <- 10 #Observations every 10 hours.
t <- 2000 #Total time span of ODE simulation.
dt <- 0.1 #Euler time step size corresponding to data generation. 
x_hat0 <- c(78.81405796, 5.66343552, 2.2209765, 1.19214226) #Originally sampled values used for Euler-Maruyama solution.
y_full <- read_csv('generated_data/SAWB-ECA-SS_no_CO2_trunc_short_2021_12_20_18_20_sample_y_t_2000_dt_0-01_sd_scale_0-25.csv')
y <- y_full %>% filter(hour <= t)
ts <- y$hour
N_t <- length(ts)
y <- y %>% select(-hour)
#y <- split(y, 1:nrow(y)) #Convert data observations to list of rows to correspond to Stan's array of vectors type.
y <- as.list(y) #Convert data observations to list of columns to correspond to Stan's array of vectors type.

#Parameter prior means
u_Q_ref_mean <- 0.25
Q_mean <- 0.001
a_MSA_mean <- 0.5
K_DE_mean <- 1000
K_UE_mean <- 0.1
V_DE_ref_mean <- 0.11
V_UE_ref_mean <- 0.0044
Ea_V_DE_mean <- 40
Ea_V_UE_mean <- 30
r_M_mean <- 0.0001667
r_E_mean <- 0.0002
r_L_mean <- 0.0004

data_list = list(
    state_dim = state_dim,
    N_t = N_t,
    ts = ts,
    y = y,
    temp_ref = temp_ref,
    temp_rise = temp_rise,
    prior_scale_factor = prior_scale_factor,
    obs_error_scale = obs_error_scale,
    x_hat0 = x_hat0,
    u_Q_ref_mean = u_Q_ref_mean,
    Q_mean = Q_mean,
    a_MSA_mean = a_MSA_mean,
    K_DE_mean = K_DE_mean,
    K_UE_mean = K_UE_mean,
    V_DE_ref_mean = V_DE_ref_mean,
    V_UE_ref_mean = V_UE_ref_mean,
    Ea_V_DE_mean = Ea_V_DE_mean,
    Ea_V_UE_mean = Ea_V_UE_mean,
    r_M_mean = r_M_mean,
    r_E_mean = r_E_mean,
    r_L_mean = r_L_mean
    )

file_path <- 'AWB-ECA_SS_no_CO2_cont_time.stan' #Read in Stan model code.
lines <- readLines(file_path, encoding = "ASCII")
for (n in 1:length(lines)) cat(lines[n],'\n')
model <- cmdstan_model(file_path)

AWB_ECA_stan_fit <- model$sample(data = data_list, iter_sampling = 1000, iter_warmup = 500, refresh = 10, chains = 3, seed = 123)

I am sure there are issues I’ve missed, but the above model presently compiles using both CmdStanR and PyStan. However, a problem arises after I initiate the sampling in R with

AWB_ECA_stan_fit <- model$sample(data = data_list, iter_sampling = 1000, iter_warmup = 500, refresh = 10, chains = 3, seed = 123)

which triggers the error message:

Chain 3 Exception: ode_ckrk: initial time is 0, but must be less than 0.000000 (in '/var/folders/mr/47dcvttx6_5dqt_p2ktdplc80000gn/T/RtmpxUsE9T/model-11ca97a709661.stan', line 143, column 2 to column 197)

The error message applies to all chains 1 - 3 that I used. It doesn’t make sense to me that the initial time would have to be negative, so it would seem that there is an upstream issue?

I will also try the other ODE solvers and report on the outcome. Thank you Stan team and forum users once more for helping me debug and diagnose these issues.

Edit: I do get the same error with ode_adams,

Chain 3 Exception: ode_adams: initial time is 0, but must be less than 0.000000 (in '/var/folders/mr/47dcvttx6_5dqt_p2ktdplc80000gn/T/RtmpeuwExU/model-134241e50d680.stan', line 143, column 2 to column 198)

so it would indeed seem like there is some upstream issue.

This is probably because your output time points include 0 but you have also set the initial time (from which integration starts) to 0. CVODES requires that the initial time needs to be strictly smaller than any output time point.

1 Like

It’s because you already know that the solution at t0 is x_hat0 so you shouldn’t include t0 in the output time points t.

1 Like

Ah, got it. Thank you!

Some divergences, of course, but presently have code that runs, so for the sake of helping others working with ODEs in the future with one more example to potentially refer to, sharing the present versions of my soil biogeochemical inference Stan and R CmdStanR interfacing code.

Stan code:

functions {

  // Temperature function for ODE forcing.
  real temp_func(real t, real temp_ref, real temp_rise) {
    return temp_ref + (temp_rise * t) / (80 * 24 * 365) + 10 * sin((2 * pi() / 24) * t) + 10 * sin((2 * pi() / (24 * 365)) * t);
  }

  // Exogenous SOC input function.
  real i_s_func(real t) {
    return 0.001 + 0.0005 * sin((2 * pi() / (24 * 365)) * t);
  }

  // Exogenous DOC input function.
  real i_d_func(real t) {
    return 0.0001 + 0.00005 * sin((2 * pi() / (24 * 365)) * t);
  }

  // Function for enforcing Arrhenius temperature dependency of ODE parameter.
  real arrhenius_temp(real input, real temp, real Ea, real temp_ref) {
    return input * exp(-Ea / 0.008314 * (1 / temp - 1 / temp_ref));
  }

  // Function for enforcing linear temperature dependency of ODE parameter.
  real linear_temp(real input, real temp, real Q, real temp_ref) {
    return input - Q * (temp - temp_ref);
  }

  /*
  AWB-ECA (equilibrium chemistry approximation) variant of AWB model.
  C[1] is soil organic carbon (SOC) density.
  C[2] is dissolved organic carbon (DOC) density.
  C[3] is microbial biomass carbon (MBC) density.
  C[4] is extracellular enzyme carbon (EEC) density.
  */

  vector AWB_ECA_ODE(real t, vector C,
                     real u_Q_ref, // Reference carbon use efficiency.
                     real Q, // Carbon use efficiency linear dependence negative slope.
                     real a_MSA, // AWB MBC-to-SOC transfer fraction.
                     real K_DE, // SOC decomposition K_m.
                     real K_UE, // DOC uptake K_m.
                     real V_DE_ref, // Reference SOC decomposition V_max.
                     real V_UE_ref, // Reference DOC uptake V_max.
                     real Ea_V_DE, // SOC V_max activation energy.
                     real Ea_V_UE, // DOC V_max activation energy.
                     real r_M, // MBC turnover rate.
                     real r_E, // Enzyme production rate.
                     real r_L, // Enzyme loss rate.
                     real temp_ref,
                     real temp_rise) {

    // Initiate exogenous input and forcing variables for future assignment.
    real temp;
    real i_s;
    real i_d;
    real u_Q; // Forced u_Q_ref.
    real V_DE; // Forced V_DE.
    real V_UE; // Forced V_UE.

    // Initiate dependent expressions.
    real F_S; // SOC decomposition
    real F_D; // DOC uptake

    // Assign input and forcing variables to appropriate value at time t.
    temp = temp_func(t, temp_ref, temp_rise); // x_r[1] is temp_ref 283.
    i_s = i_s_func(t);
    i_d = i_d_func(t);

    // Force temperature dependent parameters.
    u_Q = linear_temp(u_Q_ref, temp, Q, temp_ref);
    V_DE = arrhenius_temp(V_DE_ref, temp, Ea_V_DE, temp_ref);
    V_UE = arrhenius_temp(V_UE_ref, temp, Ea_V_UE, temp_ref);

    // Assign dependent expressions.
    F_S = V_DE * C[4] * C[1] / (K_DE + C[4] + C[1]);
    F_D = V_UE * C[3] * C[2] / (K_UE + C[3] + C[2]);

    // Initiate vector for storing derivatives.
    vector[4] dCdt;

    // Compute derivatives.
    dCdt[1] = i_s + a_MSA * r_M * C[3] - F_S;
    dCdt[2] = i_d + (1 - a_MSA) * r_M * C[3] + F_S + r_L * C[4] - F_D;
    dCdt[3] = u_Q * F_D * - (r_M + r_E) * C[3];
    dCdt[4] = r_E * C[3] - r_L * C[4];
    return dCdt;
  }

  // From https://discourse.mc-stan.org/t/rng-for-truncated-distributions/3122/12.
  real normal_lb_ub_rng(real mu, real sigma, real lb, real ub) {
      real p1 = normal_cdf(lb, mu, sigma);  // cdf with lower bound
      real p2 = normal_cdf(ub, mu, sigma);  // cdf with upper bound
      real u = uniform_rng(p1, p2);
      return (sigma * inv_Phi(u)) + mu;  // inverse cdf 
  }
}

data {
  int<lower=1> state_dim; // Number of state dimensions (4 for AWB).
  int<lower=1> N_t; // Number of observations.
  array[N_t] real<lower=0> ts; // Univariate array of observation time steps.
  array[state_dim] vector<lower=0>[N_t] y; // Multidimensional array of state observations bounded at 0. y in [state_dim, N_t] shape to facilitate likelihood sampling.
  real<lower=0> temp_ref; // Reference temperature for temperature forcing.
  real<lower=0> temp_rise; // Assumed increase in temperature (°C/K) over next 80 years.
  real<lower=0> prior_scale_factor; // Factor multiplying parameter means to obtain prior standard deviations.
  real<lower=0> obs_error_scale; // Observation noise factor multiplying observations of model output x_hat.
  vector<lower=0>[state_dim] x_hat0; // Initial ODE conditions.
  // [1] is prior mean, [2] is prior lower bound, [3] is prior upper bound.
  array[3] real<lower=0> u_Q_ref_prior_dist_params;
  array[3] real<lower=0> Q_prior_dist_params;
  array[3] real<lower=0> a_MSA_prior_dist_params;
  array[3] real<lower=0> K_DE_prior_dist_params;
  array[3] real<lower=0> K_UE_prior_dist_params;
  array[3] real<lower=0> V_DE_ref_prior_dist_params;
  array[3] real<lower=0> V_UE_ref_prior_dist_params;
  array[3] real<lower=0> Ea_V_DE_prior_dist_params;
  array[3] real<lower=0> Ea_V_UE_prior_dist_params;
  array[3] real<lower=0> r_M_prior_dist_params;
  array[3] real<lower=0> r_E_prior_dist_params;
  array[3] real<lower=0> r_L_prior_dist_params;
}

transformed data {
  real t0 = 0; // Initial time.
}

parameters {
  real<lower = u_Q_ref_prior_dist_params[2], upper = u_Q_ref_prior_dist_params[3]> u_Q_ref; // Reference carbon use efficiency.
  real<lower = Q_prior_dist_params[2], upper = Q_prior_dist_params[3]> Q; // Carbon use efficiency linear dependence negative slope.
  real<lower = a_MSA_prior_dist_params[2], upper = a_MSA_prior_dist_params[3]> a_MSA; // AWB MBC-to-SOC transfer fraction.
  real<lower = K_DE_prior_dist_params[2], upper = K_DE_prior_dist_params[3]> K_DE; // SOC decomposition K_m.
  real<lower = K_UE_prior_dist_params[2], upper = K_UE_prior_dist_params[3]> K_UE; // DOC uptake K_m.
  real<lower = V_DE_ref_prior_dist_params[2], upper = V_DE_ref_prior_dist_params[3]> V_DE_ref; // Reference SOC decomposition V_max.
  real<lower = V_UE_ref_prior_dist_params[2], upper = V_UE_ref_prior_dist_params[3]> V_UE_ref; // Reference DOC uptake V_max.
  real<lower = Ea_V_DE_prior_dist_params[2], upper = Ea_V_DE_prior_dist_params[3]> Ea_V_DE; // SOC V_max activation energy.
  real<lower = Ea_V_UE_prior_dist_params[2], upper = Ea_V_UE_prior_dist_params[3]> Ea_V_UE; // DOC V_max activation energy.
  real<lower = r_M_prior_dist_params[2], upper = r_M_prior_dist_params[3]> r_M; // MBC turnover rate.
  real<lower = r_E_prior_dist_params[2], upper = r_E_prior_dist_params[3]> r_E; // Enzyme production rate.
  real<lower = r_L_prior_dist_params[2], upper = r_L_prior_dist_params[3]> r_L; // Enzyme loss rate.
}

transformed parameters {
  array[N_t] vector<lower=0>[state_dim] x_hat_intmd = ode_ckrk(AWB_ECA_ODE, x_hat0, t0, ts, u_Q_ref, Q, a_MSA, K_DE, K_UE, V_DE_ref, V_UE_ref, Ea_V_DE, Ea_V_UE, r_M, r_E, r_L, temp_ref, temp_rise);
  // Transform model output to match observations y in shape, [state_dim, N_t].
  array[state_dim] vector<lower=0>[N_t] x_hat;
  for (i in 1:N_t) {
    for (j in 1:state_dim) {
      x_hat[j, i] = x_hat_intmd[i, j];
    }
  }
}

model {
  u_Q_ref ~ normal(u_Q_ref_prior_dist_params[1], u_Q_ref_prior_dist_params[1] * prior_scale_factor) T[u_Q_ref_prior_dist_params[2], u_Q_ref_prior_dist_params[3]];
  Q ~ normal(Q_prior_dist_params[1], Q_prior_dist_params[1] * prior_scale_factor) T[Q_prior_dist_params[2], Q_prior_dist_params[3]];
  a_MSA ~ normal(a_MSA_prior_dist_params[1], a_MSA_prior_dist_params[1] * prior_scale_factor) T[a_MSA_prior_dist_params[2], a_MSA_prior_dist_params[3]];
  K_DE ~ normal(K_DE_prior_dist_params[1], K_DE_prior_dist_params[1] * prior_scale_factor) T[K_DE_prior_dist_params[2], K_DE_prior_dist_params[3]];
  K_UE ~ normal(K_UE_prior_dist_params[1], K_UE_prior_dist_params[1] * prior_scale_factor) T[K_UE_prior_dist_params[2], K_UE_prior_dist_params[3]];
  V_DE_ref ~ normal(V_DE_ref_prior_dist_params[1], V_DE_ref_prior_dist_params[1] * prior_scale_factor) T[V_DE_ref_prior_dist_params[2], V_DE_ref_prior_dist_params[3]];
  V_UE_ref ~ normal(V_UE_ref_prior_dist_params[1], V_UE_ref_prior_dist_params[1] * prior_scale_factor) T[V_UE_ref_prior_dist_params[2], V_UE_ref_prior_dist_params[3]];
  Ea_V_DE ~ normal(Ea_V_DE_prior_dist_params[1], Ea_V_DE_prior_dist_params[1] * prior_scale_factor) T[Ea_V_DE_prior_dist_params[2], Ea_V_DE_prior_dist_params[3]];
  Ea_V_UE ~ normal(Ea_V_UE_prior_dist_params[1], Ea_V_UE_prior_dist_params[1] * prior_scale_factor) T[Ea_V_UE_prior_dist_params[2], Ea_V_UE_prior_dist_params[3]];
  r_M ~ normal(r_M_prior_dist_params[1], r_M_prior_dist_params[1] * prior_scale_factor) T[r_M_prior_dist_params[2], r_M_prior_dist_params[3]];
  r_E ~ normal(r_E_prior_dist_params[1], r_E_prior_dist_params[1] * prior_scale_factor) T[r_E_prior_dist_params[2], r_E_prior_dist_params[3]];
  r_L ~ normal(r_L_prior_dist_params[1], r_L_prior_dist_params[1] * prior_scale_factor) T[r_L_prior_dist_params[2], r_L_prior_dist_params[3]];

  // Likelihood evaluation.
  for (i in 1:state_dim) {
    y[i,] ~ normal(x_hat[i,], obs_error_scale * x_hat[i,]);
  }
}

generated quantities {
    print("Sampled theta values: ", "u_Q_ref = ", u_Q_ref, ", Q = ", Q, ", a_MSA = ", a_MSA, ", K_DE = ", K_DE, ", K_UE = ", K_UE, ", V_DE_ref = ", V_DE_ref, ", V_UE_ref = ", V_UE_ref, ", Ea_V_DE = ", Ea_V_DE, ", Ea_V_UE = ", Ea_V_UE, ", r_M = ", r_M, ", r_E = ", r_E, ", r_L = ", r_L);
  array[N_t] vector<lower=0>[state_dim] x_hat_post_pred_intmd;
  array[state_dim] vector<lower=0>[N_t] x_hat_post_pred;  
  array[state_dim, N_t] real<lower=0> y_hat_post_pred;
  x_hat_post_pred_intmd = ode_ckrk(AWB_ECA_ODE, x_hat0, t0, ts, u_Q_ref, Q, a_MSA, K_DE, K_UE, V_DE_ref, V_UE_ref, Ea_V_DE, Ea_V_UE, r_M, r_E, r_L, temp_ref, temp_rise);
  // Transform posterior predictive model output to match observations y in dimensions, [state_dim, N_t].
  for (i in 1:N_t) {
    for (j in 1:state_dim) {
      x_hat_post_pred[j, i] = x_hat_post_pred_intmd[i, j];
    }
  }
  // Add observation noise to posterior predictive model output.
  for (i in 1:state_dim) {
    y_hat_post_pred[i,] = normal_rng(x_hat_post_pred[i,], obs_error_scale * x_hat_post_pred[i,]);
  }
  print("Posterior predictive ODE output: ", y_hat_post_pred);  
  real u_Q_ref_prior_pred = normal_lb_ub_rng(u_Q_ref_prior_dist_params[1], u_Q_ref_prior_dist_params[1] * prior_scale_factor, u_Q_ref_prior_dist_params[2], u_Q_ref_prior_dist_params[3]);
  real Q_prior_pred = normal_lb_ub_rng(Q_prior_dist_params[1], Q_prior_dist_params[1] * prior_scale_factor, Q_prior_dist_params[2], Q_prior_dist_params[3]);
  real a_MSA_prior_pred = normal_lb_ub_rng(a_MSA_prior_dist_params[1], a_MSA_prior_dist_params[1] * prior_scale_factor, a_MSA_prior_dist_params[2], a_MSA_prior_dist_params[3]);
  real K_DE_prior_pred = normal_lb_ub_rng(K_DE_prior_dist_params[1], K_DE_prior_dist_params[1] * prior_scale_factor, K_DE_prior_dist_params[2], K_DE_prior_dist_params[3]);
  real K_UE_prior_pred = normal_lb_ub_rng(K_UE_prior_dist_params[1], K_UE_prior_dist_params[1] * prior_scale_factor, K_UE_prior_dist_params[2], K_UE_prior_dist_params[3]);
  real V_DE_ref_prior_pred = normal_lb_ub_rng(V_DE_ref_prior_dist_params[1], V_DE_ref_prior_dist_params[1] * prior_scale_factor, V_DE_ref_prior_dist_params[2], V_DE_ref_prior_dist_params[3]);
  real V_UE_ref_prior_pred = normal_lb_ub_rng(V_UE_ref_prior_dist_params[1], V_UE_ref_prior_dist_params[1] * prior_scale_factor, V_UE_ref_prior_dist_params[2], V_UE_ref_prior_dist_params[3]);
  real Ea_V_DE_prior_pred = normal_lb_ub_rng(Ea_V_DE_prior_dist_params[1], Ea_V_DE_prior_dist_params[1] * prior_scale_factor, Ea_V_DE_prior_dist_params[2], Ea_V_DE_prior_dist_params[3]);
  real Ea_V_UE_prior_pred = normal_lb_ub_rng(Ea_V_UE_prior_dist_params[1], Ea_V_UE_prior_dist_params[1] * prior_scale_factor, Ea_V_UE_prior_dist_params[2], Ea_V_UE_prior_dist_params[3]);
  real r_M_prior_pred = normal_lb_ub_rng(r_M_prior_dist_params[1], r_M_prior_dist_params[1] * prior_scale_factor, r_M_prior_dist_params[2], r_M_prior_dist_params[3]);
  real r_E_prior_pred = normal_lb_ub_rng(r_E_prior_dist_params[1], r_E_prior_dist_params[1] * prior_scale_factor, r_E_prior_dist_params[2], r_E_prior_dist_params[3]);
  real r_L_prior_pred = normal_lb_ub_rng(r_L_prior_dist_params[1], r_L_prior_dist_params[1] * prior_scale_factor, r_L_prior_dist_params[2], r_L_prior_dist_params[3]);
}

R code:

library(cmdstanr)
library(posterior)
library(bayesplot)
library(tidyverse)

rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

color_scheme_set("viridisA")

#Data to be passed to Stan.
state_dim <- 4
temp_ref <- 283
temp_rise <- 5 #High estimate of 5 celsius temperature rise by 2100.
prior_scale_factor <- 0.25
obs_error_scale <- 0.1
obs_every <- 10 #Observations every 10 hours.
t <- 2000 #Total time span of ODE simulation.
dt <- 0.1 #Euler time step size corresponding to data generation. 
x_hat0 <- c(78.81405796, 5.66343552, 2.2209765, 1.19214226) #Originally sampled values used for Euler-Maruyama solution.
y_full <- read_csv('generated_data/SAWB-ECA-SS_no_CO2_trunc_short_2021_12_20_18_20_sample_y_t_2000_dt_0-01_sd_scale_0-25.csv')
y <- y_full %>% filter(hour <= t) %>% tail(-1)
ts <- y$hour
N_t <- length(ts)
y <- y %>% select(-hour)
#y <- split(y, 1:nrow(y)) #Convert data observations to list of rows to correspond to Stan's array of vectors type.
y <- as.list(y) #Convert data observations to list of columns to correspond to Stan's array of vectors type.

#Parameter prior means
u_Q_ref_prior_dist_params <- c(0.25, 0, 1)
Q_prior_dist_params <- c(0.001, 0, 0.1)
a_MSA_prior_dist_params <- c(0.5, 0, 1)
K_DE_prior_dist_params <- c(1000, 0, 5000)
K_UE_prior_dist_params <- c(0.1, 0, 1)
V_DE_ref_prior_dist_params <- c(0.11, 0, 1)
V_UE_ref_prior_dist_params <- c(0.0044, 0, 0.1)
Ea_V_DE_prior_dist_params <- c(40, 5, 80)
Ea_V_UE_prior_dist_params <- c(30, 5, 80)
r_M_prior_dist_params <- c(0.0001667, 0, 0.1)
r_E_prior_dist_params <- c(0.0002, 0, 0.1)
r_L_prior_dist_params <- c(0.0004, 0, 0.1)

data_list = list(
    state_dim = state_dim,
    N_t = N_t,
    ts = ts,
    y = y,
    temp_ref = temp_ref,
    temp_rise = temp_rise,
    prior_scale_factor = prior_scale_factor,
    obs_error_scale = obs_error_scale,
    x_hat0 = x_hat0,
    u_Q_ref_prior_dist_params = u_Q_ref_prior_dist_params,
    Q_prior_dist_params = Q_prior_dist_params,
    a_MSA_prior_dist_params = a_MSA_prior_dist_params,
    K_DE_prior_dist_params = K_DE_prior_dist_params,
    K_UE_prior_dist_params = K_UE_prior_dist_params,
    V_DE_ref_prior_dist_params = V_DE_ref_prior_dist_params,
    V_UE_ref_prior_dist_params = V_UE_ref_prior_dist_params,
    Ea_V_DE_prior_dist_params = Ea_V_DE_prior_dist_params,
    Ea_V_UE_prior_dist_params = Ea_V_UE_prior_dist_params,
    r_M_prior_dist_params = r_M_prior_dist_params,
    r_E_prior_dist_params = r_E_prior_dist_params,
    r_L_prior_dist_params = r_L_prior_dist_params
    )

file_path <- 'AWB-ECA_SS_no_CO2_cont_time.stan' #Read in Stan model code.
lines <- readLines(file_path, encoding = "ASCII")
for (n in 1:length(lines)) cat(lines[n],'\n')
model <- cmdstan_model(file_path)

AWB_ECA_stan_fit <- model$sample(data = data_list, iter_sampling = 20, iter_warmup = 10, refresh = 1, chains = 3, parallel_chains = 3, seed = 123, adapt_delta = 0.95)

The Stan model presently calls the ode_ckrk function because this non-linear AWB-ECA system can be relatively unstable, but I found the following total run times using chains = 3, iter_sampling = 20, and iter_warmup = 10:

  • ode_bdf: 59.9 seconds
  • ode_adams: 38.6 seconds
  • ode_ckrk: 6.5 seconds
  • ode_rk45: 4.8 seconds

For the 60 total samples, there were 22 divergences detected for all of the ODE solvers. I’m proceeding with ckrk for now, despite rk45 having a substantial speed advantage, just in case ckrk can navigate some of the stiffer regimes in this non-linear AWB model slightly better (though this may just be burning carbon in the end). Still, ckrk is much closer to rk45 in speed than bdf and adams relatively speaking, so the trade-off for speed is hopefully not that bad.

1 Like