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"]

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,78 @@
+# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).
+
+
+_to_realvec_and_back(V::AbstractVector{<:Real}) = V, identity
+_to_realvec_and_back(x::Real) = [x], V -> V[1]
+_to_realvec_and_back(x::Complex) = [real(x), imag(x)], V -> Complex(V[1], V[2])
+_to_realvec_and_back(x::NTuple{N}) where N = [x...], V -> ntuple(i -> V[i], Val(N))
+
+function _to_realvec_and_back(x::Ref)
+    xval = x[]
+    V, to_xval = _to_realvec_and_back(xval)
+    back_to_ref(V) = Ref(to_xval(V))
+    return (V, back_to_ref)
+end
+
+_to_realvec_and_back(A::AbstractArray{<:Real}) = vec(A), V -> reshape(V, size(A))
+
+function _to_realvec_and_back(A::AbstractArray{Complex{T}, N}) where {T<:Real, N}
+    RA = cat(real.(A), imag.(A), dims = N+1) 
+    V, to_array = _to_realvec_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
+
+
+_to_realvec(x) = _to_realvec_and_back(x)[1]
+
+
+function _auto_with_logabsdet_jacobian(f, x, getjacobian)
+    y = f(x)
+    V, to_x = _to_realvec_and_back(x)
+    vf(V) = _to_realvec(f(to_x(V)))
+    ladj = logabsdet(getjacobian(vf, V))[1]
+    return (y, ladj)
+end
+
+
+"""
+    ChangesOfVariables.test_with_logabsdet_jacobian(
+        f, x, getjacobian;
+        test_inferred::Bool = true, 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 `x` or `f(x)` are real-valued scalars or complex-valued scalars or arrays,
+the test will try to reshape them automatically, to account for limitations
+of (e.g.) `ForwardDiff` and to ensure the result of `getjacobian` is a real
+matrix.
+
+If `test_inferred == true` will test type inference on
+`with_logabsdet_jacobian`.
+
+`kwargs...` are forwarded to `isapprox`.
+"""
+function test_with_logabsdet_jacobian(f, x, getjacobian; compare=isapprox, test_inferred::Bool = true, kwargs...)
+    @testset "test_with_logabsdet_jacobian: $f with input $x" begin
+        y, ladj = if test_inferred
+            @inferred with_logabsdet_jacobian(f, x)
+        else
+            with_logabsdet_jacobian(f, x)
+        end
+        ref_y, ref_ladj = _auto_with_logabsdet_jacobian(f, x, getjacobian)
+        @test compare(y, ref_y; kwargs...)
+        @test compare(ladj, ref_ladj; kwargs...)
+    end
+    return nothing
+end

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,36 @@
+# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).
+
+using ChangesOfVariables
+using Test
+
+using LinearAlgebra
+import ForwardDiff
+
+using ChangesOfVariables: test_with_logabsdet_jacobian
+
+
+@testset "test_with_logabsdet_jacobian" begin
+    @testset "ChangesOfVariables._to_realvec_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 = ChangesOfVariables._to_realvec_and_back(x)
+            @test V isa AbstractVector{<:Real}
+            @test V == ChangesOfVariables._to_realvec(x)
+            @test x isa Ref ? to_x(V)[] == x[] : to_x(V) == x
+        end
+    end
+
+    x = Complex(0.2, -0.7)
+    y, ladj_y = ChangesOfVariables._auto_with_logabsdet_jacobian(inv, x, ForwardDiff.jacobian)
+    @test y == inv(x)
+    @test ladj_y ≈ -4 * log(abs(x))
+
+    X = Complex.(randn(3,3), randn(3,3))
+    Y, ladj_Y = ChangesOfVariables._auto_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian)
+    @test Y == inv(X)
+    @test ladj_Y == -4 * 3 * logabsdet(X)[1]
+
+    myisapprox(a, b; kwargs...) = isapprox(a, b; kwargs...)
+    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, test_inferred = true)
+    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, test_inferred = false)
+    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, compare = myisapprox, atol = 10^-5)
+end

test/test_with_ladj.jl

@@ -4,66 +4,58 @@ using ChangesOfVariables
 using Test
 
 using LinearAlgebra
-using ForwardDiff: derivative, jacobian
+import ForwardDiff
 
+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)
+@testset "with_logabsdet_jacobian" begin
+    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
+    function ChangesOfVariables.with_logabsdet_jacobian(::typeof(foo), x)
+        y = foo(x)
+        ladj = -x + 2 * log(y)
+        (y, ladj)
+    end
 
 
-@testset "with_logabsdet_jacobian" begin
     x = 4.2
     X = rand(10)
     A = rand(5, 5)
-    CA = rand(10, 5)
+    CA = Complex.(rand(5, 5), rand(5, 5))
 
-    @test isaprx(@inferred(with_logabsdet_jacobian(foo, x)), fwddiff_with_ladj(foo, x))
+
+    getjacobian = ForwardDiff.jacobian
+
+    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_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
+
+    @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
 
-    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))
-
-    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