elementwise_loglikelihoods improvement for multiple chains (#171)

In hindsight, I'm not super-happy with how we handle multiple chains. Right now, the structure will be completely flattened.

This PR makes it so that the returned dict is `vn => Vector{Vector{Float64}}` rather than `Vector{Float64}`. Thoughts?

Author: torfjelde

src/loglikelihoods.jl

@@ -29,6 +29,42 @@ function Base.push!(
     ctx.loglikelihoods[vn] = logp
 end
 
+function Base.push!(
+    ctx::ElementwiseLikelihoodContext{Dict{String, Vector{Float64}}},
+    vn::VarName,
+    logp::Real
+)
+    lookup = ctx.loglikelihoods
+    ℓ = get!(lookup, string(vn), Float64[])
+    push!(ℓ, logp)
+end
+
+function Base.push!(
+    ctx::ElementwiseLikelihoodContext{Dict{String, Float64}},
+    vn::VarName,
+    logp::Real
+)
+    ctx.loglikelihoods[string(vn)] = logp
+end
+
+function Base.push!(
+    ctx::ElementwiseLikelihoodContext{Dict{String, Vector{Float64}}},
+    vn::String,
+    logp::Real
+)
+    lookup = ctx.loglikelihoods
+    ℓ = get!(lookup, vn, Float64[])
+    push!(ℓ, logp)
+end
+
+function Base.push!(
+    ctx::ElementwiseLikelihoodContext{Dict{String, Float64}},
+    vn::String,
+    logp::Real
+)
+    ctx.loglikelihoods[vn] = logp
+end
+
 
 function tilde_assume(rng, ctx::ElementwiseLikelihoodContext, sampler, right, vn, inds, vi)
     return tilde_assume(rng, ctx.ctx, sampler, right, vn, inds, vi)
@@ -56,11 +92,14 @@ end
 
 
 """
-    elementwise_loglikelihoods(model::Model, chain::Chains)
+    elementwise_loglikelihoods(model::Model, chain::Chains, keytype = String)
 
-Runs `model` on each sample in `chain` returning an array of arrays with
-the i-th element inner arrays corresponding to the the likelihood of the i-th
-observation for that particular sample in `chain`.
+Runs `model` on each sample in `chain` returning a `Dict{String, Matrix{Float64}}`
+with keys corresponding to symbols of the observations, and values being matrices
+of shape `(num_chains, num_samples)`.
+
+`keytype` specifies what the type of the keys used in the returned `Dict` are.
+Currently, only `String` and `VarName` are supported.
 
 # Notes
 Say `y` is a `Vector` of `n` i.i.d. `Normal(μ, σ)` variables, with `μ` and `σ`
@@ -75,10 +114,10 @@ or
 ```julia
 y ~ MvNormal(fill(μ, n), fill(σ, n))
 ```
-Unfortunately, just by looking at the latter statement, it's impossible to tell whether or
-not this is one *single* observation which is `n` dimensional OR if we have *multiple*
-1-dimensional observations. Therefore, `loglikelihoods` will only work with the first
-example.
+Unfortunately, just by looking at the latter statement, it's impossible to tell 
+whether or not this is one *single* observation which is `n` dimensional OR if we
+have *multiple* 1-dimensional observations. Therefore, `loglikelihoods` will only
+work with the first example.
 
 # Examples
 ```julia-repl
@@ -99,19 +138,37 @@ julia> model = demo(randn(3), randn());
 
 julia> chain = sample(model, MH(), 10);
 
-julia> DynamicPPL.elementwise_loglikelihoods(model, chain)
-Dict{String,Array{Float64,1}} with 4 entries:
-  "xs[3]" => [-1.02616, -1.26931, -1.05003, -5.05458, -1.33825, -1.02904, -1.23761, -1.30128, -1.04872, -2.03716]
-  "xs[1]" => [-2.08205, -2.51387, -3.03175, -2.5981, -2.31322, -2.62284, -2.70874, -1.18617, -1.36281, -4.39839]
-  "xs[2]" => [-2.20604, -2.63495, -3.22802, -2.48785, -2.40941, -2.78791, -2.85013, -1.24081, -1.46019, -4.59025]
-  "y"     => [-1.36627, -1.21964, -1.03342, -7.46617, -1.3234, -1.14536, -1.14781, -2.48912, -2.23705, -1.26267]
+julia> elementwise_loglikelihoods(model, chain)
+Dict{String,Array{Float64,2}} with 4 entries:
+  "xs[3]" => [-1.42862; -2.67573; … ; -1.66251; -1.66251]
+  "xs[1]" => [-1.42932; -2.68123; … ; -1.66333; -1.66333]
+  "xs[2]" => [-1.6724; -0.861339; … ; -1.62359; -1.62359]
+  "y"     => [-1.51265; -0.914129; … ; -1.5499; -1.5499]
+
+julia> elementwise_loglikelihoods(model, chain, String)
+Dict{String,Array{Float64,2}} with 4 entries:
+  "xs[3]" => [-1.42862; -2.67573; … ; -1.66251; -1.66251]
+  "xs[1]" => [-1.42932; -2.68123; … ; -1.66333; -1.66333]
+  "xs[2]" => [-1.6724; -0.861339; … ; -1.62359; -1.62359]
+  "y"     => [-1.51265; -0.914129; … ; -1.5499; -1.5499]
+
+julia> elementwise_loglikelihoods(model, chain, VarName)
+Dict{VarName,Array{Float64,2}} with 4 entries:
+  xs[2] => [-1.6724; -0.861339; … ; -1.62359; -1.62359]
+  y     => [-1.51265; -0.914129; … ; -1.5499; -1.5499]
+  xs[1] => [-1.42932; -2.68123; … ; -1.66333; -1.66333]
+  xs[3] => [-1.42862; -2.67573; … ; -1.66251; -1.66251]
 ```
 """
-function elementwise_loglikelihoods(model::Model, chain)
+function elementwise_loglikelihoods(
+    model::Model,
+    chain,
+    keytype::Type{T} = String
+) where {T}
     # Get the data by executing the model once
-    ctx = ElementwiseLikelihoodContext()
     spl = SampleFromPrior()
     vi = VarInfo(model)
+    ctx = ElementwiseLikelihoodContext(Dict{T, Vector{Float64}}())
 
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     for (sample_idx, chain_idx) in iters
@@ -121,7 +178,14 @@ function elementwise_loglikelihoods(model::Model, chain)
         # Execute model
         model(vi, spl, ctx)
     end
-    return ctx.loglikelihoods
+
+    niters = size(chain, 1)
+    nchains = size(chain, 3)
+    loglikelihoods = Dict(
+        varname => reshape(logliks, niters, nchains)
+        for (varname, logliks) in ctx.loglikelihoods
+    )
+    return loglikelihoods
 end
 
 function elementwise_loglikelihoods(model::Model, varinfo::AbstractVarInfo)

test/loglikelihoods.jl

@@ -14,19 +14,20 @@ using .Turing
     xs = randn(3);
     y = randn();
     model = demo(xs, y);
-    chain = sample(model, MH(), 100);
-    results = elementwise_loglikelihoods(model, chain)
-    var_to_likelihoods = Dict(string(varname) => logliks for (varname, logliks) in results)
+    chain = sample(model, MH(), MCMCThreads(), 100, 2);
+    var_to_likelihoods = elementwise_loglikelihoods(model, chain)
     @test haskey(var_to_likelihoods, "xs[1]")
     @test haskey(var_to_likelihoods, "xs[2]")
     @test haskey(var_to_likelihoods, "xs[3]")
     @test haskey(var_to_likelihoods, "y")
 
-    for (i, (s, m)) in enumerate(zip(chain[:s], chain[:m]))
-        @test logpdf(Normal(m, √s), xs[1]) == var_to_likelihoods["xs[1]"][i]
-        @test logpdf(Normal(m, √s), xs[2]) == var_to_likelihoods["xs[2]"][i]
-        @test logpdf(Normal(m, √s), xs[3]) == var_to_likelihoods["xs[3]"][i]
-        @test logpdf(Normal(m, √s), y) == var_to_likelihoods["y"][i]
+    for chain_idx in MCMCChains.chains(chain)
+        for (i, (s, m)) in enumerate(zip(chain[:, :s, chain_idx], chain[:, :m, chain_idx]))
+            @test logpdf(Normal(m, √s), xs[1]) == var_to_likelihoods["xs[1]"][i, chain_idx]
+            @test logpdf(Normal(m, √s), xs[2]) == var_to_likelihoods["xs[2]"][i, chain_idx]
+            @test logpdf(Normal(m, √s), xs[3]) == var_to_likelihoods["xs[3]"][i, chain_idx]
+            @test logpdf(Normal(m, √s), y) == var_to_likelihoods["y"][i, chain_idx]
+        end
     end
 
     var_info = VarInfo(model)