Kalman filter in Stan

globaltemp.R

library(KFAS)
library(rstan)
data(GlobalTemp)

model_dlm1a <- stan_model("../stan/dlm1a.stan")

y <- as.matrix(GlobalTemp)
data <-
    within(list(),
       {
           y <- y
           n <- nrow(y)
           p <- ncol(y)
           m <- 2
           r <- 2
           W <- rep(1, length(y))
           Z <- matrix(c(1, 1, 0, 0), 2, 2)
           T <- matrix(c(1, 0, 1, 1), 2, 2)
           R <- diag(2)
           a_1 <- rep(0, m)
           P_1 <- diag(m)
       })

system.time(fit1a <- sampling(model_dlm1a, data, chains=1, iter=5000))
# values of h
apply(extract(fit1a, "h")[[1]], 2, mean)
# values of q
apply(extract(fit1a, "q")[[1]], 2, mean)

modelTemp <- SSModel(y = data$y,
                     Z = data$Z,
                     T = data$T,
                     R = data$R,
                     a1=data$a1,
                     P1=data$P1,
                     H=matrix(c(NA, 0, 0, NA), 2, 2),
                     Q=matrix(c(NA, 0, 0, NA), 2, 2))
fit<-fitSSM(inits=rep(0.1, 4), model=modelTemp)
diag(fit$model$H[, , 1])
diag(fit$model$Q[, , 1])

kalman.stan

/*
  Multivariate Dynamic linear model
  
  estimated with

  - Covariance filter (no square root, or sequential processing)
  - time invariant parameters
  - no missing observations
  - known initial value

  Only parameters are the measurement error and system error, both 
  of which are assumed to be diagonal.

  .. math::

     y_t &= Z_t \alpha_t + \epsilon_t \\
     \epsilon_t &\sim N(0, H_t) \\
     \alpha_{t+1} = T_t \alpha_t + R_t \eta_t \\
     \eta_t &\sim N(0, Q_t) \\
     \alpha_1 &\sim N(a_1, P_1) \\
     t &= 1, \dots, n

  Dimensions
  
  ------------------ --------
  variable           dim
  ================== ========
  :math:`y_t`        p, 1
  :math:`\alpha_t`   m, 1
  :math:`epsilon_t`  p, 1
  :math:`eta_t`      r, 1
  :math:`a_1`        m, 1
  :math:`Z_t`        p, m
  :math:`T_t`        m, m
  :math:`H_t`        p, p
  :math:`R_t`        m, r
  :math:`Q_t`        r, r
  :math:`P_1`        m, m

  - :math:`p` dimension of data
  - :math:`n` number of observation
  - :math:`m` number of states
  - :math:`r` number of states with non-zero variance.

  See Durbin & Koopman, Ch 4.2, pp 65-67.

  .. math::
  
     v_t &= y_t - Z_t a_t \\
     F_t &= Z_t P_t Z_t' + H_t \\
     K_t &= T_t P_t Z_t' F_t^{-1} \\
     L_T &= T_t - K_t Z_t \\
     a_{t+1} &= T_t a_t + K_t v_t \\
     P_{t+1} &= T_t P_t L_t' + R_t Q_t R_t'

  For loglikelihood, Ch. 7.2, pp. 138-139.
  
  .. math::

     \log L(y) = \sum_{t=1}^n p(y_t | Y_{t-1})
    
  where :math:`p(y_1 | Y_0) = p(y_1)`.
  Supposing normal distributions,

  .. math::

     \log L(y) &= \sum_{t=1}^n \phi(Z_t a_t, F_t) \\
     &= -\frac{n p}{2} - \frac{1}{2} \sum_{t=1}^n (\log | F_t | + v_t' F_t^{-1} v_t)
  
*/
data {
  // number of observations
  int n;
  // number of variable == 1
  int p;
  // number of states
  int m;
  // size of system error
  int r;
  // observed data
  vector[p] y[n];
  // observation eq
  matrix[p, m] Z;
  // system eq
  matrix[m, m] T;
  matrix[m, r] R;
  // initial values
  vector[m] a_1;
  cov_matrix[m] P_1;
  // measurement error
  // vector<lower=0.0>[p] h;
  // vector<lower=0.0>[r] q; 
}
parameters {
  vector<lower=0.0>[p] h;
  vector<lower=0.0>[r] q;
}
transformed parameters {
  matrix[p, p] H;
  matrix[r, r] Q;
  H <- diag_matrix(h);
  Q <- diag_matrix(q);
}
model {
  vector[m] a[n + 1];
  matrix[m, m] P[n + 1];
  real llik_obs[n];
  real llik;
  // 1st observation
  a[1] <- a_1;
  P[1] <- P_1;
  for (i in 1:n) {
    vector[p] v;
    matrix[p, p] F;
    matrix[p, p] Finv;
    matrix[m, p] K;
    matrix[m, m] L;
    v <- y[i] - Z * a[i];
    F <- Z * P[i] * Z' + H;
    Finv <- inverse(F);
    K <- T * P[i] * Z' * Finv;
    L <- T - K * Z;
    // // manual update of multivariate normal
    llik_obs[i] <- -0.5 * (p * log(2 * pi()) + log(determinant(F)) + v' * Finv * v);
    //llik_obs[i] <- multi_normal_log(y[i], Z * a[i], F);
    a[i + 1] <- T * a[i] + K * v;
    P[i + 1] <- T * P[i] * L' + R * Q * R';
  }
  llik <- sum(llik_obs);
  lp__ <- lp__ + llik;
}