Fix for generated_quantities

ext/DynamicPPLMCMCChainsExt.jl

@@ -9,16 +9,20 @@ else
 end
 
 _has_varname_to_symbol(info::NamedTuple{names}) where {names} = :varname_to_symbol in names
-function DynamicPPL.supports_varname_indexing(chain::MCMCChains.Chains)
-    return _has_varname_to_symbol(chain.info)
+function _check_varname_indexing(c::MCMCChains.Chains)
+    DynamicPPL.supports_varname_indexing(c) || error("Chains do not support indexing using $vn.")
 end
 
-# TODO: Add proper overload of `Base.getindex` to Turing.jl?
-function _getindex(c::MCMCChains.Chains, sample_idx, vn::DynamicPPL.VarName, chain_idx)
-    DynamicPPL.supports_varname_indexing(c) ||
-        error("Chains do not support indexing using $vn.")
+# A few methods needed.
+DynamicPPL.supports_varname_indexing(chain::MCMCChains.Chains) = _has_varname_to_symbol(chain.info)
+function DynamicPPL.getindex_varname(c::MCMCChains.Chains, sample_idx, vn::DynamicPPL.VarName, chain_idx)
+    _check_varname_indexing(c)
     return c[sample_idx, c.info.varname_to_symbol[vn], chain_idx]
 end
+function DynamicPPL.varnames(c::MCMCChains.Chains)
+    _check_varname_indexing(c)
+    return keys(c.info.varname_to_symbol)
+end
 
 function DynamicPPL.generated_quantities(
     model::DynamicPPL.Model, chain_full::MCMCChains.Chains
@@ -27,28 +31,10 @@ function DynamicPPL.generated_quantities(
     varinfo = DynamicPPL.VarInfo(model)
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     return map(iters) do (sample_idx, chain_idx)
-        if DynamicPPL.supports_varname_indexing(chain)
-            # First we need to set every variable to be resampled.
-            for vn in keys(varinfo)
-                DynamicPPL.set_flag!(varinfo, vn, "del")
-            end
-            # Then we set the variables in `varinfo` from `chain`.
-            for vn in keys(chain.info.varname_to_symbol)
-                vn_updated = DynamicPPL.nested_setindex_maybe!(
-                    varinfo, _getindex(chain, sample_idx, vn, chain_idx), vn
-                )
+        # Update the varinfo with the current sample and make variables not present in `chain`
+        # to be sampled.
+        DynamicPPL.setval_and_resample!(varinfo, chain, sample_idx, chain_idx)
 
-                # Unset the `del` flag if we found something.
-                if vn_updated !== nothing
-                    # NOTE: This will be triggered even if only a subset of a variable has been set!
-                    DynamicPPL.unset_flag!(varinfo, vn_updated, "del")
-                end
-            end
-        else
-            # NOTE: This can be quite unreliable (but will warn the uesr in that case).
-            # Hence the above path is much more preferable.
-            DynamicPPL.setval_and_resample!(varinfo, chain, sample_idx, chain_idx)
-        end
         # TODO: Some of the variables can be a view into the `varinfo`, so we need to
         # `deepcopy` the `varinfo` before passing it to `model`.
         model(deepcopy(varinfo))

src/DynamicPPL.jl

@@ -154,6 +154,7 @@
 # Necessary forward declarations
 include("utils.jl")
 include("selector.jl")
+include("chains.jl")
 include("model.jl")
 include("sampler.jl")
 include("varname.jl")
@@ -180,8 +181,8 @@
 end
 
 @static if !isdefined(Base, :get_extension)
     function __init__()
         @require MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" include(
             "../ext/DynamicPPLMCMCChainsExt.jl"
         )
     end

src/chains.jl

@@ -0,0 +1,25 @@
+"""
+    supports_chain_indexing(chain::AbstractChains)
+
+Return `true` if `chain` supports indexing using `VarName` in place of the
+variable name index.
+"""
+supports_varname_indexing(::AbstractChains) = false
+
+"""
+    getindex_varname(chain::AbstractChains, sample_idx, varname::VarName, chain_idx)
+
+Return the value of `varname` in `chain` at `sample_idx` and `chain_idx`.
+
+Whether this method is implemented for `chains` is indicated by [`supports_varname_indexing`](@ref).
+"""
+function getindex_varname end
+
+"""
+    varnames(chains::AbstractChains)
+
+Return an iterator over the varnames present in `chains`.
+
+Whether this method is implemented for `chains` is indicated by [`supports_varname_indexing`](@ref).
+"""
+function varnames end

src/model.jl

@@ -1257,9 +1257,6 @@ function generated_quantities(model::Model, chain::AbstractChains)
     end
 end
 
-# HACK: Be better.
-supports_varname_indexing(chain::AbstractChains) = false
-
 """
     generated_quantities(model::Model, parameters::NamedTuple)
     generated_quantities(model::Model, values, keys)

src/varinfo.jl

@@ -1064,8 +1064,8 @@
     return Expr(:||, false, out...)
 end
 
 function nested_setindex_maybe!(vi::UntypedVarInfo, val, vn::VarName)
     return _nested_setindex_maybe!(vi, getmetadata(vi, vn), val, vn)
 end
 function nested_setindex_maybe!(
     vi::VarInfo{<:NamedTuple{names}}, val, vn::VarName{sym}
@@ -1080,8 +1080,8 @@
     # If `vn` is in `vns`, then we can just use the standard `setindex!`.
     vns = md.vns
     if vn in vns
         setindex!(vi, val, vn)
         return vn
     end
 
     # Otherwise, we need to check if either of the `vns` subsumes `vn`.
@@ -1636,7 +1636,26 @@
 function setval_and_resample!(
     vi::VarInfoOrThreadSafeVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int
 )
-    return setval_and_resample!(vi, chains.value[sample_idx, :, chain_idx], keys(chains))
+    if supports_varname_indexing(chains)
+        # First we need to set every variable to be resampled.
+        for vn in keys(vi)
+            set_flag!(vi, vn, "del")
+        end
+        # Then we set the variables in `varinfo` from `chain`.
+        for vn in varnames(chains)
+            vn_updated = nested_setindex_maybe!(
+                vi, getindex_varname(chains, sample_idx, vn, chain_idx), vn
+            )
+
+            # Unset the `del` flag if we found something.
+            if vn_updated !== nothing
+                # NOTE: This will be triggered even if only a subset of a variable has been set!
+                unset_flag!(vi, vn_updated, "del")
+            end
+        end
+    else
+        setval_and_resample!(vi, chains.value[sample_idx, :, chain_idx], keys(chains))
+    end
 end
 
 function _setval_and_resample_kernel!(