Fixes to correlation bijectors

[sethaxen]

The current bijectors for correlation matrices and their Cholesky factors when used in Turing are still quite brittle. A model just drawing from a uniform LKJCholesky regularly errors:

julia> using Turing

julia> @model function model(K, eta)
           x ~ LKJCholesky(K, eta, 'U')
       end;

julia> chns = sample(model(5, 1.0), NUTS(), 10_000)
┌ Info: Found initial step size
└   ϵ = 1.6
Sampling 100%|██████████████████████████████████████| Time: 0:00:04
ERROR: DomainError with -1.0798579785920026e-7:
sqrt was called with a negative real argument but will only return a complex result if called with a complex argument. Try sqrt(Complex(x)).
DomainError detected in the user `f` function. This occurs when the domain of a function is violated.
For example, `log(-1.0)` is undefined because `log` of a real number is defined to only output real
numbers, but `log` of a negative number is complex valued and therefore Julia throws a DomainError
by default. Cases to be aware of include:

* `log(x)`, `sqrt(x)`, `cbrt(x)`, etc. where `x<0`
* `x^y` for `x<0` floating point `y` (example: `(-1.0)^(1/2) == im`)

Within the context of SciML, this error can occur within the solver process even if the domain constraint
would not be violated in the solution due to adaptivity. For example, an ODE solver or optimization
routine may check a step at `new_u` which violates the domain constraint, and if violated reject the
step and use a smaller `dt`. However, the throwing of this error will have halted the solving process.

Thus the recommended fix is to replace this function with the equivalent ones from NaNMath.jl
(https://github.com/JuliaMath/NaNMath.jl) which returns a NaN instead of an error. The solver will then
effectively use the NaN within the error control routines to reject the out of bounds step. Additionally,
one could perform a domain transformation on the variables so that such an issue does not occur in the
definition of `f`.

For more information, check out the following FAQ page:
https://docs.sciml.ai/Optimization/stable/API/FAQ/#The-Solver-Seems-to-Violate-Constraints-During-the-Optimization,-Causing-DomainErrors,-What-Can-I-Do-About-That?

Stacktrace:
  [1] throw_complex_domainerror(f::Symbol, x::Float64)
    @ Base.Math ./math.jl:33
  [2] sqrt
    @ ./math.jl:686 [inlined]
  [3] sqrt
    @ ~/.julia/packages/ForwardDiff/PcZ48/src/dual.jl:240 [inlined]
  [4] _link_chol_lkj_from_upper(W::LinearAlgebra.UpperTriangular{ForwardDiff.Dual{…}, Matrix{…}})
    @ Bijectors ~/.julia/packages/Bijectors/gR4QV/src/bijectors/corr.jl:319
  [5] transform(b::Bijectors.VecCholeskyBijector, X::LinearAlgebra.Cholesky{ForwardDiff.Dual{…}, Matrix{…}})
    @ Bijectors ~/.julia/packages/Bijectors/gR4QV/src/bijectors/corr.jl:218
  [6] Transform
    @ ~/.julia/packages/Bijectors/gR4QV/src/interface.jl:82 [inlined]
  [7] logabsdetjac
    @ ~/.julia/packages/Bijectors/gR4QV/src/bijectors/corr.jl:225 [inlined]
  [8] with_logabsdet_jacobian(ib::Inverse{Bijectors.VecCholeskyBijector}, y::Vector{ForwardDiff.Dual{…}})
    @ Bijectors ~/.julia/packages/Bijectors/gR4QV/src/interface.jl:215

As the error message shows, this is coming from hitting the method

Bijectors.jl/src/interface.jl

Lines 213 to 216 in 2402be2

	function with_logabsdet_jacobian(ib::Inverse{<:Transform}, y)
	x = transform(ib, y)
	return x, -logabsdetjac(inverse(ib), x)
	end

, where the transform is performed, then the inverse transform is performed, then part of the forward transform is performed to compute the logdetjac. This hits

Bijectors.jl/src/bijectors/corr.jl

Line 319 in 2402be2

tmp *= sqrt(1 - p^2)

, which is numerically unstable. I locally have a fix I will push.

At the same time, I would like to implement all with_logabsdet_jacobian methods for these bijectors, as the logdetjacs are trivially computed as part of the transform.

Likewise, there are some naming issues here. A number of the internal functions have lkj in the name. LKJ is one distribution of correlation matrices, but it's not the only one. This should be changed to corr or corrchol. Similarly, VecCholeskyBijector is incorrectly named. This bijector does not map from any Cholesky factor to a vector. Rather, it maps only from Cholesky factors of correlation matrices, so it should be renamed to VecCorrCholeskyBijector.