tests

Author: vpuri3

test/runtests.jl

@@ -22,8 +22,7 @@ if GROUP == "All"
     @time @safetestset "Integration test" begin include("test_integration.jl") end
     @time @safetestset "Special matrices" begin include("test_specialmatrices.jl") end
     @time @safetestset "Jac Vecs and Hes Vecs" begin include("test_jaches_products.jl") end
-    @time @safetestset "Operator tests" begin include("test_ops.jl") end
-
+    @time @safetestset "Vec Jac Products" begin include("test_vecjac_products.jl") end
 end
 
 if GROUP == "GPU"

test/test_jaches_products.jl

@@ -3,12 +3,13 @@ using LinearAlgebra, Test
 
 using Random
 Random.seed!(123)
-
-const A = rand(300, 300)
+N = 300
+const A = rand(N, N)
 f(y, x) = mul!(y, A, x)
 f(x) = A * x
-x = rand(300)
-v = rand(300)
+x = rand(N)
+v = rand(N)
+a, b = rand(2)
 dy = similar(x)
 g(x) = sum(abs2, x)
 function h(x)
@@ -20,8 +21,7 @@ end
 
 cache1 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag(), eltype(x))),
                           eltype(x), 1}.(x, ForwardDiff.Partials.(Tuple.(v)))
-cache2 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag(), eltype(x))),
-                          eltype(x), 1}.(x, ForwardDiff.Partials.(Tuple.(v)))
+cache2 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag(), eltype(x))), eltype(x), 1}.(x, ForwardDiff.Partials.(Tuple.(v)))
 @test num_jacvec!(dy, f, x, v)≈ForwardDiff.jacobian(f, similar(x), x) * v rtol=1e-6
 @test num_jacvec!(dy, f, x, v, similar(v),
                   similar(v))≈ForwardDiff.jacobian(f, similar(x), x) * v rtol=1e-6
