tmbstanによるMCMCサンプリング – 伊東生態統計研究室

TMBでは、負の対数尤度を最小化することで最尤推定値を得ますが、TMBのオブジェクトを利用してMCMCのサンプリングをおこなうtmbstanというパッケージがあったので、ためしてみました。

準備

「TMBによる混合効果モデルのあてはめ」で使用したデータとTMBのモデルを再利用します。

まずはライブラリの読み込みと擬似乱数系列の固定です。

library(TMB)
library(ggplot2)
library(tmbstan)
## 要求されたパッケージ rstan をロード中です
## 要求されたパッケージ StanHeaders をロード中です
## 
## rstan version 2.32.6 (Stan version 2.32.2)
## For execution on a local, multicore CPU with excess RAM we recommend calling
## options(mc.cores = parallel::detectCores()).
## To avoid recompilation of unchanged Stan programs, we recommend calling
## rstan_options(auto_write = TRUE)
## For within-chain threading using `reduce_sum()` or `map_rect()` Stan functions,
## change `threads_per_chain` option:
## rstan_options(threads_per_chain = 1)
set.seed(123)

データ

ゼロ過剰ポアソン分布にしたがい、群ごとの変量効果のあるデータです。

N <- 400
N_group <- 10
group <- rep(1:N_group, each = N / N_group)
p <- 0.3
alpha <- -1
beta <- 0.8
log_sigma <- -1.6 # sigma = 0.202
ranef <- rnorm(N_group, 0, exp(log_sigma))

x <- runif(N, 0, 5)
y <- rbinom(N, 1, 1 - p) *
     rpois(N, exp(alpha + beta * x + ranef[group]))
example <- data.frame(x = x, y = y, group = group)

ggplot(example, aes(x = x, y = y)) +
  geom_point()

モデル

TMBのC++のコードです。これも「TMBによる混合効果モデルのあてはめ」のものとおなじです。

zipmix.cpp

// Zero-Inflated Poisson
#include <TMB.hpp>

template<class Type>
Type objective_function<Type>::operator() ()
{
  DATA_VECTOR(Y);
  DATA_VECTOR(X);
  DATA_IVECTOR(group);
  PARAMETER(p);
  PARAMETER(alpha);
  PARAMETER(beta);
  PARAMETER_VECTOR(epsilon);
  PARAMETER(log_sigma);
  Type lp = 0;
  lp += -sum(dnorm(epsilon, Type(0.0), exp(log_sigma), true));
  for (int i = 0; i < Y.size(); i++) {
    Type lambda = exp(alpha + beta * X(i) + epsilon(group(i)));
    lp += -dzipois(Y(i), lambda, p, true);
  }
  return lp;
}

コンパイルと最適化

モデルをコンパイルして、できたライブラリをロードし、最適化をおこないます。ここまでは前の記事と同じです。

model_name <- "zipmix"
file.path("models", paste(model_name, "cpp", sep = ".")) |>
  compile()
file.path("models", dynlib(model_name)) |>
  dyn.load()

data <- list(Y = example$y, X = example$x,
             group = example$group - 1)
parameters <- list(p = 0.5, alpha = 1, beta = 1,
                   epsilon = rep(0, N_group), log_sigma = 0)
obj_zipmix <- MakeADFun(data, parameters, DLL = model_name,
                        random = "epsilon")
opt_zipmix <- nlminb(obj_zipmix$par, obj_zipmix$fn, obj_zipmix$gr)

MakeADFun関数で、最適化関数に渡すオブジェクトを作成して、nlminb関数で最適化します。MakeADFun関数のrandom引数に、モデル中の変量効果であるepsilonを指定します。

結果

TMBによる最尤推定の結果です。

print(opt_zipmix)
## $par
##          p      alpha       beta  log_sigma 
##  0.2992974 -1.0778807  0.8271015 -1.5772257 
## 
## $objective
## [1] 678.5374
## 
## $convergence
## [1] 0
## 
## $iterations
## [1] 24
## 
## $evaluations
## function gradient 
##       32       25 
## 
## $message
## [1] "relative convergence (4)"

tmbstanでMCMC

つづいて、このモデルを利用してtmbstanでMCMCのサンプリングをおこないます。tmbstanはrstanを使っており、chainsなどの引数はrstan::samplingに渡されます。また返り値はstanfitクラスのオブジェクトになります。

stanfit_zipmix <- tmbstan(obj_zipmix,
                          chain = 4, iter = 5000, warmup = 1000,
                          cores = min(4, parallel::detectCores()))

結果の要約です。

print(stanfit_zipmix)
## Inference for Stan model: zipmix.
## 4 chains, each with iter=5000; warmup=1000; thin=1; 
## post-warmup draws per chain=4000, total post-warmup draws=16000.
## 
##                mean se_mean   sd    2.5%     25%     50%     75%   97.5% n_eff
## p              0.30    0.00 0.03    0.24    0.28    0.30    0.32    0.36 10405
## alpha         -1.08    0.00 0.14   -1.35   -1.17   -1.08   -0.99   -0.82  4412
## beta           0.83    0.00 0.03    0.77    0.81    0.83    0.85    0.88  6629
## epsilon[1]    -0.04    0.00 0.11   -0.25   -0.11   -0.04    0.03    0.18  5727
## epsilon[2]    -0.02    0.00 0.11   -0.23   -0.09   -0.02    0.05    0.19  5516
## epsilon[3]     0.33    0.00 0.10    0.15    0.27    0.33    0.40    0.54  4677
## epsilon[4]    -0.01    0.00 0.12   -0.25   -0.09   -0.02    0.06    0.22  6000
## epsilon[5]    -0.05    0.00 0.12   -0.28   -0.13   -0.05    0.02    0.17  6402
## epsilon[6]     0.33    0.00 0.11    0.13    0.26    0.32    0.40    0.54  5203
## epsilon[7]     0.05    0.00 0.11   -0.16   -0.02    0.05    0.13    0.27  5521
## epsilon[8]    -0.31    0.00 0.12   -0.55   -0.38   -0.30   -0.23   -0.08  6246
## epsilon[9]    -0.21    0.00 0.11   -0.44   -0.29   -0.21   -0.14    0.00  6259
## epsilon[10]   -0.05    0.00 0.11   -0.26   -0.12   -0.05    0.02    0.16  5333
## log_sigma     -1.48    0.00 0.28   -2.00   -1.68   -1.50   -1.30   -0.91  8342
## lp__        -669.59    0.03 2.71 -675.81 -671.21 -669.24 -667.65 -665.28  6019
##             Rhat
## p              1
## alpha          1
## beta           1
## epsilon[1]     1
## epsilon[2]     1
## epsilon[3]     1
## epsilon[4]     1
## epsilon[5]     1
## epsilon[6]     1
## epsilon[7]     1
## epsilon[8]     1
## epsilon[9]     1
## epsilon[10]    1
## log_sigma      1
## lp__           1
## 
## Samples were drawn using NUTS(diag_e) at Sun Sep 22 09:52:08 2024.
## For each parameter, n_eff is a crude measure of effective sample size,
## and Rhat is the potential scale reduction factor on split chains (at 
## convergence, Rhat=1).

pとalphaの事後平均は、最尤推定値と近い値となっています。

TMBの最尤推定で使っているnlminb関数による最適化では、局所最適にはまってしまう可能性もありそうですので、そのような場合にはMCMCによる推定が役に立ちそうです。