Make corr_cholesky_factor deal with large inputs correctly.

Author: tpapp

Project.toml

@@ -1,7 +1,7 @@
 name = "TransformVariables"
 uuid = "84d833dd-6860-57f9-a1a7-6da5db126cff"
 authors = ["Tamas K. Papp <tkpapp@gmail.com>"]
-version = "0.8.9"
+version = "0.8.10"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

src/special_arrays.jl

@@ -7,36 +7,53 @@ export UnitVector, UnitSimplex, CorrCholeskyFactor, corr_cholesky_factor
 """
 $(SIGNATURES)
 
-`log(abs(…))` of the derivative of `tanh`, calculated accurately.
+Return a `NamedTuple` of
+
+- `log_l2_rem`, for `log(1 - tanh(x)^2)`,
+
+- `logjac`, for `log(abs( ∂(log(abs(tanh(x))) / ∂x ))`
+
+Caller ensures that `x ≥ 0`. `x == 0` is handled correctly, but results in infinities.
 """
-function _tanh_logabsderiv(x)
+function tanh_helpers(x)
     d = 2*x
-    log(4) + d - 2 * log1pexp(d)
+    log_denom = log1pexp(d)             # log(exp(2x) + 1)
+    logjac = log(4) + d - 2 * log_denom # log(ab
+    log_l2_rem = 2*(log(2) + x - log_denom)   # log(2exp(x) / (exp(2x) + 1))
+    (; logjac, log_l2_rem)
 end
 
 """
-    (y, r, ℓ) = $SIGNATURES
+    (y, log_r, ℓ) = $SIGNATURES
 
-Given ``x ∈ ℝ`` and ``0 ≤ r ≤ 1``, return `(y, r′)` such that
+Given ``x ∈ ℝ`` and ``0 ≤ r ≤ 1``, we define `(y, r′)` such that
 
 1. ``y² + (r′)² = r²``,
 
-2. ``y: |y| ≤ r`` is mapped with a bijection from `x`.
+2. ``y: |y| ≤ r`` is mapped with a bijection from `x`, with the sign depending on `x`,
+
+but use `log(r)` for actual calculations so that large `y`s still give nonsingular results.
 
 `ℓ` is the log Jacobian (whether it is evaluated depends on `flag`).
 """
-@inline function l2_remainder_transform(flag::LogJacFlag, x, r)
+@inline function l2_remainder_transform(flag::LogJacFlag, x, log_r)
+    (; logjac, log_l2_rem) = tanh_helpers(x)
     # note that 1-tanh(x)^2 = sech(x)^2
-    (tanh(x) * √r, r*sech(x)^2,
-     flag isa NoLogJac ? flag : _tanh_logabsderiv(x) + 0.5*log(r))
+    (tanh(x) * exp(log_r / 2),
+     log_r + log_l2_rem,
+     flag isa NoLogJac ? flag : logjac + 0.5*log_r)
 end
 
 """
     (x, r′) = $SIGNATURES
 
 Inverse of [`l2_remainder_transform`](@ref) in `x` and `y`.
 """
-@inline l2_remainder_inverse(y, r) = atanh(y/√r), r-y^2
+@inline function l2_remainder_inverse(y, log_r)
+    x = atanh(y / exp(log_r / 2))
+    log_r′ = logsubexp(log_r, 2 * log(abs(y)))
+    x, log_r′
+end
 
 ####
 #### UnitVector
@@ -65,16 +82,16 @@ end
 function transform_with(flag::LogJacFlag, t::UnitVector, x::AbstractVector, index)
     @unpack n = t
     T = robust_eltype(x)
-    r = one(T)
+    log_r = zero(T)
     y = Vector{T}(undef, n)
     ℓ = logjac_zero(flag, T)
     @inbounds for i in 1:(n - 1)
         xi = x[index]
         index += 1
-        y[i], r, ℓi = l2_remainder_transform(flag, xi, r)
+        y[i], log_r, ℓi = l2_remainder_transform(flag, xi, log_r)
         ℓ += ℓi
     end
-    y[end] = √r
+    y[end] = exp(log_r / 2)
     y, ℓ, index
 end
 
@@ -83,9 +100,9 @@ inverse_eltype(t::UnitVector, y::AbstractVector) = robust_eltype(y)
 function inverse_at!(x::AbstractVector, index, t::UnitVector, y::AbstractVector)
     @unpack n = t
     @argcheck length(y) == n
-    r = one(eltype(y))
+    log_r = zero(eltype(y))
     @inbounds for yi in axes(y, 1)[1:(end-1)]
