format

Author: odunbar

src/MarkovChainMonteCarlo.jl

@@ -29,7 +29,7 @@ export EmulatorPosteriorModel,
     infMALASampling,
     infHMCSampling,
     infmMALASampling,
-    infmHMCSampling,    
+    infmHMCSampling,
     MCMCWrapper,
     accept_ratio,
     optimize_stepsize,
@@ -166,7 +166,8 @@ AdvancedMH.logratio_proposal_density(
     candidate,
 ) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate)
 
-MetropolisHastingsSampler(::MALASampling, prior::ParameterDistribution) = MetropolisAdjustedLangevin(_get_proposal(prior))
+MetropolisHastingsSampler(::MALASampling, prior::ParameterDistribution) =
+    MetropolisAdjustedLangevin(_get_proposal(prior))
 """
 $(DocStringExtensions.TYPEDEF)
 
@@ -185,7 +186,8 @@ AdvancedMH.logratio_proposal_density(
     candidate,
 ) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate)
 
-MetropolisHastingsSampler(::BarkerSampling, prior::ParameterDistribution) = BarkerMetropolisHastings(_get_proposal(prior))
+MetropolisHastingsSampler(::BarkerSampling, prior::ParameterDistribution) =
+    BarkerMetropolisHastings(_get_proposal(prior))
 """
 $(DocStringExtensions.TYPEDEF)
 
@@ -223,7 +225,8 @@ AdvancedMH.logratio_proposal_density(
     candidate,
 ) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate)
 
-MetropolisHastingsSampler(::infMALASampling, prior::ParameterDistribution) = infMALAMetropolisHastings(_get_proposal(prior))
+MetropolisHastingsSampler(::infMALASampling, prior::ParameterDistribution) =
+    infMALAMetropolisHastings(_get_proposal(prior))
 """
 $(DocStringExtensions.TYPEDEF)
 
@@ -242,7 +245,8 @@ AdvancedMH.logratio_proposal_density(
     candidate,
 ) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate)
 
-MetropolisHastingsSampler(::infHMCSampling, prior::ParameterDistribution) = infHMCMetropolisHastings(_get_proposal(prior))
+MetropolisHastingsSampler(::infHMCSampling, prior::ParameterDistribution) =
+    infHMCMetropolisHastings(_get_proposal(prior))
 """
 $(DocStringExtensions.TYPEDEF)
 
@@ -261,7 +265,8 @@ AdvancedMH.logratio_proposal_density(
     candidate,
 ) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate)
 
