From ba8f8931ac99d290d6decb3e5d2f7d7d6eb5a959 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 28 Oct 2024 22:02:59 -0400 Subject: [PATCH] refactor: move stuff around a bit --- Project.toml | 16 +- lib/NonlinearSolveBase/src/abstract_types.jl | 7 +- .../src/descent/damped_newton.jl | 2 +- lib/NonlinearSolveBase/src/descent/dogleg.jl | 6 +- .../src/descent/geodesic_acceleration.jl | 2 +- lib/NonlinearSolveBase/src/descent/newton.jl | 2 +- .../src/descent/steepest.jl | 2 +- lib/NonlinearSolveBase/src/solve.jl | 8 +- lib/NonlinearSolveBase/src/utils.jl | 3 +- lib/NonlinearSolveFirstOrder/Project.toml | 1 + .../src/NonlinearSolveFirstOrder.jl | 8 +- .../src/trust_region.jl | 458 ++++++++++++++- .../src/SimpleNonlinearSolve.jl | 2 +- src/NonlinearSolve.jl | 153 ++--- src/algorithms/extension_algs.jl | 469 --------------- src/algorithms/trust_region.jl | 36 -- src/default.jl | 531 +---------------- src/extension_algs.jl | 469 +++++++++++++++ src/{internal => }/forward_diff.jl | 0 src/globalization/trust_region.jl | 452 --------------- src/{timer_outputs.jl => helpers.jl} | 0 src/polyalg.jl | 543 ++++++++++++++++++ src/utils.jl | 50 -- 23 files changed, 1566 insertions(+), 1654 deletions(-) delete mode 100644 src/algorithms/extension_algs.jl delete mode 100644 src/algorithms/trust_region.jl create mode 100644 src/extension_algs.jl rename src/{internal => }/forward_diff.jl (100%) delete mode 100644 src/globalization/trust_region.jl rename src/{timer_outputs.jl => helpers.jl} (100%) create mode 100644 src/polyalg.jl delete mode 100644 src/utils.jl diff --git a/Project.toml b/Project.toml index 1c6485b43..6d907607e 100644 --- a/Project.toml +++ b/Project.toml @@ -9,31 +9,25 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" -MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" NonlinearSolveFirstOrder = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d" NonlinearSolveQuasiNewton = "9a2c21bd-3a47-402d-9113-8faf9a0ee114" NonlinearSolveSpectralMethods = "26075421-4e9a-44e1-8bd1-420ed7ad02b2" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Preferences = "21216c6a-2e73-6563-6e65-726566657250" -RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" -SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" -TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [weakdeps] FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce" @@ -71,7 +65,6 @@ CUDA = "5.5" CommonSolve = "0.2.4" ConcreteStructs = "0.2.3" DiffEqBase = "6.155.3" -DifferentiationInterface = "0.6.16" Enzyme = "0.13.11" ExplicitImports = "1.5" FastClosures = "0.3.2" @@ -84,11 +77,10 @@ InteractiveUtils = "<0.0.1, 1" LeastSquaresOptim = "0.8.5" LineSearch = "0.1.4" LineSearches = "7.3" -LinearAlgebra = "1.10" -LinearSolve = "2.35" +LinearAlgebra = "1.11.0" +LinearSolve = "2.36.1" MINPACK = "1.2" MPI = "0.20.22" -MaybeInplace = "0.1.4" NLSolvers = "0.5" NLsolve = "4.5" NaNMath = "1" @@ -101,12 +93,9 @@ PrecompileTools = "1.2" Preferences = "1.4" Random = "1.10" ReTestItems = "1.24" -RecursiveArrayTools = "3.27" Reexport = "1.2" SIAMFANLEquations = "1.0.1" SciMLBase = "2.54.0" -SciMLJacobianOperators = "0.1" -SciMLOperators = "0.3.10" SimpleNonlinearSolve = "2" SparseArrays = "1.10" SparseConnectivityTracer = "0.6.5" @@ -118,7 +107,6 @@ StaticArraysCore = "1.4" Sundials = "4.23.1" SymbolicIndexingInterface = "0.3.31" Test = "1.10" -TimerOutputs = "0.5.23" Zygote = "0.6.69" julia = "1.10" diff --git a/lib/NonlinearSolveBase/src/abstract_types.jl b/lib/NonlinearSolveBase/src/abstract_types.jl index 3cdb59a1d..91afc2901 100644 --- a/lib/NonlinearSolveBase/src/abstract_types.jl +++ b/lib/NonlinearSolveBase/src/abstract_types.jl @@ -18,7 +18,7 @@ function reinit_self!(x::Any; kwargs...) return end -function reinit_self!(stats::NLStats) +function reinit!(stats::NLStats) stats.nf = 0 stats.nsteps = 0 stats.nfactors = 0 @@ -257,8 +257,8 @@ Abstract Type for all NonlinearSolveBase Caches. - `SciMLBase.set_u!(cache, u)`: set the current state. - `SciMLBase.reinit!(cache, u0; kwargs...)`: reinitialize the cache with the initial state `u0` and any additional keyword arguments. - - `SciMLBase.step!(cache; kwargs...)`: See [`SciMLBase.step!`](@ref) for more details. - `SciMLBase.isinplace(cache)`: whether or not the solver is inplace. + - `CommonSolve.step!(cache; kwargs...)`: See [`CommonSolve.step!`](@ref) for more details. Additionally implements `SymbolicIndexingInterface` interface Functions. @@ -577,7 +577,8 @@ macro internal_caches(cType, internal_cache_names...) return end function NonlinearSolveBase.InternalAPI.reinit!( - cache::$(cType), args...; kwargs...) + cache::$(cType), args...; kwargs... + ) $(reinit_caches...) $(InternalAPI.reinit_self!)(cache, args...; kwargs...) return diff --git a/lib/NonlinearSolveBase/src/descent/damped_newton.jl b/lib/NonlinearSolveBase/src/descent/damped_newton.jl index a9921ba47..47ed56207 100644 --- a/lib/NonlinearSolveBase/src/descent/damped_newton.jl +++ b/lib/NonlinearSolveBase/src/descent/damped_newton.jl @@ -42,7 +42,7 @@ supports_trust_region(::DampedNewtonDescent) = true mode <: Union{Val{:normal_form}, Val{:least_squares}, Val{:simple}} end -NonlinearSolveBase.@internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache +@internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache function InternalAPI.init( prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, u; stats, diff --git a/lib/NonlinearSolveBase/src/descent/dogleg.jl b/lib/NonlinearSolveBase/src/descent/dogleg.jl index 133b5fabb..f446e9cf6 100644 --- a/lib/NonlinearSolveBase/src/descent/dogleg.jl +++ b/lib/NonlinearSolveBase/src/descent/dogleg.jl @@ -44,7 +44,7 @@ end normal_form <: Union{Val{false}, Val{true}} end -NonlinearSolveBase.@internal_caches DoglegCache :newton_cache :cauchy_cache +@internal_caches DoglegCache :newton_cache :cauchy_cache function InternalAPI.init( prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; @@ -111,7 +111,7 @@ function InternalAPI.solve!( l_grad = cache.internalnorm(δu_cauchy) @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) - δuJᵀJδu = Utils.dot(cache.δu_cache_mul, cache.δu_cache_mul) + δuJᵀJδu = Utils.safe_dot(cache.δu_cache_mul, cache.δu_cache_mul) else δu_cauchy = InternalAPI.solve!( cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs... @@ -119,7 +119,7 @@ function InternalAPI.solve!( J_ = preinverted_jacobian(cache) ? inv(J) : J l_grad = cache.internalnorm(δu_cauchy) @bb cache.Jᵀδu_cache = J_ × vec(δu_cauchy) - δuJᵀJδu = Utils.dot(cache.Jᵀδu_cache, cache.Jᵀδu_cache) + δuJᵀJδu = Utils.safe_dot(cache.Jᵀδu_cache, cache.Jᵀδu_cache) end d_cauchy = (l_grad^3) / δuJᵀJδu diff --git a/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl b/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl index 486308f09..409ed9ce9 100644 --- a/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl +++ b/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl @@ -54,7 +54,7 @@ function InternalAPI.reinit_self!(cache::GeodesicAccelerationCache; p = cache.p, cache.last_step_accepted = false end -NonlinearSolveBase.@internal_caches GeodesicAccelerationCache :descent_cache +@internal_caches GeodesicAccelerationCache :descent_cache function get_velocity(cache::GeodesicAccelerationCache) return SciMLBase.get_du(cache.descent_cache, Val(1)) diff --git a/lib/NonlinearSolveBase/src/descent/newton.jl b/lib/NonlinearSolveBase/src/descent/newton.jl index e453597a1..810661507 100644 --- a/lib/NonlinearSolveBase/src/descent/newton.jl +++ b/lib/NonlinearSolveBase/src/descent/newton.jl @@ -24,7 +24,7 @@ supports_line_search(::NewtonDescent) = true normal_form <: Union{Val{false}, Val{true}} end -NonlinearSolveBase.@internal_caches NewtonDescentCache :lincache +@internal_caches NewtonDescentCache :lincache function InternalAPI.init( prob::AbstractNonlinearProblem, alg::NewtonDescent, J, fu, u; stats, diff --git a/lib/NonlinearSolveBase/src/descent/steepest.jl b/lib/NonlinearSolveBase/src/descent/steepest.jl index b93c727c5..deb1a9817 100644 --- a/lib/NonlinearSolveBase/src/descent/steepest.jl +++ b/lib/NonlinearSolveBase/src/descent/steepest.jl @@ -21,7 +21,7 @@ supports_line_search(::SteepestDescent) = true preinverted_jacobian <: Union{Val{false}, Val{true}} end -NonlinearSolveBase.@internal_caches SteepestDescentCache :lincache +@internal_caches SteepestDescentCache :lincache function InternalAPI.init( prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; diff --git a/lib/NonlinearSolveBase/src/solve.jl b/lib/NonlinearSolveBase/src/solve.jl index 3c50521a0..08b60e4db 100644 --- a/lib/NonlinearSolveBase/src/solve.jl +++ b/lib/NonlinearSolveBase/src/solve.jl @@ -2,13 +2,13 @@ function SciMLBase.__solve( prob::AbstractNonlinearProblem, alg::AbstractNonlinearSolveAlgorithm, args...; kwargs... ) - cache = SciMLBase.init(prob, alg, args...; kwargs...) - return SciMLBase.solve!(cache) + cache = SciMLBase.__init(prob, alg, args...; kwargs...) + return CommonSolve.solve!(cache) end function CommonSolve.solve!(cache::AbstractNonlinearSolveCache) while not_terminated(cache) - SciMLBase.step!(cache) + CommonSolve.step!(cache) end # The solver might have set a different `retcode` @@ -47,7 +47,7 @@ Performs one step of the nonlinear solver. respectively. For algorithms that don't use jacobian information, this keyword is ignored with a one-time warning. """ -function SciMLBase.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) +function CommonSolve.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) not_terminated(cache) || return has_time_limit(cache) && (time_start = time()) diff --git a/lib/NonlinearSolveBase/src/utils.jl b/lib/NonlinearSolveBase/src/utils.jl index 6d7a66aa0..6e739c0f8 100644 --- a/lib/NonlinearSolveBase/src/utils.jl +++ b/lib/NonlinearSolveBase/src/utils.jl @@ -2,7 +2,7 @@ module Utils using ArrayInterface: ArrayInterface using FastClosures: @closure -using LinearAlgebra: LinearAlgebra, Diagonal, Symmetric, norm, dot, cond, diagind, pinv, inv +using LinearAlgebra: LinearAlgebra, Diagonal, Symmetric, norm, dot, cond, diagind, pinv using MaybeInplace: @bb using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition using SciMLOperators: AbstractSciMLOperator @@ -117,6 +117,7 @@ is_default_value(::Any, ::Symbol, ::Nothing) = true is_default_value(::Any, ::Symbol, ::Missing) = true is_default_value(::Any, ::Symbol, val::Int) = val == typemax(typeof(val)) is_default_value(::Any, ::Symbol, ::Any) = false +is_default_value(::Any, ::Any, ::Any) = false maybe_symmetric(x) = Symmetric(x) maybe_symmetric(x::Number) = x diff --git a/lib/NonlinearSolveFirstOrder/Project.toml b/lib/NonlinearSolveFirstOrder/Project.toml index ef949cb73..5be145a6e 100644 --- a/lib/NonlinearSolveFirstOrder/Project.toml +++ b/lib/NonlinearSolveFirstOrder/Project.toml @@ -19,6 +19,7 @@ NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" diff --git a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl index 9fdfc379a..64b23aab6 100644 --- a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl +++ b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl @@ -22,6 +22,7 @@ using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm, NewtonDescent, DampedNewtonDescent, GeodesicAcceleration, Dogleg using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode, NonlinearFunction +using SciMLJacobianOperators: VecJacOperator, JacVecOperator, StatefulJacobianOperator using Setfield: @set! using StaticArraysCore: SArray @@ -41,9 +42,8 @@ include("solve.jl") @__DIR__, "..", "..", "..", "common", "nlls_problem_workloads.jl" )) - # XXX: TrustRegion - nlp_algs = [NewtonRaphson(), LevenbergMarquardt()] - nlls_algs = [GaussNewton(), LevenbergMarquardt()] + nlp_algs = [NewtonRaphson(), TrustRegion(), LevenbergMarquardt()] + nlls_algs = [GaussNewton(), TrustRegion(), LevenbergMarquardt()] @compile_workload begin for prob in nonlinear_problems, alg in nlp_algs @@ -61,6 +61,8 @@ end export NewtonRaphson, PseudoTransient export GaussNewton, LevenbergMarquardt, TrustRegion +export RadiusUpdateSchemes + export GeneralizedFirstOrderAlgorithm end diff --git a/lib/NonlinearSolveFirstOrder/src/trust_region.jl b/lib/NonlinearSolveFirstOrder/src/trust_region.jl index 5883d23b7..249d4fa72 100644 --- a/lib/NonlinearSolveFirstOrder/src/trust_region.jl +++ b/lib/NonlinearSolveFirstOrder/src/trust_region.jl @@ -34,7 +34,459 @@ function TrustRegion(; descent = Dogleg(; linsolve, precs) trustregion = GenericTrustRegionScheme(; method = radius_update_scheme, step_threshold, shrink_threshold, expand_threshold, - shrink_factor, expand_factor, initial_trust_radius, max_trust_radius) - return GeneralizedFirstOrderAlgorithm{concrete_jac, :TrustRegion}(; - trustregion, descent, autodiff, vjp_autodiff, jvp_autodiff, max_shrink_times) + shrink_factor, expand_factor, initial_trust_radius, max_trust_radius + ) + return GeneralizedFirstOrderAlgorithm(; + trustregion, descent, autodiff, vjp_autodiff, jvp_autodiff, max_shrink_times, + concrete_jac, name = :TrustRegion + ) +end + +# Don't Pollute the namespace +""" + RadiusUpdateSchemes + +`RadiusUpdateSchemes` is provides different types of radius update schemes implemented in +the Trust Region method. These schemes specify how the radius of the so-called trust region +is updated after each iteration of the algorithm. The specific role and caveats associated +with each scheme are provided below. + +## Using `RadiusUpdateSchemes` + +Simply put the desired scheme as follows: +`sol = solve(prob, alg = TrustRegion(radius_update_scheme = RadiusUpdateSchemes.Hei))`. +""" +module RadiusUpdateSchemes +# The weird definitions here are needed to main compatibility with the older enum variants + +abstract type AbstractRadiusUpdateScheme end + +function Base.show(io::IO, rus::AbstractRadiusUpdateScheme) + print(io, "RadiusUpdateSchemes.$(string(nameof(typeof(rus)))[3:end])") +end + +const T = AbstractRadiusUpdateScheme + +struct __Simple <: AbstractRadiusUpdateScheme end +""" + RadiusUpdateSchemes.Simple + +The simple or conventional radius update scheme. This scheme is chosen by default and +follows the conventional approach to update the trust region radius, i.e. if the trial +step is accepted it increases the radius by a fixed factor (bounded by a maximum radius) +and if the trial step is rejected, it shrinks the radius by a fixed factor. +""" +const Simple = __Simple() + +struct __NLsolve <: AbstractRadiusUpdateScheme end +""" + RadiusUpdateSchemes.NLsolve + +The same updating scheme as in NLsolve's (https://github.com/JuliaNLSolvers/NLsolve.jl) +trust region dogleg implementation. +""" +const NLsolve = __NLsolve() + +struct __NocedalWright <: AbstractRadiusUpdateScheme end +""" + RadiusUpdateSchemes.NocedalWright + +Trust region updating scheme as in Nocedal and Wright [see Alg 11.5, page 291]. +""" +const NocedalWright = __NocedalWright() + +struct __Hei <: AbstractRadiusUpdateScheme end +""" + RadiusUpdateSchemes.Hei + +This scheme is proposed in [hei2003self](@citet). The trust region radius depends on the +size (norm) of the current step size. The hypothesis is to let the radius converge to zero +as the iterations progress, which is more reliable and robust for ill-conditioned as well +as degenerate problems. +""" +const Hei = __Hei() + +struct __Yuan <: AbstractRadiusUpdateScheme end +""" + RadiusUpdateSchemes.Yuan + +This scheme is proposed by [yuan2015recent](@citet). Similar to Hei's scheme, the +trust region is updated in a way so that it converges to zero, however here, the radius +depends on the size (norm) of the current gradient of the objective (merit) function. The +hypothesis is that the step size is bounded by the gradient size, so it makes sense to let +the radius depend on the gradient. +""" +const Yuan = __Yuan() + +struct __Bastin <: AbstractRadiusUpdateScheme end +""" + RadiusUpdateSchemes.Bastin + +This scheme is proposed by [bastin2010retrospective](@citet). The scheme is called a +retrospective update scheme as it uses the model function at the current iteration to +compute the ratio of the actual reduction and the predicted reduction in the previous trial +step, and use this ratio to update the trust region radius. The hypothesis is to exploit the +information made available during the optimization process in order to vary the accuracy +of the objective function computation. +""" +const Bastin = __Bastin() + +struct __Fan <: AbstractRadiusUpdateScheme end +""" + RadiusUpdateSchemes.Fan + +This scheme is proposed by [fan2006convergence](@citet). It is very much similar to Hei's +and Yuan's schemes as it lets the trust region radius depend on the current size (norm) of +the objective (merit) function itself. These new update schemes are known to improve local +convergence. +""" +const Fan = __Fan() + +end + +const RUS = RadiusUpdateSchemes + +""" + GenericTrustRegionScheme(; + method = RadiusUpdateSchemes.Simple, + max_trust_radius = nothing, initial_trust_radius = nothing, + step_threshold = nothing, shrink_threshold = nothing, expand_threshold = nothing, + shrink_factor = nothing, expand_factor = nothing + ) + +Trust Region Method that updates and stores the current trust region radius in +`trust_region`. For any of the keyword arguments, if the value is `nothing`, then we use +the value used in the respective paper. + +### Keyword Arguments + + - `radius_update_scheme`: the choice of radius update scheme to be used. Defaults to + `RadiusUpdateSchemes.Simple` which follows the conventional approach. Other available + schemes are documented in [`RadiusUpdateSchemes`](@ref),. These schemes have the trust + region radius converging to zero that is seen to improve convergence. For more details, + see [1]. + - `max_trust_radius`: the maximal trust region radius. Defaults to + `max(norm(fu), maximum(u) - minimum(u))`, except for `RadiusUpdateSchemes.NLsolve` + where it defaults to `Inf`. + - `initial_trust_radius`: the initial trust region radius. Defaults to + `max_trust_radius / 11`, except for `RadiusUpdateSchemes.NLsolve` where it defaults + to `u0_norm > 0 ? u0_norm : 1`. + - `step_threshold`: the threshold for taking a step. In every iteration, the threshold is + compared with a value `r`, which is the actual reduction in the objective function + divided by the predicted reduction. If `step_threshold > r` the model is not a good + approximation, and the step is rejected. Defaults to `nothing`. + - `shrink_threshold`: the threshold for shrinking the trust region radius. In every + iteration, the threshold is compared with a value `r` which is the actual reduction in + the objective function divided by the predicted reduction. If `shrink_threshold > r` the + trust region radius is shrunk by `shrink_factor`. Defaults to `nothing`. + - `expand_threshold`: the threshold for expanding the trust region radius. If a step is + taken, i.e `step_threshold < r` (with `r` defined in `shrink_threshold`), a check is + also made to see if `expand_threshold < r`. If that is true, the trust region radius is + expanded by `expand_factor`. Defaults to `nothing`. + - `shrink_factor`: the factor to shrink the trust region radius with if + `shrink_threshold > r` (with `r` defined in `shrink_threshold`). Defaults to `0.25`. + - `expand_factor`: the factor to expand the trust region radius with if + `expand_threshold < r` (with `r` defined in `shrink_threshold`). Defaults to `2.0`. +""" +@kwdef @concrete struct GenericTrustRegionScheme <: AbstractTrustRegionMethod + method <: RUS.AbstractRadiusUpdateScheme = RUS.Simple + step_threshold = nothing + shrink_threshold = nothing + shrink_factor = nothing + expand_factor = nothing + expand_threshold = nothing + max_trust_radius = nothing + initial_trust_radius = nothing +end + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, f, fu, u, p, + args...; stats, internalnorm::F = L2_NORM, vjp_autodiff = nothing, + jvp_autodiff = nothing, kwargs... +) where {F} + T = promote_type(eltype(u), eltype(fu)) + u0_norm = internalnorm(u) + fu_norm = internalnorm(fu) + + # Common Setup + mtr = max_trust_radius(alg.max_trust_radius, T, alg.method, u, fu_norm) + itr = initial_trust_radius( + alg.initial_trust_radius, T, alg.method, mtr, u0_norm, fu_norm + ) + stt = step_threshold(alg.step_threshold, T, alg.method) + sht = shrink_threshold(alg.shrink_threshold, T, alg.method) + shf = shrink_factor(alg.shrink_factor, T, alg.method) + et = expand_threshold(alg.expand_threshold, T, alg.method) + ef = expand_factor(alg.expand_factor, T, alg.method) + + # Scheme Specific Setup + p1, p2, p3, p4 = get_parameters(T, alg.method) + ϵ = T(1e-8) + + vjp_operator = alg.method isa RUS.__Yuan || alg.method isa RUS.__Bastin ? + VecJacOperator(prob, fu, u; autodiff = vjp_autodiff) : nothing + + jvp_operator = alg.method isa RUS.__Bastin ? + JacVecOperator(prob, fu, u; autodiff = jvp_autodiff) : nothing + + if alg.method isa RUS.__Yuan + Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * Utils.safe_vec(fu) + itr = T(p1 * internalnorm(Jᵀfu_cache)) + elseif u isa Number + Jᵀfu_cache = u + else + @bb Jᵀfu_cache = similar(u) + end + + if alg.method isa RUS.__Bastin + @bb δu_cache = similar(u) + else + δu_cache = nothing + end + + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + @bb Jδu_cache = similar(fu) + + return GenericTrustRegionSchemeCache( + alg.method, f, p, mtr, itr, itr, stt, sht, et, shf, ef, + p1, p2, p3, p4, ϵ, T(0), vjp_operator, jvp_operator, Jᵀfu_cache, Jδu_cache, + δu_cache, internalnorm, u_cache, fu_cache, false, 0, stats, alg + ) +end + +@concrete mutable struct GenericTrustRegionSchemeCache <: AbstractTrustRegionMethodCache + method + f + p + max_trust_radius + initial_trust_radius + trust_region + step_threshold + shrink_threshold + expand_threshold + shrink_factor + expand_factor + p1 + p2 + p3 + p4 + ϵ + ρ + vjp_operator + jvp_operator + Jᵀfu_cache + Jδu_cache + δu_cache + internalnorm + u_cache + fu_cache + last_step_accepted::Bool + shrink_counter::Int + stats::NLStats + alg +end + +function InternalAPI.reinit!( + cache::GenericTrustRegionSchemeCache; p = cache.p, u0 = nothing, kwargs... +) + cache.p = p + if u0 !== nothing + u0_norm = cache.internalnorm(u0) + end + cache.last_step_accepted = false + cache.shrink_counter = 0 +end + +# Defaults +for func in ( + :max_trust_radius, :initial_trust_radius, :step_threshold, :shrink_threshold, + :shrink_factor, :expand_threshold, :expand_factor +) + @eval function $(func)(val, ::Type{T}, args...) where {T} + iszero(val) && return $(func)(nothing, T, args...) + return T(val) + end +end + +max_trust_radius(::Nothing, ::Type{T}, method, u, fu_norm) where {T} = T(Inf) +function max_trust_radius(::Nothing, ::Type{T}, ::Union{RUS.__Simple, RUS.__NocedalWright}, + u, fu_norm) where {T} + u_min, u_max = extrema(u) + return max(T(fu_norm), u_max - u_min) +end + +function initial_trust_radius( + ::Nothing, ::Type{T}, method, max_tr, u0_norm, fu_norm +) where {T} + method isa RUS.__NLsolve && return T(ifelse(u0_norm > 0, u0_norm, 1)) + (method isa RUS.__Hei || method isa RUS.__Bastin) && return T(1) + method isa RUS.__Fan && return T((fu_norm^0.99) / 10) + return T(max_tr / 11) +end + +function step_threshold(::Nothing, ::Type{T}, method) where {T} + method isa RUS.__Hei && return T(0) + method isa RUS.__Yuan && return T(1 // 1000) + method isa RUS.__Bastin && return T(1 // 20) + return T(1 // 10000) +end + +function shrink_threshold(::Nothing, ::Type{T}, method) where {T} + method isa RUS.__Hei && return T(0) + (method isa RUS.__NLsolve || method isa RUS.__Bastin) && return T(1 // 20) + return T(1 // 4) +end + +function expand_threshold(::Nothing, ::Type{T}, method) where {T} + method isa RUS.__NLsolve && return T(9 // 10) + method isa RUS.__Hei && return T(0) + method isa RUS.__Bastin && return T(9 // 10) + return T(3 // 4) +end + +function shrink_factor(::Nothing, ::Type{T}, method) where {T} + method isa RUS.__NLsolve && return T(1 // 2) + method isa RUS.__Hei && return T(0) + method isa RUS.__Bastin && return T(1 // 20) + return T(1 // 4) +end + +function get_parameters(::Type{T}, method) where {T} + method isa RUS.__NLsolve && return (T(1 // 2), T(0), T(0), T(0)) + method isa RUS.__Hei && return (T(5), T(1 // 10), T(15 // 100), T(15 // 100)) + method isa RUS.__Yuan && return (T(2), T(1 // 6), T(6), T(0)) + method isa RUS.__Fan && return (T(1 // 10), T(1 // 4), T(12), T(1e18)) + method isa RUS.__Bastin && return (T(5 // 2), T(1 // 4), T(0), T(0)) + return (T(0), T(0), T(0), T(0)) +end + +expand_factor(::Nothing, ::Type{T}, method) where {T} = T(2) + +function rfunc_adaptive_trust_region(r::R, c2::R, M::R, γ1::R, β::R) where {R <: Real} + return ifelse( + r ≥ c2, + (2 * (M - 1 - γ2) * atan(r - c2) + (1 + γ2)) / R(π), + (1 - γ1 - β) * (exp(r - c2) + β / (1 - γ1 - β)) + ) +end + +function InternalAPI.solve!( + cache::GenericTrustRegionSchemeCache, J, fu, u, δu, descent_stats +) + T = promote_type(eltype(u), eltype(fu)) + @bb @. cache.u_cache = u + δu + cache.fu_cache = Utils.evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + cache.stats.nf += 1 + + if hasfield(typeof(descent_stats), :δuJᵀJδu) && !isnan(descent_stats.δuJᵀJδu) + δuJᵀJδu = descent_stats.δuJᵀJδu + else + @bb cache.Jδu_cache = J × vec(δu) + δuJᵀJδu = Utils.safe_dot(cache.Jδu_cache, cache.Jδu_cache) + end + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + num = (cache.internalnorm(cache.fu_cache)^2 - cache.internalnorm(fu)^2) / 2 + denom = Utils.safe_dot(δu, cache.Jᵀfu_cache) + δuJᵀJδu / 2 + cache.ρ = num / denom + + if cache.ρ > cache.step_threshold + cache.last_step_accepted = true + else + cache.last_step_accepted = false + end + + if cache.method isa RUS.__Simple + if cache.ρ < cache.shrink_threshold + cache.trust_region *= cache.shrink_factor + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + if cache.ρ > cache.expand_threshold && cache.ρ > cache.step_threshold + cache.trust_region = cache.expand_factor * cache.trust_region + end + end + elseif cache.method isa RUS.__NLsolve + if cache.ρ < cache.shrink_threshold + cache.trust_region *= cache.shrink_factor + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + if cache.ρ ≥ cache.expand_threshold + cache.trust_region = cache.expand_factor * cache.internalnorm(δu) + elseif cache.ρ ≥ cache.p1 + cache.trust_region = max( + cache.trust_region, cache.expand_factor * cache.internalnorm(δu) + ) + end + end + elseif cache.method isa RUS.__NocedalWright + if cache.ρ < cache.shrink_threshold + cache.trust_region = cache.shrink_factor * cache.internalnorm(δu) + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + if cache.ρ > cache.expand_threshold && + abs(cache.internalnorm(δu) - cache.trust_region) < 1e-6 * cache.trust_region + cache.trust_region = cache.expand_factor * cache.trust_region + end + end + elseif cache.method isa RUS.__Hei + tr_new = rfunc_adaptive_trust_region( + cache.ρ, cache.shrink_threshold, cache.p1, cache.p3, cache.p4 + ) + if tr_new < cache.trust_region + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + end + cache.trust_region = tr_new + elseif cache.method isa RUS.__Yuan + if cache.ρ < cache.shrink_threshold + cache.p1 = cache.p2 * cache.p1 + cache.shrink_counter += 1 + else + if cache.ρ ≥ cache.expand_threshold && + 2 * cache.internalnorm(δu) > cache.trust_region + cache.p1 = cache.p3 * cache.p1 + end + cache.shrink_counter = 0 + end + operator = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) + @bb cache.Jᵀfu_cache = operator × vec(cache.fu_cache) + cache.trust_region = cache.p1 * cache.internalnorm(cache.Jᵀfu_cache) + elseif cache.method isa RUS.__Fan + if cache.ρ < cache.shrink_threshold + cache.p1 *= cache.p2 + cache.shrink_counter += 1 + else + cache.shrink_counter = 0 + cache.ρ > cache.expand_threshold && + (cache.p1 = min(cache.p1 * cache.p3, cache.p4)) + end + cache.trust_region = cache.p1 * (cache.internalnorm(cache.fu_cache)^T(0.99)) + elseif cache.method isa RUS.__Bastin + if cache.ρ > cache.step_threshold + jvp_op = StatefulJacobianOperator(cache.jvp_operator, cache.u_cache, cache.p) + vjp_op = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) + @bb cache.Jδu_cache = jvp_op × vec(cache.δu_cache) + @bb cache.Jᵀfu_cache = vjp_op × vec(cache.fu_cache) + denom_1 = dot(_vec(cache.Jᵀfu_cache), cache.Jᵀfu_cache) + @bb cache.Jᵀfu_cache = vjp_op × vec(cache.Jδu_cache) + denom_2 = dot(_vec(cache.Jᵀfu_cache), cache.Jᵀfu_cache) + denom = denom_1 + denom_2 / 2 + ρ = num / denom + if ρ ≥ cache.expand_threshold + cache.trust_region = cache.p1 * cache.internalnorm(cache.δu_cache) + end + cache.shrink_counter = 0 + else + cache.trust_region *= cache.p2 + cache.shrink_counter += 1 + end + end + + cache.trust_region = min(cache.trust_region, cache.max_trust_radius) + + return cache.last_step_accepted, cache.u_cache, cache.fu_cache end diff --git a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl index 858970a0c..f51064000 100644 --- a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl +++ b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl @@ -130,7 +130,7 @@ function solve_adjoint_internal end #!format: on @compile_workload begin - for prob in (prob_scalar, prob_iip, prob_oop) + for prob in (prob_scalar, prob_iip, prob_oop), alg in algs CommonSolve.solve(prob, alg; abstol = 1e-2) end end diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index eaae0ee6a..e352f4891 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -3,137 +3,80 @@ module NonlinearSolve using Reexport: @reexport using PrecompileTools: @compile_workload, @setup_workload -using ArrayInterface: ArrayInterface, can_setindex, restructure, fast_scalar_indexing, - ismutable -using CommonSolve: solve, init, solve! +using ArrayInterface: ArrayInterface +using CommonSolve: CommonSolve, solve, solve! using ConcreteStructs: @concrete using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches using FastClosures: @closure -using LinearAlgebra: LinearAlgebra, Diagonal, I, LowerTriangular, Symmetric, - UpperTriangular, axpy!, cond, diag, diagind, dot, issuccess, istril, - istriu, lu, mul!, norm, pinv, tril!, triu! -using LineSearch: LineSearch, AbstractLineSearchCache, LineSearchesJL, NoLineSearch, - RobustNonMonotoneLineSearch, BackTracking, LiFukushimaLineSearch -using LinearSolve: LinearSolve -using MaybeInplace: @bb -using NonlinearSolveBase: NonlinearSolveBase, - nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution, - nonlinearsolve_∂f_∂p, nonlinearsolve_∂f_∂u, - L2_NORM, - AbsNormTerminationMode, AbstractNonlinearTerminationMode, - AbstractSafeBestNonlinearTerminationMode, - select_forward_mode_autodiff, select_reverse_mode_autodiff, - select_jacobian_autodiff, - construct_jacobian_cache, - DescentResult, - SteepestDescent, NewtonDescent, DampedNewtonDescent, Dogleg, - GeodesicAcceleration, - reset_timer!, @static_timeit, - init_nonlinearsolve_trace, update_trace!, reset! - +using LinearAlgebra: LinearAlgebra, norm +using LineSearch: BackTracking +using NonlinearSolveBase: NonlinearSolveBase, InternalAPI, AbstractNonlinearSolveAlgorithm, + AbstractNonlinearSolveCache, Utils, L2_NORM + +using Preferences: set_preferences! +using SciMLBase: SciMLBase, NLStats, ReturnCode, AbstractNonlinearProblem, NonlinearProblem, + NonlinearLeastSquaresProblem +using SymbolicIndexingInterface: SymbolicIndexingInterface +using StaticArraysCore: StaticArray + +# Default Algorithm +using NonlinearSolveFirstOrder: NewtonRaphson, TrustRegion, LevenbergMarquardt, GaussNewton, + RUS using NonlinearSolveQuasiNewton: Broyden, Klement +using SimpleNonlinearSolve: SimpleBroyden, SimpleKlement -using Preferences: Preferences, set_preferences! -using RecursiveArrayTools: recursivecopy! -using SciMLBase: SciMLBase, AbstractNonlinearAlgorithm, AbstractNonlinearProblem, - _unwrap_val, isinplace, NLStats, NonlinearFunction, - NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode, get_du, step!, - set_u!, LinearProblem, IdentityOperator -using SciMLOperators: AbstractSciMLOperator -using SimpleNonlinearSolve: SimpleNonlinearSolve -using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix - -# AD Support -using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff, - AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, AutoSparse -using DifferentiationInterface: DifferentiationInterface -using FiniteDiff: FiniteDiff -using ForwardDiff: ForwardDiff, Dual -using SciMLJacobianOperators: VecJacOperator, JacVecOperator, StatefulJacobianOperator +# Default AD Support +using FiniteDiff: FiniteDiff # Default Finite Difference Method +using ForwardDiff: ForwardDiff # Default Forward Mode AD -## Sparse AD Support -using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC -using SparseMatrixColorings: SparseMatrixColorings # NOTE: This triggers an extension in NonlinearSolveBase +# Sparse AD Support: Implemented via extensions +using SparseArrays: SparseArrays +using SparseMatrixColorings: SparseMatrixColorings -const DI = DifferentiationInterface +const SII = SymbolicIndexingInterface -include("timer_outputs.jl") -include("internal/helpers.jl") +include("helpers.jl") -include("algorithms/extension_algs.jl") +include("polyalg.jl") +# include("extension_algs.jl") -include("utils.jl") include("default.jl") -const ALL_SOLVER_TYPES = [ - Nothing, AbstractNonlinearSolveAlgorithm, GeneralizedDFSane, - GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, - LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL, - SpeedMappingJL, FixedPointAccelerationJL, SIAMFANLEquationsJL, - CMINPACK, PETScSNES, - NonlinearSolvePolyAlgorithm{:NLLS, <:Any}, NonlinearSolvePolyAlgorithm{:NLS, <:Any} -] +# const ALL_SOLVER_TYPES = [ +# Nothing, AbstractNonlinearSolveAlgorithm, GeneralizedDFSane, +# GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, +# LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL, +# SpeedMappingJL, FixedPointAccelerationJL, SIAMFANLEquationsJL, +# CMINPACK, PETScSNES, +# NonlinearSolvePolyAlgorithm{:NLLS, <:Any}, NonlinearSolvePolyAlgorithm{:NLS, <:Any} +# ] -include("internal/forward_diff.jl") # we need to define after the algorithms +# include("internal/forward_diff.jl") # we need to define after the algorithms @setup_workload begin - nlfuncs = ( - (NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), - (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1]) - ) - probs_nls = NonlinearProblem[] - for (fn, u0) in nlfuncs - push!(probs_nls, NonlinearProblem(fn, u0, 2.0)) - end - - nls_algs = ( - nothing - ) - - probs_nlls = NonlinearLeastSquaresProblem[] - nlfuncs = ( - (NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]), - (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]), - ( - NonlinearFunction{true}( - (du, u, p) -> du[1] = u[1] * u[1] - p, resid_prototype = zeros(1)), - [0.1, 0.0]), - ( - NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), - resid_prototype = zeros(4)), - [0.1, 0.1] - ) - ) - for (fn, u0) in nlfuncs - push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0)) - end - - nlls_algs = ( - LevenbergMarquardt(), - GaussNewton(), - TrustRegion(), - nothing - ) + include(joinpath(@__DIR__, "..", "common", "nonlinear_problem_workloads.jl")) + include(joinpath(@__DIR__, "..", "common", "nlls_problem_workloads.jl")) @compile_workload begin - for prob in probs_nls, alg in nls_algs - solve(prob, alg; abstol = 1e-2, verbose = false) + for prob in nonlinear_problems + CommonSolve.solve(prob, nothing; abstol = 1e-2, verbose = false) end - for prob in probs_nlls, alg in nlls_algs - solve(prob, alg; abstol = 1e-2, verbose = false) + + for prob in nlls_problems + CommonSolve.solve(prob, nothing; abstol = 1e-2, verbose = false) end end end # Rexexports -@reexport using SciMLBase, NonlinearSolveBase +@reexport using SciMLBase, NonlinearSolveBase, LineSearch, ADTypes @reexport using NonlinearSolveFirstOrder, NonlinearSolveSpectralMethods, NonlinearSolveQuasiNewton, SimpleNonlinearSolve -@reexport using LineSearch -@reexport using ADTypes +@reexport using LinearSolve -export NonlinearSolvePolyAlgorithm, RobustMultiNewton, - FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg +# Poly Algorithms +export NonlinearSolvePolyAlgorithm, + RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg # Extension Algorithms export LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL, diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl deleted file mode 100644 index a5e52cab4..000000000 --- a/src/algorithms/extension_algs.jl +++ /dev/null @@ -1,469 +0,0 @@ -# This file only include the algorithm struct to be exported by NonlinearSolve.jl. The main -# functionality is implemented as package extensions -""" - LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) - -Wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl) -for solving `NonlinearLeastSquaresProblem`. - -### Arguments - - - `alg`: Algorithm to use. Can be `:lm` or `:dogleg`. - -### Keyword Arguments - - - `linsolve`: Linear solver to use. Can be `:qr`, `:cholesky` or `:lsmr`. If `nothing`, - then `LeastSquaresOptim.jl` will choose the best linear solver based on the Jacobian - structure. - - `autodiff`: Automatic differentiation / Finite Differences. Can be `:central` or - `:forward`. - -!!! note - - This algorithm is only available if `LeastSquaresOptim.jl` is installed and loaded. -""" -struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm - autodiff -end - -function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff = :central) - @assert alg in (:lm, :dogleg) - @assert linsolve === nothing || linsolve in (:qr, :cholesky, :lsmr) - autodiff isa Symbol && @assert autodiff in (:central, :forward) - - if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing - error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded") - end - - return LeastSquaresOptimJL{alg, linsolve}(autodiff) -end - -""" - FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6, - factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, - minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, - autodiff = nothing) - -Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl) -for solving `NonlinearLeastSquaresProblem`. For details about the other keyword arguments -see the documentation for `FastLevenbergMarquardt.jl`. - -!!! warning - - This is not really the fastest solver. It is called that since the original package - is called "Fast". `LevenbergMarquardt()` is almost always a better choice. - -### Arguments - - - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`. - -### 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 - `nothing` which means that a default is selected according to the problem specification! - -!!! note - - This algorithm is only available if `FastLevenbergMarquardt.jl` is installed and loaded. -""" -@concrete struct FastLevenbergMarquardtJL{linsolve} <: - AbstractNonlinearSolveExtensionAlgorithm - autodiff - factor - factoraccept - factorreject - factorupdate::Symbol - minscale - maxscale - minfactor - maxfactor -end - -function FastLevenbergMarquardtJL( - linsolve::Symbol = :cholesky; factor = 1e-6, factoraccept = 13.0, - factorreject = 3.0, factorupdate = :marquardt, minscale = 1e-12, - maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, autodiff = nothing) - @assert linsolve in (:qr, :cholesky) - @assert factorupdate in (:marquardt, :nielson) - - if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing - error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded") - end - - return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject, - factorupdate, minscale, maxscale, minfactor, maxfactor) -end - -""" - CMINPACK(; method::Symbol = :auto, autodiff = missing) - -### Keyword Arguments - - - `method`: the choice of method for the solver. - - `autodiff`: Defaults to `missing`, which means we will default to letting `MINPACK` - construct the jacobian if `f.jac` is not provided. In other cases, we use it to generate - a jacobian similar to other NonlinearSolve solvers. - -### Submethod Choice - -The keyword argument `method` can take on different value depending on which method of -`fsolve` you are calling. The standard choices of `method` are: - - - `:hybr`: Modified version of Powell's algorithm. Uses MINPACK routine - [`hybrd1`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/hybrd1.c) - - `:lm`: Levenberg-Marquardt. Uses MINPACK routine - [`lmdif1`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/lmdif1.c) - - `:lmdif`: Advanced Levenberg-Marquardt (more options available with `; kwargs...`). See - MINPACK routine [`lmdif`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/lmdif.c) - for more information - - `:hybrd`: Advanced modified version of Powell's algorithm (more options available with - `; kwargs...`). See MINPACK routine - [`hybrd`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/hybrd.c) - for more information - -If a Jacobian is supplied as part of the [`NonlinearFunction`](@ref nonlinearfunctions), -then the following methods are allowed: - - - `:hybr`: Advanced modified version of Powell's algorithm with user supplied Jacobian. - Additional arguments are available via `; kwargs...`. See MINPACK routine - [`hybrj`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/hybrj.c) - for more information - - `:lm`: Advanced Levenberg-Marquardt with user supplied Jacobian. Additional arguments - are available via `; kwargs...`. See MINPACK routine - [`lmder`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/lmder.c) - for more information - -The default choice of `:auto` selects `:hybr` for NonlinearProblem and `:lm` for -NonlinearLeastSquaresProblem. - -!!! note - - This algorithm is only available if `MINPACK.jl` is installed and loaded. -""" -@concrete struct CMINPACK <: AbstractNonlinearSolveExtensionAlgorithm - method::Symbol - autodiff -end - -function CMINPACK(; method::Symbol = :auto, autodiff = missing) - if Base.get_extension(@__MODULE__, :NonlinearSolveMINPACKExt) === nothing - error("CMINPACK requires MINPACK.jl to be loaded") - end - return CMINPACK(method, autodiff) -end - -""" - NLsolveJL(; method = :trust_region, autodiff = :central, linesearch = Static(), - linsolve = (x, A, b) -> copyto!(x, A\\b), factor = one(Float64), autoscale = true, - m = 10, beta = one(Float64)) - -### Keyword Arguments - - - `method`: the choice of method for solving the nonlinear system. - - `autodiff`: the choice of method for generating the Jacobian. Defaults to `:central` or - central differencing via FiniteDiff.jl. The other choices are `:forward` or `ADTypes` - similar to other solvers in NonlinearSolve. - - `linesearch`: the line search method to be used within the solver method. The choices - are line search types from - [LineSearches.jl](https://github.com/JuliaNLSolvers/LineSearches.jl). - - `linsolve`: a function `linsolve(x, A, b)` that solves `Ax = b`. - - `factor`: determines the size of the initial trust region. This size is set to the - product of factor and the euclidean norm of `u0` if nonzero, or else to factor itself. - - `autoscale`: if true, then the variables will be automatically rescaled. The scaling - factors are the norms of the Jacobian columns. - - `m`: the amount of history in the Anderson method. Naive "Picard"-style iteration can be - achieved by setting m=0, but that isn't advisable for contractions whose Lipschitz - constants are close to 1. If convergence fails, though, you may consider lowering it. - - `beta`: It is also known as DIIS or Pulay mixing, this method is based on the - acceleration of the fixed-point iteration xₙ₊₁ = xₙ + beta*f(xₙ), where by default - beta = 1. - -!!! warning - - Line Search Algorithms from [`LineSearch.jl`](https://github.com/SciML/LineSearch.jl) - aren't supported by `NLsolveJL`. Instead, use the line search algorithms from - [`LineSearches.jl`](https://github.com/JuliaNLSolvers/LineSearches.jl). - -### Submethod Choice - -Choices for methods in `NLsolveJL`: - - - `:anderson`: Anderson-accelerated fixed-point iteration - - `:broyden`: Broyden's quasi-Newton method - - `:newton`: Classical Newton method with an optional line search - - `:trust_region`: Trust region Newton method (the default choice) - -For more information on these arguments, consult the -[NLsolve.jl documentation](https://github.com/JuliaNLSolvers/NLsolve.jl). - -!!! note - - This algorithm is only available if `NLsolve.jl` is installed and loaded. -""" -@concrete struct NLsolveJL <: AbstractNonlinearSolveExtensionAlgorithm - method::Symbol - autodiff - linesearch - linsolve - factor - autoscale::Bool - m::Int - beta -end - -function NLsolveJL(; method = :trust_region, autodiff = :central, linesearch = missing, - linsolve = (x, A, b) -> copyto!(x, A \ b), factor = 1.0, - autoscale = true, m = 10, beta = one(Float64)) - if Base.get_extension(@__MODULE__, :NonlinearSolveNLsolveExt) === nothing - error("NLsolveJL requires NLsolve.jl to be loaded") - end - - if autodiff isa Symbol && autodiff !== :central && autodiff !== :forward - error("`autodiff` must be `:central` or `:forward`.") - end - - return NLsolveJL(method, autodiff, linesearch, linsolve, factor, autoscale, m, beta) -end - -""" - NLSolversJL(method; autodiff = nothing) - NLSolversJL(; method, autodiff = nothing) - -Wrapper over NLSolvers.jl Nonlinear Equation Solvers. We automatically construct the -jacobian function and supply it to the solver. - -### Arguments - - - `method`: the choice of method for solving the nonlinear system. See the documentation - for NLSolvers.jl for more information. - - `autodiff`: the choice of method for generating the Jacobian. Defaults to `nothing` - which means that a default is selected according to the problem specification. Can be - any valid ADTypes.jl autodiff type (conditional on that backend being supported in - NonlinearSolve.jl). -""" -struct NLSolversJL{M, AD} <: AbstractNonlinearSolveExtensionAlgorithm - method::M - autodiff::AD - - function NLSolversJL(method, autodiff) - if Base.get_extension(@__MODULE__, :NonlinearSolveNLSolversExt) === nothing - error("NLSolversJL requires NLSolvers.jl to be loaded") - end - - return new{typeof(method), typeof(autodiff)}(method, autodiff) - end -end - -NLSolversJL(method; autodiff = nothing) = NLSolversJL(method, autodiff) -NLSolversJL(; method, autodiff = nothing) = NLSolversJL(method, autodiff) - -""" - SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, - orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) - -Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl/) for solving -Fixed Point Problems. We allow using this algorithm to solve root finding problems as well. - -### Keyword Arguments - - - `σ_min`: Setting to `1` may avoid stalling (see [lepage2021alternating](@cite)). - - `stabilize`: performs a stabilization mapping before extrapolating. Setting to `true` - may improve the performance for applications like accelerating the EM or MM algorithms - (see [lepage2021alternating](@cite)). - - `check_obj`: In case of NaN or Inf values, the algorithm restarts at the best past - iterate. - - `orders`: determines ACX's alternating order. Must be between `1` and `3` (where `1` - means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the - latter being potentially better for highly non-linear applications (see - [lepage2021alternating](@cite)). - - `time_limit`: time limit for the algorithm. - -!!! note - - This algorithm is only available if `SpeedMapping.jl` is installed and loaded. -""" -@concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm - σ_min - stabilize::Bool - check_obj::Bool - orders::Vector{Int} -end - -function SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, - orders::Vector{Int} = [3, 3, 2]) - if Base.get_extension(@__MODULE__, :NonlinearSolveSpeedMappingExt) === nothing - error("SpeedMappingJL requires SpeedMapping.jl to be loaded") - end - - return SpeedMappingJL(σ_min, stabilize, check_obj, orders) -end - -""" - FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, - condition_number_threshold = missing, extrapolation_period = missing, - replace_invalids = :NoAction) - -Wrapper over [FixedPointAcceleration.jl](https://s-baumann.github.io/FixedPointAcceleration.jl/) -for solving Fixed Point Problems. We allow using this algorithm to solve root finding -problems as well. - -### Keyword Arguments - - - `algorithm`: The algorithm to use. Can be `:Anderson`, `:MPE`, `:RRE`, `:VEA`, `:SEA`, - `:Simple`, `:Aitken` or `:Newton`. - - `m`: The number of previous iterates to use for the extrapolation. Only valid for - `:Anderson`. - - `condition_number_threshold`: The condition number threshold for Least Squares Problem. - Only valid for `:Anderson`. - - `extrapolation_period`: The number of iterates between extrapolations. Only valid for - `:MPE`, `:RRE`, `:VEA` and `:SEA`. Defaults to `7` for `:MPE` & `:RRE`, and `6` for - `:SEA` and `:VEA`. For `:SEA` and `:VEA`, this must be a multiple of `2`. - - `replace_invalids`: The method to use for replacing invalid iterates. Can be - `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`. - -!!! note - - This algorithm is only available if `FixedPointAcceleration.jl` is installed and loaded. -""" -@concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm - algorithm::Symbol - extrapolation_period::Int - replace_invalids::Symbol - dampening - m::Int - condition_number_threshold -end - -function FixedPointAccelerationJL(; - algorithm = :Anderson, m = missing, condition_number_threshold = missing, - extrapolation_period = missing, replace_invalids = :NoAction, dampening = 1.0) - if Base.get_extension(@__MODULE__, :NonlinearSolveFixedPointAccelerationExt) === nothing - error("FixedPointAccelerationJL requires FixedPointAcceleration.jl to be loaded") - end - - @assert algorithm in (:Anderson, :MPE, :RRE, :VEA, :SEA, :Simple, :Aitken, :Newton) - @assert replace_invalids in (:ReplaceInvalids, :ReplaceVector, :NoAction) - - if algorithm !== :Anderson - if condition_number_threshold !== missing - error("`condition_number_threshold` is only valid for Anderson acceleration") - end - if m !== missing - error("`m` is only valid for Anderson acceleration") - end - end - condition_number_threshold === missing && (condition_number_threshold = 1e3) - m === missing && (m = 10) - - if algorithm !== :MPE && algorithm !== :RRE && algorithm !== :VEA && algorithm !== :SEA - if extrapolation_period !== missing - error("`extrapolation_period` is only valid for MPE, RRE, VEA and SEA") - end - end - if extrapolation_period === missing - extrapolation_period = algorithm === :SEA || algorithm === :VEA ? 6 : 7 - else - if (algorithm === :SEA || algorithm === :VEA) && extrapolation_period % 2 != 0 - error("`extrapolation_period` must be multiples of 2 for SEA and VEA") - end - end - - return FixedPointAccelerationJL(algorithm, extrapolation_period, replace_invalids, - dampening, m, condition_number_threshold) -end - -""" - SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, - autodiff = missing) - -### Keyword Arguments - - - `method`: the choice of method for solving the nonlinear system. - - `delta`: initial pseudo time step, default is 1e-3. - - `linsolve` : JFNK linear solvers, choices are `gmres` and `bicgstab`. - - `m`: Depth for Anderson acceleration, default as 0 for Picard iteration. - - `beta`: Anderson mixing parameter, change f(x) to (1-beta)x+beta*f(x), - equivalent to accelerating damped Picard iteration. - - `autodiff`: Defaults to `missing`, which means we will default to letting - `SIAMFANLEquations` construct the jacobian if `f.jac` is not provided. In other cases, - we use it to generate a jacobian similar to other NonlinearSolve solvers. - -### Submethod Choice - - - `:newton`: Classical Newton method. - - `:pseudotransient`: Pseudo transient method. - - `:secant`: Secant method for scalar equations. - - `:anderson`: Anderson acceleration for fixed point iterations. - -!!! note - - This algorithm is only available if `SIAMFANLEquations.jl` is installed and loaded. -""" -@concrete struct SIAMFANLEquationsJL <: AbstractNonlinearSolveExtensionAlgorithm - method::Symbol - delta - linsolve <: Union{Symbol, Nothing} - m::Int - beta - autodiff -end - -function SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, - m = 0, beta = 1.0, autodiff = missing) - if Base.get_extension(@__MODULE__, :NonlinearSolveSIAMFANLEquationsExt) === nothing - error("SIAMFANLEquationsJL requires SIAMFANLEquations.jl to be loaded") - end - return SIAMFANLEquationsJL(method, delta, linsolve, m, beta, autodiff) -end - -""" - PETScSNES(; petsclib = missing, autodiff = nothing, mpi_comm = missing, kwargs...) - -Wrapper over [PETSc.jl](https://github.com/JuliaParallel/PETSc.jl) SNES solvers. - -### Keyword Arguments - - - `petsclib`: PETSc library to use. If set to `missing`, then we will use the first - available PETSc library in `PETSc.petsclibs` based on the problem element type. - - `autodiff`: the choice of method for generating the Jacobian. Defaults to `nothing` - which means that a default is selected according to the problem specification. Can be - any valid ADTypes.jl autodiff type (conditional on that backend being supported in - NonlinearSolve.jl). If set to `missing`, then PETSc computes the Jacobian using finite - differences. - - `mpi_comm`: MPI communicator to use. If set to `missing`, then we will use - `MPI.COMM_SELF`. - - `kwargs`: Keyword arguments to be passed to the PETSc SNES solver. See [PETSc SNES - documentation](https://petsc.org/release/manual/snes/) and - [SNESSetFromOptions](https://petsc.org/release/manualpages/SNES/SNESSetFromOptions) - for more information. - -### Options via `CommonSolve.solve` - -These options are forwarded from `solve` to the PETSc SNES solver. If these are provided to -`kwargs`, then they will be ignored. - -| `solve` option | PETSc SNES option | -|:-------------- |:----------------- | -| `atol` | `snes_atol` | -| `rtol` | `snes_rtol` | -| `maxiters` | `snes_max_it` | -| `show_trace` | `snes_monitor` | - -!!! note - - This algorithm is only available if `PETSc.jl` is installed and loaded. -""" -@concrete struct PETScSNES <: AbstractNonlinearSolveExtensionAlgorithm - petsclib - mpi_comm - autodiff <: Union{Missing, Nothing, ADTypes.AbstractADType} - snes_options -end - -function PETScSNES(; petsclib = missing, autodiff = nothing, mpi_comm = missing, kwargs...) - if Base.get_extension(@__MODULE__, :NonlinearSolvePETScExt) === nothing - error("PETScSNES requires PETSc.jl to be loaded") - end - return PETScSNES(petsclib, mpi_comm, autodiff, kwargs) -end diff --git a/src/algorithms/trust_region.jl b/src/algorithms/trust_region.jl deleted file mode 100644 index b42099d46..000000000 --- a/src/algorithms/trust_region.jl +++ /dev/null @@ -1,36 +0,0 @@ -""" - TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = nothing, - radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, - initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, - shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, - shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1, - max_shrink_times::Int = 32, - vjp_autodiff = nothing, autodiff = nothing, jvp_autodiff = nothing) - -An advanced TrustRegion implementation with support for efficient handling of sparse -matrices via colored automatic differentiation and preconditioned linear solvers. Designed -for large-scale and numerically-difficult nonlinear systems. - -### Keyword Arguments - - - `radius_update_scheme`: the scheme used to update the trust region radius. Defaults to - `RadiusUpdateSchemes.Simple`. See [`RadiusUpdateSchemes`](@ref) for more details. For a - review on trust region radius update schemes, see [yuan2015recent](@citet). - -For the remaining arguments, see [`NonlinearSolve.GenericTrustRegionScheme`](@ref) -documentation. -""" -function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = nothing, - radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, - initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, - shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, - shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1, - max_shrink_times::Int = 32, - autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing) - descent = Dogleg(; linsolve, precs) - trustregion = GenericTrustRegionScheme(; - method = radius_update_scheme, step_threshold, shrink_threshold, expand_threshold, - shrink_factor, expand_factor, initial_trust_radius, max_trust_radius) - return GeneralizedFirstOrderAlgorithm{concrete_jac, :TrustRegion}(; - trustregion, descent, autodiff, vjp_autodiff, jvp_autodiff, max_shrink_times) -end diff --git a/src/default.jl b/src/default.jl index 7302d95a9..6021a98b1 100644 --- a/src/default.jl +++ b/src/default.jl @@ -1,496 +1,3 @@ -# Poly Algorithms -""" - NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS); - start_index = 1) where {pType} - -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. - -### 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. - -### Keyword Arguments - - - `start_index`: the index to start at. Defaults to `1`. - -### Example - -```julia -using NonlinearSolve - -alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden())) -``` -""" -struct NonlinearSolvePolyAlgorithm{pType, N, A} <: AbstractNonlinearSolveAlgorithm{:PolyAlg} - algs::A - start_index::Int - - function NonlinearSolvePolyAlgorithm( - algs, ::Val{pType} = Val(:NLS); start_index::Int = 1) where {pType} - @assert pType ∈ (:NLS, :NLLS) - @assert 0 < start_index ≤ length(algs) - algs = Tuple(algs) - return new{pType, length(algs), typeof(algs)}(algs, start_index) - end -end - -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 - num = " [$(i)]: " - print(io, num) - __show_algorithm(io, alg.algs[i], get_name(alg.algs[i]), length(num)) - i == N || println(io) - end -end - -@concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N, timeit} <: - AbstractNonlinearSolveCache{iip, timeit} - caches - alg - best::Int - current::Int - nsteps::Int - stats::NLStats - total_time::Float64 - maxtime - retcode::ReturnCode.T - force_stop::Bool - maxiters::Int - internalnorm - u0 - u0_aliased - alias_u0::Bool -end - -function SymbolicIndexingInterface.symbolic_container(cache::NonlinearSolvePolyAlgorithmCache) - cache.caches[cache.current] -end -SymbolicIndexingInterface.state_values(cache::NonlinearSolvePolyAlgorithmCache) = cache.u0 - -function Base.show( - io::IO, cache::NonlinearSolvePolyAlgorithmCache{pType, N}) where {pType, N} - problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem") - println(io, "NonlinearSolvePolyAlgorithmCache for $(problem_kind) with $(N) algorithms") - best_alg = ifelse(cache.best == -1, "nothing", cache.best) - println(io, "Best algorithm: $(best_alg)") - println(io, "Current algorithm: $(cache.current)") - println(io, "nsteps: $(cache.nsteps)") - println(io, "retcode: $(cache.retcode)") - __show_cache(io, cache.caches[cache.current], 0) -end - -function reinit_cache!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) - foreach(c -> reinit_cache!(c, args...; kwargs...), cache.caches) - cache.current = cache.alg.start_index - __reinit_internal!(cache.stats) - cache.nsteps = 0 - cache.total_time = 0.0 -end - -for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) - algType = NonlinearSolvePolyAlgorithm{pType} - @eval begin - function SciMLBase.__init( - prob::$probType, alg::$algType{N}, args...; stats = empty_nlstats(), - maxtime = nothing, maxiters = 1000, internalnorm = L2_NORM, - alias_u0 = false, verbose = true, kwargs...) where {N} - if (alias_u0 && !ismutable(prob.u0)) - verbose && @warn "`alias_u0` has been set to `true`, but `u0` is \ - immutable (checked using `ArrayInterface.ismutable`)." - alias_u0 = false # If immutable don't care about aliasing - end - u0 = prob.u0 - if alias_u0 - u0_aliased = copy(u0) - else - u0_aliased = u0 # Irrelevant - end - alias_u0 && (prob = remake(prob; u0 = u0_aliased)) - return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N, maxtime !== nothing}( - map( - solver -> SciMLBase.__init(prob, solver, args...; stats, maxtime, - internalnorm, alias_u0, verbose, kwargs...), - alg.algs), - alg, - -1, - alg.start_index, - 0, - stats, - 0.0, - maxtime, - ReturnCode.Default, - false, - maxiters, - internalnorm, - u0, - u0_aliased, - alias_u0) - end - end -end - -@generated function SciMLBase.solve!(cache::NonlinearSolvePolyAlgorithmCache{ - iip, N}) where {iip, N} - calls = [quote - 1 ≤ cache.current ≤ length(cache.caches) || - error("Current choices shouldn't get here!") - end] - - cache_syms = [gensym("cache") for i in 1:N] - sol_syms = [gensym("sol") for i in 1:N] - u_result_syms = [gensym("u_result") for i in 1:N] - for i in 1:N - push!(calls, - quote - $(cache_syms[i]) = cache.caches[$(i)] - if $(i) == cache.current - cache.alias_u0 && copyto!(cache.u0_aliased, cache.u0) - $(sol_syms[i]) = SciMLBase.solve!($(cache_syms[i])) - if SciMLBase.successful_retcode($(sol_syms[i])) - stats = $(sol_syms[i]).stats - if cache.alias_u0 - copyto!(cache.u0, $(sol_syms[i]).u) - $(u_result_syms[i]) = cache.u0 - else - $(u_result_syms[i]) = $(sol_syms[i]).u - end - fu = get_fu($(cache_syms[i])) - return __build_solution_less_specialize( - $(sol_syms[i]).prob, cache.alg, $(u_result_syms[i]), - fu; retcode = $(sol_syms[i]).retcode, stats, - original = $(sol_syms[i]), trace = $(sol_syms[i]).trace) - elseif cache.alias_u0 - # For safety we need to maintain a copy of the solution - $(u_result_syms[i]) = copy($(sol_syms[i]).u) - end - cache.current = $(i + 1) - end - end) - end - - resids = map(x -> Symbol("$(x)_resid"), cache_syms) - for (sym, resid) in zip(cache_syms, resids) - push!(calls, :($(resid) = @isdefined($(sym)) ? get_fu($(sym)) : nothing)) - end - push!(calls, quote - fus = tuple($(Tuple(resids)...)) - minfu, idx = __findmin(cache.internalnorm, fus) - end) - for i in 1:N - push!(calls, quote - if idx == $(i) - if cache.alias_u0 - u = $(u_result_syms[i]) - else - u = get_u(cache.caches[$i]) - end - end - end) - end - push!(calls, - quote - retcode = cache.caches[idx].retcode - if cache.alias_u0 - copyto!(cache.u0, u) - u = cache.u0 - end - return __build_solution_less_specialize( - cache.caches[idx].prob, cache.alg, u, fus[idx]; - retcode, stats = cache.stats, cache.caches[idx].trace) - end) - - return Expr(:block, calls...) -end - -@generated function __step!( - cache::NonlinearSolvePolyAlgorithmCache{iip, N}, args...; kwargs...) where {iip, N} - calls = [] - cache_syms = [gensym("cache") for i in 1:N] - for i in 1:N - push!(calls, - quote - $(cache_syms[i]) = cache.caches[$(i)] - if $(i) == cache.current - __step!($(cache_syms[i]), args...; kwargs...) - $(cache_syms[i]).nsteps += 1 - if !not_terminated($(cache_syms[i])) - if SciMLBase.successful_retcode($(cache_syms[i]).retcode) - cache.best = $(i) - cache.force_stop = true - cache.retcode = $(cache_syms[i]).retcode - else - cache.current = $(i + 1) - end - end - return - end - end) - end - - push!(calls, quote - if !(1 ≤ cache.current ≤ length(cache.caches)) - minfu, idx = __findmin(cache.internalnorm, cache.caches) - cache.best = idx - cache.retcode = cache.caches[cache.best].retcode - cache.force_stop = true - return - end - end) - - return Expr(:block, calls...) -end - -for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) - algType = NonlinearSolvePolyAlgorithm{pType} - @eval begin - @generated function SciMLBase.__solve( - prob::$probType, alg::$algType{N}, args...; stats = empty_nlstats(), - alias_u0 = false, verbose = true, kwargs...) where {N} - sol_syms = [gensym("sol") for _ in 1:N] - prob_syms = [gensym("prob") for _ in 1:N] - u_result_syms = [gensym("u_result") for _ in 1:N] - calls = [quote - current = alg.start_index - if (alias_u0 && !ismutable(prob.u0)) - verbose && @warn "`alias_u0` has been set to `true`, but `u0` is \ - immutable (checked using `ArrayInterface.ismutable`)." - alias_u0 = false # If immutable don't care about aliasing - end - u0 = prob.u0 - u0_aliased = alias_u0 ? zero(u0) : u0 - end] - for i in 1:N - cur_sol = sol_syms[i] - push!(calls, - quote - if current == $i - if alias_u0 - copyto!(u0_aliased, u0) - $(prob_syms[i]) = remake(prob; u0 = u0_aliased) - else - $(prob_syms[i]) = prob - end - $(cur_sol) = SciMLBase.__solve( - $(prob_syms[i]), alg.algs[$(i)], args...; - stats, alias_u0, verbose, kwargs...) - if SciMLBase.successful_retcode($(cur_sol)) - if alias_u0 - copyto!(u0, $(cur_sol).u) - $(u_result_syms[i]) = u0 - else - $(u_result_syms[i]) = $(cur_sol).u - end - return __build_solution_less_specialize( - prob, alg, $(u_result_syms[i]), $(cur_sol).resid; - $(cur_sol).retcode, $(cur_sol).stats, - original = $(cur_sol), trace = $(cur_sol).trace) - elseif alias_u0 - # For safety we need to maintain a copy of the solution - $(u_result_syms[i]) = copy($(cur_sol).u) - end - current = $(i + 1) - end - end) - end - - resids = map(x -> Symbol("$(x)_resid"), sol_syms) - for (sym, resid) in zip(sol_syms, resids) - push!(calls, :($(resid) = @isdefined($(sym)) ? $(sym).resid : nothing)) - end - - push!(calls, quote - resids = tuple($(Tuple(resids)...)) - minfu, idx = __findmin(L2_NORM, resids) - end) - - for i in 1:N - push!(calls, - quote - if idx == $i - if alias_u0 - copyto!(u0, $(u_result_syms[i])) - $(u_result_syms[i]) = u0 - else - $(u_result_syms[i]) = $(sol_syms[i]).u - end - return __build_solution_less_specialize( - prob, alg, $(u_result_syms[i]), $(sol_syms[i]).resid; - $(sol_syms[i]).retcode, $(sol_syms[i]).stats, - $(sol_syms[i]).trace, original = $(sol_syms[i])) - end - end) - end - push!(calls, :(error("Current choices shouldn't get here!"))) - - return Expr(:block, calls...) - end - end -end - -""" - RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = nothing, autodiff = nothing) - -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). - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. -""" -function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = nothing, autodiff = nothing) where {T} - if __is_complex(T) - # Let's atleast have something here for complex numbers - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) - else - algs = (TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff, - radius_update_scheme = RadiusUpdateSchemes.Bastin), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff)) - end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) -end - -""" - FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = nothing, must_use_jacobian::Val = Val(false), - prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing, - u0_len::Union{Int, Nothing} = nothing) where {T} - -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. - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. - -### Keyword Arguments - - - `u0_len`: The length of the initial guess. If this is `nothing`, then the length of the - initial guess is not checked. If this is an integer and it is less than `25`, we use - jacobian based methods. -""" -function FastShortcutNonlinearPolyalg( - ::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = nothing, must_use_jacobian::Val{JAC} = Val(false), - prefer_simplenonlinearsolve::Val{SA} = Val(false), - u0_len::Union{Int, Nothing} = nothing, autodiff = nothing) where {T, JAC, SA} - start_index = 1 - if JAC - if __is_complex(T) - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) - else - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end - else - # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians - # and thus are not included in the polyalgorithm - if SA - if __is_complex(T) - algs = (SimpleBroyden(), - Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - SimpleKlement(), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) - else - start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1 - algs = (SimpleBroyden(), - Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - SimpleKlement(), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end - else - if __is_complex(T) - algs = ( - Broyden(), Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - Klement(; linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) - else - # TODO: This number requires a bit rigorous testing - start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1 - algs = ( - Broyden(; autodiff), - Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - Klement(; linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end - end - end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLS); start_index) -end - -""" - FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = nothing, autodiff = nothing, kwargs...) - -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. - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. -""" -function FastShortcutNLLSPolyalg( - ::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = nothing, autodiff = nothing, kwargs...) where {T} - if __is_complex(T) - algs = (GaussNewton(; concrete_jac, linsolve, precs, autodiff, kwargs...), - LevenbergMarquardt(; - linsolve, precs, autodiff, disable_geodesic = Val(true), kwargs...), - LevenbergMarquardt(; linsolve, precs, autodiff, kwargs...)) - else - algs = (GaussNewton(; concrete_jac, linsolve, precs, autodiff, kwargs...), - LevenbergMarquardt(; - linsolve, precs, disable_geodesic = Val(true), autodiff, kwargs...), - TrustRegion(; concrete_jac, linsolve, precs, autodiff, kwargs...), - GaussNewton(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff, kwargs...), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff, kwargs...), - LevenbergMarquardt(; linsolve, precs, autodiff, kwargs...)) - end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS)) -end - -## Defaults - ## TODO: In the long run we want to use an `Assumptions` API like LinearSolve to specify ## the conditioning of the Jacobian and such @@ -503,31 +10,43 @@ end ## the trouble of specifying a custom jacobian function, we should use algorithms that ## can use that! function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...) - must_use_jacobian = Val(prob.f.jac !== nothing) - return SciMLBase.__init(prob, + must_use_jacobian = Val(SciMLBase.has_jac(prob.f)) + return SciMLBase.__init( + prob, FastShortcutNonlinearPolyalg( - eltype(prob.u0); must_use_jacobian, u0_len = length(prob.u0)), + eltype(prob.u0); must_use_jacobian, u0_len = length(prob.u0) + ), args...; - kwargs...) + kwargs... + ) end function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...) - must_use_jacobian = Val(prob.f.jac !== nothing) - prefer_simplenonlinearsolve = Val(prob.u0 isa SArray) - return SciMLBase.__solve(prob, - FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian, - prefer_simplenonlinearsolve, u0_len = length(prob.u0)), + must_use_jacobian = Val(SciMLBase.has_jac(prob.f)) + prefer_simplenonlinearsolve = Val(prob.u0 isa StaticArray) + return SciMLBase.__solve( + prob, + FastShortcutNonlinearPolyalg( + eltype(prob.u0); + must_use_jacobian, + prefer_simplenonlinearsolve, + u0_len = length(prob.u0) + ), args...; - kwargs...) + kwargs... + ) end function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) return SciMLBase.__init( - prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; kwargs...) + prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; kwargs... + ) end function SciMLBase.__solve( - prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) + prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs... +) return SciMLBase.__solve( - prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; kwargs...) + prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; kwargs... + ) end diff --git a/src/extension_algs.jl b/src/extension_algs.jl new file mode 100644 index 000000000..fb922fb14 --- /dev/null +++ b/src/extension_algs.jl @@ -0,0 +1,469 @@ +# # This file only include the algorithm struct to be exported by NonlinearSolve.jl. The main +# # functionality is implemented as package extensions +# """ +# LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff::Symbol = :central) + +# Wrapper over [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl) +# for solving `NonlinearLeastSquaresProblem`. + +# ### Arguments + +# - `alg`: Algorithm to use. Can be `:lm` or `:dogleg`. + +# ### Keyword Arguments + +# - `linsolve`: Linear solver to use. Can be `:qr`, `:cholesky` or `:lsmr`. If `nothing`, +# then `LeastSquaresOptim.jl` will choose the best linear solver based on the Jacobian +# structure. +# - `autodiff`: Automatic differentiation / Finite Differences. Can be `:central` or +# `:forward`. + +# !!! note + +# This algorithm is only available if `LeastSquaresOptim.jl` is installed and loaded. +# """ +# struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm +# autodiff +# end + +# function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff = :central) +# @assert alg in (:lm, :dogleg) +# @assert linsolve === nothing || linsolve in (:qr, :cholesky, :lsmr) +# autodiff isa Symbol && @assert autodiff in (:central, :forward) + +# if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing +# error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded") +# end + +# return LeastSquaresOptimJL{alg, linsolve}(autodiff) +# end + +# """ +# FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6, +# factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, +# minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, +# autodiff = nothing) + +# Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl) +# for solving `NonlinearLeastSquaresProblem`. For details about the other keyword arguments +# see the documentation for `FastLevenbergMarquardt.jl`. + +# !!! warning + +# This is not really the fastest solver. It is called that since the original package +# is called "Fast". `LevenbergMarquardt()` is almost always a better choice. + +# ### Arguments + +# - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`. + +# ### 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 +# `nothing` which means that a default is selected according to the problem specification! + +# !!! note + +# This algorithm is only available if `FastLevenbergMarquardt.jl` is installed and loaded. +# """ +# @concrete struct FastLevenbergMarquardtJL{linsolve} <: +# AbstractNonlinearSolveExtensionAlgorithm +# autodiff +# factor +# factoraccept +# factorreject +# factorupdate::Symbol +# minscale +# maxscale +# minfactor +# maxfactor +# end + +# function FastLevenbergMarquardtJL( +# linsolve::Symbol = :cholesky; factor = 1e-6, factoraccept = 13.0, +# factorreject = 3.0, factorupdate = :marquardt, minscale = 1e-12, +# maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, autodiff = nothing) +# @assert linsolve in (:qr, :cholesky) +# @assert factorupdate in (:marquardt, :nielson) + +# if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing +# error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded") +# end + +# return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject, +# factorupdate, minscale, maxscale, minfactor, maxfactor) +# end + +# """ +# CMINPACK(; method::Symbol = :auto, autodiff = missing) + +# ### Keyword Arguments + +# - `method`: the choice of method for the solver. +# - `autodiff`: Defaults to `missing`, which means we will default to letting `MINPACK` +# construct the jacobian if `f.jac` is not provided. In other cases, we use it to generate +# a jacobian similar to other NonlinearSolve solvers. + +# ### Submethod Choice + +# The keyword argument `method` can take on different value depending on which method of +# `fsolve` you are calling. The standard choices of `method` are: + +# - `:hybr`: Modified version of Powell's algorithm. Uses MINPACK routine +# [`hybrd1`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/hybrd1.c) +# - `:lm`: Levenberg-Marquardt. Uses MINPACK routine +# [`lmdif1`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/lmdif1.c) +# - `:lmdif`: Advanced Levenberg-Marquardt (more options available with `; kwargs...`). See +# MINPACK routine [`lmdif`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/lmdif.c) +# for more information +# - `:hybrd`: Advanced modified version of Powell's algorithm (more options available with +# `; kwargs...`). See MINPACK routine +# [`hybrd`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/hybrd.c) +# for more information + +# If a Jacobian is supplied as part of the [`NonlinearFunction`](@ref nonlinearfunctions), +# then the following methods are allowed: + +# - `:hybr`: Advanced modified version of Powell's algorithm with user supplied Jacobian. +# Additional arguments are available via `; kwargs...`. See MINPACK routine +# [`hybrj`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/hybrj.c) +# for more information +# - `:lm`: Advanced Levenberg-Marquardt with user supplied Jacobian. Additional arguments +# are available via `; kwargs...`. See MINPACK routine +# [`lmder`](https://github.com/devernay/cminpack/blob/d1f5f5a273862ca1bbcf58394e4ac060d9e22c76/lmder.c) +# for more information + +# The default choice of `:auto` selects `:hybr` for NonlinearProblem and `:lm` for +# NonlinearLeastSquaresProblem. + +# !!! note + +# This algorithm is only available if `MINPACK.jl` is installed and loaded. +# """ +# @concrete struct CMINPACK <: AbstractNonlinearSolveExtensionAlgorithm +# method::Symbol +# autodiff +# end + +# function CMINPACK(; method::Symbol = :auto, autodiff = missing) +# if Base.get_extension(@__MODULE__, :NonlinearSolveMINPACKExt) === nothing +# error("CMINPACK requires MINPACK.jl to be loaded") +# end +# return CMINPACK(method, autodiff) +# end + +# """ +# NLsolveJL(; method = :trust_region, autodiff = :central, linesearch = Static(), +# linsolve = (x, A, b) -> copyto!(x, A\\b), factor = one(Float64), autoscale = true, +# m = 10, beta = one(Float64)) + +# ### Keyword Arguments + +# - `method`: the choice of method for solving the nonlinear system. +# - `autodiff`: the choice of method for generating the Jacobian. Defaults to `:central` or +# central differencing via FiniteDiff.jl. The other choices are `:forward` or `ADTypes` +# similar to other solvers in NonlinearSolve. +# - `linesearch`: the line search method to be used within the solver method. The choices +# are line search types from +# [LineSearches.jl](https://github.com/JuliaNLSolvers/LineSearches.jl). +# - `linsolve`: a function `linsolve(x, A, b)` that solves `Ax = b`. +# - `factor`: determines the size of the initial trust region. This size is set to the +# product of factor and the euclidean norm of `u0` if nonzero, or else to factor itself. +# - `autoscale`: if true, then the variables will be automatically rescaled. The scaling +# factors are the norms of the Jacobian columns. +# - `m`: the amount of history in the Anderson method. Naive "Picard"-style iteration can be +# achieved by setting m=0, but that isn't advisable for contractions whose Lipschitz +# constants are close to 1. If convergence fails, though, you may consider lowering it. +# - `beta`: It is also known as DIIS or Pulay mixing, this method is based on the +# acceleration of the fixed-point iteration xₙ₊₁ = xₙ + beta*f(xₙ), where by default +# beta = 1. + +# !!! warning + +# Line Search Algorithms from [`LineSearch.jl`](https://github.com/SciML/LineSearch.jl) +# aren't supported by `NLsolveJL`. Instead, use the line search algorithms from +# [`LineSearches.jl`](https://github.com/JuliaNLSolvers/LineSearches.jl). + +# ### Submethod Choice + +# Choices for methods in `NLsolveJL`: + +# - `:anderson`: Anderson-accelerated fixed-point iteration +# - `:broyden`: Broyden's quasi-Newton method +# - `:newton`: Classical Newton method with an optional line search +# - `:trust_region`: Trust region Newton method (the default choice) + +# For more information on these arguments, consult the +# [NLsolve.jl documentation](https://github.com/JuliaNLSolvers/NLsolve.jl). + +# !!! note + +# This algorithm is only available if `NLsolve.jl` is installed and loaded. +# """ +# @concrete struct NLsolveJL <: AbstractNonlinearSolveExtensionAlgorithm +# method::Symbol +# autodiff +# linesearch +# linsolve +# factor +# autoscale::Bool +# m::Int +# beta +# end + +# function NLsolveJL(; method = :trust_region, autodiff = :central, linesearch = missing, +# linsolve = (x, A, b) -> copyto!(x, A \ b), factor = 1.0, +# autoscale = true, m = 10, beta = one(Float64)) +# if Base.get_extension(@__MODULE__, :NonlinearSolveNLsolveExt) === nothing +# error("NLsolveJL requires NLsolve.jl to be loaded") +# end + +# if autodiff isa Symbol && autodiff !== :central && autodiff !== :forward +# error("`autodiff` must be `:central` or `:forward`.") +# end + +# return NLsolveJL(method, autodiff, linesearch, linsolve, factor, autoscale, m, beta) +# end + +# """ +# NLSolversJL(method; autodiff = nothing) +# NLSolversJL(; method, autodiff = nothing) + +# Wrapper over NLSolvers.jl Nonlinear Equation Solvers. We automatically construct the +# jacobian function and supply it to the solver. + +# ### Arguments + +# - `method`: the choice of method for solving the nonlinear system. See the documentation +# for NLSolvers.jl for more information. +# - `autodiff`: the choice of method for generating the Jacobian. Defaults to `nothing` +# which means that a default is selected according to the problem specification. Can be +# any valid ADTypes.jl autodiff type (conditional on that backend being supported in +# NonlinearSolve.jl). +# """ +# struct NLSolversJL{M, AD} <: AbstractNonlinearSolveExtensionAlgorithm +# method::M +# autodiff::AD + +# function NLSolversJL(method, autodiff) +# if Base.get_extension(@__MODULE__, :NonlinearSolveNLSolversExt) === nothing +# error("NLSolversJL requires NLSolvers.jl to be loaded") +# end + +# return new{typeof(method), typeof(autodiff)}(method, autodiff) +# end +# end + +# NLSolversJL(method; autodiff = nothing) = NLSolversJL(method, autodiff) +# NLSolversJL(; method, autodiff = nothing) = NLSolversJL(method, autodiff) + +# """ +# SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, +# orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) + +# Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl/) for solving +# Fixed Point Problems. We allow using this algorithm to solve root finding problems as well. + +# ### Keyword Arguments + +# - `σ_min`: Setting to `1` may avoid stalling (see [lepage2021alternating](@cite)). +# - `stabilize`: performs a stabilization mapping before extrapolating. Setting to `true` +# may improve the performance for applications like accelerating the EM or MM algorithms +# (see [lepage2021alternating](@cite)). +# - `check_obj`: In case of NaN or Inf values, the algorithm restarts at the best past +# iterate. +# - `orders`: determines ACX's alternating order. Must be between `1` and `3` (where `1` +# means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the +# latter being potentially better for highly non-linear applications (see +# [lepage2021alternating](@cite)). +# - `time_limit`: time limit for the algorithm. + +# !!! note + +# This algorithm is only available if `SpeedMapping.jl` is installed and loaded. +# """ +# @concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm +# σ_min +# stabilize::Bool +# check_obj::Bool +# orders::Vector{Int} +# end + +# function SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, +# orders::Vector{Int} = [3, 3, 2]) +# if Base.get_extension(@__MODULE__, :NonlinearSolveSpeedMappingExt) === nothing +# error("SpeedMappingJL requires SpeedMapping.jl to be loaded") +# end + +# return SpeedMappingJL(σ_min, stabilize, check_obj, orders) +# end + +# """ +# FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, +# condition_number_threshold = missing, extrapolation_period = missing, +# replace_invalids = :NoAction) + +# Wrapper over [FixedPointAcceleration.jl](https://s-baumann.github.io/FixedPointAcceleration.jl/) +# for solving Fixed Point Problems. We allow using this algorithm to solve root finding +# problems as well. + +# ### Keyword Arguments + +# - `algorithm`: The algorithm to use. Can be `:Anderson`, `:MPE`, `:RRE`, `:VEA`, `:SEA`, +# `:Simple`, `:Aitken` or `:Newton`. +# - `m`: The number of previous iterates to use for the extrapolation. Only valid for +# `:Anderson`. +# - `condition_number_threshold`: The condition number threshold for Least Squares Problem. +# Only valid for `:Anderson`. +# - `extrapolation_period`: The number of iterates between extrapolations. Only valid for +# `:MPE`, `:RRE`, `:VEA` and `:SEA`. Defaults to `7` for `:MPE` & `:RRE`, and `6` for +# `:SEA` and `:VEA`. For `:SEA` and `:VEA`, this must be a multiple of `2`. +# - `replace_invalids`: The method to use for replacing invalid iterates. Can be +# `:ReplaceInvalids`, `:ReplaceVector` or `:NoAction`. + +# !!! note + +# This algorithm is only available if `FixedPointAcceleration.jl` is installed and loaded. +# """ +# @concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm +# algorithm::Symbol +# extrapolation_period::Int +# replace_invalids::Symbol +# dampening +# m::Int +# condition_number_threshold +# end + +# function FixedPointAccelerationJL(; +# algorithm = :Anderson, m = missing, condition_number_threshold = missing, +# extrapolation_period = missing, replace_invalids = :NoAction, dampening = 1.0) +# if Base.get_extension(@__MODULE__, :NonlinearSolveFixedPointAccelerationExt) === nothing +# error("FixedPointAccelerationJL requires FixedPointAcceleration.jl to be loaded") +# end + +# @assert algorithm in (:Anderson, :MPE, :RRE, :VEA, :SEA, :Simple, :Aitken, :Newton) +# @assert replace_invalids in (:ReplaceInvalids, :ReplaceVector, :NoAction) + +# if algorithm !== :Anderson +# if condition_number_threshold !== missing +# error("`condition_number_threshold` is only valid for Anderson acceleration") +# end +# if m !== missing +# error("`m` is only valid for Anderson acceleration") +# end +# end +# condition_number_threshold === missing && (condition_number_threshold = 1e3) +# m === missing && (m = 10) + +# if algorithm !== :MPE && algorithm !== :RRE && algorithm !== :VEA && algorithm !== :SEA +# if extrapolation_period !== missing +# error("`extrapolation_period` is only valid for MPE, RRE, VEA and SEA") +# end +# end +# if extrapolation_period === missing +# extrapolation_period = algorithm === :SEA || algorithm === :VEA ? 6 : 7 +# else +# if (algorithm === :SEA || algorithm === :VEA) && extrapolation_period % 2 != 0 +# error("`extrapolation_period` must be multiples of 2 for SEA and VEA") +# end +# end + +# return FixedPointAccelerationJL(algorithm, extrapolation_period, replace_invalids, +# dampening, m, condition_number_threshold) +# end + +# """ +# SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, +# autodiff = missing) + +# ### Keyword Arguments + +# - `method`: the choice of method for solving the nonlinear system. +# - `delta`: initial pseudo time step, default is 1e-3. +# - `linsolve` : JFNK linear solvers, choices are `gmres` and `bicgstab`. +# - `m`: Depth for Anderson acceleration, default as 0 for Picard iteration. +# - `beta`: Anderson mixing parameter, change f(x) to (1-beta)x+beta*f(x), +# equivalent to accelerating damped Picard iteration. +# - `autodiff`: Defaults to `missing`, which means we will default to letting +# `SIAMFANLEquations` construct the jacobian if `f.jac` is not provided. In other cases, +# we use it to generate a jacobian similar to other NonlinearSolve solvers. + +# ### Submethod Choice + +# - `:newton`: Classical Newton method. +# - `:pseudotransient`: Pseudo transient method. +# - `:secant`: Secant method for scalar equations. +# - `:anderson`: Anderson acceleration for fixed point iterations. + +# !!! note + +# This algorithm is only available if `SIAMFANLEquations.jl` is installed and loaded. +# """ +# @concrete struct SIAMFANLEquationsJL <: AbstractNonlinearSolveExtensionAlgorithm +# method::Symbol +# delta +# linsolve <: Union{Symbol, Nothing} +# m::Int +# beta +# autodiff +# end + +# function SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, +# m = 0, beta = 1.0, autodiff = missing) +# if Base.get_extension(@__MODULE__, :NonlinearSolveSIAMFANLEquationsExt) === nothing +# error("SIAMFANLEquationsJL requires SIAMFANLEquations.jl to be loaded") +# end +# return SIAMFANLEquationsJL(method, delta, linsolve, m, beta, autodiff) +# end + +# """ +# PETScSNES(; petsclib = missing, autodiff = nothing, mpi_comm = missing, kwargs...) + +# Wrapper over [PETSc.jl](https://github.com/JuliaParallel/PETSc.jl) SNES solvers. + +# ### Keyword Arguments + +# - `petsclib`: PETSc library to use. If set to `missing`, then we will use the first +# available PETSc library in `PETSc.petsclibs` based on the problem element type. +# - `autodiff`: the choice of method for generating the Jacobian. Defaults to `nothing` +# which means that a default is selected according to the problem specification. Can be +# any valid ADTypes.jl autodiff type (conditional on that backend being supported in +# NonlinearSolve.jl). If set to `missing`, then PETSc computes the Jacobian using finite +# differences. +# - `mpi_comm`: MPI communicator to use. If set to `missing`, then we will use +# `MPI.COMM_SELF`. +# - `kwargs`: Keyword arguments to be passed to the PETSc SNES solver. See [PETSc SNES +# documentation](https://petsc.org/release/manual/snes/) and +# [SNESSetFromOptions](https://petsc.org/release/manualpages/SNES/SNESSetFromOptions) +# for more information. + +# ### Options via `CommonSolve.solve` + +# These options are forwarded from `solve` to the PETSc SNES solver. If these are provided to +# `kwargs`, then they will be ignored. + +# | `solve` option | PETSc SNES option | +# |:-------------- |:----------------- | +# | `atol` | `snes_atol` | +# | `rtol` | `snes_rtol` | +# | `maxiters` | `snes_max_it` | +# | `show_trace` | `snes_monitor` | + +# !!! note + +# This algorithm is only available if `PETSc.jl` is installed and loaded. +# """ +# @concrete struct PETScSNES <: AbstractNonlinearSolveExtensionAlgorithm +# petsclib +# mpi_comm +# autodiff <: Union{Missing, Nothing, ADTypes.AbstractADType} +# snes_options +# end + +# function PETScSNES(; petsclib = missing, autodiff = nothing, mpi_comm = missing, kwargs...) +# if Base.get_extension(@__MODULE__, :NonlinearSolvePETScExt) === nothing +# error("PETScSNES requires PETSc.jl to be loaded") +# end +# return PETScSNES(petsclib, mpi_comm, autodiff, kwargs) +# end diff --git a/src/internal/forward_diff.jl b/src/forward_diff.jl similarity index 100% rename from src/internal/forward_diff.jl rename to src/forward_diff.jl diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl deleted file mode 100644 index 813068e7e..000000000 --- a/src/globalization/trust_region.jl +++ /dev/null @@ -1,452 +0,0 @@ - -# Don't Pollute the namespace -""" - RadiusUpdateSchemes - -`RadiusUpdateSchemes` is provides different types of radius update schemes implemented in -the Trust Region method. These schemes specify how the radius of the so-called trust region -is updated after each iteration of the algorithm. The specific role and caveats associated -with each scheme are provided below. - -## Using `RadiusUpdateSchemes` - -Simply put the desired scheme as follows: -`sol = solve(prob, alg = TrustRegion(radius_update_scheme = RadiusUpdateSchemes.Hei))`. -""" -module RadiusUpdateSchemes -# The weird definitions here are needed to main compatibility with the older enum variants - -abstract type AbstractRadiusUpdateScheme end - -function Base.show(io::IO, rus::AbstractRadiusUpdateScheme) - print(io, "RadiusUpdateSchemes.$(string(nameof(typeof(rus)))[3:end])") -end - -const T = AbstractRadiusUpdateScheme - -struct __Simple <: AbstractRadiusUpdateScheme end -""" - RadiusUpdateSchemes.Simple - -The simple or conventional radius update scheme. This scheme is chosen by default and -follows the conventional approach to update the trust region radius, i.e. if the trial -step is accepted it increases the radius by a fixed factor (bounded by a maximum radius) -and if the trial step is rejected, it shrinks the radius by a fixed factor. -""" -const Simple = __Simple() - -struct __NLsolve <: AbstractRadiusUpdateScheme end -""" - RadiusUpdateSchemes.NLsolve - -The same updating scheme as in NLsolve's (https://github.com/JuliaNLSolvers/NLsolve.jl) -trust region dogleg implementation. -""" -const NLsolve = __NLsolve() - -struct __NocedalWright <: AbstractRadiusUpdateScheme end -""" - RadiusUpdateSchemes.NocedalWright - -Trust region updating scheme as in Nocedal and Wright [see Alg 11.5, page 291]. -""" -const NocedalWright = __NocedalWright() - -struct __Hei <: AbstractRadiusUpdateScheme end -""" - RadiusUpdateSchemes.Hei - -This scheme is proposed in [hei2003self](@citet). The trust region radius depends on the -size (norm) of the current step size. The hypothesis is to let the radius converge to zero -as the iterations progress, which is more reliable and robust for ill-conditioned as well -as degenerate problems. -""" -const Hei = __Hei() - -struct __Yuan <: AbstractRadiusUpdateScheme end -""" - RadiusUpdateSchemes.Yuan - -This scheme is proposed by [yuan2015recent](@citet). Similar to Hei's scheme, the -trust region is updated in a way so that it converges to zero, however here, the radius -depends on the size (norm) of the current gradient of the objective (merit) function. The -hypothesis is that the step size is bounded by the gradient size, so it makes sense to let -the radius depend on the gradient. -""" -const Yuan = __Yuan() - -struct __Bastin <: AbstractRadiusUpdateScheme end -""" - RadiusUpdateSchemes.Bastin - -This scheme is proposed by [bastin2010retrospective](@citet). The scheme is called a -retrospective update scheme as it uses the model function at the current iteration to -compute the ratio of the actual reduction and the predicted reduction in the previous trial -step, and use this ratio to update the trust region radius. The hypothesis is to exploit the -information made available during the optimization process in order to vary the accuracy -of the objective function computation. -""" -const Bastin = __Bastin() - -struct __Fan <: AbstractRadiusUpdateScheme end -""" - RadiusUpdateSchemes.Fan - -This scheme is proposed by [fan2006convergence](@citet). It is very much similar to Hei's -and Yuan's schemes as it lets the trust region radius depend on the current size (norm) of -the objective (merit) function itself. These new update schemes are known to improve local -convergence. -""" -const Fan = __Fan() - -end - -const RUS = RadiusUpdateSchemes - -""" - GenericTrustRegionScheme(; method = RadiusUpdateSchemes.Simple, - max_trust_radius = nothing, initial_trust_radius = nothing, - step_threshold = nothing, shrink_threshold = nothing, expand_threshold = nothing, - shrink_factor = nothing, expand_factor = nothing) - -Trust Region Method that updates and stores the current trust region radius in -`trust_region`. For any of the keyword arguments, if the value is `nothing`, then we use -the value used in the respective paper. - -### Keyword Arguments - - - `radius_update_scheme`: the choice of radius update scheme to be used. Defaults to - `RadiusUpdateSchemes.Simple` which follows the conventional approach. Other available - schemes are documented in [`RadiusUpdateSchemes`](@ref),. These schemes have the trust - region radius converging to zero that is seen to improve convergence. For more details, - see [1]. - - `max_trust_radius`: the maximal trust region radius. Defaults to - `max(norm(fu), maximum(u) - minimum(u))`, except for `RadiusUpdateSchemes.NLsolve` - where it defaults to `Inf`. - - `initial_trust_radius`: the initial trust region radius. Defaults to - `max_trust_radius / 11`, except for `RadiusUpdateSchemes.NLsolve` where it defaults - to `u0_norm > 0 ? u0_norm : 1`. - - `step_threshold`: the threshold for taking a step. In every iteration, the threshold is - compared with a value `r`, which is the actual reduction in the objective function - divided by the predicted reduction. If `step_threshold > r` the model is not a good - approximation, and the step is rejected. Defaults to `nothing`. - - `shrink_threshold`: the threshold for shrinking the trust region radius. In every - iteration, the threshold is compared with a value `r` which is the actual reduction in - the objective function divided by the predicted reduction. If `shrink_threshold > r` the - trust region radius is shrunk by `shrink_factor`. Defaults to `nothing`. - - `expand_threshold`: the threshold for expanding the trust region radius. If a step is - taken, i.e `step_threshold < r` (with `r` defined in `shrink_threshold`), a check is - also made to see if `expand_threshold < r`. If that is true, the trust region radius is - expanded by `expand_factor`. Defaults to `nothing`. - - `shrink_factor`: the factor to shrink the trust region radius with if - `shrink_threshold > r` (with `r` defined in `shrink_threshold`). Defaults to `0.25`. - - `expand_factor`: the factor to expand the trust region radius with if - `expand_threshold < r` (with `r` defined in `shrink_threshold`). Defaults to `2.0`. -""" -@kwdef @concrete struct GenericTrustRegionScheme{M <: - RadiusUpdateSchemes.AbstractRadiusUpdateScheme} - method::M = RadiusUpdateSchemes.Simple - step_threshold = nothing - shrink_threshold = nothing - shrink_factor = nothing - expand_factor = nothing - expand_threshold = nothing - max_trust_radius = nothing - initial_trust_radius = nothing -end - -function Base.show(io::IO, alg::GenericTrustRegionScheme) - print(io, "GenericTrustRegionScheme(method = $(alg.method))") -end - -@concrete mutable struct GenericTrustRegionSchemeCache <: AbstractTrustRegionMethodCache - method - f - p - max_trust_radius - initial_trust_radius - trust_region - step_threshold - shrink_threshold - expand_threshold - shrink_factor - expand_factor - p1 - p2 - p3 - p4 - ϵ - ρ - vjp_operator - jvp_operator - Jᵀfu_cache - Jδu_cache - δu_cache - internalnorm - u_cache - fu_cache - last_step_accepted::Bool - shrink_counter::Int - stats::NLStats - alg -end - -function reinit_cache!( - cache::GenericTrustRegionSchemeCache, args...; u0 = nothing, p = cache.p, kwargs...) - T = eltype(cache.u_cache) - cache.p = p - if u0 !== nothing - u0_norm = cache.internalnorm(u0) - cache.trust_region = __initial_trust_radius( - cache.alg.initial_trust_radius, T, cache.alg.method, - cache.max_trust_radius, u0_norm, u0_norm) # FIXME: scheme specific - end - cache.last_step_accepted = false - cache.shrink_counter = 0 -end - -# Defaults -for func in (:__max_trust_radius, :__initial_trust_radius, :__step_threshold, - :__shrink_threshold, :__shrink_factor, :__expand_threshold, :__expand_factor) - @eval begin - @inline function $(func)(val, ::Type{T}, args...) where {T} - val_T = T(val) - iszero(val_T) && return $(func)(nothing, T, args...) - return val_T - end - end -end - -@inline __max_trust_radius(::Nothing, ::Type{T}, method, u, fu_norm) where {T} = T(Inf) -@inline function __max_trust_radius( - ::Nothing, ::Type{T}, ::Union{RUS.__Simple, RUS.__NocedalWright}, - u, fu_norm) where {T} - u_min, u_max = extrema(u) - return max(T(fu_norm), u_max - u_min) -end - -@inline function __initial_trust_radius( - ::Nothing, ::Type{T}, method, max_tr, u0_norm, fu_norm) where {T} - method isa RUS.__NLsolve && return T(ifelse(u0_norm > 0, u0_norm, 1)) - (method isa RUS.__Hei || method isa RUS.__Bastin) && return T(1) - method isa RUS.__Fan && return T((fu_norm^0.99) / 10) - return T(max_tr / 11) -end - -@inline function __step_threshold(::Nothing, ::Type{T}, method) where {T} - method isa RUS.__Hei && return T(0) - method isa RUS.__Yuan && return T(1 // 1000) - method isa RUS.__Bastin && return T(1 // 20) - return T(1 // 10000) -end - -@inline function __shrink_threshold(::Nothing, ::Type{T}, method) where {T} - method isa RUS.__Hei && return T(0) - (method isa RUS.__NLsolve || method isa RUS.__Bastin) && return T(1 // 20) - return T(1 // 4) -end - -@inline function __expand_threshold(::Nothing, ::Type{T}, method) where {T} - method isa RUS.__NLsolve && return T(9 // 10) - method isa RUS.__Hei && return T(0) - method isa RUS.__Bastin && return T(9 // 10) - return T(3 // 4) -end - -@inline function __shrink_factor(::Nothing, ::Type{T}, method) where {T} - method isa RUS.__NLsolve && return T(1 // 2) - method isa RUS.__Hei && return T(0) - method isa RUS.__Bastin && return T(1 // 20) - return T(1 // 4) -end - -@inline function __get_parameters(::Type{T}, method) where {T} - method isa RUS.__NLsolve && return (T(1 // 2), T(0), T(0), T(0)) - method isa RUS.__Hei && return (T(5), T(1 // 10), T(15 // 100), T(15 // 100)) - method isa RUS.__Yuan && return (T(2), T(1 // 6), T(6), T(0)) - method isa RUS.__Fan && return (T(1 // 10), T(1 // 4), T(12), T(1e18)) - method isa RUS.__Bastin && return (T(5 // 2), T(1 // 4), T(0), T(0)) - return (T(0), T(0), T(0), T(0)) -end - -@inline __expand_factor(::Nothing, ::Type{T}, method) where {T} = T(2) - -function __internal_init( - prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, f::F, fu, u, - p, args...; stats, internalnorm::IF = L2_NORM, vjp_autodiff = nothing, - jvp_autodiff = nothing, kwargs...) where {F, IF} - T = promote_type(eltype(u), eltype(fu)) - u0_norm = internalnorm(u) - fu_norm = internalnorm(fu) - - # Common Setup - max_trust_radius = __max_trust_radius(alg.max_trust_radius, T, alg.method, u, fu_norm) - initial_trust_radius = __initial_trust_radius( - alg.initial_trust_radius, T, alg.method, max_trust_radius, u0_norm, fu_norm) - step_threshold = __step_threshold(alg.step_threshold, T, alg.method) - shrink_threshold = __shrink_threshold(alg.shrink_threshold, T, alg.method) - expand_threshold = __expand_threshold(alg.expand_threshold, T, alg.method) - shrink_factor = __shrink_factor(alg.shrink_factor, T, alg.method) - expand_factor = __expand_factor(alg.expand_factor, T, alg.method) - - # Scheme Specific Setup - p1, p2, p3, p4 = __get_parameters(T, alg.method) - ϵ = T(1e-8) - - vjp_operator = alg.method isa RUS.__Yuan || alg.method isa RUS.__Bastin ? - VecJacOperator(prob, fu, u; autodiff = vjp_autodiff) : nothing - - jvp_operator = alg.method isa RUS.__Bastin ? - JacVecOperator(prob, fu, u; autodiff = jvp_autodiff) : nothing - - if alg.method isa RUS.__Yuan - Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * _vec(fu) - initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) - else - if u isa Number - Jᵀfu_cache = u - else - @bb Jᵀfu_cache = similar(u) - end - end - - if alg.method isa RUS.__Bastin - @bb δu_cache = similar(u) - else - δu_cache = nothing - end - - @bb u_cache = similar(u) - @bb fu_cache = similar(fu) - @bb Jδu_cache = similar(fu) - - return GenericTrustRegionSchemeCache( - alg.method, f, p, max_trust_radius, initial_trust_radius, initial_trust_radius, - step_threshold, shrink_threshold, expand_threshold, shrink_factor, - expand_factor, p1, p2, p3, p4, ϵ, T(0), vjp_operator, jvp_operator, Jᵀfu_cache, - Jδu_cache, δu_cache, internalnorm, u_cache, fu_cache, false, 0, stats, alg) -end - -function __internal_solve!( - cache::GenericTrustRegionSchemeCache, J, fu, u, δu, descent_stats) - T = promote_type(eltype(u), eltype(fu)) - @bb @. cache.u_cache = u + δu - cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) - cache.stats.nf += 1 - - if hasfield(typeof(descent_stats), :δuJᵀJδu) && !isnan(descent_stats.δuJᵀJδu) - δuJᵀJδu = descent_stats.δuJᵀJδu - else - @bb cache.Jδu_cache = J × vec(δu) - δuJᵀJδu = __dot(cache.Jδu_cache, cache.Jδu_cache) - end - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - num = (cache.internalnorm(cache.fu_cache)^2 - cache.internalnorm(fu)^2) / 2 - denom = __dot(δu, cache.Jᵀfu_cache) + δuJᵀJδu / 2 - cache.ρ = num / denom - - if cache.ρ > cache.step_threshold - cache.last_step_accepted = true - else - cache.last_step_accepted = false - end - - if cache.method isa RUS.__Simple - if cache.ρ < cache.shrink_threshold - cache.trust_region *= cache.shrink_factor - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ > cache.expand_threshold && cache.ρ > cache.step_threshold - cache.trust_region = cache.expand_factor * cache.trust_region - end - end - elseif cache.method isa RUS.__NLsolve - if cache.ρ < cache.shrink_threshold - cache.trust_region *= cache.shrink_factor - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ ≥ cache.expand_threshold - cache.trust_region = cache.expand_factor * cache.internalnorm(δu) - elseif cache.ρ ≥ cache.p1 - cache.trust_region = max( - cache.trust_region, cache.expand_factor * cache.internalnorm(δu)) - end - end - elseif cache.method isa RUS.__NocedalWright - if cache.ρ < cache.shrink_threshold - cache.trust_region = cache.shrink_factor * cache.internalnorm(δu) - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - if cache.ρ > cache.expand_threshold && - abs(cache.internalnorm(δu) - cache.trust_region) < 1e-6 * cache.trust_region - cache.trust_region = cache.expand_factor * cache.trust_region - end - end - elseif cache.method isa RUS.__Hei - tr_new = __rfunc( - cache.ρ, cache.shrink_threshold, cache.p1, cache.p3, cache.p4, cache.p2) * - cache.internalnorm(δu) - if tr_new < cache.trust_region - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - end - cache.trust_region = tr_new - elseif cache.method isa RUS.__Yuan - if cache.ρ < cache.shrink_threshold - cache.p1 = cache.p2 * cache.p1 - cache.shrink_counter += 1 - else - if cache.ρ ≥ cache.expand_threshold && - 2 * cache.internalnorm(δu) > cache.trust_region - cache.p1 = cache.p3 * cache.p1 - end - cache.shrink_counter = 0 - end - operator = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) - @bb cache.Jᵀfu_cache = operator × vec(cache.fu_cache) - cache.trust_region = cache.p1 * cache.internalnorm(cache.Jᵀfu_cache) - elseif cache.method isa RUS.__Fan - if cache.ρ < cache.shrink_threshold - cache.p1 *= cache.p2 - cache.shrink_counter += 1 - else - cache.shrink_counter = 0 - cache.ρ > cache.expand_threshold && - (cache.p1 = min(cache.p1 * cache.p3, cache.p4)) - end - cache.trust_region = cache.p1 * (cache.internalnorm(cache.fu_cache)^T(0.99)) - elseif cache.method isa RUS.__Bastin - if cache.ρ > cache.step_threshold - jvp_op = StatefulJacobianOperator(cache.jvp_operator, cache.u_cache, cache.p) - vjp_op = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) - @bb cache.Jδu_cache = jvp_op × vec(δu) - @bb cache.Jᵀfu_cache = vjp_op × vec(cache.fu_cache) - denom_1 = dot(_vec(δu), cache.Jᵀfu_cache) - @bb cache.Jᵀfu_cache = vjp_op × vec(cache.Jδu_cache) - denom_2 = dot(_vec(δu), cache.Jᵀfu_cache) - denom = denom_1 + denom_2 / 2 - ρ = num / denom - if ρ ≥ cache.expand_threshold - cache.trust_region = cache.p1 * cache.internalnorm(δu) - end - cache.shrink_counter = 0 - else - cache.trust_region *= cache.p2 - cache.shrink_counter += 1 - end - end - - cache.trust_region = min(cache.trust_region, cache.max_trust_radius) - - return cache.last_step_accepted, cache.u_cache, cache.fu_cache -end - -# R-function for adaptive trust region method -function __rfunc(r::R, c2::R, M::R, γ1::R, γ2::R, β::R) where {R <: Real} - return ifelse(r ≥ c2, (2 * (M - 1 - γ2) * atan(r - c2) + (1 + γ2)) / R(π), - (1 - γ1 - β) * (exp(r - c2) + β / (1 - γ1 - β))) -end diff --git a/src/timer_outputs.jl b/src/helpers.jl similarity index 100% rename from src/timer_outputs.jl rename to src/helpers.jl diff --git a/src/polyalg.jl b/src/polyalg.jl new file mode 100644 index 000000000..5cf21d6e1 --- /dev/null +++ b/src/polyalg.jl @@ -0,0 +1,543 @@ +""" + NonlinearSolvePolyAlgorithm(algs; start_index::Int = 1) + +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. + +### Arguments + + - `algs`: a tuple of algorithms to try in-order! (If this is not a Tuple, then the + returned algorithm is not type-stable). + +### Keyword Arguments + + - `start_index`: the index to start at. Defaults to `1`. + +### Example + +```julia +using NonlinearSolve + +alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden())) +``` +""" +@concrete struct NonlinearSolvePolyAlgorithm <: AbstractNonlinearSolveAlgorithm + static_length <: Val + algs <: Tuple + start_index::Int +end + +function NonlinearSolvePolyAlgorithm(algs; start_index::Int = 1) + @assert 0 < start_index ≤ length(algs) + algs = Tuple(algs) + return NonlinearSolvePolyAlgorithm(Val(length(algs)), algs, start_index) +end + +@concrete mutable struct NonlinearSolvePolyAlgorithmCache <: AbstractNonlinearSolveCache + static_length <: Val + prob <: AbstractNonlinearProblem + + caches <: Tuple + alg <: NonlinearSolvePolyAlgorithm + + best::Int + current::Int + nsteps::Int + + stats::NLStats + total_time::Float64 + maxtime + + retcode::ReturnCode.T + force_stop::Bool + + maxiters::Int + internalnorm + + u0 + u0_aliased + alias_u0::Bool +end + +function SII.symbolic_container(cache::NonlinearSolvePolyAlgorithmCache) + return cache.caches[cache.current] +end +SII.state_values(cache::NonlinearSolvePolyAlgorithmCache) = cache.u0 + +function Base.show(io::IO, ::MIME"text/plain", cache::NonlinearSolvePolyAlgorithmCache) + println(io, "NonlinearSolvePolyAlgorithmCache with \ + $(Utils.unwrap_val(cache.static_length)) algorithms:") + best_alg = ifelse(cache.best == -1, "nothing", cache.best) + println(io, " Best Algorithm: $(best_alg)") + println( + io, " Current Algorithm: [$(cache.current) / $(Utils.unwrap_val(cache.static_length))]" + ) + println(io, " nsteps: $(cache.nsteps)") + println(io, " retcode: $(cache.retcode)") + print(io, " Current Cache: ") + NonlinearSolveBase.show_nonlinearsolve_cache(io, cache.caches[cache.current], 4) +end + +function InternalAPI.reinit!( + cache::NonlinearSolvePolyAlgorithmCache, args...; p = cache.p, u0 = cache.u0 +) + foreach(cache.caches) do cache + InternalAPI.reinit!(cache, args...; p, u0) + end + cache.current = cache.alg.start_index + InternalAPI.reinit!(cache.stats) + cache.nsteps = 0 + cache.total_time = 0.0 +end + +function SciMLBase.__init( + prob::AbstractNonlinearProblem, alg::NonlinearSolvePolyAlgorithm, args...; + stats = NLStats(0, 0, 0, 0, 0), maxtime = nothing, maxiters = 1000, + internalnorm = L2_NORM, alias_u0 = false, verbose = true, kwargs... +) + if alias_u0 && !ArrayInterface.ismutable(prob.u0) + verbose && @warn "`alias_u0` has been set to `true`, but `u0` is \ + immutable (checked using `ArrayInterface.ismutable`)." + alias_u0 = false # If immutable don't care about aliasing + end + + u0 = prob.u0 + u0_aliased = alias_u0 ? copy(u0) : u0 + alias_u0 && (prob = SciMLBase.remake(prob; u0 = u0_aliased)) + + return NonlinearSolvePolyAlgorithmCache( + alg.static_length, prob, + map(alg.algs) do solver + SciMLBase.__init( + prob, solver, args...; + stats, maxtime, internalnorm, alias_u0, verbose, kwargs... + ) + end, + alg, -1, alg.start_index, 0, stats, 0.0, maxtime, + ReturnCode.Default, false, maxiters, internalnorm, + u0, u0_aliased, alias_u0 + ) +end + +@generated function CommonSolve.solve!(cache::NonlinearSolvePolyAlgorithmCache{Val{N}}) where {N} + calls = [quote + 1 ≤ cache.current ≤ $(N) || error("Current choices shouldn't get here!") + end] + + cache_syms = [gensym("cache") for i in 1:N] + sol_syms = [gensym("sol") for i in 1:N] + u_result_syms = [gensym("u_result") for i in 1:N] + + for i in 1:N + push!(calls, + quote + $(cache_syms[i]) = cache.caches[$(i)] + if $(i) == cache.current + cache.alias_u0 && copyto!(cache.u0_aliased, cache.u0) + $(sol_syms[i]) = CommonSolve.solve!($(cache_syms[i])) + if SciMLBase.successful_retcode($(sol_syms[i])) + stats = $(sol_syms[i]).stats + if cache.alias_u0 + copyto!(cache.u0, $(sol_syms[i]).u) + $(u_result_syms[i]) = cache.u0 + else + $(u_result_syms[i]) = $(sol_syms[i]).u + end + fu = NonlinearSolveBase.get_fu($(cache_syms[i])) + return build_solution_less_specialize( + cache.prob, cache.alg, $(u_result_syms[i]), fu; + retcode = $(sol_syms[i]).retcode, stats, + original = $(sol_syms[i]), trace = $(sol_syms[i]).trace + ) + elseif cache.alias_u0 + # For safety we need to maintain a copy of the solution + $(u_result_syms[i]) = copy($(sol_syms[i]).u) + end + cache.current = $(i + 1) + end + end) + end + + resids = map(Base.Fix2(Symbol, :resid), cache_syms) + for (sym, resid) in zip(cache_syms, resids) + push!(calls, :($(resid) = @isdefined($(sym)) ? $(sym).resid : nothing)) + end + push!(calls, quote + fus = tuple($(Tuple(resids)...)) + minfu, idx = findmin_caches(cache.prob, fus) + end) + for i in 1:N + push!(calls, + quote + if idx == $(i) + u = cache.alias_u0 ? $(u_result_syms[i]) : + NonlinearSolveBase.get_u(cache.caches[$(i)]) + end + end) + end + push!(calls, + quote + retcode = cache.caches[idx].retcode + if cache.alias_u0 + copyto!(cache.u0, u) + u = cache.u0 + end + return build_solution_less_specialize( + cache.prob, cache.alg, u, fus[idx]; + retcode, cache.stats, cache.caches[idx].trace + ) + end) + + return Expr(:block, calls...) +end + +@generated function InternalAPI.step!( + cache::NonlinearSolvePolyAlgorithmCache{Val{N}}, args...; kwargs... +) where {N} + calls = [] + cache_syms = [gensym("cache") for i in 1:N] + for i in 1:N + push!(calls, + quote + $(cache_syms[i]) = cache.caches[$(i)] + if $(i) == cache.current + InternalAPI.step!($(cache_syms[i]), args...; kwargs...) + $(cache_syms[i]).nsteps += 1 + if !NonlinearSolveBase.not_terminated($(cache_syms[i])) + if SciMLBase.successful_retcode($(cache_syms[i]).retcode) + cache.best = $(i) + cache.force_stop = true + cache.retcode = $(cache_syms[i]).retcode + else + cache.current = $(i + 1) + end + end + return + end + end) + end + + push!(calls, quote + if !(1 ≤ cache.current ≤ length(cache.caches)) + minfu, idx = findmin_caches(cache.prob, cache.caches) + cache.best = idx + cache.retcode = cache.caches[idx].retcode + cache.force_stop = true + return + end + end) + + return Expr(:block, calls...) +end + +@generated function SciMLBase.__solve( + prob::AbstractNonlinearProblem, alg::NonlinearSolvePolyAlgorithm{Val{N}}, args...; + stats = NLStats(0, 0, 0, 0, 0), alias_u0 = false, verbose = true, kwargs... +) where {N} + sol_syms = [gensym("sol") for _ in 1:N] + prob_syms = [gensym("prob") for _ in 1:N] + u_result_syms = [gensym("u_result") for _ in 1:N] + calls = [quote + current = alg.start_index + if alias_u0 && !ArrayInterface.ismutable(prob.u0) + verbose && @warn "`alias_u0` has been set to `true`, but `u0` is \ + immutable (checked using `ArrayInterface.ismutable`)." + alias_u0 = false # If immutable don't care about aliasing + end + u0 = prob.u0 + u0_aliased = alias_u0 ? zero(u0) : u0 + end] + for i in 1:N + cur_sol = sol_syms[i] + push!(calls, + quote + if current == $(i) + if alias_u0 + copyto!(u0_aliased, u0) + $(prob_syms[i]) = SciMLBase.remake(prob; u0 = u0_aliased) + else + $(prob_syms[i]) = prob + end + $(cur_sol) = SciMLBase.__solve( + $(prob_syms[i]), alg.algs[$(i)], args...; + stats, alias_u0, verbose, kwargs... + ) + if SciMLBase.successful_retcode($(cur_sol)) + if alias_u0 + copyto!(u0, $(cur_sol).u) + $(u_result_syms[i]) = u0 + else + $(u_result_syms[i]) = $(cur_sol).u + end + return build_solution_less_specialize( + prob, alg, $(u_result_syms[i]), $(cur_sol).resid; + $(cur_sol).retcode, $(cur_sol).stats, + $(cur_sol).trace, original = $(cur_sol) + ) + elseif alias_u0 + # For safety we need to maintain a copy of the solution + $(u_result_syms[i]) = copy($(cur_sol).u) + end + current = $(i + 1) + end + end) + end + + resids = map(Base.Fix2(Symbol, :resid), sol_syms) + for (sym, resid) in zip(sol_syms, resids) + push!(calls, :($(resid) = @isdefined($(sym)) ? $(sym).resid : nothing)) + end + + push!(calls, quote + resids = tuple($(Tuple(resids)...)) + minfu, idx = findmin_resids(prob, resids) + end) + + for i in 1:N + push!(calls, + quote + if idx == $(i) + if alias_u0 + copyto!(u0, $(u_result_syms[i])) + $(u_result_syms[i]) = u0 + else + $(u_result_syms[i]) = $(sol_syms[i]).u + end + return build_solution_less_specialize( + prob, alg, $(u_result_syms[i]), $(sol_syms[i]).resid; + $(sol_syms[i]).retcode, $(sol_syms[i]).stats, + $(sol_syms[i]).trace, original = $(sol_syms[i]) + ) + end + end) + end + push!(calls, :(error("Current choices shouldn't get here!"))) + + return Expr(:block, calls...) +end + +""" + RobustMultiNewton( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing + ) + +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). + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. +""" +function RobustMultiNewton( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) where {T} + common_kwargs = (; concrete_jac, linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + if T <: Complex # Let's atleast have something here for complex numbers + algs = (NewtonRaphson(; common_kwargs...),) + else + algs = ( + TrustRegion(; common_kwargs...), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Bastin), + NewtonRaphson(; common_kwargs...), + NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.NLsolve), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Fan) + ) + end + return NonlinearSolvePolyAlgorithm(algs) +end + +""" + FastShortcutNonlinearPolyalg( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + must_use_jacobian::Val = Val(false), + prefer_simplenonlinearsolve::Val = Val(false), + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, + u0_len::Union{Int, Nothing} = nothing + ) where {T} + +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. + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. + +### Keyword Arguments + + - `u0_len`: The length of the initial guess. If this is `nothing`, then the length of the + initial guess is not checked. If this is an integer and it is less than `25`, we use + jacobian based methods. +""" +function FastShortcutNonlinearPolyalg( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + must_use_jacobian::Val = Val(false), + prefer_simplenonlinearsolve::Val = Val(false), + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, + u0_len::Union{Int, Nothing} = nothing +) where {T} + start_index = 1 + common_kwargs = (; concrete_jac, linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + if must_use_jacobian isa Val{true} + if T <: Complex + algs = (NewtonRaphson(; common_kwargs...),) + else + algs = ( + NewtonRaphson(; common_kwargs...), + NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), + TrustRegion(; common_kwargs...), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Bastin) + ) + end + else + # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians + # and thus are not included in the polyalgorithm + if prefer_simplenonlinearsolve isa Val{true} + if T <: Complex + algs = ( + SimpleBroyden(), + Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + SimpleKlement(), + NewtonRaphson(; common_kwargs...) + ) + else + start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1 + algs = ( + SimpleBroyden(), + Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + SimpleKlement(), + NewtonRaphson(; common_kwargs...), + NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), + TrustRegion(; common_kwargs...), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Bastin) + ) + end + else + if T <: Complex + algs = ( + Broyden(; autodiff), + Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + Klement(; linsolve, precs, autodiff), + NewtonRaphson(; common_kwargs...) + ) + else + # TODO: This number requires a bit rigorous testing + start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1 + algs = ( + Broyden(; autodiff), + Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + Klement(; linsolve, precs, autodiff), + NewtonRaphson(; common_kwargs...), + NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), + TrustRegion(; common_kwargs...), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Bastin) + ) + end + end + end + return NonlinearSolvePolyAlgorithm(algs; start_index) +end + +""" + FastShortcutNLLSPolyalg( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing + ) + +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. + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. +""" +function FastShortcutNLLSPolyalg( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) where {T} + common_kwargs = (; linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + if T <: Complex + algs = ( + GaussNewton(; common_kwargs..., concrete_jac), + LevenbergMarquardt(; common_kwargs..., disable_geodesic = Val(true)), + LevenbergMarquardt(; common_kwargs...) + ) + else + algs = ( + GaussNewton(; common_kwargs..., concrete_jac), + LevenbergMarquardt(; common_kwargs..., disable_geodesic = Val(true)), + TrustRegion(; common_kwargs..., concrete_jac), + GaussNewton(; common_kwargs..., linesearch = BackTracking(), concrete_jac), + TrustRegion(; + common_kwargs..., radius_update_scheme = RUS.Bastin, concrete_jac + ), + LevenbergMarquardt(; common_kwargs...) + ) + end + return NonlinearSolvePolyAlgorithm(algs) +end + +# Original is often determined on runtime information especially for PolyAlgorithms so it +# is best to never specialize on that +function build_solution_less_specialize( + prob::AbstractNonlinearProblem, alg, u, resid; + retcode = ReturnCode.Default, original = nothing, left = nothing, + right = nothing, stats = nothing, trace = nothing, kwargs... +) + return SciMLBase.NonlinearSolution{ + eltype(eltype(u)), ndims(u), typeof(u), typeof(resid), typeof(prob), + typeof(alg), Any, typeof(left), typeof(stats), typeof(trace) + }( + u, resid, prob, alg, retcode, original, left, right, stats, trace + ) +end + +function findmin_caches(prob::AbstractNonlinearProblem, caches) + resids = map(caches) do cache + cache === nothing && return nothing + return NonlinearSolveBase.get_fu(cache) + end + return findmin_resids(prob, resids) +end + +@views function findmin_resids(prob::AbstractNonlinearProblem, caches) + norm_fn = prob isa NonlinearLeastSquaresProblem ? Base.Fix2(norm, 2) : + Base.Fix2(norm, Inf) + idx = findfirst(Base.Fix2(!==, nothing), caches) + # This is an internal function so we assume that inputs are consistent and there is + # atleast one non-`nothing` value + fx_idx = norm_fn(caches[idx]) + idx == length(caches) && return fx_idx, idx + fmin = @closure xᵢ -> begin + xᵢ === nothing && return oftype(fx_idx, Inf) + fx = norm_fn(xᵢ) + return ifelse(isnan(fx), oftype(fx, Inf), fx) + end + x_min, x_min_idx = findmin(fmin, caches[(idx + 1):length(caches)]) + x_min < fx_idx && return x_min, x_min_idx + idx + return fx_idx, idx +end diff --git a/src/utils.jl b/src/utils.jl deleted file mode 100644 index 0a2375ef9..000000000 --- a/src/utils.jl +++ /dev/null @@ -1,50 +0,0 @@ -@inline __is_complex(::Type{ComplexF64}) = true -@inline __is_complex(::Type{ComplexF32}) = true -@inline __is_complex(::Type{Complex}) = true -@inline __is_complex(::Type{T}) where {T} = false - -@inline __findmin_caches(f::F, caches) where {F} = __findmin(f ∘ get_fu, caches) -# FIXME: L2_NORM makes an Array of NaNs not a NaN (atleast according to `isnan`) -@generated function __findmin(f::F, x) where {F} - # JET shows dynamic dispatch if this is not written as a generated function - F === typeof(L2_NORM) && return :(return __findmin_impl(Base.Fix1(maximum, abs), x)) - return :(return __findmin_impl(f, x)) -end -@inline @views function __findmin_impl(f::F, x) where {F} - idx = findfirst(Base.Fix2(!==, nothing), x) - # This is an internal function so we assume that inputs are consistent and there is - # atleast one non-`nothing` value - fx_idx = f(x[idx]) - idx == length(x) && return fx_idx, idx - fmin = @closure xᵢ -> begin - xᵢ === nothing && return oftype(fx_idx, Inf) - fx = f(xᵢ) - return ifelse(isnan(fx), oftype(fx, Inf), fx) - end - x_min, x_min_idx = findmin(fmin, x[(idx + 1):length(x)]) - x_min < fx_idx && return x_min, x_min_idx + idx - return fx_idx, idx -end - -""" - pickchunksize(x) = pickchunksize(length(x)) - pickchunksize(x::Int) - -Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the input length. -""" -@inline pickchunksize(x) = pickchunksize(length(x)) -@inline pickchunksize(x::Int) = ForwardDiff.pickchunksize(x) - -# Original is often determined on runtime information especially for PolyAlgorithms so it -# is best to never specialize on that -function __build_solution_less_specialize(prob::AbstractNonlinearProblem, alg, u, resid; - retcode = ReturnCode.Default, original = nothing, left = nothing, - right = nothing, stats = nothing, trace = nothing, kwargs...) - T = eltype(eltype(u)) - N = ndims(u) - - return SciMLBase.NonlinearSolution{ - T, N, typeof(u), typeof(resid), typeof(prob), typeof(alg), - Any, typeof(left), typeof(stats), typeof(trace)}( - u, resid, prob, alg, retcode, original, left, right, stats, trace) -end