rework the code; still not type stable

Author: sunxd3

Project.toml

@@ -6,7 +6,6 @@ desc = "A lightweight interface for common MCMC methods."
 version = "5.3.0"
 
 [deps]
-AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 ConsoleProgressMonitor = "88cd18e8-d9cc-4ea6-8889-5259c0d15c8b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
@@ -28,18 +27,21 @@ ConsoleProgressMonitor = "0.1"
 FillArrays = "1"
 LogDensityProblems = "2"
 LoggingExtras = "0.4, 0.5, 1"
+MCMCChains = "6"
 ProgressLogging = "0.1"
 StatsBase = "0.32, 0.33, 0.34"
 TerminalLoggers = "0.1"
 Transducers = "0.4.30"
 julia = "1.6"
 
 [extras]
+AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 IJulia = "7073ff75-c697-5162-941a-fcdaad2a7d2a"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AbstractPPL","FillArrays", "Distributions", "IJulia", "Statistics", "Test"]
+test = ["AbstractPPL","FillArrays", "Distributions", "IJulia", "MCMCChains", "Statistics", "Test"]

src/AbstractMCMC.jl

@@ -87,6 +87,5 @@ include("sample.jl")
 include("stepper.jl")
 include("transducer.jl")
 include("logdensityproblems.jl")
-include("gibbs.jl")
 
 end # module AbstractMCMC

test/gibbs_example/gibbs.jl

@@ -1,5 +1,7 @@
-using AbstractMCMC, AbstractPPL
-using BangBang.ConstructorBase: ConstructorBase
+using AbstractMCMC: AbstractMCMC
+using AbstractPPL: AbstractPPL 
+using MCMCChains: Chains
+using Random
 
 """
     Gibbs(sampler_map::NamedTuple)
@@ -99,27 +101,34 @@ function update_trace(trace::NamedTuple, gibbs_state::GibbsState)
     return trace
 end
 
+function error_if_not_fully_initialized(
+    initial_params::NamedTuple{ParamNames}, sampler::Gibbs{<:NamedTuple{SamplerNames}}
+) where {ParamNames,SamplerNames}
+    if Set(ParamNames) != Set(SamplerNames)
+        throw(
+            ArgumentError(
+                "initial_params must contain all parameters in the model, expected $(SamplerNames), got $(ParamNames)",
+            ),
+        )
+    end
+end
+
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     logdensity_model::AbstractMCMC.LogDensityModel,
-    sampler::Gibbs,
+    sampler::Gibbs{Tsamplingmap},
     args...;
     initial_params::NamedTuple,
     kwargs...,
-)
-    if Set(keys(initial_params)) != Set(keys(sampler.sampler_map))
-        throw(
-            ArgumentError(
-                "initial_params must contain all parameters in the model, expected $(keys(sampler.sampler_map)), got $(keys(initial_params))",
-            ),
-        )
-    end
+) where {Tsamplingmap}
+    error_if_not_fully_initialized(initial_params, sampler)
 
-    mcmc_states, variable_sizes = map(keys(sampler.sampler_map)) do parameter_variable
+    model_parameter_names = fieldnames(Tsamplingmap)
+    results = map(model_parameter_names) do parameter_variable
         sub_sampler = sampler.sampler_map[parameter_variable]
 
         variables_to_be_conditioned_on = setdiff(
-            keys(sampler.sampler_map), (parameter_variable,)
+            model_parameter_names, (parameter_variable,)
         )
         conditioning_variables_values = NamedTuple{Tuple(variables_to_be_conditioned_on)}(
             Tuple([initial_params[g] for g in variables_to_be_conditioned_on])
@@ -137,7 +146,7 @@ function AbstractMCMC.step(
             AbstractMCMC.step(
                 rng,
                 AbstractMCMC.LogDensityModel(
-                    AbstractMCMC.condition(
+                    AbstractPPL.condition(
                         logdensity_model.logdensity, conditioning_variables_values
                     ),
                 ),
@@ -150,40 +159,46 @@ function AbstractMCMC.step(
         (sub_state, Tuple(size(initial_params[parameter_variable])))
     end
 
+    mcmc_states = first.(results)
+    variable_sizes = last.(results)
+
     gibbs_state = GibbsState(
         initial_params,
-        NamedTuple{Tuple(keys(sampler.sampler_map))}(mcmc_states),
-        NamedTuple{Tuple(keys(sampler.sampler_map))}(variable_sizes),
+        NamedTuple{Tuple(model_parameter_names)}(mcmc_states),
+        NamedTuple{Tuple(model_parameter_names)}(variable_sizes),
     )
+
     trace = update_trace(NamedTuple(), gibbs_state)
     return GibbsTransition(trace), gibbs_state
 end
 
+# subsequent steps
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     logdensity_model::AbstractMCMC.LogDensityModel,
-    sampler::Gibbs,
+    sampler::Gibbs{Tsamplingmap},
     gibbs_state::GibbsState,
     args...;
     kwargs...,
-)
+) where {Tsamplingmap}
     (; trace, mcmc_states, variable_sizes) = gibbs_state
 
-    mcmc_states = map(keys(sampler.sampler_map)) do parameter_variable
+    model_parameter_names = fieldnames(Tsamplingmap)
+    mcmc_states = map(model_parameter_names) do parameter_variable
         sub_sampler = sampler.sampler_map[parameter_variable]
         sub_state = mcmc_states[parameter_variable]
         variables_to_be_conditioned_on = setdiff(
-            sampler.parameter_names, (parameter_variable,)
+            model_parameter_names, (parameter_variable,)
         )
         conditioning_variables_values = NamedTuple{Tuple(variables_to_be_conditioned_on)}(
             Tuple([trace[g] for g in variables_to_be_conditioned_on])
         )
-        cond_logdensity = AbstractMCMC.condition(
+        cond_logdensity = AbstractPPL.condition(
             logdensity_model.logdensity, conditioning_variables_values
         )
 
-        logp = LogDensityProblems.logdensity_and_state(cond_logdensity, sub_state)
-        sub_state = constructorof(typeof(sub_state))(; logp=logp)
+        logp = LogDensityProblems.logdensity(cond_logdensity, sub_state)
+        sub_state = (sub_state)(logp)
         sub_state = last(
             AbstractMCMC.step(
                 rng,
@@ -197,7 +212,7 @@ function AbstractMCMC.step(
         trace = update_trace(trace, gibbs_state)
         sub_state
     end
-    mcmc_states = NamedTuple{Tuple(keys(sampler.sampler_map))}(mcmc_states)
+    mcmc_states = NamedTuple{Tuple(model_parameter_names)}(mcmc_states)
 
     return GibbsTransition(trace), GibbsState(trace, mcmc_states, variable_sizes)
 end

test/gibbs_example/hier_normal.jl

@@ -1,3 +1,5 @@
+using AbstractPPL: AbstractPPL
+
 abstract type AbstractHierNormal end
 
 struct HierNormal{Tdata<:NamedTuple} <: AbstractHierNormal
@@ -7,6 +9,8 @@ end
 struct ConditionedHierNormal{Tdata<:NamedTuple,Tconditioned_vars<:NamedTuple} <:
        AbstractHierNormal
     data::Tdata
+
+    " The variable to be conditioned on and its value"
     conditioned_values::Tconditioned_vars
 end
 
@@ -36,14 +40,15 @@ function log_joint(; mu, tau2, x)
     return logp
 end
 
-function AbstractMCMC.condition(hn::HierNormal, conditioned_values::NamedTuple)
+function AbstractPPL.condition(hn::HierNormal, conditioned_values::NamedTuple)
     return ConditionedHierNormal(hn.data, conditioned_values)
 end
 
 function LogDensityProblems.logdensity(
-    hier_normal_model::ConditionedHierNormal{names}, params::AbstractVector
-) where {names}
-    variable_to_condition = only(names)
+    hier_normal_model::ConditionedHierNormal{Tdata,Tconditioned_vars},
+    params::AbstractVector,
+) where {Tdata,Tconditioned_vars}
+    variable_to_condition = only(fieldnames(Tconditioned_vars))
     data = hier_normal_model.data
     conditioned_values = hier_normal_model.conditioned_values
 

test/gibbs_example/mh.jl

@@ -1,12 +1,18 @@
+using AbstractMCMC: AbstractMCMC, LogDensityProblems
 using Distributions
-
+using Random
 abstract type AbstractMHSampler <: AbstractMCMC.AbstractSampler end
 
 struct MHState{T}
     params::Vector{T}
     logp::Float64
 end
 
+# Interface 3: (state::MHState)(logp::Float64)
+# This function allows the state to be updated with a new log probability.
+# ! this makes state into a Julia functor
+(state::MHState)(logp::Float64) = MHState(state.params, logp)
+
 struct MHTransition{T}
     params::Vector{T}
 end
@@ -15,7 +21,7 @@ end
 # 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(
+function LogDensityProblems.logdensity(
     logdensity_function, state::MHState; recompute_logp=true
 )
     return if recompute_logp
@@ -28,9 +34,7 @@ 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
+Base.vec(state::MHState) = state.params
 
 """
     RandomWalkMH{T} <: AbstractMCMC.AbstractSampler
