Some small `pushfwd` issues

[cscherrer]

I've been exploring our pushforward basics in MeasureBase.jl...

Say we have a uniform measure on $(-\pi/2, \pi/2)$ and want to push that through asin. So we do

# We'll use the others soon
using MeasureBase, ForwardDiff, ChangesOfVariables, IrrationalConstants

μ = Lebesgue(MeasureBase.BoundedReals(-halfπ, halfπ))

function f(x)
    asin(x)
end

function finv(y)
    if -halfπ ≤ y ≤ halfπ
        return sin(y)
    else
        @error "finv is only defined on [-π/2, π/2]"
    end
end

Now, we want to be able to do

ν = pushfwd(f, finv, μ)
logdensityof(ν, π/4)

We can't use InverseFunctions.jl to get the inverse automatically, because sin only has a one-sided inverse (not currently handled by that package). This three-argument form of pushfwd lets use say "this inverse works on the domain of $\mu$".

This won't quite work yet, because it doesn't know how to get the logjac. But we can define

function withlogjac(f, x)
    dx = ForwardDiff.Dual{ForwardDiff.Tag{typeof(f)}}(x, 1.0)
    dy = f(dx)
    value = dy.value
    deriv = first(dy.partials)
    (value, log(abs(deriv)))   
end

function ChangesOfVariables.with_logabsdet_jacobian(::typeof(f), x)
    withlogjac(f, x)
end

function ChangesOfVariables.with_logabsdet_jacobian(::typeof(finv), x)
    withlogjac(finv, x)
end

And we can test this implementation:

julia> ChangesOfVariables.test_with_logabsdet_jacobian(f, 2 * rand() - 1, ForwardDiff.derivative)
Test Summary:                                                | Pass  Total  Time
test_with_logabsdet_jacobian: f with input 0.359581995576044 |    2      2  0.0s

julia> ChangesOfVariables.test_with_logabsdet_jacobian(finv, π * rand() - halfπ, ForwardDiff.derivative)
Test Summary:                                                    | Pass  Total  Time
test_with_logabsdet_jacobian: finv with input 1.0808047560079297 |    2      2  0.0s

So now our pushforward works:

julia> ν = pushfwd(f, finv, μ)
PushforwardMeasure(
    f,
    finv,
    Lebesgue(MeasureBase.BoundedReals{Float64, Irrational{:halfπ}}(-1.5707963267948966, halfπ)))

julia> logdensityof(ν, π/4)
-0.3465735902799726

But let's make a small update to our functions to see what they're really doing:

function f(x)
    @info "calling f($x)"
    asin(x)
end

function finv(y)
    @info "calling finv($y)"
    if -halfπ ≤ y ≤ halfπ
        return sin(y)
    else
        @error "finv is only defined on [-π/2, π/2]"
    end
end

Our final call now looks like this:

julia> logdensityof(ν, π/4)
[ Info: calling finv(Dual{ForwardDiff.Tag{typeof(finv)}}(0.7853981633974483,1.0))
[ Info: calling finv(0.7853981633974483)
[ Info: calling finv(0.7853981633974483)
-0.3465735902799726

So this is making three calls to finv. These are

1. To check insupport
2. Calling logdensity_def
3. Calling logdensity_def on the base measure

The base measure is

julia> basemeasure(ν)
PushforwardMeasure(f, finv, MeasureBase.LebesgueBase())

So this raises a few questions/comments:

1. Can we update this to only call finv once? It's not so bad in this case, but it could sometimes get very expensive. Previously I had a MapsTo type for this sort of thing, maybe we need to bring that back?
2. The fact that it only calls finv (and never f) gets me back to thinking it's much more natural in lots of cases to work in terms of a pullback
3. Currently there's no way to get a density of nu with respect to, say, Lebesgue(). I think the right way to do this is to push Lebesgue() through finv (or pull back through f) and compare the result with mu.