Feature/adjoint odes

ode-adjoint/linked-mass-flow-buggy.stan

@@ -0,0 +1,154 @@
+functions {
+
+  vector linked_mass_flow(real t, vector y,
+                          vector kt, vector e50, vector k12, vector k21) {
+
+    int num_main_states = num_elements(kt);
+    int num_states = 2 * num_main_states;
+
+    vector[num_states] dydt;
+    vector[num_main_states] ksat;
+
+    for (i in 1:num_main_states) {
+      int m = 2 * (i-1) + 1;  // main state
+      int a = m + 1;  // auxilary state
+      ksat[i] = kt[i] * y[m] / (y[m] + e50[i]);
+
+      dydt[m] = -ksat[i] * y[m] - k12[i] * y[m] + k21[i] * y[a];
+      dydt[a] = +k12[i] * y[m] - k21[i] * y[a];
+
+      if (i != 1) {
+        dydt[m] += ksat[i-1] * y[2 * (i - 1) + 1];
+      }
+    }
+    return dydt;
+
+  }
+
+  vector sample_vector_rng(real m, real s, int N) {
+    vector[N] sample;
+    for(i in 1:N)
+      sample[i] = normal_rng(m, s);
+    return sample;
+  }
+
+  vector[] simulate_mean(real t0, vector log_a0, real[] ts, int adjoint_integrator,
+                         real rel_tol, real abs_tol, int max_num_steps,
+                         vector log_kt, vector log_e50, vector log_k12, vector log_k21) {
+    int num_sim = num_elements(ts);
+    int num_states = num_elements(log_a0);
+    vector[num_states] y[num_sim];
+
+    if(adjoint_integrator) {
+      y = ode_adams_tol(linked_mass_flow, exp(log_a0), t0, ts, rel_tol, abs_tol, max_num_steps,
+                        exp(log_kt), exp(log_e50), exp(log_k12), exp(log_k21));
+    } else {
+      y = ode_bdf_tol(linked_mass_flow, exp(log_a0), t0, ts, rel_tol, abs_tol, max_num_steps,
+                      exp(log_kt), exp(log_e50), exp(log_k12), exp(log_k21));
+    }
+
+    return y;
+  }
+
+}
+data {
+  real<lower=0> rel_tol;
+  real<lower=0> abs_tol;
+  int<lower=0,upper=1> adjoint_integrator;
+  int<lower=1> max_num_steps;
+  int<lower=1> system_size;
+  int<lower=1> num_obs;
+  real<lower=0> sigma_sim;
+  real<lower=0> sigma_y;
+}
+transformed data {
+  int num_states = 2 * system_size;
+
+  vector[system_size] log_kt_ = sample_vector_rng(0.0, sigma_sim, system_size);
+  vector[system_size] log_e50_ = sample_vector_rng(0.0, sigma_sim, system_size);
+  vector[system_size] log_k12_ = sample_vector_rng(0.0, sigma_sim, system_size);
+  vector[system_size] log_k21_ = sample_vector_rng(0.0, sigma_sim, system_size);
+  vector[system_size] log_sigma_y_ = sample_vector_rng(0.0, sigma_y, system_size);
+  vector[num_states] log_a0_ = sample_vector_rng(0.0, sigma_sim, num_states);
+  vector[num_states] y_[num_obs];
+  real ts[num_obs];
+  real t0 = 0.0;
+
+  ts[1] = 1.0;
+  for(i in 2:num_obs) {
+    ts[i] = 1.5 * ts[i-1];
+  }
+
+  y_ = simulate_mean(t0, log_a0_, ts, 0, rel_tol, abs_tol, max_num_steps,
+                     log_kt_, log_e50_, log_k12_, log_k21_);
+
+  for(i in 1:num_obs) {
+    for(j in 1:system_size) {
+      int m = 2*(j-1) + 1;
+      y_[i,m] = lognormal_rng(log(y_[i,m]+1E-3), exp(log_sigma_y_[j]));
+    }
+  }
+
+  if(adjoint_integrator) {
+    print("Using bdf_adjoint integrator.");
+  } else {
+    print("Using bdf integrator.");
+  }
+  print("relative tolerance: ", rel_tol);
+  print("absolute tolerance: ", abs_tol);
+  print("maximum number of steps: ", max_num_steps);
+  print("number of time points: ", num_obs);
+  print("system size: ", system_size);
+  print("time points: ", ts);
+  print("y_: ", y_);
+}
+parameters {
+  vector[system_size] log_kt;
+  vector[system_size] log_e50;
+  vector[system_size] log_k12;
+  vector[system_size] log_k21;
+  vector[system_size] log_sigma_y;
+  vector[num_states] log_a0;
+}
+transformed parameters {
+}
+model {
+  vector[num_states] mu[num_obs] = simulate_mean(t0, log_a0, ts, adjoint_integrator, rel_tol, abs_tol, max_num_steps,
+                                                 log_kt, log_e50, log_k12, log_k21);
+
+  target += normal_lpdf(log_kt| 0.0, sigma_sim);
+  target += normal_lpdf(log_e50| 0.0, sigma_sim);
+  target += normal_lpdf(log_k12| 0.0, sigma_sim);
+  target += normal_lpdf(log_k21| 0.0, sigma_sim);
+  target += normal_lpdf(log_a0| 0.0, sigma_sim);
+  target += normal_lpdf(log_sigma_y| 0.0, sigma_y);
+
+  for(j in 1:system_size) {
+    int m = 2*(j-1) + 1;
+    target += lognormal_lpdf(to_vector(y_[:,m])| log(to_vector(mu[:,m])+1E-3), exp(log_sigma_y[j]));
+  }
+}
+generated quantities {
+  int rank_log_kt[system_size];
+  int rank_log_e50[system_size];
+  int rank_log_k12[system_size];
+  int rank_log_k21[system_size];
+  int rank_log_sigma_y[system_size];
+  int rank_log_a0[system_size];
+  vector[system_size] bias_log_kt = log_kt - log_kt_;
+  vector[system_size] bias_log_e50 = log_e50 - log_e50_;
+  vector[system_size] bias_log_k12 = log_k12 - log_k12_;
+  vector[system_size] bias_log_k21 = log_k21 - log_k21_;
+  vector[system_size] bias_log_sigma_y = log_sigma_y - log_sigma_y_;
+  vector[system_size] bias_log_a0 = log_a0 - log_a0_;
+
+  for(i in 1:system_size) {
+    rank_log_kt[i] = (log_kt[i] > log_kt_[i] ? 1 : 0);
+    rank_log_e50[i] = (log_e50[i] > log_e50_[i] ? 1 : 0);
+    rank_log_k12[i] = (log_k12[i] > log_k12_[i] ? 1 : 0);
+    rank_log_k21[i] = (log_k21[i] > log_k21_[i] ? 1 : 0);
+    rank_log_sigma_y[i] = (log_sigma_y[i] > log_sigma_y_[i] ? 1 : 0);
+    rank_log_a0[i] = (log_a0[i] > log_a0_[i] ? 1 : 0);
+  }
+
+}

