# モデルの書き方

BayesRTMB では、`rtmb_lm()` や `rtmb_glmer()` のようなラッパー関数を使うだけでなく、`rtmb_code()` を使って自分でモデルを書くことができます。

自分でモデルを書くと、既存のラッパー関数では表現しにくい潜在変数モデル、混合分布モデル、独自の階層モデル、生成量を含むモデルなどを柔軟に定義できます。

このページでは、最小の二項モデルから始めて、回帰モデル、階層モデル、順序モデル、混合分布モデルへ進みながら、`rtmb_code()` の書き方を説明します。

# 最小モデル

最初に、もっとも単純な二項分布モデルを書いてみます。10回の試行のうち6回成功した、というデータを考えます。

    library(BayesRTMB)

    Trial <- 10
    Y <- 6

    data_list <- list(
      Trial = Trial,
      Y = Y
    )

モデルは次のように書けます。

    code_binom <- rtmb_code(
      parameters = {
        theta <- Dim(lower = 0, upper = 1)
      },
      model = {
        Y ~ binomial(Trial, theta)
        theta ~ beta(1, 1)
      }
    )

ここで、`theta` は成功確率です。確率なので `lower = 0, upper = 1` という制約をつけています。

`model` ブロックには、データの尤度とパラメータの事前分布を書きます。

    Y ~ binomial(Trial, theta)
    theta ~ beta(1, 1)

このように、BayesRTMB では Stan に近い感覚で、確率モデルをそのまま書くことができます。

モデルオブジェクトは `rtmb_model()` で作成します。

    mdl_binom <- rtmb_model(data = data_list, code = code_binom)

作成したモデルに対して、MAP推定、MCMC、変分推論を実行できます。

    fit_map  <- mdl_binom$optimize()
    fit_mcmc <- mdl_binom$sample()
    fit_vb   <- mdl_binom$variational()

# rtmb_code の構造

`rtmb_code()` は、いくつかのブロックに分けてモデルを書きます。

    code <- rtmb_code(
      setup = {
        # データから定数や前処理済みオブジェクトを作る
      },
      parameters = {
        # 推定するパラメータを宣言する
      },
      transform = {
        # パラメータから派生量を作る
      },
      model = {
        # 尤度と事前分布を書く
      },
      generate = {
        # 推定後に計算したい量を書く
      }
    )

すべてのブロックが必須ではありません。最小モデルでは `parameters` と `model` だけでも十分です。ただし、モデルが少し複雑になると、`setup`、`transform`、`generate` を使い分けることで、コードが読みやすくなります。

## setup

`setup` は、モデル評価の前に一度だけ実行される前処理ブロックです。データの次元、平均、標準偏差、デザイン行列など、推定中に何度も変わらないものをここで作ります。

    setup = {
      N <- length(Y)
      P <- ncol(X)
      X_mean <- apply(X, 2, mean)
      X_sd <- apply(X, 2, sd)
    }

`setup` に書くとよいものは、主に次のようなものです。

- `N`, `P`, `K` などの次元
- デザイン行列やグループ番号
- 中心化済みの説明変数
- 事前分布のスケール
- データから計算される十分統計量

最近の BayesRTMB では、`model` ブロックをできるだけ読みやすくするために、前処理は `setup` に寄せる書き方を推奨しています。

## parameters

`parameters` は、推定するパラメータを宣言するブロックです。`Dim()` を使って、次元や制約を指定します。

    parameters = {
      alpha <- Dim()
      beta <- Dim(P)
      sigma <- Dim(lower = 0)
    }

`Dim()` の基本は次のとおりです。

書き方 | 意味
--- | ---
`Dim()` | スカラー
`Dim(P)` | 長さ `P` のベクトル
`Dim(c(N, P))` | `N x P` 行列
`Dim(c(I, J, K))` | 3次元以上の配列
`Dim(lower = 0)` | 正の値に制約
`Dim(lower = 0, upper = 1)` | 0から1の範囲に制約
`Dim(G, random = TRUE)` | ランダム効果として扱う

## transform

`transform` は、パラメータやデータから派生量を作るブロックです。線形予測子、平均構造、標準化効果量、相関行列などを書く場所です。

    transform = {
      mu <- alpha + X %*% beta
    }

