Addressing comments

Author: cpfiffer

src/AdvancedMH.jl

@@ -8,7 +8,8 @@ using Requires
 import Random
 
 # Exports
-export MetropolisHastings, DensityModel, RWMH, StaticMH, StaticProposal, RandomWalkProposal
+export MetropolisHastings, DensityModel, RWMH, StaticMH, StaticProposal, 
+    RandomWalkProposal, Ensemble, StretchProposal
 
 # Reexports
 export sample, MCMCThreads, MCMCDistributed

src/emcee.jl

@@ -47,21 +47,16 @@ function propose(rng::Random.AbstractRNG, spl::Ensemble, model::DensityModel)
     return [propose(rng, mh_spl, model) for _ in 1:spl.n_walkers]
 end
 
-
-
 #
 # Every other proposal
 # 
 function propose(rng::Random.AbstractRNG, spl::Ensemble, model::DensityModel, walkers::Vector{W}) where {W<:Transition}
-    new_walkers = Vector{W}(undef, spl.n_walkers)
-    interval = 1:spl.n_walkers
-
+    new_walkers = similar(walkers)
 
-    nwalkers = spl.n_walkers
-    others = 1:(nwalkers - 1)
-    for i in interval
+    others = 1:(spl.n_walkers - 1)
+    for i in 1:spl.n_walkers
         walker = walkers[i]
-        idx = mod1(i + rand(rng, others), nwalkers)
+        idx = mod1(i + rand(rng, others), spl.n_walkers)
         other_walker = walkers[idx]
         new_walkers[i] = move(rng, spl, model, walker, other_walker)
     end
@@ -95,7 +90,7 @@ function move(
     alphamult = (n - 1) * log(z)
 
     # Make new parameters
-    y = walker.params + z .* (other_walker.params - walker.params)
+    y = @. walker.params + z * (other_walker.params - walker.params)
 
     # Construct a new walker
     new_walker = Transition(model, y)

src/mcmcchains-connect.jl

@@ -39,17 +39,20 @@ function AbstractMCMC.bundle_samples(
     param_names=missing,
     kwargs...
 )
-    # return ts
-    vals = mapreduce(
-        t -> map(i -> vcat(ts[t][i].params, 
-                 ts[t][i].lp, t, i),
-                 1:length(ts[t])), 
-        vcat, 
-        1:length(ts))
-    
-    vals = Array(reduce(hcat, vals)')
+    # 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
 
-    # return vals
+    for n in 1:N
+        for i in 1:s.n_walkers
+            walker = ts[n][i]
+            for j in 1:n_params
+                vals[n, j, i] = walker.params[j]
+            end
+            vals[n, n_params + 1, i] = walker.lp
+        end
+    end
 
     # Check if we received any parameter names.
     if ismissing(param_names)
@@ -60,8 +63,8 @@ function AbstractMCMC.bundle_samples(
     end
 
     # Add the log density field to the parameter names.
-    push!(param_names, "lp", "iteration", "walker")
+    push!(param_names, "lp")
 
     # Bundle everything up and return a Chains struct.
-    return Chains(vals, param_names, (internals=["lp", "iteration", "walker"],))
+    return Chains(vals, param_names, (internals=["lp"],))
 end

src/structarray-connect.jl

@@ -4,7 +4,7 @@ import .StructArrays: StructArray
 function AbstractMCMC.bundle_samples(
     rng::Random.AbstractRNG, 
     model::DensityModel, 
-    s::Metropolis, 
+    s::MHSampler, 
     N::Integer, 
     ts::Vector,
     chain_type::Type{StructArray};