Update to AbstractMCMC 2

Author: devmotion

Project.toml

@@ -9,7 +9,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
 [compat]
-AbstractMCMC = "1"
+AbstractMCMC = "2"
 Distributions = "0.20, 0.21, 0.22, 0.23"
 Requires = "1.0"
 julia = "1"

src/AdvancedMH.jl

@@ -35,7 +35,7 @@ struct DensityModel{F<:Function} <: AbstractMCMC.AbstractModel
     logdensity :: F
 end
 
-# Create a very basic Transition type, only stores the 
+# Create a very basic Transition type, only stores the
 # parameter draws and the log probability of the draw.
 struct Transition{T<:Union{Vector, Real, NamedTuple}, L<:Real}
     params :: T
@@ -51,25 +51,11 @@ logdensity(model::DensityModel, t::Transition) = t.lp
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    rng::Random.AbstractRNG, 
-    model::DensityModel, 
-    s::MHSampler, 
-    N::Integer,
-    ts::Vector,
-    chain_type::Type{Any}; 
-    param_names=missing,
-    kwargs...
-)
-    return ts
-end
-
-function AbstractMCMC.bundle_samples(
-    rng::Random.AbstractRNG, 
-    model::DensityModel, 
-    s::MHSampler, 
-    N::Integer, 
-    ts::Vector,
-    chain_type::Type{Vector{NamedTuple}}; 
+    ts::Vector{<:Transition},
+    model::DensityModel,
+    sampler::MHSampler,
+    state,
+    chain_type::Type{Vector{NamedTuple}};
     param_names=missing,
     kwargs...
 )

src/emcee.jl

@@ -3,43 +3,44 @@ struct Ensemble{D} <: MHSampler
     proposal::D
 end
 
-# Define the first step! function, which is called at the 
-# beginning of sampling. Return the initial parameter used
-# to define the sampler.
-function AbstractMCMC.step!(
+# Define the first sampling step.
+# Return a 2-tuple consisting of the initial sample and the initial state.
+# In this case they are identical.
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DensityModel,
-    spl::Ensemble,
-    N::Integer,
-    ::Nothing;
+    spl::Ensemble;
     init_params = nothing,
     kwargs...,
 )
     if init_params === nothing
-        return propose(rng, spl, model)
+        transitions = propose(rng, spl, model)
     else
-        return Transition(model, init_params)
+        transitions = [Transition(model, x) for x in init_params]
     end
+
+    return transitions, transitions
 end
 
-# Define the other step functions. Returns a Transition containing
-# either a new proposal (if accepted) or the previous proposal 
-# (if not accepted).
-function AbstractMCMC.step!(
+# Define the other sampling steps.
+# Return a 2-tuple consisting of the next sample and the the next state.
+# In this case they are identical, and for each walker they are either a new proposal
+# (if accepted) or the previous proposal (if not accepted).
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DensityModel,
     spl::Ensemble,
-    ::Integer,
-    params_prev;
+    params_prev::Vector{<:Transition};
     kwargs...,
 )
     # Generate a new proposal. Accept/reject happens at proposal level.
-    return propose(rng, spl, model, params_prev)
+    transitions = propose(rng, spl, model, params_prev)
+    return transitions, transitions
 end
 
 #
 # Initial proposal
-# 
+#
 function propose(rng::Random.AbstractRNG, spl::Ensemble, model::DensityModel)
     # Make the first proposal with a static draw from the prior.
     static_prop = StaticProposal(spl.proposal.proposal)
@@ -49,8 +50,13 @@ end
 
 #
 # Every other proposal
-# 
-function propose(rng::Random.AbstractRNG, spl::Ensemble, model::DensityModel, walkers::Vector{W}) where {W<:Transition}
+#
+function propose(
+    rng::Random.AbstractRNG,
+    spl::Ensemble,
+    model::DensityModel,
+    walkers::Vector{<:Transition},
+)
     new_walkers = similar(walkers)
 
     others = 1:(spl.n_walkers - 1)
@@ -64,7 +70,6 @@ function propose(rng::Random.AbstractRNG, spl::Ensemble, model::DensityModel, wa
     return new_walkers
 end
 
-
 #####################################
 # Basic stretch move implementation #
 #####################################

src/mcmcchains-connect.jl

@@ -2,12 +2,11 @@ import .MCMCChains: Chains
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    rng::Random.AbstractRNG, 
-    model::DensityModel, 
-    s::MHSampler, 
-    N::Integer, 
-    ts,
-    chain_type::Type{Chains}; 
+    ts::Vector{<:Transition},
+    model::DensityModel,
+    sampler::MHSampler,
+    state,
+    chain_type::Type{Chains};
     param_names=missing,
     kwargs...
 )
