Remove MCMCCHains by default

Author: cpfiffer

src/AdvancedMH.jl

@@ -1,13 +1,12 @@
 module AdvancedMH
 
 # Import the relevant libraries.
-using Reexport
 using AbstractMCMC
-using Distributions
 using Random
+using Requires
+using Distributions
 
 # Import specific functions and types to use or overload.
-import MCMCChains: Chains
 import AbstractMCMC: step!, AbstractSampler, AbstractTransition, transition_type, bundle_samples
 
 # Exports
@@ -22,7 +21,7 @@ abstract type ProposalStyle end
 """
     DensityModel{F<:Function} <: AbstractModel
 
-`DensityModel` wraps around a self-contained log-liklihood function `ℓπ`.
+`DensityModel` wraps around a self-contained log-liklihood function `logdensity`.
 
 Example:
 
@@ -32,52 +31,69 @@ DensityModel
 ```
 """
 struct DensityModel{F<:Function} <: AbstractModel
-    ℓπ :: F
+    logdensity :: F
 end
 
 # Create a very basic Transition type, only stores the 
 # parameter draws and the log probability of the draw.
-struct Transition{T<:Union{Vector{<:Real}, <:Real}, L<:Real} <: AbstractTransition
-    θ :: T
+struct Transition{T<:Union{Vector, Real, NamedTuple}, L<:Real} <: AbstractTransition
+    params :: T
     lp :: L
 end
 
 # Store the new draw and its log density.
-Transition(model::M, θ::T) where {M<:DensityModel, T} = Transition(θ, ℓπ(model, θ))
+Transition(model::M, params::T) where {M<:DensityModel, T} = Transition(params, logdensity(model, params))
 
 # Tell the interface what transition type we would like to use.
-transition_type(model::DensityModel, spl::Metropolis) = typeof(Transition(spl.init_θ, ℓπ(model, spl.init_θ)))
+transition_type(model::DensityModel, spl::Metropolis) = typeof(Transition(spl.init_params, logdensity(model, spl.init_params)))
 
 # Calculate the density of the model given some parameterization.
-ℓπ(model::DensityModel, θ::T) where T = model.ℓπ(θ)
-ℓπ(model::DensityModel, t::Transition) = t.lp
+logdensity(model::DensityModel, params) = model.logdensity(params)
+logdensity(model::DensityModel, t::Transition) = t.lp
 
 # A basic chains constructor that works with the Transition struct we defined.
 function bundle_samples(
     rng::AbstractRNG, 
-    ℓ::DensityModel, 
+    model::DensityModel, 
     s::Metropolis, 
     N::Integer, 
     ts::Vector{T}; 
     param_names=missing,
     kwargs...
 ) where {ModelType<:AbstractModel, T<:AbstractTransition}
-    # Turn all the transitions into a vector-of-vectors.
-    vals = copy(reduce(hcat,[vcat(t.θ, t.lp) for t in ts])')
+    return ts
+end
 
+function bundle_samples(
+    rng::AbstractRNG, 
+    model::DensityModel, 
+    s::Metropolis, 
+    N::Integer, 
+    ts::Vector{T},
+    chain_type::Type{NamedTuple}; 
+    param_names=missing,
+    kwargs...
+) where {ModelType<:AbstractModel, T<:AbstractTransition}
     # Check if we received any parameter names.
     if ismissing(param_names)
-        param_names = ["Parameter $i" for i in 1:length(s.init_θ)]
+        param_names = ["param_$i" for i in 1:length(s.init_params)]
     else
         # Deepcopy to be thread safe.
         param_names = deepcopy(param_names)
     end
 
-    # Add the log density field to the parameter names.
     push!(param_names, "lp")
 
-    # Bundle everything up and return a Chains struct.
-    return Chains(vals, param_names, (internals=["lp"],))
+    # Turn all the transitions into a vector-of-NamedTuple.
+    keys = tuple(Symbol.(param_names)...)
+    nts = [NamedTuple{keys}(tuple(t.params..., t.lp)) for t in ts]
+
+    return nts
+end
+
+function __init__()
+    @require MCMCChains="c7f686f2-ff18-58e9-bc7b-31028e88f75d" include("mcmcchains-connect.jl")
+    @require StructArrays="09ab397b-f2b6-538f-b94a-2f83cf4a842a" include("structarray-connect.jl")
 end
 
 # Include inference methods.

src/mcmcchains-connect.jl

@@ -0,0 +1,30 @@
+import MCMCChains: Chains
+
+# A basic chains constructor that works with the Transition struct we defined.
+function bundle_samples(
+    rng::AbstractRNG, 
+    model::DensityModel, 
+    s::Metropolis, 
+    N::Integer, 
+    ts::Vector{T},
+    chain_type::Type{Chains}; 
+    param_names=missing,
+    kwargs...
+) where {ModelType<:AbstractModel, T<:AbstractTransition}
+    # Turn all the transitions into a vector-of-vectors.
+    vals = copy(reduce(hcat,[vcat(t.params, t.lp) for t in ts])')
+
+    # Check if we received any parameter names.
+    if ismissing(param_names)
+        param_names = ["Parameter $i" for i in 1:length(s.init_params)]
+    else
+        # Deepcopy to be thread safe.
+        param_names = deepcopy(param_names)
+    end
+
+    # Add the log density field to the parameter names.
+    push!(param_names, "lp")
+
+    # Bundle everything up and return a Chains struct.
+    return Chains(vals, param_names, (internals=["lp"],))
+end