From 8d40ab40e2ccafa04e17ee75515622c476f97cea Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 28 Oct 2024 00:45:41 -0400 Subject: [PATCH] refactor: delete more code --- Project.toml | 3 +- lib/NonlinearSolveBase/src/solve.jl | 6 +- lib/NonlinearSolveBase/src/utils.jl | 4 +- .../src/NonlinearSolveFirstOrder.jl | 37 ++ .../src/gauss_newton.jl | 0 .../src/levenberg_marquardt.jl | 0 .../src/pseudo_transient.jl | 0 lib/NonlinearSolveFirstOrder/src/raphson.jl | 0 lib/NonlinearSolveFirstOrder/src/solve.jl | 303 ++++++++++++++++ .../src/trust_region.jl | 0 .../src/initialization.jl | 7 +- lib/NonlinearSolveQuasiNewton/src/solve.jl | 12 +- .../src/solve.jl | 1 + src/NonlinearSolve.jl | 15 - src/core/generalized_first_order.jl | 326 ------------------ src/core/generic.jl | 66 ---- src/core/noinit.jl | 37 -- src/utils.jl | 54 --- 18 files changed, 359 insertions(+), 512 deletions(-) create mode 100644 lib/NonlinearSolveFirstOrder/src/gauss_newton.jl create mode 100644 lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl create mode 100644 lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl create mode 100644 lib/NonlinearSolveFirstOrder/src/raphson.jl create mode 100644 lib/NonlinearSolveFirstOrder/src/solve.jl create mode 100644 lib/NonlinearSolveFirstOrder/src/trust_region.jl delete mode 100644 src/core/generalized_first_order.jl delete mode 100644 src/core/generic.jl delete mode 100644 src/core/noinit.jl diff --git a/Project.toml b/Project.toml index e08da6b42..1c6485b43 100644 --- a/Project.toml +++ b/Project.toml @@ -13,12 +13,12 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" -LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" 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" @@ -81,7 +81,6 @@ FixedPointAcceleration = "0.3" ForwardDiff = "0.10.36" Hwloc = "3" InteractiveUtils = "<0.0.1, 1" -LazyArrays = "1.8.2, 2" LeastSquaresOptim = "0.8.5" LineSearch = "0.1.4" LineSearches = "7.3" diff --git a/lib/NonlinearSolveBase/src/solve.jl b/lib/NonlinearSolveBase/src/solve.jl index 7316fd168..3c50521a0 100644 --- a/lib/NonlinearSolveBase/src/solve.jl +++ b/lib/NonlinearSolveBase/src/solve.jl @@ -32,8 +32,10 @@ function CommonSolve.solve!(cache::AbstractNonlinearSolveCache) end """ - step!(cache::AbstractNonlinearSolveCache; - recompute_jacobian::Union{Nothing, Bool} = nothing) + step!( + cache::AbstractNonlinearSolveCache; + recompute_jacobian::Union{Nothing, Bool} = nothing + ) Performs one step of the nonlinear solver. diff --git a/lib/NonlinearSolveBase/src/utils.jl b/lib/NonlinearSolveBase/src/utils.jl index a727bed6d..6d7a66aa0 100644 --- a/lib/NonlinearSolveBase/src/utils.jl +++ b/lib/NonlinearSolveBase/src/utils.jl @@ -248,7 +248,7 @@ function clean_sprint_struct(x) name = nameof(typeof(x)) for field in fieldnames(typeof(x)) val = getfield(x, field) - if field === :name + if field === :name && val isa Symbol && val !== :unknown name = val continue end @@ -268,7 +268,7 @@ function clean_sprint_struct(x, indent::Int) name = nameof(typeof(x)) for field in fieldnames(typeof(x)) val = getfield(x, field) - if field === :name + if field === :name && val isa Symbol && val !== :unknown name = val continue end diff --git a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl index f07882cc2..2cf9560e3 100644 --- a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl +++ b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl @@ -1,3 +1,40 @@ module NonlinearSolveFirstOrder +using Reexport: @reexport +using PrecompileTools: @compile_workload, @setup_workload + +using ArrayInterface: ArrayInterface +using CommonSolve: CommonSolve +using ConcreteStructs: @concrete +using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches +using LinearAlgebra: LinearAlgebra, Diagonal, dot, inv, diag +using LinearSolve: LinearSolve # Trigger Linear Solve extension in NonlinearSolveBase +using MaybeInplace: @bb +using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm, + AbstractNonlinearSolveCache, AbstractResetCondition, + AbstractResetConditionCache, AbstractApproximateJacobianStructure, + AbstractJacobianCache, AbstractJacobianInitialization, + AbstractApproximateJacobianUpdateRule, AbstractDescentDirection, + AbstractApproximateJacobianUpdateRuleCache, + Utils, InternalAPI, get_timer_output, @static_timeit, + update_trace!, L2_NORM, NewtonDescent +using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode +using SciMLOperators: AbstractSciMLOperator +using StaticArraysCore: StaticArray, Size, MArray + +include("raphson.jl") +include("gauss_newton.jl") +include("levenberg_marquardt.jl") +include("trust_region.jl") +include("pseudo_transient.jl") + +include("solve.jl") + +@reexport using SciMLBase, NonlinearSolveBase + +export NewtonRaphson, PseudoTransient +export GaussNewton, LevenbergMarquardt, TrustRegion + +export GeneralizedFirstOrderAlgorithm + end diff --git a/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl b/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl new file mode 100644 index 000000000..e69de29bb diff --git a/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl b/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl new file mode 100644 index 000000000..e69de29bb diff --git a/lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl b/lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl new file mode 100644 index 000000000..e69de29bb diff --git a/lib/NonlinearSolveFirstOrder/src/raphson.jl b/lib/NonlinearSolveFirstOrder/src/raphson.jl new file mode 100644 index 000000000..e69de29bb diff --git a/lib/NonlinearSolveFirstOrder/src/solve.jl b/lib/NonlinearSolveFirstOrder/src/solve.jl new file mode 100644 index 000000000..0d60aa3d7 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/src/solve.jl @@ -0,0 +1,303 @@ +""" + GeneralizedFirstOrderAlgorithm(; + descent, linesearch = missing, + trustregion = missing, autodiff = nothing, vjp_autodiff = nothing, + jvp_autodiff = nothing, max_shrink_times::Int = typemax(Int), + concrete_jac = Val(false), name::Symbol = :unknown + ) + +This is a Generalization of First-Order (uses Jacobian) Nonlinear Solve Algorithms. The most +common example of this is Newton-Raphson Method. + +First Order here refers to the order of differentiation, and should not be confused with the +order of convergence. + +### Keyword Arguments + + - `trustregion`: Globalization using a Trust Region Method. This needs to follow the + [`NonlinearSolve.AbstractTrustRegionMethod`](@ref) interface. + - `descent`: The descent method to use to compute the step. This needs to follow the + [`NonlinearSolve.AbstractDescentAlgorithm`](@ref) interface. + - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk + before the algorithm terminates. +""" +@concrete struct GeneralizedFirstOrderAlgorithm <: AbstractNonlinearSolveAlgorithm + linesearch + trustregion + descent + max_shrink_times::Int + + autodiff + vjp_autodiff + jvp_autodiff + + concrete_jac <: Union{Val{false}, Val{true}} + name::Symbol +end + +function GeneralizedFirstOrderAlgorithm(; + descent, linesearch = missing, trustregion = missing, autodiff = nothing, + vjp_autodiff = nothing, jvp_autodiff = nothing, max_shrink_times::Int = typemax(Int), + concrete_jac = Val(false), name::Symbol = :unknown) + return GeneralizedFirstOrderAlgorithm( + linesearch, trustregion, descent, max_shrink_times, + autodiff, vjp_autodiff, jvp_autodiff, + concrete_jac, name + ) +end + +@concrete mutable struct GeneralizedFirstOrderAlgorithmCache <: AbstractNonlinearSolveCache + # Basic Requirements + fu + u + u_cache + p + du # Aliased to `get_du(descent_cache)` + J # Aliased to `jac_cache.J` + alg <: GeneralizedFirstOrderAlgorithm + prob <: AbstractNonlinearProblem + globalization <: Union{Val{:LineSearch}, Val{:TrustRegion}, Val{:None}} + + # Internal Caches + jac_cache + descent_cache + linesearch_cache + trustregion_cache + + # Counters + stats::NLStats + nsteps::Int + maxiters::Int + maxtime + max_shrink_times::Int + + # Timer + timer + total_time::Float64 + + # State Affect + make_new_jacobian::Bool + + # Termination & Tracking + termination_cache + trace + retcode::ReturnCode.T + force_stop::Bool + kwargs +end + +# XXX: Implement +# function __reinit_internal!( +# cache::GeneralizedFirstOrderAlgorithmCache{iip}, args...; p = cache.p, u0 = cache.u, +# alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs...) where {iip} +# if iip +# recursivecopy!(cache.u, u0) +# cache.prob.f(cache.fu, cache.u, p) +# else +# cache.u = __maybe_unaliased(u0, alias_u0) +# set_fu!(cache, cache.prob.f(cache.u, p)) +# end +# cache.p = p + +# __reinit_internal!(cache.stats) +# cache.nsteps = 0 +# cache.maxiters = maxiters +# cache.maxtime = maxtime +# cache.total_time = 0.0 +# cache.force_stop = false +# cache.retcode = ReturnCode.Default +# cache.make_new_jacobian = true + +# reset!(cache.trace) +# reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...) +# reset_timer!(cache.timer) +# end + +NonlinearSolveBase.@internal_caches(GeneralizedFirstOrderAlgorithmCache, + :jac_cache, :descent_cache, :linesearch_cache, :trustregion_cache) + +# function SciMLBase.__init( +# prob::AbstractNonlinearProblem{uType, iip}, alg::GeneralizedFirstOrderAlgorithm, +# args...; stats = empty_nlstats(), alias_u0 = false, maxiters = 1000, +# abstol = nothing, reltol = nothing, maxtime = nothing, +# termination_condition = nothing, internalnorm = L2_NORM, +# linsolve_kwargs = (;), kwargs...) where {uType, iip} +# autodiff = select_jacobian_autodiff(prob, alg.autodiff) +# jvp_autodiff = if alg.jvp_autodiff === nothing && alg.autodiff !== nothing && +# (ADTypes.mode(alg.autodiff) isa ADTypes.ForwardMode || +# ADTypes.mode(alg.autodiff) isa ADTypes.ForwardOrReverseMode) +# select_forward_mode_autodiff(prob, alg.autodiff) +# else +# select_forward_mode_autodiff(prob, alg.jvp_autodiff) +# end +# vjp_autodiff = if alg.vjp_autodiff === nothing && alg.autodiff !== nothing && +# (ADTypes.mode(alg.autodiff) isa ADTypes.ReverseMode || +# ADTypes.mode(alg.autodiff) isa ADTypes.ForwardOrReverseMode) +# select_reverse_mode_autodiff(prob, alg.autodiff) +# else +# select_reverse_mode_autodiff(prob, alg.vjp_autodiff) +# end + +# timer = get_timer_output() +# @static_timeit timer "cache construction" begin +# (; f, u0, p) = prob +# u = __maybe_unaliased(u0, alias_u0) +# fu = evaluate_f(prob, u) +# @bb u_cache = copy(u) + +# linsolve = get_linear_solver(alg.descent) + +# abstol, reltol, termination_cache = NonlinearSolveBase.init_termination_cache( +# prob, abstol, reltol, fu, u, termination_condition, Val(:regular)) +# linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) + +# jac_cache = construct_jacobian_cache( +# prob, alg, f, fu, u, p; stats, autodiff, linsolve, jvp_autodiff, vjp_autodiff) +# J = jac_cache(nothing) + +# descent_cache = __internal_init(prob, alg.descent, J, fu, u; stats, abstol, +# reltol, internalnorm, linsolve_kwargs, timer) +# du = get_du(descent_cache) + +# has_linesearch = alg.linesearch !== missing && alg.linesearch !== nothing +# has_trustregion = alg.trustregion !== missing && alg.trustregion !== nothing + +# if has_trustregion && has_linesearch +# error("TrustRegion and LineSearch methods are algorithmically incompatible.") +# end + +# GB = :None +# linesearch_cache = nothing +# trustregion_cache = nothing + +# if has_trustregion +# supports_trust_region(alg.descent) || error("Trust Region not supported by \ +# $(alg.descent).") +# trustregion_cache = __internal_init( +# prob, alg.trustregion, f, fu, u, p; stats, internalnorm, kwargs..., +# autodiff, jvp_autodiff, vjp_autodiff) +# GB = :TrustRegion +# end + +# if has_linesearch +# supports_line_search(alg.descent) || error("Line Search not supported by \ +# $(alg.descent).") +# linesearch_cache = init( +# prob, alg.linesearch, fu, u; stats, autodiff = jvp_autodiff, kwargs...) +# GB = :LineSearch +# end + +# trace = init_nonlinearsolve_trace( +# prob, alg, u, fu, ApplyArray(__zero, J), du; kwargs...) + +# return GeneralizedFirstOrderAlgorithmCache{iip, GB, maxtime !== nothing}( +# fu, u, u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, +# trustregion_cache, stats, 0, maxiters, maxtime, alg.max_shrink_times, timer, +# 0.0, true, termination_cache, trace, ReturnCode.Default, false, kwargs) +# end +# end + +# function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; +# recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} +# @static_timeit cache.timer "jacobian" begin +# if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian +# J = cache.jac_cache(cache.u) +# new_jacobian = true +# else +# J = cache.jac_cache(nothing) +# new_jacobian = false +# end +# end + +# @static_timeit cache.timer "descent" begin +# if cache.trustregion_cache !== nothing && +# hasfield(typeof(cache.trustregion_cache), :trust_region) +# descent_result = __internal_solve!( +# cache.descent_cache, J, cache.fu, cache.u; new_jacobian, +# trust_region = cache.trustregion_cache.trust_region, cache.kwargs...) +# else +# descent_result = __internal_solve!( +# cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...) +# end +# end + +# if !descent_result.linsolve_success +# if new_jacobian +# # Jacobian Information is current and linear solve failed terminate the solve +# cache.retcode = ReturnCode.InternalLinearSolveFailed +# cache.force_stop = true +# return +# else +# # Jacobian Information is not current and linear solve failed, recompute +# # Jacobian +# if !haskey(cache.kwargs, :verbose) || cache.kwargs[:verbose] +# @warn "Linear Solve Failed but Jacobian Information is not current. \ +# Retrying with updated Jacobian." +# end +# # In the 2nd call the `new_jacobian` is guaranteed to be `true`. +# cache.make_new_jacobian = true +# __step!(cache; recompute_jacobian = true, kwargs...) +# return +# end +# end + +# δu, descent_intermediates = descent_result.δu, descent_result.extras + +# if descent_result.success +# cache.make_new_jacobian = true +# if GB === :LineSearch +# @static_timeit cache.timer "linesearch" begin +# linesearch_sol = solve!(cache.linesearch_cache, cache.u, δu) +# linesearch_failed = !SciMLBase.successful_retcode(linesearch_sol.retcode) +# α = linesearch_sol.step_size +# end +# if linesearch_failed +# cache.retcode = ReturnCode.InternalLineSearchFailed +# cache.force_stop = true +# end +# @static_timeit cache.timer "step" begin +# @bb axpy!(α, δu, cache.u) +# evaluate_f!(cache, cache.u, cache.p) +# end +# elseif GB === :TrustRegion +# @static_timeit cache.timer "trustregion" begin +# tr_accepted, u_new, fu_new = __internal_solve!( +# cache.trustregion_cache, J, cache.fu, +# cache.u, δu, descent_intermediates) +# if tr_accepted +# @bb copyto!(cache.u, u_new) +# @bb copyto!(cache.fu, fu_new) +# α = true +# else +# α = false +# cache.make_new_jacobian = false +# end +# if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && +# cache.trustregion_cache.shrink_counter > cache.max_shrink_times +# cache.retcode = ReturnCode.ShrinkThresholdExceeded +# cache.force_stop = true +# end +# end +# elseif GB === :None +# @static_timeit cache.timer "step" begin +# @bb axpy!(1, δu, cache.u) +# evaluate_f!(cache, cache.u, cache.p) +# end +# α = true +# else +# error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ +# :TrustRegion, :None)") +# end +# check_and_update!(cache, cache.fu, cache.u, cache.u_cache) +# else +# α = false +# cache.make_new_jacobian = false +# end + +# update_trace!(cache, α) +# @bb copyto!(cache.u_cache, cache.u) + +# callback_into_cache!(cache) + +# return nothing +# end diff --git a/lib/NonlinearSolveFirstOrder/src/trust_region.jl b/lib/NonlinearSolveFirstOrder/src/trust_region.jl new file mode 100644 index 000000000..e69de29bb diff --git a/lib/NonlinearSolveQuasiNewton/src/initialization.jl b/lib/NonlinearSolveQuasiNewton/src/initialization.jl index 23d9a19cd..6c22e0240 100644 --- a/lib/NonlinearSolveQuasiNewton/src/initialization.jl +++ b/lib/NonlinearSolveQuasiNewton/src/initialization.jl @@ -1,6 +1,7 @@ """ - InitializedApproximateJacobianCache(J, structure, alg, cache, initialized::Bool, - internalnorm) + InitializedApproximateJacobianCache( + J, structure, alg, cache, initialized::Bool, internalnorm + ) A cache for Approximate Jacobian. @@ -22,7 +23,7 @@ A cache for Approximate Jacobian. Returns the current Jacobian `cache.J` with the proper `structure`. ```julia -__internal_solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) +InternalAPI.solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) ``` Solves for the Jacobian `cache.J` and returns it. If `reinit` is `true`, then the Jacobian diff --git a/lib/NonlinearSolveQuasiNewton/src/solve.jl b/lib/NonlinearSolveQuasiNewton/src/solve.jl index 98ef6d9ee..5b987a10d 100644 --- a/lib/NonlinearSolveQuasiNewton/src/solve.jl +++ b/lib/NonlinearSolveQuasiNewton/src/solve.jl @@ -30,9 +30,11 @@ examples include [`Broyden`](@ref)'s Method. descent <: AbstractDescentDirection update_rule <: AbstractApproximateJacobianUpdateRule reinit_rule <: AbstractResetCondition + initialization <: AbstractJacobianInitialization + max_resets::Int max_shrink_times::Int - initialization + concrete_jac <: Union{Val{false}, Val{true}} name::Symbol end @@ -43,8 +45,8 @@ function QuasiNewtonAlgorithm(; max_shrink_times::Int = typemax(Int), concrete_jac = Val(false) ) return QuasiNewtonAlgorithm( - linesearch, trustregion, descent, update_rule, reinit_rule, - max_resets, max_shrink_times, initialization, concrete_jac, name + linesearch, trustregion, descent, update_rule, reinit_rule, initialization, + max_resets, max_shrink_times, concrete_jac, name ) end @@ -94,7 +96,7 @@ end end # XXX: Implement -# function __reinit_internal!(cache::ApproximateJacobianSolveCache{INV, GB, iip}, +# function __reinit_internal!(cache::QuasiNewtonCache{INV, GB, iip}, # args...; p = cache.p, u0 = cache.u, alias_u0::Bool = false, # maxiters = 1000, maxtime = nothing, kwargs...) where {INV, GB, iip} # if iip @@ -122,7 +124,7 @@ end # reset_timer!(cache.timer) # end -NonlinearSolveBase.@internal_caches(ApproximateJacobianSolveCache, +NonlinearSolveBase.@internal_caches(QuasiNewtonCache, :initialization_cache, :descent_cache, :linesearch_cache, :trustregion_cache, :update_rule_cache, :reinit_rule_cache) diff --git a/lib/NonlinearSolveSpectralMethods/src/solve.jl b/lib/NonlinearSolveSpectralMethods/src/solve.jl index 297b9bf69..518e9a453 100644 --- a/lib/NonlinearSolveSpectralMethods/src/solve.jl +++ b/lib/NonlinearSolveSpectralMethods/src/solve.jl @@ -23,6 +23,7 @@ Method. σ_min σ_max σ_1 + name::Symbol end diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 5e1f39b34..950710932 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -9,7 +9,6 @@ using CommonSolve: solve, init, solve! using ConcreteStructs: @concrete using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches using FastClosures: @closure -using LazyArrays: LazyArrays, ApplyArray, cache using LinearAlgebra: LinearAlgebra, Diagonal, I, LowerTriangular, Symmetric, UpperTriangular, axpy!, cond, diag, diagind, dot, issuccess, istril, istriu, lu, mul!, norm, pinv, tril!, triu! @@ -34,14 +33,6 @@ using NonlinearSolveBase: NonlinearSolveBase, using NonlinearSolveQuasiNewton: Broyden, Klement -# XXX: Remove -import NonlinearSolveBase: InternalAPI, concrete_jac, supports_line_search, - supports_trust_region, last_step_accepted, get_linear_solver, - AbstractDampingFunction, AbstractDampingFunctionCache, - requires_normal_form_jacobian, requires_normal_form_rhs, - returns_norm_form_damping, get_timer_output, get_u, get_fu, - set_fu! - using Preferences: Preferences, set_preferences! using RecursiveArrayTools: recursivecopy! using SciMLBase: SciMLBase, AbstractNonlinearAlgorithm, AbstractNonlinearProblem, @@ -51,9 +42,6 @@ using SciMLBase: SciMLBase, AbstractNonlinearAlgorithm, AbstractNonlinearProblem using SciMLOperators: AbstractSciMLOperator using SimpleNonlinearSolve: SimpleNonlinearSolve using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix -using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy, - symbolic_container, parameter_values, state_values, getu, - setu # AD Support using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff, @@ -72,15 +60,12 @@ const DI = DifferentiationInterface const True = Val(true) const False = Val(false) -include("abstract_types.jl") include("timer_outputs.jl") include("internal/helpers.jl") include("globalization/trust_region.jl") -include("core/generic.jl") include("core/generalized_first_order.jl") -include("core/noinit.jl") include("algorithms/raphson.jl") include("algorithms/pseudo_transient.jl") diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl deleted file mode 100644 index 261adbfb0..000000000 --- a/src/core/generalized_first_order.jl +++ /dev/null @@ -1,326 +0,0 @@ -""" - GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; descent, linesearch = missing, - trustregion = missing, autodiff = nothing, vjp_autodiff = nothing, - jvp_autodiff = nothing, max_shrink_times::Int = typemax(Int)) - GeneralizedFirstOrderAlgorithm(; concrete_jac = nothing, name::Symbol = :unknown, - kwargs...) - -This is a Generalization of First-Order (uses Jacobian) Nonlinear Solve Algorithms. The most -common example of this is Newton-Raphson Method. - -First Order here refers to the order of differentiation, and should not be confused with the -order of convergence. - -!!! danger - - `trustregion` and `linesearch` cannot be specified together. - -### Keyword Arguments - - - `trustregion`: Globalization using a Trust Region Method. This needs to follow the - [`NonlinearSolve.AbstractTrustRegionMethod`](@ref) interface. - - `descent`: The descent method to use to compute the step. This needs to follow the - [`NonlinearSolve.AbstractDescentAlgorithm`](@ref) interface. - - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk - before the algorithm terminates. -""" -@concrete struct GeneralizedFirstOrderAlgorithm{concrete_jac, name} <: - AbstractNonlinearSolveAlgorithm{name} - linesearch - trustregion - descent - max_shrink_times::Int - - autodiff - vjp_autodiff - jvp_autodiff -end - -function __show_algorithm(io::IO, alg::GeneralizedFirstOrderAlgorithm, name, indent) - modifiers = String[] - __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") - __is_present(alg.trustregion) && push!(modifiers, "trustregion = $(alg.trustregion)") - push!(modifiers, "descent = $(alg.descent)") - __is_present(alg.autodiff) && push!(modifiers, "autodiff = $(alg.autodiff)") - __is_present(alg.vjp_autodiff) && push!(modifiers, "vjp_autodiff = $(alg.vjp_autodiff)") - __is_present(alg.jvp_autodiff) && push!(modifiers, "jvp_autodiff = $(alg.jvp_autodiff)") - spacing = " "^indent * " " - spacing_last = " "^indent - print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))") -end - -function GeneralizedFirstOrderAlgorithm(; - concrete_jac = nothing, name::Symbol = :unknown, kwargs...) - return GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; kwargs...) -end - -function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; - descent, linesearch = missing, trustregion = missing, - autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing, - max_shrink_times::Int = typemax(Int)) where {concrete_jac, name} - return GeneralizedFirstOrderAlgorithm{concrete_jac, name}( - linesearch, trustregion, descent, max_shrink_times, - autodiff, vjp_autodiff, jvp_autodiff) -end - -# XXX: Remove -concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = concrete_jac(CJ) - -@concrete mutable struct GeneralizedFirstOrderAlgorithmCache{iip, GB, timeit} <: - AbstractNonlinearSolveCache{iip, timeit} - # Basic Requirements - fu - u - u_cache - p - du # Aliased to `get_du(descent_cache)` - J # Aliased to `jac_cache.J` - alg - prob - - # Internal Caches - jac_cache - descent_cache - linesearch_cache - trustregion_cache - - # Counters - stats::NLStats - nsteps::Int - maxiters::Int - maxtime - max_shrink_times::Int - - # Timer - timer - total_time::Float64 # Simple Counter which works even if TimerOutput is disabled - - # State Affect - make_new_jacobian::Bool - - # Termination & Tracking - termination_cache - trace - retcode::ReturnCode.T - force_stop::Bool - kwargs -end - -SymbolicIndexingInterface.state_values(cache::GeneralizedFirstOrderAlgorithmCache) = cache.u -function SymbolicIndexingInterface.parameter_values(cache::GeneralizedFirstOrderAlgorithmCache) - cache.p -end - -function __reinit_internal!( - cache::GeneralizedFirstOrderAlgorithmCache{iip}, args...; p = cache.p, u0 = cache.u, - alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs...) where {iip} - if iip - recursivecopy!(cache.u, u0) - cache.prob.f(cache.fu, cache.u, p) - else - cache.u = __maybe_unaliased(u0, alias_u0) - set_fu!(cache, cache.prob.f(cache.u, p)) - end - cache.p = p - - __reinit_internal!(cache.stats) - cache.nsteps = 0 - cache.maxiters = maxiters - cache.maxtime = maxtime - cache.total_time = 0.0 - cache.force_stop = false - cache.retcode = ReturnCode.Default - cache.make_new_jacobian = true - - reset!(cache.trace) - reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...) - reset_timer!(cache.timer) -end - -@internal_caches GeneralizedFirstOrderAlgorithmCache :jac_cache :descent_cache :linesearch_cache :trustregion_cache - -function SciMLBase.__init( - prob::AbstractNonlinearProblem{uType, iip}, alg::GeneralizedFirstOrderAlgorithm, - args...; stats = empty_nlstats(), alias_u0 = false, maxiters = 1000, - abstol = nothing, reltol = nothing, maxtime = nothing, - termination_condition = nothing, internalnorm = L2_NORM, - linsolve_kwargs = (;), kwargs...) where {uType, iip} - autodiff = select_jacobian_autodiff(prob, alg.autodiff) - jvp_autodiff = if alg.jvp_autodiff === nothing && alg.autodiff !== nothing && - (ADTypes.mode(alg.autodiff) isa ADTypes.ForwardMode || - ADTypes.mode(alg.autodiff) isa ADTypes.ForwardOrReverseMode) - select_forward_mode_autodiff(prob, alg.autodiff) - else - select_forward_mode_autodiff(prob, alg.jvp_autodiff) - end - vjp_autodiff = if alg.vjp_autodiff === nothing && alg.autodiff !== nothing && - (ADTypes.mode(alg.autodiff) isa ADTypes.ReverseMode || - ADTypes.mode(alg.autodiff) isa ADTypes.ForwardOrReverseMode) - select_reverse_mode_autodiff(prob, alg.autodiff) - else - select_reverse_mode_autodiff(prob, alg.vjp_autodiff) - end - - timer = get_timer_output() - @static_timeit timer "cache construction" begin - (; f, u0, p) = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - @bb u_cache = copy(u) - - linsolve = get_linear_solver(alg.descent) - - abstol, reltol, termination_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fu, u, termination_condition, Val(:regular)) - linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) - - jac_cache = construct_jacobian_cache( - prob, alg, f, fu, u, p; stats, autodiff, linsolve, jvp_autodiff, vjp_autodiff) - J = jac_cache(nothing) - - descent_cache = __internal_init(prob, alg.descent, J, fu, u; stats, abstol, - reltol, internalnorm, linsolve_kwargs, timer) - du = get_du(descent_cache) - - has_linesearch = alg.linesearch !== missing && alg.linesearch !== nothing - has_trustregion = alg.trustregion !== missing && alg.trustregion !== nothing - - if has_trustregion && has_linesearch - error("TrustRegion and LineSearch methods are algorithmically incompatible.") - end - - GB = :None - linesearch_cache = nothing - trustregion_cache = nothing - - if has_trustregion - supports_trust_region(alg.descent) || error("Trust Region not supported by \ - $(alg.descent).") - trustregion_cache = __internal_init( - prob, alg.trustregion, f, fu, u, p; stats, internalnorm, kwargs..., - autodiff, jvp_autodiff, vjp_autodiff) - GB = :TrustRegion - end - - if has_linesearch - supports_line_search(alg.descent) || error("Line Search not supported by \ - $(alg.descent).") - linesearch_cache = init( - prob, alg.linesearch, fu, u; stats, autodiff = jvp_autodiff, kwargs...) - GB = :LineSearch - end - - trace = init_nonlinearsolve_trace( - prob, alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - - return GeneralizedFirstOrderAlgorithmCache{iip, GB, maxtime !== nothing}( - fu, u, u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, - trustregion_cache, stats, 0, maxiters, maxtime, alg.max_shrink_times, timer, - 0.0, true, termination_cache, trace, ReturnCode.Default, false, kwargs) - end -end - -function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; - recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} - @static_timeit cache.timer "jacobian" begin - if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian - J = cache.jac_cache(cache.u) - new_jacobian = true - else - J = cache.jac_cache(nothing) - new_jacobian = false - end - end - - @static_timeit cache.timer "descent" begin - if cache.trustregion_cache !== nothing && - hasfield(typeof(cache.trustregion_cache), :trust_region) - descent_result = __internal_solve!( - cache.descent_cache, J, cache.fu, cache.u; new_jacobian, - trust_region = cache.trustregion_cache.trust_region, cache.kwargs...) - else - descent_result = __internal_solve!( - cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...) - end - end - - if !descent_result.linsolve_success - if new_jacobian - # Jacobian Information is current and linear solve failed terminate the solve - cache.retcode = ReturnCode.InternalLinearSolveFailed - cache.force_stop = true - return - else - # Jacobian Information is not current and linear solve failed, recompute - # Jacobian - if !haskey(cache.kwargs, :verbose) || cache.kwargs[:verbose] - @warn "Linear Solve Failed but Jacobian Information is not current. \ - Retrying with updated Jacobian." - end - # In the 2nd call the `new_jacobian` is guaranteed to be `true`. - cache.make_new_jacobian = true - __step!(cache; recompute_jacobian = true, kwargs...) - return - end - end - - δu, descent_intermediates = descent_result.δu, descent_result.extras - - if descent_result.success - cache.make_new_jacobian = true - if GB === :LineSearch - @static_timeit cache.timer "linesearch" begin - linesearch_sol = solve!(cache.linesearch_cache, cache.u, δu) - linesearch_failed = !SciMLBase.successful_retcode(linesearch_sol.retcode) - α = linesearch_sol.step_size - end - if linesearch_failed - cache.retcode = ReturnCode.InternalLineSearchFailed - cache.force_stop = true - end - @static_timeit cache.timer "step" begin - @bb axpy!(α, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) - end - elseif GB === :TrustRegion - @static_timeit cache.timer "trustregion" begin - tr_accepted, u_new, fu_new = __internal_solve!( - cache.trustregion_cache, J, cache.fu, - cache.u, δu, descent_intermediates) - if tr_accepted - @bb copyto!(cache.u, u_new) - @bb copyto!(cache.fu, fu_new) - α = true - else - α = false - cache.make_new_jacobian = false - end - if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && - cache.trustregion_cache.shrink_counter > cache.max_shrink_times - cache.retcode = ReturnCode.ShrinkThresholdExceeded - cache.force_stop = true - end - end - elseif GB === :None - @static_timeit cache.timer "step" begin - @bb axpy!(1, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) - end - α = true - else - error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ - :TrustRegion, :None)") - end - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - else - α = false - cache.make_new_jacobian = false - end - - update_trace!(cache, α) - @bb copyto!(cache.u_cache, cache.u) - - callback_into_cache!(cache) - - return nothing -end diff --git a/src/core/generic.jl b/src/core/generic.jl deleted file mode 100644 index df6f2ecff..000000000 --- a/src/core/generic.jl +++ /dev/null @@ -1,66 +0,0 @@ -function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, - alg::AbstractNonlinearSolveAlgorithm, args...; stats = empty_nlstats(), kwargs...) - cache = SciMLBase.__init(prob, alg, args...; stats, kwargs...) - return solve!(cache) -end - -function not_terminated(cache::AbstractNonlinearSolveCache) - return !cache.force_stop && cache.nsteps < cache.maxiters -end - -function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) - while not_terminated(cache) - step!(cache) - end - - # The solver might have set a different `retcode` - if cache.retcode == ReturnCode.Default - cache.retcode = ifelse( - cache.nsteps ≥ cache.maxiters, ReturnCode.MaxIters, ReturnCode.Success) - end - - update_from_termination_cache!(cache.termination_cache, cache) - - update_trace!(cache.trace, cache.nsteps, get_u(cache), - get_fu(cache), nothing, nothing, nothing; last = True) - - return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode, cache.stats, cache.trace) -end - -""" - step!(cache::AbstractNonlinearSolveCache; - recompute_jacobian::Union{Nothing, Bool} = nothing) - -Performs one step of the nonlinear solver. - -### Keyword Arguments - - - `recompute_jacobian`: allows controlling whether the jacobian is recomputed at the - current step. If `nothing`, then the algorithm determines whether to recompute the - jacobian. If `true` or `false`, then the jacobian is recomputed or not recomputed, - respectively. For algorithms that don't use jacobian information, this keyword is - ignored with a one-time warning. -""" -function SciMLBase.step!(cache::AbstractNonlinearSolveCache{iip, timeit}, - args...; kwargs...) where {iip, timeit} - not_terminated(cache) || return - timeit && (time_start = time()) - res = @static_timeit cache.timer "solve" begin - __step!(cache, args...; kwargs...) - end - - cache.stats.nsteps += 1 - cache.nsteps += 1 - - if timeit - cache.total_time += time() - time_start - if !cache.force_stop && cache.retcode == ReturnCode.Default && - cache.total_time ≥ cache.maxtime - cache.retcode = ReturnCode.MaxTime - cache.force_stop = true - end - end - - return res -end diff --git a/src/core/noinit.jl b/src/core/noinit.jl deleted file mode 100644 index b51c09c23..000000000 --- a/src/core/noinit.jl +++ /dev/null @@ -1,37 +0,0 @@ -# Some algorithms don't support creating a cache and doing `solve!`, this unfortunately -# makes it difficult to write generic code that supports caching. For the algorithms that -# don't have a `__init` function defined, we create a "Fake Cache", which just calls -# `__solve` from `solve!` -@concrete mutable struct NonlinearSolveNoInitCache{iip, timeit} <: - AbstractNonlinearSolveCache{iip, timeit} - prob - alg - args - kwargs::Any -end - -function SciMLBase.reinit!( - cache::NonlinearSolveNoInitCache, u0 = cache.prob.u0; p = cache.prob.p, kwargs...) - prob = remake(cache.prob; u0, p) - cache.prob = prob - cache.kwargs = merge(cache.kwargs, kwargs) - return cache -end - -function Base.show(io::IO, cache::NonlinearSolveNoInitCache) - print(io, "NonlinearSolveNoInitCache(alg = $(cache.alg))") -end - -function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, - alg::Union{AbstractNonlinearSolveAlgorithm, - SimpleNonlinearSolve.AbstractSimpleNonlinearSolveAlgorithm}, - args...; - maxtime = nothing, - kwargs...) where {uType, iip} - return NonlinearSolveNoInitCache{iip, maxtime !== nothing}( - prob, alg, args, merge((; maxtime), kwargs)) -end - -function SciMLBase.solve!(cache::NonlinearSolveNoInitCache) - return solve(cache.prob, cache.alg, cache.args...; cache.kwargs...) -end diff --git a/src/utils.jl b/src/utils.jl index cb346da0a..20602d2db 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,53 +1,3 @@ -# Helper Functions -@generated function __getproperty(s::S, ::Val{X}) where {S, X} - hasfield(S, X) && return :(s.$X) - return :(missing) -end - -@inline @generated function _vec(v) - hasmethod(vec, Tuple{typeof(v)}) || return :(vec(v)) - return :(v) -end -@inline _vec(v::Number) = v -@inline _vec(v::AbstractVector) = v - -@inline _restructure(y, x) = restructure(y, x) -@inline _restructure(y::Number, x::Number) = x - -@inline __maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x -@inline function __maybe_unaliased(x::AbstractArray, alias::Bool) - # Spend time coping iff we will mutate the array - (alias || !__can_setindex(typeof(x))) && return x - return deepcopy(x) -end -@inline __maybe_unaliased(x::AbstractSciMLOperator, ::Bool) = x - -@inline __copy(x::AbstractArray) = copy(x) -@inline __copy(x::Number) = x -@inline __copy(x) = x - -# LazyArrays for tracing -__zero(x::AbstractArray) = zero(x) -__zero(x) = x -LazyArrays.applied_eltype(::typeof(__zero), x) = eltype(x) -LazyArrays.applied_ndims(::typeof(__zero), x) = ndims(x) -LazyArrays.applied_size(::typeof(__zero), x) = size(x) -LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) - -# Use Symmetric Matrices if known to be efficient -@inline __maybe_symmetric(x) = Symmetric(x) -@inline __maybe_symmetric(x::Number) = x -## LinearSolve with `nothing` doesn't dispatch correctly here -@inline __maybe_symmetric(x::StaticArray) = x -@inline __maybe_symmetric(x::AbstractSparseMatrix) = x -@inline __maybe_symmetric(x::AbstractSciMLOperator) = x - -# Simple Checks -@inline __is_present(::Nothing) = false -@inline __is_present(::Missing) = false -@inline __is_present(::Any) = true -@inline __is_present(::NoLineSearch) = false - @inline __is_complex(::Type{ComplexF64}) = true @inline __is_complex(::Type{ComplexF32}) = true @inline __is_complex(::Type{Complex}) = true @@ -76,10 +26,6 @@ end return fx_idx, idx end -@inline __can_setindex(x) = can_setindex(x) -@inline __can_setindex(::Number) = false - -@inline __dot(x, y) = dot(_vec(x), _vec(y)) """ pickchunksize(x) = pickchunksize(length(x))