-MetropolisHastingsSampler(::infmMALASampling, prior::ParameterDistribution) = infmMALAMetropolisHastings(_get_proposal(prior))
+MetropolisHastingsSampler(::infmMALASampling, prior::ParameterDistribution) =
+    infmMALAMetropolisHastings(_get_proposal(prior))
 """
 $(DocStringExtensions.TYPEDEF)
 
@@ -280,7 +285,8 @@ AdvancedMH.logratio_proposal_density(
     candidate,
 ) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate)
 
-MetropolisHastingsSampler(::infmHMCSampling, prior::ParameterDistribution) = infmHMCMetropolisHastings(_get_proposal(prior))
+MetropolisHastingsSampler(::infmHMCSampling, prior::ParameterDistribution) =
+    infmHMCMetropolisHastings(_get_proposal(prior))
 
 # ------------------------------------------------------------------------------------------
 # Use emulated model in sampler
@@ -396,10 +402,10 @@ function AdvancedMH.propose(
     current_state::MCMCState;
     stepsize::FT = 1.0,
 ) where {FT <: AbstractFloat}
-        # Compute the gradient of the log-density at the current state
-        log_gradient = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
-        proposed_state = current_state.params .+ (stepsize^2 / 2) .* log_gradient .+ stepsize * rand(rng, sampler.proposal)
-        return proposed_state
+    # Compute the gradient of the log-density at the current state
+    log_gradient = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
+    proposed_state = current_state.params .+ (stepsize^2 / 2) .* log_gradient .+ stepsize * rand(rng, sampler.proposal)
+    return proposed_state
 end
 
 # method extending AdvancedMH.propose() for the Barker proposal
@@ -410,13 +416,13 @@ function AdvancedMH.propose(
     current_state::MCMCState;
     stepsize::FT = 1.0,
 ) where {FT <: AbstractFloat}
-# Livingstone and Zanella (2022)
+    # Livingstone and Zanella (2022)
     # Compute the gradient of the log-density at the current state
     log_gradient = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
     n = length(current_state.params)
     u = rand(rng, n)
     xi = rand(rng, sampler.proposal)
-    b = u .< 1 ./ (1 .+ exp.(- log_gradient .* xi))
+    b = u .< 1 ./ (1 .+ exp.(-log_gradient .* xi))
     return current_state.params .+ b .* xi
 end
 
@@ -441,9 +447,16 @@ function AdvancedMH.propose(
         log_gradient = log_grad_proposed_state
         proposed_state .+= sqrt(stepsize) .* proposed_aux - (stepsize / 2) .* log_grad_proposed_state
         log_grad_proposed_state = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), proposed_state)
-        proposed_aux .+= - (sqrt(stepsize) / 2) .* log_gradient .- (sqrt(stepsize) / 2) .* log_grad_proposed_state
+        proposed_aux .+= -(sqrt(stepsize) / 2) .* log_gradient .- (sqrt(stepsize) / 2) .* log_grad_proposed_state
     end
-    println("L: ", L, " stepsize: ", round(stepsize, digits = 6), " proposed_state: ", round.(proposed_state, digits = 5))
+    println(
+        "L: ",
+        L,
+        " stepsize: ",
+        round(stepsize, digits = 6),
+        " proposed_state: ",
+        round.(proposed_state, digits = 5),
+    )
     return proposed_state
 end
 
@@ -455,11 +468,13 @@ function AdvancedMH.propose(
     current_state::MCMCState;
     stepsize::FT = 1.0,
 ) where {FT <: AbstractFloat}
-        # Compute the gradient of the log-density at the current state
-        ρ = (1 - stepsize / 4) / (1 + stepsize / 4)
-        log_gradient = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
-        proposed_state = ρ * current_state.params .- sqrt(1 - ρ^2) * (sqrt(stepsize) / 2) .* log_gradient .+ sqrt(1 - ρ^2) * rand(rng, sampler.proposal)
-        return proposed_state
+    # Compute the gradient of the log-density at the current state
+    ρ = (1 - stepsize / 4) / (1 + stepsize / 4)
+    log_gradient = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
+    proposed_state =
+        ρ * current_state.params .- sqrt(1 - ρ^2) * (sqrt(stepsize) / 2) .* log_gradient .+
+        sqrt(1 - ρ^2) * rand(rng, sampler.proposal)
+    return proposed_state
 end
 
 # method extending AdvancedMH.propose() for the  ∞-HMC proposal
@@ -483,9 +498,16 @@ function AdvancedMH.propose(
         log_gradient = log_grad_proposed_state
         proposed_state .+= sqrt(stepsize) .* proposed_aux - (stepsize / 2) .* log_grad_proposed_state
         log_grad_proposed_state = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), proposed_state)
-        proposed_aux .+= - (sqrt(stepsize) / 2) .* log_gradient .- (sqrt(stepsize) / 2) .* log_grad_proposed_state
+        proposed_aux .+= -(sqrt(stepsize) / 2) .* log_gradient .- (sqrt(stepsize) / 2) .* log_grad_proposed_state
     end
-    println("L: ", L, " stepsize: ", round(stepsize, digits = 8), " proposed_state: ", round.(proposed_state, digits = 5))
+    println(
+        "L: ",
+        L,
+        " stepsize: ",
+        round(stepsize, digits = 8),
+        " proposed_state: ",
+        round.(proposed_state, digits = 5),
+    )
     return proposed_state
 end
 
@@ -497,16 +519,16 @@ function AdvancedMH.propose(
     current_state::MCMCState;
     stepsize::FT = 1.0,
 ) where {FT <: AbstractFloat}
-        # Compute the gradient of the log-density at the current state
-        ρ = (1 - stepsize / 4) / (1 + stepsize / 4)
-        log_gradient = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
-        hessian = Symmetric(ForwardDiff.hessian(x -> AdvancedMH.logdensity(model, x), current_state.params))
-        K = Symmetric(inv(- hessian))
-        C_inv = I(size(K, 1))
-        xi = cholesky(K, check=false).L * randn(size(K, 1))
-        # xi = rand(rng, MvNormal(zeros(size(K, 1)), K))# or cholesky(K_u).L * randn(size(K_u, 1))
-        nu =  xi .- (stepsize / 2) .* K * ((C_inv + hessian) * current_state.params .+ log_gradient)
-        return ρ * current_state.params .+ sqrt(1 - ρ^2) * nu
+    # Compute the gradient of the log-density at the current state
+    ρ = (1 - stepsize / 4) / (1 + stepsize / 4)
+    log_gradient = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
+    hessian = Symmetric(ForwardDiff.hessian(x -> AdvancedMH.logdensity(model, x), current_state.params))
+    K = Symmetric(inv(-hessian))
+    C_inv = I(size(K, 1))
+    xi = cholesky(K, check = false).L * randn(size(K, 1))
+    # xi = rand(rng, MvNormal(zeros(size(K, 1)), K))# or cholesky(K_u).L * randn(size(K_u, 1))
+    nu = xi .- (stepsize / 2) .* K * ((C_inv + hessian) * current_state.params .+ log_gradient)
+    return ρ * current_state.params .+ sqrt(1 - ρ^2) * nu
 end
 
 # method extending AdvancedMH.propose() for the  ∞-mHMC proposal
@@ -526,13 +548,13 @@ function AdvancedMH.propose(
     proposed_state = proposed_state_init
     log_grad_proposed_state = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
 
-    for t in 1:L-1
+    for t in 1:(L - 1)
         println("Iteration t = ", t)
         println("Before update, proposed_state: ", proposed_state)
         log_gradient = log_grad_proposed_state
         proposed_state .+= stepsize .* proposed_aux - (stepsize^2 / 2) .* log_grad_proposed_state
         log_grad_proposed_state = ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), proposed_state)
-        proposed_aux .+= - (stepsize / 2) .* log_gradient .- (stepsize / 2) .* log_grad_proposed_state
+        proposed_aux .+= -(stepsize / 2) .* log_gradient .- (stepsize / 2) .* log_grad_proposed_state
         println("After update, proposed_state: ", proposed_state)
         println("proposed_aux: ", proposed_aux)
     end
@@ -939,18 +961,18 @@ $(DocStringExtensions.TYPEDSIGNATURES)
 Computes the expected squared jump distance of the chain.
 """
 function esjd(chain::MCMCChains.Chains)
