Add adaptive Jacobian reuse to first-order solvers

Project.toml

@@ -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
 
@@ -104,6 +108,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

lib/NonlinearSolveBase/src/utils.jl

@@ -259,7 +259,7 @@ function linsolve_workspace(A::AbstractMatrix)
         stats = SciMLBase.NLStats(0, 0, 0, 0, 0),
         alias = SciMLBase.LinearAliasSpecifier(alias_A = true, alias_b = true)
     )
-    return (; lincache, rhs), A
+    return (; lincache, rhs, A_buf), A
 end
 
 # scalar analog of the default solver's least-squares rescue: a singular (zero) entry
@@ -300,8 +300,32 @@ function linsolve_identity!!(workspace, A::AbstractMatrix)
     # buffer, so A is never mutated and passing the previously returned inverse as A is
     # safe. On singular A the default algorithm's pivoted-QR rescue returns a finite
     # least-squares generalized inverse (not the SVD `pinv`).
+    # The triangular solve returns `rhs` itself, so a nested reinversion can pass that
+    # buffer back as `A`. Preserve its contents before refilling the identity RHS.
+    A_solve = if A === workspace.rhs
+        copyto!(workspace.A_buf, A)
+        workspace.A_buf
+    else
+        A
+    end
     make_identity!!(workspace.rhs, true)
-    return workspace.lincache(; A, b = workspace.rhs).u
+    if A_solve isa StridedMatrix
+        diagonal = @view A_solve[LinearAlgebra.diagind(A_solve)]
+        nonsingular = !any(iszero, diagonal)
+        # Preserve exact triangular structure instead of allowing pivoted LU roundoff to
+        # alter sensitive quasi-Newton trajectories. Singular matrices still need the
+        # default solver's pivoted-QR rescue below.
+        if nonsingular && LinearAlgebra.istriu(A_solve)
+            return LinearAlgebra.ldiv!(
+                LinearAlgebra.UpperTriangular(A_solve), workspace.rhs
+            )
+        elseif nonsingular && LinearAlgebra.istril(A_solve)
+            return LinearAlgebra.ldiv!(
+                LinearAlgebra.LowerTriangular(A_solve), workspace.rhs
+            )
+        end
+    end
+    return workspace.lincache(; A = A_solve, b = workspace.rhs).u
 end
 
 function initial_jacobian_scaling_alpha(α, u, fu, ::Any)

lib/NonlinearSolveBase/src/verbosity.jl

@@ -9,7 +9,8 @@ diagnostic messages, warnings, and errors during nonlinear system solution.
 ## Error Control Group
 - `non_enclosing_interval`: Messages when interval doesn't enclose root (bracketing methods)
 - `alias_u0_immutable`: Messages when aliasing u0 with immutable array
-- `linsolve_failed_noncurrent`: Messages when linear solve fails on non-current iteration
+- `linsolve_failed_noncurrent`: Messages when a linear solve or line search retries with
+  non-current Jacobian information
 - `termination_condition`: Messages about termination conditions
 - `unsupported_postcondition`: Messages when a `postcondition` corrector is supplied to a
   solver that cannot apply it (the corrector is then ignored)

lib/NonlinearSolveBase/test/linsolve_workspace.jl

@@ -46,6 +46,30 @@ end
     @test Utils.linsolve_identity!!(workspace, A) ≈ pinv(A.data)
 end
 
