Juliaを使ってみる(20/22)：更なる統計モデル

今回やりたいこと

　より発展的な統計モデルをJuliaで実装してみる。

Juliaで作って学ぶベイズ統計学 (KS情報科学専門書)

を参照する。

前回

power-of-awareness.com

1.　ポアソン回帰

　ハミルトニアン・モンテカルロ法を用いたいくつかの統計モデルを扱う。

1.1　ポアソン回帰

　ポアソン回帰は

で定式化する。ここでは平均ベクトルがで分散共分散行列がであるような多変量正規分布の確率密度関数である。または母数がであるようなポアソン分布の確率密度関数である。

using Distributions, PyPlot, ForwardDiff, LinearAlgebra

# n次元単位行列
eye(n) = Diagonal{Float64}(I, n)

# グラフの諸設定を行う関数
function set_options(ax, xlabel, ylabel, title;
                      grid = true, gridy = false, legend = false)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    if grid
        if gridy
            ax.grid(axis = "y")
        else
            ax.grid()
        end
    end
    legend && ax.legend()
end

# Hamiltonian Monte Carlo method
function HMC(log_p_tilde, μ₀; maxiter::Int64, L::Int64, ε::Float64)
    # leapfrog による値の更新
    function leapflog(grad, p_in,μ_in, L, ε)
        μ=μ_in
        p = p_in + 0.5 *ε*grad(μ)
        for l in 1: L-1
            μ +=ε*p
            p += ε*grad(μ)
        end
        μ += ε*p
        p += 0.5 *ε*grad(μ)
        p, μ
    end
    
    # 非正規化対数事後分布の勾配関数
    grad(μ) = ForwardDiff.gradient(log_p_tilde, μ)
    
    # サンプルを格納する配列
    D = length(μ₀)
    μ_samples = Array{typeof(μ₀[1]), 2}(undef, D, maxiter)
    
    # 初期サンプル
    μ_samples[:, 1] = μ₀
    
    # 受容されたサンプル数
    num_accepted = 1
    
    for i in 2:maxiter
        # 運動量の生成
        p_in = randn(size(μ₀))
        
        # リープフロッグ
        p_out, μ_out = leapflog(grad, p_in, μ_samples[:, i-1], L, ε)
        
        # 比率rの対数を計算
        μ_in = μ_samples[:, i-1]
        log_r = (log_p_tilde(μ_out) + 
                 logpdf(MvNormal(zeros(D),eye(D)), vec(p_out))) -
                (log_p_tilde(μ_in) + 
                 logpdf(MvNormal(zeros(D),eye(D)),vec(p_in)))
                
        # 確率rでサンプル受容
        is_accepted = min(1, exp(log_r)) > rand()
        new_sample = is_accepted ? μ_out : μ_in
        
        # 新サンプルを格納
        μ_samples[:,i] = new_sample
        
        # 受容された場合、合計を加算する
        num_accepted += is_accepted
    end
    
    μ_samples, num_accepted
end

function inference_wrapper_HMC(log_joint, params, w_init; maxiter::Int64 = 100_000, L::Int64 = 100, ε::Float64 = 1e-1)
    ulp(w) = log_joint(w, params...)
    HMC(ulp, w_init; maxiter = maxiter, L = L, ε = ε)
end

### 入力データ
X_obs = [-2,-1.5,0.5,0.7,1.2]
Y_obs = [0,0,2,1,8]

### 対数同時分布
log_joint(w, X, Y) = 
    sum(logpdf.(Poisson.(exp.(w[1] .* X .+ w[2])), Y)) +
    logpdf(Normal(0, 1.0), w[1]) + logpdf(Normal(0, 1.0), w[2])
parameters = (X_obs, Y_obs)

# 非正規化対数事後分布
ulp(x) = log_joint(w, parameters...)

# 初期値
w_init = randn(2)

# 標本サイズ
maxiter = 10000

# Hamiltonian Monte Carlo methodによるSampling
param_posterior, num_accepted = 
    inference_wrapper_HMC(log_joint, parameters, w_init,
                          maxiter = maxiter)

### trace plot
fig, axes = subplots(2,1,figsize=(8,4))
axes[1].plot(param_posterior[1,:])
set_options(axes[1], "iteration", "w₁", "w₁ sequence")
axes[2].plot(param_posterior[2,:])
set_options(axes[2], "iteration", "w₂", "w₂ sequence")
tight_layout()
println("acceptance rate = $(num_accepted/maxiter)")

### 事後分布からのサンプリング
fig, ax = subplots()
ax.scatter(param_posterior[1,:], param_posterior[2,:], alpha = 0.1)
set_options(ax, "w₁", "w₂", "samples from posterior")
tight_layout()

### 予測関数の可視化
xs = range(-4,4,length = 1000)

fig,ax = subplots()
for i in 1:size(param_posterior, 2)
    w₁,w₂ = param_posterior[:,i]
    
    # 指数関数を可視化
    f(x) = exp.(w₁*x + w₂)
    ax.plot(xs, f.(xs), "g", alpha = 0.1)
end

ax.plot(X_obs, Y_obs, "ko")
ax.set_xlim(extrema(xs))
ax.set_ylim([-1,15])
set_options(ax, "x", "y", "predictive distribution")

tight_layout()

次回

power-of-awareness.com