@@ -16,7 +15,7 @@ function AbstractMCMC.bundle_samples(
 
     # Check if we received any parameter names.
     if ismissing(param_names)
-        param_names = [Symbol(:param_, i) for i in 1:length(s.init_params)]
+        param_names = [Symbol(:param_, i) for i in 1:length(keys(ts[1].params))]
     else
         # Generate new array to be thread safe.
         param_names = Symbol.(param_names)
@@ -30,22 +29,23 @@ function AbstractMCMC.bundle_samples(
 end
 
 function AbstractMCMC.bundle_samples(
-    rng::Random.AbstractRNG, 
-    model::DensityModel, 
-    s::Ensemble, 
-    N::Integer, 
-    ts::Vector,
-    chain_type::Type{Chains}; 
+    ts::Vector{<:Vector{<:Transition}},
+    model::DensityModel,
+    sampler::Ensemble,
+    state,
+    chain_type::Type{Chains};
     param_names=missing,
     kwargs...
 )
     # Preallocate return array
     # NOTE: requires constant dimensionality.
     n_params = length(ts[1][1].params)
-    vals = Array{Float64, 3}(undef, N, n_params + 1, s.n_walkers)  # add 1 parameter for lp
+    nsamples = length(ts)
+    # add 1 parameter for lp
+    vals = Array{Float64, 3}(undef, nsamples, n_params + 1, sampler.n_walkers)
 
-    for n in 1:N
-        for i in 1:s.n_walkers
+    for n in 1:nsamples
+        for i in 1:sampler.n_walkers
             walker = ts[n][i]
             for j in 1:n_params
                 vals[n, j, i] = walker.params[j]
@@ -56,7 +56,7 @@ function AbstractMCMC.bundle_samples(
 
     # Check if we received any parameter names.
     if ismissing(param_names)
-        param_names = [Symbol(:param_, i) for i in 1:length(ts[1][1].params)]
+        param_names = [Symbol(:param_, i) for i in 1:length(keys(ts[1][1].params))]
     else
         # Generate new array to be thread safe.
         param_names = Symbol.(param_names)

src/mh-core.jl

@@ -169,34 +169,34 @@ function q(
     end
 end
 
-# Define the first step! function, which is called at the 
-# beginning of sampling. Return the initial parameter used
-# to define the sampler.
-function AbstractMCMC.step!(
+# Define the first sampling step.
+# Return a 2-tuple consisting of the initial sample and the initial state.
+# In this case they are identical.
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DensityModel,
-    spl::MetropolisHastings,
-    N::Integer,
-    ::Nothing;
+    spl::MetropolisHastings;
     init_params=nothing,
     kwargs...
 )
     if init_params === nothing
-        return propose(rng, spl, model)
+        transition = propose(rng, spl, model)
     else
-        return Transition(model, init_params)
+        transition = Transition(model, init_params)
     end
+
+    return transition, transition
 end
 
-# Define the other step functions. Returns a Transition containing
-# either a new proposal (if accepted) or the previous proposal 
-# (if not accepted).
-function AbstractMCMC.step!(
+# Define the other sampling steps.
+# Return a 2-tuple consisting of the next sample and the the next state.
+# In this case they are identical, and either a new proposal (if accepted)
+# or the previous proposal (if not accepted).
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DensityModel,
     spl::MetropolisHastings,
-    ::Integer,
-    params_prev;
+    params_prev::Transition;
     kwargs...
 )
     # Generate a new proposal.
@@ -208,8 +208,8 @@ function AbstractMCMC.step!(
 
     # Decide whether to return the previous params or the new one.
     if -Random.randexp(rng) < logα
-        return params
+        return params, params
     else
-        return params_prev
+        return params_prev, params_prev
     end
 end

src/structarray-connect.jl

@@ -2,18 +2,18 @@ import .StructArrays: StructArray
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    rng::Random.AbstractRNG, 
-    model::DensityModel, 
-    s::MHSampler, 
-    N::Integer, 
-    ts::Vector,
-    chain_type::Type{StructArray}; 
-    param_names=missing,
+    ts,
+    model::DensityModel,
+    sampler::MHSampler,
+    state,
+    chain_type::Type{StructArray};
     kwargs...
 )
-    samples = AbstractMCMC.bundle_samples(rng, model, s, N, ts, Vector{NamedTuple};
-                                          param_names=param_names, kwargs...)
+    samples = AbstractMCMC.bundle_samples(
+        ts, model, sampler, state, Vector{NamedTuple};
+        kwargs...
+    )
     return StructArray(samples)
 end
 
-AbstractMCMC.chainscat(c::StructArray, cs::StructArray...) = vcat(c, cs...)
+AbstractMCMC.chainscat(c::StructArray, cs::StructArray...) = vcat(c, cs...)