Fix for generated_quantities

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.23.16"
+version = "0.23.17"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -16,19 +16,11 @@ LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
-[weakdeps]
-MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-
-[extensions]
-DynamicPPLMCMCChainsExt = ["MCMCChains"]
-
-[extras]
-MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-
 [compat]
 AbstractMCMC = "2, 3.0, 4"
 AbstractPPL = "0.6"
@@ -39,9 +31,19 @@ ConstructionBase = "1"
 Distributions = "0.23.8, 0.24, 0.25"
 DocStringExtensions = "0.8, 0.9"
 LogDensityProblems = "2"
-MacroTools = "0.5.6"
 MCMCChains = "6"
+MacroTools = "0.5.6"
 OrderedCollections = "1"
+Requires = "1"
 Setfield = "0.7.1, 0.8, 1"
 ZygoteRules = "0.2"
 julia = "1.6"
+
+[extensions]
+DynamicPPLMCMCChainsExt = ["MCMCChains"]
+
+[extras]
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+
+[weakdeps]
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"

ext/DynamicPPLMCMCChainsExt.jl

@@ -1,15 +1,57 @@
 module DynamicPPLMCMCChainsExt
 
-using DynamicPPL: DynamicPPL
-using MCMCChains: MCMCChains
+if isdefined(Base, :get_extension)
+    using DynamicPPL: DynamicPPL
+    using MCMCChains: MCMCChains
+else
+    using ..DynamicPPL: DynamicPPL
+    using ..MCMCChains: MCMCChains
+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)
+end
 
-function DynamicPPL.generated_quantities(model::DynamicPPL.Model, chain::MCMCChains.Chains)
-    chain_parameters = MCMCChains.get_sections(chain, :parameters)
+# 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.")
+    return c[sample_idx, c.info.varname_to_symbol[vn], chain_idx]
+end
+
+function DynamicPPL.generated_quantities(
+    model::DynamicPPL.Model, chain_full::MCMCChains.Chains
+)
+    chain = MCMCChains.get_sections(chain_full, :parameters)
     varinfo = DynamicPPL.VarInfo(model)
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     return map(iters) do (sample_idx, chain_idx)
-        DynamicPPL.setval_and_resample!(varinfo, chain_parameters, sample_idx, chain_idx)
-        model(varinfo)
+        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
+                )
+
+                # 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))
     end
 end
 

src/DynamicPPL.jl

@@ -175,4 +175,16 @@ include("logdensityfunction.jl")
 include("model_utils.jl")
 include("extract_priors.jl")
 
+if !isdefined(Base, :get_extension)
+    using Requires
+end
+
+function __init__()
+    @static if !isdefined(Base, :get_extension)
+        @require MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" include(
+            "../ext/DynamicPPLMCMCChainsExt.jl"
+        )
+    end
+end
+
 end # module

src/model.jl

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

src/utils.jl

@@ -501,6 +501,42 @@
     return current_parent, current_child, condition(current_parent)
 end
 
+"""
+    remove_parent_lens(vn_parent::VarName, vn_child::VarName)
+
+Remove the parent lens `vn_parent` from `vn_child`.
+
+# Examples
+```jldoctest
+julia> DynamicPPL.remove_parent_lens(@varname(x), @varname(x.a))
+(@lens _.a)
+
+julia> DynamicPPL.remove_parent_lens(@varname(x), @varname(x.a[1]))
+(@lens _.a[1])
+
+julia> DynamicPPL.remove_parent_lens(@varname(x.a), @varname(x.a[1]))
+(@lens _[1])
+
+julia> DynamicPPL.remove_parent_lens(@varname(x.a), @varname(x.a[1].b))
+(@lens _[1].b)
+
+julia> DynamicPPL.remove_parent_lens(@varname(x.a), @varname(x.a))
+ERROR: Could not find x.a in x.a
+
+julia> DynamicPPL.remove_parent_lens(@varname(x.a[2]), @varname(x.a[1]))
+ERROR: Could not find x.a[2] in x.a[1]
+```
+"""
+function remove_parent_lens(vn_parent::VarName{sym}, vn_child::VarName{sym}) where {sym}
+    _, child, issuccess = splitlens(getlens(vn_child)) do lens
+        l = lens === nothing ? Setfield.IdentityLens() : lens
+        VarName(vn_child, l) == vn_parent
+    end
+
+    issuccess || error("Could not find $vn_parent in $vn_child")
+    return child
+end
+
 # HACK: All of these are related to https://github.com/JuliaFolds/BangBang.jl/issues/233
 # and https://github.com/JuliaFolds/BangBang.jl/pull/238.
 # HACK(torfjelde): Avoids type-instability in `dot_assume` for `SimpleVarInfo`.
@@ -1045,3 +1081,26 @@
         for I in CartesianIndices(x) if I[1] <= I[2]
     )
 end
