diff --git a/Project.toml b/Project.toml index 8ca9df8..9e363ce 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ version = "0.1.0" [deps] LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] julia = "1" @@ -11,7 +12,6 @@ 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"] diff --git a/src/ChangesOfVariables.jl b/src/ChangesOfVariables.jl index 9315fb7..b2b034d 100644 --- a/src/ChangesOfVariables.jl +++ b/src/ChangesOfVariables.jl @@ -10,7 +10,9 @@ transformations). module ChangesOfVariables using LinearAlgebra +using Test include("with_ladj.jl") +include("test.jl") end # module diff --git a/src/test.jl b/src/test.jl new file mode 100644 index 0000000..d7fe106 --- /dev/null +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 3c63814..035b048 100644 --- a/test/runtests.jl +++ b/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 + diff --git a/test/test_test.jl b/test/test_test.jl new file mode 100644 index 0000000..dee7698 --- /dev/null +++ b/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, 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) +end diff --git a/test/test_with_ladj.jl b/test/test_with_ladj.jl index 4c210ea..115d57f 100644 --- a/test/test_with_ladj.jl +++ b/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