Skip to content

Commit

Permalink
Use abstract types in MarkovChainMonteCarlo
Browse files Browse the repository at this point in the history
  • Loading branch information
@tsj5
tsj5 committed Jan 28, 2022
1 parent 0cfeebc commit 3d67c96
Showing 1 changed file with 79 additions and 68 deletions.
147 changes: 79 additions & 68 deletions src/MarkovChainMonteCarlo.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,27 +28,27 @@ Structure to organize MCMC parameters and data
# Fields
$(DocStringExtensions.FIELDS)
"""
struct MCMC{FT<:AbstractFloat, IT<:Int}
mutable struct MCMC{FT<:AbstractFloat, IT<:Int}
"a single sample from the observations. Can e.g. be picked from an Obs struct using get_obs_sample"
obs_sample::Vector{FT}
obs_sample::AbstractVector{FT}
"covariance of the observational noise"
obs_noise_cov::Array{FT, 2}
obs_noise_cov::AbstractMatrix{FT}
"array of length N_parameters with the parameters' prior distributions"
prior::ParameterDistribution
"MCMC step size"
step::Array{FT}
step::FT
"Number of MCMC steps that are considered burnin"
burnin::IT
"the current parameters"
param::Vector{FT}
param::AbstractVector{FT}
"Array of accepted MCMC parameter samples. The histogram of these samples gives an approximation of the posterior distribution of the parameters. param_dim x n_samples"
posterior::Array{FT, 2}
posterior::AbstractMatrix{FT}
"the (current) value of the logarithm of the posterior (= log_likelihood + log_prior of the current parameters)"
log_posterior::Array{Union{FT, Nothing}}
log_posterior::Union{FT, Nothing}
"iteration/step of the MCMC"
iter::Array{IT}
iter::IT
"number of accepted proposals"
accept::Array{IT}
accept::IT
"MCMC algorithm to use - currently implemented: 'rmw' (random walk Metropolis)"
algtype::String
end
Expand All @@ -67,20 +67,19 @@ where max_iter is the number of MCMC steps to perform (e.g., 100_000)
"""
function MCMC(
obs_sample::Vector{FT},
obs_noise_cov::Array{FT, 2},
obs_sample::AbstractVector{FT},
obs_noise_cov::AbstractMatrix{FT},
prior::ParameterDistribution,
step::FT,
param_init::Vector{FT},
param_init::AbstractVector{FT},
max_iter::IT,
algtype::String,
burnin::IT;
svdflag=true,
standardize=false,
norm_factor::Union{Array{FT, 1}, Nothing}=nothing,
truncate_svd=1.0) where {FT<:AbstractFloat, IT<:Int}


svdflag = true,
standardize = false,
norm_factor::Union{AbstractVector{FT}, Nothing} = nothing,
truncate_svd = 1.0
) where {FT<:AbstractFloat, IT<:Int}
param_init_copy = deepcopy(param_init)

