State space / kalman filtering advice for inverse model

Modeling
12

Here a Stan function for doing the Kalman filtering, using notation mostly based on https://users.aalto.fi/~ssarkka/pub/cup_book_online_20131111.pdf:

dX/dt = A_c X_t + B_c U_t + N(0, Q)
y_t = HX_t + N(0, R),

real kf(
  matrix Ac,  // q*q, continuous time state transition matrix
  matrix Bc,  // q*m, continuous time input matrix
  matrix H,   // p*q, measurement matrix
  matrix U,   // N*m, input vector
  matrix y,   // N*p, measurement vector
  vector w,   // q, process noise
  vector v,   // q, measurement noise
  real dt,    // time interval
  vector X1,  // q,  initital X  at t=1
  matrix P1   // q*q, covariance at t=1
) {
  int N = dims(y)[1];  // number of time steps
  int p = dims(y)[2];  // number of observation variables
  int m = dims(Bc)[2]; // number of input variables
  int q = dims(A)[1];  // number of states
  
  vector[q] X[N];
  vector[p] v;
  matrix[q, q] Q = diag_matrix(w);   // state covariance
  matrix[p, p] R = diag_matrix(v);   // observation covariance
  matrix[q, q] P;     // covariance
  matrix[q, p] K;     // Kalman gain
  matrix[p, p] S;     // aka F 
  matrix[p, q] H;     // aka C
  matrix[p, p] S_inv; 
  matrix[q, q] A;
  matrix[q, m] B;
  matrix[q, q] I = diag_matrix(rep_vector(1., q));
  real ll = 0;        // logliklihood
  
  // continuous time to discrete time
  A = matrix_exp(A_c*dt);
  B = (A - diag_matrix(rep_vector(1., q))) / A_c * B_c;
  
  // Initialization
  X[1] = rep_vector(0, q);  
  P = rep_matrix(0, q, q);
  
  for (k in 1:N) {
    k__1 = max(k-1, 1);        // k-1
  
    /* Time update */
    X[k] = A*X[k__1] + B*U[k]; // predict state mean values
    P = A*P*A' + Q;            // predict state covariance matrix
  
    /* measurement update */  
    v = y[k]' - H*X[k];        // output mean residual
    S = H*P*H' + R;            // output covariance
    S_inv = inverse(S); 
    K = P*H' * S_inv;          // Kalman gain
  
    // update state covariance
    P = (I - K*H)*P*(I - K*H)' + K*R*K';  
    //P = (I-K*H)*P;      // less stable
  
    X[k] += (K*v);            // update states
  
    // log-likelihood
    ll -= (log_determinant(S) + v'*S_inv*v); 
  }
  return 0.5*(ll - p*N*log(2*pi()));
}

I just simplified the code above from a more specific code of mine. I haven’t tested it and probably missed a parenthesis or something. But, it should give you the general idea of how to implement a Kalman filter in Stan. @Charles_Driver post above shows how to get your model into general state-space representation.

A great resource on Bayesian filters is found here https://github.com/rlabbe/Kalman-and-Bayesian-Filters-in-Python.

2 Likes