Revert "temporarily removed VarNameVector completely"

This reverts commit 95dc8e3.

Author: torfjelde

docs/make.jl

@@ -20,7 +20,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,

docs/src/internals/varinfo.md

@@ -44,6 +44,7 @@ export AbstractVarInfo,
     UntypedVarInfo,
     TypedVarInfo,
     SimpleVarInfo,
+    VarNameVector,
     push!!,
     empty!!,
     subset,
@@ -158,6 +159,7 @@ include("sampler.jl")
 include("varname.jl")
 include("distribution_wrappers.jl")
 include("contexts.jl")
+include("varnamevector.jl")
 include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")

src/DynamicPPL.jl

@@ -55,6 +55,8 @@ function setlogp!!(vi::ThreadSafeVarInfoWithRef, logp)
     return ThreadSafeVarInfo(setlogp!!(vi.varinfo, logp), vi.logps)
 end
 
+has_varnamevector(vi::DynamicPPL.ThreadSafeVarInfo) = has_varnamevector(vi.varinfo)
+
 function BangBang.push!!(
     vi::ThreadSafeVarInfo, vn::VarName, r, dist::Distribution, gidset::Set{Selector}
 )

src/threadsafe.jl

@@ -101,6 +101,7 @@ struct VarInfo{Tmeta,Tlogp} <: AbstractVarInfo
     logp::Base.RefValue{Tlogp}
     num_produce::Base.RefValue{Int}
 end
+const VectorVarInfo = VarInfo{<:VarNameVector}
 const UntypedVarInfo = VarInfo{<:Metadata}
 const TypedVarInfo = VarInfo{<:NamedTuple}
 const VarInfoOrThreadSafeVarInfo{Tmeta} = Union{
@@ -125,6 +126,46 @@ function VarInfo(old_vi::TypedVarInfo, spl, x::AbstractVector)
     )
 end
 
+# No-op if we're already working with a `VarNameVector`.
+metadata_to_varnamevector(vnv::VarNameVector) = vnv
+function metadata_to_varnamevector(md::Metadata)
+    idcs = copy(md.idcs)
+    vns = copy(md.vns)
+    ranges = copy(md.ranges)
+    vals = copy(md.vals)
+    transforms = map(md.dists) do dist
+        # TODO: Handle linked distributions.
+        from_vec_transform(dist)
+    end
+
+    return VarNameVector(
+        OrderedDict{eltype(keys(idcs)),Int}(idcs), vns, ranges, vals, transforms
+    )
+end
+
+function VectorVarInfo(vi::UntypedVarInfo)
+    md = metadata_to_varnamevector(vi.metadata)
+    lp = getlogp(vi)
+    return VarInfo(md, Base.RefValue{eltype(lp)}(lp), Ref(get_num_produce(vi)))
+end
+
+function VectorVarInfo(vi::TypedVarInfo)
+    md = map(metadata_to_varnamevector, vi.metadata)
+    lp = getlogp(vi)
+    return VarInfo(md, Base.RefValue{eltype(lp)}(lp), Ref(get_num_produce(vi)))
+end
+
+"""
+    has_varnamevector(varinfo::VarInfo)
+
+Returns `true` if `varinfo` uses `VarNameVector` as metadata.
+"""
+has_varnamevector(vi) = false
+function has_varnamevector(vi::VarInfo)
+    return vi.metadata isa VarNameVector ||
+           (vi isa TypedVarInfo && any(Base.Fix2(isa, VarNameVector), values(vi.metadata)))
+end
+
 function untyped_varinfo(
     rng::Random.AbstractRNG,
     model::Model,
@@ -321,6 +362,10 @@ function _merge(varinfo_left::VarInfo, varinfo_right::VarInfo)
     )
 end
 
+function merge_metadata(vnv_left::VarNameVector, vnv_right::VarNameVector)
+    return merge(vnv_left, vnv_right)
+end
+
 @generated function merge_metadata(
     metadata_left::NamedTuple{names_left}, metadata_right::NamedTuple{names_right}
 ) where {names_left,names_right}
