diff --git a/docs/src/devdocs/operators.md b/docs/src/devdocs/operators.md index 15d00093a..fcbadc778 100644 --- a/docs/src/devdocs/operators.md +++ b/docs/src/devdocs/operators.md @@ -1,11 +1,5 @@ # Custom SciML Operators -## Abstract Operators - -```@docs -NonlinearSolve.AbstractNonlinearSolveOperator -``` - ## Low-Rank Jacobian Operators ```@docs diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index df76e09a4..5a8cd1f34 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -17,11 +17,13 @@ FunctionProperties = "f62d2435-5019-4c03-9749-2d4c77af0cbc" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" +TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [weakdeps] DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" @@ -57,6 +59,7 @@ LinearAlgebra = "1.10" LinearSolve = "2.36.1" Markdown = "1.10" MaybeInplace = "0.1.4" +Preferences = "1.4" RecursiveArrayTools = "3" SciMLBase = "2.50" SciMLJacobianOperators = "0.1.1" @@ -65,6 +68,7 @@ SparseArrays = "1.10" SparseMatrixColorings = "0.4.8" StaticArraysCore = "1.4" Test = "1.10" +TimerOutputs = "0.5.23" julia = "1.10" [extras] diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl index 3b6ae2935..fca3793ad 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl @@ -1,6 +1,6 @@ module NonlinearSolveBaseSparseArraysExt -using NonlinearSolveBase: NonlinearSolveBase +using NonlinearSolveBase: NonlinearSolveBase, Utils using SparseArrays: AbstractSparseMatrix, AbstractSparseMatrixCSC, nonzeros function NonlinearSolveBase.NAN_CHECK(x::AbstractSparseMatrixCSC) @@ -9,4 +9,6 @@ end NonlinearSolveBase.sparse_or_structured_prototype(::AbstractSparseMatrix) = true +Utils.maybe_symmetric(x::AbstractSparseMatrix) = x + end diff --git a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl index 585ac1c5a..5007763b2 100644 --- a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl +++ b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl @@ -11,9 +11,10 @@ using DifferentiationInterface: DifferentiationInterface, Constant using EnzymeCore: EnzymeCore using FastClosures: @closure using FunctionProperties: hasbranching -using LinearAlgebra: Diagonal, norm, ldiv! +using LinearAlgebra: LinearAlgebra, Diagonal, norm, ldiv!, diagind using Markdown: @doc_str using MaybeInplace: @bb +using Preferences: @load_preference using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinearProblem, AbstractNonlinearAlgorithm, AbstractNonlinearFunction, @@ -37,6 +38,14 @@ include("termination_conditions.jl") include("autodiff.jl") include("jacobian.jl") include("linear_solve.jl") +include("timer_outputs.jl") + +include("descent/common.jl") +include("descent/newton.jl") +include("descent/steepest.jl") +include("descent/damped_newton.jl") +include("descent/dogleg.jl") +include("descent/geodesic_acceleration.jl") # Unexported Public API @compat(public, (L2_NORM, Linf_NORM, NAN_CHECK, UNITLESS_ABS2, get_tolerance)) @@ -55,4 +64,7 @@ export RelTerminationMode, AbsTerminationMode, RelNormSafeTerminationMode, AbsNormSafeTerminationMode, RelNormSafeBestTerminationMode, AbsNormSafeBestTerminationMode +export DescentResult, SteepestDescent, NewtonDescent, DampedNewtonDescent, Dogleg, + GeodesicAcceleration + end diff --git a/lib/NonlinearSolveBase/src/abstract_types.jl b/lib/NonlinearSolveBase/src/abstract_types.jl index 4dba65a14..42c57842a 100644 --- a/lib/NonlinearSolveBase/src/abstract_types.jl +++ b/lib/NonlinearSolveBase/src/abstract_types.jl @@ -6,7 +6,65 @@ function reinit! end end -abstract type AbstractDescentDirection end +abstract type AbstractNonlinearSolveBaseAPI end # Mostly used for pretty-printing + +function Base.show(io::IO, ::MIME"text/plain", alg::AbstractNonlinearSolveBaseAPI) + main_name = nameof(typeof(alg)) + modifiers = String[] + for field in fieldnames(typeof(alg)) + val = getfield(alg, field) + Utils.is_default_value(val, field, getfield(alg, field)) && continue + push!(modifiers, "$(field) = $(val)") + end + print(io, "$(main_name)($(join(modifiers, ", ")))") + return +end + +""" + AbstractDescentDirection + +Abstract Type for all Descent Directions used in NonlinearSolveBase. Given the Jacobian +`J` and the residual `fu`, these algorithms compute the descent direction `δu`. + +For non-square Jacobian problems, if we need to solve a linear solve problem, we use a +least squares solver by default, unless the provided `linsolve` can't handle non-square +matrices, in which case we use the normal form equations ``JᵀJ δu = Jᵀ fu``. Note that +this factorization is often the faster choice, but it is not as numerically stable as +the least squares solver. + +### `InternalAPI.init` specification + +```julia +InternalAPI.init( + prob::AbstractNonlinearProblem, alg::AbstractDescentDirection, J, fu, u; + pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, alias_J::Bool = true, + shared::Val = Val(1), kwargs... +)::AbstractDescentCache +``` + + - `pre_inverted`: whether or not the Jacobian has been pre_inverted. + - `linsolve_kwargs`: keyword arguments to pass to the linear solver. + - `abstol`: absolute tolerance for the linear solver. + - `reltol`: relative tolerance for the linear solver. + - `alias_J`: whether or not to alias the Jacobian. + - `shared`: Store multiple descent directions in the cache. Allows efficient and + correct reuse of factorizations if needed. + +Some of the algorithms also allow additional keyword arguments. See the documentation for +the specific algorithm for more information. + +### Interface Functions + + - `supports_trust_region(alg)`: whether or not the algorithm supports trust region + methods. Defaults to `false`. + - `supports_line_search(alg)`: whether or not the algorithm supports line search + methods. Defaults to `false`. + +See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), +[`DampedNewtonDescent`](@ref). +""" +abstract type AbstractDescentDirection <: AbstractNonlinearSolveBaseAPI end supports_line_search(::AbstractDescentDirection) = false supports_trust_region(::AbstractDescentDirection) = false @@ -15,7 +73,46 @@ function get_linear_solver(alg::AbstractDescentDirection) return Utils.safe_getproperty(alg, Val(:linsolve)) end -abstract type AbstractDescentCache end +""" + AbstractDescentCache + +Abstract Type for all Descent Caches. + +### `InternalAPI.solve!` specification + +```julia +InternalAPI.solve!( + cache::AbstractDescentCache, J, fu, u, idx::Val; + skip_solve::Bool = false, new_jacobian::Bool = true, kwargs... +)::DescentResult +``` + + - `J`: Jacobian or Inverse Jacobian (if `pre_inverted = Val(true)`). + - `fu`: residual. + - `u`: current state. + - `idx`: index of the descent problem to solve and return. Defaults to `Val(1)`. + - `skip_solve`: Skip the direction computation and return the previous direction. + Defaults to `false`. This is useful for Trust Region Methods where the previous + direction was rejected and we want to try with a modified trust region. + - `new_jacobian`: Whether the Jacobian has been updated. Defaults to `true`. + - `kwargs`: keyword arguments to pass to the linear solver if there is one. + +#### Returned values + + - `descent_result`: Result in a [`DescentResult`](@ref). + +### Interface Functions + + - `get_du(cache)`: get the descent direction. + - `get_du(cache, ::Val{N})`: get the `N`th descent direction. + - `set_du!(cache, δu)`: set the descent direction. + - `set_du!(cache, δu, ::Val{N})`: set the `N`th descent direction. + - `last_step_accepted(cache)`: whether or not the last step was accepted. Checks if the + cache has a `last_step_accepted` field and returns it if it does, else returns `true`. + - `preinverted_jacobian(cache)`: whether or not the Jacobian has been preinverted. + - `normal_form(cache)`: whether or not the linear solver uses normal form. +""" +abstract type AbstractDescentCache <: AbstractNonlinearSolveBaseAPI end SciMLBase.get_du(cache::AbstractDescentCache) = cache.δu SciMLBase.get_du(cache::AbstractDescentCache, ::Val{1}) = SciMLBase.get_du(cache) @@ -29,6 +126,79 @@ function last_step_accepted(cache::AbstractDescentCache) return true end +for fname in (:preinverted_jacobian, :normal_form) + @eval function $(fname)(alg::AbstractDescentCache) + res = Utils.unwrap_val(Utils.safe_getproperty(alg, Val($(QuoteNode(fname))))) + res === missing && return false + return res + end +end + +""" + AbstractDampingFunction + +Abstract Type for Damping Functions in DampedNewton. + +### `InternalAPI.init` specification + +```julia +InternalAPI.init( + prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, + J, fu, u, args...; + internalnorm::F = L2_NORM, kwargs... +)::AbstractDampingFunctionCache +``` + +Returns a [`NonlinearSolveBase.AbstractDampingFunctionCache`](@ref). +""" +abstract type AbstractDampingFunction <: AbstractNonlinearAlgorithm end + +""" + AbstractDampingFunctionCache + +Abstract Type for the Caches created by AbstractDampingFunctions + +### Interface Functions + + - `requires_normal_form_jacobian(alg)`: whether or not the Jacobian is needed in normal + form. No default. + - `requires_normal_form_rhs(alg)`: whether or not the residual is needed in normal form. + No default. + - `returns_norm_form_damping(alg)`: whether or not the damping function returns the + damping factor in normal form. Defaults to + `requires_normal_form_jacobian(alg) || requires_normal_form_rhs(alg)`. + - `(cache::AbstractDampingFunctionCache)(::Nothing)`: returns the damping factor. The type + of the damping factor returned from `solve!` is guaranteed to be the same as this. + +### `InternalAPI.solve!` specification + +```julia +InternalAPI.solve!( + cache::AbstractDampingFunctionCache, J, fu, u, δu, descent_stats +) +``` + +Returns the damping factor. +""" +abstract type AbstractDampingFunctionCache <: AbstractNonlinearAlgorithm end + +function requires_normal_form_jacobian end +function requires_normal_form_rhs end +function returns_norm_form_damping(f::F) where {F} + return requires_normal_form_jacobian(f) || requires_normal_form_rhs(f) +end + +""" + AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm + +Abstract Type for all NonlinearSolveBase Algorithms. + +### Interface Functions + + - `concrete_jac(alg)`: whether or not the algorithm uses a concrete Jacobian. Defaults + to `nothing`. + - `get_name(alg)`: get the name of the algorithm. +""" abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end get_name(alg::AbstractNonlinearSolveAlgorithm) = Utils.safe_getproperty(alg, Val(:name)) @@ -47,20 +217,20 @@ concrete_jac(v::Bool) = v concrete_jac(::Val{false}) = false concrete_jac(::Val{true}) = true -abstract type AbstractNonlinearSolveCache end +abstract type AbstractNonlinearSolveCache <: AbstractNonlinearSolveBaseAPI end """ AbstractLinearSolverCache -Abstract Type for all Linear Solvers used in NonlinearSolve. Subtypes of these are +Abstract Type for all Linear Solvers used in NonlinearSolveBase. Subtypes of these are meant to be constructured via [`construct_linear_solver`](@ref). """ -abstract type AbstractLinearSolverCache end +abstract type AbstractLinearSolverCache <: AbstractNonlinearSolveBaseAPI end """ AbstractJacobianCache -Abstract Type for all Jacobian Caches used in NonlinearSolve. Subtypes of these are +Abstract Type for all Jacobian Caches used in NonlinearSolveBase. Subtypes of these are meant to be constructured via [`construct_jacobian_cache`](@ref). """ -abstract type AbstractJacobianCache end +abstract type AbstractJacobianCache <: AbstractNonlinearSolveBaseAPI end diff --git a/src/descent/common.jl b/lib/NonlinearSolveBase/src/descent/common.jl similarity index 74% rename from src/descent/common.jl rename to lib/NonlinearSolveBase/src/descent/common.jl index 3d53573ea..b757a990a 100644 --- a/src/descent/common.jl +++ b/lib/NonlinearSolveBase/src/descent/common.jl @@ -1,6 +1,8 @@ """ - DescentResult(; δu = missing, u = missing, success::Bool = true, - linsolve_success::Bool = true, extras = (;)) + DescentResult(; + δu = missing, u = missing, success::Bool = true, linsolve_success::Bool = true, + extras = (;) + ) Construct a `DescentResult` object. @@ -23,8 +25,10 @@ Construct a `DescentResult` object. extras end -function DescentResult(; δu = missing, u = missing, success::Bool = true, - linsolve_success::Bool = true, extras = (;)) +function DescentResult(; + δu = missing, u = missing, success::Bool = true, linsolve_success::Bool = true, + extras = (;) +) @assert δu !== missing || u !== missing return DescentResult(δu, u, success, linsolve_success, extras) end diff --git a/lib/NonlinearSolveBase/src/descent/damped_newton.jl b/lib/NonlinearSolveBase/src/descent/damped_newton.jl new file mode 100644 index 000000000..3ab507065 --- /dev/null +++ b/lib/NonlinearSolveBase/src/descent/damped_newton.jl @@ -0,0 +1,268 @@ +""" + DampedNewtonDescent(; + linsolve = nothing, precs = nothing, initial_damping, damping_fn + ) + +A Newton descent algorithm with damping. The damping factor is computed using the +`damping_fn` function. The descent direction is computed as ``(JᵀJ + λDᵀD) δu = -fu``. For +non-square Jacobians, we default to solving for `Jδx = -fu` and `√λ⋅D δx = 0` +simultaneously. If the linear solver can't handle non-square matrices, we use the normal +form equations ``(JᵀJ + λDᵀD) δu = Jᵀ fu``. Note that this factorization is often the faster +choice, but it is not as numerically stable as the least squares solver. + +The damping factor returned must be a non-negative number. + +### Keyword Arguments + + - `initial_damping`: the initial damping factor to use + - `damping_fn`: the function to use to compute the damping factor. This must satisfy the + [`NonlinearSolveBase.AbstractDampingFunction`](@ref) interface. +""" +@kwdef @concrete struct DampedNewtonDescent <: AbstractDescentDirection + linsolve = nothing + precs = nothing + initial_damping + damping_fn <: AbstractDampingFunction +end + +supports_line_search(::DampedNewtonDescent) = true +supports_trust_region(::DampedNewtonDescent) = true + +@concrete mutable struct DampedNewtonDescentCache <: AbstractDescentCache + J + δu + δus + lincache + JᵀJ_cache + Jᵀfu_cache + rhs_cache + damping_fn_cache + timer + preinverted_jacobian <: Union{Val{false}, Val{true}} + mode <: Union{Val{:normal_form}, Val{:least_squares}, Val{:simple}} +end + +# XXX: Implement +# @internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, u; stats, + pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + timer = get_timer_output(), + alias_J::Bool = true, shared::Val = Val(1), + kwargs... +) + length(fu) != length(u) && + @assert pre_inverted isa Val{false} "Precomputed Inverse for Non-Square Jacobian doesn't make sense." + + @bb δu = similar(u) + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i + @bb δu_ = similar(u) + end + + normal_form_damping = returns_norm_form_damping(alg.damping_fn) + normal_form_linsolve = needs_square_A(alg.linsolve, u) + + mode = if u isa Number + :simple + elseif prob isa NonlinearProblem + if normal_form_damping + ifelse(normal_form_linsolve, :normal_form, :least_squares) + else + :simple + end + else + if normal_form_linsolve & !normal_form_damping + throw(ArgumentError("Linear Solver expects Normal Form but returned Damping is \ + not Normal Form. This is not supported.")) + end + if normal_form_damping & !normal_form_linsolve + :least_squares + else + ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form) + end + end + + if mode === :least_squares + if requires_normal_form_jacobian(alg.damping_fn) + JᵀJ = transpose(J) * J # Needed to compute the damping factor + jac_damp = JᵀJ + else + JᵀJ = nothing + jac_damp = J + end + if requires_normal_form_rhs(alg.damping_fn) + Jᵀfu = transpose(J) * Utils.safe_vec(fu) + rhs_damp = Jᵀfu + else + Jᵀfu = nothing + rhs_damp = fu + end + + damping_fn_cache = InternalAPI.init( + prob, alg.damping_fn, alg.initial_damping, jac_damp, rhs_damp, u, Val(false); + stats, kwargs... + ) + D = damping_fn_cache(nothing) + + D isa Number && (D = D * LinearAlgebra.I) + rhs_cache = vcat(Utils.safe_vec(fu), Utils.safe_vec(u)) + J_cache = Utils.faster_vcat(J, D) + A, b = J_cache, rhs_cache + elseif mode === :simple + damping_fn_cache = InternalAPI.init( + prob, alg.damping_fn, alg.initial_damping, J, fu, u, Val(false); kwargs... + ) + J_cache = Utils.maybe_unaliased(J, alias_J) + D = damping_fn_cache(nothing) + + J_damped = dampen_jacobian!!(J_cache, J, D) + J_cache = J_damped + A, b = J_damped, Utils.safe_vec(fu) + JᵀJ, Jᵀfu, rhs_cache = nothing, nothing, nothing + elseif mode === :normal_form + JᵀJ = transpose(J) * J + Jᵀfu = transpose(J) * Utils.safe_vec(fu) + jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J + rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu + + damping_fn_cache = InternalAPI.init( + prob, alg.damping_fn, alg.initial_damping, jac_damp, rhs_damp, u, Val(true); + stats, kwargs... + ) + D = damping_fn_cache(nothing) + + @bb J_cache = similar(JᵀJ) + @bb @. J_cache = 0 + J_damped = dampen_jacobian!!(J_cache, JᵀJ, D) + A, b = Utils.maybe_symmetric(J_damped), Utils.safe_vec(Jᵀfu) + rhs_cache = nothing + end + + lincache = construct_linear_solver( + alg, alg.linsolve, A, b, Utils.safe_vec(u); + stats, abstol, reltol, linsolve_kwargs... + ) + + return DampedNewtonDescentCache( + J_cache, δu, δus, lincache, JᵀJ, Jᵀfu, rhs_cache, + damping_fn_cache, timer, pre_inverted, Val(mode) + ) +end + +function InternalAPI.solve!( + cache::DampedNewtonDescentCache, J, fu, u, idx::Val = Val(1); + skip_solve::Bool = false, new_jacobian::Bool = true, kwargs... +) + δu = SciMLBase.get_du(cache, idx) + skip_solve && return DescentResult(; δu) + + recompute_A = idx === Val(1) + + @static_timeit cache.timer "dampen" begin + if cache.mode isa Val{:least_squares} + if (J !== nothing || new_jacobian) && recompute_A + preinverted_jacobian(cache) && (J = inv(J)) + if requires_normal_form_jacobian(cache.damping_fn_cache) + @bb cache.JᵀJ_cache = transpose(J) × J + jac_damp = cache.JᵀJ_cache + else + jac_damp = J + end + if requires_normal_form_rhs(cache.damping_fn_cache) + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + rhs_damp = cache.Jᵀfu_cache + else + rhs_damp = fu + end + D = InternalAPI.solve!( + cache.damping_fn_cache, jac_damp, rhs_damp, Val(false) + ) + if Utils.can_setindex(cache.J) + copyto!(@view(cache.J[1:size(J, 1), :]), J) + cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) + else + cache.J = Utils.faster_vcat(J, sqrt.(D)) + end + end + A = cache.J + if Utils.can_setindex(cache.rhs_cache) + cache.rhs_cache[1:length(fu)] .= Utils.safe_vec(fu) + cache.rhs_cache[(length(fu) + 1):end] .= false + else + cache.rhs_cache = vcat(Utils.safe_vec(fu), zero(Utils.safe_vec(u))) + end + b = cache.rhs_cache + elseif cache.mode isa Val{:simple} + if (J !== nothing || new_jacobian) && recompute_A + preinverted_jacobian(cache) && (J = inv(J)) + D = InternalAPI.solve!(cache.damping_fn_cache, J, fu, Val(false)) + cache.J = dampen_jacobian!!(cache.J, J, D) + end + A, b = cache.J, Utils.safe_vec(fu) + elseif cache.mode isa Val{:normal_form} + if (J !== nothing || new_jacobian) && recompute_A + preinverted_jacobian(cache) && (J = inv(J)) + @bb cache.JᵀJ_cache = transpose(J) × J + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + D = InternalAPI.solve!( + cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, Val(true) + ) + cache.J = dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) + A = Utils.maybe_symmetric(cache.J) + elseif !recompute_A + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + A = Utils.maybe_symmetric(cache.J) + else + A = nothing + end + b = cache.Jᵀfu_cache + else + error("Unknown Mode: $(cache.mode).") + end + end + + @static_timeit cache.timer "linear solve" begin + linres = cache.lincache(; + A, b, + reuse_A_if_factorization = !new_jacobian && !recompute_A, + kwargs..., + linu = Utils.safe_vec(δu) + ) + δu = Utils.restructure(SciMLBase.get_du(cache, idx), linres.u) + if !linres.success + set_du!(cache, δu, idx) + return DescentResult(; δu, success = false, linsolve_success = false) + end + end + + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return DescentResult(; δu) +end + +dampen_jacobian!!(::Any, J::Union{AbstractSciMLOperator, Number}, D) = J + D +function dampen_jacobian!!(J_cache, J::AbstractMatrix, D::Union{AbstractMatrix, Number}) + ArrayInterface.can_setindex(J_cache) || return J .+ D + J_cache !== J && copyto!(J_cache, J) + if ArrayInterface.fast_scalar_indexing(J_cache) + if D isa Number + @simd ivdep for i in axes(J_cache, 1) + @inbounds J_cache[i, i] += D + end + else + @simd ivdep for i in axes(J_cache, 1) + @inbounds J_cache[i, i] += D[i, i] + end + end + else + idxs = diagind(J_cache) + if D isa Number + J_cache[idxs] .+= D + else + J_cache[idxs] .+= @view(D[idxs]) + end + end + return J_cache +end diff --git a/src/descent/dogleg.jl b/lib/NonlinearSolveBase/src/descent/dogleg.jl similarity index 56% rename from src/descent/dogleg.jl rename to lib/NonlinearSolveBase/src/descent/dogleg.jl index 679f33c26..c138adbde 100644 --- a/src/descent/dogleg.jl +++ b/lib/NonlinearSolveBase/src/descent/dogleg.jl @@ -1,5 +1,5 @@ """ - Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS) + Dogleg(; linsolve = nothing, precs = nothing) Switch between Newton's method and the steepest descent method depending on the size of the trust region. The trust region is specified via keyword argument `trust_region` to @@ -7,82 +7,94 @@ trust region. The trust region is specified via keyword argument `trust_region` See also [`SteepestDescent`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). """ -@concrete struct Dogleg <: AbstractDescentAlgorithm - newton_descent - steepest_descent -end - -function Base.show(io::IO, d::Dogleg) - print(io, - "Dogleg(newton_descent = $(d.newton_descent), steepest_descent = $(d.steepest_descent))") +@concrete struct Dogleg <: AbstractDescentDirection + newton_descent <: Union{NewtonDescent, DampedNewtonDescent} + steepest_descent <: SteepestDescent end supports_trust_region(::Dogleg) = true get_linear_solver(alg::Dogleg) = get_linear_solver(alg.newton_descent) -function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, damping = False, +function Dogleg(; linsolve = nothing, precs = nothing, damping = Val(false), damping_fn = missing, initial_damping = missing, kwargs...) - if damping === False + if !Utils.unwrap_val(damping) return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent(; linsolve, precs)) end if damping_fn === missing || initial_damping === missing throw(ArgumentError("`damping_fn` and `initial_damping` must be supplied if \ `damping = Val(true)`.")) end - return Dogleg(DampedNewtonDescent(; linsolve, precs, damping_fn, initial_damping), - SteepestDescent(; linsolve, precs)) + return Dogleg( + DampedNewtonDescent(; linsolve, precs, damping_fn, initial_damping), + SteepestDescent(; linsolve, precs) + ) end -@concrete mutable struct DoglegCache{pre_inverted, normalform} <: AbstractDescentCache +@concrete mutable struct DoglegCache <: AbstractDescentCache δu δus - newton_cache - cauchy_cache + newton_cache <: Union{NewtonDescentCache, DampedNewtonDescentCache} + cauchy_cache <: SteepestDescentCache internalnorm - JᵀJ_cache + Jᵀδu_cache δu_cache_1 δu_cache_2 δu_cache_mul + preinverted_jacobian <: Union{Val{false}, Val{true}} + normal_form <: Union{Val{false}, Val{true}} end -@internal_caches DoglegCache :newton_cache :cauchy_cache +# XXX: Implement +# @internal_caches DoglegCache :newton_cache :cauchy_cache -function __internal_init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; - pre_inverted::Val{INV} = False, linsolve_kwargs = (;), +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; + pre_inverted::Val = Val(false), linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = L2_NORM, - shared::Val{N} = Val(1), kwargs...) where {F, INV, N} - newton_cache = __internal_init(prob, alg.newton_descent, J, fu, u; pre_inverted, - linsolve_kwargs, abstol, reltol, shared, kwargs...) - cauchy_cache = __internal_init(prob, alg.steepest_descent, J, fu, u; pre_inverted, - linsolve_kwargs, abstol, reltol, shared, kwargs...) + shared::Val = Val(1), kwargs... +) where {F} + newton_cache = InternalAPI.init( + prob, alg.newton_descent, J, fu, u; + pre_inverted, linsolve_kwargs, abstol, reltol, shared, kwargs... + ) + cauchy_cache = InternalAPI.init( + prob, alg.steepest_descent, J, fu, u; + pre_inverted, linsolve_kwargs, abstol, reltol, shared, kwargs... + ) + @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i @bb δu_ = similar(u) end @bb δu_cache_1 = similar(u) @bb δu_cache_2 = similar(u) @bb δu_cache_mul = similar(u) - T = promote_type(eltype(u), eltype(fu)) - normal_form = prob isa NonlinearLeastSquaresProblem && - NonlinearSolveBase.needs_square_A(alg.newton_descent.linsolve, u) - JᵀJ_cache = !normal_form ? J * _vec(δu) : nothing # TODO: Rename + needs_square_A(alg.newton_descent.linsolve, u) + + Jᵀδu_cache = !normal_form ? J * Utils.safe_vec(δu) : nothing - return DoglegCache{INV, normal_form}(δu, δus, newton_cache, cauchy_cache, internalnorm, - JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul) + return DoglegCache( + δu, δus, newton_cache, cauchy_cache, internalnorm, Jᵀδu_cache, + δu_cache_1, δu_cache_2, δu_cache_mul, pre_inverted, Val(normal_form) + ) end -# If TrustRegion is not specified, then use a Gauss-Newton step -function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = Val(1); - trust_region = nothing, skip_solve::Bool = false, kwargs...) where {INV, NF, N} +# If trust_region is not specified, then use a Gauss-Newton step +function InternalAPI.solve!( + cache::DoglegCache, J, fu, u, idx::Val = Val(1); + trust_region = nothing, skip_solve::Bool = false, kwargs... +) @assert trust_region!==nothing "Trust Region must be specified for Dogleg. Use \ `NewtonDescent` or `SteepestDescent` if you don't \ want to use a Trust Region." - δu = get_du(cache, idx) + + δu = SciMLBase.get_du(cache, idx) T = promote_type(eltype(u), eltype(fu)) - δu_newton = __internal_solve!( - cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...).δu + δu_newton = InternalAPI.solve!( + cache.newton_cache, J, fu, u, idx; skip_solve, kwargs... + ).δu # Newton's Step within the trust region if cache.internalnorm(δu_newton) ≤ trust_region @@ -91,23 +103,24 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = return DescentResult(; δu, extras = (; δuJᵀJδu = T(NaN))) end - # Take intersection of steepest descent direction and trust region if Cauchy point lies - # outside of trust region - if NF + # Take intersection of steepest descent direction and trust region if Cauchy point + # lies outside of trust region + if normal_form(cache) δu_cauchy = cache.newton_cache.Jᵀfu_cache JᵀJ = cache.newton_cache.JᵀJ_cache @bb @. δu_cauchy *= -1 l_grad = cache.internalnorm(δu_cauchy) @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) - δuJᵀJδu = __dot(δu_cauchy, cache.δu_cache_mul) + δuJᵀJδu = Utils.dot(cache.δu_cache_mul, cache.δu_cache_mul) else - δu_cauchy = __internal_solve!( - cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...).δu - J_ = INV ? inv(J) : J + δu_cauchy = InternalAPI.solve!( + cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs... + ).δu + J_ = preinverted_jacobian(cache) ? inv(J) : J l_grad = cache.internalnorm(δu_cauchy) - @bb cache.JᵀJ_cache = J × vec(δu_cauchy) # TODO: Rename - δuJᵀJδu = __dot(cache.JᵀJ_cache, cache.JᵀJ_cache) + @bb cache.Jᵀδu_cache = J_ × vec(δu_cauchy) + δuJᵀJδu = Utils.dot(cache.Jᵀδu_cache, cache.Jᵀδu_cache) end d_cauchy = (l_grad^3) / δuJᵀJδu @@ -125,8 +138,8 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = # trust region @bb @. cache.δu_cache_1 = (d_cauchy / l_grad) * δu_cauchy @bb @. cache.δu_cache_2 = δu_newton - cache.δu_cache_1 - a = dot(cache.δu_cache_2, cache.δu_cache_2) - b = 2 * dot(cache.δu_cache_1, cache.δu_cache_2) + a = Utils.safe_dot(cache.δu_cache_2, cache.δu_cache_2) + b = 2 * Utils.safe_dot(cache.δu_cache_1, cache.δu_cache_2) c = d_cauchy^2 - trust_region^2 aux = max(0, b^2 - 4 * a * c) τ = (-b + sqrt(aux)) / (2 * a) diff --git a/src/descent/geodesic_acceleration.jl b/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl similarity index 55% rename from src/descent/geodesic_acceleration.jl rename to lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl index 8e8a305f0..7b7645072 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl @@ -24,18 +24,12 @@ for other methods are not theorectically or experimentally verified. `α_geodesic = 0.1` is an effective choice. Defaults to `0.75`. See Section 3 of [transtrum2012improvements](@citet). """ -@concrete struct GeodesicAcceleration <: AbstractDescentAlgorithm +@concrete struct GeodesicAcceleration <: AbstractDescentDirection descent finite_diff_step_geodesic α end -function Base.show(io::IO, alg::GeodesicAcceleration) - print( - io, "GeodesicAcceleration(descent = $(alg.descent), finite_diff_step_geodesic = ", - "$(alg.finite_diff_step_geodesic), α = $(alg.α))") -end - supports_trust_region(::GeodesicAcceleration) = true get_linear_solver(alg::GeodesicAcceleration) = get_linear_solver(alg.descent) @@ -55,78 +49,86 @@ get_linear_solver(alg::GeodesicAcceleration) = get_linear_solver(alg.descent) last_step_accepted::Bool end -function __reinit_internal!( - cache::GeodesicAccelerationCache, args...; p = cache.p, kwargs...) - cache.p = p - cache.last_step_accepted = false -end +# XXX: Implement +# function __reinit_internal!( +# cache::GeodesicAccelerationCache, args...; p = cache.p, kwargs...) +# cache.p = p +# cache.last_step_accepted = false +# end -@internal_caches GeodesicAccelerationCache :descent_cache +# @internal_caches GeodesicAccelerationCache :descent_cache -get_velocity(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(1)) -function set_velocity!(cache::GeodesicAccelerationCache, δv) - set_du!(cache.descent_cache, δv, Val(1)) +function get_velocity(cache::GeodesicAccelerationCache) + return SciMLBase.get_du(cache.descent_cache, Val(1)) end function get_velocity(cache::GeodesicAccelerationCache, ::Val{N}) where {N} - get_du(cache.descent_cache, Val(2N - 1)) + return SciMLBase.get_du(cache.descent_cache, Val(2N - 1)) end -function set_velocity!(cache::GeodesicAccelerationCache, δv, ::Val{N}) where {N} - set_du!(cache.descent_cache, δv, Val(2N - 1)) -end -get_acceleration(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(2)) -function set_acceleration!(cache::GeodesicAccelerationCache, δa) - set_du!(cache.descent_cache, δa, Val(2)) +function get_acceleration(cache::GeodesicAccelerationCache) + return SciMLBase.get_du(cache.descent_cache, Val(2)) end function get_acceleration(cache::GeodesicAccelerationCache, ::Val{N}) where {N} - get_du(cache.descent_cache, Val(2N)) -end -function set_acceleration!(cache::GeodesicAccelerationCache, δa, ::Val{N}) where {N} - set_du!(cache.descent_cache, δa, Val(2N)) + return SciMLBase.get_du(cache.descent_cache, Val(2N)) end -function __internal_init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleration, J, - fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, - linsolve_kwargs = (;), abstol = nothing, reltol = nothing, - internalnorm::F = L2_NORM, kwargs...) where {INV, N, F} +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::GeodesicAcceleration, J, fu, u; + shared::Val = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + internalnorm::F = L2_NORM, kwargs... +) where {F} T = promote_type(eltype(u), eltype(fu)) @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i @bb δu_ = similar(u) end - descent_cache = __internal_init(prob, alg.descent, J, fu, u; shared = Val(N * 2), - pre_inverted, linsolve_kwargs, abstol, reltol, kwargs...) + descent_cache = InternalAPI.init( + prob, alg.descent, J, fu, u; + shared = Val(2 * Utils.unwrap_val(shared)), pre_inverted, linsolve_kwargs, + abstol, reltol, + kwargs... + ) @bb Jv = similar(fu) @bb fu_cache = copy(fu) @bb u_cache = similar(u) return GeodesicAccelerationCache( δu, δus, descent_cache, prob.f, prob.p, T(alg.α), internalnorm, - T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache, false) + T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache, false + ) end -function __internal_solve!( - cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N} = Val(1); - skip_solve::Bool = false, kwargs...) where {N} - a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx) +function InternalAPI.solve!( + cache::GeodesicAccelerationCache, J, fu, u, idx::Val = Val(1); + skip_solve::Bool = false, kwargs... +) + a = get_acceleration(cache, idx) + v = get_velocity(cache, idx) + δu = SciMLBase.get_du(cache, idx) skip_solve && return DescentResult(; δu, extras = (; a, v)) - v = __internal_solve!( - cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...).δu + + v = InternalAPI.solve!( + cache.descent_cache, J, fu, u, Val(2 * Utils.unwrap_val(idx) - 1); + skip_solve, kwargs... + ).δu @bb @. cache.u_cache = u + cache.h * v - cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + cache.fu_cache = Utils.evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) J !== nothing && @bb(cache.Jv=J × vec(v)) - Jv = _restructure(cache.fu_cache, cache.Jv) + Jv = Utils.restructure(cache.fu_cache, cache.Jv) @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv) - a = __internal_solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N); - skip_solve, kwargs..., reuse_A_if_factorization = true).δu + a = InternalAPI.solve!( + cache.descent_cache, J, cache.fu_cache, u, Val(2 * Utils.unwrap_val(idx)); + skip_solve, kwargs..., reuse_A_if_factorization = true + ).δu norm_v = cache.internalnorm(v) norm_a = cache.internalnorm(a) if 2 * norm_a ≤ norm_v * cache.α @bb @. δu = v + a / 2 - set_du!(cache, δu, idx) + SciMLBase.set_du!(cache, δu, idx) cache.last_step_accepted = true else cache.last_step_accepted = false diff --git a/lib/NonlinearSolveBase/src/descent/newton.jl b/lib/NonlinearSolveBase/src/descent/newton.jl new file mode 100644 index 000000000..1a7acf177 --- /dev/null +++ b/lib/NonlinearSolveBase/src/descent/newton.jl @@ -0,0 +1,132 @@ +""" + NewtonDescent(; linsolve = nothing, precs = nothing) + +Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is +commonly referred to as the Gauss-Newton Descent. + +See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@ref). +""" +@kwdef @concrete struct NewtonDescent <: AbstractDescentDirection + linsolve = nothing + precs = nothing +end + +supports_line_search(::NewtonDescent) = true + +@concrete mutable struct NewtonDescentCache <: AbstractDescentCache + δu + δus + lincache + JᵀJ_cache # For normal form else nothing + Jᵀfu_cache + timer + preinverted_jacobian <: Union{Val{false}, Val{true}} + normal_form <: Union{Val{false}, Val{true}} +end + +# XXX: Implement +# @internal_caches NewtonDescentCache :lincache + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::NewtonDescent, J, fu, u; stats, + shared = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + timer = get_timer_output(), kwargs... +) + @bb δu = similar(u) + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i + @bb δu_ = similar(u) + end + if Utils.unwrap_val(pre_inverted) + lincache = nothing + else + lincache = construct_linear_solver( + alg, alg.linsolve, J, Utils.safe_vec(fu), Utils.safe_vec(u); + stats, abstol, reltol, linsolve_kwargs... + ) + end + return NewtonDescentCache( + δu, δus, lincache, nothing, nothing, timer, pre_inverted, Val(false) + ) +end + +function InternalAPI.init( + prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; stats, + shared = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + timer = get_timer_output(), kwargs... +) + length(fu) != length(u) && + @assert !Utils.unwrap_val(pre_inverted) "Precomputed Inverse for Non-Square Jacobian doesn't make sense." + + @bb δu = similar(u) + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end + + normal_form = needs_square_A(alg.linsolve, u) + if normal_form + JᵀJ = transpose(J) * J + Jᵀfu = transpose(J) * Utils.safe_vec(fu) + A, b = Utils.maybe_symmetric(JᵀJ), Jᵀfu + else + JᵀJ, Jᵀfu = nothing, nothing + A, b = J, Utils.safe_vec(fu) + end + + lincache = construct_linear_solver( + alg, alg.linsolve, A, b, Utils.safe_vec(u); + stats, abstol, reltol, linsolve_kwargs... + ) + + return NewtonDescentCache( + δu, δus, lincache, JᵀJ, Jᵀfu, timer, pre_inverted, Val(normal_form) + ) +end + +function InternalAPI.solve!( + cache::NewtonDescentCache, J, fu, u, idx::Val = Val(1); + skip_solve::Bool = false, new_jacobian::Bool = true, kwargs... +) + δu = SciMLBase.get_du(cache, idx) + skip_solve && return DescentResult(; δu) + + if preinverted_jacobian(cache) && !normal_form(cache) + @assert J!==nothing "`J` must be provided when `preinverted_jacobian = Val(true)`." + @bb δu = J × vec(fu) + else + if normal_form(cache) + @assert !preinverted_jacobian(cache) + if idx === Val(1) + @bb cache.JᵀJ_cache = transpose(J) × J + end + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + @static_timeit cache.timer "linear solve" begin + linres = cache.lincache(; + A = Utils.maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, + kwargs..., linu = Utils.safe_vec(δu), du = Utils.safe_vec(δu), + reuse_A_if_factorization = !new_jacobian || (idx !== Val(1)) + ) + end + else + @static_timeit cache.timer "linear solve" begin + linres = cache.lincache(; + A = J, b = Utils.safe_vec(fu), + kwargs..., linu = Utils.safe_vec(δu), du = Utils.safe_vec(δu), + reuse_A_if_factorization = !new_jacobian || idx !== Val(1) + ) + end + end + @static_timeit cache.timer "linear solve" begin + δu = Utils.restructure(SciMLBase.get_du(cache, idx), linres.u) + if !linres.success + set_du!(cache, δu, idx) + return DescentResult(; δu, success = false, linsolve_success = false) + end + end + end + + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return DescentResult(; δu) +end diff --git a/lib/NonlinearSolveBase/src/descent/steepest.jl b/lib/NonlinearSolveBase/src/descent/steepest.jl new file mode 100644 index 000000000..b29045bd5 --- /dev/null +++ b/lib/NonlinearSolveBase/src/descent/steepest.jl @@ -0,0 +1,76 @@ +""" + SteepestDescent(; linsolve = nothing, precs = nothing) + +Compute the descent direction as ``δu = -Jᵀfu``. The linear solver and preconditioner are +only used if `J` is provided in the inverted form. + +See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). +""" +@kwdef @concrete struct SteepestDescent <: AbstractDescentDirection + linsolve = nothing + precs = nothing +end + +supports_line_search(::SteepestDescent) = true + +@concrete mutable struct SteepestDescentCache <: AbstractDescentCache + δu + δus + lincache + timer + preinverted_jacobian <: Union{Val{false}, Val{true}} +end + +# XXX: Implement +# @internal_caches SteepestDescentCache :lincache + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; + stats, shared = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + timer = get_timer_output(), + kwargs... +) + if Utils.unwrap_val(pre_inverted) + @assert length(fu)==length(u) "Non-Square Jacobian Inverse doesn't make sense." + end + @bb δu = similar(u) + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i + @bb δu_ = similar(u) + end + if Utils.unwrap_val(pre_inverted) + lincache = construct_linear_solver( + alg, alg.linsolve, transpose(J), Utils.safe_vec(fu), Utils.safe_vec(u); + stats, abstol, reltol, linsolve_kwargs... + ) + else + lincache = nothing + end + return SteepestDescentCache(δu, δus, lincache, timer, pre_inverted) +end + +function InternalAPI.solve!( + cache::SteepestDescentCache, J, fu, u, idx::Val = Val(1); + new_jacobian::Bool = true, kwargs... +) + δu = SciMLBase.get_du(cache, idx) + if Utils.unwrap_val(cache.preinverted_jacobian) + A = J === nothing ? nothing : transpose(J) + linres = cache.lincache(; + A, b = Utils.safe_vec(fu), kwargs..., linu = Utils.safe_vec(δu), + du = Utils.safe_vec(δu), + reuse_A_if_factorization = !new_jacobian || idx !== Val(1) + ) + δu = Utils.restructure(SciMLBase.get_du(cache, idx), linres.u) + if !linres.success + set_du!(cache, δu, idx) + return DescentResult(; δu, success = false, linsolve_success = false) + end + else + @assert J!==nothing "`J` must be provided when `preinverted_jacobian = Val(false)`." + @bb δu = transpose(J) × vec(fu) + end + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return DescentResult(; δu) +end diff --git a/lib/NonlinearSolveBase/src/timer_outputs.jl b/lib/NonlinearSolveBase/src/timer_outputs.jl new file mode 100644 index 000000000..6a65f1bab --- /dev/null +++ b/lib/NonlinearSolveBase/src/timer_outputs.jl @@ -0,0 +1,33 @@ +# Timer Outputs has some overhead, so we only use it if we are debugging +# Even `@timeit` has overhead so we write our custom version of that using Preferences +const TIMER_OUTPUTS_ENABLED = @load_preference("enable_timer_outputs", false) + +@static if TIMER_OUTPUTS_ENABLED + using TimerOutputs: TimerOutput, timer_expr, reset_timer! +end + +function get_timer_output() + @static if TIMER_OUTPUTS_ENABLED + return TimerOutput() + else + return nothing + end +end + +""" + @static_timeit to name expr + +Like `TimerOutputs.@timeit_debug` but has zero overhead if `TimerOutputs` is disabled via +[`NonlinearSolve.disable_timer_outputs()`](@ref). +""" +macro static_timeit(to, name, expr) + @static if TIMER_OUTPUTS_ENABLED + return timer_expr(__module__, false, to, name, expr) + else + return esc(expr) + end +end + +@static if !TIMER_OUTPUTS_ENABLED + @inline reset_timer!(::Nothing) = nothing +end diff --git a/lib/NonlinearSolveBase/src/utils.jl b/lib/NonlinearSolveBase/src/utils.jl index 3f14de9ea..628a4e873 100644 --- a/lib/NonlinearSolveBase/src/utils.jl +++ b/lib/NonlinearSolveBase/src/utils.jl @@ -2,8 +2,11 @@ module Utils using ArrayInterface: ArrayInterface using FastClosures: @closure -using LinearAlgebra: norm +using LinearAlgebra: Symmetric, norm, dot using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition +using SciMLOperators: AbstractSciMLOperator +using SciMLBase: SciMLBase, AbstractNonlinearProblem, NonlinearFunction +using StaticArraysCore: StaticArray, SArray using ..NonlinearSolveBase: L2_NORM, Linf_NORM @@ -93,8 +96,52 @@ safe_reshape(x::Number, args...) = x safe_reshape(x, args...) = reshape(x, args...) @generated function safe_getproperty(s::S, ::Val{X}) where {S, X} - hasfield(S, X) && return :(getproperty(s, $(X))) + hasfield(S, X) && return :(getproperty(s, $(Meta.quot(X)))) return :(missing) end +@generated function safe_vec(v) + hasmethod(vec, Tuple{typeof(v)}) || return :(vec(v)) + return :(v) +end +safe_vec(v::Number) = v +safe_vec(v::AbstractVector) = v + +safe_dot(x, y) = dot(safe_vec(x), safe_vec(y)) + +unwrap_val(x) = x +unwrap_val(::Val{x}) where {x} = unwrap_val(x) + +is_default_value(::Any, ::Symbol, ::Nothing) = true +is_default_value(::Any, ::Symbol, ::Missing) = true +is_default_value(::Any, ::Symbol, ::Any) = false + +maybe_symmetric(x) = Symmetric(x) +maybe_symmetric(x::Number) = x +## LinearSolve with `nothing` doesn't dispatch correctly here +maybe_symmetric(x::StaticArray) = x # XXX: Can we remove this? +maybe_symmetric(x::AbstractSciMLOperator) = x + +# Define special concatenation for certain Array combinations +faster_vcat(x, y) = vcat(x, y) + +maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x +function maybe_unaliased(x::AbstractArray, alias::Bool) + (alias || !ArrayInterface.can_setindex(typeof(x))) && return x + return copy(x) +end +maybe_unaliased(x::AbstractSciMLOperator, ::Bool) = x + +can_setindex(x) = ArrayInterface.can_setindex(x) +can_setindex(::Number) = false + +evaluate_f!!(prob::AbstractNonlinearProblem, fu, u, p) = evaluate_f!!(prob.f, fu, u, p) +function evaluate_f!!(f::NonlinearFunction, fu, u, p) + if SciMLBase.isinplace(f) + f(fu, u, p) + return fu + end + return f(u, p) +end + end diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index b5bfa76bc..6a5af1004 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -18,20 +18,37 @@ using LineSearch: LineSearch, AbstractLineSearchCache, LineSearchesJL, NoLineSea using LinearSolve: LinearSolve using MaybeInplace: @bb using NonlinearSolveBase: NonlinearSolveBase, + # ForwardDiff integration support nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution, nonlinearsolve_∂f_∂p, nonlinearsolve_∂f_∂u, + # faster norms L2_NORM, + # termination conditions AbsNormTerminationMode, AbstractNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode, + # autodiff selection select_forward_mode_autodiff, select_reverse_mode_autodiff, select_jacobian_autodiff, - construct_linear_solver, construct_jacobian_cache + # helpers for constructing caches + construct_linear_solver, construct_jacobian_cache, + # Descent Directions + DescentResult, + SteepestDescent, NewtonDescent, DampedNewtonDescent, Dogleg, + GeodesicAcceleration, + # Timer Outputs + reset_timer!, @static_timeit + # XXX: Remove -import NonlinearSolveBase: concrete_jac +import NonlinearSolveBase: InternalAPI, concrete_jac, supports_line_search, + supports_trust_region, set_du!, last_step_accepted, + get_linear_solver, + AbstractDampingFunction, AbstractDampingFunctionCache, + requires_normal_form_jacobian, requires_normal_form_rhs, + returns_norm_form_damping, get_timer_output using Printf: @printf -using Preferences: Preferences, @load_preference, @set_preferences! +using Preferences: Preferences, set_preferences! using RecursiveArrayTools: recursivecopy! using SciMLBase: SciMLBase, AbstractNonlinearAlgorithm, AbstractNonlinearProblem, _unwrap_val, isinplace, NLStats, NonlinearFunction, @@ -66,13 +83,6 @@ include("abstract_types.jl") include("timer_outputs.jl") include("internal/helpers.jl") -include("descent/common.jl") -include("descent/newton.jl") -include("descent/steepest.jl") -include("descent/dogleg.jl") -include("descent/damped_newton.jl") -include("descent/geodesic_acceleration.jl") - include("internal/termination.jl") include("internal/tracing.jl") include("internal/approximate_initialization.jl") @@ -180,9 +190,6 @@ export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, NLSol # Advanced Algorithms -- Without Bells and Whistles export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, GeneralizedDFSane -# Descent Algorithms -export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, GeodesicAcceleration - # Globalization ## Line Search Algorithms export LineSearch, BackTracking, NoLineSearch, RobustNonMonotoneLineSearch, diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 0d60ed3bb..a94b2b50e 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -1,110 +1,5 @@ -function __internal_init end -function __internal_solve! end - -""" - AbstractDescentAlgorithm - -Given the Jacobian `J` and the residual `fu`, this type of algorithm computes the descent -direction `δu`. - -For non-square Jacobian problems, if we need to solve a linear solve problem, we use a least -squares solver by default, unless the provided `linsolve` can't handle non-square matrices, -in which case we use the normal form equations ``JᵀJ δu = Jᵀ fu``. Note that this -factorization is often the faster choice, but it is not as numerically stable as the least -squares solver. - -### `__internal_init` specification - -```julia -__internal_init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, - fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, alias_J::Bool = true, - shared::Val{N} = Val(1), kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache - -__internal_init( - prob::NonlinearLeastSquaresProblem{uType, iip}, alg::AbstractDescentAlgorithm, - J, fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, alias_J::Bool = true, - shared::Val{N} = Val(1), kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache -``` - - - `pre_inverted`: whether or not the Jacobian has been pre_inverted. Defaults to `False`. - Note that for most algorithms except `NewtonDescent` setting it to `Val(true)` is - generally a bad idea. - - `linsolve_kwargs`: keyword arguments to pass to the linear solver. Defaults to `(;)`. - - `abstol`: absolute tolerance for the linear solver. Defaults to `nothing`. - - `reltol`: relative tolerance for the linear solver. Defaults to `nothing`. - - `alias_J`: whether or not to alias the Jacobian. Defaults to `true`. - - `shared`: Store multiple descent directions in the cache. Allows efficient and correct - reuse of factorizations if needed, - -Some of the algorithms also allow additional keyword arguments. See the documentation for -the specific algorithm for more information. - -### Interface Functions - - - `supports_trust_region(alg)`: whether or not the algorithm supports trust region - methods. Defaults to `false`. - - `supports_line_search(alg)`: whether or not the algorithm supports line search - methods. Defaults to `false`. - -See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), -[`DampedNewtonDescent`](@ref). -""" -abstract type AbstractDescentAlgorithm end - -supports_trust_region(::AbstractDescentAlgorithm) = false -supports_line_search(::AbstractDescentAlgorithm) = false - -get_linear_solver(alg::AbstractDescentAlgorithm) = __getproperty(alg, Val(:linsolve)) - -""" - AbstractDescentCache - -Abstract Type for all Descent Caches. - -### `__internal_solve!` specification - -```julia -descent_result = __internal_solve!( - cache::AbstractDescentCache, J, fu, u, idx::Val; skip_solve::Bool = false, kwargs...) -``` - - - `J`: Jacobian or Inverse Jacobian (if `pre_inverted = Val(true)`). - - `fu`: residual. - - `u`: current state. - - `idx`: index of the descent problem to solve and return. Defaults to `Val(1)`. - - `skip_solve`: Skip the direction computation and return the previous direction. - Defaults to `false`. This is useful for Trust Region Methods where the previous - direction was rejected and we want to try with a modified trust region. - - `kwargs`: keyword arguments to pass to the linear solver if there is one. - -#### Returned values - - - `descent_result`: Result in a [`DescentResult`](@ref). - -### Interface Functions - - - `get_du(cache)`: get the descent direction. - - `get_du(cache, ::Val{N})`: get the `N`th descent direction. - - `set_du!(cache, δu)`: set the descent direction. - - `set_du!(cache, δu, ::Val{N})`: set the `N`th descent direction. - - `last_step_accepted(cache)`: whether or not the last step was accepted. Checks if the - cache has a `last_step_accepted` field and returns it if it does, else returns `true`. -""" -abstract type AbstractDescentCache end - -SciMLBase.get_du(cache::AbstractDescentCache) = cache.δu -SciMLBase.get_du(cache::AbstractDescentCache, ::Val{1}) = get_du(cache) -SciMLBase.get_du(cache::AbstractDescentCache, ::Val{N}) where {N} = cache.δus[N - 1] -set_du!(cache::AbstractDescentCache, δu) = (cache.δu = δu) -set_du!(cache::AbstractDescentCache, δu, ::Val{1}) = set_du!(cache, δu) -set_du!(cache::AbstractDescentCache, δu, ::Val{N}) where {N} = (cache.δus[N - 1] = δu) - -function last_step_accepted(cache::AbstractDescentCache) - hasfield(typeof(cache), :last_step_accepted) && return cache.last_step_accepted - return true -end +const __internal_init = InternalAPI.init +const __internal_solve! = InternalAPI.solve! """ AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorithm @@ -215,63 +110,7 @@ function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache, u0; kwargs...) return reinit_cache!(cache; u0, kwargs...) end -""" - AbstractDampingFunction - -Abstract Type for Damping Functions in DampedNewton. - -### `__internal_init` specification - -```julia -__internal_init( - prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, - J, fu, u, args...; internal_norm = L2_NORM, kwargs...) --> AbstractDampingFunctionCache -``` - -Returns a [`AbstractDampingFunctionCache`](@ref). -""" -abstract type AbstractDampingFunction end - -""" - AbstractDampingFunctionCache - -Abstract Type for the Caches created by AbstractDampingFunctions - -### Interface Functions - - - `requires_normal_form_jacobian(f)`: whether or not the Jacobian is needed in normal - form. No default. - - `requires_normal_form_rhs(f)`: whether or not the residual is needed in normal form. - No default. - - `returns_norm_form_damping(f)`: whether or not the damping function returns the - damping factor in normal form. Defaults to `requires_normal_form_jacobian(f) || requires_normal_form_rhs(f)`. - - `(cache::AbstractDampingFunctionCache)(::Nothing)`: returns the damping factor. The type - of the damping factor returned from `solve!` is guaranteed to be the same as this. - -### `__internal_solve!` specification - -```julia -__internal_solve!(cache::AbstractDampingFunctionCache, J, fu, args...; kwargs...) -``` - -Returns the damping factor. -""" -abstract type AbstractDampingFunctionCache end - -function requires_normal_form_jacobian end -function requires_normal_form_rhs end -function returns_norm_form_damping(f::F) where {F} - return requires_normal_form_jacobian(f) || requires_normal_form_rhs(f) -end - -""" - AbstractNonlinearSolveOperator <: AbstractSciMLOperator - -NonlinearSolve.jl houses a few custom operators. These will eventually be moved out but till -then this serves as the abstract type for them. -""" -abstract type AbstractNonlinearSolveOperator{T} <: AbstractSciMLOperator{T} end - +# XXX: Move to NonlinearSolveQuasiNewton # Approximate Jacobian Algorithms """ AbstractApproximateJacobianStructure @@ -430,15 +269,6 @@ abstract type AbstractTrustRegionMethodCache end last_step_accepted(cache::AbstractTrustRegionMethodCache) = cache.last_step_accepted -""" - AbstractNonlinearSolveJacobianCache{iip} <: Function - -Abstract Type for all Jacobian Caches used in NonlinearSolve.jl. -""" -abstract type AbstractNonlinearSolveJacobianCache{iip} <: Function end - -SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = iip - """ AbstractNonlinearSolveTraceLevel @@ -454,12 +284,3 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii `store_trace == Val(true)`. """ abstract type AbstractNonlinearSolveTraceLevel end - -# Default Printing -for aType in (AbstractTrustRegionMethod, AbstractResetCondition, - AbstractApproximateJacobianUpdateRule, AbstractDampingFunction, - AbstractNonlinearSolveExtensionAlgorithm) - @eval function Base.show(io::IO, alg::$(aType)) - print(io, "$(nameof(typeof(alg)))()") - end -end diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl index b6ce5f9b9..bbbc55b28 100644 --- a/src/algorithms/lbroyden.jl +++ b/src/algorithms/lbroyden.jl @@ -72,7 +72,7 @@ end Low Rank Approximation of the Jacobian Matrix. Currently only used for [`LimitedMemoryBroyden`](@ref). This computes the Jacobian as ``U \\times V^T``. """ -@concrete mutable struct BroydenLowRankJacobian{T} <: AbstractNonlinearSolveOperator{T} +@concrete mutable struct BroydenLowRankJacobian{T} <: AbstractSciMLOperator{T} U Vᵀ idx::Int diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl deleted file mode 100644 index f1d2996c9..000000000 --- a/src/descent/damped_newton.jl +++ /dev/null @@ -1,256 +0,0 @@ -""" - DampedNewtonDescent(; linsolve = nothing, precs = DEFAULT_PRECS, initial_damping, - damping_fn) - -A Newton descent algorithm with damping. The damping factor is computed using the -`damping_fn` function. The descent direction is computed as ``(JᵀJ + λDᵀD) δu = -fu``. For -non-square Jacobians, we default to solving for `Jδx = -fu` and `√λ⋅D δx = 0` -simultaneously. If the linear solver can't handle non-square matrices, we use the normal -form equations ``(JᵀJ + λDᵀD) δu = Jᵀ fu``. Note that this factorization is often the faster -choice, but it is not as numerically stable as the least squares solver. - -The damping factor returned must be a non-negative number. - -### Keyword Arguments - - - `initial_damping`: The initial damping factor to use - - `damping_fn`: The function to use to compute the damping factor. This must satisfy the - [`NonlinearSolve.AbstractDampingFunction`](@ref) interface. -""" -@kwdef @concrete struct DampedNewtonDescent <: AbstractDescentAlgorithm - linsolve = nothing - precs = DEFAULT_PRECS - initial_damping - damping_fn -end - -function Base.show(io::IO, d::DampedNewtonDescent) - modifiers = String[] - d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") - d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") - push!(modifiers, "initial_damping = $(d.initial_damping)") - push!(modifiers, "damping_fn = $(d.damping_fn)") - print(io, "DampedNewtonDescent($(join(modifiers, ", ")))") -end - -supports_line_search(::DampedNewtonDescent) = true -supports_trust_region(::DampedNewtonDescent) = true - -@concrete mutable struct DampedNewtonDescentCache{pre_inverted, mode} <: - AbstractDescentCache - J - δu - δus - lincache - JᵀJ_cache - Jᵀfu_cache - rhs_cache - damping_fn_cache - timer -end - -@internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache - -function __internal_init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, - u; stats, pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - abstol = nothing, timer = get_timer_output(), reltol = nothing, - alias_J = true, shared::Val{N} = Val(1), kwargs...) where {INV, N} - length(fu) != length(u) && - @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - - normal_form_damping = returns_norm_form_damping(alg.damping_fn) - normal_form_linsolve = NonlinearSolveBase.needs_square_A(alg.linsolve, u) - if u isa Number - mode = :simple - elseif prob isa NonlinearProblem - mode = ifelse(!normal_form_damping, :simple, - ifelse(normal_form_linsolve, :normal_form, :least_squares)) - else - if normal_form_linsolve & !normal_form_damping - throw(ArgumentError("Linear Solver expects Normal Form but returned Damping is \ - not Normal Form. This is not supported.")) - end - mode = ifelse(normal_form_damping & !normal_form_linsolve, :least_squares, - ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form)) - end - - if mode === :least_squares - if requires_normal_form_jacobian(alg.damping_fn) - JᵀJ = transpose(J) * J # Needed to compute the damping factor - jac_damp = JᵀJ - else - JᵀJ = nothing - jac_damp = J - end - if requires_normal_form_rhs(alg.damping_fn) - Jᵀfu = transpose(J) * _vec(fu) - rhs_damp = Jᵀfu - else - Jᵀfu = nothing - rhs_damp = fu - end - damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, - jac_damp, rhs_damp, u, False; stats, kwargs...) - D = damping_fn_cache(nothing) - D isa Number && (D = D * I) - rhs_cache = vcat(_vec(fu), _vec(u)) - J_cache = _vcat(J, D) - A, b = J_cache, rhs_cache - elseif mode === :simple - damping_fn_cache = __internal_init( - prob, alg.damping_fn, alg.initial_damping, J, fu, u, False; kwargs...) - J_cache = __maybe_unaliased(J, alias_J) - D = damping_fn_cache(nothing) - J_damped = __dampen_jacobian!!(J_cache, J, D) - J_cache = J_damped - A, b = J_damped, _vec(fu) - JᵀJ, Jᵀfu, rhs_cache = nothing, nothing, nothing - elseif mode === :normal_form - JᵀJ = transpose(J) * J - Jᵀfu = transpose(J) * _vec(fu) - jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J - rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu - damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, - jac_damp, rhs_damp, u, True; stats, kwargs...) - D = damping_fn_cache(nothing) - @bb J_cache = similar(JᵀJ) - @bb @. J_cache = 0 - J_damped = __dampen_jacobian!!(J_cache, JᵀJ, D) - A, b = __maybe_symmetric(J_damped), _vec(Jᵀfu) - rhs_cache = nothing - end - - lincache = construct_linear_solver( - alg, alg.linsolve, A, b, _vec(u); stats, abstol, reltol, linsolve_kwargs...) - - return DampedNewtonDescentCache{INV, mode}( - J_cache, δu, δus, lincache, JᵀJ, Jᵀfu, rhs_cache, damping_fn_cache, timer) -end - -function __internal_solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, - u, idx::Val{N} = Val(1); skip_solve::Bool = false, - new_jacobian::Bool = true, kwargs...) where {INV, N, mode} - δu = get_du(cache, idx) - skip_solve && return DescentResult(; δu) - - recompute_A = idx === Val(1) - - @static_timeit cache.timer "dampen" begin - if mode === :least_squares - if (J !== nothing || new_jacobian) && recompute_A - INV && (J = inv(J)) - if requires_normal_form_jacobian(cache.damping_fn_cache) - @bb cache.JᵀJ_cache = transpose(J) × J - jac_damp = cache.JᵀJ_cache - else - jac_damp = J - end - if requires_normal_form_rhs(cache.damping_fn_cache) - @bb cache.Jᵀfu_cache = transpose(J) × fu - rhs_damp = cache.Jᵀfu_cache - else - rhs_damp = fu - end - D = __internal_solve!(cache.damping_fn_cache, jac_damp, rhs_damp, False) - if __can_setindex(cache.J) - copyto!(@view(cache.J[1:size(J, 1), :]), J) - cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) - else - cache.J = _vcat(J, sqrt.(D)) - end - end - A = cache.J - if __can_setindex(cache.rhs_cache) - cache.rhs_cache[1:length(fu)] .= _vec(fu) - cache.rhs_cache[(length(fu) + 1):end] .= false - else - cache.rhs_cache = vcat(_vec(fu), zero(_vec(u))) - end - b = cache.rhs_cache - elseif mode === :simple - if (J !== nothing || new_jacobian) && recompute_A - INV && (J = inv(J)) - D = __internal_solve!(cache.damping_fn_cache, J, fu, False) - cache.J = __dampen_jacobian!!(cache.J, J, D) - end - A, b = cache.J, _vec(fu) - elseif mode === :normal_form - if (J !== nothing || new_jacobian) && recompute_A - INV && (J = inv(J)) - @bb cache.JᵀJ_cache = transpose(J) × J - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - D = __internal_solve!( - cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) - cache.J = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) - A = __maybe_symmetric(cache.J) - elseif !recompute_A - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - A = __maybe_symmetric(cache.J) - else - A = nothing - end - b = _vec(cache.Jᵀfu_cache) - else - error("Unknown mode: $(mode)") - end - end - - @static_timeit cache.timer "linear solve" begin - linres = cache.lincache(; - A, b, reuse_A_if_factorization = !new_jacobian && !recompute_A, - kwargs..., linu = _vec(δu)) - δu = _restructure(get_du(cache, idx), linres.u) - if !linres.success - set_du!(cache, δu, idx) - return DescentResult(; δu, success = false, linsolve_success = false) - end - end - - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return DescentResult(; δu) -end - -# Define special concatenation for certain Array combinations -@inline _vcat(x, y) = vcat(x, y) - -# J_cache is allowed to alias J -## Compute ``J + D`` -@inline __dampen_jacobian!!(J_cache, J::AbstractSciMLOperator, D) = J + D -@inline __dampen_jacobian!!(J_cache, J::Number, D) = J + D -@inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::AbstractMatrix) - if __can_setindex(J_cache) - copyto!(J_cache, J) - if fast_scalar_indexing(J_cache) - @simd for i in axes(J_cache, 1) - @inbounds J_cache[i, i] += D[i, i] - end - else - idxs = diagind(J_cache) - J_cache[idxs] .= @view(J[idxs]) .+ @view(D[idxs]) - end - return J_cache - else - return @. J + D - end -end -@inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::Number) - if __can_setindex(J_cache) - copyto!(J_cache, J) - if fast_scalar_indexing(J_cache) - @simd for i in axes(J_cache, 1) - @inbounds J_cache[i, i] += D - end - else - idxs = diagind(J_cache) - J_cache[idxs] .= @view(J[idxs]) .+ D - end - return J_cache - else - return @. J + D - end -end diff --git a/src/descent/newton.jl b/src/descent/newton.jl deleted file mode 100644 index f0964c7dd..000000000 --- a/src/descent/newton.jl +++ /dev/null @@ -1,122 +0,0 @@ -""" - NewtonDescent(; linsolve = nothing, precs = DEFAULT_PRECS) - -Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is -commonly referred to as the Gauss-Newton Descent. - -See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@ref). -""" -@kwdef @concrete struct NewtonDescent <: AbstractDescentAlgorithm - linsolve = nothing - precs = DEFAULT_PRECS -end - -function Base.show(io::IO, d::NewtonDescent) - modifiers = String[] - d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") - d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") - print(io, "NewtonDescent($(join(modifiers, ", ")))") -end - -supports_line_search(::NewtonDescent) = true - -@concrete mutable struct NewtonDescentCache{pre_inverted, normalform} <: - AbstractDescentCache - δu - δus - lincache - JᵀJ_cache # For normal form else nothing - Jᵀfu_cache - timer -end - -@internal_caches NewtonDescentCache :lincache - -function __internal_init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; stats, - shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, - linsolve_kwargs = (;), abstol = nothing, reltol = nothing, - timer = get_timer_output(), kwargs...) where {INV, N} - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - INV && return NewtonDescentCache{true, false}(δu, δus, nothing, nothing, nothing, timer) - lincache = construct_linear_solver( - alg, alg.linsolve, J, _vec(fu), _vec(u); stats, abstol, reltol, linsolve_kwargs...) - return NewtonDescentCache{false, false}(δu, δus, lincache, nothing, nothing, timer) -end - -function __internal_init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, - u; stats, pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - shared::Val{N} = Val(1), abstol = nothing, reltol = nothing, - timer = get_timer_output(), kwargs...) where {INV, N} - length(fu) != length(u) && - @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." - - normal_form = NonlinearSolveBase.needs_square_A(alg.linsolve, u) - if normal_form - JᵀJ = transpose(J) * J - Jᵀfu = transpose(J) * _vec(fu) - A, b = __maybe_symmetric(JᵀJ), Jᵀfu - else - JᵀJ, Jᵀfu = nothing, nothing - A, b = J, _vec(fu) - end - lincache = construct_linear_solver( - alg, alg.linsolve, A, b, _vec(u); stats, abstol, reltol, linsolve_kwargs...) - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - return NewtonDescentCache{false, normal_form}(δu, δus, lincache, JᵀJ, Jᵀfu, timer) -end - -function __internal_solve!( - cache::NewtonDescentCache{INV, false}, J, fu, u, idx::Val = Val(1); - skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) where {INV} - δu = get_du(cache, idx) - skip_solve && return DescentResult(; δu) - if INV - @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." - @bb δu = J × vec(fu) - else - @static_timeit cache.timer "linear solve" begin - linres = cache.lincache(; - A = J, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu), - reuse_A_if_factorization = !new_jacobian || (idx !== Val(1))) - δu = _restructure(get_du(cache, idx), linres.u) - if !linres.success - set_du!(cache, δu, idx) - return DescentResult(; δu, success = false, linsolve_success = false) - end - end - end - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return DescentResult(; δu) -end - -function __internal_solve!( - cache::NewtonDescentCache{false, true}, J, fu, u, idx::Val = Val(1); - skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) - δu = get_du(cache, idx) - skip_solve && return DescentResult(; δu) - if idx === Val(1) - @bb cache.JᵀJ_cache = transpose(J) × J - end - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - @static_timeit cache.timer "linear solve" begin - linres = cache.lincache(; - A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, - kwargs..., linu = _vec(δu), du = _vec(δu), - reuse_A_if_factorization = !new_jacobian || (idx !== Val(1))) - δu = _restructure(get_du(cache, idx), linres.u) - if !linres.success - set_du!(cache, δu, idx) - return DescentResult(; δu, success = false, linsolve_success = false) - end - end - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return DescentResult(; δu) -end diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl deleted file mode 100644 index 82b02552d..000000000 --- a/src/descent/steepest.jl +++ /dev/null @@ -1,74 +0,0 @@ -""" - SteepestDescent(; linsolve = nothing, precs = DEFAULT_PRECS) - -Compute the descent direction as ``δu = -Jᵀfu``. The linear solver and preconditioner are -only used if `J` is provided in the inverted form. - -See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). -""" -@kwdef @concrete struct SteepestDescent <: AbstractDescentAlgorithm - linsolve = nothing - precs = DEFAULT_PRECS -end - -function Base.show(io::IO, d::SteepestDescent) - modifiers = String[] - d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") - d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") - print(io, "SteepestDescent($(join(modifiers, ", ")))") -end - -supports_line_search(::SteepestDescent) = true - -@concrete mutable struct SteepestDescentCache{pre_inverted} <: AbstractDescentCache - δu - δus - lincache - timer -end - -@internal_caches SteepestDescentCache :lincache - -@inline function __internal_init( - prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; - stats, shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, - linsolve_kwargs = (;), abstol = nothing, reltol = nothing, - timer = get_timer_output(), kwargs...) where {INV, N} - INV && @assert length(fu)==length(u) "Non-Square Jacobian Inverse doesn't make sense." - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - if INV - lincache = construct_linear_solver( - alg, alg.linsolve, transpose(J), _vec(fu), _vec(u); - stats, abstol, reltol, linsolve_kwargs...) - else - lincache = nothing - end - return SteepestDescentCache{INV}(δu, δus, lincache, timer) -end - -function __internal_solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val = Val(1); - new_jacobian::Bool = true, kwargs...) where {INV} - δu = get_du(cache, idx) - if INV - A = J === nothing ? nothing : transpose(J) - @static_timeit cache.timer "linear solve" begin - linres = cache.lincache(; - A, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu), - reuse_A_if_factorization = !new_jacobian || idx !== Val(1)) - δu = _restructure(get_du(cache, idx), linres.u) - if !linres.success - set_du!(cache, δu, idx) - return DescentResult(; δu, success = false, linsolve_success = false) - end - end - else - @assert J!==nothing "`J` must be provided when `pre_inverted = Val(false)`." - @bb δu = transpose(J) × vec(fu) - end - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return DescentResult(; δu) -end diff --git a/src/timer_outputs.jl b/src/timer_outputs.jl index 510e1d5ed..b28ef01dd 100644 --- a/src/timer_outputs.jl +++ b/src/timer_outputs.jl @@ -1,12 +1,3 @@ -# Timer Outputs has some overhead, so we only use it if we are debugging -# Even `@static_timeit` has overhead so we write our custom version of that using -# Preferences -const TIMER_OUTPUTS_ENABLED = @load_preference("enable_timer_outputs", false) - -@static if TIMER_OUTPUTS_ENABLED - using TimerOutputs -end - """ enable_timer_outputs() @@ -14,7 +5,7 @@ Enable `TimerOutput` for all `NonlinearSolve` algorithms. This is useful for deb but has some overhead, so it is disabled by default. """ function enable_timer_outputs() - @set_preferences!("enable_timer_outputs"=>true) + set_preferences!(NonlinearSolveBase, "enable_timer_outputs" => true; force = true) @info "Timer Outputs Enabled. Restart the Julia session for this to take effect." end @@ -25,32 +16,6 @@ Disable `TimerOutput` for all `NonlinearSolve` algorithms. This should be used w `NonlinearSolve` is being used in performance-critical code. """ function disable_timer_outputs() - @set_preferences!("enable_timer_outputs"=>false) + set_preferences!(NonlinearSolveBase, "enable_timer_outputs" => false; force = true) @info "Timer Outputs Disabled. Restart the Julia session for this to take effect." end - -function get_timer_output() - @static if TIMER_OUTPUTS_ENABLED - return TimerOutput() - else - return nothing - end -end - -""" - @static_timeit to name expr - -Like `TimerOutputs.@timeit_debug` but has zero overhead if `TimerOutputs` is disabled via -[`NonlinearSolve.disable_timer_outputs()`](@ref). -""" -macro static_timeit(to, name, expr) - @static if TIMER_OUTPUTS_ENABLED - return TimerOutputs.timer_expr(__module__, false, to, name, expr) - else - return esc(expr) - end -end - -@static if !TIMER_OUTPUTS_ENABLED - @inline reset_timer!(::Nothing) = nothing -end diff --git a/src/utils.jl b/src/utils.jl index af657b0b8..9d44fbaf0 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -23,8 +23,7 @@ end (alias || !__can_setindex(typeof(x))) && return x return deepcopy(x) end -@inline __maybe_unaliased(x::AbstractNonlinearSolveOperator, alias::Bool) = x -@inline __maybe_unaliased(x::AbstractJacobianOperator, alias::Bool) = x +@inline __maybe_unaliased(x::AbstractSciMLOperator, ::Bool) = x @inline __cond(J::AbstractMatrix) = cond(J) @inline __cond(J::SVector) = __cond(Diagonal(MVector(J)))