Immutable versions of link and invlink

src/abstract_varinfo.jl

@@ -368,7 +368,8 @@
     link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
     link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
 
-Transforms the variables in `vi` to their linked space, using the transformation `t`.
+Transform the variables in `vi` to their linked space, using the transformation `t`,
+mutating `vi` if possible.
 
 If `t` is not provided, `default_transformation(model, vi)` will be used.
 
@@ -383,12 +384,31 @@
     return link!!(default_transformation(model, vi), vi, spl, model)
 end
 
+"""
+    link([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
+    link([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+
+Transform the variables in `vi` to their linked space, 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!!(deepcopy(vi), SampleFromPrior(), model)
+function link(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
+    return link!!(t, deepcopy(vi), SampleFromPrior(), 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)
+end
+
 """
     invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
     invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
 
 Transform the variables in `vi` to their constrained space, using the (inverse of) 
-transformation `t`.
+transformation `t`, mutating `vi` if possible.
 
 If `t` is not provided, `default_transformation(model, vi)` will be used.
 
@@ -434,6 +454,25 @@
     return settrans!!(vi_new, NoTransformation())
 end
 
+"""
+    invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
+    invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+
+Transform the variables in `vi` to their constrained space, using the (inverse of)
+transformation `t`.
+
+If `t` is not provided, `default_transformation(model, vi)` will be used.
+
+See also: [`default_transformation`](@ref), [`link`](@ref).
+"""
+invlink(vi::AbstractVarInfo, model::Model) = invlink(vi, SampleFromPrior(), model)
+function invlink(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
+    return invlink(t, vi, SampleFromPrior(), model)
+end
+function invlink(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    return invlink(transformation(vi), vi, spl, model)
+end
+
 """
     maybe_invlink_before_eval!!([t::Transformation,] vi, context, model)
 

src/varinfo.jl

@@ -885,7 +885,6 @@
 end
 
 function _inner_transform!(vi::VarInfo, vn::VarName, dist, f)
-    @debug "X -> ℝ for $(vn)..."
     # TODO: Use inplace versions to avoid allocations
     y, logjac = with_logabsdet_jacobian_and_reconstruct(f, dist, getval(vi, vn))
     yvec = vectorize(dist, y)
@@ -899,6 +898,144 @@
     return vi
 end
 
+function link(
+    ::DynamicTransformation, varinfo::VarInfo, spl::AbstractSampler, model::Model
+)
+    return _link(varinfo)
+end
+
+function _link(varinfo::UntypedVarInfo)
+    varinfo = deepcopy(varinfo)
+    return VarInfo(
+        _link_metadata!(varinfo, varinfo.metadata),
+        Base.Ref(getlogp(varinfo)),
+        Ref(get_num_produce(varinfo)),
+    )
+end
+
+function _link(varinfo::TypedVarInfo)
+    varinfo = deepcopy(varinfo)
+    md = map(Base.Fix1(_link_metadata!, varinfo), varinfo.metadata)
+    # TODO: Update logp, etc.
+    return VarInfo(md, Base.Ref(getlogp(varinfo)), Ref(get_num_produce(varinfo)))
+end
+
+function _link_metadata!(varinfo::VarInfo, metadata::Metadata)
+    vns = metadata.vns
+
+    # Construct the new transformed values, and keep track of their lengths.
+    vals_new = map(vns) do vn
+        # Return early if we're already in unconstrained space.
+        if istrans(varinfo, vn)
+            return metadata.vals[getrange(metadata, vn)]
+        end
+
+        # Transform to constrained space.
+        x = getval(varinfo, vn)
+        dist = getdist(varinfo, vn)
+        f = link_transform(dist)
+        y, logjac = with_logabsdet_jacobian_and_reconstruct(f, dist, x)
+        # Vectorize value.
+        yvec = vectorize(dist, y)
+        # Accumulate the log-abs-det jacobian correction.
+        acclogp!!(varinfo, -logjac)
+        # Mark as no longer transformed.
+        settrans!!(varinfo, true, vn)
+        # Return the vectorized transformed value.
+        return yvec
+    end
+
+    # Determine new ranges.
+    ranges_new = similar(metadata.ranges)
+    offset = 0
+    for (i, v) in enumerate(vals_new)
+        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 Metadata(
+        metadata.idcs,
+        metadata.vns,
+        ranges_new,
+        reduce(vcat, vals_new),
+        metadata.dists,
+        metadata.gids,
+        metadata.orders,
+        metadata.flags,
+    )
+end
+
+function invlink(
+    ::DynamicTransformation, varinfo::VarInfo, spl::AbstractSampler, model::Model
+)
+    return _invlink(varinfo)
+end
+
+function _invlink(varinfo::UntypedVarInfo)
+    varinfo = deepcopy(varinfo)
+    return VarInfo(
+        _invlink_metadata!(varinfo, varinfo.metadata),
+        Base.Ref(getlogp(varinfo)),
+        Ref(get_num_produce(varinfo)),
+    )
+end
+
+function _invlink(varinfo::TypedVarInfo)
+    varinfo = deepcopy(varinfo)
+    md = map(Base.Fix1(_invlink_metadata!, varinfo), varinfo.metadata)
+    # TODO: Update logp, etc.
+    return VarInfo(md, Base.Ref(getlogp(varinfo)), Ref(get_num_produce(varinfo)))
+end
+
+function _invlink_metadata!(varinfo::VarInfo, metadata::Metadata)
+    vns = metadata.vns
+
+    # Construct the new transformed values, and keep track of their lengths.
+    vals_new = map(vns) do vn
+        # Return early if we're already in constrained space.
+        if !istrans(varinfo, vn)
+            return metadata.vals[getrange(metadata, vn)]
+        end
+
+        # Transform to constrained space.
+        y = getval(varinfo, vn)
+        dist = getdist(varinfo, vn)
+        f = invlink_transform(dist)
+        x, logjac = with_logabsdet_jacobian_and_reconstruct(f, dist, y)
+        # Vectorize value.
+        xvec = vectorize(dist, x)
+        # Accumulate the log-abs-det jacobian correction.
+        acclogp!!(varinfo, -logjac)
+        # Mark as no longer transformed.
+        settrans!!(varinfo, false, vn)
+        # Return the vectorized transformed value.
+        return xvec
+    end
+
+    # Determine new ranges.
+    ranges_new = similar(metadata.ranges)
+    offset = 0
+    for (i, v) in enumerate(vals_new)
+        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 Metadata(
+        metadata.idcs,
+        metadata.vns,
+        ranges_new,
+        reduce(vcat, vals_new),
+        metadata.dists,
+        metadata.gids,
+        metadata.orders,
+        metadata.flags,
+    )
+end
+
 """
     islinked(vi::VarInfo, spl::Union{Sampler, SampleFromPrior})