`transform` で作った量は、`model` ブロックの中で使えます。また、推定後の出力にも含まれます。MAP推定では、デルタ法によって標準誤差や信頼区間も計算されます。

ただし、`transform` ブロックの中で作ったすべての途中変数を保存したいとは限りません。その場合は、`report()` を使って、明示的に保存したい量だけを指定できます。

    transform = {
      eta <- alpha + X %*% beta
      mu <- inv_logit(eta)
      report(mu, beta)
    }

この例では、`eta` は計算途中で使うだけなので出力には含めず、`mu` と `beta` だけを保存対象にしています。

## model

`model` は、尤度と事前分布を書くブロックです。

    model = {
      Y ~ normal(mu, sigma)
      alpha ~ normal(0, 10)
      beta ~ normal(0, 2.5)
      sigma ~ exponential(1)
    }

`~` の左側に確率変数、右側に分布を書きます。`Y ~ normal(mu, sigma)` は、観測値 `Y` が平均 `mu`、標準偏差 `sigma` の正規分布に従う、という意味です。

同じ内容は、対数確率を明示的に足し合わせる形でも書けます。

    model = {
      lp <- lp + normal_lpdf(Y, mu, sigma)
      lp <- lp + normal_lpdf(alpha, 0, 10)
      lp <- lp + normal_lpdf(beta, 0, 2.5)
      lp <- lp + exponential_lpdf(sigma, 1)
    }

つまり、`Y ~ normal(mu, sigma)` と `lp <- lp + normal_lpdf(Y, mu, sigma)` は、同じモデルを表す2つの書き方です。通常は `~` を使ったほうが読みやすいですが、特殊な尤度を自分で組み立てたいときには `lp <- lp + ..._lpdf(...)` の形が便利です。Stanと違って、`BayesRTMB`パッケージでは`Y ~ normal(mu, sigma)`と書いても正規化定数は計算されます。なのでbridgesamplingをしたいときでも`~`を使った記法を使って問題ありません。

ただし、`lp` は対数確率を足し合わせるための予約語です。`parameters` ブロックで `lp = Dim()` のように宣言したり、一般の作業用変数として使ったりしないでください。

## generate

`generate` は、推定後に計算したい量を書くブロックです。事後予測チェックや予測値の生成に便利です。

    generate = {
      Y_rep <- rnorm(length(Y), mu, sigma)
    }

`generate` に書いた計算は、尤度評価には入りません。そのため、重い計算を書いても推定速度には直接影響しません。

`generate` ブロックは、事後予測チェックや予測値の計算に便利です。

    code_ppc <- rtmb_code(
      setup = {
        N <- length(Y)
        P <- ncol(X)
      },
      parameters = {
        alpha <- Dim()
        beta <- Dim(P)
        sigma <- Dim(lower = 0)
      },
      transform = {
        mu <- alpha + X %*% beta
      },
      model = {
        Y ~ normal(mu, sigma)
        alpha ~ normal(0, 10)
        beta ~ normal(0, 2.5)
        sigma ~ exponential(1)
      },
      generate = {
        Y_rep <- rnorm(N, mu, sigma)
        residual <- Y - mu
      }
    )

`Y_rep` は事後予測分布からの複製データ、`residual` は残差です。これらは推定後に取り出して、モデル診断に使えます。

モデル定義時に書き忘れた場合でも、推定後に `generated_quantities()` で追加できます。

    new_generate <- rtmb_code(
      generate = {
        abs_residual <- abs(Y - mu)
      }
    )

    fit_reg <- mdl_reg$optimize()
    fit_reg$generated_quantities(new_generate)

# 自作分布を使う

`setup` ブロックには、データの前処理だけでなく、モデル内で使う補助関数も定義できます。連続分布の場合は、関数名を `_lpdf` で終わる形にしておくと、`model` ブロックでサンプリング記法として使えます。

たとえば、`exgaussian_lpdf()` を `setup` で定義すると、`model` ブロックでは次のように書けます。

    Y ~ exgaussian(mu, sigma, lambda)

