From c5c0176664dc48fd91df51bb86042ad2e3908e3a Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 5 Nov 2023 20:10:15 -0500 Subject: [PATCH] Use a proper polyalgorithm algorithm type --- Project.toml | 2 +- src/NonlinearSolve.jl | 2 +- src/default.jl | 370 ++++++++++++++++++++---------------------- test/polyalgs.jl | 7 + 4 files changed, 186 insertions(+), 195 deletions(-) diff --git a/Project.toml b/Project.toml index b2e25a90d..ee073a893 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "NonlinearSolve" uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" authors = ["SciML"] -version = "2.7.0" +version = "2.8.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 34d8d869b..39210cd3a 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -147,7 +147,7 @@ export RadiusUpdateSchemes export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient, GeneralBroyden, GeneralKlement, LimitedMemoryBroyden export LeastSquaresOptimJL, FastLevenbergMarquardtJL -export RobustMultiNewton, FastShortcutNonlinearPolyalg +export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPolyalg export LineSearch, LiFukushimaLineSearch diff --git a/src/default.jl b/src/default.jl index 178d72f1c..758e049f3 100644 --- a/src/default.jl +++ b/src/default.jl @@ -1,216 +1,66 @@ """ - RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - adkwargs...) + NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} -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 -method that is able to adequately solve the minimization problem. +A general way to define PolyAlgorithms for `NonlinearProblem` and +`NonlinearLeastSquaresProblem`. This is a container for a tuple of algorithms that will be +tried in order until one succeeds. If none succeed, then the algorithm with the lowest +residual is returned. -Basically, if this algorithm fails, then "most" good ways of solving your problem fail and -you may need to think about reformulating the model (either there is an issue with the model, -or more precision / more stable linear solver choice is required). +## Arguments -### Keyword Arguments + - `algs`: a tuple of algorithms to try in-order! (If this is not a Tuple, then the + returned algorithm is not type-stable). + - `pType`: the problem type. Defaults to `:NLS` for `NonlinearProblem` and `:NLLS` for + `NonlinearLeastSquaresProblem`. This is used to determine the correct problem type to + dispatch on. - - `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, - 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 - 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 - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -""" -@concrete struct RobustMultiNewton{CJ} <: AbstractNewtonAlgorithm{CJ, Nothing} - adkwargs - linsolve - precs -end - -# When somethin's strange, and numerical -# who you gonna call? -# Robusters! -const Robusters = RobustMultiNewton - -function RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, adkwargs...) - return RobustMultiNewton{_unwrap_val(concrete_jac)}(adkwargs, linsolve, precs) -end - -@concrete mutable struct RobustMultiNewtonCache{iip, N} <: AbstractNonlinearSolveCache{iip} - caches - alg - current::Int -end +## Example -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...)) - - # Partially Type Unstable but can't do much since some upstream caches -- LineSearches - # and SparseDiffTools cause the instability - return RobustMultiNewtonCache{iip, length(algs)}(map(solver -> SciMLBase.__init(prob, - solver, args...; kwargs...), algs), alg, 1) -end +```julia +using NonlinearSolve +alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), GeneralBroyden())) +``` """ - 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 +struct NonlinearSolvePolyAlgorithm{pType, N, A} <: AbstractNonlinearSolveAlgorithm + algs::A - - `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, - 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 - 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 - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). -""" -@concrete struct FastShortcutNonlinearPolyalg{CJ} <: AbstractNewtonAlgorithm{CJ, Nothing} - adkwargs - linsolve - precs + function NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType} + @assert pType ∈ (:NLS, :NLLS) + algs = Tuple(algs) + return new{pType, length(algs), typeof(algs)}(algs) + end end -function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, adkwargs...) - return FastShortcutNonlinearPolyalg{_unwrap_val(concrete_jac)}(adkwargs, linsolve, - precs) +function Base.show(io::IO, alg::NonlinearSolvePolyAlgorithm{pType, N}) where {pType, N} + problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem") + println(io, "NonlinearSolvePolyAlgorithm for $(problem_kind) with $(N) algorithms") + for i in 1:(N - 1) + println(io, " $(i): $(alg.algs[i])") + end + print(io, " $(N): $(alg.algs[N])") end -@concrete mutable struct FastShortcutNonlinearPolyalgCache{iip, N} <: +@concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N} <: 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) -end - -function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, - alg::FastShortcutNonlinearPolyalg, args...; kwargs...) where {uType, iip} - @unpack adkwargs, linsolve, precs = alg - - algs = (GeneralKlement(; linsolve, precs), - GeneralBroyden(), - 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, length(algs)}(map(solver -> SciMLBase.__init(prob, - solver, args...; kwargs...), algs), alg, 1) -end - -# This version doesn't allocate all the caches! -@generated function SciMLBase.__solve(prob::NonlinearProblem{uType, iip}, - alg::Union{FastShortcutNonlinearPolyalg, RobustMultiNewton}, args...; - kwargs...) where {uType, iip} - calls = [:(@unpack adkwargs, linsolve, precs = alg)] - - algs = if parameterless_type(alg) == RobustMultiNewton - [ - :(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...)), - ] - else - [ - :(GeneralKlement(; linsolve, precs)), - :(GeneralBroyden()), - :(NewtonRaphson(; linsolve, precs, adkwargs...)), - :(NewtonRaphson(; linsolve, precs, linesearch = BackTracking(), adkwargs...)), - :(TrustRegion(; linsolve, precs, adkwargs...)), - :(TrustRegion(; linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...)), - ] - end - filter!(!isnothing, algs) - sol_syms = [gensym("sol") for i in 1:length(algs)] - for i in 1:length(algs) - cur_sol = sol_syms[i] - push!(calls, - quote - $(cur_sol) = SciMLBase.__solve(prob, $(algs[i]), args...; kwargs...) - if SciMLBase.successful_retcode($(cur_sol)) - return SciMLBase.build_solution(prob, alg, $(cur_sol).u, - $(cur_sol).resid; $(cur_sol).retcode, $(cur_sol).stats, - original = $(cur_sol)) - end - end) - end - - resids = map(x -> Symbol("$(x)_resid"), sol_syms) - for (sym, resid) in zip(sol_syms, resids) - push!(calls, :($(resid) = $(sym).resid)) - end - - push!(calls, - quote - resids = tuple($(Tuple(resids)...)) - minfu, idx = __findmin(DEFAULT_NORM, resids) - end) - - for i in 1:length(algs) - push!(calls, - quote - if idx == $i - return SciMLBase.build_solution(prob, alg, $(sol_syms[i]).u, - $(sol_syms[i]).resid; $(sol_syms[i]).retcode, $(sol_syms[i]).stats) - end - end) +for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) + algType = NonlinearSolvePolyAlgorithm{pType} + @eval begin + function SciMLBase.__init(prob::$probType, alg::$algType{N}, args...; + kwargs...) where {N} + return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N}(map(solver -> SciMLBase.__init(prob, + solver, args...; kwargs...), alg.algs), alg, 1) + end end - push!(calls, :(error("Current choices shouldn't get here!"))) - - return Expr(:block, calls...) end -## General shared polyalg functions - -@generated function SciMLBase.solve!(cache::Union{RobustMultiNewtonCache{iip, N}, - FastShortcutNonlinearPolyalgCache{iip, N}}) where {iip, N} +@generated function SciMLBase.solve!(cache::NonlinearSolvePolyAlgorithmCache{iip, + N}) where {iip, N} calls = [ quote 1 ≤ cache.current ≤ length(cache.caches) || @@ -259,13 +109,147 @@ end return Expr(:block, calls...) end -function SciMLBase.reinit!(cache::Union{RobustMultiNewtonCache, - FastShortcutNonlinearPolyalgCache}, args...; kwargs...) +for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) + algType = NonlinearSolvePolyAlgorithm{pType} + @eval begin + @generated function SciMLBase.__solve(prob::$probType, alg::$algType{N}, args...; + kwargs...) where {N} + calls = [] + sol_syms = [gensym("sol") for _ in 1:N] + for i in 1:N + cur_sol = sol_syms[i] + push!(calls, + quote + $(cur_sol) = SciMLBase.__solve(prob, alg.algs[$(i)], args...; + kwargs...) + if SciMLBase.successful_retcode($(cur_sol)) + return SciMLBase.build_solution(prob, alg, $(cur_sol).u, + $(cur_sol).resid; $(cur_sol).retcode, $(cur_sol).stats, + original = $(cur_sol)) + end + end) + end + + resids = map(x -> Symbol("$(x)_resid"), sol_syms) + for (sym, resid) in zip(sol_syms, resids) + push!(calls, :($(resid) = $(sym).resid)) + end + + push!(calls, + quote + resids = tuple($(Tuple(resids)...)) + minfu, idx = __findmin(DEFAULT_NORM, resids) + end) + + for i in 1:N + push!(calls, + quote + if idx == $i + return SciMLBase.build_solution(prob, alg, $(sol_syms[i]).u, + $(sol_syms[i]).resid; $(sol_syms[i]).retcode, + $(sol_syms[i]).stats) + end + end) + end + push!(calls, :(error("Current choices shouldn't get here!"))) + + return Expr(:block, calls...) + end + end +end + +function SciMLBase.reinit!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) for c in cache.caches SciMLBase.reinit!(c, args...; kwargs...) end end +""" + 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 +method that is able to adequately solve the minimization problem. + +Basically, if this algorithm fails, then "most" good ways of solving your problem fail and +you may need to think about reformulating the model (either there is an issue with the model, +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 + 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, + 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 + 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 + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +""" +function RobustMultiNewton(; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, adkwargs...) + algs = (TrustRegion(; concrete_jac, linsolve, precs, adkwargs...), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...), + NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), + adkwargs...), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.NLsolve, adkwargs...), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Fan, adkwargs...)) + return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) +end + +""" + 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 + 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, + 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 + 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 + preconditioners. For more information on specifying preconditioners for LinearSolve + algorithms, consult the + [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). +""" +function FastShortcutNonlinearPolyalg(; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, adkwargs...) + algs = (GeneralKlement(; linsolve, precs), + GeneralBroyden(), + NewtonRaphson(; concrete_jac, linsolve, precs, adkwargs...), + NewtonRaphson(; concrete_jac, linsolve, precs, linesearch = BackTracking(), + adkwargs...), + TrustRegion(; concrete_jac, linsolve, precs, adkwargs...), + TrustRegion(; concrete_jac, linsolve, precs, + radius_update_scheme = RadiusUpdateSchemes.Bastin, adkwargs...)) + return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) +end + ## Defaults function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::Nothing, args...; diff --git a/test/polyalgs.jl b/test/polyalgs.jl index 4b5377159..0a4e599b3 100644 --- a/test/polyalgs.jl +++ b/test/polyalgs.jl @@ -4,6 +4,8 @@ f(u, p) = u .* u .- 2 u0 = [1.0, 1.0] probN = NonlinearProblem{false}(f, u0) +custom_polyalg = NonlinearSolvePolyAlgorithm((GeneralBroyden(), LimitedMemoryBroyden())) + # Uses the `__solve` function @time solver = solve(probN; abstol = 1e-9) @test SciMLBase.successful_retcode(solver) @@ -11,6 +13,8 @@ probN = NonlinearProblem{false}(f, u0) @test SciMLBase.successful_retcode(solver) @time solver = solve(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9) @test SciMLBase.successful_retcode(solver) +@time solver = solve(probN, custom_polyalg; abstol = 1e-9) +@test SciMLBase.successful_retcode(solver) # Test the caching interface cache = init(probN; abstol = 1e-9); @@ -22,6 +26,9 @@ cache = init(probN, RobustMultiNewton(); abstol = 1e-9); cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9); @time solver = solve!(cache) @test SciMLBase.successful_retcode(solver) +cache = init(probN, custom_polyalg; abstol = 1e-9); +@time solver = solve!(cache) +@test SciMLBase.successful_retcode(solver) # https://github.com/SciML/NonlinearSolve.jl/issues/153 function f(du, u, p)