Fix UnitVector implementation

[sethaxen]

This PR implements the proposed changes to UnitVector in #66 (i.e. Stan's approach). If x == zeros(n), the transformation is technically undefined. Stan handles this by always initializing with jitter. Since we can't control how users initialize, zeros(n) is instead mapped deterministically to a valid unit vector.

Along the way, a few tests had to be changed to handle cases like this where the "inverse" is only a right inverse. I also had to add a compat entry for Flux, since Tracker was removed from Flux in v0.10.

Here's a worked example, sampling from the uniform distribution on the circle and sphere:

using TransformVariables, LogDensityProblems, DynamicHMC, Random, LinearAlgebra
using Plots; plotlyjs()

function sample_post(rng, n, nsamples)
    f(θ) = zero(eltype(θ.x))
    trans = as((x = UnitVector(n),))
    P = TransformedLogDensity(trans, f)
    ∇P = ADgradient(:ForwardDiff, P)
    results = mcmc_with_warmup(rng, ∇P, nsamples; reporter = NoProgressReport())
    posterior = transform.(trans, results.chain)
    x = hcat(first.(posterior)...)
    return x
end

julia> rng = MersenneTwister(42);

julia> x = sample_post(rng, 2, 100000); # sample from circle

julia> all(norm.(eachcol(x)) .≈ 1)
true

julia> θ = (xy -> atan(xy[2], xy[1])).(eachcol(x));  # angle around the circle

julia> histogram(θ);

julia> png("histogram2d.png");

julia> x = sample_post(rng, 3, 100000); # sample from sphere

julia> all(norm.(eachcol(x)) .≈ 1)
true

julia> scatter3d(eachrow(x)...; markersize = 1, markeralpha = 0.2);

julia> png("sample3d.png");

This PR produces a uniform distribution on the circle (right), while the existing implementation (left) does not:

This PR likewise produces a uniform distribution on the sphere (right), while the existing implementation (left) gives a distribution only on the hemisphere:

[codecov]

Codecov Report

Merging #67 into master will increase coverage by 0.35%.
The diff coverage is 100%.

@@            Coverage Diff             @@
##           master      #67      +/-   ##
==========================================
+ Coverage   93.08%   93.44%   +0.35%     
==========================================
  Files           7        7              
  Lines         246      244       -2     
==========================================
- Hits          229      228       -1     
+ Misses         17       16       -1

Impacted Files	Coverage Δ
src/aggregation.jl	100% <ø> (ø)	⬆️
src/special_arrays.jl	100% <100%> (ø)	⬆️
src/generic.jl	95.45% <0%> (+4.97%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 44d9dd7...809c8dd. Read the comment docs.

[sethaxen]

This PR is ready for review.

[tpapp]

Thanks for the very nice contribution! I am happy to include it, but I would prefer not to make a breaking change to UnitVector.

So I would ask that you leave the existing UnitVector code in place (ie revert those changes), and introduce a new type/transformation — I would suggest SphericalUnitVector. For this, I am not sure that an inverse should be defined, as it is not invertible; but if you need it, feel free to define a method, just please emphasize in the docstring that it is not a bijection so strictly speaking it is not invertible.

[sethaxen]

I'm happy to make those changes. But I also think it's necessary to document that UnitVector is broken in the sense that while it does produce unit vectors, it does not cover the whole sphere or use a uniform measure, which will likely be unexpected for users (it certainly was for me).

[tpapp]

Perhaps I would not call it "broken", for some applications it can be fine. There are trade-offs between the two mappings, eg the spherical one is not invertible. It is best if the properties are documented so that the users can pick the appropriate one.

[sethaxen]

Hmm I don't see a good way to revert such that I keep UnitVector as it was with its tests but with SphericalUnitVector. I'll integrate these changes in a new PR.

[sethaxen]

Closing because stale and I don't intend to work on this further.

[tpapp]

@sethaxen, I apologize for not pursuing this.

[tpapp]

@sethaxen, I apologize for not pursuing this. I fully understand that you don't want to work on this further, but would you give permission for me to use your code and fix the issue?

[sethaxen]

@sethaxen, I apologize for not pursuing this. I fully understand that you don't want to work on this further, but would you give permission for me to use your code and fix the issue?

No problem, after all, I had originally planned to do it myself. Yes, I'm very happy for you to use the code.