これは内部的には、次のように対数密度を足すのと同じです。

    lp <- lp + exgaussian_lpdf(Y, mu, sigma, lambda)

次の例では、秒単位の反応時間を ex-Gaussian 分布から生成し、同じ分布で推定します。ここでは `lambda` を指数分布成分の rate として使います。

    library(BayesRTMB)

    simulate_exgaussian <- function(N, mu, sigma, lambda) {
      rnorm(N, mean = mu, sd = sigma) + rexp(N, rate = lambda)
    }

    set.seed(1234)

    true <- c(mu = 0.45, sigma = 0.06, lambda = 8)
    rt <- simulate_exgaussian(
      N = 500,
      mu = true["mu"],
      sigma = true["sigma"],
      lambda = true["lambda"]
    )

    hist(rt)

モデルコードは次のように書けます。

    code_exgaussian <- rtmb_code(
      setup = {
        Y <- rt

        exgaussian_lpdf <- function(x, mu, sigma, lambda, sum = TRUE) {
          z <- (x - mu) / sigma - lambda * sigma

          res <-
            log(lambda) +
            lambda * (mu - x) +
            0.5 * (lambda * sigma)^2 +
            RTMB::pnorm(z, log.p = TRUE)

          if (sum) sum(res) else res
        }
      },

      parameters = {
        mu     <- Dim(lower = 0)
        sigma  <- Dim(lower = 0)
        lambda <- Dim(lower = 0)
      },

      model = {
        Y ~ exgaussian(mu, sigma, lambda)

        mu     ~ exponential(1)
        sigma  ~ exponential(1)
        lambda ~ exponential(1 / 10)
      },

      generate = {
        log_lik <- exgaussian_lpdf(Y, mu, sigma, lambda, sum = FALSE)
        report(log_lik)
      }
    )

このモデルでは、初期値を観測データのスケールに合わせて指定しておくと安定しやすくなります。

    init_exgaussian <- list(
      mu = median(rt) - sd(rt) / 2,
      sigma = sd(rt) / 2,
      lambda = 1 / (sd(rt) / 2)
    )

    mdl_exgaussian <- rtmb_model(
      data = list(rt = rt),
      code = code_exgaussian,
      init = init_exgaussian
    )

まず MAP 推定で、既知の真値がどの程度回復されるかを確認できます。

    fit_map <- mdl_exgaussian$optimize(num_estimate = 5)

    est_map <- unlist(fit_map$estimate(pars = "parameters"))[names(true)]
    data.frame(
      parameter = names(true),
      true = as.numeric(true),
      estimate = as.numeric(est_map),
      error = as.numeric(est_map - true)
    )

同じモデルから MCMC や変分推論も実行できます。

    fit_mcmc <- mdl_exgaussian$sample(parallel = TRUE)
    fit_mcmc
    fit_mcmc$WAIC()

    fit_vb <- mdl_exgaussian$variational()
    fit_vb

WAIC を使いたい場合は、`generate` ブロックで観測ごとの `log_lik` を保存します。そのため、自作した `_lpdf` 関数を `sum = FALSE` で呼び出します。

# 回帰モデル

次に、回帰モデルを書きます。ここでは `debate` データを使い、満足度 `sat` を発言量 `talk`、パフォーマンス `perf`、スキル `skill` で説明します。

    data(debate)

    Y <- debate$sat
    X_names <- c("talk", "perf", "skill")
    X <- as.matrix(debate[, X_names])

    data_reg <- list(
      Y = Y,
      X = X
    )

`X` は `as.matrix()` で数値行列にしておくのが安全です。`data.frame` のまま `%*%` を使うと、列の型によってエラーになることがあります。

    code_reg <- rtmb_code(
      setup = {
        N <- length(Y)
        P <- ncol(X)
      },
      parameters = {
        alpha <- Dim()
        beta <- Dim(P)
        sigma <- Dim(lower = 0)
      },
      transform = {
        mu <- alpha + X %*% beta
      },
      model = {
        Y ~ normal(mu, sigma)
        alpha ~ normal(0, 10)
        beta ~ normal(0, 2.5)
        sigma ~ exponential(1)
      }
    )

