diff --git a/Project.toml b/Project.toml index 93643123b..3b0172faa 100644 --- a/Project.toml +++ b/Project.toml @@ -89,7 +89,7 @@ NLSolvers = "0.5" NLsolve = "4.5" NaNMath = "1" NonlinearProblemLibrary = "0.1.2" -NonlinearSolveBase = "1.2" +NonlinearSolveBase = "1.3" NonlinearSolveFirstOrder = "1" NonlinearSolveQuasiNewton = "1" NonlinearSolveSpectralMethods = "1" diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index e7ed5851f..ae16dfd93 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -1,7 +1,7 @@ name = "NonlinearSolveBase" uuid = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" authors = ["Avik Pal and contributors"] -version = "1.2.0" +version = "1.3.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl index 412f6a748..649ac79d2 100644 --- a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl +++ b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl @@ -47,6 +47,7 @@ include("linear_solve.jl") include("timer_outputs.jl") include("tracing.jl") include("wrappers.jl") +include("polyalg.jl") include("descent/common.jl") include("descent/newton.jl") @@ -81,4 +82,6 @@ export RelTerminationMode, AbsTerminationMode, export DescentResult, SteepestDescent, NewtonDescent, DampedNewtonDescent, Dogleg, GeodesicAcceleration +export NonlinearSolvePolyAlgorithm + end diff --git a/lib/NonlinearSolveBase/src/polyalg.jl b/lib/NonlinearSolveBase/src/polyalg.jl new file mode 100644 index 000000000..54b61998f --- /dev/null +++ b/lib/NonlinearSolveBase/src/polyalg.jl @@ -0,0 +1,202 @@ +""" + 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 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 + +# 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/lib/NonlinearSolveBase/src/solve.jl b/lib/NonlinearSolveBase/src/solve.jl index 08b60e4db..76d765d7b 100644 --- a/lib/NonlinearSolveBase/src/solve.jl +++ b/lib/NonlinearSolveBase/src/solve.jl @@ -31,6 +31,164 @@ function CommonSolve.solve!(cache::AbstractNonlinearSolveCache) ) 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 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 + """ step!( cache::AbstractNonlinearSolveCache; diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index a3e8bb6ca..3aa4cac22 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -13,7 +13,8 @@ using LinearAlgebra: LinearAlgebra, norm using LineSearch: BackTracking using NonlinearSolveBase: NonlinearSolveBase, InternalAPI, AbstractNonlinearSolveAlgorithm, AbstractNonlinearSolveCache, Utils, L2_NORM, - enable_timer_outputs, disable_timer_outputs + enable_timer_outputs, disable_timer_outputs, + NonlinearSolvePolyAlgorithm using Preferences: set_preferences! using SciMLBase: SciMLBase, NLStats, ReturnCode, AbstractNonlinearProblem, @@ -47,7 +48,7 @@ using SimpleNonlinearSolve: SimpleNonlinearSolve const SII = SymbolicIndexingInterface -include("polyalg.jl") +include("poly_algs.jl") include("extension_algs.jl") include("default.jl") diff --git a/src/poly_algs.jl b/src/poly_algs.jl new file mode 100644 index 000000000..31c16917f --- /dev/null +++ b/src/poly_algs.jl @@ -0,0 +1,184 @@ +""" + RobustMultiNewton( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = 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, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) where {T} + common_kwargs = (; concrete_jac, linsolve, 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, + 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, + 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, 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, 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, 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, + 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, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) where {T} + common_kwargs = (; linsolve, 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 diff --git a/src/polyalg.jl b/src/polyalg.jl deleted file mode 100644 index f16dd7e0e..000000000 --- a/src/polyalg.jl +++ /dev/null @@ -1,545 +0,0 @@ -""" - 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, - 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, - autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing -) where {T} - common_kwargs = (; concrete_jac, linsolve, 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, - 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, - 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, 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, 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, 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, - 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, - autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing -) where {T} - common_kwargs = (; linsolve, 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