Skip to content

Commit 99348eb

Browse files
Add adaptive Jacobian reuse policy
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
1 parent 3f1ea6e commit 99348eb

17 files changed

Lines changed: 472 additions & 32 deletions

File tree

Project.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
name = "NonlinearSolve"
22
uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
3-
version = "4.21.0"
3+
version = "4.22.0"
44
authors = ["SciML"]
55

66
[deps]
@@ -99,7 +99,7 @@ NLsolve = "4.5"
9999
NaNMath = "1"
100100
NonlinearProblemLibrary = "0.1.2"
101101
NonlinearSolveBase = "2.31"
102-
NonlinearSolveFirstOrder = "2"
102+
NonlinearSolveFirstOrder = "2.3"
103103
NonlinearSolveQuasiNewton = "1.12"
104104
NonlinearSolveSpectralMethods = "1.6"
105105
PETSc = "0.4.2"

docs/src/native/solvers.md

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ documentation.
3939
iterative linear solver (Krylov method), this controls how accurately the linear system
4040
is solved at each Newton iteration. Defaults to `nothing` (fixed tolerance). See
4141
[Forcing Term Strategies](@ref forcing_strategies) for available options.
42+
- `jacobian_reuse`: controls whether a Jacobian can be reused across accepted nonlinear
43+
iterations. `nothing` or `false` (the default) uses a fresh Jacobian after every accepted
44+
step. `true` selects [`JacobianReuse()`](@ref), or a configured `JacobianReuse` policy can
45+
be supplied directly. An unchanged concrete linear system also reuses its factorization.
4246

4347
## Nonlinear Solvers
4448

