SciML / NonlinearSolve.jl

High-performance and differentiation-enabled nonlinear solvers (Newton methods), bracketed rootfinding (bisection, Falsi), with sparsity and Newton-Krylov support.
https://docs.sciml.ai/NonlinearSolve/stable/
MIT License
216 stars 39 forks source link

Add forward mode to line search #446

Closed tansongchen closed 3 weeks ago

tansongchen commented 4 weeks ago

Checklist

Additional context

Added forward mode support to line searches. We prefer forward AD for better performance, however, reverse AD is also supported if user explicitly requests it.

  1. If jvp is available, we use forward AD;
  2. If reverse type is requested, we use reverse AD;
  3. Otherwise, we use forward AD.
tansongchen commented 4 weeks ago

Getting significant perf improvement for in-place functions like

quadratic_f!(du, u, p) = (du .= u .* u .- p)
u0 = fill(1.0, 1000)
p = 2.0
prob = NonlinearProblem(quadratic_f!, u0, p)
algfwd = NewtonRaphson(; linesearch=LineSearchesJL(; method=LineSearches.HagerZhang()))
algrev = NewtonRaphson(; linesearch=LineSearchesJL(; method=LineSearches.HagerZhang(), autodiff=AutoZygote())) # fallback to finitediff
julia> @btime sol = solve(prob, algfwd)
  27.795 ms (265 allocations: 15.62 MiB)
Julia> @btime sol = solve(prob, algrev)
  50.685 ms (15732 allocations: 133.12 MiB)
codecov[bot] commented 4 weeks ago

Codecov Report

Attention: Patch coverage is 77.77778% with 4 lines in your changes missing coverage. Please review.

Project coverage is 85.54%. Comparing base (a74c321) to head (82944d0).

Files Patch % Lines
src/globalization/line_search.jl 77.77% 4 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## master #446 +/- ## ========================================== - Coverage 86.45% 85.54% -0.91% ========================================== Files 47 47 Lines 2872 2878 +6 ========================================== - Hits 2483 2462 -21 - Misses 389 416 +27 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

tansongchen commented 3 weeks ago

Updated strategy:

  1. If FiniteDiff is requested, we use finite diff on $\phi$;
  2. If jvp is available, we use jvp;
  3. If vjp is available, we use vjp;
  4. If reverse type is requested, we use reverse AD;
  5. Finally, we use forward AD.