ベータ分布のサンプリングアルゴリズムを読みとく

乱数とかサンプリングは好きですか？
私はといえば一様乱数（一様分布のサンプリング）ならメルセンヌ・ツイスタ、他の確率分布はうまいこと一様分布を歪ませてるのだろう。くらいのイメージだったのですが、先日ベイズ推定を実装していて、そういえばこのベータ分布ってどうやってサンプリングしてるんだろうと、 Apache Commons Math の実装を覗いてみたのですね。
ここで（第一種）ベータ分布の確率密度関数は $0 \leq x \leq 1$ の範囲で、

```math
f_X(x) = \frac{x^{a - 1} (1 - x)^{b - 1}}{B(a, b)} \\
B(a, b) = \int^1_0 x^{a - 1} (1 - x)^{b - 1} dx
```

というものでした。
なので、一様乱数にそんな感じの数式変換をしているのかなと思ったのですが、実際のところは……

```java
private static double algorithmBB(RandomGenerator random,
                                  final double a0,
                                  final double a,
                                  final double b) {
    final double alpha = a + b;
    final double beta = FastMath.sqrt((alpha - 2.) / (2. * a * b - alpha));
    final double gamma = a + 1. / beta;

    double r;
    double w;
    double t;
    do {
        final double u1 = random.nextDouble();
        final double u2 = random.nextDouble();
        final double v = beta * (FastMath.log(u1) - FastMath.log1p(-u1));
        w = a * FastMath.exp(v);
        final double z = u1 * u1 * u2;
        r = gamma * v - 1.3862944;
        final double s = a + r - w;
        if (s + 2.609438 >= 5 * z) {
            break;
        }

        t = FastMath.log(z);
        if (s >= t) {
            break;
        }
    } while (r + alpha * (FastMath.log(alpha) - FastMath.log(b + w)) < t);

    w = FastMath.min(w, Double.MAX_VALUE);
    return Precision.equals(a, a0) ? w / (b + w) : b / (b + w);
}
```
[commons-math/src/main/java/org/apache/commons/math3/distribution/BetaDistribution.java](https://github.com/apache/commons-math/blob/3.6.1-release/src/main/java/org/apache/commons/math3/distribution/BetaDistribution.java#L321)

となっていて、こんなわけのわからないものがベータ分布？？？？？
という気持ちになったので、アルゴリズムの元論文『[Generating Beta Variates With Nonintegral Shape Parameters](https://dl.acm.org/citation.cfm?id=359482)』を眺めながら、書かれていない数式展開も追ってみたよという記事です。

# 基本アルゴリズム

まずベータ分布からサンプリングする基本アルゴリズムは、

---

Algorithm BA ($a, b > 0$)
Initialization: Set $\alpha = a + b$.
If $\min(a, b) \leq 1$, set $\beta =
\max (\frac{1}{a}, \frac{1}{b})$; otherwise set $\beta = \sqrt{\frac{\alpha - 2}{2ab - \alpha}}$.
Set $\gamma = a + \frac{1}{\beta}$.

1. Generate random numbers $U_1, U_2$ and set $V = \beta \log \frac{U_1}{1 - U_1}, W = ae ^V$.
2. If $\alpha \log \frac{\alpha}{b + W} + \gamma V - 1.3862944 < log (U_1^2U_2)$, go to 1. (The constant is $\log 4$.)
3. Deliver $X = \frac{W}{b + W}$. 

---

というもので、これを改良したアルゴリズム BB を実装したのが、上記の引用コードになります。
で、このアルゴリズム BA を読んでもベータ分布の気配は感じられないと思うのですが、これは次のような手順になっています。

- 一様乱数 $U_1$ から $W$ を算出する。
- $W$ について、一様乱数 $U_2$ により棄却サンプリングを行う
    - 受理された場合、 $\frac{W}{b}$ は第二種ベータ分布に従う。
- $W$ から $X$ を算出する。 $X$ は第一種ベータ分布に従う。

なるほどなー？
それでは数式を一つ一つ追ってみませう。

# 基本アルゴリズムの数式展開

まず $Y=\frac{W}{b}$ として、

```math
Y = \frac{a}{b} (\frac{U_1}{1 - U_1})^\beta \\
U_1 = \frac{Y^\frac{1}{\beta}}{Y^\frac{1}{\beta} + (\frac{a}{b})^\frac{1}{\beta}}
```

$\lambda=\frac{1}{\beta}, \mu = (\frac{a}{b})^\lambda$ とおくと、

```math
U_1 = \frac{Y^\lambda}{Y^\lambda + \mu} \\
\frac{\partial U_1}{\partial Y} = \frac{\lambda \mu Y^{\lambda - 1}}{(Y^\lambda + \mu)^2}
```

$Y$ が成す分布の確率密度関数は、

```math
f_Y(x) = f_{U_1}(\frac{x^\lambda}{x^\lambda + \mu}) \frac{\lambda \mu x^{\lambda - 1}}{(x^\lambda + \mu)^2} =  \frac{\lambda \mu x^{\lambda - 1}}{(x^\lambda + \mu)^2}
```

これが第二種ベータ分布の提案分布となります。
さて第二種ベータ分布の確率密度関数は、[確率密度関数の変換](http://www5e.biglobe.ne.jp/~emm386/2016/distribution/parameter01.html)をすると、

```math
f_{Y'}(x) = \frac{x^{a - 1}}{B(a, b) (1 + x)^{a + b}}
```

ここで $A = \frac{4 a^a b^b}{\lambda B(a, b)(a + b)^{a + b}}$ として第二種ベータ分布の受理確率を、

```math
\frac{f_{Y'}(Y)}{A f_Y(Y)} = \frac{(a + b)^{a + b} Y^{a - \lambda} (Y^\lambda + \mu)^2}{4 a^a b^b \mu (1 + Y)^{a + b}} = (\frac{a + b}{b + W})^{a + b} (\frac{W}{a})^{a + \lambda} \frac{1}{4 U_1^2}
```

とすると棄却判定が、

```math
(\frac{a + b}{b + W})^{a + b} (\frac{W}{a})^{a + \lambda} \frac{1}{4 U_1^2} < U_2 \\
(a + b)\log{\frac{a + b}{b + W}} + (a + \lambda) \log{\frac{W}{a}} - \log{4} - \log{(U_1^2)} < \log{U_2} \\
\alpha \log{\frac{\alpha}{b + W}} + \gamma V - \log{4} < \log{(U_1^2 U_2)} \\
```

となってアルゴリズム BA の 2 式を導くことができました。
よって $X$ が成す分布の確率密度関数は、

```math
f_X(x) = f_{Y'}(\frac{x}{1 - x})\frac{1}{(1 - x)^2} = \frac{x^{a - 1} (1 - x)^{b - 1}}{B(a, b)}
```

めでたく第一種ベータ関数になることが分かります。

# 基本アルゴリズムの試行回数

さて、棄却サンプリングの計算効率を上げるためには、なるべく棄却される確率を減らすことが重要になってきます。
ここでは $A f_Y(Y)$ からサンプリングして $f_{Y'}(Y)$ に入っていれば受理されるもので、 $A f_Y(Y), f_{Y'}(Y)$ どちらも積分値 $1$ なので、受理・棄却の試行回数の期待値は $A$ となります。
そのため $A$ をなるべく小さくしたいのですが、棄却サンプリングを成立させるためには受理確率 $\frac{f_{Y'}(X)}{A f_Y(x)}$ が $1$ 以下でないとならないため、そこが腕の見せ所となるわけですね。このアルゴリズムでは $A$ が依存してる

```math
\beta =
\begin{cases}
  \max (\frac{1}{a}, \frac{1}{b}) & \min(a, b) \leq 1 \\
  \sqrt{\frac{\alpha - 2}{2ab - \alpha}} & \min(a, b) > 1
\end{cases}
```

がポイントになっていて $A$ は $4$ 以下に抑えられているみたいなのですが、この式自体をどうやって導出したのかまでは分かりませんでした。。。
なにかしらご存知な方はコメントいただけると＞＜

# 改良アルゴリズム

さて、上記の引用コードが実装していた改良アルゴリズムは、

---

Algorithm BB ($\min(a_o, b_0) > 1$)
Initialization: Set $a = \min(a_0, b_0), b = \max(a_0, b_0),$
$\alpha = a + b, \beta = \sqrt{\frac{\alpha - 2}{2ab - \alpha}}, \gamma = a + \frac{1}{\beta}$.

1. Generate random numbers $U_1, U_2$ and
set $V = \beta \log\frac{U_1}{1 - U_1}, W = a e^v, Z = U_1^2 U_2,$
$R = \gamma V - 1.3862944, S = a + R - W$. (The constant is $\log 4$.)
2. If $S + 2.609438 \geq 5Z$, go to 5. (The constant is $1 + \log 5$.)
3. Set $T= \log Z$. If $S \geq T$, go to 5.
4. If $R + \alpha \log{\frac{\alpha}{b + W}} < T$, go to 1.
5. If $a = a_0$ deliver $X = \frac{W}{b + W}$ ; otherwise deliver $X= \frac{b}{b + W}$.

---

となっていて $a \leq b$ という制限をかけることで $A$ をなるべく小さくしたり、対数関数や指数関数の呼び出し回数をなるべく減らすため受理の十分条件による判定を挟んだり、という工夫をしています。涙ぐましいですね。
そしてアルゴリズム BB に当てはまらない場合は、

---

Algorithm BC ($\min(a_0, b_0) \leq 1$)
Initialization: Set $a = \min(a_0, b_0), b = \max(a_0, b_0),$
$\alpha = a + b, \beta = \frac{1}{b}, \delta = 1 + a - b,$
$\kappa_1 = \delta \frac{0.0138889+0.0416667b}{a \beta-0.777778}, \kappa_2 = 0.25 + (0.5 + \frac{0.258}{\delta})b$.

1. Generate random numbers $U_1, U_2$, and if $U_1 \geq \frac{1}{2}$ go to 3.
2. Set $Y= U_1 U_2, Z = U_1 Y$. If $0.25 U_2 + Z - Y \geq \kappa_1$, go to 1; otherwise go to 5.
3. Set $Z = U_1^2 U_2$. If $Z \leq 0.25$, set $V = \beta \log \frac{U_1}{1-U_1}, W = a e^v$, and go to 6.
4. If $Z_\geq \kappa_2$ go to 1.
5. Set $V = \beta \log \frac{U_I}{1-U_1}, W = ae^v$.
If $\alpha (\log \frac{\alpha}{b + W} + V) - 1.3862944 < \log Z$, go to 1.
6. If $a = a_0$ deliver $X = \frac{W}{b + W}$ ; otherwise deliver $X= \frac{b}{b + W}$.

---

# その他のアルゴリズム

論文『[Evaluation of Beta Generation Algorithms](https://www.researchgate.net/publication/220504775_Evaluation_of_Beta_Generation_Algorithms)』では、様々なベータ分布のサンプリングアルゴリズムを性能比較しています。
その中で驚いたのが、一様乱数を $a + b -1$ 個生成して $a$ 番目に大きい値を取ると、それがベータ分布に従うというものです。 $a, b$ が整数の時にしか使えないですが、すごく綺麗ですね。