[Rust] MCMC 法の実装

本記事では Rust で MCMC 法 (特に Metropolis 法) を実装し, イテレータとして使えるようにします.

MCMC 法とは

$n$ 次元空間 $\mathbb{R}^n$ 上の任意の確率分布 $p: \mathbb{R}^n \to \mathbb{R}$ が与えられたとします. Markov 連鎖モンテカルロ法 (MCMC 法) はこの確率分布 $p$ に従ってサンプル列 $\{ x_n \}_{n \in \mathbb{N}}$ を数値的に得るアルゴリズム (のクラス) を指す名称です. 基本的にはモンテカルロ法の一種ですが, サンプル列を Markov 連鎖として生成することでこの名称がつきました.

MCMC 法の特徴

MCMC 法は性質のよくわからない確率分布からのサンプル列を得ることで, 平均や分散, 最頻値, 相関係数といった興味ある量を数値的に求めるために用いられます. この手法には次の利点があります.

• 求める確率分布 $p$ が規格化されている必要はない. このため, 統計力学における Boltzmann 分布, あるいは Bayes 確率分布からのサンプリングに威力を発揮します.
• 次元の呪いに耐性がある. つまり, 通常のモンテカルロ法では空間次元 $n$ が増大すると指数関数的に効率が悪化するような場合でも, 比較的効率よくサンプリングを行うことができます.

ただし, 状況によってアルゴリズムの選択や収束性の判定, サンプルの相関など気を遣う点がいくつもあります. 教科書としては 伊庭「計算統計II」 などがあります.

Metropolis 法

MCMC 法と一口に言っても多様なアルゴリズムがありますが, ここでは最も基礎的な Metropolis 法を扱います. これは次のアルゴリズムに従って Markov 連鎖を生成すると, それが求める確率分布 $p$ のサンプル列になっているというものです. 証明は Markov 連鎖のエルゴード定理によります.

1. 遷移確率 $q ( x' | x )$ に従って, 現在の点 $x \in \mathbb{R}^n$ から次の点 の候補 $x' \in \mathbb{R}^n$ を生成する.
2. 比 $r := p ( x' ) / p ( x )$ を計算し, 確率 $\min \{ 1, r \}$ で候補 $x'$ を採用する. そうでなければ候補 $x'$ は棄却し, もとの点 $x$ を次の点として採用する.

遷移確率 $q$ としては点 $x$ を平均とする Gauss 分布や一様分布などが可能です. また, 点 $x \in \mathbb{R}^n$ の $n$ 成分を一括で更新する必要はなく, 各成分を順番に更新すればよいです.

Rust 実装

既に rmcmc, emcee といったクレートがありますが, 自前で実装します.

方針

MCMC 法は確率分布からのサンプリングを行うものですから, rand のように PRNG を用意して .gen() メソッドを叩くようにする (つまり rand::Rng のようなものを実装する) のが素直な方針かもしれません. ただ, 実用上サンプル列の統計処理が主眼となることが多いと思うので, むしろイテレータとして実装する方が便利でしょう. そうするとこんな感じに使えます:

let mcmc = Metropolis::new(
    |x: [f64; 3]| -> f64 { x.iter().map(|i| -i * i/2.).sum::<f64>().exp() }
);

for x in mcmc.step_by(2).take(8) {
    println!("{:?}", x);
}

あるいは, .collect() しておいて平均や分散を計算することも簡単です. なので本記事ではイテレータとして MCMC 法を実現する方針でいきます.

なお, MCMC はモンテカルロ法ですから当然 PRNG が必要です. ここでは rand = 0.7.2 を使用します.

実装

まずはイテレータとなるべき struct を用意します. これは, 現在の Markov 連鎖の状態 x, 目標の確率分布関数 prob, PRNG rng を保持すべきです.なお状態空間の次元 DIM はコンパイル時定数として処理しています.

const DIM: usize = 3;

pub struct Metropolis<T, F> {
    rng: rand::rngs::ThreadRng,
    x: [T; DIM],
    prob: F,
}

Metropolis インスタンスを初期化する new() 関数くらいは用意しておきます. 引数として求める確率分布を渡すようにします.

impl<T, F> Metropolis<T, F>
where
    T: Default + Copy,
    F: Fn([T; DIM]) -> T,
{
    pub fn new(prob: F) -> Self {
        Metropolis {
            rng: rand::thread_rng(),
            x: [T::default(); DIM],
            prob,
        }
    }
}

最後に本題の std::iter::Iterator の実装です. 型 T としては f32 または f64 が想定されていて, num-traits = "0.2.8" から必要なトレイトを導入しています.

use rand::Rng;
use std::ops::{AddAssign, Neg, Div};
use num_traits::{Zero, One};
use std::iter::Iterator;

impl<T, F> Iterator for Metropolis<T, F>
where
    T: Copy + PartialOrd + AddAssign + Neg<Output=T> + Div<Output=T> 
        + Zero + One + rand::distributions::uniform::SampleUniform,
    F: Fn([T; DIM]) -> T,
{
    type Item = [T; DIM];

    fn next(&mut self) -> Option<Self::Item> {
        for i in 0..DIM {
            let mut y = self.x;
            y[i] += self.rng.gen_range::<T, T, T>(-T::one(), T::one());

            let r = (self.prob)(y) / (self.prob)(self.x);

            if  r > self.rng.gen_range::<T, T, T>(T::zero(), T::one()) {
                self.x = y;
            }
        }

        Some(self.x)
    }
}

これで完成です.