このモデルでは、`beta <- Dim(P)` として、説明変数の数に応じた回帰係数ベクトルを宣言しています。

`par_names` を使うと、出力の名前を読みやすくできます。

    mdl_reg <- rtmb_model(
      data = data_reg,
      code = code_reg,
      par_names = list(beta = X_names)
    )

    mdl_reg$optimize()

`beta[1]`, `beta[2]` ではなく、`beta[talk]`, `beta[perf]`, `beta[skill]` のように表示されます。

## 中心化した回帰モデル

回帰モデルでは、説明変数を中心化しておくと、切片の推定が安定しやすくなります。切片が「説明変数がすべて0のときの平均」ではなく、「説明変数が平均的な値のときの平均」に近い意味を持つためです。

    data_reg2 <- list(
      Y = debate$sat,
      X = as.matrix(debate[, c("talk", "perf", "skill")])
    )

    code_reg2 <- rtmb_code(
      setup = {
        N <- length(Y)
        P <- ncol(X)
        X_mean <- apply(X, 2, mean)
        X_c <- X - rep(1, N) %*% t(X_mean)
      },
      parameters = {
        alpha_c <- Dim()
        beta <- Dim(P)
        sigma <- Dim(lower = 0)
      },
      transform = {
        alpha <- alpha_c - sum(X_mean * beta)
        mu <- alpha_c + X_c %*% beta
      },
      model = {
        Y ~ normal(mu, sigma)
        alpha_c ~ normal(0, 10)
        beta ~ normal(0, 2.5)
        sigma ~ exponential(1)
      },
      generate = {
        Y_rep <- rnorm(length(Y), mu, sigma)
      }
    )

ここでは、`setup` で中心化済みの `X_c` を作っています。`alpha_c` は中心化後の切片です。`transform` では、元の尺度での切片 `alpha` も計算しています。

    alpha <- alpha_c - sum(X_mean * beta)

このようにしておくと、推定では安定しやすい中心化切片 `alpha_c` を使いながら、必要に応じて元の尺度での切片 `alpha` も確認できます。

# 事前分布を書く

自分で `rtmb_code()` を書く場合、事前分布は `model` ブロックに明示的に書きます。

    model = {
      Y ~ normal(mu, sigma)
      alpha ~ normal(0, 10)
      beta ~ normal(0, 2.5)
      sigma ~ exponential(1)
    }

基本的な考え方は次のとおりです。

- 位置パラメータには `normal(0, scale)` を使うことが多い
- 正のスケールパラメータには `exponential(rate)` や `gamma()` を使う
- 確率パラメータには `beta(a, b)` を使う
- 制約と事前分布のサポートを合わせる

たとえば、`sigma` は正の値なので、宣言では `Dim(lower = 0)` とし、事前分布も `exponential()` のような正の分布を使います。

    parameters = {
      sigma <- Dim(lower = 0)
    }
    model = {
      sigma ~ exponential(1)
    }

# 制約つきパラメータ

`Dim()` の `type` を使うと、特殊な制約を持つパラメータを定義できます。

型 | 使い道
--- | ---
`type = "ordered"` | 昇順に並ぶカットポイント
`type = "simplex"` | 混合率やカテゴリ確率
`type = "centered"` | 和が0になる効果
`type = "CF_corr"` | 相関行列のコレスキー因子
`type = "centered_tri"` | 識別制約つきの座標行列

たとえば、順序ロジスティックモデルでは、カットポイントが昇順である必要があります。

    parameters = {
      cutpoints <- Dim(K - 1, type = "ordered")
    }

混合分布モデルの混合率には、和が1になる正のベクトルが必要です。

    parameters = {
      theta <- Dim(K, type = "simplex")
    }

# 階層モデル

階層モデルでは、ランダム効果を `random = TRUE` として宣言します。

