From 3ca7a5482c5990e4d6e4099d2f5bf421b39d9c9c Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 5 Dec 2023 21:57:11 -0500 Subject: [PATCH] Allow specifying an alpha for Broyden --- src/broyden.jl | 21 +++++++++++++-------- src/utils.jl | 43 ++++++++++++++++++++++++++++--------------- 2 files changed, 41 insertions(+), 23 deletions(-) diff --git a/src/broyden.jl b/src/broyden.jl index 785ac0221..0884e07a8 100644 --- a/src/broyden.jl +++ b/src/broyden.jl @@ -1,7 +1,7 @@ # Sadly `Broyden` is taken up by SimpleNonlinearSolve.jl """ GeneralBroyden(; max_resets = 3, linesearch = nothing, reset_tolerance = nothing, - init_jacobian::Val = Val(:identity), autodiff = nothing) + init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = 1.0) An implementation of `Broyden` with resetting and line search. @@ -16,6 +16,8 @@ An implementation of `Broyden` with resetting and line search. used here directly, and they will be converted to the correct `LineSearch`. It is recommended to use [LiFukushimaLineSearch](@ref) -- a derivative free linesearch specifically designed for Broyden's method. + - `alpha_initial`: If `init_jacobian` is set to `Val(:identity)`, then the initial + Jacobian inverse is set to be `(αI)⁻¹`. Defaults to `1.0`. - `init_jacobian`: the method to use for initializing the jacobian. Defaults to using the identity matrix (`Val(:identitiy)`). Alternatively, can be set to `Val(:true_jacobian)` to use the true jacobian as initialization. @@ -33,22 +35,24 @@ An implementation of `Broyden` with resetting and line search. max_resets::Int reset_tolerance linesearch + inv_alpha end -function __alg_print_modifiers(::GeneralBroyden{IJ, UR}) where {IJ, UR} +function __alg_print_modifiers(alg::GeneralBroyden{IJ, UR}) where {IJ, UR} modifiers = String[] IJ !== :identity && push!(modifiers, "init_jacobian = :$(IJ)") UR !== :good_broyden && push!(modifiers, "update_rule = :$(UR)") + alg.inv_alpha != 1 && push!(modifiers, "alpha = $(1 / alg.inv_alpha)") return modifiers end function set_ad(alg::GeneralBroyden{IJ, UR, CJ}, ad) where {IJ, UR, CJ} return GeneralBroyden{IJ, UR, CJ}(ad, alg.max_resets, alg.reset_tolerance, - alg.linesearch) + alg.linesearch, alg.inv_alpha) end function GeneralBroyden(; max_resets = 3, linesearch = nothing, reset_tolerance = nothing, - init_jacobian::Val = Val(:identity), autodiff = nothing, + init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = 1.0, update_rule = Val(:good_broyden)) UR = _unwrap_val(update_rule) @assert UR ∈ (:good_broyden, :bad_broyden) @@ -56,7 +60,8 @@ function GeneralBroyden(; max_resets = 3, linesearch = nothing, reset_tolerance @assert IJ ∈ (:identity, :true_jacobian) linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch) CJ = IJ === :true_jacobian - return GeneralBroyden{IJ, UR, CJ}(autodiff, max_resets, reset_tolerance, linesearch) + return GeneralBroyden{IJ, UR, CJ}(autodiff, max_resets, reset_tolerance, linesearch, + 1 / alpha) end @concrete mutable struct GeneralBroydenCache{iip, IJ, UR} <: @@ -109,7 +114,7 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralBroyd alg = alg_ @bb du = similar(u) uf, fu_cache, jac_cache = nothing, nothing, nothing - J⁻¹ = __init_identity_jacobian(u, fu) + J⁻¹ = __init_identity_jacobian(u, fu, alg.inv_alpha) end reset_tolerance = alg.reset_tolerance === nothing ? sqrt(eps(real(eltype(u)))) : @@ -162,7 +167,7 @@ function perform_step!(cache::GeneralBroydenCache{iip, IJ, UR}) where {iip, IJ, if IJ === :true_jacobian cache.J⁻¹ = inv(jacobian!!(cache.J⁻¹, cache)) else - cache.J⁻¹ = __reinit_identity_jacobian!!(cache.J⁻¹) + cache.J⁻¹ = __reinit_identity_jacobian!!(cache.J⁻¹, cache.alg.inv_alpha) end cache.resets += 1 else @@ -189,7 +194,7 @@ function perform_step!(cache::GeneralBroydenCache{iip, IJ, UR}) where {iip, IJ, end function __reinit_internal!(cache::GeneralBroydenCache; kwargs...) - cache.J⁻¹ = __reinit_identity_jacobian!!(cache.J⁻¹) + cache.J⁻¹ = __reinit_identity_jacobian!!(cache.J⁻¹, cache.alg.inv_alpha) cache.resets = 0 return nothing end diff --git a/src/utils.jl b/src/utils.jl index a31cd659e..59535aea1 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -315,34 +315,47 @@ function check_and_update!(tc_cache, cache, fu, u, uprev, end end -@inline __init_identity_jacobian(u::Number, _) = one(u) -@inline function __init_identity_jacobian(u, fu) +@inline __init_identity_jacobian(u::Number, fu, α = true) = oftype(u, α) +@inline @views function __init_identity_jacobian(u, fu, α = true) J = similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) fill!(J, zero(eltype(J))) - J[diagind(J)] .= one(eltype(J)) + if fast_scalar_indexing(J) + @inbounds for i in axes(J, 1) + J[i, i] = α + end + else + J[diagind(J)] .= α + end return J end -@inline function __init_identity_jacobian(u::StaticArray, fu::StaticArray) +@inline function __init_identity_jacobian(u::StaticArray, fu::StaticArray, α = true) T = promote_type(eltype(fu), eltype(u)) - return MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) + return MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) end -@inline function __init_identity_jacobian(u::SArray, fu::SArray) +@inline function __init_identity_jacobian(u::SArray, fu::SArray, α = true) T = promote_type(eltype(fu), eltype(u)) - return SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I) + return SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) end -@inline __reinit_identity_jacobian!!(J::Number) = one(J) -@inline __reinit_identity_jacobian!!(J::AbstractVector) = fill!(J, one(eltype(J))) -@inline function __reinit_identity_jacobian!!(J::AbstractMatrix) +@inline __reinit_identity_jacobian!!(J::Number, α = true) = oftype(J, α) +@inline __reinit_identity_jacobian!!(J::AbstractVector, α = true) = fill!(J, α) +@inline @views function __reinit_identity_jacobian!!(J::AbstractMatrix, α = true) fill!(J, zero(eltype(J))) - J[diagind(J)] .= one(eltype(J)) + if fast_scalar_indexing(J) + @inbounds for i in axes(J, 1) + J[i, i] = α + end + else + J[diagind(J)] .= α + end return J end -@inline __reinit_identity_jacobian!!(J::SVector) = ones(SArray{ - Tuple{Size(J)[1]}, eltype(J)}) -@inline function __reinit_identity_jacobian!!(J::SMatrix) +@inline function __reinit_identity_jacobian!!(J::SVector, α = true) + return ones(SArray{Tuple{Size(J)[1]}, eltype(J)}) .* α +end +@inline function __reinit_identity_jacobian!!(J::SMatrix, α = true) S = Size(J) - return SArray{Tuple{S[1], S[2]}, eltype(J)}(I) + return SArray{Tuple{S[1], S[2]}, eltype(J)}(I) .* α end function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2},