Add tests for the caching interface

SciML · Oct 17, 2023 · ff604ee · ff604ee
1 parent a1c97a6
commit ff604ee
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/src/default.jl b/src/default.jl
@@ -46,7 +46,7 @@ function RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing,
     return RobustMultiNewton{_unwrap_val(concrete_jac)}(adkwargs, linsolve, precs)
 end
 
-@concrete mutable struct RobustMultiNewtonCache{iip} <: AbstractNonlinearSolveCache{iip}
+@concrete mutable struct RobustMultiNewtonCache{iip, N} <: AbstractNonlinearSolveCache{iip}
     caches
     alg
     current::Int
@@ -67,8 +67,8 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::RobustMultiNe
 
     # Partially Type Unstable but can't do much since some upstream caches -- LineSearches
     # and SparseDiffTools cause the instability
-    return RobustMultiNewtonCache{iip}(map(solver -> SciMLBase.__init(prob, solver, args...;
-                kwargs...), algs), alg, 1)
+    return RobustMultiNewtonCache{iip, length(algs)}(map(solver -> SciMLBase.__init(prob,
+                solver, args...; kwargs...), algs), alg, 1)
 end
 
 """
@@ -110,7 +110,7 @@ function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothi
         precs)
 end
 
-@concrete mutable struct FastShortcutNonlinearPolyalgCache{iip} <:
+@concrete mutable struct FastShortcutNonlinearPolyalgCache{iip, N} <:
                          AbstractNonlinearSolveCache{iip}
     caches
     alg
@@ -136,7 +136,7 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip},
         TrustRegion(; linsolve, precs,
             radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...))
 
-    return FastShortcutNonlinearPolyalgCache{iip}(map(solver -> SciMLBase.__init(prob,
+    return FastShortcutNonlinearPolyalgCache{iip, length(algs)}(map(solver -> SciMLBase.__init(prob,
                 solver, args...; kwargs...), algs), alg, 1)
 end
 
@@ -159,6 +159,8 @@ end
         ]
     else
         [
+            # FIXME: Broyden and Klement are type unstable
+            #        (upstream SimpleNonlinearSolve.jl issue)
             !iip ? :(Klement()) : nothing, # Klement not yet implemented for IIP
             !iip ? :(Broyden()) : nothing, # Broyden not yet implemented for IIP
             :(NewtonRaphson(; linsolve, precs, adkwargs...)),

diff --git a/test/polyalgs.jl b/test/polyalgs.jl
@@ -3,9 +3,25 @@ using NonlinearSolve, Test
 f(u, p) = u .* u .- 2
 u0 = [1.0, 1.0]
 probN = NonlinearProblem(f, u0)
+
+# Uses the `__solve` function
 @time solver = solve(probN; abstol = 1e-9)
+@test SciMLBase.successful_retcode(solver)
 @time solver = solve(probN, RobustMultiNewton(); abstol = 1e-9)
+@test SciMLBase.successful_retcode(solver)
 @time solver = solve(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9)
+@test SciMLBase.successful_retcode(solver)
+
+# Test the caching interface
+cache = init(probN; abstol = 1e-9)
+@time solver = solve!(cache)
+@test SciMLBase.successful_retcode(solver)
+cache = init(probN, RobustMultiNewton(); abstol = 1e-9)
+@time solver = solve!(cache)
+@test SciMLBase.successful_retcode(solver)
+cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9)
+@time solver = solve!(cache)
+@test SciMLBase.successful_retcode(solver)
 
 # https://github.com/SciML/NonlinearSolve.jl/issues/153
 function f(du, u, p)