ここでは、グループごとに切片が異なる線形モデルを書きます。

    Y <- debate$sat
    X_names <- c("talk", "perf", "skill")
    X <- as.matrix(debate[, X_names])
    group <- as.integer(as.factor(debate$group))
    G <- length(unique(group))

    data_hlm <- list(
      Y = Y,
      X = X,
      group = group,
      G = G
    )

    code_hlm <- rtmb_code(
      setup = {
        N <- length(Y)
        P <- ncol(X)
      },
      parameters = {
        alpha <- Dim()
        beta <- Dim(P)
        tau <- Dim(lower = 0)
        sigma <- Dim(lower = 0)
        r <- Dim(G, random = TRUE)
      },
      transform = {
        mu <- alpha + X %*% beta + tau * r[group]
      },
      model = {
        Y ~ normal(mu, sigma)
        r ~ normal(0, 1)
        alpha ~ normal(0, 10)
        beta ~ normal(0, 2.5)
        tau ~ exponential(1)
        sigma ~ exponential(1)
      }
    )

ここでは、`r` が標準正規分布に従うランダム効果で、`tau` がその標準偏差です。

    r <- Dim(G, random = TRUE)
    r ~ normal(0, 1)
    mu <- alpha + X %*% beta + tau * r[group]

`random = TRUE` にしたパラメータは、MAP推定ではラプラス近似の対象になります。

    mdl_hlm <- rtmb_model(
      data = data_hlm,
      code = code_hlm,
      par_names = list(beta = X_names)
    )

    fit_hlm <- mdl_hlm$optimize()
    fit_hlm$random_effects

ランダム効果の推定値は、`random_effects` に保存されます。

# 順序モデル

順序カテゴリカルデータには、`ordered_logistic()` を使えます。

    Y <- as.integer(debate$sat)
    X <- as.matrix(debate[, c("talk", "perf", "skill")])
    K <- length(sort(unique(Y)))

    data_ord <- list(
      Y = Y,
      X = X,
      K = K
    )

    code_ord <- rtmb_code(
      setup = {
        P <- ncol(X)
      },
      parameters = {
        beta <- Dim(P)
        cutpoints <- Dim(K - 1, type = "ordered")
      },
      transform = {
        eta <- X %*% beta
      },
      model = {
        Y ~ ordered_logistic(eta, cutpoints)
        beta ~ normal(0, 2.5)
        cutpoints ~ normal(0, 5)
      }
    )

`cutpoints` は必ず順序を持つ必要があるので、`type = "ordered"` としています。

二値データやカウントデータでは、確率そのものではなく、線形予測子を直接受け取る分布が便利です。

    Y ~ bernoulli_logit(eta)
    Y ~ poisson(exp(eta))

ロジットスケールや対数スケールで計算するため、数値的に安定しやすくなります。

# 混合分布モデル

混合分布モデルは、制約つきパラメータと複数初期値を使うモデリングのいい例です。

    set.seed(123)
    N <- 300
    K <- 3

    theta_true <- c(0.3, 0.2, 0.5)
    mu_true <- c(-3, 0, 4)
    sigma_true <- c(0.6, 1.0, 0.8)

    z <- sample(seq_len(K), size = N, replace = TRUE, prob = theta_true)
    Y <- rnorm(N, mu_true[z], sigma_true[z])

    data_mix <- list(
      Y = Y,
      K = K
    )

    code_mix <- rtmb_code(
      parameters = {
        theta <- Dim(K, type = "simplex")
        mu <- Dim(K)
        sigma <- Dim(K, lower = 0)
      },
      model = {
        Y ~ normal_mixture(theta, mu, sigma)
        mu ~ normal(0, 5)
        sigma ~ exponential(1)
      }
    )

`theta` は混合率なので、`type = "simplex"` を使います。

混合分布は局所解が生じやすいため、MAP推定では `num_estimate` を多めにして、複数の初期値から最もよい解を選ぶのが有効です。通常は、自分で初期値を細かく指定するよりも、まずは `init` を指定せずに `num_estimate` を増やすほうが手軽です。

    mdl_mix <- rtmb_model(data_mix, code_mix)
    fit_mix <- mdl_mix$optimize(num_estimate = 20)

MCMCを行う場合は、MAP推定で得られたパラメータ推定値を初期値として使うと安定しやすくなります。必要に応じて `init_jitter` で少しだけ揺らします。

    fit_mix_mcmc <- mdl_mix$sample(
      init = fit_mix$estimate("parameters"),
      init_jitter = 0.2
    )