@@ -94,6 +98,33 @@ GeneralizedFirstOrderAlgorithm
9498
GeneralizedDFSane
9599
```
96100

101+
## Jacobian Reuse
102+
103+
```@docs
104+
JacobianReuse
105+
```
106+
107+
Jacobian reuse is most useful when constructing or factorizing the Jacobian dominates the
108+
cost of evaluating the residual. It changes exact Newton iteration into a modified-Newton
109+
iteration, which can require more nonlinear steps, so it is opt-in. For example:
110+
111+
```julia
112+
sol = solve(prob, NewtonRaphson(jacobian_reuse = JacobianReuse()))
113+
```
114+
115+
The same policy works with `TrustRegion`, `GaussNewton`, `LevenbergMarquardt`, and
116+
`PseudoTransient`. Damped and matrix-free systems retain their normal linear-solver update
117+
behavior. Rejected trust-region steps reuse a fresh Jacobian because the nonlinear state did
118+
not change; a rejected step based on stale Jacobian information requests a refresh.
119+
120+
The policy is local to one nonlinear cache lifecycle and is reset by `reinit!`. An outer
121+
solver that owns a related but distinct operator should keep using the explicit
122+
`step!(cache; recompute_jacobian = ...)` interface. In particular,
123+
OrdinaryDiffEqNonlinearSolve distinguishes the ODE Jacobian `J` from
124+
the iteration matrix `W` assembled from `J`, the mass matrix, `γ`, and `dt`; it decides
125+
independently when each must be rebuilt and retains convergence information across time
126+
steps. Its explicit decision takes precedence over this standalone policy.
127+
97128
## [Forcing Term Strategies](@id forcing_strategies)
98129

99130
Forcing term strategies control how accurately the linear system is solved at each Newton

lib/NonlinearSolveBase/src/verbosity.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,8 @@ diagnostic messages, warnings, and errors during nonlinear system solution.
99
## Error Control Group
1010
- `non_enclosing_interval`: Messages when interval doesn't enclose root (bracketing methods)
1111
- `alias_u0_immutable`: Messages when aliasing u0 with immutable array
12-
- `linsolve_failed_noncurrent`: Messages when linear solve fails on non-current iteration
12+
- `linsolve_failed_noncurrent`: Messages when a linear solve or line search retries with
13+
non-current Jacobian information
1314
- `termination_condition`: Messages about termination conditions
1415
1516
## Numerical Group

lib/NonlinearSolveFirstOrder/Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "NonlinearSolveFirstOrder"
22
uuid = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d"
33
authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
4-
version = "2.2.1"
4+
version = "2.3.0"
55

66
[deps]
77
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ using SciMLJacobianOperators: VecJacOperator, JacVecOperator, StatefulJacobianOp
3131
using FiniteDiff: FiniteDiff # Default Finite Difference Method
3232
using ForwardDiff: ForwardDiff, Dual # Default Forward Mode AD
3333

34+
include("jacobian_reuse.jl")
3435
include("solve.jl")
3536
include("raphson.jl")
3637
include("eisenstat_walker.jl")
@@ -106,6 +107,7 @@ export NewtonRaphson, PseudoTransient
106107
export GaussNewton, LevenbergMarquardt, TrustRegion
107108

108109
export EisenstatWalkerForcing2
110+
export JacobianReuse
109111

110112
export RadiusUpdateSchemes
111113

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,28 @@
11
"""
22
GaussNewton(;
33
concrete_jac = nothing, linsolve = nothing, linesearch = missing,
4-
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
4+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
5+
jacobian_reuse = nothing
56
)
67
78
An advanced GaussNewton implementation with support for efficient handling of sparse
89
matrices via colored automatic differentiation and preconditioned linear solvers. Designed
910
for large-scale and numerically-difficult nonlinear systems.
11+
12+
Set `jacobian_reuse = JacobianReuse()` (or `true`) to adaptively reuse the Jacobian and
13+
factorization across accepted steps. Reuse is disabled by default.
1014
"""
1115
function GaussNewton(;
1216
concrete_jac = nothing, linsolve = nothing, linesearch = missing,
13-
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
17+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
18+
jacobian_reuse = nothing
1419
)
1520
return GeneralizedFirstOrderAlgorithm(;
1621
linesearch,
1722
descent = NewtonDescent(; linsolve),
1823
autodiff, vjp_autodiff, jvp_autodiff,
1924
concrete_jac,
25+
jacobian_reuse,
2026
name = :GaussNewton
2127
)
2228
end
Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
"""
2+
JacobianReuse(; max_age::Int = 10, max_residual_ratio::Real = 1)
3+
4+
Reuse a Jacobian across accepted nonlinear iterations. This turns a first-order method into
5+
an adaptive modified-Newton method: the current Jacobian is reused while the residual norm
6+
continues to improve, subject to a maximum Jacobian age. Solvers of an unchanged concrete
7+
linear system also reuse its factorization; damped and matrix-free systems retain their own
8+
linear-solver update behavior.
9+
10+
The Jacobian is refreshed when either of these conditions holds:
11+
12+
- `max_age` accepted steps have used the current Jacobian;
13+
- the new residual norm is not strictly less than `max_residual_ratio` times the previous
14+
residual norm;
15+
- a linear solve or globalization step fails with stale Jacobian information.
16+
17+
`max_age = 1` recomputes the Jacobian after every accepted step. Setting
18+
`max_residual_ratio = Inf` selects purely periodic refreshes. The reuse state is reset by
19+
`reinit!`; retaining a Jacobian across separate nonlinear solves requires the manual
20+
`step!(cache; recompute_jacobian = false)` interface.
21+
22+
Pass `jacobian_reuse = JacobianReuse()` (or `jacobian_reuse = true`) to
23+
[`NewtonRaphson`](@ref), [`TrustRegion`](@ref), or another first-order solver to enable the
24+
policy. Jacobian reuse is disabled by default.
25+
"""
26+
struct JacobianReuse{R <: Real}
27+
max_age::Int
28+
max_residual_ratio::R
29+
30+
function JacobianReuse(max_age::Int, max_residual_ratio::R) where {R <: Real}
31+
max_age > 0 || throw(ArgumentError("`max_age` must be positive, got $max_age."))
32+
max_residual_ratio >= 0 || throw(
33+
ArgumentError(
34+
"`max_residual_ratio` must be nonnegative, got $max_residual_ratio."
35+
)
36+
)
37+
return new{R}(max_age, max_residual_ratio)
38+
end
39+
end
40+
41+
function JacobianReuse(; max_age::Int = 10, max_residual_ratio::Real = 1)
42+
return JacobianReuse(max_age, max_residual_ratio)
43+
end
44+
45+
normalize_jacobian_reuse(::Nothing) = nothing
46+
normalize_jacobian_reuse(reuse::JacobianReuse) = reuse
47+
normalize_jacobian_reuse(reuse::Bool) = reuse ? JacobianReuse() : nothing
48+
function normalize_jacobian_reuse(reuse)
49+
throw(
50+
ArgumentError(
51+
"`jacobian_reuse` must be `nothing`, a `Bool`, or a `JacobianReuse`, got $(typeof(reuse))."
52+
)
53+
)
54+
end
55+
56+
@concrete mutable struct JacobianReuseCache
57+
residual_norm
58+
age::Int
59+
internalnorm
60+
end
61+
62+
init_jacobian_reuse_cache(::Nothing, fu, internalnorm) = nothing
63+
function init_jacobian_reuse_cache(::JacobianReuse, fu, internalnorm)
64+
return JacobianReuseCache(internalnorm(fu), 0, internalnorm)
65+
end
66+
67+
reset_jacobian_reuse!(::Nothing, fu) = nothing
68+
function reset_jacobian_reuse!(cache::JacobianReuseCache, fu)
69+
cache.residual_norm = cache.internalnorm(fu)
70+
cache.age = 0
71+
return nothing
72+
end
73+
74+
mark_jacobian_refresh!(cache, fu) = reset_jacobian_reuse!(cache, fu)
75+
76+
jacobian_is_stale(::Nothing) = false
77+
jacobian_is_stale(cache::JacobianReuseCache) = cache.age > 0
78+
79+
function prepare_next_jacobian!(::Nothing, ::Nothing, fu)
80+
return true
81+
end
82+
function prepare_next_jacobian!(cache::JacobianReuseCache, policy::JacobianReuse, fu)
83+
residual_norm = cache.internalnorm(fu)
84+
cache.age += 1
85+
residual_improved = isfinite(residual_norm) && isfinite(cache.residual_norm) &&
86+
residual_norm < policy.max_residual_ratio * cache.residual_norm
87+
cache.residual_norm = residual_norm
88+
return !(residual_improved && cache.age < policy.max_age)
89+
end

lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,8 @@
44
damping_initial::Real = 1.0, α_geodesic::Real = 0.75, disable_geodesic = Val(false),
55
damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0,
66
finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, min_damping_D::Real = 1e-8,
7-
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
7+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
8+
jacobian_reuse = nothing
89
)
910
1011
An advanced Levenberg-Marquardt implementation with the improvements suggested in
@@ -28,6 +29,8 @@ nonlinear systems.
2829
- `disable_geodesic`: Disables Geodesic Acceleration if set to `Val(true)`. It provides
2930
a way to trade-off robustness for speed, though in most situations Geodesic Acceleration
3031
should not be disabled.
32+
- `jacobian_reuse`: a [`JacobianReuse`](@ref) policy, `true` for the default policy, or
33+
`nothing`/`false` to recompute after every accepted step. Defaults to `nothing`.
3134
3235
For the remaining arguments, see [`GeodesicAcceleration`](@ref) and
3336
[`NonlinearSolveFirstOrder.LevenbergMarquardtTrustRegion`](@ref) documentations.
@@ -37,7 +40,8 @@ function LevenbergMarquardt(;
3740
damping_initial::Real = 1.0, α_geodesic::Real = 0.75, disable_geodesic = Val(false),
3841
damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0,
3942
finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, min_damping_D::Real = 1.0e-8,
40-
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
43+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
44+
jacobian_reuse = nothing
4145
)
4246
descent = DampedNewtonDescent(;
4347
linsolve,
@@ -56,6 +60,7 @@ function LevenbergMarquardt(;
5660
autodiff,
5761
vjp_autodiff,
5862
jvp_autodiff,
63+
jacobian_reuse,
5964
name = :LevenbergMarquardt,
6065
concrete_jac = Val(true)
6166
)

lib/NonlinearSolveFirstOrder/src/poly_algs.jl

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@
33
::Type{T} = Float64;
44
concrete_jac = nothing,
55
linsolve = nothing,
6-
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
6+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
7+
jacobian_reuse = nothing
78
)
89
910
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).
1819
1920
- `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
2021
are compatible with the problem type. Defaults to `Float64`.
22+
- `jacobian_reuse`: forwarded to each first-order method in the polyalgorithm.
2123
"""
2224
function RobustMultiNewton(
2325
::Type{T} = Float64;
2426
concrete_jac = nothing,
2527
linsolve = nothing,
26-
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
28+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
29+
jacobian_reuse = nothing
2730
) where {T}
28-
common_kwargs = (; concrete_jac, linsolve, autodiff, vjp_autodiff, jvp_autodiff)
31+
common_kwargs = (;
32+
concrete_jac, linsolve, autodiff, vjp_autodiff, jvp_autodiff, jacobian_reuse,
33+
)
2934
if T <: Complex # Let's atleast have something here for complex numbers
3035
algs = (
3136
NewtonRaphson(; common_kwargs...),
@@ -48,7 +53,8 @@ end
4853
::Type{T} = Float64;
4954
concrete_jac = nothing,
5055
linsolve = nothing,
51-
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
56+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
57+
jacobian_reuse = nothing
5258
)
5359
5460
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
5864
5965
- `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
6066
are compatible with the problem type. Defaults to `Float64`.
67+
- `jacobian_reuse`: forwarded to each first-order method in the polyalgorithm.
6168
"""
6269
function FastShortcutNLLSPolyalg(
6370
::Type{T} = Float64;
6471
concrete_jac = nothing,
6572
linsolve = nothing,
66-
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
73+
autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing,
74+
jacobian_reuse = nothing
6775
) where {T}
68-
common_kwargs = (; linsolve, autodiff, vjp_autodiff, jvp_autodiff)
76+
common_kwargs = (; linsolve, autodiff, vjp_autodiff, jvp_autodiff, jacobian_reuse)
6977
if T <: Complex
7078
algs = (
7179
GaussNewton(; common_kwargs..., concrete_jac),

lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
PseudoTransient(;
33
concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3,
44
linsolve = nothing, mass_matrix = nothing,
5-
autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing
5+
autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing,
6+
jacobian_reuse = nothing
67
)
78
89
An implementation of PseudoTransient Method [coffey2003pseudotransient](@cite) that is used
@@ -31,11 +32,15 @@ are treated consistently with the DAE structure.
3132
`NonlinearFunction` carries a non-identity `mass_matrix`, that mass matrix is used
3233
automatically. A diagonal `M` (e.g. `Diagonal(...)`) uses an efficient diagonal update;
3334
a general sparse/dense `M` is supported as well. Intended for square DAE-derived systems.
35+
- `jacobian_reuse`: a [`JacobianReuse`](@ref) policy, `true` for the default policy, or
36+
`nothing`/`false` to recompute after every accepted step. The damped system is still
37+
rebuilt when the pseudo-timestep changes. Defaults to `nothing`.
3438
"""
3539
function PseudoTransient(;
3640
concrete_jac = nothing, linesearch = missing, alpha_initial = 1.0e-3,
3741
linsolve = nothing, mass_matrix = nothing,
38-
autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing
42+
autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing,
43+
jacobian_reuse = nothing
3944
)
4045
return GeneralizedFirstOrderAlgorithm(;
4146
linesearch,
@@ -47,6 +52,7 @@ function PseudoTransient(;
4752
jvp_autodiff,
4853
vjp_autodiff,
4954
concrete_jac,
55+
jacobian_reuse,
5056
name = :PseudoTransient
5157
)
5258
end

0 commit comments

Comments
 (0)