『ゼロからできるMCMC』 Chapter6 HMCをJuliaで

using Plots
using Random:seed!
using Distributions
##
# 函数を定義: 作用を計算 ( 定義 ) する
function calc_action(x)
    action = 0.5*x*x
end

# 函数を定義：ハミルトニアンHを計算する
function calc_h(x, p)
    ham = calc_action(x)
    ham = ham + 0.5*p*p
end

# 函数を定義：ハミルトニアンの微分dH/dxを計算する
function calc_dhdx(x)
    delh = x
end


function hmc1d(;niter=10^5, nτ=40, dτ=1.0)

    # ================= 準備 =================
    seed!(12345)
    x       = 0.0             # 初期配位
    naccept = 0               # 受理回数カウンター
    res     = zeros(niter+1)  # サンプルを格納する
    res[1]  = x               # 初期配位を格納する
    # ========================================

    for iSample = 1:niter # niter回サンプル

        backup_x = x
        p        = randn()      # 共役運動量のサンプル ( サボってライブラリを使う ) 
        h_init   = calc_h(x, p) # LF前のハミルトニアンを計算

        # ============= Leap Frog ================
        # 1st step
        x += 0.5*p*dτ
        # 2nd - nτ step
        for iLF in 1:nτ
            dhdx =  calc_dhdx(x)
            p   -= dhdx*dτ
            x   += p*dτ
        end
        # last step
        dhdx =  calc_dhdx(x)
        p   -= dhdx*dτ
        x   += 0.5*p*dτ
        # ========================================
        h_fin = calc_h(x, p)  # LF後のハミルトニアンを計算

        # =========== Metropolis test ============
        metropolis = rand()
        exp(h_init-h_fin)>metropolis ? naccept+=1 : x = backup_x
        # ========================================
        res[iSample+1] = x
    end
    return res
end

# 結果
pyplot(fmt=:svg, size=(400, 250))
res = hmc1d();
histogram(res)
PyPlot.savefig("HMC_1.png", dpi=300)

# 作用にでててくるAᵢⱼを定義
A = [1.0 1.0 1.0;
     1.0 2.0 1.0;
     1.0 1.0 2.0];

# 函数を定義：作用を計算 ( 定義 ) する
function calc_action(x, A)
    action = 0.5*transpose(x)*A*x
end

# 函数を定義：ハミルトニアンHを計算する
function calc_h(x, A, p)
    ham = calc_action(x, A)
    ham = ham + 0.5*transpose(p)*p
end

# 函数を定義：ハミルトニアンの微分dH/dxを計算する
function calc_dhdx(x, A)
    dhdx = A*x
end


function hmc2d(;niter=10^3, nτ=20, dτ=0.5, ndim=3)

    # ================= 準備 =================
    seed!(12345)
    x         = zeros(ndim)           # 初期配位
    naccept   = 0                     # 受理回数カウンター
    res       = zeros(niter+1, ndim)  # サンプルを格納する
    res[1, :] = x
    # ========================================

    for iSample = 1:niter # niter回サンプル

        backup_x = x
        p        = randn(ndim)  # 共役運動量のサンプル ( サボってライブラリを使う ) 
        h_init   = calc_h(x, A, p) # LF前のハミルトニアンを計算

        # ============= Leap Frog ================
        # 1st step
        @. x = x + 0.5*p*dτ
        # 2nd - nτ step
        for iLF in 1:nτ
            dhdx =  calc_dhdx(x, A)
            @. p = p - dhdx*dτ
            @. x = x + p*dτ
        end
        # last step
        dhdx =  calc_dhdx(x, A)
        @. p = p - dhdx*dτ
        @. x = x + 0.5*p*dτ
        # ========================================
        h_fin = calc_h(x, A, p)  # LF後のハミルトニアンを計算

        # =========== Metropolis test ============
        metropolis = rand()
        exp(h_init-h_fin)>metropolis ? naccept+=1 : x = backup_x
        # ========================================
        res[iSample+1, :] = x
    end
    return res
end

# 結果
pyplot(fmt=:svg, size=(400, 400))
res = hmc2d();
scatter3d(res[:, 1], res[:, 2], res[:, 3])
PyPlot.savefig("HMC_2.png", dpi=300)