OOP NSA W-reuse: pass jac_prototype + refresh J each should_update

Fixes two defects in the OOP W-reuse path: (1) closure handed raw WOperator to NonlinearSolve whose jac cache is Matrix, causing convert(Matrix, WOperator) MethodError; (2) J stayed at build-time since update_coefficients!(W, ...) is a no-op for concrete Matrix J. Now stores W in a Ref, passes jac_prototype to inner NonlinearFunction, and recomputes J via calc_J semantics inside should_update. Tests cover Vector and StaticArrays paths.

Author: singhharsh1708

lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl

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

lib/OrdinaryDiffEqNonlinearSolve/src/newton.jl

@@ -54,7 +54,7 @@ function initialize!(
             nlsolver.status === Divergence ||
             abs(inv(dtgamma) / inv(W_γdt) - 1) > alg.new_W_dt_cutoff
         if should_update
-            update_coefficients!(cache.W, uprev, p, tstep; gamma = dtgamma)
+            _update_nlsolvealg_W_oop!(cache, integrator, dtgamma, tstep)
             cache.W_γdt = dtgamma
             integrator.stats.nw += 1
             cache.new_W = true
@@ -155,6 +155,31 @@ function _update_nlsolvealg_W!(nlcache, integrator, dtgamma, tstep)
     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)

lib/OrdinaryDiffEqNonlinearSolve/src/utils.jl

@@ -508,20 +508,31 @@ function build_nlsolver(
             else
                 (tmp, γ, α, tstep, invγdt, DIRK, p, dt, f)
             end
+            W_ref = use_w_reuse ? Ref{typeof(W)}(W) : nothing
             prob = if use_w_reuse
-                nlf_jac = if W isa StaticWOperator
-                    let Ws = W
-                        (z, p) -> Ws.W
+                if W isa StaticWOperator
+                    nlf_jac = let W_ref = W_ref
+                        (z, p) -> W_ref[].W
                     end
+                    jac_prototype = W.W
+                    NonlinearProblem(
+                        NonlinearFunction{false, SciMLBase.FullSpecialize}(
+                            nlf; jac = nlf_jac, jac_prototype = jac_prototype
+                        ),
+                        copy(ztmp), nlp_params
+                    )
                 else
-                    let Ww = W
-                        (z, p) -> Ww
+                    nlf_jac = let W_ref = W_ref
+                        (z, p) -> convert(AbstractMatrix, W_ref[])
                     end
+                    jac_prototype = copy(convert(AbstractMatrix, W))
+                    NonlinearProblem(
+                        NonlinearFunction{false, SciMLBase.FullSpecialize}(
+                            nlf; jac = nlf_jac, jac_prototype = jac_prototype
+                        ),
+                        copy(ztmp), nlp_params
+                    )
                 end
-                NonlinearProblem(
-                    NonlinearFunction{false, SciMLBase.FullSpecialize}(nlf; jac = nlf_jac),
-                    copy(ztmp), nlp_params
-                )
             else
                 NonlinearProblem(
                     NonlinearFunction{false, SciMLBase.FullSpecialize}(nlf),
@@ -533,8 +544,8 @@ function build_nlsolver(
             nlcache = NonlinearSolveCache(
                 nothing, tstep, nothing, nothing, invγdt, prob, cache,
                 use_w_reuse ? J : nothing,
-                use_w_reuse ? W : nothing,
-                nothing, nothing, nothing,
+                W_ref,
+                use_w_reuse ? uf : 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, NonlinearSolve, SciMLBase
 Random.seed!(123)
 
 A = 0.01 * rand(3, 3)
@@ -95,6 +95,46 @@ sol = @test_nowarn solve(
 );
 @test length(sol.t) < 20
 
+# NonlinearSolveAlg OOP W-reuse: plain Vector u0 with concrete_jac=true must
+# not crash on the first Jacobian evaluation. Regression for two defects
+# flagged on #3798: (a) the (z, p) -> W closure handed a raw WOperator to
+# NonlinearSolve whose jac cache is a dense Matrix, causing a
+# `convert(Matrix, WOperator)` MethodError, and (b) J was frozen at
+# build_nlsolver-time since `update_coefficients!(W, ...)` is a no-op when
+# W.J is a concrete Matrix.
+using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
+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 = 1e-8, abstol = 1e-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 = 1e-8, abstol = 1e-10
+    )
+    @test SciMLBase.successful_retcode(sol.retcode)
+    @test sol.stats.njacs > 0
+end
+
 sol = @test_nowarn solve(prob, TRBDF2(autodiff = AutoFiniteDiff()));
 @test length(sol.t) < 20
 sol = @test_nowarn solve(