Add precs_simplified.jl test

Project.toml

@@ -12,18 +12,20 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
+SparsityTracing = "06eadbd4-12ad-4cbc-ab6e-10f8370940a5"
 Sundials_jll = "fb77eaff-e24c-56d4-86b1-d163f2edb164"
-SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 
 [compat]
 CEnum = "0.2, 0.3, 0.4"
 DataStructures = "0.18"
 DiffEqBase = "6.44"
 Reexport = "0.2, 1.0"
 SciMLBase = "1.71"
-Sundials_jll = "5.2"
 SnoopPrecompile = "1"
+Sundials_jll = "5.2"
 julia = "1.6"
 
 [extras]

src/nvector_wrapper.jl

@@ -63,7 +63,6 @@ Base.convert(::Type{NVector}, nv::N_Vector) = NVector(nv)
 Base.convert(::Type{Vector{realtype}}, nv::NVector) = nv.v
 Base.convert(::Type{Vector}, nv::NVector) = nv.v
 
-
 """ 
     Base.cconvert(::Type{N_Vector}, v::Vector{realtype}) -> NVector
     Base.unsafe_convert(::Type{N_Vector}, nv::NVector) -> N_Vector
@@ -79,7 +78,6 @@ Conversion happens in two steps within ccall:
 Base.cconvert(::Type{N_Vector}, v::Vector{realtype}) = convert(NVector, v) # will just return v if v is an NVector
 Base.unsafe_convert(::Type{N_Vector}, nv::NVector) = nv.n_v
 
-
 Base.similar(nv::NVector) = NVector(similar(nv.v))
 
 nvlength(x::N_Vector) = unsafe_load(unsafe_load(convert(Ptr{Ptr{Clong}}, x)))

test/arkstep_Roberts_dns.jl

@@ -11,7 +11,8 @@ function f(t, y_nv, ydot_nv, user_data)
     return Sundials.ARK_SUCCESS
 end
 
-f_C = @cfunction(f, Cint, (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Ptr{Cvoid}))
+f_C = @cfunction(f, Cint,
+                 (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Ptr{Cvoid}))
 
 neq = 3
 

test/common_interface/precs.jl

@@ -1,7 +1,6 @@
-using Sundials, Test, ModelingToolkit, LinearAlgebra, IncompleteLU
+using Sundials, Test, LinearAlgebra, IncompleteLU
 import AlgebraicMultigrid
-
-println("precs.jl start tests..."); flush(stdout)
+import SparsityTracing, SparseDiffTools
 
 const N = 32
 const xyd_brusselator = range(0; stop = 1, length = N)
@@ -25,6 +24,14 @@ function brusselator_2d_loop(du, u, p, t)
                       A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2]
     end
 end
+
+# vector arguments to simplify Jacobian setup
+function brusselator_2d_vec(du_vec, u_vec, p, t)
+    du = reshape(du_vec, N, N, 2)
+    u = reshape(u_vec, N, N, 2)
+    return brusselator_2d_loop(du, u, p, t)
+end
+
 p = (3.4, 1.0, 10.0, step(xyd_brusselator))
 
 function init_brusselator_2d(xyd)
@@ -38,25 +45,33 @@ function init_brusselator_2d(xyd)
     end
     u
 end
-u0 = init_brusselator_2d(xyd_brusselator)
-prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
-                                     u0, (0.0, 11.5), p)
+u0 = vec(init_brusselator_2d(xyd_brusselator))
 
-println("rob_ode_brusselator_2d_mtk"); flush(stdout)
-@time prob_ode_brusselator_2d_mtk = ODEProblem(modelingtoolkitize(prob_ode_brusselator_2d), [],
-                                         (0.0, 11.5); jac = true, sparse = true);
+prob_ode_brusselator_2d = ODEProblem(brusselator_2d_vec,
+                                     u0, (0.0, 11.5), p)
 
-u0 = prob_ode_brusselator_2d_mtk.u0
-p = prob_ode_brusselator_2d_mtk.p
-const jaccache = prob_ode_brusselator_2d_mtk.f.jac(u0, p, 0.0)
+# find Jacobian sparsity pattern
+u0_st = SparsityTracing.create_advec(u0)
+du_st = similar(u0_st)
+brusselator_2d_vec(du_st, u0_st, p, 0.0)
+const jaccache = SparsityTracing.jacobian(du_st, length(du_st))
 const W = I - 1.0 * jaccache
 
+# setup sparse AD for Jacobian
+colors = SparseDiffTools.matrix_colors(jaccache)
+const jaccache_fc = SparseDiffTools.ForwardColorJacCache(nothing, # don't use f to create unique Tag
+                                                         u0,
+                                                         colorvec = colors,
+                                                         sparsity = jaccache)
+
 prectmp = ilu(W; τ = 50.0)
 const preccache = Ref(prectmp)
 
 function psetupilu(p, t, u, du, jok, jcurPtr, gamma)
     if jok
-        prob_ode_brusselator_2d_mtk.f.jac(jaccache, u, p, t)
+        SparseDiffTools.forwarddiff_color_jacobian!(jaccache,
+                                                    (y, x) -> brusselator_2d_vec(y, x, p, t),
+                                                    u, jaccache_fc)
         jcurPtr[] = true
 
         # W = I - gamma*J
@@ -73,16 +88,17 @@ function precilu(z, r, p, t, y, fy, gamma, delta, lr)
     ldiv!(z, preccache[], r)
 end
 
-println("aspreconditioner"); flush(stdout)
-@time prectmp2 = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben(W;
+prectmp2 = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben(W;
                                                                               presmoother = AlgebraicMultigrid.Jacobi(rand(size(W,
                                                                                                                                 1))),
                                                                               postsmoother = AlgebraicMultigrid.Jacobi(rand(size(W,
                                                                                                                                  1)))))
 const preccache2 = Ref(prectmp2)
 function psetupamg(p, t, u, du, jok, jcurPtr, gamma)
     if jok
-        prob_ode_brusselator_2d_mtk.f.jac(jaccache, u, p, t)
+        SparseDiffTools.forwarddiff_color_jacobian!(jaccache,
+                                                    (y, x) -> brusselator_2d_vec(y, x, p, t),
+                                                    u, jaccache_fc)
         jcurPtr[] = true
 
         # W = I - gamma*J
@@ -103,16 +119,13 @@ function precamg(z, r, p, t, y, fy, gamma, delta, lr)
     ldiv!(z, preccache2[], r)
 end
 
-println("sol1:"); flush(stdout)
-@time sol1 = solve(prob_ode_brusselator_2d, CVODE_BDF(; linear_solver = :GMRES);
+sol1 = solve(prob_ode_brusselator_2d, CVODE_BDF(; linear_solver = :GMRES);
              save_everystep = false);
-println("sol2:"); flush(stdout)
-@time sol2 = solve(prob_ode_brusselator_2d,
-        CVODE_BDF(; linear_solver = :GMRES, prec = precilu, psetup = psetupilu,
-                prec_side = 1); save_everystep = false);
-println("sol3:"); flush(stdout)
-@time sol3 = solve(prob_ode_brusselator_2d,
-        CVODE_BDF(; linear_solver = :GMRES, prec = precamg, psetup = psetupamg,
-                prec_side = 1); save_everystep = false);
+sol2 = solve(prob_ode_brusselator_2d,
+             CVODE_BDF(; linear_solver = :GMRES, prec = precilu, psetup = psetupilu,
+                       prec_side = 1); save_everystep = false);
+sol3 = solve(prob_ode_brusselator_2d,
+             CVODE_BDF(; linear_solver = :GMRES, prec = precamg, psetup = psetupamg,
+                       prec_side = 1); save_everystep = false);
 @test sol1.destats.nf > sol2.destats.nf
 @test sol1.destats.nf > sol3.destats.nf

test/cvode_Roberts_dns.jl

@@ -22,7 +22,7 @@ function g(t, y_nv, gout_ptr, user_data)
 end
 
 g_C = @cfunction(g, Cint,
-                (Sundials.realtype, Sundials.N_Vector, Ptr{Sundials.realtype}, Ptr{Cvoid}))
+                 (Sundials.realtype, Sundials.N_Vector, Ptr{Sundials.realtype}, Ptr{Cvoid}))
 
 ## Jacobian routine. Compute J(t,y) = df/dy.
 # broken -- needs a wrapper from Sundials._DlsMat to Matrix and Jac user function wrapper
@@ -64,7 +64,8 @@ function getcfunrob(userfun::T) where {T}
                (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Ref{T}))
 end
 
-Sundials.@checkflag Sundials.CVodeInit(cvode_mem, getcfunrob(userfun), t1, convert(Sundials.NVector, y0))
+Sundials.@checkflag Sundials.CVodeInit(cvode_mem, getcfunrob(userfun), t1,
+                                       convert(Sundials.NVector, y0))
 Sundials.@checkflag Sundials.CVodeInit(cvode_mem, getcfunrob(userfun), t0, y0)
 Sundials.@checkflag Sundials.CVodeSVtolerances(cvode_mem, reltol, abstol)
 Sundials.@checkflag Sundials.CVodeRootInit(cvode_mem, 2, g_C)

test/erkstep_nonlin.jl

@@ -45,7 +45,8 @@ function f(t, y, ydot, user_data)
     return Sundials.ARK_SUCCESS
 end
 
-f_C = @cfunction(f, Cint, (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Ptr{Cvoid}))
+f_C = @cfunction(f, Cint,
+                 (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Ptr{Cvoid}))
 
 mem_ptr = Sundials.ERKStepCreate(f_C, t0, y0)
 erkStep_mem = Sundials.Handle(mem_ptr)

test/ida_Roberts_dns.jl

@@ -45,7 +45,8 @@ function resrob(tres, yy_nv, yp_nv, rr_nv, user_data)
 end
 
 resrob_C = @cfunction(resrob, Cint,
-                (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Sundials.N_Vector, Ptr{Cvoid}))
+                      (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector,
+                       Sundials.N_Vector, Ptr{Cvoid}))
 
 ## Root function routine. Compute functions g_i(t,y) for i = 0,1.
 function grob(t, yy_nv, yp_nv, gout_ptr, user_data)
@@ -57,7 +58,8 @@ function grob(t, yy_nv, yp_nv, gout_ptr, user_data)
 end
 
 grob_C = @cfunction(grob, Cint,
-                (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Ptr{Sundials.realtype}, Ptr{Cvoid}))
+                    (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector,
+                     Ptr{Sundials.realtype}, Ptr{Cvoid}))
 
 ## Define the Jacobian function. BROKEN - JJ is wrong
 function jacrob(Neq, tt, cj, yy, yp, resvec, JJ, user_data, tempv1, tempv2, tempv3)

test/runtests.jl

@@ -39,9 +39,5 @@ end
     @testset "Iterator" begin include("common_interface/iterators.jl") end
     @testset "Errors" begin include("common_interface/errors.jl") end
     @testset "Mass Matrix" begin include("common_interface/mass_matrix.jl") end
-
-    # Don't run, CI failure before Sundials is involved
-    if VERSION < v"1.8"
-        @testset "Preconditioners" begin include("common_interface/precs.jl") end
-    end
+    @testset "Preconditioners" begin include("common_interface/precs.jl") end
 end