Add step! for polyalgorithms

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.6.0"
+version = "3.6.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/adtypes.jl

@@ -83,13 +83,13 @@ AutoSparseForwardDiff
 Uses [`PolyesterForwardDiff.jl`](https://github.com/JuliaDiff/PolyesterForwardDiff.jl)
 to compute the jacobian. This is essentially parallelized `ForwardDiff.jl`.
 
-    - Supports both inplace and out-of-place functions
+  - Supports both inplace and out-of-place functions
 
 ### Keyword Arguments
 
-    - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting
-      this number to a high value will lead to slowdowns. Use
-      [`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
+  - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting
+    this number to a high value will lead to slowdowns. Use
+    [`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
 """
 AutoPolyesterForwardDiff
 

src/core/generic.jl

@@ -51,7 +51,8 @@ function SciMLBase.step!(cache::AbstractNonlinearSolveCache{iip, timeit},
     res = @static_timeit cache.timer "solve" begin
         __step!(cache, args...; kwargs...)
     end
-    cache.nsteps += 1
+
+    hasfield(typeof(cache), :nsteps) && (cache.nsteps += 1)
 
     if timeit
         cache.total_time += time() - time_start

src/default.jl

@@ -44,26 +44,35 @@ function Base.show(io::IO, alg::NonlinearSolvePolyAlgorithm{pType, N}) where {pT
     end
 end
 
-@concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N} <:
-                         AbstractNonlinearSolveCache{iip, false}
+@concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N, timeit} <:
+                         AbstractNonlinearSolveCache{iip, timeit}
     caches
     alg
+    best::Int
     current::Int
+    nsteps::Int
+    total_time::Float64
+    maxtime
+    retcode::ReturnCode.T
+    force_stop::Bool
 end
 
 function reinit_cache!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...)
     foreach(c -> reinit_cache!(c, args...; kwargs...), cache.caches)
     cache.current = 1
+    cache.nsteps = 0
+    cache.total_time = 0.0
 end
 
 for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS))
     algType = NonlinearSolvePolyAlgorithm{pType}
     @eval begin
-        function SciMLBase.__init(
-                prob::$probType, alg::$algType{N}, args...; kwargs...) where {N}
-            return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N}(
-                map(solver -> SciMLBase.__init(prob, solver, args...; kwargs...), alg.algs),
-                alg, 1)
+        function SciMLBase.__init(prob::$probType, alg::$algType{N}, args...;
+                maxtime = nothing, kwargs...) where {N}
+            return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N, maxtime !== nothing}(
+                map(solver -> SciMLBase.__init(prob, solver, args...; maxtime, kwargs...),
+                    alg.algs), alg, -1, 1, 0, 0.0, maxtime,
+                ReturnCode.Default, false)
         end
     end
 end
@@ -87,9 +96,8 @@ end
                         stats = $(sol_syms[i]).stats
                         u = $(sol_syms[i]).u
                         fu = get_fu($(cache_syms[i]))
-                        return SciMLBase.build_solution(
-                            $(sol_syms[i]).prob, cache.alg, u, fu;
-                            retcode = ReturnCode.Success, stats,
+                        return SciMLBase.build_solution($(sol_syms[i]).prob, cache.alg, u,
+                            fu; retcode = $(sol_syms[i]).retcode, stats,
                             original = $(sol_syms[i]), trace = $(sol_syms[i]).trace)
                     end
                     cache.current = $(i + 1)
@@ -103,12 +111,11 @@ end
     end
     push!(calls,
         quote
-            retcode = ReturnCode.MaxIters
-
             fus = tuple($(Tuple(resids)...))
             minfu, idx = __findmin(cache.caches[1].internalnorm, fus)
             stats = cache.caches[idx].stats
-            u = cache.caches[idx].u
+            u = get_u(cache.caches[idx])
+            retcode = cache.caches[idx].retcode
 
             return SciMLBase.build_solution(cache.caches[idx].prob, cache.alg, u, fus[idx];
                 retcode, stats, cache.caches[idx].trace)
@@ -117,6 +124,45 @@ end
     return Expr(:block, calls...)
 end
 
+@generated function __step!(
+        cache::NonlinearSolvePolyAlgorithmCache{iip, N}, args...; kwargs...) where {iip, N}
+    calls = []
+    cache_syms = [gensym("cache") for i in 1:N]
+    for i in 1:N
+        push!(calls,
+            quote
+                $(cache_syms[i]) = cache.caches[$(i)]
+                if $(i) == cache.current
+                    __step!($(cache_syms[i]), args...; kwargs...)
+                    if !not_terminated($(cache_syms[i]))
+                        if SciMLBase.successful_retcode($(cache_syms[i]).retcode)
+                            cache.best = $(i)
+                            cache.force_stop = true
+                            cache.retcode = $(cache_syms[i]).retcode
+                        else
+                            cache.current = $(i + 1)
+                        end
+                    end
+                    return
+                end
+            end)
+    end
+
+    push!(calls,
+        quote
+            if !(1 ≤ cache.current ≤ length(cache.caches))
+                minfu, idx = __findmin(first(cache.caches).internalnorm, cache.caches)
+                cache.best = idx
+                cache.retcode = cache.caches[cache.best].retcode
+                cache.force_stop = true
+                return
+            end
+        end
+    )
+
+    return Expr(:block, calls...)
+end
+
 for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS))
     algType = NonlinearSolvePolyAlgorithm{pType}
     @eval begin

src/utils.jl

@@ -94,6 +94,9 @@ LazyArrays.applied_axes(::typeof(__zero), x) = axes(x)
 @inline __is_complex(::Type{Complex}) = true
 @inline __is_complex(::Type{T}) where {T} = false
 
+function __findmin_caches(f, caches)
+    return __findmin(f ∘ get_fu, caches)
+end
 function __findmin(f, x)
     return findmin(x) do xᵢ
         fx = f(xᵢ)

test/misc/polyalg_tests.jl

@@ -28,6 +28,31 @@
     cache = init(probN, custom_polyalg; abstol = 1e-9)
     solver = solve!(cache)
     @test SciMLBase.successful_retcode(solver)
+
+    # Test the step interface
+    cache = init(probN; abstol = 1e-9)
+    for i in 1:10000
+        step!(cache)
+        cache.force_stop && break
+    end
+    @test SciMLBase.successful_retcode(cache.retcode)
+    cache = init(probN, RobustMultiNewton(); abstol = 1e-9)
+    for i in 1:10000
+        step!(cache)
+        cache.force_stop && break
+    end
+    @test SciMLBase.successful_retcode(cache.retcode)
+    cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9)
+    for i in 1:10000
+        step!(cache)
+        cache.force_stop && break
+    end
+    @test SciMLBase.successful_retcode(cache.retcode)
+    cache = init(probN, custom_polyalg; abstol = 1e-9)
+    for i in 1:10000
+        step!(cache)
+        cache.force_stop && break
+    end
 end
 
 @testitem "Testing #153 Singular Exception" begin

testing.jl

@@ -0,0 +1,14 @@
+using NonlinearSolve
+
+f(u, p) = u .* u .- 2
+
+u0 = [1.0, 1.0]
+
+prob = NonlinearProblem(f, u0)
+
+nlcache = init(prob);
+
+for i in 1:10
+    step!(nlcache)
+    @show nlcache.retcode
+end