Add reusable homotopy solver caches

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

Project.toml

@@ -1,6 +1,6 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
-version = "4.22.0"
+version = "4.23.0"
 authors = ["SciML"]
 
 [deps]
@@ -98,7 +98,7 @@ NLSolvers = "0.5"
 NLsolve = "4.5"
 NaNMath = "1"
 NonlinearProblemLibrary = "0.1.2"
-NonlinearSolveBase = "2.36"
+NonlinearSolveBase = "2.37"
 NonlinearSolveFirstOrder = "2"
 NonlinearSolveQuasiNewton = "1.12"
 NonlinearSolveSpectralMethods = "1.6"

lib/NonlinearSolveBase/Project.toml

@@ -1,6 +1,6 @@
 name = "NonlinearSolveBase"
 uuid = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
-version = "2.36.1"
+version = "2.37.0"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]

lib/NonlinearSolveBase/src/homotopy_sweep.jl

@@ -60,7 +60,8 @@ residual cost from one evaluation per subalgorithm to one per subalgorithm actua
 attempted (usually just the retained one). A consequence is that the shared `stats`
 of a solution produced after mid-solve escalation only reflect the subalgorithms run
 since the last deferred reinitialization, i.e. effectively the winning subalgorithm's
-own effort.
+own effort. Updating solver options in the same `reinit!` eagerly updates every
+subcache so that later escalation observes the new values.
 """
 @concrete struct NonlinearSolvePolyAlgorithm <: AbstractNonlinearSolveAlgorithm
     static_length <: Val
@@ -175,7 +176,7 @@ const RETAIN_REPROBE_INTERVAL = 8
 
 function InternalAPI.reinit!(
         cache::NonlinearSolvePolyAlgorithmCache, args...; p = cache.prob.p, u0 = cache.u0,
-        retain_best::Bool = false
+        retain_best::Bool = false, kwargs...
     )
     cache.retain_best = retain_best
     cache.retain_count = retain_best ? cache.retain_count + 1 : 0
@@ -186,19 +187,21 @@ function InternalAPI.reinit!(
     cache.start_current = cache.current
     cache.wrapped = false
     if retain_best
+        cache.deferred_u0 = u0
+        cache.deferred_p = p
+    end
+    if retain_best && isempty(kwargs)
         # Lazy subcache reinitialization: every subcache `reinit!` evaluates the
         # residual at the new `u0` (`Utils.reinit_common!`), so eagerly
         # reinitializing all N subcaches costs N residual calls per warm-started
         # solve even though a retained solve usually runs only the starting
         # subalgorithm. Reinitialize just that one here and defer the rest to the
         # moment escalation actually reaches them (`deferred_subcache_reinit!`).
-        cache.deferred_u0 = u0
-        cache.deferred_p = p
         subcache = cache.caches[cache.current]
-        InternalAPI.reinit!(subcache, args...; u = get_u(subcache), p, u0)
+        InternalAPI.reinit!(subcache, args...; u = get_u(subcache), p, u0, kwargs...)
     else
         foreach(cache.caches) do cache
-            InternalAPI.reinit!(cache, args...; u = get_u(cache), p, u0)
+            InternalAPI.reinit!(cache, args...; u = get_u(cache), p, u0, kwargs...)
         end
     end
     InternalAPI.reinit!(cache.stats)
@@ -225,6 +228,10 @@ reinit_retaining!(cache, u0) = SciMLBase.reinit!(cache, u0)
 function reinit_retaining!(cache::NonlinearSolvePolyAlgorithmCache, u0)
     return SciMLBase.reinit!(cache, u0; retain_best = true)
 end
+reinit_retaining!(cache, u0, p) = SciMLBase.reinit!(cache, u0; p)
+function reinit_retaining!(cache::NonlinearSolvePolyAlgorithmCache, u0, p)
+    return SciMLBase.reinit!(cache, u0; p, retain_best = true)
+end
 
 function SciMLBase.__init(
         prob::AbstractNonlinearProblem, alg::NonlinearSolvePolyAlgorithm, args...;

lib/NonlinearSolveBase/src/polyalg.jl

@@ -1,5 +1,6 @@
 using NonlinearSolve
 
+using CommonSolve
 using SciMLBase
 
 import NonlinearSolveBase as NLB
@@ -53,6 +54,26 @@ cache_result, same_cache, cache_solve_allocs = cache_only_solve_allocations(inne
 @test SciMLBase.successful_retcode(cache_result.retcode)
 VERSION >= v"1.11" && @test cache_solve_allocs == 0
 
+function cached_sweep_retcode!(cache, u0, p, abstol)
+    SciMLBase.reinit!(cache, u0; p, abstol)
+    return CommonSolve.solve!(cache).retcode
+end
+
+function cached_sweep_allocations(prob)
+    alg = HomotopySweep(; inner = NewtonRaphson(), adaptive = false, nsteps = 10)
+    cache = SciMLBase.init(prob, alg; abstol = 0.0)
+    abstol = 1.0e-10
+    cached_sweep_retcode!(cache, prob.u0, prob.p, abstol)
+    cached_sweep_retcode!(cache, prob.u0, prob.p, abstol)
+    GC.gc()
+    allocs = @allocated cached_sweep_retcode!(cache, prob.u0, prob.p, abstol)
+    return cache, allocs
+end
+
+sweep_cache, sweep_cache_allocs = cached_sweep_allocations(prob)
+@test SciMLBase.successful_retcode(CommonSolve.solve!(sweep_cache))
+VERSION >= v"1.11" && @test sweep_cache_allocs == 0
+
 # --- end-to-end per-step allocation with a clean `NewtonRaphson` inner. The slope between two
 # fixed-step runs cancels compile/one-time-init cost, leaving pure per-step allocation. Gated
 # on Julia 1.11+: the LinearSolve factorization-workspace reuse that removes the last

test/Core/homotopy_alloc_tests__item1.jl

@@ -14,6 +14,7 @@ sol = solve(prob, HomotopySweep(; initial_step_factor = 1.0e-4, expand_factor =
 @test !SciMLBase.successful_retcode(sol)
 @test isfinite(sol.u[1])
 @test 0 < sol.u[1] < 1
+@test sol.resid !== nothing
 
 Hs(u, p, λ) = SA[u[1] - λ]
 probs = HomotopyProblem(Hs, SA[0.0]; λspan = (0.0, 1.0))

test/Core/homotopy_effort_tests__item2.jl

@@ -12,4 +12,5 @@ prob = HomotopyProblem(H, [0.0]; λspan = (1.0e9, 2.0e9))
 sol = solve(prob, HomotopySweep(; inner = NewtonRaphson()); maxiters = 5)
 @test sol.retcode == SciMLBase.ReturnCode.Stalled
 @test !SciMLBase.successful_retcode(sol)
-@test sol.resid === nothing
+@test sol.resid !== nothing
+@test all(isfinite, sol.resid)

test/Core/homotopy_sweep_tests__item10.jl

@@ -0,0 +1,33 @@
+using NonlinearSolve
+using CommonSolve
+using SciMLBase
+using Test
+
+function Hcache!(du, u, p, λ)
+    du[1] = (1 - λ) * (u[1] - p[1]) + λ * (u[1]^2 - p[1])
+    return nothing
+end
+
+prob = HomotopyProblem(Hcache!, [4.0], [4.0])
+algs = (
+    HomotopySweep(; inner = NewtonRaphson(), adaptive = false, nsteps = 10),
+    KantorovichHomotopy(; inner = NewtonRaphson(), nsteps = 10, strict = false),
+    HomotopySweep(; adaptive = false, nsteps = 10),
+)
+
+for alg in algs
+    cache = init(prob, alg; abstol = 0.0)
+    @test reinit!(cache, prob.u0; p = prob.p, abstol = 1.0e-10) === cache
+    sol = solve!(cache)
+    @test SciMLBase.successful_retcode(sol)
+    @test sol.u[1] ≈ 2.0 atol = 1.0e-10
+
+    u0 = [9.0]
+    p = [9.0]
+    @test reinit!(cache, u0; p, abstol = 1.0e-10) === cache
+    sol = solve!(cache)
+    @test SciMLBase.successful_retcode(sol)
+    @test sol.u[1] ≈ 3.0 atol = 1.0e-10
+    @test sol.prob.u0 === u0
+    @test sol.prob.p === p
+end

test/Core/homotopy_sweep_tests__item23.jl

@@ -3,6 +3,7 @@
 # subcache reinitialization, upward escalation, wrap-around floored at the
 # algorithm's `start_index`, and exact status-quo behavior when retention is off.
 using NonlinearSolve, SciMLBase
+using NonlinearSolveBase: get_abstol
 
 using Test
 
@@ -123,3 +124,16 @@ end
     @test SciMLBase.successful_retcode(sol2)
     @test sol2.u[1] ≈ ROOT atol = 1.0e-6
 end
+
+@testset "solver option updates reach every retained subcache" begin
+    alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden()))
+    prob = NonlinearProblem(fcubic, [0.0])
+    cache = init(prob, alg; abstol = 0.0)
+    cache.best = 2
+
+    NF[] = 0
+    reinit!(cache, [1.2]; retain_best = true, abstol = 1.0e-8)
+    @test cache.current == 2
+    @test NF[] == 2
+    @test all(subcache -> get_abstol(subcache) == 1.0e-8, cache.caches)
+end

test/Core/polyalg_retention_tests__item1.jl

@@ -87,6 +87,7 @@ else
             @time @safetestset "HomotopySweep secant predictor reduces corrector work" include("Core/homotopy_sweep_tests__item20.jl")
             @time @safetestset "HomotopySweep regrows the step after bisecting a hard region" include("Core/homotopy_sweep_tests__item21.jl")
             @time @safetestset "Continuation drivers return concretely-typed solutions (no Any-typed original)" include("Core/homotopy_sweep_tests__item22.jl")
+            @time @safetestset "HomotopySweep init/reinit!/solve! cache reuse" include("Core/homotopy_sweep_tests__item23.jl")
             @time @safetestset "Kantorovich homotopy controller" include("Core/kantorovich_homotopy_tests__item1.jl")
             @time @safetestset "ArcLengthContinuation construction + defaults" include("Core/arclength_tests__item1.jl")
             @time @safetestset "ArcLengthContinuation happy path (fold-free, matches sweep)" include("Core/arclength_tests__item2.jl")