Add support for ChangesOfVariables.with_logabsdet_jacobian (#30)

* Support ChangesOfVariables.with_logabsdet_jacobian

* Document support for InverseFunctions and ChangesOfVariables

* Neatify with_logabdet_jacobian defs and tests

Co-authored-by: David Widmann <[email protected]>

* Fix with_logabsdet_jacobian

Co-authored-by: David Widmann <[email protected]>

Author: oschulz

Project.toml

@@ -5,13 +5,15 @@ version = "0.3.4"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
 ChainRulesCore = "1"
+ChangesOfVariables = "0.1"
 DocStringExtensions = "0.8"
 InverseFunctions = "0.1"
 IrrationalConstants = "0.1"

docs/src/index.md

@@ -4,6 +4,8 @@ Various special functions based on `log` and `exp` moved from [StatsFuns.jl](htt
 
 The original authors of these functions are the StatsFuns.jl contributors.
 
+LogExpFunctions supports [`InverseFunctions.inverse`](https://github.com/JuliaMath/InverseFunctions.jl) and [`ChangesOfVariables.test_with_logabsdet_jacobian`](https://github.com/JuliaMath/ChangesOfVariables.jl) for `log1mexp`, `log1pexp`, `log1pexp`, `log2mexp`, `log2mexp`, `logexpm1`, `logistic`, `logistic` and `logit`.
+
 ```@docs
 xlogx
 xlogy

src/LogExpFunctions.jl

@@ -4,6 +4,7 @@ using DocStringExtensions: SIGNATURES
 using Base: Math.@horner
 
 import ChainRulesCore
+import ChangesOfVariables
 import InverseFunctions
 import IrrationalConstants
 import LinearAlgebra
@@ -16,5 +17,6 @@ include("basicfuns.jl")
 include("logsumexp.jl")
 include("chainrules.jl")
 include("inverse.jl")
+include("with_logabsdet_jacobian.jl")
 
 end # module

src/with_logabsdet_jacobian.jl

@@ -0,0 +1,29 @@
+function ChangesOfVariables.with_logabsdet_jacobian(::typeof(log1pexp), x::Real)
+    y = log1pexp(x)
+    return y, x - y
+end
+
+function ChangesOfVariables.with_logabsdet_jacobian(::typeof(logexpm1), x::Real)
+    y = logexpm1(x)
+    return y, x - y
+end
+
+function ChangesOfVariables.with_logabsdet_jacobian(::typeof(log1mexp), x::Real)
+    y = log1mexp(x)
+    return y, x - y
+end
+
+function ChangesOfVariables.with_logabsdet_jacobian(::typeof(log2mexp), x::Real)
+    y = log2mexp(x)
+    return y, x - y
+end
+
+function ChangesOfVariables.with_logabsdet_jacobian(::typeof(logit), x::Real)
+    y = logit(x)
+    y, -log(x * (1 - x))
+end
+
+function ChangesOfVariables.with_logabsdet_jacobian(::typeof(logistic), x::Real)
+    y = logistic(x)
+    y, log(y * (1 - y))
+end

test/runtests.jl

@@ -1,5 +1,6 @@
 using LogExpFunctions
 using ChainRulesTestUtils
+using ChangesOfVariables
 using InverseFunctions
 using OffsetArrays
 
@@ -11,3 +12,4 @@ Random.seed!(1234)
 include("basicfuns.jl")
 include("chainrules.jl")
 include("inverse.jl")
+include("with_logabsdet_jacobian.jl")

test/with_logabsdet_jacobian.jl

@@ -0,0 +1,19 @@
+@testset "with_logabsdet_jacobian" begin
+    derivative(f, x) = ChainRulesTestUtils.frule((ChainRulesTestUtils.NoTangent(), 1), f, x)[2]
+
+    x = randexp()
+
+    ChangesOfVariables.test_with_logabsdet_jacobian(log1pexp, x, derivative)
+    ChangesOfVariables.test_with_logabsdet_jacobian(log1pexp, -x, derivative)
+
+    ChangesOfVariables.test_with_logabsdet_jacobian(logexpm1, x, derivative)
+
+    ChangesOfVariables.test_with_logabsdet_jacobian(log1mexp, -x, derivative)
+
+    ChangesOfVariables.test_with_logabsdet_jacobian(log2mexp, log(2) - x, derivative)
+
+    ChangesOfVariables.test_with_logabsdet_jacobian(logistic, -x, derivative)
+    ChangesOfVariables.test_with_logabsdet_jacobian(logistic, x, derivative)
+
+    ChangesOfVariables.test_with_logabsdet_jacobian(logit, rand(), derivative)
+end