-    samples = chain.value[:,:,1] # N_samples x N_params x n_chains
+    samples = chain.value[:, :, 1] # N_samples x N_params x n_chains
     n_samples, n_params = size(samples)
     esjd = zeros(Float64, n_params)
     for i in 2:n_samples
-        esjd .+= (samples[i, :] .- samples[i - 1, :]).^ 2 ./ n_samples
+        esjd .+= (samples[i, :] .- samples[i - 1, :]) .^ 2 ./ n_samples
     end
     return esjd
-    
+
 end
 
 
 
 
-    
+
 end # module MarkovChainMonteCarlo

test/MarkovChainMonteCarlo/runtests.jl

@@ -122,7 +122,7 @@ function mcmc_test_template(
 
     # Now begin the actual MCMC, sample is multiply exported so we qualify
     chain = MCMC.sample(rng, mcmc, 100_000; stepsize = new_step, discard_initial = 1000)
-    
+
     posterior_distribution = get_posterior(mcmc, chain)
     #post_mean = mean(posterior, dims=1)[1]
     posterior_mean = mean(posterior_distribution)
@@ -158,8 +158,8 @@ end
 
     @testset "Sine GP & RW Metropolis" begin
         em_1 = test_gp_1(y, σ2_y, iopairs)
-        new_step, posterior_mean_1,chain1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params...)
-        esjd1=esjd(chain1)
+        new_step, posterior_mean_1, chain1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params...)
+        esjd1 = esjd(chain1)
         @info "ESJD = $esjd1"
         @test isapprox(new_step, 0.5; atol = 0.5)
         # difference between mean_1 and ground truth comes from MCMC convergence and GP sampling
@@ -172,12 +172,12 @@ end
         _, posterior_mean_2, chain2 = mcmc_test_template(prior, σ2_y, em_2; mcmc_params...)
         # difference between mean_1 and mean_2 only from MCMC convergence
         @test isapprox(posterior_mean_2, posterior_mean_1; atol = 0.1)
-       # test diagnostic functions on the chain
-        esjd2=esjd(chain2)
+        # test diagnostic functions on the chain
+        esjd2 = esjd(chain2)
         @info "ESJD = $esjd2"
         # approx [0.04190683285347798, 0.1685296224916364, 0.4129400000002722]
-        @test all(isapprox.(esjd1,esjd2, rtol=0.1))
-                  
+        @test all(isapprox.(esjd1, esjd2, rtol = 0.1))
+
     end
 
     @testset "Sine GP & pCN" begin
@@ -190,8 +190,8 @@ end
         )
 
         em_1 = test_gp_1(y, σ2_y, iopairs)
-        new_step, posterior_mean_1,chain1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params...)
-        esjd1=esjd(chain1)
+        new_step, posterior_mean_1, chain1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params...)
+        esjd1 = esjd(chain1)
         @info "ESJD = $esjd1"
         @test isapprox(new_step, 0.75; atol = 0.6)
         # difference between mean_1 and ground truth comes from MCMC convergence and GP sampling
@@ -201,15 +201,15 @@ end
         norm_factor = 10.0
         norm_factor = fill(norm_factor, size(y[:, 1])) # must be size of output dim
         em_2 = test_gp_2(y, σ2_y, iopairs; norm_factor = norm_factor)
-        _, posterior_mean_2,chain2 = mcmc_test_template(prior, σ2_y, em_2; mcmc_params...)
+        _, posterior_mean_2, chain2 = mcmc_test_template(prior, σ2_y, em_2; mcmc_params...)
         # difference between mean_1 and mean_2 only from MCMC convergence
         @test isapprox(posterior_mean_2, posterior_mean_1; atol = 0.1)
 
-        esjd2=esjd(chain2)
+        esjd2 = esjd(chain2)
         @info "ESJD = $esjd2"
         # approx [0.03470825350663073, 0.161606734823579, 0.38970000000024896]
 
-        @test all(isapprox.(esjd1,esjd2, rtol=0.1))
-      
+        @test all(isapprox.(esjd1, esjd2, rtol = 0.1))
+
     end
 end