subset and merge for VarInfo

[torfjelde]

To implement a Gibbs sampler based on condition a la TuringLang/Turing.jl#2099 we need two missing operations for VarInfo:

• subset(varinfo, varnames): creates a VarInfo from varinfo based only on varnames.
• merge(varinfo_left, varinfo_right): attempts to merge two instances of VarInfo into a single VarInfo. These can include both shared and not-shared variables.

Here's an example:

julia> using Distributions, DynamicPPL

julia> @model function demo()
           s ~ InverseGamma(2, 3)
           m ~ Normal(0, sqrt(s))
           x = Vector{Float64}(undef, 2)
           x[1] ~ Normal(m, sqrt(s))
           x[2] ~ Normal(m, sqrt(s))
       end

demo (generic function with 2 methods)

julia> model = demo();

julia> varinfo = VarInfo(model);

julia> keys(varinfo)
4-element Vector{VarName}:
 s
 m
 x[1]
 x[2]

julia> varinfo[@varname(s)] = 1;

julia> varinfo[@varname(m)] = 2;

julia> varinfo[@varname(x[1])] = 3;

julia> varinfo[@varname(x[2])] = 4;

julia> varinfo[[@varname(s), @varname(m), @varname(x[1]), @varname(x[2])]]
4-element Vector{Float64}:
 1.0
 2.0
 3.0
 4.0

julia> # Extract one with only `m`.
       varinfo_subset1 = subset(varinfo, [@varname(m),]);


julia> keys(varinfo_subset1)
1-element Vector{VarName{:m, Setfield.IdentityLens}}:
 m

julia> varinfo_subset1[@varname(m)]
2.0

julia> # Extract one with both `s` and `x[2]`.
       varinfo_subset2 = subset(varinfo, [@varname(s), @varname(x[2])]);

julia> keys(varinfo_subset2)
2-element Vector{VarName}:
 s
 x[2]

julia> varinfo_subset2[[@varname(s), @varname(x[2])]]
2-element Vector{Float64}:
 1.0
 4.0

julia> # Merge the two.
       varinfo_subset_merged = merge(varinfo_subset1, varinfo_subset2);

julia> keys(varinfo_subset_merged)
3-element Vector{VarName}:
 m
 s
 x[2]

julia> varinfo_subset_merged[[@varname(s), @varname(m), @varname(x[2])]]
3-element Vector{Float64}:
 1.0
 2.0
 4.0

julia> # Merge the two with the original.
       varinfo_merged = merge(varinfo, varinfo_subset_merged);

julia> keys(varinfo_merged)
4-element Vector{VarName}:
 s
 m
 x[1]
 x[2]

julia> varinfo_merged[[@varname(s), @varname(m), @varname(x[1]), @varname(x[2])]]
4-element Vector{Float64}:
 1.0
 2.0
 3.0
 4.0

Note that none of these are, at the moment, implemented for SimpleVarInfo. This is both because (1) we don't need it yet, and (2) this will support a much smaller set of combinations.

[torfjelde]

Note that this is based on #542 (because I wanted to test the Gibbs PR with this)

[github-actions]

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

	push!(
	vals.args,
	:(merge_metadata(metadata_left.$sym, metadata_right.$sym))
	)
	push!(vals.args, :(merge_metadata(metadata_left.$sym, metadata_right.$sym)))

[torfjelde]

Replaced by #544 to avoid dependence on #542