+@testset "dense triangular input preserves triangular solve semantics" begin
+    A = [1.1 0.0 0.0; 3.2 2.2 0.0; -4.1 5.3 3.3]
+    expected = Matrix{Float64}(I, 3, 3)
+    ldiv!(LowerTriangular(A), expected)
+    workspace, _ = Utils.linsolve_workspace(A)
+    @test Utils.linsolve_identity!!(workspace, A) == expected
+
+    A_reset = [0.1 0.0 0.0; 100.3 0.2 0.0; -44.1 55.3 0.3]
+    expected_reset = Matrix{Float64}(I, 3, 3)
+    ldiv!(LowerTriangular(A_reset), expected_reset)
+    @test Utils.linsolve_identity!!(workspace, A_reset) == expected_reset
+
+    A_general = [2.0 1.0 0.0; 0.0 3.0 1.0; 1.0 0.0 4.0]
+    copyto!(workspace.rhs, A_general)
+    @test Utils.linsolve_identity!!(workspace, workspace.rhs) ≈ inv(A_general)
+end
+
+@testset "arrays without fast scalar indexing use pinv" begin
+    A = NoFastScalarMatrix(rand(5, 5))
+    workspace, A_ret = Utils.linsolve_workspace(A)
+    @test workspace === nothing && A_ret === A
+    @test Utils.linsolve_identity!!(workspace, A) ≈ pinv(A.data)
+end
+
 @testset "singular input takes the pivoted-QR rescue" begin
     # The result is the LinearSolve default algorithm's least-squares generalized
     # inverse from its singular-LU → pivoted-QR rescue, NOT the SVD `pinv` (an

lib/NonlinearSolveBase/test/runtests.jl