@@ -65,61 +65,90 @@ cache4 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(Nothing, eltype(x))), eltype(x)
 @test auto_hesvecgrad!(dy, h, x, v, cache1, cache2)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
 @test auto_hesvecgrad(h, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
 
+# JacVec
+
 L = JacVec(f, x)
 @test L * x ≈ auto_jacvec(f, x, x)
 @test L * v ≈ auto_jacvec(f, x, v)
 @test mul!(dy, L, v) ≈ auto_jacvec(f, x, v)
-L.x .= v
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*auto_jacvec(f,x,v) + b*_dy
+update_coefficients!(L, v, nothing, 0.0)
 @test mul!(dy, L, v) ≈ auto_jacvec(f, v, v)
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*auto_jacvec(f,x,v) + b*_dy
 
 L = JacVec(f, x, autodiff = false)
 @test L * x ≈ num_jacvec(f, x, x)
 @test L * v ≈ num_jacvec(f, x, v)
-L.x == x
 @test mul!(dy, L, v)≈num_jacvec(f, x, v) rtol=1e-6
-L.x .= v
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*num_jacvec(f,x,v) + b*_dy rtol=1e-6
+update_coefficients!(L, v, nothing, 0.0)
 @test mul!(dy, L, v)≈num_jacvec(f, v, v) rtol=1e-6
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*num_jacvec(f,x,v) + b*_dy rtol=1e-6
 
-### Integration test with IterativeSolvers
 out = similar(v)
 gmres!(out, L, v)
 
-x = rand(300)
-v = rand(300)
+#=
+ff1 = ODEFunction(lorenz, jac_prototype = JacVec{Float64}(lorenz, u0))
+ff2 = ODEFunction(lorenz, jac_prototype = JacVec{Float64}(lorenz, u0, autodiff=false))
+
+for ff in [ff1, ff2]
+    prob = ODEProblem(ff, u0, tspan)
+    @test solve(prob, TRBDF2()).retcode == :Success
+    @test solve(prob, TRBDF2(linsolve = KrylovJL_GMRES())).retcode == :Success
+    @test solve(prob, Exprb32()).retcode == :Success
+    @test solve(prob, Rosenbrock23()).retcode == :Success
+    @test solve(prob, Rosenbrock23(linsolve = KrylovJL_GMRES())).retcode == :Success
+end
+=#
+
+# HesVec
+
+x = rand(N)
+v = rand(N)
 L = HesVec(g, x, autodiff = false)
 @test L * x ≈ num_hesvec(g, x, x)
 @test L * v ≈ num_hesvec(g, x, v)
 @test mul!(dy, L, v)≈num_hesvec(g, x, v) rtol=1e-2
-L.x .= v
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*num_hesvec(g,x,v) + b*_dy rtol=1e-2
+update_coefficients!(L, v, nothing, 0.0)
 @test mul!(dy, L, v)≈num_hesvec(g, v, v) rtol=1e-2
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*num_hesvec(g,x,v) + b*_dy rtol=1e-2
 
 L = HesVec(g, x)
 @test L * x ≈ numauto_hesvec(g, x, x)
 @test L * v ≈ numauto_hesvec(g, x, v)
 @test mul!(dy, L, v)≈numauto_hesvec(g, x, v) rtol=1e-8
-L.x .= v
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy rtol=1e-8
+update_coefficients!(L, v, nothing, 0.0)
 @test mul!(dy, L, v)≈numauto_hesvec(g, v, v) rtol=1e-8
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy rtol=1e-8
 
-### Integration test with IterativeSolvers
 out = similar(v)
 gmres!(out, L, v)
 
-x = rand(300)
-v = rand(300)
+# HesVecGrad
+
+x = rand(N)
+v = rand(N)
 L = HesVecGrad(h, x, autodiff = false)
 @test L * x ≈ num_hesvec(g, x, x)
 @test L * v ≈ num_hesvec(g, x, v)
 @test mul!(dy, L, v)≈num_hesvec(g, x, v) rtol=1e-2
-L.x .= v
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*num_hesvec(g,x,v)+b*_dy rtol=1e-2
+update_coefficients!(L, v, nothing, 0.0)
 @test mul!(dy, L, v)≈num_hesvec(g, v, v) rtol=1e-2
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*num_hesvec(g,x,v)+b*_dy rtol=1e-2
 
 L = HesVecGrad(h, x, autodiff = true)
 @test L * x ≈ autonum_hesvec(g, x, x)
 @test L * v ≈ numauto_hesvec(g, x, v)
 @test mul!(dy, L, v)≈numauto_hesvec(g, x, v) rtol=1e-8
-L.x .= v
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy rtol=1e-8
+update_coefficients!(L, v, nothing, 0.0)
 @test mul!(dy, L, v)≈numauto_hesvec(g, v, v) rtol=1e-8
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy rtol=1e-8
 
-### Integration test with IterativeSolvers
 out = similar(v)
 gmres!(out, L, v)
+#

test/test_ops.jl

@@ -0,0 +1,27 @@
+using SparseDiffTools, ForwardDiff, FiniteDiff, Zygote, IterativeSolvers
+using LinearAlgebra, Test
+
+using Random
+Random.seed!(123)
+N = 300
+const A = rand(N, N)
+a, b = rand(2)
+
+x = rand(Float32, N)
+v = rand(Float32, N)
+
+f(du,u,p,t) = mul!(du, A, u)
+f(u,p,t) = A * u
+
+# VecJac
+
+L = VecJac(f, x)
+actual_vjp = Zygote.jacobian(x -> f(x, nothing, 0.0), x)[1]' * v
+update_coefficients!(L, v, nothing, 0.0)
+@test L * v ≈ actual_vjp
+L = VecJac(f, x; autodiff = false)
+update_coefficients!(L, v, nothing, 0.0)
+@test L * v ≈ actual_vjp
+#dy=rand(N); @test mul!(dy, L, v) ≈ actual_vjp
+#dy=rand(N); _dy=copy(dy); @test mul!(dy,L,v,a,b) ≈ a * actual_vjp + b * _dy
+#