ode-adjoint/linked-mass-flow.R

@@ -0,0 +1,84 @@
+##install.packages("cmdstanr", repos = c("https://mc-stan.org/r-packages/", getOption("repos")))
+
+library(cmdstanr)
+library(posterior)
+library(bayesplot)
+color_scheme_set("brightblue")
+
+set_cmdstan_path("~/work/cmdstan")
+
+stan_model_file <- "linked-mass-flow.stan"
+
+mod <- cmdstan_model(stan_model_file)
+
+base_data <- list(rel_tol=1e-6,
+                  abs_tol=1e-6,
+                  max_num_steps=10000,
+                  adjoint_integrator=0,
+                  system_size=2,
+                  num_obs=8,
+                  sigma_sim=2.0,
+                  sigma_y=0.2)
+
+run_benchmark <- function(..., num_iter=250, adapt_delta=0.8) {
+    fit_args <- list(...)
+    fit_data <- modifyList(base_data, fit_args)
+    print(str(fit_data))
+    fit <- mod$sample(
+                   data = fit_data,
+                   iter_warmup = num_iter,
+                   iter_sampling = num_iter,
+                   seed = 123,
+                   chains = 1,
+                   parallel_chains = 1,
+                   refresh = round(num_iter/5),
+                   save_warmup=TRUE,
+                   adapt_delta=adapt_delta
+                   ##verbose=TRUE
+               )
+    fit
+}
+
+bdf_fit <- run_benchmark(adjoint_integrator=0)
+adjoint_fit <- run_benchmark(adjoint_integrator=1)
+
+if(FALSE) {
+    adjoint_fit_2 <- run_benchmark(adjoint_integrator=1, adapt_delta=0.65)
+    adjoint_fit_3 <- run_benchmark(adjoint_integrator=1, adapt_delta=0.95)
+    adjoint_fit_4 <- run_benchmark(adjoint_integrator=1, abs_tol=1E-10, rel_tol=1E-10)
+}
+
+summarize_benchmark <- function(fit) {
+    sdiag_all  <- fit$sampler_diagnostics(inc_warmup=TRUE)
+    sdiag_main  <- fit$sampler_diagnostics(inc_warmup=FALSE)
+
+    num_divergent <- sum(sdiag_main[,,"divergent__"])
+    accept_stat <- mean(sdiag_main[,,"accept_stat__"])
+    treedepth <- mean(sdiag_main[,,"treedepth__"])
+    num_leaps <- sum(sdiag_all[,,"n_leapfrog__"])
+    fit_time <- fit$time()$total
+
+    ess_speed <- as.vector(fit$summary("lp__")[1,c("ess_bulk", "ess_tail")]) / fit_time
+    lp_est <- as.vector(fit$summary("lp__"))
+
+    list(num_leaps=num_leaps, time=fit_time, time_per_leap=1E3 * fit_time / num_leaps,
+         step_size=fit$metadata()$step_size_adaptation,
+         num_divergent=num_divergent, accept_stat=accept_stat, treedepth=treedepth,
+         ess_lp_speed=ess_speed, lp_est=lp_est)
+}
+
+mcmc_rhat(rhat = rhat(bdf_fit))  + yaxis_text(hjust = 0)
+
+str(summarize_benchmark(bdf_fit))
+str(summarize_benchmark(adjoint_fit))
+
+
+str(summarize_benchmark(adjoint_fit_2))
+str(summarize_benchmark(adjoint_fit_3))
+
+bdf_fit$metadata()
+
+bdf_fit$time()
+
+
+bdf_fit$summary("lp__")