@@ -513,9 +558,13 @@ Return the distribution from which `vn` was sampled in `vi`.
 """
 getdist(vi::VarInfo, vn::VarName) = getdist(getmetadata(vi, vn), vn)
 getdist(md::Metadata, vn::VarName) = md.dists[getidx(md, vn)]
+# HACK: we shouldn't need this
+getdist(::VarNameVector, ::VarName) = nothing
 
 getindex_internal(vi::VarInfo, vn::VarName) = getindex_internal(getmetadata(vi, vn), vn)
 getindex_internal(md::Metadata, vn::VarName) = view(md.vals, getrange(md, vn))
+# HACK: We shouldn't need this
+getindex_internal(vnv::VarNameVector, vn::VarName) = view(vnv.vals, getrange(vnv, vn))
 
 function getindex_internal(vi::VarInfo, vns::Vector{<:VarName})
     return mapreduce(Base.Fix1(getindex_internal, vi), vcat, vns)
@@ -535,6 +584,9 @@ end
 function setval!(md::Metadata, val, vn::VarName)
     return md.vals[getrange(md, vn)] = vectorize(getdist(md, vn), val)
 end
+function setval!(vnv::VarNameVector, val, vn::VarName)
+    return setindex_raw!(vnv, tovec(val), vn)
+end
 
 """
     getall(vi::VarInfo)
@@ -552,6 +604,7 @@ function getall(md::Metadata)
         Base.Fix1(getindex_internal, md), vcat, md.vns; init=similar(md.vals, 0)
     )
 end
