In a recent project I was missing a look-up functionality in stan, meaning I wanted to define some problem-specific function (of one variable) and use it in my model specification. An example would be if my model included some material property that depends on another variable in the model. Defining and deploying functions like this is doable in R, python, Matlab, C++, etc, so I was a little surprised to not find it in Stan as well. Should it maybe be added? (How and where…?)
After some procrastrination, I made ChatGPT code up an example in the form of PCHIP interpolation (see, e.g., PCHIP interpolation in boost) as stan code, and a dummy model to demonstrate it.
y_i = V_x(a+bx_i)+\sigma \xi_i,\; \xi_i \sim N(0,1)\; iid.
Here, V_x is the function to tabulate. For my synthetic example I took V_x(x)=0.2x+x^5.
require(tidyverse)
library(cmdstanr)
require(posterior)
# generate tabular values for interpolation
Vfun <- \(x)(0.2*x+x^5)
Vx_data <- tibble( x=seq(-3,3,0.1) ) |> mutate(Vx=Vfun(x))
# synthetic data
a <- 0.2
b <- 1.5
sig <- 0.25
set.seed(123) # reproducible random numbers
xi <- runif(30)
yi <- Vfun(a+b*xi)+sig*rnorm(length(xi))
train_data <- tibble( x=xi, y=yi)
stan_data2 <- list(
N = nrow(train_data),
x = train_data$x,
y = train_data$y,
K = nrow(Vx_data),
xk = Vx_data$x,
Vk = Vx_data$Vx
)
mod2 <- cmdstan_model("Vx_pchip_chatgpt_linexp.stan")
fit2 <- mod2$sample(
data = stan_data2,
chains = 4,
parallel_chains = 4,
iter_warmup = 1000,
iter_sampling = 2000,
thin=20,
seed = 123,
refresh = 0
)
fit2 |>
as_draws() |>
subset_draws(c("a","b","sigma")) |>
summarise_draws("mean", "mcse_mean", "sd","mcse_sd", "rhat")
where the stan model Vx_pchip_chatgpt_linexp.stan is
functions {
#include pchip.stan
}
data {
int<lower=1> N; // data used for inference
vector[N] x;
vector[N] y;
int<lower=2> K; // data used to parameterise the lookup table
vector[K] xk;
vector[K] Vk;
}
transformed data {
// xk and Vk are data, so the PCHIP representation can be computed once
// to speed up evaluation.
matrix[K - 1, 4] W = pchip_setup(xk, Vk);
}
parameters {
real a;
real b;
real<lower=0> sigma;
}
transformed parameters {
vector[N] mu;
for (n in 1:N) {
real z = a + b * x[n];
mu[n] = pchip_eval(z, xk, W);
}
}
model {
a ~ normal(0,1);
b ~ normal(0,1);
sigma ~ gamma(2,20);
y ~ normal(mu, sigma);
}
The pchip_setup() and pchip_eval() functions and defined in the imported pchip.stan :
// Fritsch-Carlson/PCHIP slopes at the knots.
vector pchip_slopes(vector x, vector y) {
int K = num_elements(x);
vector[K] d;
vector[K - 1] h;
vector[K - 1] delta;
for (i in 1:(K - 1)) {
h[i] = x[i + 1] - x[i];
delta[i] = (y[i + 1] - y[i]) / h[i];
}
if (K == 2) {
d[1] = delta[1];
d[2] = delta[1];
return d;
}
// Interior slopes: weighted harmonic mean when adjacent secants
// have the same sign; zero otherwise.
for (i in 2:(K - 1)) {
if (delta[i - 1] * delta[i] <= 0) {
d[i] = 0;
} else {
real w1 = 2 * h[i] + h[i - 1];
real w2 = h[i] + 2 * h[i - 1];
d[i] = (w1 + w2) /
(w1 / delta[i - 1] + w2 / delta[i]);
}
}
// One-sided endpoint slopes, with the usual PCHIP limiting rule.
{
real d1 = ((2 * h[1] + h[2]) * delta[1] - h[1] * delta[2]) /
(h[1] + h[2]);
if (d1 * delta[1] <= 0)
d[1] = 0;
else if ((delta[1] * delta[2] < 0) && (abs(d1) > 3 * abs(delta[1])))
d[1] = 3 * delta[1];
else
d[1] = d1;
}
{
int i = K - 1;
real dK = ((2 * h[i] + h[i - 1]) * delta[i] - h[i] * delta[i - 1]) /
(h[i] + h[i - 1]);
if (dK * delta[i] <= 0)
d[K] = 0;
else if ((delta[i] * delta[i - 1] < 0) && (abs(dK) > 3 * abs(delta[i])))
d[K] = 3 * delta[i];
else
d[K] = dK;
}
return d;
}
// Precompute cubic coefficients for each interval.
// Row i contains [a, b, c, d] such that
// f(x) = a + b*t + c*t^2 + d*t^3, t = x - xk[i].
matrix pchip_setup(vector xk, vector Vk) {
int K = num_elements(xk);
vector[K] slope = pchip_slopes(xk, Vk);
matrix[K - 1, 4] W;
for (i in 1:(K - 1)) {
real h = xk[i + 1] - xk[i];
real dy = Vk[i + 1] - Vk[i];
real h2 = h * h;
real h3 = h2 * h;
W[i, 1] = Vk[i];
W[i, 2] = slope[i];
W[i, 3] = 3 * dy / h2 - (2 * slope[i] + slope[i + 1]) / h;
W[i, 4] = -2 * dy / h3 + (slope[i] + slope[i + 1]) / h2;
}
return W;
}
real pchip_eval(real x, vector xk, matrix W) {
int K = num_elements(xk);
int i;
int lo = 1;
int hi = K;
real t;
// Linear extrapolation using the left endpoint slope.
if (x <= xk[1]) {
return W[1, 1] + W[1, 2] * (x - xk[1]);
}
// Linear extrapolation using the right endpoint slope.
if (x >= xk[K]) {
real h = xk[K] - xk[K - 1];
real yK = ((W[K - 1, 4] * h + W[K - 1, 3]) * h + W[K - 1, 2]) * h + W[K - 1, 1];
real dK = (3 * W[K - 1, 4] * h + 2 * W[K - 1, 3]) * h + W[K - 1, 2];
return yK + dK * (x - xk[K]);
}
// Binary search for xk[i] <= x < xk[i+1].
while (hi - lo > 1) {
int mid = (lo + hi) %/% 2;
if (x < xk[mid])
hi = mid;
else
lo = mid;
}
i = lo;
t = x - xk[i];
return ((W[i, 4] * t + W[i, 3]) * t + W[i, 2]) * t + W[i, 1];
}
pchip.stan is almost entirely written by ChatGPT. I “checked” it by also fitting the data with a reference stan model that uses the true formula for V_x instead, and finding close enough posterior distriutions for the model parameters a,b, sigma.