Add test_with_logabsdet_jacobian

Project.toml

@@ -4,14 +4,14 @@ version = "0.1.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 julia = "1"
 
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Documenter", "ForwardDiff", "Test"]
+test = ["Documenter", "ForwardDiff"]

docs/Project.toml

@@ -1,5 +1,7 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 
 [compat]

docs/src/api.md

@@ -5,3 +5,9 @@
 ```@docs
 with_logabsdet_jacobian
 ```
+
+## Test utility
+
+```@docs
+ChangesOfVariables.test_with_logabsdet_jacobian
+```

src/ChangesOfVariables.jl

@@ -10,7 +10,9 @@ transformations).
 module ChangesOfVariables
 
 using LinearAlgebra
+using Test
 
 include("with_ladj.jl")
+include("test.jl")
 
 end # module

src/test.jl

@@ -0,0 +1,48 @@
+# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).
+
+
+"""
+    ChangesOfVariables.test_with_logabsdet_jacobian(f, x, getjacobian; compare = isapprox, kwargs...)
+
+Test if [`with_logabsdet_jacobian(f, x)`](@ref) is implemented correctly.
+
+Checks if the result of `with_logabsdet_jacobian(f, x)` is approximately
+equal to `(f(x), logabsdet(getjacobian(f, x)))`
+
+So the test uses `getjacobian(f, x)` to calculate a reference Jacobian for
+`f` at `x`. Passing `ForwardDiff.jabobian`, `Zygote.jacobian` or similar as
+the `getjacobian` function will do fine in most cases. If input and output
+of `f` are real scalar values, use `ForwardDiff.derivative`.
+
+Note that the result of `getjacobian(f, x)` must be a real-valued matrix
+or a real scalar, so you may need to use a custom `getjacobian` function
+that transforms the shape of `x` and `f(x)` internally, in conjunction
+with automatic differentiation.
+
+`kwargs...` are forwarded to `compare`.
+"""
+function test_with_logabsdet_jacobian(f, x, getjacobian; compare = isapprox, kwargs...)
+    @testset "test_with_logabsdet_jacobian: $f with input $x" begin
+        ref_y, test_type_inference = try
+            @inferred(f(x)), true
+        catch err
+            f(x), false
+        end
+
+        y, ladj = if test_type_inference
+            @inferred with_logabsdet_jacobian(f, x)
+        else
+            with_logabsdet_jacobian(f, x)
+        end
+
+        ref_ladj = _generalized_logabsdet(getjacobian(f, x))[1]
+
+        @test compare(y, ref_y; kwargs...)
+        @test compare(ladj, ref_ladj; kwargs...)
+    end
+    return nothing
+end
+
+
+_generalized_logabsdet(A) = logabsdet(A)
+_generalized_logabsdet(x::Real) = log(abs(x))

src/with_ladj.jl

@@ -14,9 +14,11 @@ For `(y, ladj) = with_logabsdet_jacobian(f, x)`, the following must hold true:
 `with_logabsdet_jacobian` comes with support for broadcasted/mapped functions
 (via `Base.Fix1`) and (Julia >=v1.6 only) `ComposedFunction`.
 
-Example:
+# Examples
+
+```jldoctest a
+using ChangesOfVariables
 
-```julia
 foo(x) = inv(exp(-x) + 1)
 
 function ChangesOfVariables.with_logabsdet_jacobian(::typeof(foo), x)
@@ -28,15 +30,39 @@ end
 x = 4.2
 y, ladj_y = with_logabsdet_jacobian(foo, x)
 
+using LinearAlgebra, ForwardDiff
+y == foo(x) && ladj_y ≈ log(abs(ForwardDiff.derivative(foo, x)))
+
+# output
+
+true
+```
+
+```jldoctest a
 X = rand(10)
 broadcasted_foo = Base.Fix1(broadcast, foo)
 Y, ladj_Y = with_logabsdet_jacobian(broadcasted_foo, X)
+Y == broadcasted_foo(X) && ladj_Y ≈ logabsdet(ForwardDiff.jacobian(broadcasted_foo, X))[1]
+
+# output
 
