Skip to content

Jacobian reuse across reinit! and for matrix-free preconditioners - #1353

Open
ParadaCarleton wants to merge 3 commits into
SciML:masterfrom
ParadaCarleton:jacobian-reuse-extensions
Open

ParadaCarleton wants to merge 3 commits into
SciML:masterfrom
ParadaCarleton:jacobian-reuse-extensions

Conversation

@ParadaCarleton

Copy link
Copy Markdown

Summary

Three extensions of the adaptive Jacobian reuse from #1072, each opt-in, defaults unchanged.

 applicable_jacobian_reuse(alg, J)
-  J isa StatefulJacobianOperator → reuse off
+  J isa StatefulJacobianOperator and the linear solver builds a preconditioner
+    → reuse kept: operator rebound every step (exact Newton), preconditioner rebuilt on refresh

 reinit!(cache, u0; p)
-  resets Jacobian, age, factorization
+reinit!(cache, u0; p, reuse_jacobian = true)
+  keeps Jacobian, age, factorization/preconditioner; reseeds the residual norm
+  (first-order solvers and polyalgorithms; quasi-Newton keeps its secant Jacobian or history)
  1. Matrix-free preconditioner reuse (b25d1a9). An explicit JacobianReuse on a matrix-free Jacobian is kept when the linear solver builds a preconditioner; NonlinearSolveBaseLinearSolveExt updates the operator with reinit!(lincache; A, reuse_precs = true). The size-based default and matrix-free solves without precs stay off (the concern in JacobianReuse default should key off Jacobian cost, not length(u0) #1216).
  2. reinit!(cache, u0; reuse_jacobian = true) (2e443be) carries the Jacobian across a sequence of related solves (continuation correctors, parameter sweeps). The existing stale-Jacobian retry still refreshes when the carried one fails.
  3. The same keyword for quasi-Newton solvers (e74098b). QuasiNewtonCache.reinit_self! zeroed the step counter and discarded the Broyden-updated Jacobian, inverse or LimitedMemoryBroyden history; with the keyword it is kept and the method's own reset rules decide refreshes.

Evidence

N = 1500 nonlinear BVP, GMRES with a dense-LU preconditioner, one solve, same Newton steps:

policy time preconditioner builds
reuse off 0.304 s 7
size default (unchanged) 0.217 s 7
JacobianReuse() 0.117 s 5
JacobianReuse(max_residual_ratio = 1) 0.082 s 3

Same problem, 21 continuation steps, one reinit! per step, same 66 Newton steps:

configuration time preconditioner builds
reuse off 2.211 s 67
explicit policy, reinit! default 1.057 s 25
explicit policy, reuse_jacobian = true 0.601 s 11
max_age = 50, max_residual_ratio = 1, reuse_jacobian = true 0.346 s 4

Quasi-Newton, 21-step continuation, N = 800, Broyden(init_jacobian = Val(:true_jacobian)):

configuration time steps Jacobian evaluations
reuse_jacobian = false 1.831 s 69 22
reuse_jacobian = true 0.197 s 152 2

Tests: lib/NonlinearSolveFirstOrder/test/misc_tests__item9.jl (matrix-free keeps precs; no-precs and auto stay off; carried Jacobian across reinit!; retry when the line search rejects it; a 21-step sequence) and lib/NonlinearSolveQuasiNewton/test/core_tests__item14.jl (Broyden good/bad, true Jacobian, Klement, LimitedMemoryBroyden, in-place and out-of-place; a forced refresh re-evaluates; no carry before any step). NonlinearSolveBase Core and NonlinearSolveQuasiNewton Core pass. FirstOrder Core has 21 Enzyme ThunkCallErrors in the TrustRegion OOP-SVector rootfind tests, identical on the base commit 06b1229.

Not covered: Broyden(update_rule = Val(:diagonal)), which does not converge on the test problem with or without the keyword.

Merge Danger

Door: two-way

Blast Radius: opt-in

Every change is behind an explicit JacobianReuse policy or reuse_jacobian = true; the defaults take the old paths.

A matrix-free Jacobian operator is rebound to the current iterate on every
step, so JacobianReuse switched itself off for it. What is expensive there is
the preconditioner: an iterative linear solver rebuilt its `precs` whenever
`A` was rebound, even when the caller asked to reuse.

An explicit JacobianReuse policy now stays active for a matrix-free operator.
Steps remain exact Newton steps (the operator is bound to the current
iterate); the preconditioner built for the last refreshed Jacobian is kept
until the policy refreshes. The size-based default still leaves matrix-free
solves untouched.

`reuse_A_if_factorization = true` now keeps the preconditioner of an
iterative linear solver (via LinearSolve's `reinit!(...; reuse_precs = true)`)
in addition to a factorization.
…oss solves

reinit! reset the Jacobian reuse policy, so a sequence of related solves
through one cache (continuation correctors) refactorized and rebuilt the
preconditioner at the start of every solve; only the manual
step!(cache; recompute_jacobian = false) interface kept a Jacobian.

With reuse_jacobian = true, the first-order cache keeps its Jacobian, the
linear solver's factorization and preconditioner, and the policy's age. The
residual reference is set to the new initial residual, so the first step's
contraction decides a refresh as usual, and the stale-Jacobian retry still
applies. Default reinit! behavior is unchanged.
@ParadaCarleton

Copy link
Copy Markdown
Author

Could a maintainer approve the CI workflows (they are waiting on first-contributor approval) and review this? BifurcationKit.jl#345 builds its continuation corrector on reinit!(cache, u0; p, reuse_jacobian = true): the Jacobian, its age, and its factorization carry from one continuation step to the next, so the JacobianReuse policy decides refreshes across steps. Its corrector tests (65/65) pass on this branch, combined with #1354.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant