update code further

Author: sunxd3

src/AbstractMCMC.jl

@@ -82,36 +82,6 @@ struct MCMCSerial <: AbstractMCMCEnsemble end
 
 function condition end
 
-function recompute_logprob!! end
-
-"""
-    get_logprob(state)
-
-Returns the log-probability of the last sampling step, stored in `state`.
-"""
-function get_logprob(state) end
-
-"""
-    set_logprob!(state, logprob)
-
-Set the log-probability of the last sampling step, stored in `state`.
-"""
-function set_logprob!!(state, logprob) end
-
-"""
-    get_params(state)
-
-Returns the values of the parameters in the state.
-"""
-function get_params(state) end
-
-"""
-    setparams!(state, params)
-
-Set the values of the parameters in the state.
-"""
-function set_params!!(state, params) end
-
 include("samplingstats.jl")
 include("logging.jl")
 include("interface.jl")

src/gibbs.jl

@@ -77,7 +77,8 @@ julia> unflatten([1.0,2.0,3.0,4.0,5.0,6.0], (x=(2,2), y=(2,)))
 function unflatten(vec::AbstractVector, variable_sizes::NamedTuple)
     result = Dict{Symbol,Array}()
     start_idx = 1
-    for (name, size) in pairs(variable_sizes)
+    for name in keys(variable_sizes)
+        size = variable_sizes[name]
         end_idx = start_idx + prod(size) - 1
         result[name] = reshape(vec[start_idx:end_idx], size...)
         start_idx = end_idx + 1
@@ -100,7 +101,7 @@ function update_trace(trace::NamedTuple, gibbs_state::GibbsState)
         trace = merge(
             trace,
             unflatten(
-                AbstractMCMC.get_params(sub_state),
+                vec(sub_state),
                 NamedTuple{(parameter_variable,)}((
                     gibbs_state.variable_sizes[parameter_variable],
                 )),
@@ -197,9 +198,14 @@ function AbstractMCMC.step(
         )
 
         # recompute the logdensity stored in the mcmc state, because the values might have been updated in other sub-problems
-        sub_state = AbstractMCMC.recompute_logprob!!(
-            cond_logdensity, AbstractMCMC.get_params(sub_state), sub_state
-        )
+        updated_log_prob = LogDensityProblems.logdensity(cond_logdensity, sub_state)
+
+        if !hasproperty(sub_state, :logp)
+            error(
+                "$(typeof(sub_state)) does not have a `:logp` field, which is required by Gibbs sampling",
+            )
+        end
+        sub_state = BangBang.setproperty!!(sub_state, :logp, updated_log_prob)
 
         sub_state = last(
             AbstractMCMC.step(

test/gibbs_example/gibbs.jl

@@ -5,7 +5,7 @@ include("hier_normal.jl")
 @testset "hierarchical normal with gibbs" begin
     # generate data
     N = 1000  # Number of data points
-    mu_true = 0.5  # True mean
+    mu_true = 5  # True mean
     tau2_true = 2.0  # True variance
     x_data = rand(Distributions.Normal(mu_true, sqrt(tau2_true)), N)
 
@@ -15,7 +15,8 @@ include("hier_normal.jl")
     samples = sample(
         hn,
         AbstractMCMC.Gibbs((
-            mu=RandomWalkMH(1), tau2=IndependentMH(product_distribution([InverseGamma(1, 1)]))
+            mu=RandomWalkMH(0.3),
+            tau2=IndependentMH(product_distribution([InverseGamma(1, 1)])),
         )),
         200000;
         initial_params=(mu=[0.0], tau2=[1.0]),

test/gibbs_example/gmm.jl

@@ -77,7 +77,3 @@ end
 function unflatten(vec::AbstractVector, group::Tuple)
     return NamedTuple((only(group) => vec,))
 end
-
-function recompute_logprob!!(gmm::ConditionedGMM, vals, state)
-    return set_logp!!(state, LogDensityProblems.logdensity(gmm, vals))
-end

test/gibbs_example/mh.jl

@@ -9,10 +9,14 @@ struct MHTransition{T}
     params::Vector{T}
 end
 
-AbstractMCMC.get_params(state::MHState) = state.params
-AbstractMCMC.set_params!!(state::MHState, params) = MHState(params, state.logp)
-AbstractMCMC.get_logprob(state::MHState) = state.logp
-AbstractMCMC.set_logprob!!(state::MHState, logp) = MHState(state.params, logp)
+function AbstractMCMC.LogDensityProblems.logdensity(logdensity_function, state::MHState)
+    # recompute the logdensity, instead of using the one stored in the state
+    return AbstractMCMC.LogDensityProblems.logdensity(logdensity_function, state.params)
+end
+
+function Base.vec(state::MHState)
+    return state.params
+end
 
 struct RandomWalkMH <: AbstractMCMC.AbstractSampler
     σ::Float64
@@ -64,7 +68,7 @@ end
 function compute_log_acceptance_ratio(
     ::RandomWalkMH, state::MHState, ::Vector{Float64}, logp_proposal::Float64
 )
-    return min(0, logp_proposal - AbstractMCMC.get_logprob(state))
+    return min(0, logp_proposal - state.logp)
 end
 
 function compute_log_acceptance_ratio(