Rewrote testing code

Author: theogf

test/utils_AD.jl

@@ -1,7 +1,9 @@
 
-FDM = FiniteDifferences.central_fdm(5, 1)
+const FDM = FiniteDifferences.central_fdm(5, 1)
 
-function gradient(::Val{:Zygote}, f::Function, args)
+gradient(f, s::Symbol, args) = gradient(f, Val(s), args)
+
+function gradient(f, ::Val{:Zygote}, args)
     g = first(Zygote.gradient(f, args))
     if isnothing(g)
         if args isa AbstractArray{<:Real}
@@ -14,18 +16,21 @@ function gradient(::Val{:Zygote}, f::Function, args)
     end
 end
 
-function gradient(::Val{:ForwardDiff}, f::Function, args)
+function gradient(f, ::Val{:ForwardDiff}, args)
     ForwardDiff.gradient(f, args)
 end
 
-function gradient(::Val{:ReverseDiff}, f::Function, args)
+function gradient(f, ::Val{:ReverseDiff}, args)
     ReverseDiff.gradient(f, args)
 end
 
-function gradient(::Val{:FiniteDiff}, f::Function, args)
+function gradient(f, ::Val{:FiniteDiff}, args)
     first(FiniteDifferences.grad(FDM, f, args))
 end
 
+function compare_gradient(f, AD::Symbol, args)
+    isapprox(gradient(f, AD, args), gradient(f, :FiniteDiff, args), atol=1e-8, rtol=1e-5)
+end
 
 testfunction(k, A, B, dim) = sum(kernelmatrix(k, A, B, obsdim = dim))
 testfunction(k, A, dim) = sum(kernelmatrix(k, A, obsdim = dim))
@@ -50,25 +55,39 @@ function test_FiniteDiff(kernelfunction, args = nothing, dims = [3, 3])
     @testset "FiniteDifferences" begin
         if k isa SimpleKernel
             for d in log.([eps(), rand(rng)])
-                @test_nowarn gradient(Val(:FiniteDiff), x -> kappa(k, exp(first(x))), [d])
+                @test_nowarn gradient(:FiniteDiff, [d]) do x
+                    kappa(k, exp(first(x)))
+                end
             end
         end
         ## Testing Kernel Functions
         x = rand(rng, dims[1])
         y = rand(rng, dims[1])
-        @test_nowarn gradient(Val(:FiniteDiff), x -> k(x, y), x)
+        @test_nowarn gradient(:FiniteDiff, x) do x
+                k(x, y)
+            end
         if !(args === nothing)
-            @test_nowarn gradient(Val(:FiniteDiff), p -> kernelfunction(p)(x, y), args)
+            @test_nowarn gradient(:FiniteDiff, args) do p
+                kernelfunction(p)(x, y)
+            end
         end
         ## Testing Kernel Matrices
         A = rand(rng, dims...)
         B = rand(rng, dims...)
         for dim in 1:2
-            @test_nowarn gradient(Val(:FiniteDiff), a -> testfunction(k, a, dim), A)
-            @test_nowarn gradient(Val(:FiniteDiff), a -> testfunction(k, a, B, dim), A)
-            @test_nowarn gradient(Val(:FiniteDiff), b -> testfunction(k, A, b, dim), B)
+            @test_nowarn gradient(:FiniteDiff, A) do a
+                testfunction(k, a, dim)
+            end
+            @test_nowarn gradient(:FiniteDiff , A) do a
+                testfunction(k, a, B, dim)
+            end
+            @test_nowarn gradient(:FiniteDiff, B) do b
+                testfunction(k, A, b, dim)
+            end
             if !(args === nothing)
-                @test_nowarn gradient(Val(:FiniteDiff), p -> testfunction(kernelfunction(p), A, B, dim), args)
+                @test_nowarn gradient(:FiniteDiff, args) do p
+                    testfunction(kernelfunction(p), A, B, dim)
+                end
             end
         end
     end
@@ -85,26 +104,42 @@ function test_AD(AD::Symbol, kernelfunction, args = nothing, dims = [3, 3])
         rng = MersenneTwister(42)
         if k isa SimpleKernel
             for d in log.([eps(), rand(rng)])
-                @test gradient(Val(AD), x -> kappa(k, exp(x[1])), [d]) ≈ gradient(Val(:FiniteDiff), x -> kappa(k, exp(x[1])), [d]) atol=1e-8 rtol=1e-5
+                @test compare_gradient(AD, [d]) do x
+                    kappa(k, exp(x[1])
+                end
             end
         end
         # Testing kernel evaluations
         x = rand(rng, dims[1])
         y = rand(rng, dims[1])
-        @test gradient(Val(AD), x -> k(x, y), x) ≈ gradient(Val(:FiniteDiff), x -> k(x, y), x) atol=1e-8 rtol=1e-5
-        @test gradient(Val(AD), y -> k(x, y), y) ≈ gradient(Val(:FiniteDiff), y -> k(x, y), y) atol=1e-8 rtol=1e-5
+        @test compare_gradient(AD, x) do x
+            k(x, y)
+        end
+        @test compare_gradient(AD, y) do y
+            k(x, y)
+        end
         if !(args === nothing)
-            @test gradient(Val(AD), p -> kernelfunction(p)(x,y), args) ≈ gradient(Val(:FiniteDiff), p -> kernelfunction(p)(x, y), args) atol=1e-8 rtol=1e-5
+            @test compare_gradient(AD, args) do p
+                kernelfunction(p)(x,y)
+            end
         end
         # Testing kernel matrices
         A = rand(rng, dims...)
         B = rand(rng, dims...)
         for dim in 1:2
-            @test gradient(Val(AD), x -> testfunction(k, x, dim), A) ≈ gradient(Val(:FiniteDiff), x -> testfunction(k, x, dim), A) atol=1e-8 rtol=1e-5
-            @test gradient(Val(AD), a -> testfunction(k, a, B, dim), A) ≈ gradient(Val(:FiniteDiff), a -> testfunction(k, a, B, dim), A) atol=1e-8 rtol=1e-5
-            @test gradient(Val(AD), b -> testfunction(k, A, b, dim), B) ≈ gradient(Val(:FiniteDiff), b -> testfunction(k, A, b, dim), B) atol=1e-8 rtol=1e-5
+            @test compare_gradient(AD, A) do a
+                testfunction(k, a, dim)
+            end
+            @test conpare_gradient(AD, A) do a
+                testfunction(k, a, B, dim)
+            end
+            @test compare_gradient(AD, B) do b
+                testfunction(k, A, b, dim)
+            end
             if !(args === nothing)
-                @test gradient(Val(AD), p -> testfunction(kernelfunction(p), A, dim), args) ≈ gradient(Val(:FiniteDiff), p -> testfunction(kernelfunction(p), A, dim), args) atol=1e-8 rtol=1e-5
+                @test compare_gradient(AD, args) do p
+                    testfunction(kernelfunction(p), AD, A, dim)
+                end
             end
         end
     end