@@ -124,9 +124,17 @@ run_tests(;
 
             # Vector{Float64} u0 — wraps.
             prob_f64 = NonlinearProblem(f, [1.0, 2.0], [0.5, 0.25])
-            @test NonlinearSolveBase.is_fw_wrapped(
-                NonlinearSolveBase.maybe_wrap_nonlinear_f(prob_f64)
-            )
+            wrapped_f64 = NonlinearSolveBase.maybe_wrap_nonlinear_f(prob_f64)
+            @test NonlinearSolveBase.is_fw_wrapped(wrapped_f64)
+
+            du_big = zeros(BigFloat, 2)
+            u_big = BigFloat[3, 5]
+            p_big = BigFloat[1, 2]
+            wrapped_f64(du_big, u_big, p_big)
+            @test du_big == u_big
+            fill!(du_big, 0)
+            wrapped_f64(du_big, u_big, p_big)
+            @test du_big == u_big
 
             # Vector{Dual} u0 — must NOT wrap.
             DualF = ForwardDiff.Dual{ForwardDiff.Tag{typeof(identity), Float64}, Float64, 2}

lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl

@@ -50,6 +50,7 @@ using SciMLJacobianOperators: VecJacOperator, JacVecOperator, StatefulJacobianOp
 using FiniteDiff: FiniteDiff    # Default Finite Difference Method
 using ForwardDiff: ForwardDiff, Dual  # Default Forward Mode AD
 
+include("jacobian_reuse.jl")
 include("solve.jl")
 include("raphson.jl")
 include("eisenstat_walker.jl")
@@ -145,6 +146,7 @@ export NewtonRaphson, PseudoTransient
 export GaussNewton, LevenbergMarquardt, TrustRegion
 
 export EisenstatWalkerForcing2
+export JacobianReuse
 
 export RadiusUpdateSchemes
 

lib/NonlinearSolveFirstOrder/src/gauss_newton.jl

@@ -1,22 +1,28 @@
 """
     GaussNewton(;
         concrete_jac = nothing, linsolve = nothing, linesearch = missing,
-        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
+        jacobian_reuse = nothing
     )
 
 An advanced GaussNewton implementation with support for efficient handling of sparse
 matrices via colored automatic differentiation and preconditioned linear solvers. Designed
 for large-scale and numerically-difficult nonlinear systems.
+
+Set `jacobian_reuse = JacobianReuse()` (or `true`) to adaptively reuse the Jacobian and
+factorization across accepted steps. Reuse is disabled by default.
 """
 function GaussNewton(;
         concrete_jac = nothing, linsolve = nothing, linesearch = missing,
-        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
+        jacobian_reuse = nothing
     )
     return GeneralizedFirstOrderAlgorithm(;
         linesearch,
         descent = NewtonDescent(; linsolve),
         autodiff, vjp_autodiff, jvp_autodiff,
         concrete_jac,
+        jacobian_reuse,
         name = :GaussNewton
     )
 end

lib/NonlinearSolveFirstOrder/src/jacobian_reuse.jl

@@ -0,0 +1,89 @@
+"""
+    JacobianReuse(; max_age::Int = 10, max_residual_ratio::Real = 1)
+
+Reuse a Jacobian across accepted nonlinear iterations. This turns a first-order method into
+an adaptive modified-Newton method: the current Jacobian is reused while the residual norm
+continues to improve, subject to a maximum Jacobian age. Solvers of an unchanged concrete
+linear system also reuse its factorization; damped and matrix-free systems retain their own
+linear-solver update behavior.
+
+The Jacobian is refreshed when either of these conditions holds:
+
+  - `max_age` accepted steps have used the current Jacobian;
+  - the new residual norm is not strictly less than `max_residual_ratio` times the previous
+    residual norm;
+  - a linear solve or globalization step fails with stale Jacobian information.
+
+`max_age = 1` recomputes the Jacobian after every accepted step. Setting
+`max_residual_ratio = Inf` selects purely periodic refreshes. The reuse state is reset by
+`reinit!`; retaining a Jacobian across separate nonlinear solves requires the manual
+`step!(cache; recompute_jacobian = false)` interface.
+
+Pass `jacobian_reuse = JacobianReuse()` (or `jacobian_reuse = true`) to
+[`NewtonRaphson`](@ref), [`TrustRegion`](@ref), or another first-order solver to enable the
+policy. Jacobian reuse is disabled by default.
+"""
+struct JacobianReuse{R <: Real}
+    max_age::Int
+    max_residual_ratio::R
+
+    function JacobianReuse(max_age::Int, max_residual_ratio::R) where {R <: Real}
+        max_age > 0 || throw(ArgumentError("`max_age` must be positive, got $max_age."))
+        max_residual_ratio >= 0 || throw(
+            ArgumentError(
+                "`max_residual_ratio` must be nonnegative, got $max_residual_ratio."
+            )
+        )
+        return new{R}(max_age, max_residual_ratio)
+    end
+end
+
+function JacobianReuse(; max_age::Int = 10, max_residual_ratio::Real = 1)
+    return JacobianReuse(max_age, max_residual_ratio)
+end
+
+normalize_jacobian_reuse(::Nothing) = nothing
+normalize_jacobian_reuse(reuse::JacobianReuse) = reuse
+normalize_jacobian_reuse(reuse::Bool) = reuse ? JacobianReuse() : nothing
+function normalize_jacobian_reuse(reuse)
+    throw(
+        ArgumentError(
+            "`jacobian_reuse` must be `nothing`, a `Bool`, or a `JacobianReuse`, got $(typeof(reuse))."
+        )
+    )
+end
+
+@concrete mutable struct JacobianReuseCache
+    residual_norm
+    age::Int
+    internalnorm
+end
+
+init_jacobian_reuse_cache(::Nothing, fu, internalnorm) = nothing
+function init_jacobian_reuse_cache(::JacobianReuse, fu, internalnorm)
+    return JacobianReuseCache(internalnorm(fu), 0, internalnorm)
+end
+
+reset_jacobian_reuse!(::Nothing, fu) = nothing
+function reset_jacobian_reuse!(cache::JacobianReuseCache, fu)
+    cache.residual_norm = cache.internalnorm(fu)
+    cache.age = 0
+    return nothing
+end
+
+mark_jacobian_refresh!(cache, fu) = reset_jacobian_reuse!(cache, fu)
+
+jacobian_is_stale(::Nothing) = false
+jacobian_is_stale(cache::JacobianReuseCache) = cache.age > 0
+
+function prepare_next_jacobian!(::Nothing, ::Nothing, fu)
+    return true
+end
+function prepare_next_jacobian!(cache::JacobianReuseCache, policy::JacobianReuse, fu)
+    residual_norm = cache.internalnorm(fu)
+    cache.age += 1
+    residual_improved = isfinite(residual_norm) && isfinite(cache.residual_norm) &&
+        residual_norm < policy.max_residual_ratio * cache.residual_norm
+    cache.residual_norm = residual_norm
+    return !(residual_improved && cache.age < policy.max_age)
+end

lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl

@@ -4,7 +4,8 @@
         damping_initial::Real = 1.0, α_geodesic::Real = 0.75, disable_geodesic = Val(false),
         damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0,
         finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, min_damping_D::Real = 1e-8,
-        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
+        jacobian_reuse = nothing
     )
 
 An advanced Levenberg-Marquardt implementation with the improvements suggested in