+getall(vnv::VarNameVector) = vnv.vals
 
 """
     setall!(vi::VarInfo, val)
@@ -567,6 +620,12 @@ function _setall!(metadata::Metadata, val)
         metadata.vals[r] .= val[r]
     end
 end
+function _setall!(vnv::VarNameVector, val)
+    # TODO: Do something more efficient here.
+    for i in 1:length(vnv)
+        vnv[i] = val[i]
+    end
+end
 @generated function _setall!(metadata::NamedTuple{names}, val) where {names}
     expr = Expr(:block)
     start = :(1)
@@ -599,6 +658,10 @@ function settrans!!(metadata::Metadata, trans::Bool, vn::VarName)
 
     return metadata
 end
+function settrans!!(vnv::VarNameVector, trans::Bool, vn::VarName)
+    settrans!(vnv, trans, vn)
+    return vnv
+end
 
 function settrans!!(vi::VarInfo, trans::Bool)
     for vn in keys(vi)
@@ -940,6 +1003,8 @@ end
 
 # X -> R for all variables associated with given sampler
 function link!!(t::DynamicTransformation, vi::VarInfo, spl::AbstractSampler, model::Model)
+    # If we're working with a `VarNameVector`, we always use immutable.
+    has_varnamevector(vi) && return link(t, vi, spl, model)
     # Call `_link!` instead of `link!` to avoid deprecation warning.
     _link!(vi, spl)
     return vi
@@ -1035,6 +1100,8 @@ end
 function invlink!!(
     t::DynamicTransformation, vi::VarInfo, spl::AbstractSampler, model::Model
 )
+    # If we're working with a `VarNameVector`, we always use immutable.
+    has_varnamevector(vi) && return invlink(t, vi, spl, model)
     # Call `_invlink!` instead of `invlink!` to avoid deprecation warning.
     _invlink!(vi, spl)
     return vi
@@ -1260,6 +1327,66 @@ function _link_metadata!(model::Model, varinfo::VarInfo, metadata::Metadata, tar
     )
 end
 
+function _link_metadata!(
+    model::Model, varinfo::VarInfo, metadata::VarNameVector, target_vns
+)
+    # HACK: We ignore `target_vns` here.
+    vns = keys(metadata)
+    # Need to extract the priors from the model.
+    dists = extract_priors(model, varinfo)
+
+    is_transformed = copy(metadata.is_transformed)
+
+    # Construct the linking transformations.
+    link_transforms = map(vns) do vn
+        # If `vn` is not part of `target_vns`, the `identity` transformation is used.
+        if (target_vns !== nothing && vn ∉ target_vns)
+            return identity
+        end
+
+        # Otherwise, we derive the transformation from the distribution.
+        is_transformed[getidx(metadata, vn)] = true
+        internal_to_linked_internal_transform(varinfo, vn, dists[vn])
+    end
+    # Compute the transformed values.
+    ys = map(vns, link_transforms) do vn, f
+        # TODO: Do we need to handle scenarios where `vn` is not in `dists`?
+        dist = dists[vn]
+        x = getindex_internal(metadata, vn)
+        y, logjac = with_logabsdet_jacobian(f, x)
+        # Accumulate the log-abs-det jacobian correction.
+        acclogp!!(varinfo, -logjac)
+        # Return the transformed value.
+        return y
+    end
+    # Extract the from-vec transformations.
+    fromvec_transforms = map(from_vec_transform, ys)
+    # Compose the transformations to form a full transformation from
+    # unconstrained vector representation to constrained space.
+    transforms = map(∘, map(inverse, link_transforms), fromvec_transforms)
+    # Convert to vector representation.
+    yvecs = map(tovec, ys)
+
+    # Determine new ranges.
+    ranges_new = similar(metadata.ranges)
+    offset = 0
+    for (i, v) in enumerate(yvecs)
+        r_start, r_end = offset + 1, length(v) + offset
+        offset = r_end
+        ranges_new[i] = r_start:r_end
+    end
+
+    # Now we just create a new metadata with the new `vals` and `ranges`.
+    return VarNameVector(
+        metadata.varname_to_index,
+        metadata.varnames,
+        ranges_new,
+        reduce(vcat, yvecs),
+        transforms,
+        is_transformed,
+    )
+end
+
 function invlink(
     ::DynamicTransformation, varinfo::VarInfo, spl::AbstractSampler, model::Model
 )
@@ -1360,6 +1487,55 @@ function _invlink_metadata!(::Model, varinfo::VarInfo, metadata::Metadata, targe
     )
 end
 
+function _invlink_metadata!(
+    model::Model, varinfo::VarInfo, metadata::VarNameVector, target_vns
+)
+    # HACK: We ignore `target_vns` here.
+    # TODO: Make use of `update!` to aovid copying values.
+    #       => Only need to allocate for transformations.
+
+    vns = keys(metadata)
+    is_transformed = copy(metadata.is_transformed)
+
+    # Compute the transformed values.
+    xs = map(vns) do vn
+        f = gettransform(metadata, vn)
+        y = getindex_internal(metadata, vn)
+        # No need to use `with_reconstruct` as `f` will include this.
+        x, logjac = with_logabsdet_jacobian(f, y)
+        # Accumulate the log-abs-det jacobian correction.
+        acclogp!!(varinfo, -logjac)
+        # Mark as no longer transformed.
+        is_transformed[getidx(metadata, vn)] = false
+        # Return the transformed value.
+        return x
+    end
+    # Compose the transformations to form a full transformation from
+    # unconstrained vector representation to constrained space.
+    transforms = map(from_vec_transform, xs)
+    # Convert to vector representation.
+    xvecs = map(tovec, xs)
+
+    # Determine new ranges.
+    ranges_new = similar(metadata.ranges)
+    offset = 0
+    for (i, v) in enumerate(xvecs)
+        r_start, r_end = offset + 1, length(v) + offset
+        offset = r_end
+        ranges_new[i] = r_start:r_end
+    end
+
+    # Now we just create a new metadata with the new `vals` and `ranges`.
+    return VarNameVector(
+        metadata.varname_to_index,
+        metadata.varnames,
+        ranges_new,
+        reduce(vcat, xvecs),
+        transforms,
+        is_transformed,
+    )
+end
+
 """
     islinked(vi::VarInfo, spl::Union{Sampler, SampleFromPrior})
 
@@ -1429,6 +1605,14 @@ function getindex(vi::VarInfo, vn::VarName, dist::Distribution)
     val = getindex_internal(vi, vn)
     return from_maybe_linked_internal(vi, vn, dist, val)
 end
