【R】Bayesian heteroskedastic linear regresssionをRで実装してみる（実装編）

はじめに

前回は以下の記事にて、Bayesian heteroskedastic linear regressionの理論について整理しました。

sunaninattahito.hatenablog.com

この記事では、前回の記事で導出したposteriorを基に、Rで自力でMCMCを実装してみることにします。また、この記事はGreenberg (2011)の教科書の練習問題がベースです。

www.amazon.co.jp

改善点や間違い等あればぜひ指摘してください。

パッケージの読み込み

まず最初に以下のパッケージを読み込みましょう。{tidylog}は{tidyverse}で行ったデータ加工の結果を表示してくれる便利な代物ですが、必要がなければ読み込まなくても全く問題ありません。

library(tidyverse)
library(tidylog)
library(mvtnorm) #多変量正規分布を発生させるやつ
set.seed(1)

データの読み込み

データは、Greenberg (2011) 8.4 Excersisesの問題8.3に基づいて著者のサイト(https://www.stat.berkeley.edu/users/statlabs/labs.html) にある``Birth weight II"というデータを使います。このデータは生まれてきた赤ちゃんの体重とそれに対して影響を与えるであろう変数が盛り込まれています。早速、読み込みましょう。

df <- read.table("https://www.stat.berkeley.edu/users/statlabs/data/babies.data", header = TRUE) %>% 
  as_tibble() # tibbleが好きなので

リンク先が.dataファイルなのでread.table関数を使用しました。少しデータを確認してみましょう。

dim(df)
head(df)
tail(df)

[1] 1236    7

# A tibble: 6 x 3
  variable   value iterations
  <chr>      <dbl>      <int>
1 intercept -103.           1
2 intercept  -76.8          2
3 intercept  -83.7          3
4 intercept  -63.4          4
5 intercept -115.           5
6 intercept  -68.5          6

# A tibble: 6 x 3
  variable value iterations
  <chr>    <dbl>      <int>
1 smoke    -8.07        995
2 smoke    -7.58        996
3 smoke    -8.25        997
4 smoke    -9.25        998
5 smoke    -7.94        999
6 smoke    -7.05       1000

また、記録されている変数は

names(df)

[1] "bwt" "gestation" "parity" "age" "height" "weight"
[7] "smoke"

• bwt: 赤ちゃんの出生時の体重 (ounces)
• gestation: 妊娠期間
• parity: 第一子かどうか (0 = first born)
• age: 出産時の年齢
• height: 母親の身長（inch）
• weight: 母親の妊娠時の体重 (pounds)
• smoke: 母親が喫煙するかどうか (0 = not now, 1 = yes now)

です。次に、データから欠損値がある観察iを削除していいとの事のなので、消しちゃいましょう。ちなみに、gestation・weightでは999、age・heightでは99、smokeでは9が欠損値とされています。

df %>% 
  filter(gestation != 999 & 
           age != 99 & 
           height != 99 & 
           weight != 999 & 
           smoke != 9) -> df

filter: removed 62 rows (5%), 1,174 rows remaining

上記のように{tidylog}は、filterなどの結果を教えてくれます。

記述統計をして、データがどのようになっているか確認します。

summary(df)

      bwt          gestation         parity            age       
 Min.   : 55.0   Min.   :148.0   Min.   :0.0000   Min.   :15.00  
 1st Qu.:108.0   1st Qu.:272.0   1st Qu.:0.0000   1st Qu.:23.00  
 Median :120.0   Median :280.0   Median :0.0000   Median :26.00  
 Mean   :119.5   Mean   :279.1   Mean   :0.2624   Mean   :27.23  
 3rd Qu.:131.0   3rd Qu.:288.0   3rd Qu.:1.0000   3rd Qu.:31.00  
 Max.   :176.0   Max.   :353.0   Max.   :1.0000   Max.   :45.00  
     height          weight          smoke      
 Min.   :53.00   Min.   : 87.0   Min.   :0.000  
 1st Qu.:62.00   1st Qu.:114.2   1st Qu.:0.000  
 Median :64.00   Median :125.0   Median :0.000  
 Mean   :64.05   Mean   :128.5   Mean   :0.391  
 3rd Qu.:66.00   3rd Qu.:139.0   3rd Qu.:1.000  
 Max.   :72.00   Max.   :250.0   Max.   :1.000  

実装

データも準備できたので、早速実装してみようと思います。また、簡単のために予めデータをattachしておきます。

attach(df)

この問題の要旨としては、赤ちゃんの出生時体重であるbwtを従属変数として、その他の変数を独立変数とする形になります。要は赤ちゃんの出生時体重の決定要因を探ろうみたいな感じですね。

まず初めに、分析で用いるデータを行列に変換します。

# 行列X (N × K)を作成
X <- cbind(1, gestation, parity, age, height, weight, smoke) 

# Yを N × 1 行列に変換
Y <- as.matrix(bwt) 

# 後ほど使用するので変数名を保存
X_name <- colnames(X) 

# 先頭をinterceptという名前にする
X_name[1] <- "intercept" 

Posteriors

次に各パラメータのposteriorを計算するための関数を以下で定義していきます。

βのposterior

ここで$\beta$のposteriorは、

where

$$ \begin{align} B_1 &= [B_0^{-1} + σ^{-2}X^{T}\Lambda X]^{-1},\\ \tilde{\beta} &= B_1[B_0^{-1}\beta_0 + σ^{-2}X^{T}\Lambda Y]. \end{align} $$

でした。これに基づいて、$\beta$のposteriorを計算する関数beta_postを以下で定義します。

beta_post <- function(Y, X, b0, B0_inv, sigma2, Lambda){
  B1 <- solve(B0_inv + t(X) %*% Lambda %*% X / sigma2)
  beta_mean <- B1 %*% (B0_inv %*% b0 + t(X) %*% Lambda %*% Y / sigma2)
  beta <- matrix(rmvnorm(n = 1, mean = beta_mean, sigma = B1), , 1)
  
  return(beta)
}

σ²のposterior

次に、σ²のposteriorを導出する関数を定義します。posteriorは、

$$ σ^2|\beta, \lambda, Y \sim \mathcal{IG}\left(\frac{\alpha_1}{2}, \frac{\delta_1}{2}\right), $$

where

$$ \begin{align} \alpha_1 &= \alpha_0 + n, \\\ \delta_1 &= \delta_0 + (Y-X\beta)^{T} \Lambda (Y-X\beta). \end{align} $$

です。よって、以下のsigma2_postを定義します。

sigma2_post　<- function(Y, X, alpha0, delta0, beta, Lambda){
  alpha1 <- alpha0 + nrow(X)
  delta1 <- delta0 + t(Y - X %*% beta) %*% Lambda %*%(Y - X %*% beta)
  sigma2 <- MCMCpack::rinvgamma(n = 1, alpha1 / 2, delta1 / 2)
  
  return(sigma2)
}

λiのposterior

最後に$\lambda_i$のposteriorを見ていきます。posteriorは以下のように表せます。

$$ \begin{align} \lambda_i | \beta, σ^{2}, Y \sim Ga\left(\frac{\nu_1}{2}, \frac{\nu_{2, i}}{2}\right),\\ \end{align} $$

where

$$ \begin{align} \nu_1 &= \nu + 1,\\ \nu_{2, i} &= \nu + σ^{-2}(Y_i - X_i\beta). \end{align} $$

また、ここで自由度$\nu$は既知として、関数を以下で定義します。

lambda_i_post <- function(Y, X, b0, B0_inv, nu, sigma2, Lambda) {
  nu1 <- nu + 1
  nu2_i <- matrix(NA, nrow(X), 1)
  lambda_i <- matrix(NA, nrow(X), 1)
    b <- beta_post(X = X, 
                             Y = Y, 
                             b0 = b0, 
                             B0_inv = B0_inv, 
                             sigma2 = sigma2, 
                             Lambda = Lambda)
    nu2_i <- nu + ((Y - X %*% b) / sigma2)
    for (i in 1:nrow(nu2_i)) {
      lambda_i[i] <- rgamma(n = 1, scale = (nu1 / 2), shape = (nu2_i[i] / 2)) 
    }
  return(lambda_i)
}

Gibbs Sampler

次に、Gibbs samplingをするために関数を定義します。

gibbs <- function(iter, warmup, thin, verbose = FALSE, Y, X, b0, B0_inv, alpha0, delta0, nu){
  
  set.seed(1)

  # 時間計測スタート
  stime <- Sys.time()
 
  # storage
  store_beta <- matrix(NA, iter / thin, ncol(X)) ### storage for beta
  store_sigma2 <- matrix(NA, iter / thin, 1) ### storage for sigma2
  store_lambda <- matrix(NA, iter / thin, nrow(X)) ### storage for lambda
  
  # initital value
  sigma2 <- 1 # initial value fot sigma2
  Lambda <- diag(1, nrow(X)) # initital value for Lambda
  
  total_iter <- iter + warmup 
  

  for (i in 1:total_iter){
    beta <- beta_post(Y, X, b0, B0_inv, sigma2, Lambda)
    sigma2 <- sigma2_post(Y, X, alpha0, delta0, beta, Lambda)
    lambda_i <- lambda_i_post(Y, X, b0, B0_inv, nu, sigma2, Lambda)
    Lambda <- diag(as.numeric(lambda_i))
    
    # save posteriors...
    if (i > warmup & ((i - warmup) %% thin == 0)){
      store_beta[(i - warmup) / thin, ] <- beta # save beta
      store_sigma2[(i - warmup) / thin, ] <- sigma2 # save sigma2
      store_lambda[(i - warmup) / thin, ] <- t(lambda_i) # save lambda
    }
    
    if (verbose == TRUE & i <= warmup & i %% 100 == 0){ 
      cat("Warmup", i,"DONE.", crayon::yellow("Total time:", Sys.time() - stime), "\n") 
    }
    if (verbose == TRUE & i > warmup & i %% 100 == 0){ 
      cat("Iterations", i - warmup,"DONE.", crayon::yellow("Total time:", Sys.time() - stime), "\n") 
    }
    
  }

  # naming
  colnames(store_beta) <- X_name 
  colnames(store_sigma2) <- "sigma2"
  colnames(store_lambda) <- paste0("observation", 1:nrow(X)) 

  # Convert into data frame
  store_beta <- as_tibble(store_beta) %>% 
    mutate(iterations = 1:(iter / thin), .before = X_name[1])
  
  store_sigma2 <- as_tibble(store_sigma2) %>%  
    mutate(iterations = 1:(iter / thin), .before = "sigma2")
  
  store_lambda <- as_tibble(store_lambda) %>% 
    mutate(iterations = 1:(iter / thin), .before = "observation1")
  
  return(list(beta_posterior = store_beta, 
              sigma2_posterior = store_sigma2, 
              lambda_posterior = store_lambda))

Greenberg (2011)の問題8.3に基づいて自由度$\nu$ = 5とします。また、このモデルをspecifyするために$\beta_0$, $B_0^{-1}$, $\alpha_0$, $\delta_0$にpriorを置きます。今回は適当にnoninformative priorを置くこととします。ここでは、

$$ \begin{align} \beta_0 &= 0, \\ B_0^{-1} &= 0, \\ \alpha_0 &= 0.001, \\ \delta_0 &= 0.001. \end{align} $$

とします。それでは、早速MCMCしていきましょう！

fit <- gibbs(Y = Y,
             X = X,
             iter = 3000, 
             warmup = 1000, 
             thin = 3, 
             verbose = TRUE,
             b0 = matrix(0, ncol(X), 1),
             B0_inv = diag(0, ncol(X), ncol(X)), 
             alpha0 = 0.001, 
             delta0 = 0.001,
             nu = 5)

Warmup 100 DONE. Total time: 9.5207188129425
Warmup 200 DONE. Total time: 16.6507158279419
Warmup 300 DONE. Total time: 22.6647148132324
Warmup 400 DONE. Total time: 28.8763778209686
Warmup 500 DONE. Total time: 34.9007859230042
Warmup 600 DONE. Total time: 41.0568199157715
Warmup 700 DONE. Total time: 46.8590228557587
Warmup 800 DONE. Total time: 52.8343329429626
Warmup 900 DONE. Total time: 58.8880109786987
Warmup 1000 DONE. Total time: 1.08210526704788
Iterations 100 DONE. Total time: 1.17846973339717
Iterations 200 DONE. Total time: 1.27412329912186
Iterations 300 DONE. Total time: 1.37088464895884
Iterations 400 DONE. Total time: 1.46519904931386
Iterations 500 DONE. Total time: 1.56177023251851
Iterations 600 DONE. Total time: 1.65910194714864
Iterations 700 DONE. Total time: 1.75826054811478
Iterations 800 DONE. Total time: 1.85463916460673
Iterations 900 DONE. Total time: 1.94738284746806
Iterations 1000 DONE. Total time: 2.04625216325124
Iterations 1100 DONE. Total time: 2.14010365009308
Iterations 1200 DONE. Total time: 2.23587293227514
Iterations 1300 DONE. Total time: 2.33102181355158
Iterations 1400 DONE. Total time: 2.42716313203176
Iterations 1500 DONE. Total time: 2.52142798105876
Iterations 1600 DONE. Total time: 2.61657724777857
Iterations 1700 DONE. Total time: 2.71183398167292
Iterations 1800 DONE. Total time: 2.80728929837545
Iterations 1900 DONE. Total time: 2.90187638203303
Iterations 2000 DONE. Total time: 2.99862943092982
Iterations 2100 DONE. Total time: 3.09441526333491
Iterations 2200 DONE. Total time: 3.18699366648992
Iterations 2300 DONE. Total time: 3.28234141667684
Iterations 2400 DONE. Total time: 3.37668879826864
Iterations 2500 DONE. Total time: 3.47107758124669
Iterations 2600 DONE. Total time: 3.56890536546707
Iterations 2700 DONE. Total time: 3.66383141676585
Iterations 2800 DONE. Total time: 3.7627916653951
Iterations 2900 DONE. Total time: 3.85598178307215
Iterations 3000 DONE. Total time: 3.95359458128611

mutate: new variable 'iterations' (integer) with 1,000 unique values and 0% NA
mutate: new variable 'iterations' (integer) with 1,000 unique values and 0% NA
mutate: new variable 'iterations' (integer) with 1,000 unique values and 0% NA

結構遅くてイライラしますが気長に待ちましょう。

結果

結果を確認してみましょう。まずはconvergenceをチェックします。$\lambda_i$については各$i$ごとにチェックしなければならず、結構面倒くさいので割愛します。

βのconvergence

{ggplot2}に含まれるfacet_wrap関数を活用して、変数ごとにconvergenceを見ていきたいと思います。まずは、プロット用に脳筋的にデータを加工していきます。

# βのposteriorを抽出
fit_df <- fit$beta_posterior 

# 各変数の名前を抽出
v_name <- colnames(fit_df[, -1]) 

# 各変数の名前を1000回ずつ入れた縦ベクトルを作る
v <- rep(v_name, each = nrow(fit_df)) %>% 
  as_tibble()


# 各変数の推定値を縦にどんどん追加していく
for (i in 2:8) {
  if (exists("value") == FALSE) {
    value <- fit_df[, i] %>% 
      as.matrix()
  } else{
      value2 <- fit_df[, i] %>% 
        as.matrix()
      value <- rbind(value, value2)
    }
}

# tibble化
value %>% 
  as_tibble() -> value

# 各変数の名前が入った縦ベクトルにvalue追加する
v %>% 
  bind_cols(value) %>% 
  rename(variable = value, 
         value = intercept) %>% 
  mutate(iterations = rep(1:1000, times = 7)) -> plot_df

# 確認
head(plot_df)
tail(plot_df)

rename: renamed one variable (variable)

mutate: new variable 'iterations' (integer) with 1,000 unique values and 0% NA

# A tibble: 6 x 3
  variable   value iterations
  <chr>      <dbl>      <int>
1 intercept -103.           1
2 intercept  -76.8          2
3 intercept  -83.7          3
4 intercept  -63.4          4
5 intercept -115.           5
6 intercept  -68.5          6

# A tibble: 6 x 3
  variable value iterations
  <chr>    <dbl>      <int>
1 smoke    -8.07        995
2 smoke    -7.58        996
3 smoke    -8.25        997
4 smoke    -9.25        998
5 smoke    -7.94        999
6 smoke    -7.05       1000

早速、プロットしてみましょう。

plot_df %>% 
  ggplot(aes(x = iterations, y = value)) + 
  geom_line() +
  facet_wrap(~ variable, scales = "free") +
  theme_light() +
  xlab("Iterations") + ylab("Estimated") +
  ggtitle("Convergence check for β")

大体収束してそうな気がします。

σ²のconvergence

次に、σ²についてのconvergenceを見てみましょう。

fit$sigma2_posterior %>% 
  ggplot(aes(x = iterations, y = sigma2)) +
  geom_line() +
  theme_light() +
  xlab("Iterations") + ylab("Estimated") +
  ggtitle(expression(paste("Convergence check for ", sigma^2)))

こちらもよさそうです。

Posterior distributions

posteriorsも確認しましょう。先ほど作ったplot_dfというデータフレームを用いて作図します。また、各変数ごとにposterior meanを図示するために、数ステップ踏みます。

# 変数ごとにposterior meanを計算
plot_df %>% 
  group_by(variable) %>% 
  summarise(mean = mean(value)) -> mean_df

group_by: one grouping variable (variable)

summarise: now 7 rows and 2 columns, ungrouped

plot_df %>% 
  ggplot(aes(x = value)) + 
  geom_density(alpha = .3, fill = "black") +
  facet_wrap(~ variable, scales = "free") +
  geom_vline(data = mean_df, aes(xintercept = mean), linetype = "dashed") +
  theme_light() +
  xlab("Estimated") + ylab("Density") +
  ggtitle("Posterior distributions with posterior mean")

density中央付近にある破線はposterior meanを表しています。interceptの値がデカすぎますが、OLSでもほとんど同じ位の推定値だったので見逃します。

また、通常のBayesian linear regressionの結果との比較をするために、以下のようなプロットを作ってみました。

  ここから見ると、Heteroskedasticityを考慮したモデルの方がposteriorの裾がやや長いことが確認できます。上手く実装できているか不安ですが（（（（（（