NonlinearSolveAlg: give inner NonlinearFunction W's structure as jac_prototype

lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl

@@ -107,7 +107,16 @@ function initialize!(
             nlp_params = (tmp, ustep, γ, α, tstep, k, invγdt, method, p, dt, f)
         end
         if length(cache.cache.u) != length(z)
-            new_prob = SciMLBase.remake(cache.prob; u0 = copy(z), p = nlp_params)
+            new_prob = if cache.W !== nothing
+                # The problem was resized: point the WReuseJac at the resized W and give
+                # the NonlinearFunction a matching jac_prototype. Both keep their types
+                # (only array sizes change), so cache.prob's concrete type is preserved.
+                cache.prob.f.jac.W[] = cache.W
+                new_f = SciMLBase.remake(cache.prob.f; jac_prototype = similar(cache.W))
+                SciMLBase.remake(cache.prob; f = new_f, u0 = copy(z), p = nlp_params)
+            else
+                SciMLBase.remake(cache.prob; u0 = copy(z), p = nlp_params)
+            end
             cache.prob = new_prob
             cache.cache = init(new_prob, cache.cache.alg)
         else

lib/OrdinaryDiffEqNonlinearSolve/src/type.jl

@@ -226,6 +226,16 @@ mutable struct NLAndersonConstantCache{uType, tType, uEltypeNoUnits} <:
     droptol::Union{Nothing, tType}
 end
 
+# Analytic-jac callback handing the ODE-side W to the inner NonlinearSolve under
+# W-reuse. Deliberately a struct holding a Ref rather than a closure: its type depends
+# only on W's type, not on a captured binding, so when `resize!` replaces the W array
+# the same NonlinearFunction/NonlinearProblem types remain valid and `initialize!` just
+# swaps the Ref target.
+struct WReuseJac{W} <: Function
+    W::Base.RefValue{W}
+end
+(j::WReuseJac)(J_out, z, p) = (copyto!(J_out, j.W[]); J_out)
+
 mutable struct NonlinearSolveCache{uType, tType, rateType, tType2, P, C, JType, WType, ufType, jcType, du1Type} <:
     AbstractNLSolverCache
     ustep::uType

lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl

@@ -345,13 +345,15 @@ function build_nlsolver(
                     (tmp, ustep, γ, α, tstep, k, invγdt, DIRK, p, dt, f)
                 end
                 if use_w_reuse
-                    nlf_jac! = let W = W_for_reuse
-                        (J_out, z, p) -> (copyto!(J_out, W); J_out)
-                    end
+                    nlf_jac! = WReuseJac(Ref(W_for_reuse))
+                    # Without a `jac_prototype`, NonlinearSolve allocates a dense J for an
+                    # analytic-jac function, so a sparse/structured W degrades to dense LU
+                    # and sparse-only linsolves (e.g. KLU) crash. Mirror W's structure.
                     NonlinearProblem(
                         NonlinearFunction{true, SciMLBase.FullSpecialize}(
                             nlf;
-                            jac = nlf_jac!
+                            jac = nlf_jac!,
+                            jac_prototype = similar(W_for_reuse)
                         ),
                         ztmp, nlp_params
                     )

lib/OrdinaryDiffEqNonlinearSolve/test/nsa_sparse_tests.jl

@@ -0,0 +1,110 @@
+using OrdinaryDiffEqBDF, OrdinaryDiffEqSDIRK
+using OrdinaryDiffEqNonlinearSolve
+using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
+using NonlinearSolve: NewtonRaphson
+using LinearSolve, LinearAlgebra, SparseArrays, ADTypes
+using SciMLBase
+using Test
+
+# 1D Brusselator with a hand-assembled sparse Jacobian pattern (tridiagonal blocks).
+const N_b = 8
+function bruss1d!(du, u, p, t)
+    A, B, alpha, dx = p
+    a = alpha / dx^2
+    n = N_b
+    @inbounds for i in 1:n
+        im1 = i == 1 ? n : i - 1
+        ip1 = i == n ? 1 : i + 1
+        x = u[i]
+        y = u[n + i]
+        du[i] = a * (u[im1] + u[ip1] - 2x) + B + x^2 * y - (A + 1) * x
+        du[n + i] = a * (u[n + im1] + u[n + ip1] - 2y) + A * x - x^2 * y
+    end
+    return nothing
+end
+function bruss1d_jacproto(n)
+    Is = Int[]
+    Js = Int[]
+    for i in 1:n
+        im1 = i == 1 ? n : i - 1
+        ip1 = i == n ? 1 : i + 1
+        for (r, c) in (
+                (i, im1), (i, i), (i, ip1), (i, n + i),
+                (n + i, n + im1), (n + i, n + i), (n + i, n + ip1), (n + i, i),
+            )
+            push!(Is, r)
+            push!(Js, c)
+        end
+    end
+    return sparse(Is, Js, ones(length(Is)), 2n, 2n)
+end
+u0 = vcat(
+    [22 * (i / (N_b + 1) * (1 - i / (N_b + 1)))^(3 / 2) for i in 1:N_b],
+    [27 * (i / (N_b + 1) * (1 - i / (N_b + 1)))^(3 / 2) for i in 1:N_b]
+)
+p = (3.4, 1.0, 10.0, 1 / N_b)
+f_sparse = ODEFunction(bruss1d!; jac_prototype = bruss1d_jacproto(N_b))
+prob = ODEProblem(f_sparse, u0, (0.0, 2.0), p)
+
+refsol = solve(prob, FBDF(); reltol = 1.0e-10, abstol = 1.0e-12)
+
+nsa = NonlinearSolveAlg(NewtonRaphson(; autodiff = AutoForwardDiff()))
+
+@testset "NSA W-reuse keeps inner Jacobian sparse" begin
+    integ = init(prob, FBDF(nlsolve = nsa); reltol = 1.0e-8, abstol = 1.0e-10)
+    nsacache = integ.cache.nlsolver.cache
+    @test nsacache.W isa SparseMatrixCSC
+    # The inner NonlinearSolve Jacobian must mirror W's structure, not densify.
+    @test integ.cache.nlsolver.cache.cache.jac_cache.J isa SparseMatrixCSC
+end
+
+@testset "NSA + sparse-only linsolve (KLU) solves" begin
+    for alg in (
+            FBDF(linsolve = KLUFactorization(), nlsolve = nsa),
+            TRBDF2(linsolve = KLUFactorization(), nlsolve = nsa),
+        )
+        sol = solve(prob, alg; reltol = 1.0e-8, abstol = 1.0e-10)
+        @test SciMLBase.successful_retcode(sol)
+        @test sol.u[end] ≈ refsol.u[end] rtol = 1.0e-5
+    end
+end
+
+# Resizing the integrator must rebuild the inner Jacobian at the new size: the jac
+# closure and jac_prototype are fixed at build time, so the length-mismatch path in
+# initialize! has to refresh both against the resized W (this crashed with
+# `DimensionMismatch: B has leading dimension 2, but needs 1` when jac_prototype
+# followed the W structure).
+@testset "resize with W-reuse rebuilds inner Jacobian" begin
+    fgrow = function (du, u, p, t)
+        for i in 1:length(u)
+            du[i] = (0.3 / length(u)) * u[i]
+        end
+        return nothing
+    end
+    condition = (u, t, integrator) -> 1 - maximum(u)
+    affect! = function (integrator)
+        u = integrator.u
+        maxidx = findmax(u)[2]
+        resize!(integrator, length(u) + 1)
+        Θ = 0.3
+        u[maxidx] = Θ
+        u[end] = 1 - Θ
+        return nothing
+    end
+    cb = ContinuousCallback(condition, affect!)
+    probg = ODEProblem(fgrow, [0.2], (0.0, 10.0))
+    # NOTE: assert the resize worked and nothing crashed (mirroring
+    # test/Integrators_II/ode_cache_tests.jl). Whether the marginally-stable
+    # grow-a-cell dynamics finishes with Success is method- and
+    # rounding-sensitive (ImplicitEuler+NLNewton is Unstable on it too), so
+    # retcode is only checked for TRBDF2.
+    for alg in (
+            TRBDF2(nlsolve = nsa), KenCarp4(nlsolve = nsa),
+            ImplicitEuler(nlsolve = nsa),
+        )
+        sol = solve(probg, alg; callback = cb, dt = 1 / 2)
+        @test length(sol.u[end]) > 1
+    end
+    sol = solve(probg, TRBDF2(nlsolve = nsa); callback = cb, dt = 1 / 2)
+    @test SciMLBase.successful_retcode(sol)
+end

lib/OrdinaryDiffEqNonlinearSolve/test/runtests.jl

@@ -36,6 +36,7 @@ if TEST_GROUP ∉ ("QA", "ModelingToolkit")
     @time @safetestset "DAE Initialization Tests" include("dae_initialization_tests.jl")
     @time @safetestset "CheckInit Tests" include("checkinit_tests.jl")
     @time @safetestset "Nested AD over NonlinearSolveAlg" include("nested_ad_nlsolvealg_tests.jl")
+    @time @safetestset "NonlinearSolveAlg Sparse Jacobian Tests" include("nsa_sparse_tests.jl")
 end
 
 # Run QA tests (JET, Aqua)