Attempt at implementation of VarNameVector (Metadata alternative)

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.29"
+version = "0.30"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -46,7 +46,7 @@ AbstractMCMC = "5"
 AbstractPPL = "0.8.4"
 Accessors = "0.1"
 BangBang = "0.4.1"
-Bijectors = "0.13.9"
+Bijectors = "0.13.18"
 ChainRulesCore = "1"
 Compat = "4"
 ConstructionBase = "1.5.4"

docs/make.jl

@@ -23,7 +23,7 @@ makedocs(;
         "Home" => "index.md",
         "API" => "api.md",
         "Tutorials" => ["tutorials/prob-interface.md"],
-        "Internals" => ["internals/transformations.md"],
+        "Internals" => ["internals/varinfo.md", "internals/transformations.md"],
     ],
     checkdocs=:exports,
     doctest=false,

ext/DynamicPPLChainRulesCoreExt.jl

@@ -24,4 +24,6 @@ ChainRulesCore.@non_differentiable DynamicPPL.updategid!(
 # No need + causes issues for some AD backends, e.g. Zygote.
 ChainRulesCore.@non_differentiable DynamicPPL.infer_nested_eltype(x)
 
+ChainRulesCore.@non_differentiable DynamicPPL.recontiguify_ranges!(ranges)
+
 end # module

ext/DynamicPPLMCMCChainsExt.jl

@@ -42,21 +42,152 @@ function DynamicPPL.varnames(c::MCMCChains.Chains)
     return keys(c.info.varname_to_symbol)
 end
 
+"""
+    generated_quantities(model::Model, chain::MCMCChains.Chains)
+
+Execute `model` for each of the samples in `chain` and return an array of the values
+returned by the `model` for each sample.
+
+# Examples
+## General
+Often you might have additional quantities computed inside the model that you want to
+inspect, e.g.
+```julia
+@model function demo(x)
+    # sample and observe
+    θ ~ Prior()
+    x ~ Likelihood()
+    return interesting_quantity(θ, x)
+end
+m = demo(data)
+chain = sample(m, alg, n)
+# To inspect the `interesting_quantity(θ, x)` where `θ` is replaced by samples
+# from the posterior/`chain`:
+generated_quantities(m, chain) # <= results in a `Vector` of returned values
+                               #    from `interesting_quantity(θ, x)`
+```
+## Concrete (and simple)
+```julia
+julia> using DynamicPPL, Turing
+
+julia> @model function demo(xs)
+           s ~ InverseGamma(2, 3)
+           m_shifted ~ Normal(10, √s)
+           m = m_shifted - 10
+
+           for i in eachindex(xs)
+               xs[i] ~ Normal(m, √s)
+           end
+
+           return (m, )
+       end
+demo (generic function with 1 method)
+
+julia> model = demo(randn(10));
+
+julia> chain = sample(model, MH(), 10);
+
+julia> generated_quantities(model, chain)
+10×1 Array{Tuple{Float64},2}:
+ (2.1964758025119338,)
+ (2.1964758025119338,)
+ (0.09270081916291417,)
+ (0.09270081916291417,)
+ (0.09270081916291417,)
+ (0.09270081916291417,)
+ (0.09270081916291417,)
+ (0.043088571494005024,)
+ (-0.16489786710222099,)
+ (-0.16489786710222099,)
+```
+"""
 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)
