Dynamic hierarchical model for transition probabilities

Modeling
1

I realize that my questions tend to be a bit niche but I hope that this may be helpful to someone else as well.

I want to allow a user to specify whether a predictor affects the general transiency of the states (\alpha), the state-specific transience (\alpha_i) or the state-state specific transience (\alpha_{i, j}) so that no additional recompilation of the model is required and the user only has to input which predictors vary at which level when she interacts with the main function. With that information, the model is supposed to populate the predictor array \Lambda'_{K\times M \times K} such that as in a previous question, the diagonal elements are 0. For example if we have three states and three predictors with \alpha, \beta, \gamma which are general, state-specific, and state-state specific respectively this would look like this and enter a multinomial specification for the probabilities with Pr(S_t = 2\vert S_{t - 1} = 1) = \frac{exp(\alpha x_1 + \beta_1 x_2 + \gamma_{1,2} x_3)}{1 + \sum_{j' \neq 1}^K exp(\alpha x_1 + \beta_1 x_2 + \gamma_{1,{j'}} x_3)}:
\begin{bmatrix}\begin{bmatrix} 0 &\alpha & \alpha \\ 0 & \beta_1 & \beta_1 \\ 0 & \gamma_{1,2} & \gamma_{1, 3} \end{bmatrix} \\ \begin{bmatrix} \alpha & 0& \alpha \\ \beta_2 & 0 & \beta_2 \\ \gamma_{2,1} & 0 & \gamma_{2, 3} \end{bmatrix} \\ \begin{bmatrix} \alpha & \alpha & 0 \\ \beta_3 & \beta_3 & 0 \\ \gamma_{3,1} & \gamma_{3, 2} & 0 \end{bmatrix}\end{bmatrix}

and for one example:

\gamma \sim \mathcal{N}(\mu, \sigma) \\ \gamma_{1} \sim \mathcal{N}(\gamma, \sigma_{\gamma}) \\ \gamma_{1,2} \sim \mathcal{N}(\gamma_{1}, \sigma_{\gamma_1})

The code below works but is very ugly (imo) and I don’t know how much overhead it causes with all the separate conditions. Hence, is there potential an easier way to do this or are there any simplifications I could implement? This takes about 1s for 1 chain of 2000 draws but obviously without actual data.

data {
  int K;
  int M;
  int pp1;
  int pp2;
  int pp3;
  int<lower = 0, upper = 1> pp_lambda[3, M];
  int<lower = 0, upper = 1> pp_lambda23[M];
  int<lower = 0, upper = 1> pp_lambda2_[pp2];
  int<lower = 0, upper = 1> pp_lambda3_[pp2];
}

parameters {
  real lambda1[M];
  real lambda2[pp2, K];
  real lambda3[K, pp3, K - 1];
}

transformed parameters {

  matrix[M, K] lambda_prime[K];
  {
    for (i in 1:K){
      int count3 = 1;
      for (j in 1:K){
        int count2 = 1;
        int count1 = 1;
        for (q in 1:M){
          if (i == j){
            lambda_prime[i, q, j] = 0.0;
          } else {
            if (pp_lambda[1, q]) {                            // same persistence/transience
              lambda_prime[i, q, j] = lambda1[q];
            } else if (pp_lambda[2, q]) {                     // state specific persistence/transience
              while (pp_lambda2_[count2] != 1) {
                count2 += 1;
              }
              lambda_prime[i, q, j] = lambda2[count2, i];
              count2 += 1;
            } else if (pp_lambda[3, q]) {                     // state-state specific persistence/transience
              lambda_prime[i, q, j] = lambda3[i, count1, count3];
              count1 += 1;
            }
            if (q == M) {
              count3 += 1;
            }
          }
        }

      }
    }
  }
}

model {
  real lambda2_mean[pp2];
  real lambda3_mean[pp3, K];
  int count = 1;
  for (i in 1:M){
    if (pp_lambda23[i] == 1) {
      lambda2_mean[count] = lambda1[i]; count += 1;
    }
  }
  count = 1;
  for (i in 1:pp2) {
    if (pp_lambda3_[i] == 1) {
      for (k in 1:K){
        lambda3_mean[count, k] = lambda2[i, k];
      }
      count += 1;
    }
  }

  target += normal_lpdf(lambda1 | 0, 1);
  print(lambda2_mean);
  for (i in 1:pp2){
    for (k in 1:K){
      target += normal_lpdf(lambda2[i, k] | lambda2_mean[i], 1);
    }
  }
    for (i in 1:pp3){
    for (k in 1:K){
      target += normal_lpdf(lambda3[k, i,:] | lambda3_mean[i, k], 1);
    }
  }
}


library(rstan)

M <- 6
K <- 3
pp_lambda <- rmultinom(N, 1, prob = rep(1/3, 3))
pp_lambda23 <- pp_lambda[2, ] + pp_lambda[3, ]
pp_lambda3_ <- pp_lambda[3, ][!pp_lambda[1,]]
pp_lambda2_ <- as.integer(!pp_lambda3_)

pp1 <- sum(pp_lambda[1, ])
pp2 <- M - sum(pp_lambda[1, ])
pp3 <- sum(pp_lambda[3, ])

setwd("~/Documents/Projects/01_MSM_TVTP_CSTS/rstanmsm_testing/01_Intercepts in call/test")

m <- stan_model("hierarchical.stan")

data = list(
  K = K,
  M = M,
  pp1 = pp1,
  pp2 = pp2,
  pp3 = pp3,
  pp_lambda = pp_lambda,
  pp_lambda23 = pp_lambda23,
  pp_lambda2_ = pp_lambda2_,
  pp_lambda3_ = pp_lambda3_
)

fit <- rstan::sampling(m, data = data)
3 Likes
2

This one might take a few replies :-) Some suggestions (without testing):

  1. I think normal_lpdf is a vectorised function. So you could get rid of the for(k in 1:K) loops and use the vectorised versions.

  2. I think you can replace lambda3_mean[count, k] = lambda2[i, k]; with lambda3_mean[count,] = lambda2[i,]; , and then get rid of the loop.

  3. Can you express the transformed parameters section as matrix algebra anyhow? If you could there’d be a direct route to translating that into efficient Stan code.

1 Like