# Standardize MCMC input?
Expand All @@ -97,42 +96,44 @@ function MCMC(
# We need to transform obs_sample into the correct space
if svdflag
println("Applying SVD to decorrelating outputs, if not required set svdflag=false")
obs_sample, unused = Emulators.svd_transform(obs_sample, obs_noise_cov; truncate_svd=truncate_svd)
obs_sample, _ = Emulators.svd_transform(obs_sample, obs_noise_cov; truncate_svd=truncate_svd)
else
println("Assuming independent outputs.")
end
println(obs_sample)

# first row is param_init
posterior = zeros(length(param_init_copy),max_iter + 1)
posterior = zeros(length(param_init_copy), max_iter + 1)
posterior[:, 1] = param_init_copy
param = param_init_copy
log_posterior = [nothing]
iter = [1]
accept = [0]
log_posterior = nothing
iter = 1
accept = 0
if algtype != "rwm"
error("only random walk metropolis 'rwm' is implemented so far")
end
MCMC{FT,IT}(obs_sample,
obs_noise_cov,
prior,
[step],
burnin,
param,
posterior,
log_posterior,
iter,
accept,
algtype)
MCMC{FT,IT}(
obs_sample,
obs_noise_cov,
prior,
step,
burnin,
param,
posterior,
log_posterior,
iter,
accept,
algtype
)
end


function reset_with_step!(mcmc::MCMC{FT}, step::FT) where {FT}
function reset_with_step!(mcmc::MCMC{FT,IT}, step::FT) where {FT<:AbstractFloat, IT<:Int}
# reset to beginning with new stepsize
mcmc.step[1] = step
mcmc.log_posterior[1] = nothing
mcmc.iter[1] = 1
mcmc.accept[1] = 0
mcmc.step = step
mcmc.log_posterior = nothing
mcmc.iter = 1
mcmc.accept = 0
mcmc.posterior[:,2:end] = zeros(size(mcmc.posterior[:,2:end]))
mcmc.param[:] = mcmc.posterior[:, 1]
end
Expand All @@ -150,45 +151,54 @@ function get_posterior(mcmc::MCMC)

end

function mcmc_sample!(mcmc::MCMC{FT}, g::Vector{FT},
gcov::Union{Matrix{FT},Diagonal{FT}}) where {FT}
function mcmc_sample!(
mcmc::MCMC{FT,IT},
g::AbstractVector{FT},
gcov::AbstractMatrix{FT}
) where {FT<:AbstractFloat, IT<:Int}
if mcmc.algtype == "rwm"
log_posterior = log_likelihood(mcmc, g, gcov) + log_prior(mcmc)
end

if mcmc.log_posterior[1] isa Nothing # do an accept step.
mcmc.log_posterior[1] = log_posterior - log(FT(0.5)) # this makes p_accept = 0.5
if mcmc.log_posterior === nothing # do an accept step.
mcmc.log_posterior = log_posterior - log(FT(0.5)) # this makes p_accept = 0.5
end
# Get new parameters by comparing likelihood_current * prior_current to
# likelihood_proposal * prior_proposal - either we accept the proposed
# parameter or we stay where we are.
p_accept = exp(log_posterior - mcmc.log_posterior[1])
p_accept = exp(log_posterior - mcmc.log_posterior)

if p_accept > rand(Distributions.Uniform(0, 1))
mcmc.posterior[:,1 + mcmc.iter[1]] = mcmc.param
mcmc.log_posterior[1] = log_posterior
mcmc.accept[1] = mcmc.accept[1] + 1
mcmc.posterior[:,1 + mcmc.iter] = mcmc.param
mcmc.log_posterior = log_posterior
mcmc.accept = mcmc.accept + 1
else
mcmc.posterior[:,1 + mcmc.iter[1]] = mcmc.posterior[:,mcmc.iter[1]]
mcmc.posterior[:, 1 + mcmc.iter[1]] = mcmc.posterior[:,mcmc.iter]
end
mcmc.param[:] = proposal(mcmc)[:]
mcmc.iter[1] = mcmc.iter[1] + 1
mcmc.param = proposal(mcmc)[:]
mcmc.iter = mcmc.iter + 1

end

function mcmc_sample!(mcmc::MCMC{FT}, g::Vector{FT}, gvar::Vector{FT}) where {FT}
function mcmc_sample!(
mcmc::MCMC{FT,IT},
g::AbstractVector{FT},
gvar::AbstractVector{FT}
) where {FT<:AbstractFloat, IT<:Int}
return mcmc_sample!(mcmc, g, Diagonal(gvar))
end

function accept_ratio(mcmc::MCMC{FT}) where {FT}
return convert(FT, mcmc.accept[1]) / mcmc.iter[1]
function accept_ratio(mcmc::MCMC{FT, IT}) where {FT<:AbstractFloat, IT<:Int}
return convert(FT, mcmc.accept) / mcmc.iter
end


function log_likelihood(mcmc::MCMC{FT},
g::Vector{FT},
gcov::Union{Matrix{FT},Diagonal{FT}}) where {FT}
log_rho = FT[0]
function log_likelihood(
mcmc::MCMC{FT,IT},
g::AbstractVector{FT},
gcov::AbstractMatrix{FT}
) where {FT<:AbstractFloat, IT<:Int}
log_rho = 0.0
#if gcov == nothing
# diff = g - mcmc.obs_sample
# log_rho[1] = -FT(0.5) * diff' * (mcmc.obs_noise_cov \ diff)
Expand All @@ -202,13 +212,13 @@ function log_likelihood(mcmc::MCMC{FT},
log_gpfidelity = -FT(0.5) * sum(log.(eigs))
# Combine got log_rho
diff = g - mcmc.obs_sample
log_rho[1] = -FT(0.5) * diff' * (gcov \ diff) + log_gpfidelity
log_rho = -FT(0.5) * diff' * (gcov \ diff) + log_gpfidelity
#end
return log_rho[1]
return log_rho
end


function log_prior(mcmc::MCMC{FT}) where {FT}
function log_prior(mcmc::MCMC)
return get_logpdf(mcmc.prior,mcmc.param)
end

Expand All @@ -221,13 +231,13 @@ function proposal(mcmc::MCMC)
prop_dist = MvNormal(zeros(length(mcmc.param)),
(mcmc.step[1]^2) * proposal_covariance)
end
sample = mcmc.posterior[:,1 + mcmc.iter[1]] .+ rand(prop_dist)
sample = mcmc.posterior[:,1 + mcmc.iter] .+ rand(prop_dist)
return sample
end


function find_mcmc_step!(mcmc_test::MCMC{FT}, em::Emulator{FT}; max_iter=2000) where {FT}
step = mcmc_test.step[1]
function find_mcmc_step!(mcmc_test::MCMC{FT,IT}, em::Emulator{FT}; max_iter::IT=2000) where {FT<:AbstractFloat, IT<:Int}
step = mcmc_test.step
mcmc_accept = false
doubled = false
halved = false
Expand Down Expand Up @@ -284,14 +294,15 @@ function find_mcmc_step!(mcmc_test::MCMC{FT}, em::Emulator{FT}; max_iter=2000) w

end

return mcmc_test.step[1]
return mcmc_test.step
end


function sample_posterior!(mcmc::MCMC{FT,IT},
em::Emulator{FT},
max_iter::IT) where {FT,IT<:Int}

function sample_posterior!(
mcmc::MCMC{FT, IT},
em::Emulator{FT},
max_iter::IT
) where {FT<:AbstractFloat, IT<:Int}
for mcmcit in 1:max_iter
param = reshape(mcmc.param, :, 1)
# test predictions (param is 1 x N_parameters)
Expand Down

0 comments on commit 3d67c96

Please sign in to comment.