From a8a39d5faaa53ff29ad03c2ad643a6e438c3f716 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 17 Oct 2023 14:12:00 -0400 Subject: [PATCH] Cleanup the PolyAlg Code --- docs/src/solvers/NonlinearSystemSolvers.md | 8 +- src/default.jl | 310 ++++++++------------- src/trustRegion.jl | 20 +- test/polyalgs.jl | 2 +- 4 files changed, 127 insertions(+), 213 deletions(-) diff --git a/docs/src/solvers/NonlinearSystemSolvers.md b/docs/src/solvers/NonlinearSystemSolvers.md index 08eb3f37d..effe39f7c 100644 --- a/docs/src/solvers/NonlinearSystemSolvers.md +++ b/docs/src/solvers/NonlinearSystemSolvers.md @@ -16,7 +16,7 @@ it will make sure to work. If one is looking for more robustness then `RobustMultiNewton` is a good choice. It attempts a set of the most robust methods in succession and only fails if -all of the methods fail to converge. Additionally, `DynamicSS` can be a good choice +all of the methods fail to converge. Additionally, `DynamicSS` can be a good choice for high stability. As a balance, `NewtonRaphson` is a good choice for most problems that aren't too @@ -25,13 +25,13 @@ but more stable. If the problem is well-conditioned, `Klement` or `Broyden` may be faster, but highly dependent on the eigenvalues of the Jacobian being sufficiently small. -`NewtonRaphson` and `TrustRegion` are designed for for large systems. +`NewtonRaphson` and `TrustRegion` are designed for for large systems. They can make use of sparsity patterns for sparse automatic differentiation and sparse linear solving of very large systems. Meanwhile, `SimpleNewtonRaphson` and `SimpleTrustRegion` are implementations which is specialized for small equations. They are non-allocating on static arrays and thus really well-optimized for small systems, thus usually outperforming the other methods when such types are -used for `u0`. +used for `u0`. ## Full List of Methods @@ -57,7 +57,7 @@ features, but have a bit of overhead on very small problems. large-scale and numerically-difficult nonlinear systems. - `RobustMultiNewton()`: A polyalgorithm that mixes highly robust methods (line searches and trust regions) in order to be as robust as possible for difficult problems. If this method - fails to converge, then one can be pretty certain that most (all?) other choices would + fails to converge, then one can be pretty certain that most (all?) other choices would likely fail. - `FastShortcutNonlinearPolyalg`: The default method. A polyalgorithm that mixes fast methods with fallbacks to robust methods to allow for solving easy problems quickly without sacrificing diff --git a/src/default.jl b/src/default.jl index 90ac0d9ec..c7092341b 100644 --- a/src/default.jl +++ b/src/default.jl @@ -1,6 +1,6 @@ """ -RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, adkwargs...) + RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, + adkwargs...) A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different globalizing techniques (trust region updates, line searches, etc.) in order to find a @@ -12,20 +12,20 @@ or more precision / more stable linear solver choice is required). ### Keyword Arguments -- `autodiff`: determines the backend used for the Jacobian. Note that this argument is + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. -- `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, then the Jacobian will not be constructed and instead direct Jacobian-vector products `J*v` are computed using forward-mode automatic differentiation or finite differencing tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, for example for a preconditioner, `concrete_jac = true` can be passed in order to force the construction of the Jacobian. -- `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the linear solves within the Newton method. Defaults to `nothing`, which means it uses the LinearSolve.jl default algorithm choice. For more information on available algorithm choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -- `precs`: the choice of preconditioners for the linear solver. Defaults to using no + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no preconditioners. For more information on specifying preconditioners for LinearSolve algorithms, consult the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). @@ -41,9 +41,8 @@ end # Robusters! const Robusters = RobustMultiNewton -function RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, adkwargs...) - +function RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, adkwargs...) return RobustMultiNewton{_unwrap_val(concrete_jac)}(adkwargs, linsolve, precs) end @@ -53,45 +52,46 @@ end current::Int end -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::RobustMultiNewton, args...; - kwargs...) where {uType, iip} - - adkwargs = alg.adkwargs - linsolve = alg.linsolve - precs = alg.precs - - RobustMultiNewtonCache{iip}(( - SciMLBase.__init(prob, TrustRegion(;linsolve, precs, adkwargs...), args...; kwargs...), - SciMLBase.__init(prob, TrustRegion(;linsolve, precs, radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...), args...; kwargs...), - SciMLBase.__init(prob, NewtonRaphson(;linsolve, precs, linesearch=BackTracking(), adkwargs...), args...; kwargs...), - SciMLBase.__init(prob, TrustRegion(;linsolve, precs, radius_update_scheme = RadiusUpdateSchemes.Fan, adkwargs...), args...; kwargs...), - ), alg, 1 - ) +function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::RobustMultiNewton, + args...; kwargs...) where {uType, iip} + @unpack adkwargs, linsolve, precs = alg + + algs = (TrustRegion(; linsolve, precs, adkwargs...), + TrustRegion(; linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...), + NewtonRaphson(; linsolve, precs, linesearch = BackTracking(), adkwargs...), + TrustRegion(; linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.NLsolve, adkwargs...), + TrustRegion(; linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Fan, adkwargs...)) + + return RobustMultiNewtonCache{iip}(map(solver -> SciMLBase.__init(prob, solver, args...; + kwargs...), algs), alg, 1) end """ -FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, adkwargs...) + FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, adkwargs...) A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods for more performance and then tries more robust techniques if the faster ones fail. ### Keyword Arguments -- `autodiff`: determines the backend used for the Jacobian. Note that this argument is + - `autodiff`: determines the backend used for the Jacobian. Note that this argument is ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. -- `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, + - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, then the Jacobian will not be constructed and instead direct Jacobian-vector products `J*v` are computed using forward-mode automatic differentiation or finite differencing tricks (without ever constructing the Jacobian). However, if the Jacobian is still needed, for example for a preconditioner, `concrete_jac = true` can be passed in order to force the construction of the Jacobian. -- `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the + - `linsolve`: the [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) used for the linear solves within the Newton method. Defaults to `nothing`, which means it uses the LinearSolve.jl default algorithm choice. For more information on available algorithm choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -- `precs`: the choice of preconditioners for the linear solver. Defaults to using no + - `precs`: the choice of preconditioners for the linear solver. Defaults to using no preconditioners. For more information on specifying preconditioners for LinearSolve algorithms, consult the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). @@ -102,212 +102,128 @@ for more performance and then tries more robust techniques if the faster ones fa precs end -function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing, +function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, adkwargs...) - - return FastShortcutNonlinearPolyalg{_unwrap_val(concrete_jac)}(adkwargs, linsolve, precs) + return FastShortcutNonlinearPolyalg{_unwrap_val(concrete_jac)}(adkwargs, linsolve, + precs) end -@concrete mutable struct FastShortcutNonlinearPolyalgCache{iip} <: AbstractNonlinearSolveCache{iip} +@concrete mutable struct FastShortcutNonlinearPolyalgCache{iip} <: + AbstractNonlinearSolveCache{iip} caches alg current::Int end -function FastShortcutNonlinearPolyalgCache(; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, adkwargs...) - - return FastShortcutNonlinearPolyalgCache{_unwrap_val(concrete_jac)}(adkwargs, linsolve, precs) +function FastShortcutNonlinearPolyalgCache(; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, adkwargs...) + return FastShortcutNonlinearPolyalgCache{_unwrap_val(concrete_jac)}(adkwargs, linsolve, + precs) end -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::FastShortcutNonlinearPolyalg, args...; - kwargs...) where {uType, iip} - - adkwargs = alg.adkwargs - linsolve = alg.linsolve - precs = alg.precs - - FastShortcutNonlinearPolyalgCache{iip}(( - #SciMLBase.__init(prob, Klement(), args...; kwargs...), - #SciMLBase.__init(prob, Broyden(), args...; kwargs...), - SciMLBase.__init(prob, NewtonRaphson(;linsolve, precs, adkwargs...), args...; kwargs...), - SciMLBase.__init(prob, NewtonRaphson(;linsolve, precs, linesearch=BackTracking(), adkwargs...), args...; kwargs...), - SciMLBase.__init(prob, TrustRegion(;linsolve, precs, adkwargs...), args...; kwargs...), - SciMLBase.__init(prob, TrustRegion(;linsolve, precs, radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...), args...; kwargs...), - ), alg, 1 - ) +function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, + alg::FastShortcutNonlinearPolyalg, args...; kwargs...) where {uType, iip} + @unpack adkwargs, linsolve, precs = alg + + algs = ( + # Klement(), + # Broyden(), + NewtonRaphson(; linsolve, precs, adkwargs...), + NewtonRaphson(; linsolve, precs, linesearch = BackTracking(), adkwargs...), + TrustRegion(; linsolve, precs, adkwargs...), + TrustRegion(; linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...)) + + return FastShortcutNonlinearPolyalgCache{iip}(map(solver -> SciMLBase.__init(prob, + solver, args...; kwargs...), algs), alg, 1) end -function SciMLBase.__solve(prob::NonlinearProblem{uType, iip}, alg::FastShortcutNonlinearPolyalg, args...; - kwargs...) where {uType, iip} - - adkwargs = alg.adkwargs - linsolve = alg.linsolve - precs = alg.precs - - sol1 = SciMLBase.__solve(prob, Klement(), args...; kwargs...) - if SciMLBase.successful_retcode(sol1) - return SciMLBase.build_solution(prob, alg, sol1.u, sol1.resid; - sol1.retcode, sol1.stats) - end - - sol2 = SciMLBase.__solve(prob, Broyden(), args...; kwargs...) - if SciMLBase.successful_retcode(sol2) - return SciMLBase.build_solution(prob, alg, sol2.u, sol2.resid; - sol2.retcode, sol2.stats) - end - - sol3 = SciMLBase.__solve(prob, NewtonRaphson(;linsolve, precs, adkwargs...), args...; kwargs...) - if SciMLBase.successful_retcode(sol3) - return SciMLBase.build_solution(prob, alg, sol3.u, sol3.resid; - sol3.retcode, sol3.stats) - end - - sol4 = SciMLBase.__solve(prob, TrustRegion(;linsolve, precs, adkwargs...), args...; kwargs...) - if SciMLBase.successful_retcode(sol4) - return SciMLBase.build_solution(prob, alg, sol4.u, sol4.resid; - sol4.retcode, sol4.stats) +# This version doesn't allocate all the caches! +function SciMLBase.__solve(prob::NonlinearProblem{uType, iip}, + alg::FastShortcutNonlinearPolyalg, args...; kwargs...) where {uType, iip} + @unpack adkwargs, linsolve, precs = alg + + algs = [ + iip ? Klement() : nothing, # Klement not yet implemented for IIP + iip ? Broyden() : nothing, # Broyden not yet implemented for IIP + NewtonRaphson(; linsolve, precs, adkwargs...), + NewtonRaphson(; linsolve, precs, linesearch = BackTracking(), adkwargs...), + TrustRegion(; linsolve, precs, adkwargs...), + TrustRegion(; linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...), + ] + filter!(!isnothing, algs) + + sols = Vector{SciMLBase.NonlinearSolution}(undef, length(algs)) + + for (i, solver) in enumerate(algs) + sols[i] = SciMLBase.__solve(prob, solver, args...; kwargs...) + if SciMLBase.successful_retcode(sols[i]) + return SciMLBase.build_solution(prob, alg, sols[i].u, sols[i].resid; + sols[i].retcode, sols[i].stats, original = sols[i]) + end end - resids = (sol1.resid, sol2.resid, sol3.resid, sol4.resid) + resids = map(Base.Fix2(getproperty, resid), sols) minfu, idx = findmin(DEFAULT_NORM, resids) - if idx == 1 - SciMLBase.build_solution(prob, alg, sol1.u, sol1.resid; - sol1.retcode, sol1.stats) - elseif idx == 2 - SciMLBase.build_solution(prob, alg, sol2.u, sol2.resid; - sol2.retcode, sol2.stats) - elseif idx == 3 - SciMLBase.build_solution(prob, alg, sol3.u, sol3.resid; - sol3.retcode, sol3.stats) - elseif idx == 4 - SciMLBase.build_solution(prob, alg, sol4.u, sol4.resid; - sol4.retcode, sol4.stats) - else - error("Unreachable reached, 박정석") - end - + return SciMLBase.build_solution(prob, alg, sols[idx].u, sols[idx].resid; + sols[idx].retcode, sols[idx].stats, original = sols[idx]) end ## General shared polyalg functions -function perform_step!(cache::Union{RobustMultiNewtonCache, FastShortcutNonlinearPolyalgCache}) +function perform_step!(cache::Union{RobustMultiNewtonCache, + FastShortcutNonlinearPolyalgCache}) current = cache.current + 1 ≤ current ≤ length(cache.caches) || error("Current choices shouldn't get here!") - while true - if current == 1 - perform_step!(cache.caches[1]) - elseif current == 2 - perform_step!(cache.caches[2]) - elseif current == 3 - perform_step!(cache.caches[3]) - elseif current == 4 - perform_step!(cache.caches[4]) - else - error("Current choices shouldn't get here!") - end + current_cache = cache.caches[current] + while not_terminated(current_cache) + perform_step!(current_cache) end return nothing end -function SciMLBase.solve!(cache::Union{RobustMultiNewtonCache, FastShortcutNonlinearPolyalgCache}) +function SciMLBase.solve!(cache::Union{RobustMultiNewtonCache, + FastShortcutNonlinearPolyalgCache}) current = cache.current - - while current < 5 && all(not_terminated, cache.caches) - if current == 1 - perform_step!(cache.caches[1]) - !not_terminated(cache.caches[1]) && (cache.current += 1) - elseif current == 2 - perform_step!(cache.caches[2]) - !not_terminated(cache.caches[2]) && (cache.current += 1) - elseif current == 3 - perform_step!(cache.caches[3]) - !not_terminated(cache.caches[3]) && (cache.current += 1) - elseif current == 4 - perform_step!(cache.caches[4]) - !not_terminated(cache.caches[4]) && (cache.current += 1) - else - error("Current choices shouldn't get here!") - end - - - - #cache.stats.nsteps += 1 + 1 ≤ current ≤ length(cache.caches) || error("Current choices shouldn't get here!") + + current_cache = cache.caches[current] + while current ≤ length(cache.caches) # && !all(terminated[current:end]) + sol_tmp = solve!(current_cache) + SciMLBase.successful_retcode(sol_tmp) && break + current += 1 + cache.current = current + current_cache = cache.caches[current] end - if current < 5 - stats = if current == 1 - cache.caches[1].stats - elseif current == 2 - cache.caches[2].stats - elseif current == 3 - cache.caches[3].stats - elseif current == 4 - cache.caches[4].stats - end - - u = if current == 1 - cache.caches[1].u - elseif current == 2 - cache.caches[2].u - elseif current == 3 - cache.caches[3].u - elseif current == 4 - cache.caches[4].u - end - - fu = if current == 1 - get_fu(cache.caches[1]) - elseif current == 2 - get_fu(cache.caches[2]) - elseif current == 3 - get_fu(cache.caches[3]) - elseif current == 4 - get_fu(cache.caches[4]) - end + if current ≤ length(cache.caches) + retcode = ReturnCode.Success - retcode = if stats.nsteps == cache.caches[1].maxiters - ReturnCode.MaxIters - else - ReturnCode.Success - end + stats = cache.caches[current].stats + u = cache.caches[current].u + fu = get_fu(cache.caches[current]) return SciMLBase.build_solution(cache.caches[1].prob, cache.alg, u, fu; retcode, stats) else retcode = ReturnCode.MaxIters - fus = (get_fu(cache.caches[1]), get_fu(cache.caches[2]), get_fu(cache.caches[3]), get_fu(cache.caches[4])) + fus = get_fu.(cache.caches) minfu, idx = findmin(cache.caches[1].internalnorm, fus) + stats = cache.caches[idx].stats + u = cache.caches[idx].u - stats = if idx == 1 - cache.caches[1].stats - elseif idx == 2 - cache.caches[2].stats - elseif idx == 3 - cache.caches[3].stats - elseif idx == 4 - cache.caches[4].stats - end - - u = if idx == 1 - cache.caches[1].u - elseif idx == 2 - cache.caches[2].u - elseif idx == 3 - cache.caches[3].u - elseif idx == 4 - cache.caches[4].u - end - - return SciMLBase.build_solution(cache.caches[1].prob, cache.alg, u, fu; - retcode, stats) + return SciMLBase.build_solution(cache.caches[idx].prob, cache.alg, u, fus[idx]; + retcode, stats) end end -function SciMLBase.reinit!(cache::Union{RobustMultiNewtonCache, FastShortcutNonlinearPolyalgCache}, args...; kwargs...) +function SciMLBase.reinit!(cache::Union{RobustMultiNewtonCache, + FastShortcutNonlinearPolyalgCache}, args...; kwargs...) for c in cache.caches SciMLBase.reinit!(c, args...; kwargs...) end @@ -315,14 +231,12 @@ end ## Defaults -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::Nothing, args...; +function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::Nothing, args...; kwargs...) where {uType, iip} - SciMLBase.__init(prob, FastShortcutNonlinearPolyalg(), args...; kwargs...) end -function SciMLBase.__solve(prob::NonlinearProblem{uType, iip}, alg::Nothing, args...; +function SciMLBase.__solve(prob::NonlinearProblem{uType, iip}, alg::Nothing, args...; kwargs...) where {uType, iip} - SciMLBase.__solve(prob, FastShortcutNonlinearPolyalg(), args...; kwargs...) -end \ No newline at end of file +end diff --git a/src/trustRegion.jl b/src/trustRegion.jl index 50d1078f3..583b1b675 100644 --- a/src/trustRegion.jl +++ b/src/trustRegion.jl @@ -28,14 +28,14 @@ states as `RadiusUpdateSchemes.T`. Simply put the desired scheme as follows: """ `RadiusUpdateSchemes.NLsolve` - The same updating scheme as in NLsolve's (https://github.com/JuliaNLSolvers/NLsolve.jl) trust region dogleg implementation. + The same updating scheme as in NLsolve's (https://github.com/JuliaNLSolvers/NLsolve.jl) trust region dogleg implementation. """ NLsolve """ `RadiusUpdateSchemes.NocedalWright` - Trust region updating scheme as in Nocedal and Wright [see Alg 11.5, page 291]. + Trust region updating scheme as in Nocedal and Wright [see Alg 11.5, page 291]. """ NocedalWright @@ -239,7 +239,7 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::TrustRegion, u_gauss_newton = zero(u) loss_new = loss - H = zero(J) + H = zero(J' * J) g = _mutable_zero(fu1) shrink_counter = 0 fu_new = zero(fu1) @@ -341,7 +341,7 @@ function perform_step!(cache::TrustRegionCache{true}) mul!(cache.g, J', fu) cache.stats.njacs += 1 - # do not use A = cache.H, b = _vec(cache.g) since it is equivalent + # do not use A = cache.H, b = _vec(cache.g) since it is equivalent # to A = cache.J, b = _vec(fu) as long as the Jacobian is non-singular linres = dolinsolve(alg.precs, linsolve, A = J, b = _vec(fu), linu = _vec(u_gauss_newton), @@ -450,12 +450,12 @@ function trust_region_step!(cache::TrustRegionCache) cache.make_new_J = false end - # trust region update - if r < 1 // 10 # cache.shrink_threshold - cache.trust_r *= 1 // 2 # cache.shrink_factor - elseif r >= 9 // 10 # cache.expand_threshold - cache.trust_r = 2 * norm(cache.du) # cache.expand_factor * norm(cache.du) - elseif r >= 1 // 2 # cache.p1 + # trust region update + if r < 1 // 10 # cache.shrink_threshold + cache.trust_r *= 1 // 2 # cache.shrink_factor + elseif r >= 9 // 10 # cache.expand_threshold + cache.trust_r = 2 * norm(cache.du) # cache.expand_factor * norm(cache.du) + elseif r >= 1 // 2 # cache.p1 cache.trust_r = max(cache.trust_r, 2 * norm(cache.du)) # cache.expand_factor * norm(cache.du)) end diff --git a/test/polyalgs.jl b/test/polyalgs.jl index 110f6228a..f514870b0 100644 --- a/test/polyalgs.jl +++ b/test/polyalgs.jl @@ -5,4 +5,4 @@ u0 = [1.0, 1.0] probN = NonlinearProblem(f, u0) @time solver = solve(probN, abstol = 1e-9) @time solver = solve(probN, RobustMultiNewton(), abstol = 1e-9) -@time solver = solve(probN, FastShortcutNonlinearPolyalg(), abstol = 1e-9) \ No newline at end of file +@time solver = solve(probN, FastShortcutNonlinearPolyalg(), abstol = 1e-9)