+# HACK: Allows us to also work with `VarNameVector` where `dist` is not used,
+# but we instead use a transformation stored with the variable.
+function getindex(vi::VarInfo, vn::VarName, ::Nothing)
+    if !haskey(vi, vn)
+        throw(KeyError(vn))
+    end
+    return getmetadata(vi, vn)[vn]
+end
 
 function getindex(vi::VarInfo, vns::Vector{<:VarName})
     vals_linked = mapreduce(vcat, vns) do vn
@@ -1621,6 +1805,11 @@ function Base.push!(meta::Metadata, vn, r, dist, gidset, num_produce)
     return meta
 end
 
+function Base.push!(vnv::VarNameVector, vn, r, dist, gidset, num_produce)
+    f = from_vec_transform(dist)
+    return push!(vnv, vn, r, f)
+end
+
 """
     setorder!(vi::VarInfo, vn::VarName, index::Int)
 
@@ -1635,6 +1824,7 @@ function setorder!(metadata::Metadata, vn::VarName, index::Int)
     metadata.orders[metadata.idcs[vn]] = index
     return metadata
 end
+setorder!(vnv::VarNameVector, ::VarName, ::Int) = vnv
 
 """
     getorder(vi::VarInfo, vn::VarName)
@@ -1660,6 +1850,8 @@ end
 function is_flagged(metadata::Metadata, vn::VarName, flag::String)
     return metadata.flags[flag][getidx(metadata, vn)]
 end
+# HACK: This is bad. Should we always return `true` here?
+is_flagged(::VarNameVector, ::VarName, flag::String) = flag == "del" ? true : false
 
 """
     unset_flag!(vi::VarInfo, vn::VarName, flag::String)
@@ -1674,6 +1866,7 @@ function unset_flag!(metadata::Metadata, vn::VarName, flag::String)
     metadata.flags[flag][getidx(metadata, vn)] = false
     return metadata
 end
+unset_flag!(vnv::VarNameVector, ::VarName, ::String) = vnv
 
 """
     set_retained_vns_del_by_spl!(vi::VarInfo, spl::Sampler)
@@ -2027,6 +2220,8 @@ function values_from_metadata(md::Metadata)
     )
 end
 
+values_from_metadata(md::VarNameVector) = pairs(md)
+
 # Transforming from internal representation to distribution representation.
 # Without `dist` argument: base on `dist` extracted from self.
 function from_internal_transform(vi::VarInfo, vn::VarName)
@@ -2035,11 +2230,17 @@ end
 function from_internal_transform(md::Metadata, vn::VarName)
     return from_internal_transform(md, vn, getdist(md, vn))
 end
+function from_internal_transform(md::VarNameVector, vn::VarName)
+    return gettransform(md, vn)
+end
 # With both `vn` and `dist` arguments: base on provided `dist`.
 function from_internal_transform(vi::VarInfo, vn::VarName, dist)
     return from_internal_transform(getmetadata(vi, vn), vn, dist)
 end
 from_internal_transform(::Metadata, ::VarName, dist) = from_vec_transform(dist)
+function from_internal_transform(::VarNameVector, ::VarName, dist)
+    return from_vec_transform(dist)
+end
 
 # Without `dist` argument: base on `dist` extracted from self.
 function from_linked_internal_transform(vi::VarInfo, vn::VarName)
@@ -2048,6 +2249,9 @@ end
 function from_linked_internal_transform(md::Metadata, vn::VarName)
     return from_linked_internal_transform(md, vn, getdist(md, vn))
 end
+function from_linked_internal_transform(md::VarNameVector, vn::VarName)
+    return gettransform(md, vn)
+end
 # With both `vn` and `dist` arguments: base on provided `dist`.
 function from_linked_internal_transform(vi::VarInfo, vn::VarName, dist)
     # Dispatch to metadata in case this alters the behavior.
@@ -2056,3 +2260,6 @@ end
 function from_linked_internal_transform(::Metadata, ::VarName, dist)
     return from_linked_vec_transform(dist)
 end
+function from_linked_internal_transform(::VarNameVector, ::VarName, dist)
+    return from_linked_vec_transform(dist)
+end