Bayesian changepoint detection in linear regression with R and Stan

changepoint-linear-regression.r

# Based on this blog post: http://nowave.it/pages/bayesian-changepoint-detection-with-r-and-stan.html
library(rstan)
rstan_options(auto_write = TRUE)
set.seed(42)

beta0 <- 3
beta1 <- 9
beta2 <- 15

set.seed(42)
x <- sort(runif(100, 0, 10))
x1 <- head(x, 41)
x2 <- tail(x, 59)
y1 <- rnorm(41, mean=beta0 + beta1 * x1, sd=4.3)
y2 <- rnorm(59, mean=beta0 + beta2 * x2, sd=6.3)

y <- c(y1, y2)
plot(x, y)


stan_code <- '
data {
  int<lower=1> N;
  real x[N];
  real y[N];
}

parameters {
  real beta0;
  real beta1;
  real beta2;
  real<lower=0> sigma;
}

transformed parameters {
  vector[N] mu1;
  vector[N] mu2;
  vector[N] log_p;
  {
    vector[N+1] log_p_e;
    vector[N+1] log_p_l;
    log_p_e[1] = 0;
    log_p_l[1] = 0;
    for (i in 1:N) {
      mu1[i] = beta0 + beta1 * x[i];
      mu2[i] = beta0 + beta2 * x[i];

      log_p_e[i + 1] = log_p_e[i] + normal_lpdf(y[i] | mu1[i], sigma);
      log_p_l[i + 1] = log_p_l[i] + normal_lpdf(y[i] | mu2[i], sigma);
    }
    log_p = rep_vector(-log(N) + log_p_l[N + 1], N) + head(log_p_e, N) - head(log_p_l, N);
  }
}
model {
  beta0 ~ normal(0, 100);
  beta1 ~ normal(0, 100);
  beta2 ~ normal(0, 100);
  sigma ~ normal(0, 100);
  
  target += log_sum_exp(log_p);
}

generated quantities {
  int<lower=1,upper=N> tau;
  // K simplex are a data type in Stan
  simplex[N] sp;
  sp = softmax(log_p);
  tau = categorical_rng(sp);
}
'

fit <- stan(
  model_code = stan_code,
  data = list(x = x, y = y, N=length(x)),
  chains = 4,
  iter = 10000,
  cores = 4,
  refresh = 500
)

plot(fit, pars=c("beta0", "beta1", "beta2", "sigma", "tau"))