混合分布では、ラベルスイッチングにも注意が必要です。成分の順序に意味がないため、MCMCでは成分ラベルが入れ替わることがあります。

# よくある注意点

## data.frame と matrix

行列積 `%*%` を使う場合、`X` は `matrix` にしておくのが安全です。

    X <- as.matrix(debate[, c("talk", "perf", "skill")])

`data.frame` のまま渡すと、列型によっては `X %*% beta` でエラーになります。

## model ブロックを複雑にしすぎない

`model` ブロックは、尤度と事前分布が見えるように書くのがおすすめです。

前処理は `setup`、派生量は `transform` に分けると、ユーザーが `print_code()` で見たときにも読みやすくなります。

## AD型に対する if 文

推定中のパラメータを使った条件分岐は、`RTMB`のAD計算で問題になることがあります。

    if (theta > 0) ...

のような書き方は避け、滑らかな関数や制約つきパラメータ化で表現するほうが安全です。

## 初期値

複雑なモデルでは、初期値が重要です。

    init <- list(
      alpha = mean(Y),
      beta = rep(0, ncol(X)),
      sigma = sd(Y)
    )

    mdl <- rtmb_model(data_reg, code_reg, init = init)

特に、混合分布モデル、因子分析、多次元展開法、項目反応理論のような潜在変数モデルでは、適切な初期値が推定の安定性に大きく影響します。

## random = TRUE

`random = TRUE` は、ランダム効果や潜在変数をラプラス近似で周辺化したいときに使います。

ただし、すべてのパラメータに付ければよいわけではありません。固定効果として推定したい回帰係数や切片には、通常 `random = TRUE` は付けません。

## AD対応の作業用コンテナ

多くの場合、もっとも速く安定しやすいのはベクトル化されたコードです。たとえば、`Y ~ normal(mu, sigma)` のように一度に尤度を書けるなら、その書き方を優先します。

一方で、`for` ループの中で要素を順番に代入しながら、あとから AD 型の値が入るベクトルや配列を作りたいことがあります。その場合は、`numeric()`、`matrix()`、`array()` の代わりに `rtmb_vector()` や `rtmb_array()` を使います。

    model = {
      eta <- rtmb_vector(0, N)
      for (i in seq_len(N)) {
        eta[i] <- alpha + X[i, ] %*% beta
      }
      Y ~ normal(eta, sigma)
    }

    model = {
      logit_x <- rtmb_array(0, dim = c(N_time, C, D))
      for (t in seq_len(N_time)) {
        for (c in seq_len(C)) {
          for (d in seq_len(D)) {
            logit_x[t, c, d] <- alpha[d] + beta[d] * X[t, c]
          }
        }
      }
    }

`rtmb_code()` の中では、通常、BayesRTMB がモデルパラメータから AD seed を自動的に用意するため、`seed` 引数を指定する必要はありません。ただし、これはあくまで代入を多用するコードのための補助です。同じ計算を行列演算やベクトル化で書ける場合は、テープ化の時間とAD評価の速度の両方の面で、ベクトル化した書き方を優先するのがおすすめです。

## 保存済み fit object の更新

古いバージョンの BayesRTMB で保存した fit object は、新しいバージョンで追加されたメソッドをそのまま持っていないことがあります。その場合は、`upgrade_fit()` を使うと、保存済みの MCMC、VB、MAP、classic の fit object を現在読み込まれているクラス定義で作り直せます。

    fit2 <- upgrade_fit(fit)

これはモデルを再推定する関数ではありません。保存されている推定結果を新しい fit object に移し替え、可能な場合は内部のモデルオブジェクトも更新します。

# 次に読む記事

`rtmb_code()` の内部的な構造をさらに詳しく知りたい場合は、内部構造 を参照してください。

既存のモデルを短く書きたい場合は、ラッパー関数の使い方 が便利です。

階層モデルや GLMM を wrapper で詳しく扱いたい場合は、階層モデル・GLMM・分散分析 を参照してください。

まずパッケージ全体をざっと試したい場合は、クイックスタート を参照してください。

分析中に fit object のメソッド、モデル比較、固定パラメータ、分布、AD テープ化の注意点を確認したい場合は、分析リファレンス を参照してください。