diff --git a/Project.toml b/Project.toml index 1d77727..a949879 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/index.md b/docs/src/index.md index f21e0db..04823ef 100644 --- a/docs/src/index.md +++ b/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: ``` diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl index 339abf7..805a04b 100644 --- a/ext/ReactantExt.jl +++ b/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 \ No newline at end of file +end diff --git a/src/TransformVariables.jl b/src/TransformVariables.jl index a9140c9..9f7c8a5 100644 --- a/src/TransformVariables.jl +++ b/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 diff --git a/src/scalar.jl b/src/scalar.jl index cbf1d49..e64ddc5 100644 --- a/src/scalar.jl +++ b/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,65 @@ 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::Real) + 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(float(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 #### diff --git a/test/runtests.jl b/test/runtests.jl index 79c3ee7..cd86446 100644 --- a/test/runtests.jl +++ b/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,34 @@ 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 + t = TVCauchy() + test_transformation(t, y -> 0 < y < 1) + @test transform(t, -1e300) ≈ inv(π * 1e300) + @test transform(t, -Inf) == 0 + @test transform(t, Inf) == 1 + @test inverse(t, 0.0) == -Inf + @test inverse(t, 1.0) == Inf + @test_throws DomainError inverse(t, 1.5) + @test @inferred(transform(t, 1f0)) isa Float32 + @test @inferred(inverse(t, 0.3f0)) isa Float32 + @test transform_and_logjac(t, 2^40)[2] ≈ transform_and_logjac(t, 2.0^40)[2] +end + ### ### Reactant compat tests ### + @testset "Reactant.jl" begin @testset "Scalar transforms" begin a = 3.1 @@ -1156,7 +1181,11 @@ 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(TVSoftPlus(), -ar)) ≈ transform(TVSoftPlus(), -a) + @test @jit(transform(TVCauchy(), ar)) ≈ transform(TVCauchy(), a) + @test @jit(transform(TVCauchy(), -ar)) ≈ transform(TVCauchy(), -a) + @test @jit(transform(TVCauchy(), ConcreteRNumber(0.0))) ≈ transform(TVCauchy(), 0.0) end @testset "Array transforms begin" begin @@ -1173,7 +1202,7 @@ end end end @testset "Test inner transformations" begin - for T in (asℝ, asℝ₊, asℝ₋, as𝕀) + for T in (asℝ, asℝ₊, asℝ₋, as𝕀, TVSoftPlus(), TVCauchy()) tr = as(Array, as(Array, T, 3), 3) a = randn(dimension(tr)) ar = Reactant.to_rarray(a)