Dynamic Linear Models with Switching in Stan
Dynamic Linear Models with Switching in Stan
2022-10-30
In this post we are trying to reproduce Example 6.22 presented in the book “Time Series Analysis and its Applications with R examples”, by Robert H. Shumway and David S. Stoffer. We will analyze the U.S. monthly pneumonia and influenza mortality using Dynamic Linear Models with Switching. This kind of models attempt to generalize state space models to include the possibility of regime changes occurring over time. The idea is to use hidden markov models (HMM) to marginalize the hidden discrete states allowing the state space model to behave differently depending on the underlying unobserved state. The example in the book utilizes the exact likelihood, but in this post we will try to express the model explicitly with Stan. For an example of state-space model with Stan see Juho Kokkala’s blog post http://www.juhokokkala.fi/blog/posts/kalman-filter-style-recursion-to-marginalize-state-variables-to-speed-up-stan-inference/ or Jeffrey B. Arnold’s library https://jrnold.github.io/ssmodels-in-stan/filtering-and-smoothing.html, while an example of HMM can be found in Stan’s manual https://mc-stan.org/docs/stan-users-guide/hmms.html.
require(tidyverse)
require(rstan)
require(mvtnorm)
require(cmdstanr)
require(posterior)
library(astsa)Let’s have a look at the date representing mortality rate from 1968 to 1978:
plot.default(flu, ylab="Mortality", xlab="", cex.axis=1.5, cex.lab=1.5, type="o")
The model consists of three structural components. The first component is an AR(2) process chosen to represent the periodic (seasonal) component of the data:
\[ x_{t1} = \alpha_{1}x_{t-1,1} + \alpha_{2}x_{t-2,1} + w_{t1} \] where \(w_{t1}\) is white noise with variance = \(\sigma_{1}^2\).
The second component is an AR(1) process with a nonzero constant term, representing the sharp rise in the data during an epidemic:
\[ x_{t2} = \beta_{0} + \beta_{1}x_{t-1,2} + w_{t2} \] where \(w_{t2}\) is white noise with variance = \(\sigma_{2}^2\). The third component is a fixed trend component:
\[ x_{t3} = x_{t-1,3} + w_{t3} \]
In the book this process is assumed to be deterministic because estimation was unstable, we will try try to estimate it nonetheless. The model can be expressed in state-space form as:
\[ \begin{pmatrix} x_{t,1} \\ x_{t-1,1} \\ x_{t,2} \\ x_{t,3} \end{pmatrix} = \begin{bmatrix} \alpha_{1} & \alpha_{2} & 0 & 0 \\ 1 & 0 & 0 & 0 \\ 0 & 0 & \beta_{1} & 0 \\ 0 & 0 & 0 & 1 \end{bmatrix} \begin{pmatrix} x_{t-1,1} \\ x_{t-2,1} \\ x_{t-1,2} \\ x_{t-1,3} \end{pmatrix} + \begin{pmatrix} 0 \\ 0 \\ \beta_{0} \\ 0 \end{pmatrix} + \begin{pmatrix} w_{t1} \\ 0 \\ w_{t2} \\ 0 \end{pmatrix} \]
The observation equation is:
\[ y_{t} = A_{t}x_{t} + v_{t} \]
We want periods of normal influenza mortality (regime 1) are modeled as:
\[ y_{t} = x_{t1} + x_{t3} + v_{t} \]
while during epidemics (regime 2) mortality is modelled as:
\[ y_{t} = x_{t1} + x_{t2} + x_{t3} + v_{t} \]
Therefore, the matrix A can take two forms depending on the regime: [1,0,0,1] for no epidemic, and [1,0,1,1] for epidemic.
We decided to model the process in Stan as follows:
// Model for the influenza mortality, mode of 4 states and 2 regimes
data {
int<lower=0> N;
array[N] real y;
vector[4] m0;
matrix[4, 4] P0;
}
transformed data {
int D = 4; // the number of states
matrix[D, D] I; // the identity matrix
I = diag_matrix(rep_vector(1, D));
}
parameters {
vector[2] alpha;
real<lower=0> beta0;
real beta1;
vector<lower=0>[2] sigma;
real<lower=0> sigmav;
array[2] simplex[2] theta; // state transistion probabilities
}
transformed parameters {
matrix[D, D] Q; // process error matrix
real R; // observation error matrix
matrix[D, D] Z; // transition matrix
matrix[D, D] Zt; // transition matrix transpose
array[2, N] vector[D] m; // filtered states
array[2, N] vector[D] m_pred; // predicted states
array[2, N] matrix[D, D] P; // filtered covariance matrix
array[2, N] matrix[D, D] P_pred; // predicted covariance matrix
array[2, N] real S; // variance of the prediction errors
array[2] row_vector[D] A; // observation model matrices, depengin on the regime
vector[D] u; // interceps of the process
array[2, N] vector[D] K; // Kalman gains
array[2, N] real v; // prediction errors
array[2, N] real mu;
array[2] real acc; // HMM temporary probabilities
array[N] vector[2] gamma; // HMM forward values
{
row_vector[D] H; // Just a variable to store the different A matrices depending on the regime
vector[D] Ht; // H transpose
Z = [[alpha[1], alpha[2], 0, 0],
[1, 0, 0, 0],
[0, 0, beta1, 0],
[0, 0, 0, 1]];
Zt = Z';
A[1] = [1, 0, 0, 1];
A[2] = [1, 0, 1, 1];
u = [0, 0, beta0, 0]';
Q = diag_matrix([sigma[1] ^ 2, 0, sigma[2] ^ 2, 0]');
R = sigmav ^ 2;
// Initialize the states, we can also set these as paramteres to estimate
m_pred[1][1] = m0;
m_pred[2][1] = m0;
P_pred[1][1] = P0;
P_pred[2][1] = P0;
// Kalman filter see https://jrnold.github.io/ssmodels-in-stan/filtering-and-smoothing.html
for (t in 1 : N) {
for (k in 1 : 2) {
if (t > 1) {
m_pred[k][t] = Z * m[k][t - 1] + u;
P_pred[k][t] = Z * P[k][t - 1] * Zt + Q;
}
H = A[k];
Ht = H';
v[k][t] = y[t] - H * m_pred[k][t];
S[k][t] = H * P_pred[k][t] * Ht + R;
K[k][t] = (P_pred[k][t] * Ht) / S[k][t];
m[k][t] = m_pred[k][t] + K[k][t] * v[k][t];
P[k][t] = (I - K[k][t] * H) * P_pred[k][t];
mu[k][t] = H * m_pred[k][t];
}
}
}
// HMM marginalization see https://mc-stan.org/docs/stan-users-guide/hmms.html
for (k in 1 : 2)
gamma[1, k] = normal_lpdf(y[1] | mu[k][1], sqrt(S[k][1]));
for (t in 2 : N) {
for (k in 1 : 2) {
for (j in 1 : 2) {
acc[j] = gamma[t - 1, j] + log(theta[j, k])
+ normal_lpdf(y[t] | mu[j][t], sqrt(S[j][t]));
}
gamma[t, k] = log_sum_exp(acc);
}
}
}
model {
sigma ~ exponential(1);
sigmav ~ exponential(1);
alpha ~ normal(0, 1);
beta0 ~ exponential(1);
beta1 ~ normal(1, 0.5);
theta[1, 1] ~ beta(10, 2);
theta[2, 2] ~ beta(10, 2);
theta[2, 1] ~ beta(2, 10);
theta[1, 2] ~ beta(2, 10);
target += log_sum_exp(gamma[N]);
}
generated quantities {
array[N] vector[2] regimes;
for (t in 1 : N) {
regimes[t] = softmax(gamma[t]);
}
}
The code seems to be very involved but you can see that apart for a large amount of variables declared, most of the code is mode of the main loop for the kalman filter and the hmm marginalization. Notice that the Kalman filter loop is run twice for both values of matrix A representing the two possible regimes, therefore we need to have two value for each variable of the Kalman filter, as you can see in the declaration at the top. Later during the HMM forward algorithm we calculate the probability of each regime (which differ only for the value of A).
m <- cmdstan_model("models/Misc/flu.stan")
fit <- m$sample(data=list(N=length(flu), y=c(0, diff(as.vector(flu))), m0=rep(0, 4), P0=diag(c(1,1,1,1))), parallel_chains = 4, cores = 4,iter_warmup = 250, iter_sampling = 250)Parameters estimations are vaguely close to the ones in the book, we additionally obtain transition probabilities estimations:
fit$print(max_rows = 12, digits = 4)## variable mean median sd mad q5 q95 rhat ess_bulk
## lp__ 179.2173 179.5520 2.1490 2.1349 175.3534 182.1791 1.0099 336
## alpha[1] 1.2787 1.2769 0.0552 0.0546 1.1880 1.3708 1.0097 323
## alpha[2] -0.3632 -0.3610 0.0449 0.0457 -0.4384 -0.2917 1.0056 313
## beta0 0.8701 0.5959 0.9020 0.6077 0.0523 2.6551 1.0008 534
## beta1 0.5458 0.5559 0.1520 0.1536 0.2873 0.7814 1.0033 557
## sigma[1] 0.0210 0.0209 0.0020 0.0021 0.0179 0.0241 1.0041 513
## sigma[2] 0.1671 0.1656 0.0244 0.0253 0.1331 0.2091 1.0014 678
## sigmav 0.0034 0.0033 0.0022 0.0027 0.0003 0.0072 1.0192 285
## theta[1,1] 0.9051 0.9095 0.0298 0.0270 0.8506 0.9455 1.0043 662
## theta[2,1] 0.2264 0.2225 0.0622 0.0619 0.1319 0.3343 1.0055 572
## theta[1,2] 0.0949 0.0905 0.0298 0.0270 0.0545 0.1494 1.0043 662
## theta[2,2] 0.7736 0.7775 0.0622 0.0619 0.6657 0.8681 1.0055 572
## ess_tail
## 583
## 282
## 336
## 337
## 483
## 374
## 635
## 388
## 477
## 434
## 477
## 434
##
## # showing 12 of 13011 rows (change via 'max_rows' argument or 'cmdstanr_max_rows' option)Regime estimation are reasonable, even if not so good as in the book:
k <- 6; y <- as.matrix(flu); num <- length(y); nstate <- 4;
s <- fit$draws(variables = "regimes") %>% merge_chains() %>% colMeans() %>% matrix(nrow=nrow(y))
Time <- as.matrix(time(flu))
regime <- ifelse(s[,1] < s[,2], 1, 2);
plot(Time, y, type="n", ylab="", cex=2, cex.axis=2, xlab="")
grid(lty=2); lines(Time, y, col=gray(.7))
text(Time, t(y), col=regime, labels=regime, cex=2)
To obtain the estimation of the final filtered states, we get the filtered states for the two regimes and we calculated the average weighted using the probability for each regimes:
a <- fit$draws(variables = "m") %>% merge_chains() %>% colMeans()
m1 <- matrix(a[seq(1, 1056, 2)], ncol=4)
m2 <- matrix(a[seq(2, 1056, 2)], ncol=4)
p1 <- matrix(rep(s[,1], 4), ncol=4)
p2 <- matrix(rep(s[,2], 4), ncol=4)
m3 <- m1 * p1 + m2 * p2
matplot(Time, m3[,-2], type="o", cex=1.5, ylab="", cex.axis=2)
Estimation for the trend process seems to be strongly affected by the epidemic oscillation, unlike the state obtained in the book. It is possible that additional precautions are needed in the definition of our model to better match the results. We finally show one-month-ahead predictions with the variability (2*innovations-variance):
mu <- fit$draws(variables = "mu") %>% merge_chains() %>% colMeans()
S <- fit$draws(variables = "S") %>% merge_chains() %>% colMeans()
mu1 <- mu[seq(1, 264, 2)]
mu2 <- mu[seq(2, 264, 2)]
mu3 <- mu1 * s[,1] + mu2 * s[,2]
S1 <- S[seq(1, 264, 2)]
S2 <- S[seq(2, 264, 2)]
S3 <- 2*sqrt(S1 * s[,1] + S2 * s[,2])
plot(Time, mu3, type="n", ylab="", xlab="", ylim=c(0,1), cex=2, cex.axis=2)
grid(lty=2);
points(Time, as.vector(y), pch=16, cex=1.5)
xx = c(Time, rev(Time))
yy = c(mu2-S3, rev(mu3+S3))
polygon(xx, yy, border=8, col=gray(.6, alpha=.3))
Comments
Post a Comment