updates

Author: sunxd3

test/gibbs_example/hier_normal.jl

@@ -1,14 +1,16 @@
 abstract type AbstractHierNormal end
 
-struct HierNormal <: AbstractHierNormal
-    data::NamedTuple
+struct HierNormal{Tdata<:NamedTuple} <: AbstractHierNormal
+    data::Tdata
 end
 
-struct ConditionedHierNormal{conditioned_vars} <: AbstractHierNormal
-    data::NamedTuple
-    conditioned_values::NamedTuple{conditioned_vars}
+struct ConditionedHierNormal{Tdata<:NamedTuple,Tconditioned_vars<:NamedTuple} <:
+       AbstractHierNormal
+    data::Tdata
+    conditioned_values::Tconditioned_vars
 end
 
+# `mu` and `tau2` are length-1 vectors to make 
 function log_joint(; mu, tau2, x)
     # mu is the mean
     # tau2 is the variance
@@ -39,14 +41,18 @@ function AbstractMCMC.condition(hn::HierNormal, conditioned_values::NamedTuple)
 end
 
 function LogDensityProblems.logdensity(
-    hn::ConditionedHierNormal{names}, params::AbstractVector
+    hier_normal_model::ConditionedHierNormal{names}, params::AbstractVector
 ) where {names}
-    if Set(names) == Set([:mu]) # conditioned on mu, so params are tau2
-        return log_joint(; mu=hn.conditioned_values.mu, tau2=params, x=hn.data.x)
-    elseif Set(names) == Set([:tau2]) # conditioned on tau2, so params are mu
-        return log_joint(; mu=params, tau2=hn.conditioned_values.tau2, x=hn.data.x)
+    variable_to_condition = only(names)
+    data = hier_normal_model.data
+    conditioned_values = hier_normal_model.conditioned_values
+
+    if variable_to_condition == :mu
+        return log_joint(; mu=conditioned_values.mu, tau2=params, x=data.x)
+    elseif variable_to_condition == :tau2
+        return log_joint(; mu=params, tau2=conditioned_values.tau2, x=data.x)
     else
-        error("Unsupported conditioning configuration.")
+        error("Unsupported conditioning variable: $variable_to_condition")
     end
 end
 

test/gibbs_example/mh.jl

@@ -1,5 +1,7 @@
 using Distributions
 
+abstract type AbstractMHSampler <: AbstractMCMC.AbstractSampler end
+
 struct MHState{T}
     params::Vector{T}
     logp::Float64
@@ -9,65 +11,89 @@ struct MHTransition{T}
     params::Vector{T}
 end
 
+# Interface 1: LogDensityProblems.logdensity
+# This function takes the logdensity function and the state (state is defined by the sampler package)
+# and returns the logdensity. It allows for optional recomputation of the log probability.
+# If recomputation is not needed, it returns the stored log probability from the state.
 function AbstractMCMC.logdensity_and_state(
-    logdensity_function, state::MHState; recompute_logp::Bool=true
+    logdensity_function, state::MHState; recompute_logp=true
 )
-    if recompute_logp
-        logp = AbstractMCMC.LogDensityProblems.logdensity(logdensity_function, state.params)
-        return logp, MHState(state.params, logp)
+    return if recompute_logp
+        AbstractMCMC.LogDensityProblems.logdensity(logdensity_function, state.params)
     else
-        return state.logp, state
+        state.logp
     end
 end
 
+# Interface 2: Base.vec
+# This function takes a state and returns a vector of the parameter values stored in the state.
+# It is part of the interface for interacting with the state object.
 function Base.vec(state::MHState)
     return state.params
 end
 
-struct RandomWalkMH{T} <: AbstractMCMC.AbstractSampler
+"""
+    RandomWalkMH{T} <: AbstractMCMC.AbstractSampler
+
+A random walk Metropolis-Hastings sampler with a normal proposal distribution. The field σ
+is the standard deviation of the proposal distribution.
+"""
+struct RandomWalkMH{T} <: AbstractMHSampler
     σ::T
 end
 
-struct IndependentMH{T} <: AbstractMCMC.AbstractSampler
+"""
+    IndependentMH{T} <: AbstractMCMC.AbstractSampler
+
+A Metropolis-Hastings sampler with an independent proposal distribution.
+"""
+struct IndependentMH{T} <: AbstractMHSampler
     proposal_dist::T
 end
 
+# the first step of the sampler
 function AbstractMCMC.step(
     rng::AbstractRNG,
     logdensity_model::AbstractMCMC.LogDensityModel,
-    sampler::Union{RandomWalkMH,IndependentMH},
+    sampler::AbstractMHSampler,
     args...;
     initial_params,
     kwargs...,
 )
-    return MHTransition(initial_params),
-    MHState(
-        initial_params,
-        only(LogDensityProblems.logdensity(logdensity_model.logdensity, initial_params)),
-    )
+    logdensity_function = logdensity_model.logdensity
+    transition = MHTransition(initial_params)
+    state = MHState(initial_params, only(logdensity_function(initial_params)))
+
+    return transition, state
 end
 
+@inline proposal_dist(sampler::RandomWalkMH, current_params::Vector{Float64}) =
+    MvNormal(current_params, sampler.σ)
+@inline proposal_dist(sampler::IndependentMH, current_params::Vector{T}) where {T} =
+    sampler.proposal_dist
+
+# the subsequent steps of the sampler
 function AbstractMCMC.step(
     rng::AbstractRNG,
     logdensity_model::AbstractMCMC.LogDensityModel,
-    sampler::Union{RandomWalkMH,IndependentMH},
+    sampler::AbstractMHSampler,
     state::MHState,
     args...;
     kwargs...,
 )
-    params = state.params
-    proposal_dist =
-        sampler isa RandomWalkMH ? MvNormal(state.params, sampler.σ) : sampler.proposal_dist
-    proposal = rand(rng, proposal_dist)
+    logdensity_function = logdensity_model.logdensity
+    current_params = state.params
+    proposal_dist = proposal_dist(sampler, current_params)
+    proposed_params = rand(rng, proposal_dist)
     logp_proposal = only(
-        LogDensityProblems.logdensity(logdensity_model.logdensity, proposal)
+        LogDensityProblems.logdensity(logdensity_function, proposed_params)
     )
 
     if log(rand(rng)) <
-        compute_log_acceptance_ratio(sampler, state, proposal, logp_proposal)
-        return MHTransition(proposal), MHState(proposal, logp_proposal)
+        compute_log_acceptance_ratio(sampler, state, proposed_params, logp_proposal)
+        return MHTransition(proposed_params), MHState(proposed_params, logp_proposal)
     else
-        return MHTransition(params), MHState(params, state.logp)
+        return MHTransition(current_params), MHState(current_params, state.logp)
     end
 end