Julia による MCMC サンプリング

新時代の確率的プログラミング環境の構築に向けて

Process
Sampling
Julia
MCMC
俺のためのJulia入門
Author

司馬 博文

Published

7/03/2024

概要
Julia に存在する MCMC 関連のパッケージをまとめ,多くの MCMC のパッケージを支える,Turing ecosystem の基盤となる抽象的なフレームワーク MCMCChainsAbstractMCMC を概観する.

1 導入

具体的なサンプラーは次の稿も参照:

1.1 依存パッケージ一覧

  • AbstractMCMC.jl (GitHub / Juliapackages / HP) は MCMC サンプリングのための抽象的なインターフェースを提供するパッケージ.後述の Turing.jl エコシステムの基盤となる.

  • LogDensityProblems.jl:対数密度を扱うフレームワーク.

1.2 MCMC パッケージ一覧

  • AdaptiveMCMC.jl (Juliapackages / Docs) は適応的な乱歩 MH アルゴリズムを提供するパッケージ.

    • 関連に AdaptiveParticleMCMC.jl (GitHub) がある.これは SequentialMonteCarlo.jl (GitHub) に基づいている.
  • Mamba.jl (Docs / GitHub) Markov chain Monte Carlo (MCMC) for Bayesian analysis in julia

  • KissMCMC (Juliapackages / GitHub) は ‘Keep it simple, stupid, MCMC’ ということで,計量な MCMC を提供するパッケージ.

  • DynamicHMC.jl (GitHub) は NUTS サンプラーを提供するパッケージ.Tamás K. Papp による.

  • SGMCMC.jl (GitHub) は Oxford CSML チームによる,Bayesian deep learning に向けたサンプラーを提供するパッケージ.

1.3 連続時間 MCMC パッケージ一覧

連続時間 MCMC
  • ZigZagBoomerang.jl (GitHub / Juliapackages)

  • PDMP.jl (GitHub / Docs) は 2018 年まで Alan Turing Institute によって開発されていたパッケージ.

  • PiecewiseDeterministicMarkovProcesses.jl (GitHub / Docs / (Veltz, 2015) / HP of Dr. Veltz / Discource) は細胞生物学におけるモデリング手法としての PDMP を提供するパッケージである.

1.4 確率的プログラミング

代表的なものは次の2つである:

ベイズ推論のためのパッケージ

ただし,現状の Soss.jl v0.21.2 (6/22/2022) は FillArrays のバージョンを 0.13.11 以下に制限してしまうため(最新は 1.11.0),これが Turing.jl の最新バージョン v0.33.1 と衝突してしまう.

インストールは1つの環境にどちらか1つのみにすることが推奨される.

1.5 Turing ecosystem

Turing ecosystem 一覧
  • AdvancedPS (GitHub / Docs) は Turing.jl による粒子フィルターベースのサンプラーを提供するパッケージ.
  • AdvancedHMC (GitHub)
  • AdvancedMH (GitHub)
  • AdvancedVI (GitHub) 変分推論を提供するパッケージ.
  • Bijectors.jl (GitHub) 正規化流などによる分布の変換を提供するパッケージ.

(Storopoli, 2021) も参照.

2 AbstractMCMC.jl の枠組み

AbstractMCMC.jl は,DensityProblem をはじめとした AbstractModelAbstractSampler のデータ構造を提供する,Turing エコシステムの根幹部分を支えるパッケージである.

2.1 LogDensityProblem

AbstractMCMC.jl は,Tamás K. Papp による LogDensityProblem.jlwrapper を提供している.

sample 関数なども,AbstractModel の他に,longdensity オブジェクトに対するメソッドも定義されている.

3 MCMCChains.jlAbstractMCMC.jl

Turing によるエコシステムは,この MCMC のデザインパターンを利用している.

using MCMCChains
using StatsPlots

# Define the experiment
n_iter = 100
n_name = 3
n_chain = 2

# experiment results
val = randn(n_iter, n_name, n_chain) .+ [1, 2, 3]'
val = hcat(val, rand(1:2, n_iter, 1, n_chain))

# construct a Chains object
chn = Chains(val, [:A, :B, :C, :D])

# visualize the MCMC simulation results
plot(chn; size=(840, 600))

3.1 sample 関数

sample 関数は,MCMCCahins での実装AbstractMCMC での実装 との2つがある.

3.1.1 使い方

AbstractModel に対するデフォルトは次のように始まる:

function mcmcsample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    N::Integer;
    progress=PROGRESS[],
    progressname="Sampling",
    callback=nothing,
    discard_initial=0,
    thinning=1,
    chain_type::Type=Any,
    initial_state=nothing,
    kwargs...,
)

inisital_params を指定した場合:

function mcmcsample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    ::MCMCThreads,
    N::Integer,
    nchains::Integer;
    progress=PROGRESS[],
    progressname="Sampling ($(min(nchains, Threads.nthreads())) threads)",
    initial_params=nothing,
    initial_state=nothing,
    kwargs...,
)

3.2 autocor 関数

Chains オブジェクトに対する autocor 関数が,次のように定義されている:

function autocor(
    chains::Chains;
    append_chains = true,
    demean::Bool = true,
    lags::AbstractVector{<:Integer} = _default_lags(chains, append_chains),
    kwargs...
)
    funs = Function[]
    func_names = @. Symbol("lag ", lags)
    for i in lags
        push!(funs, x -> autocor(x, [i], demean=demean)[1])
    end

    return summarize(
        chains, funs...;
        func_names = func_names,
        append_chains = append_chains,
        name = "Autocorrelation",
        kwargs...
    )
end

References

Ge, H., Xu, K., and Ghahramani, Z. (2018). Turing: A language for flexible probabilistic inference. In International conference on artificial intelligence and statistics, AISTATS 2018, 9-11 april 2018, playa blanca, lanzarote, canary islands, spain, pages 1682–1690.
Storopoli, J. (2021). Bayesian statistics with julia and turing.
Veltz, R. (2015). A new twist for the simulation of hybrid systems using the true jump method.