@@ -28,6 +29,8 @@ nonlinear systems.
   - `disable_geodesic`: Disables Geodesic Acceleration if set to `Val(true)`. It provides
     a way to trade-off robustness for speed, though in most situations Geodesic Acceleration
     should not be disabled.
+  - `jacobian_reuse`: a [`JacobianReuse`](@ref) policy, `true` for the default policy, or
+    `nothing`/`false` to recompute after every accepted step. Defaults to `nothing`.
 
 For the remaining arguments, see [`GeodesicAcceleration`](@ref) and
 [`NonlinearSolveFirstOrder.LevenbergMarquardtTrustRegion`](@ref) documentations.
@@ -37,7 +40,8 @@ function LevenbergMarquardt(;
         damping_initial::Real = 1.0, α_geodesic::Real = 0.75, disable_geodesic = Val(false),
         damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0,
         finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, min_damping_D::Real = 1.0e-8,
-        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
+        jacobian_reuse = nothing
     )
     descent = DampedNewtonDescent(;
         linsolve,
@@ -56,6 +60,7 @@ function LevenbergMarquardt(;
         autodiff,
         vjp_autodiff,
         jvp_autodiff,
+        jacobian_reuse,
         name = :LevenbergMarquardt,
         concrete_jac = Val(true)
     )

lib/NonlinearSolveFirstOrder/src/poly_algs.jl

@@ -3,7 +3,8 @@
         ::Type{T} = Float64;
         concrete_jac = nothing,
         linsolve = nothing,
-        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
+        jacobian_reuse = nothing
     )
 
 A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different
@@ -18,14 +19,18 @@ or more precision / more stable linear solver choice is required).
 
   - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
     are compatible with the problem type. Defaults to `Float64`.
+  - `jacobian_reuse`: forwarded to each first-order method in the polyalgorithm.
 """
 function RobustMultiNewton(
         ::Type{T} = Float64;
         concrete_jac = nothing,
         linsolve = nothing,
-        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
+        jacobian_reuse = nothing
     ) where {T}
-    common_kwargs = (; concrete_jac, linsolve, autodiff, vjp_autodiff, jvp_autodiff)
+    common_kwargs = (;
+        concrete_jac, linsolve, autodiff, vjp_autodiff, jvp_autodiff, jacobian_reuse,
+    )
     if T <: Complex # Let's atleast have something here for complex numbers
         algs = (
             NewtonRaphson(; common_kwargs...),
@@ -48,7 +53,8 @@ end
         ::Type{T} = Float64;
         concrete_jac = nothing,
         linsolve = nothing,
-        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
+        jacobian_reuse = nothing
     )
 
 A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods
@@ -58,14 +64,16 @@ for more performance and then tries more robust techniques if the faster ones fa
 
   - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
     are compatible with the problem type. Defaults to `Float64`.
+  - `jacobian_reuse`: forwarded to each first-order method in the polyalgorithm.
 """
 function FastShortcutNLLSPolyalg(
         ::Type{T} = Float64;
         concrete_jac = nothing,
         linsolve = nothing,
-        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
+        jacobian_reuse = nothing
     ) where {T}
-    common_kwargs = (; linsolve, autodiff, vjp_autodiff, jvp_autodiff)
+    common_kwargs = (; linsolve, autodiff, vjp_autodiff, jvp_autodiff, jacobian_reuse)
     if T <: Complex
         algs = (
             GaussNewton(; common_kwargs..., concrete_jac),