Add rv_and_back argument to test_with_logabsdet_jacobian

Author: oschulz

src/test.jl

@@ -1,47 +1,22 @@
 # 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))
+_generalized_logabsdet(A) = logabsdet(A)
+_generalized_logabsdet(x::Real) = log(abs(x))
 
-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)
+function _auto_with_logabsdet_jacobian(f, x, getjacobian, rv_and_back)
     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]
+    V, to_x = rv_and_back(x)
+    vf(V) = rv_and_back(f(to_x(V)))[1]
+    ladj = _generalized_logabsdet(getjacobian(vf, V))[1]
     return (y, ladj)
 end
 
 
 """
     ChangesOfVariables.test_with_logabsdet_jacobian(
-        f, x, getjacobian;
-        test_inferred::Bool = true, kwargs...
+        f, x, getjacobian, rv_and_back = x -> (x, identity);
+        compare = isapprox, test_inferred::Bool = true, kwargs...
     )
 
 Test if [`with_logabsdet_jacobian(f, x)`](@ref) is implemented correctly.
@@ -51,26 +26,35 @@ 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.
+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` of `f(x)` because they
+are not real-valued vectors, use the `rv_and_back` argument to pass a
+function with the following behavior
 
-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.
+```julia
+v, back = rv_and_back(x)
+v isa AbstractVector{<:Real}
+back(v) == x
+```
 
-If `test_inferred == true` will test type inference on
-`with_logabsdet_jacobian`.
+If `test_inferred == true`, type inference on `with_logabsdet_jacobian` will
+be tested.
 
-`kwargs...` are forwarded to `isapprox`.
+`kwargs...` are forwarded to `compare`.
 """
-function test_with_logabsdet_jacobian(f, x, getjacobian; compare=isapprox, test_inferred::Bool = true, kwargs...)
+function test_with_logabsdet_jacobian(
+    f, x, getjacobian, rv_and_back = x -> (x, identity);
+    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)
+        ref_y, ref_ladj = _auto_with_logabsdet_jacobian(f, x, getjacobian, rv_and_back)
         @test compare(y, ref_y; kwargs...)
         @test compare(ladj, ref_ladj; kwargs...)
     end

test/rv_and_back.jl

@@ -0,0 +1,26 @@
+# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).
+
+
+rv_and_back(V::AbstractVector{<:Real}) = V, identity
+rv_and_back(x::Real) = [x], V -> V[1]
+rv_and_back(x::Complex) = [real(x), imag(x)], V -> Complex(V[1], V[2])
+rv_and_back(x::NTuple{N}) where N = [x...], V -> ntuple(i -> V[i], Val(N))
+
+function rv_and_back(x::Ref)
+    xval = x[]
+    V, to_xval = rv_and_back(xval)
+    back_to_ref(V) = Ref(to_xval(V))
+    return (V, back_to_ref)
+end
+
+rv_and_back(A::AbstractArray{<:Real}) = vec(A), V -> reshape(V, size(A))
+
+function rv_and_back(A::AbstractArray{Complex{T}, N}) where {T<:Real, N}
+    RA = cat(real.(A), imag.(A), dims = N+1) 
+    V, to_array = rv_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

test/test_test.jl

@@ -9,28 +9,33 @@ 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
+include("rv_and_back.jl")
+
+@testset "rv_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)
+        @test V isa AbstractVector{<:Real}
+        @test V == rv_and_back(x)[1]
+        @test x isa Ref ? to_x(V)[] == x[] : to_x(V) == x
     end
+end
 
+
+@testset "test_with_logabsdet_jacobian" begin
     x = Complex(0.2, -0.7)
-    y, ladj_y = ChangesOfVariables._auto_with_logabsdet_jacobian(inv, x, ForwardDiff.jacobian)
+    y, ladj_y = ChangesOfVariables._auto_with_logabsdet_jacobian(inv, x, ForwardDiff.jacobian, rv_and_back)
     @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)
+    Y, ladj_Y = ChangesOfVariables._auto_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, rv_and_back)
     @test Y == inv(X)
-    @test ladj_Y == -4 * 3 * logabsdet(X)[1]
+    @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, atol = 10^-6)
-    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, test_inferred = false, atol = 10^-6)
-    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, compare = myisapprox, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, 0.5, ForwardDiff.derivative, test_inferred = true, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, rand(2,2), ForwardDiff.jacobian, test_inferred = true, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, rv_and_back, test_inferred = true, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, rv_and_back, test_inferred = false, atol = 10^-6)
+    test_with_logabsdet_jacobian(inv, X, ForwardDiff.jacobian, rv_and_back, compare = myisapprox, atol = 10^-6)
 end

test/test_with_ladj.jl

@@ -8,6 +8,8 @@ import ForwardDiff
 
 using ChangesOfVariables: test_with_logabsdet_jacobian
 
+include("rv_and_back.jl")
+
 
 @testset "with_logabsdet_jacobian" begin
     foo(x) = inv(exp(-x) + 1)
@@ -31,31 +33,31 @@ using ChangesOfVariables: test_with_logabsdet_jacobian
     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)
+    test_with_logabsdet_jacobian(foo, x, getjacobian, rv_and_back)
 
     @static if VERSION >= v"1.6"
-        test_with_logabsdet_jacobian(log ∘ foo, x, getjacobian)
+        test_with_logabsdet_jacobian(log ∘ foo, x, getjacobian, rv_and_back)
     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)
+                test_with_logabsdet_jacobian(f, arg, getjacobian, rv_and_back, 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)
+                test_with_logabsdet_jacobian(f, arg, getjacobian, rv_and_back)
             end
         end
     end
 
     @testset "getjacobian on inv" begin
         for arg in (x, A, CA)
-            test_with_logabsdet_jacobian(inv, arg, getjacobian)
+            test_with_logabsdet_jacobian(inv, arg, getjacobian, rv_and_back)
         end
     end
 end