diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl index 3775ff134d1..07dbc91d3cd 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/type.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/type.jl index 45120561f25..dbb5c7288b0 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/type.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl index 63fcf4ea67a..c8b9cb3d3cb 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl +++ b/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 ) diff --git a/lib/OrdinaryDiffEqNonlinearSolve/test/nsa_sparse_tests.jl b/lib/OrdinaryDiffEqNonlinearSolve/test/nsa_sparse_tests.jl new file mode 100644 index 00000000000..f8952dcb7ca --- /dev/null +++ b/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 diff --git a/lib/OrdinaryDiffEqNonlinearSolve/test/runtests.jl b/lib/OrdinaryDiffEqNonlinearSolve/test/runtests.jl index e770ad84d62..ba11623361e 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/test/runtests.jl +++ b/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)