More work on VarNameVector

docs/src/internals/varinfo.md

@@ -229,7 +229,7 @@ This does incur a runtime cost as it requires re-allocation of the `ranges` in a
 
 !!! note
 
-    Higher-dimensional arrays, e.g. `Matrix`, are handled by simply vectorizing them before storing them in the `Vector{T}`, and composing he `VarName`'s transformation with a `DynamicPPL.FromVec`.
+    Higher-dimensional arrays, e.g. `Matrix`, are handled by simply vectorizing them before storing them in the `Vector{T}`, and composing the `VarName`'s transformation with a `DynamicPPL.ReshapeTransform`.
 
 Continuing from the example from the previous section, we can use a `VarInfo` with a `VarNameVector` as the `metadata` field:
 

ext/DynamicPPLMCMCChainsExt.jl

@@ -11,7 +11,7 @@ end
 _has_varname_to_symbol(info::NamedTuple{names}) where {names} = :varname_to_symbol in names
 function _check_varname_indexing(c::MCMCChains.Chains)
     return DynamicPPL.supports_varname_indexing(c) ||
-           error("Chains do not support indexing using $vn.")
+           error("$(typeof(c)) do not support indexing using varnmes.")
 end
 
 # Load state from a `Chains`: By convention, it is stored in `:samplerstate` metadata
@@ -41,21 +41,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

ext/DynamicPPLReverseDiffExt.jl

@@ -9,12 +9,12 @@ else
 end
 
 function LogDensityProblemsAD.ADgradient(
-    ad::ADTypes.AutoReverseDiff{Tcompile}, ℓ::DynamicPPL.LogDensityFunction
-) where {Tcompile}
+    ad::ADTypes.AutoReverseDiff, ℓ::DynamicPPL.LogDensityFunction
+)
     return LogDensityProblemsAD.ADgradient(
         Val(:ReverseDiff),
         ℓ;
-        compile=Val(Tcompile),
+        compile=Val(ad.compile),
         # `getparams` can return `Vector{Real}`, in which case, `ReverseDiff` will initialize the gradients to Integer 0
         # because at https://github.com/JuliaDiff/ReverseDiff.jl/blob/c982cde5494fc166965a9d04691f390d9e3073fd/src/tracked.jl#L473
         # `zero(D)` will return 0 when D is Real.

src/DynamicPPL.jl

@@ -45,8 +45,9 @@ export AbstractVarInfo,
     VarInfo,
     UntypedVarInfo,
     TypedVarInfo,
+    VectorVarInfo,
     SimpleVarInfo,
-    VarNameVector,
+    VarNamedVector,
     push!!,
     empty!!,
     subset,
@@ -176,7 +177,7 @@ include("sampler.jl")
 include("varname.jl")
 include("distribution_wrappers.jl")
 include("contexts.jl")
-include("varnamevector.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})

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)

src/model.jl

@@ -1201,74 +1201,6 @@ function Distributions.loglikelihood(model::Model, chain::AbstractMCMC.AbstractC
     end
 end
 
-"""
-    generated_quantities(model::Model, chain::AbstractChains)
-
-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 generated_quantities(model::Model, chain::AbstractChains)
-    varinfo = VarInfo(model)
-    iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
-    return map(iters) do (sample_idx, chain_idx)
-        setval_and_resample!(varinfo, chain, sample_idx, chain_idx)
-        model(varinfo)
-    end
-end
-
 """
     generated_quantities(model::Model, parameters::NamedTuple)
     generated_quantities(model::Model, values, keys)
@@ -1295,7 +1227,7 @@ demo (generic function with 2 methods)
 
 julia> model = demo(randn(10));
 
-julia> parameters = (; s = 1.0, m_shifted=10);
+julia> parameters = (; s = 1.0, m_shifted=10.0);
 
 julia> generated_quantities(model, parameters)
 (0.0,)
@@ -1305,13 +1237,10 @@ julia> generated_quantities(model, values(parameters), keys(parameters))
 ```
 """
 function generated_quantities(model::Model, parameters::NamedTuple)
-    varinfo = VarInfo(model)
-    setval_and_resample!(varinfo, values(parameters), keys(parameters))
-    return model(varinfo)
+    fixed_model = fix(model, parameters)
+    return fixed_model()
 end
 
 function generated_quantities(model::Model, values, keys)
-    varinfo = VarInfo(model)
-    setval_and_resample!(varinfo, values, keys)
-    return model(varinfo)
+    return generated_quantities(model, NamedTuple{keys}(values))
 end

src/sampler.jl

@@ -150,7 +150,7 @@ function set_values!!(
     flattened_param_vals = varinfo[spl]
     length(flattened_param_vals) == length(initial_params) || throw(
         DimensionMismatch(
-            "Provided initial value size ($(length(initial_params))) doesn't match the model size ($(length(theta)))",
+            "Provided initial value size ($(length(initial_params))) doesn't match the model size ($(length(flattened_param_vals)))",
         ),
     )