-        # 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)
+        if DynamicPPL.supports_varname_indexing(chain)
+            varname_pairs = _varname_pairs_with_varname_indexing(
+                chain, varinfo, sample_idx, chain_idx
+            )
+        else
+            varname_pairs = _varname_pairs_without_varname_indexing(
+                chain, varinfo, sample_idx, chain_idx
+            )
+        end
+        fixed_model = DynamicPPL.fix(model, Dict(varname_pairs))
+        return fixed_model()
+    end
+end
+
+"""
+    _varname_pairs_with_varname_indexing(
+        chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx
+    )
 
-        # 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))
+Get pairs of `VarName => value` for all the variables in the `varinfo`, picking the values
+from the chain.
+
+This implementation assumes `chain` can be indexed using variable names, and is the
+preffered implementation.
+"""
+function _varname_pairs_with_varname_indexing(
+    chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx
+)
+    vns = DynamicPPL.varnames(chain)
+    vn_parents = Iterators.map(vns) do vn
+        # The call nested_setindex_maybe! is used to handle cases where vn is not
+        # the variable name used in the model, but rather subsumed by one. Except
+        # for the subsumption part, this could be
+        # vn => getindex_varname(chain, sample_idx, vn, chain_idx)
+        # TODO(mhauru) This call to nested_setindex_maybe! is unintuitive.
+        DynamicPPL.nested_setindex_maybe!(
+            varinfo, DynamicPPL.getindex_varname(chain, sample_idx, vn, chain_idx), vn
+        )
     end
+    varname_pairs = Iterators.map(Iterators.filter(!isnothing, vn_parents)) do vn_parent
+        vn_parent => varinfo[vn_parent]
+    end
+    return varname_pairs
+end
+
+"""
+Check which keys in `key_strings` are subsumed by `vn_string` and return the their values.
+
+The subsumption check is done with `DynamicPPL.subsumes_string`, which is quite weak, and
+won't catch all cases. We should get rid of this if we can.
+"""
+# TODO(mhauru) See docstring above.
+function _vcat_subsumed_values(vn_string, values, key_strings)
+    indices = findall(Base.Fix1(DynamicPPL.subsumes_string, vn_string), key_strings)
+    return !isempty(indices) ? reduce(vcat, values[indices]) : nothing
+end
+
+"""
+    _varname_pairs_without_varname_indexing(
+        chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx
+    )
+
+Get pairs of `VarName => value` for all the variables in the `varinfo`, picking the values
+from the chain.
+
+This implementation does not assume that `chain` can be indexed using variable names. It is
+thus not guaranteed to work in cases where the variable names have complex subsumption
+patterns, such as if the model has a variable `x` but the chain stores `x.a[1]`.
+"""
+function _varname_pairs_without_varname_indexing(
+    chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx
+)
+    values = chain.value[sample_idx, :, chain_idx]
+    keys = Base.keys(chain)
+    keys_strings = map(string, keys)
+    varname_pairs = [
+        vn => _vcat_subsumed_values(string(vn), values, keys_strings) for
+        vn in Base.keys(varinfo)
+    ]
+    return varname_pairs
 end
 
 end

src/DynamicPPL.jl

@@ -45,7 +45,9 @@ export AbstractVarInfo,
     VarInfo,
     UntypedVarInfo,
     TypedVarInfo,
+    VectorVarInfo,
     SimpleVarInfo,
+    VarNamedVector,
     push!!,
     empty!!,
     subset,
@@ -175,6 +177,7 @@ include("sampler.jl")
 include("varname.jl")
 include("distribution_wrappers.jl")
 include("contexts.jl")
+include("varnamedvector.jl")
 include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")

src/abstract_varinfo.jl

@@ -295,7 +295,7 @@ julia> # Just use an example model to construct the `VarInfo` because we're lazy
 julia> vi[@varname(s)] = 1.0; vi[@varname(m)] = 2.0;
 
 julia> # For the sake of brevity, let's just check the type.
-       md = values_as(vi); md.s isa DynamicPPL.Metadata
+       md = values_as(vi); md.s isa Union{DynamicPPL.Metadata, DynamicPPL.VarNamedVector}
 true
 
 julia> values_as(vi, NamedTuple)
@@ -321,7 +321,7 @@ julia> # Just use an example model to construct the `VarInfo` because we're lazy
 julia> vi[@varname(s)] = 1.0; vi[@varname(m)] = 2.0;
 
 julia> # For the sake of brevity, let's just check the type.
-       values_as(vi) isa DynamicPPL.Metadata
+       values_as(vi) isa Union{DynamicPPL.Metadata, Vector}
 true
 
 julia> values_as(vi, NamedTuple)
