subset and merge for VarInfo (clean version)

docs/src/api.md

@@ -255,6 +255,8 @@ DynamicPPL.reconstruct
 #### Utils
 
 ```@docs
+Base.merge(::VarInfo, ::VarInfo)
+DynamicPPL.subset
 DynamicPPL.unflatten
 DynamicPPL.tonamedtuple
 DynamicPPL.varname_leaves

src/DynamicPPL.jl

@@ -47,6 +47,7 @@ export AbstractVarInfo,
     SimpleVarInfo,
     push!!,
     empty!!,
+    subset,
     getlogp,
     setlogp!!,
     acclogp!!,

src/abstract_varinfo.jl

@@ -53,6 +53,27 @@ struct StaticTransformation{F} <: AbstractTransformation
     bijector::F
 end
 
+"""
+    merge_transformations(transformation_left, transformation_right)
+
+Merge two transformations.
+
+The main use of this is in [`merge(::AbstractVarInfo, ::AbstractVarInfo)`](@ref).
+"""
+function merge_transformations(::NoTransformation, ::NoTransformation)
+    return NoTransformation()
+end
+function merge_transformations(::DynamicTransformation, ::DynamicTransformation)
+    return DynamicTransformation()
+end
+function merge_transformations(left::StaticTransformation, right::StaticTransformation)
+    return StaticTransformation(merge_bijectors(left.bijector, right.bijector))
+end
+
+function merge_bijectors(left::Bijectors.NamedTransform, right::Bijectors.NamedTransform)
+    return Bijectors.NamedTransform(merge_bijector(left.bs, right.bs))
+end
+
 """
     default_transformation(model::Model[, vi::AbstractVarInfo])
 
@@ -337,6 +358,144 @@ function Base.eltype(vi::AbstractVarInfo, spl::Union{AbstractSampler,SampleFromP
     return eltype(Core.Compiler.return_type(getindex, Tuple{typeof(vi),typeof(spl)}))
 end
 
+# 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}`.
+"""
+    subset(varinfo::AbstractVarInfo, vns::AbstractVector{<:VarName})
+
+Subset a `varinfo` to only contain the variables `vns`.
+
+!!! warning
+    The ordering of the variables in the resulting `varinfo` is _not_
+    guaranteed to follow the ordering of the variables in `varinfo`.
+    Hence care must be taken, in particular when used in conjunction with
+    other methods which uses the vector-representation of the `varinfo`,
+    e.g. `getindex(varinfo, sampler)`.
+
+# Examples
+```jldoctest varinfo-subset; setup = :(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> for (i, vn) in enumerate(keys(varinfo))
+           varinfo[vn] = i
+       end
+
+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
+```
+
+`subset` is particularly useful when combined with [`merge(varinfo_left::VarInfo, varinfo_right::VarInfo)`](@ref)
+
+```jldoctest varinfo-subset
+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
+```
+
+# Notes
+
+## Type-stability
+
+!!! warning
+    This function is only type-stable when `vns` contains only varnames
+    with the same symbol. For exmaple, `[@varname(m[1]), @varname(m[2])]` will
+    be type-stable, but `[@varname(m[1]), @varname(x)]` will not be.
+"""
+function subset end
+
+"""
+    merge(varinfo, other_varinfos...)
+
+Merge varinfos into one, giving precedence to the right-most varinfo when sensible.
+
+This is particularly useful when combined with [`subset(varinfo, vns)`](@ref).
+
+See docstring of [`subset(varinfo, vns)`](@ref) for examples.
+"""
+function Base.merge(varinfo::AbstractVarInfo, varinfo_others::AbstractVarInfo...)
+    return merge(Base.merge(varinfo, first(varinfo_others)), Base.tail(varinfo_others)...)
+end
+
+# Avoid `StackoverFlowError` if implementation is missing.
+function Base.merge(varinfo::AbstractVarInfo, varinfo_other::AbstractVarInfo)
+    throw(MethodError(Base.merge, (varinfo, varinfo_other)))
+end
+
 # Transformations
 """
     istrans(vi::AbstractVarInfo[, vns::Union{VarName, AbstractVector{<:Varname}}])

src/simple_varinfo.jl

@@ -419,6 +419,48 @@ function Base.eltype(
     return V
 end
 
+# `subset`
+function subset(varinfo::SimpleVarInfo, vns::AbstractVector{<:VarName})
+    return Setfield.@set varinfo.values = _subset(varinfo.values, vns)
+end
+
+function _subset(x::AbstractDict, vns)
+    # NOTE: This requires `vns` to be explicitly present in `x`.
+    if any(!Base.Fix1(haskey, x), vns)
+        error(
+            "Cannot subset `AbstractDict` with `VarName` that is not an explicit key. " *
+            "For example, if `keys(x) == [@varname(x[1])]`, then subsetting with " *
+            "`@varname(x[1])` is allowed, but subsetting with `@varname(x)` is not."
+        )
+    end
+    C = ConstructionBase.constructorof(typeof(x))
+    return C(vn => x[vn] for vn in vns)
+end
+
+function _subset(x::NamedTuple, vns)
+    # NOTE: Here we can only handle `vns` that contain the `IdentityLens`.
+    if any(!==(Setfield.IdentityLens()) ∘ getlens, vns)
+        error(
+            "Cannot subset `NamedTuple` with non-`IdentityLens` `VarName`. " *
+            "For example, `@varname(x)` is allowed, but `@varname(x[1])` is not."
+        )
+    end
+
+    syms = map(getsym, vns)
+    return NamedTuple{(syms...,)}((map(Base.Fix2(getindex, x), syms)...,))
+end
+
+# `merge`
+function Base.merge(varinfo_left::SimpleVarInfo, varinfo_right::SimpleVarInfo)
+    values = merge(varinfo_left.values, varinfo_right.values)
+    logp = getlogp(varinfo_right)
+    transformation = merge_transformations(
+        varinfo_left.transformation,
+        varinfo_right.transformation,
+    )
+    return SimpleVarInfo(values, logp, transformation)
+end
+
 # Context implementations
 # NOTE: Evaluations, i.e. those without `rng` are shared with other
 # implementations of `AbstractVarInfo`.