add softplus and Cauchy transformation

Project.toml

@@ -13,6 +13,7 @@ CompositionsBase = "a33af91c-f02d-484b-be07-31d278c5ca2b"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -40,6 +41,7 @@ ConstructionBase = "1.6.0"
 DocStringExtensions = "0.8, 0.9"
 ForwardDiff = "0.10, 1"
 InverseFunctions = "0.1"
+IrrationalConstants = "0.2"
 LinearAlgebra = "1.6"
 LogExpFunctions = "0.3, 1"
 Random = "1.6"

docs/src/index.md

@@ -118,8 +118,8 @@ t = as((; e = Identity(), foo = tf))
 ```
 where each instance of `Identity()` here could be replaced with an arbitrary scalar transform.
 
-This relies on [`constructorof` from ConstructionBase](https://juliaobjects.github.io/ConstructionBase.jl/stable/#ConstructionBase.constructorof). If a `Tuple` transform is wrapped in a type this way, transform results will be unpacked and passed to the constructor, hence in the same order as the transform. 
-If a `NamedTuple` transform is used, it will be unpacked as keyword arguments to the constructor of the type. 
+This relies on [`constructorof` from ConstructionBase](https://juliaobjects.github.io/ConstructionBase.jl/stable/#ConstructionBase.constructorof). If a `Tuple` transform is wrapped in a type this way, transform results will be unpacked and passed to the constructor, hence in the same order as the transform.
+If a `NamedTuple` transform is used, it will be unpacked as keyword arguments to the constructor of the type.
 `Tuple` or `NamedTuple` transforms that do not provide proper arguments to the constructor of a given type will simply result in `MethodError`s.
 
 For structs that are constructed from a single argument, wrap a single transform in a `Tuple`, like so:
@@ -128,8 +128,8 @@ struct Baz; f; end
 tbaz = as(Baz, (Identity(),))
 ```
 
-Inverting these transforms would, in the most general case, require inverting the constructor of the given type, which may itself have several valid dispatches. 
-Rather than conduct a close inspection of user types and their available constructors, inverting a `Tuple`-based struct transform will check that the struct has exactly `n` fields, where `n` is the length of the `Tuple` transform, and use those fields in their struct order; if there are not `n` fields the inversion will error. This will work for any structs with only the default constructor. Inverting a `NamedTuple`-based struct transform will attempt to use struct fields with names matching the names in the `NamedTuple` transform, which fails if the struct does not have matching fields. 
+Inverting these transforms would, in the most general case, require inverting the constructor of the given type, which may itself have several valid dispatches.
+Rather than conduct a close inspection of user types and their available constructors, inverting a `Tuple`-based struct transform will check that the struct has exactly `n` fields, where `n` is the length of the `Tuple` transform, and use those fields in their struct order; if there are not `n` fields the inversion will error. This will work for any structs with only the default constructor. Inverting a `NamedTuple`-based struct transform will attempt to use struct fields with names matching the names in the `NamedTuple` transform, which fails if the struct does not have matching fields.
 
 ## Scalar transforms
 
@@ -162,8 +162,12 @@ TVLogistic
 TVScale
 TVShift
 TVNeg
+TVSoftPlus
+TVCauchy
 ```
 
+The above include additional transformations that don't have a short syntax with `as`.
+
 Consistent with common notation, transforms are applied right-to-left; for example, `as(Real, -∞, 3)` is equivalent to `TVShift(3) ∘ TVNeg() ∘ TVExp()`.
 If you are working in an editor where typing Unicode is difficult, `TransformVariables.compose` is also available, as in `TransformVariables.compose(TVScale(5.0), TVNeg(), TVExp())`.
 
@@ -173,16 +177,16 @@ This is useful especially for making sure that values near 0, when transformed,
 In addition, the `TVScale` transform accepts arbitrary types. It can be used as the outermost transform (so leftmost in the composition) to add, for example, `Unitful` units to a number (or to create other exotic number types which can be constructed by multiplying, such as a `ForwardDiff.Dual`).
 
 However, note that calculating log Jacobian determinants may error for types that are not real numbers.
-For example, 
+For example,
 
 ```julia
 using Unitful
 t = TVScale(5u"m") ∘ TVExp()
 ```
-produces positive quantities with the dimension of length. 
+produces positive quantities with the dimension of length.
 !!! note
     Because the log-Jacobian of a transform that adds units is not defined, `transform_and_logjac` and `inverse_and_logjac`
-    only have methods defined for `TVScale{T} where {T<:Real}`. 
+    only have methods defined for `TVScale{T} where {T<:Real}`.
 !!! note
     The inverse transform of `TVScale(scale)` divides by `scale`, which is the correct inverse for adding units to a number, but may be inappropriate for other custom number types. A transform that doesn't just multiply or an inverse that extracts a float from an exotic number type could be defined by adding methods to `transform` and `inverse` like the following:
     ```