@@ -62,14 +66,17 @@ function AbstractMCMC.step(
 )
     logdensity_function = logdensity_model.logdensity
     transition = MHTransition(initial_params)
-    state = MHState(initial_params, only(logdensity_function(initial_params)))
+    state = MHState(
+        initial_params,
+        only(LogDensityProblems.logdensity(logdensity_function, initial_params)),
+    )
 
     return transition, state
 end
 
-@inline proposal_dist(sampler::RandomWalkMH, current_params::Vector{Float64}) =
+@inline get_proposal_dist(sampler::RandomWalkMH, current_params::Vector{Float64}) =
     MvNormal(current_params, sampler.σ)
-@inline proposal_dist(sampler::IndependentMH, current_params::Vector{T}) where {T} =
+@inline get_proposal_dist(sampler::IndependentMH, current_params::Vector{T}) where {T} =
     sampler.proposal_dist
 
 # the subsequent steps of the sampler
@@ -83,7 +90,7 @@ function AbstractMCMC.step(
 )
     logdensity_function = logdensity_model.logdensity
     current_params = state.params
-    proposal_dist = proposal_dist(sampler, current_params)
+    proposal_dist = get_proposal_dist(sampler, current_params)
     proposed_params = rand(rng, proposal_dist)
     logp_proposal = only(
         LogDensityProblems.logdensity(logdensity_function, proposed_params)

test/runtests.jl

@@ -24,5 +24,5 @@ include("utils.jl")
     include("stepper.jl")
     include("transducer.jl")
     include("logdensityproblems.jl")
-    include("gibbs_example/gibbs.jl")
+    include("gibbs_example/gibbs_test.jl")
 end