NonlinearSolveAlg OOP: extend W-reuse to WOperator/StaticWOperator

lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl

@@ -14,7 +14,8 @@ import DiffEqBase: OrdinaryDiffEqTag, calculate_residuals, calculate_residuals!,
     BrownFullBasicInit, ShampineCollocationInit
 import ConstructionBase
 import PreallocationTools: DiffCache, get_tmp
-using SimpleNonlinearSolve: SimpleTrustRegion, SimpleGaussNewton
+using SimpleNonlinearSolve: SimpleTrustRegion, SimpleGaussNewton,
+    AbstractSimpleNonlinearSolveAlgorithm
 using NonlinearSolve: FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg, NewtonRaphson,
     step!
 using MuladdMacro: @muladd
@@ -56,7 +57,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

@@ -41,11 +41,29 @@ function initialize!(
     cache.tstep = integrator.t + nlsolver.c * dt
 
     (; ustep, tstep, k, invγdt) = cache
-    if SciMLBase.has_stats(integrator)
+    if SciMLBase.has_stats(integrator) && hasfield(typeof(cache.cache), :stats)
         integrator.stats.nf += cache.cache.stats.nf
         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,22 @@ 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
+    J_new = if SciMLBase.has_jac(f)
+        f.jac(uprev, p, tstep)
+    else
+        uf.f = nlsolve_f(f, alg)
+        uf.p = p
+        uf.t = tstep
+        jacobian(uf, uprev, integrator)
+    end
+    nlcache.W[] = StaticWOperator(J_new - mass_matrix * inv(dtgamma))
+    return nothing
+end
+
 ## compute_step!
 
 @muladd function compute_step!(nlsolver::NLSolver{<:NonlinearSolveAlg, false}, integrator)
@@ -701,7 +735,8 @@ function Base.resize!(nlcache::NonlinearSolveCache, ::AbstractNLSolver, integrat
     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)
+    # Ref{StaticWOperator} = OOP path, fixed-size, not resizable.
+    nlcache.W === nothing || nlcache.W isa Ref || resize_J_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,38 @@ function build_nlsolver(
             γ = tTypeNoUnits(γ)
             α = tTypeNoUnits(α)
             dt = tTypeNoUnits(dt)
+            use_w_reuse = !isdae && f.nlstep_data === nothing &&
+                W isa StaticWOperator &&
+                nlalg.alg isa AbstractSimpleNonlinearSolveAlgorithm
             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
+                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
+                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,
+                nothing, W_ref, use_w_reuse ? uf : nothing,
+                nothing, nothing, nothing,
                 zero(tstep), true
             )
         else

lib/OrdinaryDiffEqNonlinearSolve/test/linear_nonlinear_tests.jl

@@ -1,7 +1,8 @@
 using OrdinaryDiffEqSDIRK, OrdinaryDiffEqBDF, OrdinaryDiffEqRosenbrock, Test, Random,
-    LinearAlgebra, LinearSolve, ADTypes
+    LinearAlgebra, LinearSolve, ADTypes, SciMLBase
 using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
 using NonlinearSolve: NewtonRaphson
+using SimpleNonlinearSolve: SimpleNewtonRaphson
 Random.seed!(123)
 
 A = 0.01 * rand(3, 3)
@@ -208,3 +209,20 @@ let integ = init(
     step!(integ)
     @test !iszero(integ.cache.nlsolver.cache.weight)
 end
+
+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(SimpleNewtonRaphson(; 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