refactor: Remove unnecessary snippet

lib/NonlinearSolveBase/src/polyalg.jl

@@ -121,6 +121,78 @@ function SciMLBase.__init(
     )
 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}
@@ -160,6 +232,92 @@ 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
+
 # Original is often determined on runtime information especially for PolyAlgorithms so it
 # is best to never specialize on that
 function build_solution_less_specialize(

src/NonlinearSolve.jl

@@ -47,7 +47,7 @@ using SimpleNonlinearSolve: SimpleNonlinearSolve
 
 const SII = SymbolicIndexingInterface
 
-include("polyalg.jl")
+include("poly_algs.jl")
 include("extension_algs.jl")
 
 include("default.jl")

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