[Rust] 高次元 Gauss 分布 (正規分布) からのサンプリング

この記事は Rustその2 Advent Calendar 2019 12/2 の記事です.

Gauss 分布 (正規分布) の重要性・必要性は改めて述べるまでもないでしょう. 本記事では Rust で Gauss 分布からのサンプリングを行います.

動作確認は Rust 1.39 (stable-x86_64-unknown-linux-gnu) で, 以下のクレートに依存します.

Cargo.toml

[dependencies]
rand = "0.7.2"
rand_distr = "0.2.2"
ndarray = "0.13.0"
ndarray-linalg = { version = "0.12", features = ["intel-mkl"] }
ndarray-rand = "0.11.0"

1 次元分布

まずは 1 次元 Gauss 分布 $N [ \mu, \sigma^2 ]$ からのサンプリングを考えます.

Box-Muller 法

自分で実装するなら Box–Muller 法が一番手っ取り早いです. 引数は PRNG &mut rng, 平均 mu, 分散 sigma2 です.

use rand::Rng;
use std::f64::consts::PI;

pub fn box_muller<R: Rng>(rng: &mut R, mu: f64, sigma2: f64) -> f64 {
    let u1 = rng.gen::<f64>();
    let u2 = rng.gen::<f64>();
    
    mu + (-2.0*u1.ln()*sigma2).sqrt() * (2.0*PI*u2).cos()
}

なお Box-Muller 法は一度に 2 個の独立なサンプルを生成する (もうひとつは上のコードで cos を sin に置き換えたもの) ので, 効率よくサンプリングしたい場合は有効に使うべきです.

rand_distr

実際のところ, rand 公式が rand_distr クレートを用意してくれているため, 自前で Box-Muller 法を実装する必要はありません. なおこれはもともと rand クレートに含まれていましたが別クレートに分離されました. この経緯のため rand_distr::Distribution と rand::distributions::Distribution は同じもので, どちらを use しても大丈夫です.

use rand::Rng;
use rand_distr::{Normal, Distribution};

pub fn use_distr<R: Rng>(rng: &mut R, mu: f64, sigma2: f64) -> f64 {
    let normal = Normal::new(mu, sigma2.sqrt()).unwrap();
    normal.sample(rng)
}

Normal::new 関数には平均と 標準偏差 を渡します. なお標準正規分布なら rand_distr::StandardNormal を使うべきです.

高次元分布

高次元ベクトルの各成分が同一の確率分布からの独立なサンプルであるならば, 1 次元の場合の方法を繰り返すだけですし, ndarray-rand クレートもあります. 問題は相関 (共分散行列の非対角成分) があるときです.

Cholesky 分解

高次元 Gauss 分布は平均 $\mu$, 共分散行列 $C$ により特定されます. ここで共分散行列は正定値対称行列である必要があります. この場合, 共分散行列 $C$ の Cholesky 分解

$$C = L L^\dagger$$

を利用して, 標準正規分布からの独立なサンプル $x$ を $L x + \mu$ へと変換すれば求める Gauss 分布からのサンプルになります.

use rand::Rng;
use rand_distr::StandardNormal;
const NORMAL: StandardNormal = StandardNormal;

use ndarray::{Array1, Array2};
use ndarray_linalg::{cholesky::Cholesky, lapack::UPLO};
use ndarray_rand::RandomExt;

pub fn use_cholesky<R: Rng>(rng: &mut R, mu: &Array1<f64>, cov: &Array2<f64>) -> Array1<f64> {
    let l = cov.cholesky(UPLO::Lower).unwrap();
    let x = Array1::<f64>::random_using(mu.len(), NORMAL, rng);

    l.dot(&x) + mu
}

必要な型・トレイト等の準備が長いですが, 関数の中身は実質 3 行です. Cholesky 分解に失敗する可能性があるなら unwrap() を ? にする (戻り値の型は Result<Array1<f64>, ndarray_linalg::error::LinalgError>) か, 次の固有値分解を用いる方法を使ってください.

固有値分解

Cholesky 分解は正定値 Hermite 行列にしか適用できないため, 共分散行列に 0 固有値が含まれると上のコードはパニックを起こします. そんな状況あるのかと思うかもしれませんが, $n$ 自由度系だと思っていたものが実は冗長な自由度を含んでいた, だとか, Gauss 過程でサンプル点を密に取り過ぎた, などの場合に起こり得ます. そのときは素直に固有値分解して正でない固有値は確率 1 で平均値に一致すると要求することでサンプリングが可能になります.

use rand::Rng;
use rand_distr::{StandardNormal, Distribution};

use ndarray::{Array1, Array2};
use ndarray_linalg::{Eigh, lapack::UPLO};

pub fn use_eigh<R: Rng>(rng: &mut R, mu: &Array1<f64>, cov: &Array2<f64>) -> Array1<f64> {
    let normal = StandardNormal;
    
    let (w, v) = cov.eigh(UPLO::Lower).unwrap();
    let x = w.iter().map(|s2| {
        if s2.is_sign_positive() {
            Distribution::<f64>::sample(&normal, rng) * s2.sqrt()
        } else {
            0.
        }
    }).collect::<Array1<_>>();

    v.dot(&x) + mu
}