+
+# TODO: Remove as soon as https://github.com/JuliaObjects/ConstructionBase.jl/pull/80 goes through.
+ConstructionBase.setproperties(C::LinearAlgebra.Cholesky, ::NamedTuple{()}) = C
+function ConstructionBase.setproperties(C::LinearAlgebra.Cholesky, patch::NamedTuple{(:L,)})
+    return LinearAlgebra.Cholesky(
+        C.uplo === 'U' ? permutedims(patch.L) : patch.L, C.uplo, C.info
+    )
+end
+function ConstructionBase.setproperties(C::LinearAlgebra.Cholesky, patch::NamedTuple{(:U,)})
+    return LinearAlgebra.Cholesky(
+        C.uplo === 'L' ? permutedims(patch.U) : patch.U, C.uplo, C.info
+    )
+end
+function ConstructionBase.setproperties(
+    C::LinearAlgebra.Cholesky, patch::NamedTuple{(:UL,)}
+)
+    return LinearAlgebra.Cholesky(patch.UL, C.uplo, C.info)
+end
+@nospecialize function ConstructionBase.setproperties(
+    C::LinearAlgebra.Cholesky, patch::NamedTuple
+)
+    return error("Can only patch one of :L, :U, :UL at the time")
+end

src/varinfo.jl

@@ -1064,6 +1064,41 @@
     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}
+) where {names,sym}
+    return if sym in names
+        _nested_setindex_maybe!(vi, getmetadata(vi, vn), val, vn)
+    else
+        nothing
+    end
+end
+function _nested_setindex_maybe!(vi::VarInfo, md::Metadata, val, vn::VarName)
+    # 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`.
+    i = findfirst(Base.Fix2(subsumes, vn), vns)
+    i === nothing && return nothing
+
+    vn_parent = vns[i]
+    dist = getdist(md, vn_parent)
+    val_parent = getindex(vi, vn_parent, dist)  # TODO: Ensure that we're working with a view here.
+    # Split the varname into its tail lens.
+    lens = remove_parent_lens(vn_parent, vn)
+    # Update the value for the parent.
+    val_parent_updated = set!!(val_parent, lens, val)
+    setindex!(vi, val_parent_updated, vn_parent)
+    return vn_parent
+end
+
 # The default getindex & setindex!() for get & set values
 # NOTE: vi[vn] will always transform the variable to its original space and Julia type
 getindex(vi::VarInfo, vn::VarName) = getindex(vi, vn, getdist(vi, vn))
@@ -1131,7 +1166,8 @@
 """
 setindex!(vi::VarInfo, val, vn::VarName) = (setval!(vi, val, vn); return vi)
 function BangBang.setindex!!(vi::VarInfo, val, vn::VarName)
-    return (setindex!(vi, val, vn); return vi)
+    setindex!(vi, val, vn)
+    return vi
 end
 
 """

test/Project.toml

@@ -2,6 +2,7 @@
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
@@ -24,6 +25,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 AbstractMCMC = "2.1, 3.0, 4"
 AbstractPPL = "0.6"
 Bijectors = "0.13"
+Compat = "4.3.0"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 Documenter = "0.26.1, 0.27"

test/model.jl

@@ -278,4 +278,55 @@ end
             @test DynamicPPL.TestUtils.posterior_mean(model) isa typeof(x)
         end
     end
+
+    @testset "generated_quantities on `LKJCholesky`" begin
+        n = 10
+        d = 2
+        model = DynamicPPL.TestUtils.demo_lkjchol(d)
+        xs = [model().x for _ in 1:n]
+
+        # Extract varnames and values.
+        vns_and_vals_xs = map(
+            collect ∘ Base.Fix1(DynamicPPL.varname_and_value_leaves, @varname(x)), xs
+        )
+        vns = map(first, first(vns_and_vals_xs))
+        vals = map(vns_and_vals_xs) do vns_and_vals
+            map(last, vns_and_vals)
+        end
+
+        # Construct the chain.
+        syms = map(Symbol, vns)
+        vns_to_syms = OrderedDict{VarName,Any}(zip(vns, syms))
+
+        chain = MCMCChains.Chains(
+            permutedims(stack(vals)), syms; info=(varname_to_symbol=vns_to_syms,)
+        )
+        display(chain)
+
+        # Test!
+        results = generated_quantities(model, chain)
+        for (x_true, result) in zip(xs, results)
+            @test x_true.UL == result.x.UL
+        end
+
+        # With variables that aren't in the `model`.
+        vns_to_syms_with_extra = let d = deepcopy(vns_to_syms)
+            d[@varname(y)] = :y
+            d
+        end
+        vals_with_extra = map(enumerate(vals)) do (i, v)
+            vcat(v, i)
+        end
+        chain_with_extra = MCMCChains.Chains(
+            permutedims(stack(vals_with_extra)),
+            vcat(syms, [:y]);
+            info=(varname_to_symbol=vns_to_syms_with_extra,),
+        )
+        display(chain_with_extra)
+        # Test!
+        results = generated_quantities(model, chain_with_extra)
+        for (x_true, result) in zip(xs, results)
+            @test x_true.UL == result.x.UL
+        end
+    end
 end

test/runtests.jl

@@ -11,6 +11,7 @@ using MCMCChains
 using Tracker
 using Zygote
 using Setfield
+using Compat
 
 using Distributed
 using LinearAlgebra