@@ -349,7 +349,7 @@ Determine the default `eltype` of the values returned by `vi[spl]`.
     This should generally not be called explicitly, as it's only used in
     [`matchingvalue`](@ref) to determine the default type to use in place of
     type-parameters passed to the model.
-    
+
     This method is considered legacy, and is likely to be deprecated in the future.
 """
 function Base.eltype(vi::AbstractVarInfo, spl::Union{AbstractSampler,SampleFromPrior})
@@ -363,6 +363,13 @@ function Base.eltype(vi::AbstractVarInfo, spl::Union{AbstractSampler,SampleFromP
     return eltype(T)
 end
 
+"""
+    has_varnamedvector(varinfo::VarInfo)
+
+Returns `true` if `varinfo` uses `VarNamedVector` as metadata.
+"""
+has_varnamedvector(vi::AbstractVarInfo) = false
+
 # TODO: Should relax constraints on `vns` to be `AbstractVector{<:Any}` and just try to convert
 # the `eltype` to `VarName`? This might be useful when someone does `[@varname(x[1]), @varname(m)]` which
 # might result in a `Vector{Any}`.
@@ -554,7 +561,7 @@ end
     link([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
     link([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
 
-Transform the variables in `vi` to their linked space without mutating `vi`, using the transformation `t`. 
+Transform the variables in `vi` to their linked space without mutating `vi`, using the transformation `t`.
 
 If `t` is not provided, `default_transformation(model, vi)` will be used.
 
@@ -573,7 +580,7 @@ end
     invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
     invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
 
-Transform the variables in `vi` to their constrained space, using the (inverse of) 
+Transform the variables in `vi` to their constrained space, using the (inverse of)
 transformation `t`, mutating `vi` if possible.
 
 If `t` is not provided, `default_transformation(model, vi)` will be used.

src/context_implementations.jl

@@ -240,7 +240,10 @@ function assume(
     if haskey(vi, vn)
         # Always overwrite the parameters with new ones for `SampleFromUniform`.
         if sampler isa SampleFromUniform || is_flagged(vi, vn, "del")
-            unset_flag!(vi, vn, "del")
+            # TODO(mhauru) Is it important to unset the flag here? The `true` allows us
+            # to ignore the fact that for VarNamedVector this does nothing, but I'm unsure if
+            # that's okay.
+            unset_flag!(vi, vn, "del", true)
             r = init(rng, dist, sampler)
             f = to_maybe_linked_internal_transform(vi, vn, dist)
             BangBang.setindex!!(vi, f(r), vn)
@@ -516,7 +519,10 @@ function get_and_set_val!(
     if haskey(vi, vns[1])
         # Always overwrite the parameters with new ones for `SampleFromUniform`.
         if spl isa SampleFromUniform || is_flagged(vi, vns[1], "del")
-            unset_flag!(vi, vns[1], "del")
+            # TODO(mhauru) Is it important to unset the flag here? The `true` allows us
+            # to ignore the fact that for VarNamedVector this does nothing, but I'm unsure if
+            # that's okay.
+            unset_flag!(vi, vns[1], "del", true)
             r = init(rng, dist, spl, n)
             for i in 1:n
                 vn = vns[i]
@@ -554,7 +560,10 @@ function get_and_set_val!(
     if haskey(vi, vns[1])
         # Always overwrite the parameters with new ones for `SampleFromUniform`.
         if spl isa SampleFromUniform || is_flagged(vi, vns[1], "del")
-            unset_flag!(vi, vns[1], "del")
+            # TODO(mhauru) Is it important to unset the flag here? The `true` allows us
+            # to ignore the fact that for VarNamedVector this does nothing, but I'm unsure if
+            # that's okay.
+            unset_flag!(vi, vns[1], "del", true)
             f = (vn, dist) -> init(rng, dist, spl)
             r = f.(vns, dists)
             for i in eachindex(vns)