Fix NonlinearSolveAlg crashing/stalling with polyalgorithm inner solvers

lib/OrdinaryDiffEqNonlinearSolve/Project.toml

@@ -51,7 +51,7 @@ DiffEqDevTools = {path = "../DiffEqDevTools"}
 SciMLTesting = "2.1"
 Pkg = "1"
 NonlinearSolve = "4.20.3"
-NonlinearSolveBase = "2.35"
+NonlinearSolveBase = "2.39"
 ForwardDiff = "1.3.3"
 Test = "<0.0.1, 1"
 FastBroadcast = "1.3"

lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl

@@ -20,7 +20,9 @@ using NonlinearSolve: FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg, New
     HomotopySweep, HomotopyPolyAlgorithm, ArcLengthContinuation, step!
 # The operator Jacobian path is implemented in NonlinearSolveBase and needs its own floor.
 import NonlinearSolveBase
-using NonlinearSolveBase: get_linear_cache
+# `get_u`/`get_fu` are the only inner-state reads that hold for every inner cache type:
+# polyalgorithm caches keep `u`/`fu` on the active branch, not as top-level fields.
+using NonlinearSolveBase: get_linear_cache, get_u, get_fu
 using MuladdMacro: @muladd
 using FastBroadcast: @..
 import FastClosures: @closure

lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl

@@ -140,7 +140,7 @@ function initialize!(
         else
             nlp_params = (tmp, ustep, γ, α, tstep, k, invγdt, method, p, dt, f)
         end
-        if length(cache.cache.u) != length(z)
+        if length(get_u(cache.cache)) != length(z)
             new_prob = if cache.W !== nothing
                 # W-reuse: re-point the inner jacobian at the resized W via the same mapping
                 # used at build time. Only array sizes change, so the jac's concrete type —
@@ -229,7 +229,7 @@ function rescale_stale_W_rhs!(nlcache, nlsolver, integrator, nlstep_data)
     isdae = nlsolve_f(integrator) isa DAEFunction
     γdt = isdae ? nlsolver.α * cache.invγdt : nlsolver.γ * integrator.dt
     W_γdt ≈ γdt && return nothing
-    rmul!(nlcache.fu, 2 / (1 + γdt / W_γdt))
+    rmul!(get_fu(nlcache), 2 / (1 + γdt / W_γdt))
     return nothing
 end
 
@@ -252,11 +252,12 @@ end
         return convert(eltype(z), Inf)
     end
     step!(nlcache; recompute_jacobian)
-    nlsolver.ztmp = nlcache.u
+    active_u = get_u(nlcache)
+    nlsolver.ztmp = active_u
 
     ustep = compute_ustep(tmp, γ, z, method)
     atmp = calculate_residuals(
-        z .- nlcache.u, uprev, ustep, opts.abstol, opts.reltol,
+        z .- active_u, uprev, ustep, opts.abstol, opts.reltol,
         opts.internalnorm, t
     )
     ndz = opts.internalnorm(atmp, t)
@@ -289,15 +290,18 @@ end
     step!(nlcache; recompute_jacobian)
 
     if nlstep_data !== nothing
+        active_fu = get_fu(nlcache)
+        # No `trace` is attached: this solution only feeds `nlprobmap` right below, and
+        # polyalgorithm caches keep the trace on the active branch with no accessor for it.
         nlstepsol = SciMLBase.build_solution(
-            nlcache.prob, nlcache.alg, nlcache.u, nlcache.fu;
-            nlcache.retcode, nlcache.stats, nlcache.trace
+            nlcache.prob, nlcache.alg, get_u(nlcache), active_fu;
+            nlcache.retcode, nlcache.stats
         )
         nlstep_data.nlprobmap(ztmp, nlstepsol)
         ustep = compute_ustep!(ustep, tmp, γ, z, method)
         atmp_sub = @view(atmp[nlstep_data.u0perm])
         calculate_residuals!(
-            atmp_sub, nlcache.fu,
+            atmp_sub, active_fu,
             @view(uprev[nlstep_data.u0perm]),
             @view(ustep[nlstep_data.u0perm]), opts.abstol,
             opts.reltol, opts.internalnorm, t
@@ -309,7 +313,8 @@ end
         # convergence check to declare success ~sqrt(n_full/n_sub) early.
         ndz = opts.internalnorm(atmp_sub, t)
     else
-        @.. broadcast = false ztmp = nlcache.u
+        active_u = get_u(nlcache)
+        @.. broadcast = false ztmp = active_u
         ustep = compute_ustep!(ustep, tmp, γ, z, method)
         @.. broadcast = false atmp = z - ztmp
         calculate_residuals!(

lib/OrdinaryDiffEqNonlinearSolve/test/nsa_polyalg_tests.jl

@@ -0,0 +1,43 @@
+using OrdinaryDiffEqBDF, OrdinaryDiffEqSDIRK
+using OrdinaryDiffEqNonlinearSolve
+using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
+using NonlinearSolve: RobustMultiNewton, NewtonRaphson
+using ADTypes, LinearAlgebra, SciMLBase
+using Test
+
+# Regression test for #3859: a NonlinearSolve polyalgorithm (e.g. RobustMultiNewton)
+# used as the NonlinearSolveAlg inner solver used to crash in `initialize!` with
+# `type NonlinearSolvePolyAlgorithmCache has no field u`, because a polyalg cache keeps
+# `u`/`fu` on its active branch rather than at top level; newton.jl now reads them
+# through `NonlinearSolveBase.get_u`/`get_fu`. A second bug lived upstream: the polyalg
+# cache's `reinit!` did not clear `force_stop`/`retcode`, so once any branch converged,
+# every later `step!` short-circuited on `not_terminated`, silently reusing a stale `u`
+# and diverging (fixed in NonlinearSolveBase 2.39). This test exercises both: it would
+# throw without the accessors and return wrong values without the `reinit!` reset.
+
+const A = [-1000.0 1.0; 1.0 -2.0]
+function lin!(du, u, p, t)
+    mul!(du, A, u)
+    return nothing
+end
+prob = ODEProblem(lin!, [1.0, 1.0], (0.0, 0.5))
+uexact = exp(A * 0.5) * prob.u0
+
+poly = NonlinearSolveAlg(RobustMultiNewton(; autodiff = AutoForwardDiff()))
+single = NonlinearSolveAlg(NewtonRaphson(; autodiff = AutoForwardDiff()))
+
+# Each run is checked against the analytic solution rather than the polyalg run against
+# the NewtonRaphson one: the outer iteration takes exactly one inner `step!` per
+# iteration, and RobustMultiNewton's active trust-region branch damps that step, so the
+# two inner solvers legitimately produce different step sequences (and step counts).
+# The stale-iterate bug produced O(1) errors, far above this bound.
+@testset "polyalg inner solver solves as accurately as single-algorithm" begin
+    for ODEAlg in (TRBDF2, FBDF)
+        ref = solve(prob, ODEAlg(nlsolve = single); reltol = 1.0e-8, abstol = 1.0e-10)
+        sol = solve(prob, ODEAlg(nlsolve = poly); reltol = 1.0e-8, abstol = 1.0e-10)
+        @test SciMLBase.successful_retcode(ref)
+        @test SciMLBase.successful_retcode(sol)
+        @test norm(ref.u[end] .- uexact) / norm(uexact) < 1.0e-4
+        @test norm(sol.u[end] .- uexact) / norm(uexact) < 1.0e-4
+    end
+end

lib/OrdinaryDiffEqNonlinearSolve/test/qa/qa.jl

@@ -40,7 +40,7 @@ run_qa(
                 # Base — owner-internal, no public alternative
                 :RefValue,
                 # NonlinearSolveBase — owner-internal, no public alternative
-                :not_terminated,
+                :not_terminated, :get_u, :get_fu,
                 # OrdinaryDiffEqDifferentiation — owner-internal, no public alternative
                 :default_krylov_warm_start,
                 # `@SciMLMessage` reached through OrdinaryDiffEqCore (owner SciMLLogging).

lib/OrdinaryDiffEqNonlinearSolve/test/runtests.jl

@@ -42,6 +42,7 @@ if TEST_GROUP ∉ ("QA", "ModelingToolkit")
     @time @safetestset "NonlinearSolveAlg Matrix-Free WOperator Tests" include("nsa_matrixfree_tests.jl")
     @time @safetestset "NonlinearSolveAlg Smoothed Error Estimate Tests" include("nsa_smooth_est_tests.jl")
     @time @safetestset "NonlinearSolveAlg Stats Tests" include("nsa_stats_tests.jl")
+    @time @safetestset "NonlinearSolveAlg Polyalgorithm Inner Solver Tests" include("nsa_polyalg_tests.jl")
 end
 
 # Run QA tests (JET, Aqua)