Redundant computation for Affine measures

[cscherrer]

Say we have a measure like

julia> d = MvNormal((σ=[3;4;;],))
Affine{(:σ,),PowerMeasure{typeof(identity), Normal{(), Tuple{}}, 1, Tuple{Base.OneTo{Int64}}},Tuple{Matrix{Int64}}}(AffineTransform{(:σ,),Tuple{Matrix{Int64}}}((σ = [3; 4;;],)), Normal() ^ 1)

Then this works just fine:

julia> x = rand(d)
2-element Vector{Float64}:
 -3.212402995915957
 -4.283203994554609

julia> logdensity(d,x)
-0.5733073893427674

But there's a problem. Along the way, this requires calculating

z = d.σ \ x

and the base measure is left as

julia> basemeasure(d)
0.398942 * Affine{(:σ,),PowerMeasure{typeof(identity), Lebesgue{ℝ}, 1, Tuple{Base.OneTo{Int64}}},Tuple{Matrix{Int64}}}(AffineTransform{(:σ,),Tuple{Matrix{Int64}}}((σ = [3; 4;;],)), Lebesgue(ℝ) ^ 1)

So computing the logdensity with respect to other measures will typically require more repetitions of the z = d.σ \ x code, which is very inefficient.

Adding some @show statements in logpdf makes this clear:

julia> logpdf(d, [0,0])
d = Affine{(:σ,),PowerMeasure{typeof(identity), Normal{(), Tuple{}}, 1, Tuple{Base.OneTo{Int64}}},Tuple{Matrix{Int64}}}(AffineTransform{(:σ,),Tuple{Matrix{Int64}}}((σ = [3; 4;;],)), Normal() ^ 1)
z = [-0.0]

d = 0.398942 * Affine{(:σ,),PowerMeasure{typeof(identity), Lebesgue{ℝ}, 1, Tuple{Base.OneTo{Int64}}},Tuple{Matrix{Int64}}}(AffineTransform{(:σ,),Tuple{Matrix{Int64}}}((σ = [3; 4;;],)), Lebesgue(ℝ) ^ 1)
z = [-0.0]

d = Affine{(:σ,),PowerMeasure{typeof(identity), Lebesgue{ℝ}, 1, Tuple{Base.OneTo{Int64}}},Tuple{Matrix{Int64}}}(AffineTransform{(:σ,),Tuple{Matrix{Int64}}}((σ = [3; 4;;],)), Lebesgue(ℝ) ^ 1)
z = [-0.0]

d = 0.2 * Lebesgue(ℝ) ^ 1
d = Lebesgue(ℝ) ^ 1
-2.528376445638773

Each z = ... is the result of solving the same linear system, three times.

I got a start on fixing this, by adding a MapsTo combinator:

struct MapsTo{X,Y}
    x::X
    y::Y
end

export ↦, mapsto

mapsto(x,y) = x ↦ y

↦(x::X, y::Y) where {X,Y} = MapsTo{X,Y}(x,y)

Base.first(t::MapsTo) = t.x
Base.last(t::MapsTo) = t.y

Base.Pair(t::MapsTo) = t.x => t.y

Base.show(io::IO, t::MapsTo) = print(t.x, " ↦ ", t.y)

It's pretty easy to make this work with logpdf, but for a more generic logdensity I think it will take some restructuring. We might need a LogdensityComputation struct that's just for internal use, and holds things we know so far about the result.