-        x[index], r = l2_remainder_inverse(y[yi], r)
+        x[index], log_r = l2_remainder_inverse(y[yi], log_r)
         index += 1
     end
     index
@@ -244,14 +261,14 @@ function calculate_corr_cholesky_factor!(U::AbstractMatrix{T}, flag::LogJacFlag,
     n = size(U, 1)
     ℓ = logjac_zero(flag, T)
     @inbounds for col_index in 1:n
-        r = one(T)
+        log_r = zero(T)
         for row_index in 1:(col_index-1)
             xi = x[index]
-            U[row_index, col_index], r, ℓi = l2_remainder_transform(flag, xi, r)
+            U[row_index, col_index], log_r, ℓi = l2_remainder_transform(flag, xi, log_r)
             ℓ += ℓi
             index += 1
         end
-        U[col_index, col_index] = √r
+        U[col_index, col_index] = exp(log_r / 2)
     end
     U, ℓ, index
 end
@@ -285,9 +302,9 @@ function inverse_at!(x::AbstractVector, index,
     n = result_size(t)
     @argcheck size(U, 1) == n
     @inbounds for col in 1:n
-        r = one(eltype(U))
+        log_r = zero(eltype(U))
         for row in 1:(col-1)
-            x[index], r = l2_remainder_inverse(U[row, col], r)
+            x[index], log_r = l2_remainder_inverse(U[row, col], log_r)
             index += 1
         end
     end

test/runtests.jl

@@ -6,7 +6,8 @@ using LogDensityProblems: logdensity, logdensity_and_gradient
 using LogDensityProblemsAD
 using TransformVariables:
     AbstractTransform, ScalarTransform, VectorTransform, ArrayTransformation,
-    unit_triangular_dimension, logistic, logistic_logjac, logit, inverse_and_logjac, NOLOGJAC, transform_with
+    unit_triangular_dimension, logistic, logistic_logjac, logit, inverse_and_logjac,
+    NOLOGJAC, transform_with
 import ChangesOfVariables, InverseFunctions
 using Enzyme: autodiff, ReverseWithPrimal, Active, Const
 
@@ -136,9 +137,18 @@ end
     end
 end
 
+@testset "tanh helpers" begin
+    for _ in 1:10000
+        x = (rand() - 0.5) * 100
+        @unpack log_l2_rem, logjac =  TransformVariables.tanh_helpers(x)
+        @test Float64(AD_logjac(tanh, BigFloat(x))) ≈ logjac atol = 1e-4
+        @test Float64(log(sech(BigFloat(x))^2)) ≈ log_l2_rem atol = 1e-4
+    end
+end
+
 @testset "to correlation cholesky factor" begin
     @testset "dimension checks" begin
-        C = CorrCholeskyFactor(3)
+        C = corr_cholesky_factor(3)
         wrong_x = zeros(dimension(C) + 1)
 
         @test_throws ArgumentError transform(C, wrong_x)
@@ -147,7 +157,7 @@ end
 
     @testset "consistency checks" begin
         for K in 1:8
-            t = CorrCholeskyFactor(K)
+            t = corr_cholesky_factor(K)
             @test dimension(t) == (K - 1)*K/2
             CIENV && @info "testing correlation cholesky K = $(K)"
             if K > 1
@@ -615,6 +625,15 @@ end
     end
 end
 
+@testset "corr cholesky factor large inputs" begin
+    t = corr_cholesky_factor(7)
+    d = dimension(t)
+    for _ in 1:100
+        x = sign.(rand(d) .- 0.5) .* 100
+        @test isfinite(logdet(transform(t, x)) )
+    end
+end
+
 @testset "pretty printing" begin
     t = as((a = asℝ₊,
             b = as(Array, asℝ₋, 3, 3),

test/utilities.jl

@@ -3,11 +3,12 @@ $(SIGNATURES)
 
 Log jacobian abs determinant via automatic differentiation. For testing.
 """
+AD_logjac(f, x) = log(abs(ForwardDiff.derivative(f, x)))
+
 AD_logjac(t::VectorTransform, x, vec_y) =
     logabsdet(ForwardDiff.jacobian(x -> vec_y(transform(t, x)), x))[1]
 
-AD_logjac(t::ScalarTransform, x) =
-    log(abs(ForwardDiff.derivative(x -> transform(t, x), x)))
+AD_logjac(t::ScalarTransform, x) = AD_logjac(x -> transform(t, x), x)
 
 """
 $(SIGNATURES)