feat: SimpleBroyden implementation

lib/SimpleNonlinearSolve/Project.toml

@@ -13,6 +13,7 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
 NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"

lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl

@@ -3,6 +3,7 @@ module SimpleNonlinearSolve
 using CommonSolve: CommonSolve, solve
 using ConcreteStructs: @concrete
 using FastClosures: @closure
+using LineSearch: LiFukushimaLineSearch
 using LinearAlgebra: dot
 using MaybeInplace: @bb
 using PrecompileTools: @compile_workload, @setup_workload
@@ -79,6 +80,7 @@ function solve_adjoint_internal end
         prob_oop = NonlinearProblem{false}((u, p) -> u .* u .- p, ones(T, 3), T(2))
 
         algs = [
+            SimpleBroyden(),
             SimpleKlement(),
             SimpleNewtonRaphson(),
             SimpleTrustRegion()
@@ -100,7 +102,7 @@ export AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff
 
 export Alefeld, Bisection, Brent, Falsi, ITP, Ridder
 
-export SimpleKlement
+export SimpleBroyden, SimpleKlement
 export SimpleGaussNewton, SimpleNewtonRaphson, SimpleTrustRegion
 
 end

lib/SimpleNonlinearSolve/src/broyden.jl

@@ -1 +1,101 @@
+"""
+    SimpleBroyden(; linesearch = Val(false), alpha = nothing)
 
+A low-overhead implementation of Broyden. This method is non-allocating on scalar and static
+array problems.
+
+### Keyword Arguments
+
+  - `linesearch`: If `linesearch` is `Val(true)`, then we use the `LiFukushimaLineSearch`
+    line search else no line search is used. For advanced customization of the line search,
+    use `Broyden` from `NonlinearSolve.jl`.
+  - `alpha`: Scale the initial jacobian initialization with `alpha`. If it is `nothing`, we
+    will compute the scaling using `2 * norm(fu) / max(norm(u), true)`.
+"""
+@concrete struct SimpleBroyden <: AbstractSimpleNonlinearSolveAlgorithm
+    linesearch <: Union{Val{false}, Val{true}}
+    alpha
+end
+
+function SimpleBroyden(;
+        linesearch::Union{Bool, Val{true}, Val{false}} = Val(false), alpha = nothing)
+    linesearch = linesearch isa Bool ? Val(linesearch) : linesearch
+    return SimpleBroyden(linesearch, alpha)
+end
+
+function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleBroyden, args...;
+        abstol = nothing, reltol = nothing, maxiters = 1000,
+        alias_u0 = false, termination_condition = nothing, kwargs...)
+    x = Utils.maybe_unaliased(prob.u0, alias_u0)
+    fx = Utils.get_fx(prob, x)
+    fx = Utils.eval_f(prob, fx, x)
+    T = promote_type(eltype(fx), eltype(x))
+
+    iszero(fx) &&
+        return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
+
+    @bb xo = copy(x)
+    @bb δx = similar(x)
+    @bb δf = copy(fx)
+    @bb fprev = copy(fx)
+
+    if alg.alpha === nothing
+        fx_norm = L2_NORM(fx)
+        x_norm = L2_NORM(x)
+        init_α = ifelse(fx_norm ≥ 1e-5, max(x_norm, T(true)) / (2 * fx_norm), T(true))
+    else
+        init_α = inv(alg.alpha)
+    end
+
+    J⁻¹ = Utils.identity_jacobian(fx, x, init_α)
+    @bb J⁻¹δf = copy(x)
+    @bb xᵀJ⁻¹ = copy(x)
+    @bb δJ⁻¹n = copy(x)
+    @bb δJ⁻¹ = copy(J⁻¹)
+
+    abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache(
+        prob, abstol, reltol, fx, x, termination_condition, Val(:simple))
+
+    if alg.linesearch === Val(true)
+        ls_alg = LiFukushimaLineSearch(; nan_maxiters = nothing)
+        ls_cache = init(prob, ls_alg, fx, x)
+    else
+        ls_cache = nothing
+    end
+
+    for _ in 1:maxiters
+        @bb δx = J⁻¹ × vec(fprev)
+        @bb δx .*= -1
+
+        if ls_cache === nothing
+            α = true
+        else
+            ls_sol = solve!(ls_cache, xo, δx)
+            α = ls_sol.step_size # Ignores the return code for now
+        end
+
+        @bb @. x = xo + α * δx
+        fx = Utils.eval_f(prob, fx, x)
+        @bb @. δf = fx - fprev
+
+        # Termination Checks
+        solved, retcode, fx_sol, x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob)
+        solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode)
+
+        @bb J⁻¹δf = J⁻¹ × vec(δf)
+        d = dot(δx, J⁻¹δf)
+        @bb xᵀJ⁻¹ = transpose(J⁻¹) × vec(δx)
+
+        @bb @. δJ⁻¹n = (δx - J⁻¹δf) / d
+
+        δJ⁻¹n_ = Utils.safe_vec(δJ⁻¹n)
+        xᵀJ⁻¹_ = Utils.safe_vec(xᵀJ⁻¹)
+        @bb δJ⁻¹ = δJ⁻¹n_ × transpose(xᵀJ⁻¹_)
+        @bb J⁻¹ .+= δJ⁻¹
+
+        @bb copyto!(xo, x)
+        @bb copyto!(fprev, fx)
+    end
+
+    return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters)
+end

lib/SimpleNonlinearSolve/src/utils.jl

@@ -62,13 +62,13 @@ function identity_jacobian(u::Number, fu::Number, α = true)
     return convert(promote_type(eltype(u), eltype(fu)), α)
 end
 function identity_jacobian(u, fu, α = true)
-    J = safe_similar(u, promote_type(eltype(u), eltype(fu)))
+    J = safe_similar(u, promote_type(eltype(u), eltype(fu)), length(fu), length(u))
     fill!(J, zero(eltype(J)))
     J[diagind(J)] .= eltype(J)(α)
     return J
 end
 function identity_jacobian(u::StaticArray, fu, α = true)
-    return SMatrix{length(fu), length(u), eltype(u)}(I * α)
+    return SMatrix{length(fu), length(u), promote_type(eltype(fu), eltype(u))}(I * α)
 end
 
 identity_jacobian!!(J::Number) = one(J)