rstan の plot のパラメータ名（縦軸のラベル）を変更する

rstan を使っていると、プロットのパラメータ名を変更したくなるときがあります。

例えば、パラメータをベクトルで定義しているが、個々のパラメータに意味がある場合などです。

具体例を見てみましょう。

以下の例は、正規分布のパラメータを配列 theta で定義しています。一方で、それぞれのパラメータは、 theta[1] は平均、theta[2] は標準偏差という意味を持っています。

library(rstan)

set.seed(314)
# 平均0,標準偏差1の正規分布から20個のデータを生成
x <- rnorm(20, mean = 0, sd = 1)

# 正規分布の推定コード
model_code <- 
"
data {
  int N;
  vector[N] x;
}
parameters {
  real theta[2];
}
model {
  // 正規分布。平均 theta[1], 標準偏差 theta[2]
  x ~ normal(theta[1], theta[2]);
}
"

# Stan に渡すデータ
data <- list(N = 20, x = x)

# 推定の実行
fit <- stan(model_code = model_code, data = data)

このとき、パラメータ分布をプロットすると、パラメータ名は theta に固定されてしまいます。

# パラメータの事後分布の描画
plot(fit)

これのプロットのパラメータ名を、無機質な theta から、意味のある mu と sigma に変えたい。

それには次のようにします。

# 変えたいパラメータ名を指定
param_names <- c("mu", "sigma")

nparams <- 2  # パラメータ数
breaks <- seq(nparams, 1, by = -1)
plot(fit) +
  ggplot2::scale_y_continuous(breaks = breaks, labels = param_names, 
                              limits = c(0.5, nparams + 1))

rstan の plot のパラメータ名を変えることができました。

ポイントは、scale_y_discrete ではなく scale_y_continuous を使うところです。

ちなみに、y軸のスケールを上書きするので次のようなメッセージが出ます。

Scale for 'y' is already present. Adding another scale for 'y', which will replace the existing scale.

これを出さないようにする（無視する）関数 quietgg というのが rstan にはあるようです。

quietgg(
  plot(fit) +
    ggplot2::scale_y_continuous(breaks = breaks, labels = param_names, 
                                limits = c(0.5, nparams + 1))
)