ext/ReactantExt.jl

@@ -1,5 +1,5 @@
 module ReactantExt
-using TransformVariables: TransformVariables, ArrayTransformation, LogJacFlag,  
+using TransformVariables: TransformVariables, ArrayTransformation, LogJacFlag,
                           logjac_zero, transform_with, _ensure_float, dimension
 using Reactant
 using Reactant: TracedRNumber, AnyTracedRArray
@@ -22,12 +22,18 @@ Base.@nospecialize function Reactant.traced_type_inner(
     return TransformVariables.ArrayTransformation{T_traced, M}
 end
 
-TransformVariables._ensure_float(x::Type{TracedRNumber{T}}) where T = TracedRNumber{_ensure_float(T)}
+function TransformVariables._ensure_float(x::Type{TracedRNumber{T}}) where T
+    TracedRNumber{_ensure_float(T)}
+end
+
+function TransformVariables._cauchy_cdf(x::TracedRNumber{T}) where {T<:AbstractFloat}
+    ifelse(x > -1, one(x) / 2 + atan(x) / π, atan(inv(-x)) / π)
+end
 
 @noinline function TransformVariables.transform_with(
-    flag::TransformVariables.LogJacFlag, 
-    transformation::TransformVariables.ArrayTransformation, 
-    x::AnyTracedRArray, 
+    flag::TransformVariables.LogJacFlag,
+    transformation::TransformVariables.ArrayTransformation,
+    x::AnyTracedRArray,
     index
 )
     (; inner_transformation, dims) = transformation
@@ -48,18 +54,18 @@ TransformVariables._ensure_float(x::Type{TracedRNumber{T}}) where T = TracedRNum
         throw(ArgumentError("only `Number` and `AbstractArray` transformations are supported in Reactant compilation mode"))
     end
 
-    # Reactant can't handle arrays of non-primitive types so we have to be more explicit about 
-    # the output and use a `@trace` for loop. 
+    # Reactant can't handle arrays of non-primitive types so we have to be more explicit about
+    # the output and use a `@trace` for loop.
     @trace track_numbers=false for i in eachindex(𝐼)
         idx = 𝐼[i]
         y, ℓ, _ = transform_with(flag, inner_transformation, x, idx)
         if !isempty(y)
-            ℓa += ℓ        
+            ℓa += ℓ
         end
         if tmp isa Number
             Reactant.@allowscalar yℓ[i] = y
         else
-            idxs = ntuple(n-> Colon(), ndims(tmp)) 
+            idxs = ntuple(n-> Colon(), ndims(tmp))
             yℓ[idxs..., i] = y
         end
     end
@@ -83,4 +89,4 @@ function TransformVariables.transform_with(flag::LogJacFlag, t::ArrayTransformat
     y, logjac_zero(flag, _ensure_float(eltype(x))), index′
 end
 
-end
+end

src/TransformVariables.jl

@@ -4,6 +4,7 @@ using ArgCheck: @argcheck
 import Compat
 using DocStringExtensions: FUNCTIONNAME, SIGNATURES, TYPEDEF
 import ForwardDiff
+import IrrationalConstants
 using LogExpFunctions
 using LinearAlgebra: UpperTriangular, logabsdet, norm, rmul!
 using Random: AbstractRNG, GLOBAL_RNG

src/scalar.jl

@@ -1,4 +1,5 @@
-export TVExp, TVScale, TVShift, TVLogistic, TVNeg
+export TVExp, TVScale, TVShift, TVLogistic, TVNeg, TVSoftPlus, TVCauchy
+
 export ∞, asℝ, asℝ₊, asℝ₋, as𝕀, as_real, as_positive_real, as_negative_real,
     as_unit_interval
 
@@ -157,6 +158,66 @@ inverse_eltype(::TVNeg, ::Type{T}) where T = typeof(-oneunit(T))
 inverse(::TVNeg, x::Number) = -x
 inverse_and_logjac(::TVNeg, x::Number) = -x, logjac_zero(LogJac(), typeof(x))
 
+"""
+$(TYPEDEF)
+
+“Softplus” transformation ``x ↦ log(1+exp(x))``, mapping the real line to positive
+numbers.
+"""
+struct TVSoftPlus <: ScalarTransform end
+
+transform(::TVSoftPlus, x::Number) = log1pexp(x)
+transform_and_logjac(t::TVSoftPlus, x::Number) = transform(t, x), -log1pexp(-x)
+
+inverse_eltype(::TVSoftPlus, ::Type{T}) where T = _ensure_float(T)
+inverse(::TVSoftPlus, y::Number) = logexpm1(y)
+inverse_and_logjac(::TVSoftPlus, y::Number) = logexpm1(y), -log1mexp(-y)
+
+"""
+$(TYPEDEF)
+
+Transformation ``x ↦ 1/2 + arctan(x)/π``, mapping the real line to `(0, 1)`, which is
+also the CDF of the Cauchy distribution.
+"""
+struct TVCauchy <: ScalarTransform end
+
+function _cauchy_cdf(x::T) where {T<:Real}
+    _x = float(x)
+    if _x > -1
+        one(_x) / 2 + atan(_x) / π
+    else
+        # avoid 1/2 - 1/2 cancellations
+        atan(inv(-_x)) / π
+    end
+end
+
+transform(::TVCauchy, x::Number) = _cauchy_cdf(float(x))
+
+function transform_and_logjac(t::TVCauchy, x::Number)
+    y = transform(t, x)
+    y,  -(log1psq(x) + IrrationalConstants.logπ)
+end
+
+inverse_eltype(::TVCauchy, ::Type{T}) where T = _ensure_float(T)
+
+function inverse(::TVCauchy, y::Number)
+    @argcheck 0 <= y <= 1 DomainError
+    first(_cauchy_quantile_and_sinpi(y))
+end
+
+function inverse_and_logjac(::TVCauchy, y::Number)
+    @argcheck 0 <= y <= 1 DomainError
+    x, s = _cauchy_quantile_and_sinpi(y)
+    x, IrrationalConstants.logπ - 2 * log(s)
+end
+
+# returns the quantile and sinpi(y)
+function _cauchy_quantile_and_sinpi(y)
+    flip = y > 0.5
+    s, c = sincospi(flip ? 1 - y : y)
+    (flip ? c / s : -c / s), s
+end
+
 ####
 #### composite scalar transforms
 ####

test/runtests.jl

@@ -470,7 +470,7 @@ end
         @test tn[1] == tn.a == inner[1]
         @test tn[2] == tn.b == inner[2]
         @test tn[3] == tn.c == inner[3]
-        @test_throws ErrorException tn.d
+        @test_throws (VERSION < v"1.12" ? ErrorException : FieldError) tn.d
         @test_throws BoundsError tn[4]
         @test propertynames(tn) == propertynames(inner)
         @test (@set tn[3] = t2) == (@set tn.c = t2) == as((a = t1, b = t2, c = t2))
@@ -1145,9 +1145,25 @@ end
     @test_throws InexactError inverse(t, fill(Complex(0, 1), 3))
 end
 
+@testset "TVSoftPlus" begin
+    t = TVSoftPlus()
+    test_transformation(t, y -> y > 0)
+    @test transform(t, 1000.0) == 1000.0
+    @test inverse(t, 1000.0) == 1000.0
+    @test transform(t, -30.0) ≈ exp(-30.0)
+    @test inverse(t, exp(-30.0)) ≈ -30.0
+    @test @inferred(transform(t, 1f0)) isa Float32
+end
+
+@testset "TVCauchy" begin
+    test_transformation(TransformVariables.TVCauchy(), y -> 0 < y < 1)
+end
+
+
 ###
 ###  Reactant compat tests
 ###
+
 @testset "Reactant.jl" begin
     @testset "Scalar transforms" begin
         a = 3.1
@@ -1156,7 +1172,8 @@ end
         @test @jit(transform(asℝ₊, ar)) ≈ transform(asℝ₊, a)
         @test @jit(transform(asℝ₋, ar)) ≈ transform(asℝ₋, a)
         @test @jit(transform(as𝕀, ar)) ≈ transform(as𝕀, a)
-
+        @test @jit(transform(TVSoftPlus(), ar)) ≈ transform(TVSoftPlus(), a)
+        @test @jit(transform(TVCauchy(), ar)) ≈ transform(TVCauchy(), a)
     end
 
     @testset "Array transforms begin" begin