Chronikis: a language for Bayesian time-series models that compiles to Stan

I’d like to announce the open-source release of Chronikis (kroh-NEE-kees), a special-purpose language for creating Bayesian time-series models, available here:

[http://opensource.adobe.com/Chronikis/](http://opensource.adobe.com/Chronikis/)

Stan is the target language for the Chronikis compiler, and use from within R is supported via rstan. The class of models currently supported are linear state-space models (a.k.a. dynamic linear models).

As a simple example, here is a Chronikis program for a random-walk model with added independent observation noise:

def main(sigma_q_scale, sigma_h_scale: real{0.0,},
         mu_a: real, sigma_a: real{0.0,})
  =
  sigma_q ~ half_cauchy(sigma_q_scale);
  sigma_h ~ half_normal(sigma_h_scale);
  accum(wn(sigma_q), mu_a, sigma_a) + wn(sigma_h)

In the above, sigma_q_scale, sigma_h_scale, mu_a, and sigma_a are all prior parameters that are supplied when a model is created and estimated. The first two lines express priors on sigma_q and sigma_h, and the last line is an expression whose “value” is a time series model, that is, a probability distribution over temporal sequences.

The subexpression

wn(sigma_h)

means “white noise” with mean 0 and variance sigma_h^2. The subexpression

accum(wn(sigma_q), mu_a, sigma_a)

(“accum” is short for “accumulate”) indicates a time series model for which

  • the initial value y[0] is drawn from a normal distribution with mean mu_a and variance sigma_a^2, and

  • the differences y[t] - y[t-1] are independent normal draws with mean 0 and variance sigma_q^2.

Note that Chronikis programs are fully declarative. It is also statically typed, but all types other than the parameters of main() are inferred. Type inference is bottom-up; I have not attempted to implement anything like the Hindley-Milner type inference used in the Haskell and the ML family of programming languages.

The Chronikis compiler translates the above into the following Stan program:

functions {
}
data {
  int<lower = 1> N;
  vector[N] y;
  real<lower = 0.0> sigma_q_scale_;
  real<lower = 0.0> sigma_h_scale_;
  real mu_a_;
  real<lower = 0.0> sigma_a_;
}
transformed data {
  matrix[1, N] yy = to_matrix(y');
  vector[1] Z_0_ = rep_vector(1.0, 1);
  matrix[1, 1] T_0_ = rep_matrix(1.0, 1, 1);
  vector[1] a0_0_ = rep_vector(mu_a_, 1);
  matrix[1, 1] P0_0_ = rep_matrix(square(sigma_a_), 1, 1);
}
parameters {
  real<lower = 0.0> raw_0_;
  real<lower = 0.0> raw_1_;
}
transformed parameters {
  real sigma_q_ = raw_0_ * sigma_q_scale_;
  real sigma_h_ = raw_1_ * sigma_h_scale_;
}
model {
  real H_0_ = square(sigma_h_);
  matrix[1, 1] Q_0_ = rep_matrix(square(sigma_q_), 1, 1);
  raw_0_ ~ cauchy(0.0, 1.0) T[0.0, ];
  raw_1_ ~ normal(0.0, 1.0) T[0.0, ];
  yy ~ gaussian_dlm_obs(to_matrix(Z_0_), T_0_, rep_vector(H_0_, 1), Q_0_, a0_0_, P0_0_);
}

A few notes:

  • Instead of requiring the user to specify the matrices defining the linear state-space model, Chronikis provides an algebra of time-series models which it transforms into a symbolic expression describing the equivalent linear state-space model.

  • The compiler infers the shapes and bounds of all variables, starting with the shapes and bounds of the parameters of main().

  • The compiler automatically reparameterizes inferred variables to use non-centered parameterization, and assigns variables to the appropriate Stan model blocks.

  • The Chronikis compiler is implemented in Haskell.

This current release is a very early version, and is still missing some obvious functionality. I welcome any feedback that (potential) users may provide.

Thanks to my employer, Adobe Inc., for allowing me to open-source this project.

5 Likes

Cool thanks.

Andrew just had a blog post on state space models.

https://statmodeling.stat.columbia.edu/2019/04/15/state-space-models-in-stan/

Thanks for posting!

Indeed, thanks to Adobe, too!