Simplify test_with_logabsdet_jacobian

Author: oschulz

src/test.jl

@@ -2,10 +2,7 @@
 
 
 """
-    ChangesOfVariables.test_with_logabsdet_jacobian(
-        f, x, getjacobian, rv_and_back = x -> (x, identity);
-        compare = isapprox, kwargs...
-    )
+    ChangesOfVariables.test_with_logabsdet_jacobian(f, x, getjacobian; compare = isapprox, kwargs...)
 
 Test if [`with_logabsdet_jacobian(f, x)`](@ref) is implemented correctly.
 
@@ -17,25 +14,14 @@ So the test uses `getjacobian(f, x)` to calculate a reference Jacobian for
 the `getjacobian` function will do fine in most cases. If input and output
 of `f` are real scalar values, use `ForwardDiff.derivative`.
 
-If `getjacobian(f, x)` can't handle the type of `x` or `f(x)` because they
-are not real-valued vectors, use the `rv_and_back` argument to pass a
-function with the following behavior
-
-```julia
-v, back = rv_and_back(x)
-v isa AbstractVector{<:Real}
-back(v) == x
-```
-
-If `test_inferred == true`, type inference on `with_logabsdet_jacobian` will
-be tested.
+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, rv_and_back = x -> (x, identity);
-    compare = isapprox, kwargs...
-)
+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
@@ -49,9 +35,7 @@ function test_with_logabsdet_jacobian(
             with_logabsdet_jacobian(f, x)
         end
 
-        V, to_x = rv_and_back(x)
-        vf(V) = rv_and_back(f(to_x(V)))[1]
-        ref_ladj = _generalized_logabsdet(getjacobian(vf, V))[1]
+        ref_ladj = _generalized_logabsdet(getjacobian(f, x))[1]
 
         @test compare(y, ref_y; kwargs...)
         @test compare(ladj, ref_ladj; kwargs...)

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/rv_and_back.jl

@@ -4,18 +4,17 @@ using ChangesOfVariables
 using Test
 
 using LinearAlgebra
-import ForwardDiff
 
 using ChangesOfVariables: test_with_logabsdet_jacobian
 
+include("getjacobian.jl")
 
-include("rv_and_back.jl")
 
-@testset "rv_and_back" begin
+@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 = rv_and_back(x)
+        V, to_x = torv_and_back(x)
         @test V isa AbstractVector{<:Real}
-        @test V == rv_and_back(x)[1]
+        @test V == torv_and_back(x)[1]
         @test x isa Ref ? to_x(V)[] == x[] : to_x(V) == x
     end
 end
@@ -34,10 +33,10 @@ end
     @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, ForwardDiff.jacobian, rv_and_back, 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, ForwardDiff.jacobian, rv_and_back, atol = 10^-6)
-    test_with_logabsdet_jacobian(inv, CX, ForwardDiff.jacobian, rv_and_back, atol = 10^-6)
-    test_with_logabsdet_jacobian(inv, CX, ForwardDiff.jacobian, rv_and_back, compare = myisapprox, atol = 10^-6)
-    test_with_logabsdet_jacobian(noninferrable_inv, CX, ForwardDiff.jacobian, rv_and_back, 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_test.jl

@@ -4,60 +4,53 @@ using ChangesOfVariables
 using Test
 
 using LinearAlgebra
-import ForwardDiff
 
 using ChangesOfVariables: test_with_logabsdet_jacobian
 
-include("rv_and_back.jl")
+include("getjacobian.jl")
 
 
 @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
 
-
     x = 4.2
     X = rand(10)
     A = rand(5, 5)
     CA = Complex.(rand(5, 5), rand(5, 5))
 
-
-    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, rv_and_back)
+    test_with_logabsdet_jacobian(foo, x, getjacobian)
 
     @static if VERSION >= v"1.6"
-        test_with_logabsdet_jacobian(log ∘ foo, x, getjacobian, rv_and_back)
+        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, rv_and_back, compare = isaprx)
+                test_with_logabsdet_jacobian(f, arg, getjacobian, compare = isaprx)
             end
         end
     end
 
     @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, rv_and_back)
+                test_with_logabsdet_jacobian(f, arg, getjacobian)
             end
         end
     end
 
     @testset "getjacobian on inv" begin
         for arg in (x, A, CA)
-            test_with_logabsdet_jacobian(inv, arg, getjacobian, rv_and_back)
+            test_with_logabsdet_jacobian(inv, arg, getjacobian)
         end
     end
 end