Transforms with purely functional state

[cscherrer]

TransformVariables.jl's transform takes a transformation and a floating-point vector x to a structured value (e.g. a named tuple) y. But transform is just a wrapper around transform_with, which does most of the work. I really like the way this is structured, and I think something like it can be useful for us.

Say our transform is

julia> tr = as(Array, as𝕀, 2, 3)
TransformVariables.ArrayTransform{TransformVariables.ScaledShiftedLogistic{Int64}, 2}(as𝕀, (2, 3))

This can take six unbounded reals and map them to a 2x3 array with each component in the unit interval. But often we have a transform that's made up of other transforms; this is where transform_with comes in.

Say I started with a vector

julia> x = rand(Float16, 10)
10-element Vector{Float16}:
 0.4092
 0.5435
 0.248
 0.7397
 0.9946
 0.8535
 0.1606
 0.746
 0.4292
 0.784

(16 bit so they take up less screen space)

Then we can call

julia> (y, ladj, index) = TV.transform_with(TV.LogJac(), tr, x, 1)
(Float16[0.601 0.5615 0.73; 0.6323 0.677 0.701], -8.9951171875, 7)

The inputs are

• A LogJacFlag indicating whether we want to compute the logjac,
• The transform,
• The input vector, and
• The starting index

And as a result we get

• The desired 2x3 array,
• The logjac, and
• The next index

The index state tells us the next index to use, a lot like you'd do with Julia's iterator protocol.

I think we can have something like this, except that the "transform" can be more general. Maybe it maps from $\mathbb{R}^n$, or maybe from a hypercube. It might be a "transform" in the TransformVariables sense, or it could be a transport.

Another thing... In Tilde, a transform will have to call runmodel. In a naive implementation, we'd have to call that again to get the log-density. Instead, we should set this up so we can get the log-density at the same time as we're doing the transform. This can be much more efficient. Note that we don't need this for transports, because in that case we can get the log-density more directly from x, and y is only needed for computing the likelihood. In fact, we may want to allow only computing those variables that will be needed for the likelihood, but wiring things up that way could be tricky, so let's come back to that.