Add adaptive Jacobian reuse to first-order solvers

Project.toml

@@ -1,6 +1,6 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
-version = "4.21.1"
+version = "4.22.0"
 authors = ["SciML"]
 
 [deps]
@@ -91,15 +91,15 @@ LeastSquaresOptim = "0.8.5"
 LineSearch = "0.1.9"
 LineSearches = "7.3"
 LinearAlgebra = "1.10"
-LinearSolve = "4.3, 5"
+LinearSolve = "5"
 MINPACK = "1.2"
 MPI = "0.20.22"
 NLSolvers = "0.5"
 NLsolve = "4.5"
 NaNMath = "1"
 NonlinearProblemLibrary = "0.1.2"
 NonlinearSolveBase = "2.35"
-NonlinearSolveFirstOrder = "2"
+NonlinearSolveFirstOrder = "2.3"
 NonlinearSolveQuasiNewton = "1.12"
 NonlinearSolveSpectralMethods = "1.6"
 PETSc = "0.4.2"
@@ -134,6 +134,7 @@ NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
+Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLLogging = "a6db7da4-7206-11f0-1eab-35f2a5dbe1d1"
@@ -146,4 +147,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["InteractiveUtils", "NaNMath", "NonlinearProblemLibrary", "Pkg", "PolyesterForwardDiff", "Random", "SafeTestsets", "SciMLLogging", "SciMLOperators", "SciMLTesting", "SparseArrays", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test"]
+test = ["InteractiveUtils", "NaNMath", "NonlinearProblemLibrary", "Pkg", "PolyesterForwardDiff", "Profile", "Random", "SafeTestsets", "SciMLLogging", "SciMLOperators", "SciMLTesting", "SparseArrays", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test"]

docs/src/native/solvers.md

@@ -39,6 +39,10 @@ documentation.
     iterative linear solver (Krylov method), this controls how accurately the linear system
     is solved at each Newton iteration. Defaults to `nothing` (fixed tolerance). See
     [Forcing Term Strategies](@ref forcing_strategies) for available options.
+  - `jacobian_reuse`: controls whether a Jacobian can be reused across accepted nonlinear
+    iterations. `nothing` or `false` (the default) uses a fresh Jacobian after every accepted
+    step. `true` selects [`JacobianReuse()`](@ref), or a configured `JacobianReuse` policy can
+    be supplied directly. An unchanged concrete linear system also reuses its factorization.
 
 ## Nonlinear Solvers
 
@@ -94,6 +98,33 @@ GeneralizedFirstOrderAlgorithm
 GeneralizedDFSane
 ```
 
+## Jacobian Reuse
+
+```@docs
+JacobianReuse
+```
+
+Jacobian reuse is most useful when constructing or factorizing the Jacobian dominates the
+cost of evaluating the residual. It changes exact Newton iteration into a modified-Newton
+iteration, which can require more nonlinear steps, so it is opt-in. For example:
+
+```julia
+sol = solve(prob, NewtonRaphson(jacobian_reuse = JacobianReuse()))
+```
+
+The same policy works with `TrustRegion`, `GaussNewton`, `LevenbergMarquardt`, and
+`PseudoTransient`. Damped and matrix-free systems retain their normal linear-solver update
+behavior. Rejected trust-region steps reuse a fresh Jacobian because the nonlinear state did
+not change; a rejected step based on stale Jacobian information requests a refresh.
+
+The policy is local to one nonlinear cache lifecycle and is reset by `reinit!`. An outer
+solver that owns a related but distinct operator should keep using the explicit
+`step!(cache; recompute_jacobian = ...)` interface. In particular,
+OrdinaryDiffEqNonlinearSolve distinguishes the ODE Jacobian `J` from
+the iteration matrix `W` assembled from `J`, the mass matrix, `γ`, and `dt`; it decides
+independently when each must be rebuilt and retains convergence information across time
+steps. Its explicit decision takes precedence over this standalone policy.
+
 ## [Forcing Term Strategies](@id forcing_strategies)
 
 Forcing term strategies control how accurately the linear system is solved at each Newton

docs/src/refs.bib

@@ -137,16 +137,6 @@ @article{yuan2015recent
   publisher = {Springer}
 }
 
-@article{yuan2015recent,
-  title     = {Recent advances in trust region algorithms},
-  author    = {Yuan, Ya-xiang},
-  journal   = {Mathematical Programming},
-  volume    = {151},
-  pages     = {249--281},
-  year      = {2015},
-  publisher = {Springer}
-}
-
 @article{ziani2008autoadaptative,
   title     = {An autoadaptative limited memory Broyden’s method to solve systems of nonlinear equations},
   author    = {Ziani, Mohammed and Guyomarc’h, Fr{\'e}d{\'e}ric},
@@ -156,4 +146,4 @@ @article{ziani2008autoadaptative
   pages     = {202--211},
   year      = {2008},
   publisher = {Elsevier}
-}
+}

lib/NonlinearSolveBase/Project.toml

@@ -1,6 +1,6 @@
 name = "NonlinearSolveBase"
 uuid = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
-version = "2.35.0"
+version = "2.35.3"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]
@@ -78,7 +78,7 @@ EnzymeCore = "0.8.16"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.36, 1"
 FunctionWrappers = "1.1.2"
-FunctionWrappersWrappers = "1"
+FunctionWrappersWrappers = "1.9.3"
 InteractiveUtils = "<0.0.1, 1"
 LineSearch = "0.1.4"
 LinearAlgebra = "1.10"

lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl

@@ -21,6 +21,20 @@ import NonlinearSolveBase: wrapfun_iip, standardize_forwarddiff_tag
 
 const DI = DifferentiationInterface
 
+_cache_storage_type(
+    ::FunctionWrappersWrappers.FunctionWrappersWrapper{FW, P, CS}
+) where {FW, P, CS} = CS
+
+# The concrete storage types are intentionally not public. Derive the type selected by the
+# public cache-mode API, then use the documented explicit-parameter constructor below.
+const FWW_SINGLE_CACHE_STORAGE_TYPE = _cache_storage_type(
+    FunctionWrappersWrappers.FunctionWrappersWrapper(
+        identity, (Tuple{Float64},), (Float64,);
+        cache = FunctionWrappersWrappers.SingleCache(),
+        policy = FunctionWrappersWrappers.AllowNonIsBits(),
+    )
+)
+
 # --- AutoSpecialize / norecompile infrastructure for ForwardDiff ---
 
 const dualT = ForwardDiff.Dual{
@@ -78,14 +92,13 @@ function _make_fww_iip(
         @nospecialize(vff), ::Type{A1}, ::Type{A2}, ::Type{A3}, ::Type{A4}
     ) where {A1, A2, A3, A4}
     FW = FunctionWrappers.FunctionWrapper
-    fwt = (
-        FW{Nothing, A1}(vff), FW{Nothing, A2}(vff),
-        FW{Nothing, A3}(vff), FW{Nothing, A4}(vff),
-    )
-    cs = FunctionWrappersWrappers.SingleCacheStorage()
+    FWT = Tuple{
+        FW{Nothing, A1}, FW{Nothing, A2},
+        FW{Nothing, A3}, FW{Nothing, A4},
+    }
     return FunctionWrappersWrappers.FunctionWrappersWrapper{
-        typeof(fwt), FunctionWrappersWrappers.AllowNonIsBits, typeof(cs),
-    }(fwt, cs)
+        FWT, FunctionWrappersWrappers.AllowNonIsBits, FWW_SINGLE_CACHE_STORAGE_TYPE,
+    }(vff)
 end
 
 # IIP wrapfun: wraps f(du, u, p) with dual-aware type combinations.

lib/NonlinearSolveBase/src/arclength.jl

@@ -829,10 +829,10 @@ function CommonSolve.solve(
         # and discovers the winner) instead of re-failing the cheaper ladder members
         # on every warm-started step.
         reinit_retaining!(corr_cache, xpred)
-        last_sol = CommonSolve.solve!(corr_cache)
+        last_sol = _solve_without_solution!(corr_cache)
 
-        if SciMLBase.successful_retcode(last_sol)
-            xnew = last_sol.u
+        if _solve_result_successful(last_sol)
+            xnew = _solve_result_u(last_sol)
             # realized chord, built in the (currently free) secant scratch
             @. tscratch = xnew - xcur
             nchord = _theta_norm(tscratch, wu, wλ, n)
@@ -892,7 +892,8 @@ function CommonSolve.solve(
             # sweep — see `_effort_growth_factor`), giving the arclength driver an effort
             # signal alongside the purely geometric bend-angle test.
             if alg.adaptive
-                nit = last_sol.stats === nothing ? -1 : Int(last_sol.stats.nsteps)
+                result_stats = _solve_result_stats(last_sol)
+                nit = result_stats === nothing ? -1 : Int(result_stats.nsteps)
                 if _effort_wants_shrink(nit, budget)
                     ds / 2 >= min_ds && (ds = ds / 2)
                     streak = 0
@@ -913,7 +914,8 @@ function CommonSolve.solve(
         else
             return build_solution_less_specialize(
                 prob, alg, u, nothing;
-                retcode = last_sol.retcode, original = last_sol,
+                retcode = _solve_result_retcode(last_sol),
+                original = _solve_result_original(last_sol),
                 store_original = alg.store_original
             )
         end

lib/NonlinearSolveBase/src/forward_diff.jl

@@ -11,6 +11,9 @@ get_u(cache::NonlinearSolveForwardDiffCache) = get_u(cache.cache)
 get_fu(cache::NonlinearSolveForwardDiffCache) = get_fu(cache.cache)
 set_fu!(cache::NonlinearSolveForwardDiffCache, fu) = set_fu!(cache.cache, fu)
 SciMLBase.set_u!(cache::NonlinearSolveForwardDiffCache, u) = SciMLBase.set_u!(cache.cache, u)
+function _solve_without_solution!(cache::NonlinearSolveForwardDiffCache)
+    return CommonSolve.solve!(cache)
+end
 function NonlinearSolveBase.get_abstol(cache::NonlinearSolveForwardDiffCache)
     return NonlinearSolveBase.get_abstol(cache.cache)
 end

lib/NonlinearSolveBase/src/homotopy_sweep.jl

@@ -553,10 +553,10 @@ function _homotopy_sweep_solve(
         # won the previous solve (the anchor's full-ladder run discovers the winner)
         # instead of re-failing the cheaper ladder members on every warm-started step.
         reinit_retaining!(cache, guess)
-        last_sol = CommonSolve.solve!(cache)
+        last_sol = _solve_without_solution!(cache)
 
         if next_λ == λend
-            if cap_active && !SciMLBase.successful_retcode(last_sol)
+            if cap_active && !_solve_result_successful(last_sol)
                 # The final landing on λspan[2] is exempt from the tracking cap: give
                 # it the full user budget before letting the failure feed the
                 # bisection logic. The prediction is recomputed (never reuse `guess`:
@@ -571,19 +571,19 @@ function _homotopy_sweep_solve(
                 last_sol = _sweep_exempt_solve(
                     prob, alg.inner, retry_guess, next_λ, args...; kwargs...
                 )
-            elseif tol_active && SciMLBase.successful_retcode(last_sol)
+            elseif tol_active && _solve_result_successful(last_sol)
                 # The landing is exempt from the loose tracking tolerance: the loose
                 # cache solve above did the bulk of the convergence, now re-polish at
                 # the user's full tolerances warm-started from it (~1–2 corrector
                 # iterations), so the returned solution's accuracy semantics are
                 # unchanged. A failed polish feeds the ordinary bisection logic.
                 last_sol = _sweep_exempt_solve(
-                    prob, alg.inner, last_sol.u, next_λ, args...; kwargs...
+                    prob, alg.inner, _solve_result_u(last_sol), next_λ, args...; kwargs...
                 )
             end
         end
 
-        if SciMLBase.successful_retcode(last_sol)
+        if _solve_result_successful(last_sol)
             # The secant prediction error θ (relative to the recent solution movement)
             # is a cheap local error estimate: it grows with the path's curvature times
             # dλ², so a large θ means the path is turning and a stale tangent would
@@ -599,24 +599,26 @@ function _homotopy_sweep_solve(
                 # have iterated in place in the guess buffer when aliasing is forwarded
                 sv = (next_λ - λ) / (λ - λ_prev)
                 virtual = _sweep_extrapolate!(virtual, u, u_prev, sv)
-                correction = Utils.norm_op(L2_NORM, -, last_sol.u, virtual)
-                disp = Utils.norm_op(L2_NORM, -, last_sol.u, u)
-                scale = max(disp, disp_prev, sqrt(eps(λT)) * (1 + L2_NORM(last_sol.u)))
+                result_u = _solve_result_u(last_sol)
+                correction = Utils.norm_op(L2_NORM, -, result_u, virtual)
+                disp = Utils.norm_op(L2_NORM, -, result_u, u)
+                scale = max(disp, disp_prev, sqrt(eps(λT)) * (1 + L2_NORM(result_u)))
                 θ = correction / scale
                 # the secant only earns its keep when it predicts at least twice as
                 # well as the trivial constant prediction (whose θ is exactly 1)
                 trust = θ < 1 / 2 ? trust + 1 : 0
                 disp_prev = disp
             else
-                disp_prev = Utils.norm_op(L2_NORM, -, last_sol.u, u)
+                disp_prev = Utils.norm_op(L2_NORM, -, _solve_result_u(last_sol), u)
             end
             # accept: swap `u`↔`u_prev` and copy the solution into `u` (no allocation).
-            u, u_prev = _sweep_accept!(u, u_prev, last_sol.u)
+            u, u_prev = _sweep_accept!(u, u_prev, _solve_result_u(last_sol))
             λ_prev = λ
             λ = next_λ
             λ == λend && break
             if alg.adaptive
-                nit = last_sol.stats === nothing ? -1 : Int(last_sol.stats.nsteps)
+                result_stats = _solve_result_stats(last_sol)
+                nit = result_stats === nothing ? -1 : Int(result_stats.nsteps)
                 if _effort_wants_shrink(nit, budget)
                     # Proactive shrink on a straining success (see
                     # `_effort_wants_shrink`); the floor guard keeps the increment
@@ -657,15 +659,16 @@ function _homotopy_sweep_solve(
         else
             # on failure: u is the last converged iterate (λ<λ1); resid is from the failed step (advisory)
             return build_solution_less_specialize(
-                    prob, alg, u, last_sol.resid;
-                    retcode = last_sol.retcode, original = last_sol,
+                    prob, alg, u, _solve_result_resid(last_sol);
+                    retcode = _solve_result_retcode(last_sol),
+                    original = _solve_result_original(last_sol),
                     store_original = alg.store_original
                 ), λ
         end
     end
 
     return SciMLBase.build_solution(
-            prob, alg, u, last_sol.resid; retcode = ReturnCode.Success
+            prob, alg, u, _solve_result_resid(last_sol); retcode = ReturnCode.Success
         ), nothing
 end
 

lib/NonlinearSolveBase/src/jacobian.jl

@@ -125,7 +125,8 @@ function construct_jacobian_cache(
         end
     end
 
-    return JacobianCache(J, f, fu, p, stats, autodiff, di_extras)
+    J_destination = jacobian_destination(J, autodiff)
+    return JacobianCache(J, J_destination, f, fu, p, stats, autodiff, di_extras)
 end
 
 function construct_jacobian_cache(
@@ -134,7 +135,7 @@ function construct_jacobian_cache(
         linsolve = missing
     )
     if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
-        return JacobianCache(fu, f, fu, p, stats, autodiff, nothing)
+        return JacobianCache(fu, fu, f, fu, p, stats, autodiff, nothing)
     end
     if autodiff === nothing
         throw(ArgumentError("`autodiff` argument to `construct_jacobian_cache` must be \
@@ -145,11 +146,33 @@ function construct_jacobian_cache(
     @assert !(autodiff isa AutoSparse) "`autodiff` cannot be `AutoSparse` for scalar \
                                         nonlinear problems."
     di_extras = DI.prepare_derivative(f, autodiff, u, Constant(prob.p))
-    return JacobianCache(u, f, fu, p, stats, autodiff, di_extras)
+    return JacobianCache(u, u, f, fu, p, stats, autodiff, di_extras)
 end
 
+struct FixedShapeJacobianDestination{T, A <: Matrix{T}} <: AbstractMatrix{T}
+    parent::A
+end
+
+Base.size(destination::FixedShapeJacobianDestination) = size(destination.parent)
+Base.axes(destination::FixedShapeJacobianDestination) = axes(destination.parent)
+Base.IndexStyle(::Type{<:FixedShapeJacobianDestination}) = IndexLinear()
+@inline Base.getindex(destination::FixedShapeJacobianDestination, indices...) =
+    getindex(destination.parent, indices...)
+@inline Base.setindex!(destination::FixedShapeJacobianDestination, value, indices...) =
+    setindex!(destination.parent, value, indices...)
+function Base.reshape(destination::FixedShapeJacobianDestination, dims::Dims)
+    size(destination) == dims || throw(DimensionMismatch("cannot reshape Jacobian destination"))
+    return destination
+end
+Base.reshape(destination::FixedShapeJacobianDestination, dims::Int...) =
+    reshape(destination, dims)
+
+jacobian_destination(J::Matrix, ::AutoForwardDiff) = FixedShapeJacobianDestination(J)
+jacobian_destination(J, autodiff) = J
+
 @concrete mutable struct JacobianCache <: AbstractJacobianCache
     J
+    J_destination
     f <: NonlinearFunction
     fu
     p
@@ -219,7 +242,8 @@ function (cache::JacobianCache)(u)
             f.jac(J, u, p)
         else
             DI.jacobian!(
-                f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p)
+                f, cache.fu, cache.J_destination, cache.di_extras, cache.autodiff, u,
+                Constant(p)
             )
         end
         return J