Remove selector stuff from VarInfo tests and link/invlink

src/abstract_varinfo.jl

@@ -539,7 +539,9 @@
 
 """
     link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
-    link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::VarName, model::Model)
+    link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::Tuple{N,VarName}, model::Model)
+    link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::AbstractVector{<:VarName}, model::Model)
 
 Transform the variables in `vi` to their linked space, using the transformation `t`,
 mutating `vi` if possible.
@@ -548,37 +550,67 @@
 
 See also: [`default_transformation`](@ref), [`invlink!!`](@ref).
 """
-link!!(vi::AbstractVarInfo, model::Model) = link!!(vi, SampleFromPrior(), model)
+# Use `default_transformation` to decide which transformation to use if none is specified.
+function link!!(vi::AbstractVarInfo, model::Model)
+    return link!!(default_transformation(model, vi), vi, model)
+end
+function link!!(vi::AbstractVarInfo, vns, model::Model)
+    return link!!(default_transformation(model, vi), vi, vns, model)
+end
+# If no variable names are provided, link all variables.
 function link!!(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
-    return link!!(t, vi, SampleFromPrior(), model)
+    vns = collect(keys(vi))
+    # In case e.g. vns = Any[].
+    # TODO(mhauru) Could we rather fix akeys` so that it would always return VarName[]?
+    if !(eltype(vns) <: VarName)
+        vns = collect(VarName, vns)
+    end
+    return link!!(t, vi, vns, model)
 end
-function link!!(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
-    # Use `default_transformation` to decide which transformation to use if none is specified.
-    return link!!(default_transformation(model, vi), vi, spl, model)
+# Wrap a single VarName in a singleton tuple.
+function link!!(t::AbstractTransformation, vi::AbstractVarInfo, vn::VarName, model::Model)
+    return link!!(t, vi, (vn,), model)
 end
 
 """
     link([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
-    link([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    link([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::VarName, model::Model)
+    link([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::Tuple{N,VarName}, model::Model)
+    link([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::AbstractVector{<:VarName}, model::Model)
 
 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.
 
 See also: [`default_transformation`](@ref), [`invlink`](@ref).
 """
-link(vi::AbstractVarInfo, model::Model) = link(vi, SampleFromPrior(), model)
+# Use `default_transformation` to decide which transformation to use if none is specified.
+function link(vi::AbstractVarInfo, model::Model)
+    return link(default_transformation(model, vi), vi, model)
+end
+function link(vi::AbstractVarInfo, vns, model::Model)
+    return link(default_transformation(model, vi), vi, vns, model)
+end
+# If no variable names are provided, link all variables.
 function link(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
-    return link(t, deepcopy(vi), SampleFromPrior(), model)
+    vns = collect(keys(vi))
+    # In case e.g. vns = Any[].
+    # TODO(mhauru) Could we rather fix akeys` so that it would always return VarName[]?
+    if !(eltype(vns) <: VarName)
+        vns = collect(VarName, vns)
+    end
+    return link(t, vi, vns, model)
 end
-function link(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
-    # Use `default_transformation` to decide which transformation to use if none is specified.
-    return link(default_transformation(model, vi), deepcopy(vi), spl, model)
+# Wrap a single VarName in a singleton tuple.
+function link(t::AbstractTransformation, vi::AbstractVarInfo, vn::VarName, model::Model)
+    return link(t, vi, (vn,), model)
 end
 
 """
     invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
-    invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::VarName, model::Model)
+    invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::Tuple{N,VarName}, model::Model)
+    invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::AbstractVector{<:VarName}, model::Model)
 
 Transform the variables in `vi` to their constrained space, using the (inverse of)
 transformation `t`, mutating `vi` if possible.
@@ -587,49 +619,81 @@
 
 See also: [`default_transformation`](@ref), [`link!!`](@ref).
 """
-invlink!!(vi::AbstractVarInfo, model::Model) = invlink!!(vi, SampleFromPrior(), model)
+# Use `default_transformation` to decide which transformation to use if none is specified.
+function invlink!!(vi::AbstractVarInfo, model::Model)
+    return invlink!!(default_transformation(model, vi), vi, model)
+end
+function invlink!!(vi::AbstractVarInfo, vns, model::Model)
+    return invlink!!(default_transformation(model, vi), vi, vns, model)
+end
+# If no variable names are provided, invlink!! all variables.
 function invlink!!(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
-    return invlink!!(t, vi, SampleFromPrior(), model)
+    vns = collect(keys(vi))
+    # In case e.g. vns = Any[].
+    if !(eltype(vns) <: VarName)
+        vns = collect(VarName, vns)
+    end
+    return invlink!!(t, vi, vns, model)
 end
-function invlink!!(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
-    # Here we extract the `transformation` from `vi` rather than using the default one.
-    return invlink!!(transformation(vi), vi, spl, model)
+# Wrap a single VarName in a singleton tuple.
+function invlink!!(
+    t::AbstractTransformation, vi::AbstractVarInfo, vn::VarName, model::Model
+)
+    return invlink!!(t, vi, (vn,), model)
 end
 
 # Vector-based ones.
 function link!!(
     t::StaticTransformation{<:Bijectors.Transform},
     vi::AbstractVarInfo,
-    spl::AbstractSampler,
-    model::Model,
+    vns::VarNameCollection,
+    ::Model,
 )
+    # TODO(mhauru) The behavior of this before the removal of indexing with samplers was a
+    # bit mixed. For TypedVarInfo you could transform only a subset of the variables, but
+    # for UntypedVarInfo and SimpleVarInfo it was silently assumed that all variables were
+    # being set. Unsure if we should support this or not, but at least it now errors
+    # loudly.
+    all_vns = Set(keys(vi))
+    if Set(vns) != all_vns
+        msg = "StaticTransforming only a subset of variables is not supported."
+        throw(ArgumentError(msg))
+    end
     b = inverse(t.bijector)
-    x = vi[spl]
+    x = vi[:]
     y, logjac = with_logabsdet_jacobian(b, x)
 
     lp_new = getlogp(vi) - logjac
-    vi_new = setlogp!!(unflatten(vi, spl, y), lp_new)
+    vi_new = setlogp!!(unflatten(vi, y), lp_new)
     return settrans!!(vi_new, t)
 end
 
 function invlink!!(
     t::StaticTransformation{<:Bijectors.Transform},
     vi::AbstractVarInfo,
-    spl::AbstractSampler,
-    model::Model,
+    vns::VarNameCollection,
+    ::Model,
 )
+    # TODO(mhauru) See comment in link!! above.
+    all_vns = Set(keys(vi))
+    if Set(vns) != all_vns
+        msg = "StaticTransforming only a subset of variables is not supported."
+        throw(ArgumentError(msg))
+    end
     b = t.bijector
-    y = vi[spl]
+    y = vi[:]
     x, logjac = with_logabsdet_jacobian(b, y)
 
     lp_new = getlogp(vi) + logjac
-    vi_new = setlogp!!(unflatten(vi, spl, x), lp_new)
+    vi_new = setlogp!!(unflatten(vi, x), lp_new)
     return settrans!!(vi_new, NoTransformation())
 end
 
 """
     invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
-    invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::VarName, model::Model)
+    invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::Tuple{N,VarName}, model::Model)
+    invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, vn::AbstractVector{<:VarName}, model::Model)
 
 Transform the variables in `vi` to their constrained space without mutating `vi`, using the (inverse of)
 transformation `t`.
@@ -638,12 +702,25 @@
 
 See also: [`default_transformation`](@ref), [`link`](@ref).
 """
-invlink(vi::AbstractVarInfo, model::Model) = invlink(vi, SampleFromPrior(), model)
+# Use `default_transformation` to decide which transformation to use if none is specified.
+function invlink(vi::AbstractVarInfo, model::Model)
+    return invlink(default_transformation(model, vi), vi, model)
+end
+function invlink(vi::AbstractVarInfo, vns, model::Model)
+    return invlink(default_transformation(model, vi), vi, vns, model)
+end
+# If no variable names are provided, invlink all variables.
 function invlink(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
-    return invlink(t, vi, SampleFromPrior(), model)
+    vns = collect(keys(vi))
+    # In case e.g. vns = Any[].
+    if !(eltype(vns) <: VarName)
+        vns = collect(VarName, vns)
+    end
+    return invlink(t, vi, vns, model)
 end
-function invlink(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
-    return invlink(transformation(vi), vi, spl, model)
+# Wrap a single VarName in a singleton tuple.
+function invlink(t::AbstractTransformation, vi::AbstractVarInfo, vn::VarName, model::Model)
+    return invlink(t, vi, (vn,), model)
 end
 
 """
@@ -715,9 +792,10 @@
     return vi
 end
 function maybe_invlink_before_eval!!(
-    t::StaticTransformation, vi::AbstractVarInfo, context::AbstractContext, model::Model
+    t::StaticTransformation, vi::AbstractVarInfo, ::AbstractContext, model::Model
 )
-    return invlink!!(t, vi, _default_sampler(context), model)
+    # TODO(mhauru) Why does this function need the context argument?
+    return invlink!!(t, vi, model)
 end
 
 function _default_sampler(context::AbstractContext)

src/simple_varinfo.jl

@@ -680,7 +680,7 @@ Distributions.loglikelihood(model::Model, θ) = loglikelihood(model, SimpleVarIn
 function link!!(
     t::StaticTransformation{<:Bijectors.NamedTransform},
     vi::SimpleVarInfo{<:NamedTuple},
-    spl::AbstractSampler,
+    ::VarNameCollection,
     model::Model,
 )
     # TODO: Make sure that `spl` is respected.
@@ -695,7 +695,7 @@ end
 function invlink!!(
     t::StaticTransformation{<:Bijectors.NamedTransform},
     vi::SimpleVarInfo{<:NamedTuple},
-    spl::AbstractSampler,
+    ::VarNameCollection,
     model::Model,
 )
     # TODO: Make sure that `spl` is respected.

src/threadsafe.jl

@@ -82,41 +82,41 @@
 islinked(vi::ThreadSafeVarInfo, spl::AbstractSampler) = islinked(vi.varinfo, spl)
 
 function link!!(
-    t::AbstractTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
+    t::AbstractTransformation, vi::ThreadSafeVarInfo, vns::VarNameCollection, model::Model
 )
-    return Accessors.@set vi.varinfo = link!!(t, vi.varinfo, spl, model)
+    return Accessors.@set vi.varinfo = link!!(t, vi.varinfo, vns, model)
 end
 
 function invlink!!(
-    t::AbstractTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
+    t::AbstractTransformation, vi::ThreadSafeVarInfo, vns::VarNameCollection, model::Model
 )
-    return Accessors.@set vi.varinfo = invlink!!(t, vi.varinfo, spl, model)
+    return Accessors.@set vi.varinfo = invlink!!(t, vi.varinfo, vns, model)
 end
 
 function link(
-    t::AbstractTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
+    t::AbstractTransformation, vi::ThreadSafeVarInfo, vns::VarNameCollection, model::Model
 )
-    return Accessors.@set vi.varinfo = link(t, vi.varinfo, spl, model)
+    return Accessors.@set vi.varinfo = link(t, vi.varinfo, vns, model)
 end
 
 function invlink(
-    t::AbstractTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
+    t::AbstractTransformation, vi::ThreadSafeVarInfo, vns::VarNameCollection, model::Model
 )
-    return Accessors.@set vi.varinfo = invlink(t, vi.varinfo, spl, model)
+    return Accessors.@set vi.varinfo = invlink(t, vi.varinfo, vns, model)
 end
 
 # Need to define explicitly for `DynamicTransformation` to avoid method ambiguity.
 # NOTE: We also can't just defer to the wrapped varinfo, because we need to ensure
 # consistency between `vi.logps` field and `getlogp(vi.varinfo)`, which accumulates
 # to define `getlogp(vi)`.
 function link!!(
-    t::DynamicTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
+    t::DynamicTransformation, vi::ThreadSafeVarInfo, ::VarNameCollection, model::Model
 )
     return settrans!!(last(evaluate!!(model, vi, DynamicTransformationContext{false}())), t)
 end
 
 function invlink!!(
-    ::DynamicTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
+    ::DynamicTransformation, vi::ThreadSafeVarInfo, ::VarNameCollection, model::Model
 )
     return settrans!!(
         last(evaluate!!(model, vi, DynamicTransformationContext{true}())),
@@ -125,15 +125,15 @@
 end
 
 function link(
-    t::DynamicTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
+    t::DynamicTransformation, vi::ThreadSafeVarInfo, vns::VarNameCollection, model::Model
 )
-    return link!!(t, deepcopy(vi), spl, model)
+    return link!!(t, deepcopy(vi), vns, model)
 end
 
 function invlink(
-    t::DynamicTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
+    t::DynamicTransformation, vi::ThreadSafeVarInfo, vns::VarNameCollection, model::Model
 )
-    return invlink!!(t, deepcopy(vi), spl, model)
+    return invlink!!(t, deepcopy(vi), vns, model)
 end
 
 function maybe_invlink_before_eval!!(

src/transforming.jl

@@ -91,14 +91,16 @@ function dot_tilde_assume(
     return r, lp, vi
 end
 
+VarNameCollection = Union{NTuple{N,VarName} where N,AbstractVector{<:VarName},NamedTuple}
+
 function link!!(
-    t::DynamicTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
+    t::DynamicTransformation, vi::AbstractVarInfo, ::VarNameCollection, model::Model
 )
     return settrans!!(last(evaluate!!(model, vi, DynamicTransformationContext{false}())), t)
 end
 
 function invlink!!(
-    ::DynamicTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
+    ::DynamicTransformation, vi::AbstractVarInfo, ::VarNameCollection, model::Model
 )
     return settrans!!(
         last(evaluate!!(model, vi, DynamicTransformationContext{true}())),
@@ -107,13 +109,13 @@ function invlink!!(
 end
 
 function link(
-    t::DynamicTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
+    t::DynamicTransformation, vi::AbstractVarInfo, vns::VarNameCollection, model::Model
 )
-    return link!!(t, deepcopy(vi), spl, model)
+    return link!!(t, deepcopy(vi), vns, model)
 end
 
 function invlink(
-    t::DynamicTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
+    t::DynamicTransformation, vi::AbstractVarInfo, vns::VarNameCollection, model::Model
 )
-    return invlink!!(t, deepcopy(vi), spl, model)
+    return invlink!!(t, deepcopy(vi), vns, model)
 end