-# Requires Julia >= v1.6:
-z, ladj_z = with_logabsdet_jacobian(log ∘ foo, x)
-z == log(foo(x))
-ladj_z == ladj_y + with_logabsdet_jacobian(log, y)[2]
+true
 ```
+
+```jldoctest a
+VERSION < v"1.6" || begin # Support for ∘ requires Julia >= v1.6
+    z, ladj_z = with_logabsdet_jacobian(log ∘ foo, x)
+    z == log(foo(x)) && ladj_z == ladj_y + with_logabsdet_jacobian(log, y)[2]
+end
+
+# output
+
+true
+```
+
+Implementations of with_logabsdet_jacobian can be tested (as a
+`Test.@testset`) using
+[`ChangesOfVariables.test_with_logabsdet_jacobian`](@ref).
 """
 function with_logabsdet_jacobian end
 export with_logabsdet_jacobian

test/getjacobian.jl

@@ -0,0 +1,36 @@
+# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).
+
+import ForwardDiff
+
+torv_and_back(V::AbstractVector{<:Real}) = V, identity
+torv_and_back(x::Real) = [x], V -> V[1]
+torv_and_back(x::Complex) = [real(x), imag(x)], V -> Complex(V[1], V[2])
+torv_and_back(x::NTuple{N}) where N = [x...], V -> ntuple(i -> V[i], Val(N))
+
+function torv_and_back(x::Ref)
+    xval = x[]
+    V, to_xval = torv_and_back(xval)
+    back_to_ref(V) = Ref(to_xval(V))
+    return (V, back_to_ref)
+end
+
+torv_and_back(A::AbstractArray{<:Real}) = vec(A), V -> reshape(V, size(A))
+
+function torv_and_back(A::AbstractArray{Complex{T}, N}) where {T<:Real, N}
+    RA = cat(real.(A), imag.(A), dims = N+1) 
+    V, to_array = torv_and_back(RA)
+    function back_to_complex(V)
+        RA = to_array(V)
+        Complex.(view(RA, map(_ -> :, size(A))..., 1), view(RA, map(_ -> :, size(A))..., 2))
+    end
+    return (V, back_to_complex)
+end
+
+
+function getjacobian(f, x)
+    V, to_x = torv_and_back(x)
+    vf(V) = torv_and_back(f(to_x(V)))[1]
+    ForwardDiff.jacobian(vf, V)
+end
+
+foo(x) = inv(exp(-x) + 1)

test/runtests.jl

@@ -5,6 +5,7 @@ import ChangesOfVariables
 import Documenter
 
 Test.@testset "Package ChangesOfVariables" begin
+    include("test_test.jl")
     include("test_with_ladj.jl")
 
     # doctests
@@ -16,3 +17,4 @@ Test.@testset "Package ChangesOfVariables" begin
     )
     Documenter.doctest(ChangesOfVariables)
 end # testset
+

test/test_test.jl

@@ -0,0 +1,42 @@
+# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).
+
+using ChangesOfVariables
+using Test
+
+using LinearAlgebra
+
+using ChangesOfVariables: test_with_logabsdet_jacobian
+
+include("getjacobian.jl")
+
+
+@testset "torv_and_back" begin
+    for x in (rand(3), 0.5, Complex(0.2,0.7), (3,5,9), Ref(42), rand(3, 4, 5), Complex.(rand(3,5), rand(3,5)))
+        V, to_x = torv_and_back(x)
+        @test V isa AbstractVector{<:Real}
+        @test V == torv_and_back(x)[1]
+        @test x isa Ref ? to_x(V)[] == x[] : to_x(V) == x
+    end
+end
+
+
+@testset "test_with_logabsdet_jacobian" begin
+    rx = 0.5
+    cx = Complex(0.2, -0.7)
+    X = rand(3, 3)
+    CX = Complex.(randn(3,3), randn(3,3))
+
+    myisapprox(a, b; kwargs...) = isapprox(a, b; kwargs...)
+
+    noninferrable_inv(x) = x!=rand(size(x)...) ? inv(x) : ""
+    ChangesOfVariables.with_logabsdet_jacobian(::typeof(noninferrable_inv), x) = noninferrable_inv(x), with_logabsdet_jacobian(inv, x)[2]
+    @test_throws ErrorException @inferred with_logabsdet_jacobian(noninferrable_inv, rand(2, 2))
+
+    test_with_logabsdet_jacobian(inv, rx, ForwardDiff.derivative, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, cx, getjacobian, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, CX, getjacobian, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, CX, getjacobian, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, CX, getjacobian, compare = myisapprox, atol = 10^-6)
+    test_with_logabsdet_jacobian(noninferrable_inv, CX, getjacobian, atol = 10^-6)
+end

test/test_with_ladj.jl

@@ -4,66 +4,53 @@ using ChangesOfVariables
 using Test
 
 using LinearAlgebra
-using ForwardDiff: derivative, jacobian
 
+using ChangesOfVariables: test_with_logabsdet_jacobian
 
-fwddiff_ladj(f, x::Real) = log(abs(derivative(f, x)))
-fwddiff_ladj(f, x::AbstractArray{<:Real}) = logabsdet(jacobian(f, x))[1]
-fwddiff_with_ladj(f, x) = (f(x), fwddiff_ladj(f, x))
-
-ascomplex(A::AbstractArray{T}) where T = reinterpret(Complex{T}, A)
-asreal(A::AbstractArray{Complex{T}}) where T = reinterpret(T, A)
-
-isaprx(a, b) = isapprox(a,b)
-isaprx(a::NTuple{N,Any}, b::NTuple{N,Any}) where N = all(map(isaprx, a, b))
-
-
-foo(x) = inv(exp(-x) + 1)
-
-function ChangesOfVariables.with_logabsdet_jacobian(::typeof(foo), x)
-    y = foo(x)
-    ladj = -x + 2 * log(y)
-    (y, ladj)
-end
+include("getjacobian.jl")
 
 
 @testset "with_logabsdet_jacobian" begin
+    function ChangesOfVariables.with_logabsdet_jacobian(::typeof(foo), x)
+        y = foo(x)
+        ladj = -x + 2 * log(y)
+        (y, ladj)
+    end
+
     x = 4.2
     X = rand(10)
     A = rand(5, 5)
-    CA = rand(10, 5)
+    CA = Complex.(rand(5, 5), rand(5, 5))
+
+    isaprx(a, b; kwargs...) = isapprox(a,b; kwargs...)
+    isaprx(a::NTuple{N,Any}, b::NTuple{N,Any}; kwargs...) where N = all(map((a,b) -> isaprx(a, b; kwargs...), a, b))
 
-    @test isaprx(@inferred(with_logabsdet_jacobian(foo, x)), fwddiff_with_ladj(foo, x))
+
+    test_with_logabsdet_jacobian(foo, x, getjacobian)
 
     @static if VERSION >= v"1.6"
-        log_foo = log ∘ foo
-        @test isaprx(@inferred(with_logabsdet_jacobian(log_foo, x)), fwddiff_with_ladj(log_foo, x))
+        test_with_logabsdet_jacobian(log ∘ foo, x, getjacobian)
     end
 
-    mapped_foo = Base.Fix1(map, foo)
-    @test isaprx(@inferred(with_logabsdet_jacobian(mapped_foo, x)), fwddiff_with_ladj(mapped_foo, x))
-    @test isaprx(@inferred(with_logabsdet_jacobian(mapped_foo, fill(x))), fwddiff_with_ladj(mapped_foo, fill(x)))
-    @test isaprx(@inferred(with_logabsdet_jacobian(mapped_foo, Ref(x))), fwddiff_with_ladj(mapped_foo, fill(x)))
-    @test isaprx(@inferred(with_logabsdet_jacobian(mapped_foo, (x,))), (mapped_foo((x,)), fwddiff_ladj(mapped_foo, x)))
-    @test isaprx(@inferred(with_logabsdet_jacobian(mapped_foo, X)), fwddiff_with_ladj(mapped_foo, X))
-
-    broadcasted_foo = Base.Fix1(broadcast, foo)
-    @test isaprx(@inferred(with_logabsdet_jacobian(broadcasted_foo, x)), fwddiff_with_ladj(broadcasted_foo, x))
-    @test isaprx(@inferred(with_logabsdet_jacobian(broadcasted_foo, fill(x))), fwddiff_with_ladj(broadcasted_foo, x))
-    @test isaprx(@inferred(with_logabsdet_jacobian(broadcasted_foo, Ref(x))), fwddiff_with_ladj(broadcasted_foo, x))
-    @test isaprx(@inferred(with_logabsdet_jacobian(broadcasted_foo, (x,))), (mapped_foo((x,)), fwddiff_ladj(mapped_foo, x)))
-    @test isaprx(@inferred(with_logabsdet_jacobian(broadcasted_foo, X)), fwddiff_with_ladj(broadcasted_foo, X))
+    @testset "getjacobian on mapped and broadcasted" begin
+        for f in (Base.Fix1(map, foo), Base.Fix1(broadcast, foo))
+            for arg in (x, fill(x,), Ref(x), (x,), X)
+                test_with_logabsdet_jacobian(f, arg, getjacobian, compare = isaprx)
+            end
+        end
+    end
 
-    for f in (identity, adjoint, transpose)
-        @test isaprx(@inferred(with_logabsdet_jacobian(f, x)), fwddiff_with_ladj(f, x))
-        @test isaprx(@inferred(with_logabsdet_jacobian(f, A)), fwddiff_with_ladj(f, A))
+    @testset "getjacobian on identity, adjoint and transpose" begin
+        for f in (identity, adjoint, transpose)
+            for arg in (x, A)
+                test_with_logabsdet_jacobian(f, arg, getjacobian)
+            end
+        end
     end
-
-    @test isaprx(@inferred(with_logabsdet_jacobian(inv, x)), fwddiff_with_ladj(inv, x))
-    @test isaprx(@inferred(with_logabsdet_jacobian(inv, A)), fwddiff_with_ladj(inv, A))
-    @test isaprx(@inferred(with_logabsdet_jacobian(inv, ascomplex(CA))), (inv(ascomplex(CA)), fwddiff_ladj(CA -> asreal(inv(ascomplex(CA))), CA)))
 
-    for f in (exp, log, exp2, log2, exp10, log10, expm1, log1p)
-        @test isaprx(@inferred(with_logabsdet_jacobian(f, x)), fwddiff_with_ladj(f, x))
+    @testset "getjacobian on inv" begin
+        for arg in (x, A, CA)
+            test_with_logabsdet_jacobian(inv, arg, getjacobian)
+        end
     end
 end