OOP NonlinearSolveAlg: extend W-reuse to WOperator/StaticWOperator + refresh J each should_update

Author: Your Name

lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl

@@ -55,7 +55,7 @@ import OrdinaryDiffEqCore: _initialize_dae!,
 import OrdinaryDiffEqDifferentiation: update_W!, is_always_new, build_uf, build_J_W,
     WOperator, StaticWOperator, wrapprecs,
     build_jac_config, dolinsolve,
-    resize_jac_config!, jacobian2W!, jacobian!
+    resize_jac_config!, jacobian2W!, jacobian!, jacobian
 
 import StaticArraysCore: StaticArray
 

lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl

@@ -46,6 +46,24 @@ function initialize!(
         integrator.stats.njacs += cache.cache.stats.njacs
         integrator.stats.nsolve += cache.cache.stats.nsolve
     end
+
+    if cache.W !== nothing
+        dtgamma = method === DIRK ? γ * dt : γ * dt / α
+        W_γdt = cache.W_γdt
+        first_call = iszero(W_γdt)
+        should_update = first_call || alg.always_new ||
+            nlsolver.status === Divergence ||
+            abs(inv(dtgamma) / inv(W_γdt) - 1) > alg.new_W_dt_cutoff
+        if should_update
+            _update_nlsolvealg_W_oop!(cache, integrator, dtgamma, tstep)
+            cache.W_γdt = dtgamma
+            integrator.stats.nw += 1
+            cache.new_W = true
+        else
+            cache.new_W = false
+        end
+    end
+
     if f isa DAEFunction
         nlp_params = (tmp, α, tstep, invγdt, p, dt, uprev, f)
     else
@@ -153,6 +171,31 @@ function _update_nlsolvealg_W!(nlcache, integrator, dtgamma, tstep, new_jac = tr
     return nothing
 end
 
+function _update_nlsolvealg_W_oop!(nlcache, integrator, dtgamma, tstep)
+    (; f, p, uprev, alg) = integrator
+    mass_matrix = f.mass_matrix
+    uf = nlcache.uf
+    if SciMLBase.has_jac(f)
+        J_new = f.jac(uprev, p, tstep)
+    else
+        uf.f = nlsolve_f(f, alg)
+        uf.p = p
+        uf.t = tstep
+        J_new = jacobian(uf, uprev, integrator)
+    end
+    W = nlcache.W[]
+    if W isa StaticWOperator
+        nlcache.W[] = StaticWOperator(J_new - mass_matrix * inv(dtgamma))
+    else
+        copyto!(W.J, J_new)
+        W.gamma = dtgamma
+        if nlcache.J !== nothing && ArrayInterface.can_setindex(nlcache.J)
+            copyto!(nlcache.J, J_new)
+        end
+    end
+    return nothing
+end
+
 ## compute_step!
 
 @muladd function compute_step!(nlsolver::NLSolver{<:NonlinearSolveAlg, false}, integrator)
@@ -694,14 +737,26 @@ function Base.resize!(nlcache::NLNewtonCache, ::AbstractNLSolver, integrator, i:
     return nothing
 end
 
+function _resize_nlsolvealg_W!(nlcache, integrator, i)
+    W = nlcache.W[]
+    W isa StaticWOperator && return nothing
+    if nlcache.J !== nothing && ArrayInterface.can_setindex(nlcache.J)
+        nlcache.J = similar(nlcache.J, i, i)
+    end
+    nlcache.W[] = WOperator{SciMLBase.isinplace(integrator.sol.prob)}(
+        integrator.f.mass_matrix, integrator.dt, nlcache.J, integrator.u
+    )
+    return nothing
+end
+
 function Base.resize!(nlcache::NonlinearSolveCache, ::AbstractNLSolver, integrator, i::Int)
     nlcache.ustep === nothing || resize!(nlcache.ustep, i)
     nlcache.k === nothing || resize!(nlcache.k, i)
     nlcache.atmp === nothing || resize!(nlcache.atmp, i)
     nlcache.du1 === nothing || resize!(nlcache.du1, i)
     nlcache.weight === nothing || resize!(nlcache.weight, i)
     nlcache.jac_config === nothing || resize_jac_config!(nlcache, integrator)
-    nlcache.W === nothing || resize_J_W!(nlcache, integrator, i)
+    nlcache.W === nothing || _resize_nlsolvealg_W!(nlcache, integrator, i)
     nlcache.W_γdt = zero(nlcache.W_γdt)
     nlcache.new_W = true
     return nothing

lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl

@@ -499,21 +499,52 @@ function build_nlsolver(
             γ = tTypeNoUnits(γ)
             α = tTypeNoUnits(α)
             dt = tTypeNoUnits(dt)
+            use_w_reuse = !isdae && f.nlstep_data === nothing && (
+                W isa StaticWOperator ||
+                    (W isa WOperator && W.J !== nothing && !(W.J isa AbstractSciMLOperator))
+            )
             nlf = isdae ? oopdaenlf : oopodenlf
             nlp_params = if isdae
                 (tmp, α, tstep, invγdt, p, dt, uprev, f)
             else
                 (tmp, γ, α, tstep, invγdt, DIRK, p, dt, f)
             end
-            prob = NonlinearProblem(
-                NonlinearFunction{false, SciMLBase.FullSpecialize}(nlf),
-                copy(ztmp), nlp_params
-            )
+            W_ref = use_w_reuse ? Ref{typeof(W)}(W) : nothing
+            prob = if use_w_reuse
+                if W isa StaticWOperator
+                    nlf_jac = let W_ref = W_ref
+                        (z, p) -> W_ref[].W
+                    end
+                    NonlinearProblem(
+                        NonlinearFunction{false, SciMLBase.FullSpecialize}(
+                            nlf; jac = nlf_jac, jac_prototype = W.W
+                        ),
+                        copy(ztmp), nlp_params
+                    )
+                else
+                    nlf_jac = let W_ref = W_ref
+                        (z, p) -> convert(AbstractMatrix, W_ref[])
+                    end
+                    NonlinearProblem(
+                        NonlinearFunction{false, SciMLBase.FullSpecialize}(
+                            nlf; jac = nlf_jac, jac_prototype = W.J
+                        ),
+                        copy(ztmp), nlp_params
+                    )
+                end
+            else
+                NonlinearProblem(
+                    NonlinearFunction{false, SciMLBase.FullSpecialize}(nlf),
+                    copy(ztmp), nlp_params
+                )
+            end
             inner_alg = _nlalg_with_linsolve(nlalg.alg, alg.linsolve)
             cache = init(prob, inner_alg, verbose = verbose.nonlinear_verbosity)
             nlcache = NonlinearSolveCache(
                 nothing, tstep, nothing, nothing, invγdt, prob, cache,
-                nothing, nothing, nothing, nothing, nothing, nothing,
+                use_w_reuse ? J : nothing,
+                W_ref,
+                use_w_reuse ? uf : nothing, nothing, nothing, nothing,
                 zero(tstep), true
             )
         else

lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl

@@ -1,5 +1,5 @@
 using OrdinaryDiffEqSDIRK, OrdinaryDiffEqBDF, OrdinaryDiffEqRosenbrock, Test, Random,
-    LinearAlgebra, LinearSolve, ADTypes
+    LinearAlgebra, LinearSolve, ADTypes, SciMLBase
 using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
 using NonlinearSolve: NewtonRaphson
 Random.seed!(123)
@@ -208,3 +208,39 @@ let integ = init(
     step!(integ)
     @test !iszero(integ.cache.nlsolver.cache.weight)
 end
+
+# NonlinearSolveAlg OOP W-reuse regression: plain `Vector` u0 with
+# `concrete_jac = true` must not crash on the first Jacobian evaluation
+# (the raw-WOperator-passed-to-Matrix-cache defect) and J must be refreshed
+# each `should_update` (not frozen at build_nlsolver-time).
+let
+    vdp(u, p, t) = [u[2], p[1] * ((1 - u[1]^2) * u[2] - u[1])]
+    prob = ODEProblem(vdp, [2.0, 0.0], (0.0, 1.0), [1.0e3])
+    sol = solve(
+        prob,
+        TRBDF2(
+            nlsolve = NonlinearSolveAlg(NewtonRaphson(; autodiff = AutoForwardDiff())),
+            concrete_jac = true
+        ); reltol = 1.0e-8, abstol = 1.0e-10
+    )
+    @test SciMLBase.successful_retcode(sol.retcode)
+    @test sol.stats.njacs > 0
+end
+
+# StaticArrays / StaticWOperator variant of the same repro.
+using StaticArrays
+let
+    vdp_static(u, p, t) = SVector(u[2], p[1] * ((1 - u[1]^2) * u[2] - u[1]))
+    prob = ODEProblem(
+        vdp_static, SVector(2.0, 0.0), (0.0, 1.0), SVector(1.0e3)
+    )
+    sol = solve(
+        prob,
+        TRBDF2(
+            nlsolve = NonlinearSolveAlg(NewtonRaphson(; autodiff = AutoForwardDiff())),
+            concrete_jac = true
+        ); reltol = 1.0e-8, abstol = 1.0e-10
+    )
+    @test SciMLBase.successful_retcode(sol.retcode)
+    @test sol.stats.njacs > 0
+end