kalman

kalman_filter_arm.stan

functions {
  vector kalman_lp(vector y, 
               real P0, 
               real Q,
               real H) {
  
    int n = size(y);
    vector[n] x_return;
    real x = y[1];
    real P = P0;
    real loglik = 0.0;
    
    for (i in 1:n) {
        real inno = y[i] - x;
        real F = P + H;
        real K = P / F;
        
        x += K * inno;
        x_return[i] = x;
        P = P * (1 - K) + Q;
        
        loglik += log(F) + inno^2 / F;
    }
    
    target += -0.5 * loglik;
    
    return x_return;
    
  }
  
}
data {
  int<lower=1> N;
  int<lower=1> K;
  int<lower=1> T;
  array[N] int<lower=1> trials;
  array[N] int<lower=0> successes;
  array[N] int<lower=1, upper=K> kk;
  array[N] int<lower=1, upper=T> tt;
}
parameters {
  vector[K] mu_0;
} 
transformed parameters {
  array[K] vector[T] theta;
  array[K] vector[T] mu;
  
  for (k in 1:K) {
    mu[k] = kalman_lp(mu_0[k], mu, 0., 1., 1.);
    theta[k] = inv_logit(mu[k]); 
  }
}
model {
  mu_0 ~ std_normal();
  
  for (n in 1:N)
    successes[n] ~ binomial(trials[n], theta[kk[n], tt[n]]);
}