Add true jacobian init for broyden and klement

Author: avik-pal

src/broyden.jl

@@ -1,6 +1,7 @@
 # Sadly `Broyden` is taken up by SimpleNonlinearSolve.jl
 """
-    GeneralBroyden(; max_resets = 3, linesearch = nothing, reset_tolerance = nothing)
+    GeneralBroyden(; max_resets = 3, linesearch = nothing, reset_tolerance = nothing,
+        init_jacobian::Val = Val(:identity), autodiff = nothing)
 
 An implementation of `Broyden` with resetting and line search.
 
@@ -14,20 +15,36 @@ 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.
+  - `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.
+  - `autodiff`: determines the backend used for the Jacobian. Note that this argument is
+    ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to
+    `nothing` which means that a default is selected according to the problem specification!
+    Valid choices are types from ADTypes.jl. (Used if `init_jacobian = Val(:true_jacobian)`)
 """
-@concrete struct GeneralBroyden <: AbstractNewtonAlgorithm{false, Nothing}
+@concrete struct GeneralBroyden{IJ, CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD}
+    ad::AD
     max_resets::Int
     reset_tolerance
     linesearch
 end
 
+function set_ad(alg::GeneralBroyden{IJ, CJ}, ad) where {IJ, CJ}
+    return GeneralBroyden{IJ, CJ}(ad, alg.max_resets, alg.reset_tolerance, alg.linesearch)
+end
+
 function GeneralBroyden(; max_resets = 3, linesearch = nothing,
-        reset_tolerance = nothing)
+        reset_tolerance = nothing, init_jacobian::Val = Val(:identity),
+        autodiff = nothing)
+    IJ = _unwrap_val(init_jacobian)
+    @assert IJ ∈ (:identity, :true_jacobian)
     linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch)
-    return GeneralBroyden(max_resets, reset_tolerance, linesearch)
+    CJ = IJ === :true_jacobian
+    return GeneralBroyden{IJ, CJ}(autodiff, max_resets, reset_tolerance, linesearch)
 end
 
-@concrete mutable struct GeneralBroydenCache{iip} <: AbstractNonlinearSolveCache{iip}
+@concrete mutable struct GeneralBroydenCache{iip, IJ} <: AbstractNonlinearSolveCache{iip}
     f
     alg
     u
@@ -37,6 +54,7 @@ end
     fu_cache
     dfu
     p
+    uf
     J⁻¹
     J⁻¹dfu
     force_stop::Bool
@@ -49,45 +67,62 @@ end
     reltol
     reset_tolerance
     reset_check
+    jac_cache
     prob
     stats::NLStats
     ls_cache
     tc_cache
     trace
 end
 
-function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::GeneralBroyden, args...;
-        alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing,
+function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralBroyden{IJ},
+        args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing,
         termination_condition = nothing, internalnorm::F = DEFAULT_NORM,
-        kwargs...) where {uType, iip, F}
+        kwargs...) where {uType, iip, F, IJ}
     @unpack f, u0, p = prob
     u = __maybe_unaliased(u0, alias_u0)
     fu = evaluate_f(prob, u)
     @bb du = copy(u)
-    J⁻¹ = __init_identity_jacobian(u, fu)
+
+    if IJ === :true_jacobian
+        alg = get_concrete_algorithm(alg_, prob)
+        uf, _, J, fu_cache, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip);
+            lininit = Val(false))
+        J⁻¹ = J
+    else
+        alg = alg_
+        @bb du = similar(u)
+        uf, fu_cache, jac_cache = nothing, nothing, nothing
+        J⁻¹ = __init_identity_jacobian(u, fu)
+    end
+
     reset_tolerance = alg.reset_tolerance === nothing ? sqrt(eps(real(eltype(u)))) :
                       alg.reset_tolerance
     reset_check = x -> abs(x) ≤ reset_tolerance
 
     @bb u_cache = copy(u)
-    @bb fu_cache = copy(fu)
-    @bb dfu = similar(fu)
+    @bb dfu = copy(fu)
     @bb J⁻¹dfu = similar(u)
 
     abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u,
         termination_condition)
     trace = init_nonlinearsolve_trace(alg, u, fu, J⁻¹, du; uses_jac_inverse = Val(true),
         kwargs...)
 
-    return GeneralBroydenCache{iip}(f, alg, u, u_cache, du, fu, fu_cache, dfu, p,
+    return GeneralBroydenCache{iip, IJ}(f, alg, u, u_cache, du, fu, fu_cache, dfu, p, uf,
         J⁻¹, J⁻¹dfu, false, 0, alg.max_resets, maxiters, internalnorm, ReturnCode.Default,
-        abstol, reltol, reset_tolerance, reset_check, prob, NLStats(1, 0, 0, 0, 0),
+        abstol, reltol, reset_tolerance, reset_check, jac_cache, prob,
+        NLStats(1, 0, 0, 0, 0),
         init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache, trace)
 end
 
-function perform_step!(cache::GeneralBroydenCache{iip}) where {iip}
+function perform_step!(cache::GeneralBroydenCache{iip, IJ}) where {iip, IJ}
     T = eltype(cache.u)
 
+    if IJ === :true_jacobian && cache.stats.nsteps == 0
+        cache.J⁻¹ = inv(jacobian!!(cache.J⁻¹, cache)) # This allocates
+    end
+
     @bb cache.du = cache.J⁻¹ × vec(cache.fu)
     α = perform_linesearch!(cache.ls_cache, cache.u, cache.du)
     @bb axpy!(-α, cache.du, cache.u)
@@ -100,15 +135,19 @@ function perform_step!(cache::GeneralBroydenCache{iip}) where {iip}
     cache.force_stop && return nothing
 
     # Update the inverse jacobian
-    @bb @. cache.dfu = cache.fu - cache.fu_cache
+    @bb @. cache.dfu = cache.fu - cache.dfu
 
     if all(cache.reset_check, cache.du) || all(cache.reset_check, cache.dfu)
         if cache.resets ≥ cache.max_resets
             cache.retcode = ReturnCode.ConvergenceFailure
             cache.force_stop = true
             return nothing
         end
-        cache.J⁻¹ = __reinit_identity_jacobian!!(cache.J⁻¹)
+        if IJ === :true_jacobian
+            cache.J⁻¹ = inv(jacobian!!(cache.J⁻¹, cache))
+        else
+            cache.J⁻¹ = __reinit_identity_jacobian!!(cache.J⁻¹)
+        end
         cache.resets += 1
     else
         @bb cache.du .*= -1
@@ -119,7 +158,7 @@ function perform_step!(cache::GeneralBroydenCache{iip}) where {iip}
         @bb cache.J⁻¹ += vec(cache.du) × transpose(_vec(cache.u_cache))
     end
 
-    @bb copyto!(cache.fu_cache, cache.fu)
+    @bb copyto!(cache.dfu, cache.fu)
     @bb copyto!(cache.u_cache, cache.u)
 
     return nothing

src/jacobian.jl

@@ -58,6 +58,7 @@ function jacobian!!(::Number, cache)
     cache.stats.njacs += 1
     return last(value_derivative(cache.uf, cache.u))
 end
+
 # Build Jacobian Caches
 function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f::F, u, p, ::Val{iip};
         linsolve_kwargs = (;), lininit::Val{linsolve_init} = Val(true),

src/klement.jl

@@ -19,33 +19,54 @@ solves.
   - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref),
     which means that no line search is performed. Algorithms from `LineSearches.jl` can be
     used here directly, and they will be converted to the correct `LineSearch`.
+  - `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. (Our tests suggest it is a good idea to
+    to initialize with an identity matrix)
+  - `autodiff`: determines the backend used for the Jacobian. Note that this argument is
+    ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to
+    `nothing` which means that a default is selected according to the problem specification!
+    Valid choices are types from ADTypes.jl. (Used if `init_jacobian = Val(:true_jacobian)`)
 """
-@concrete struct GeneralKlement <: AbstractNewtonAlgorithm{false, Nothing}
+@concrete struct GeneralKlement{IJ, CJ, AD} <: AbstractNewtonAlgorithm{CJ, AD}
+    ad::AD
     max_resets::Int
     linsolve
     precs
     linesearch
 end
 
-function set_linsolve(alg::GeneralKlement, linsolve)
-    return GeneralKlement(alg.max_resets, linsolve, alg.precs, alg.linesearch)
+function set_linsolve(alg::GeneralKlement{IJ, CS}, linsolve) where {IJ, CS}
+    return GeneralKlement{IJ, CS}(alg.ad, alg.max_resets, linsolve, alg.precs,
+        alg.linesearch)
+end
+
+function set_ad(alg::GeneralKlement{IJ, CS}, ad) where {IJ, CS}
+    return GeneralKlement{IJ, CS}(ad, alg.max_resets, alg.linsolve, alg.precs,
+        alg.linesearch)
 end
 
 function GeneralKlement(; max_resets::Int = 5, linsolve = nothing,
-        linesearch = nothing, precs = DEFAULT_PRECS)
+        linesearch = nothing, precs = DEFAULT_PRECS, init_jacobian::Val = Val(:identity),
+        autodiff = nothing)
+    IJ = _unwrap_val(init_jacobian)
+    @assert IJ ∈ (:identity, :true_jacobian)
     linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch)
-    return GeneralKlement(max_resets, linsolve, precs, linesearch)
+    CJ = IJ === :true_jacobian
+    return GeneralKlement{IJ, CJ}(autodiff, max_resets, linsolve, precs, linesearch)
 end
 
-@concrete mutable struct GeneralKlementCache{iip} <: AbstractNonlinearSolveCache{iip}
+@concrete mutable struct GeneralKlementCache{iip, IJ} <: AbstractNonlinearSolveCache{iip}
     f
     alg
     u
     u_cache
     fu
     fu_cache
+    fu_cache_2
     du
     p
+    uf
     linsolve
     J
     J_cache
@@ -60,73 +81,96 @@ end
     abstol
     reltol
     prob
+    jac_cache
     stats::NLStats
     ls_cache
     tc_cache
     trace
 end
 
-function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralKlement, args...;
-        alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing,
+function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralKlement{IJ},
+        args...; alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing,
         termination_condition = nothing, internalnorm::F = DEFAULT_NORM,
-        linsolve_kwargs = (;), kwargs...) where {uType, iip, F}
+        linsolve_kwargs = (;), kwargs...) where {uType, iip, F, IJ}
     @unpack f, u0, p = prob
     u = __maybe_unaliased(u0, alias_u0)
     fu = evaluate_f(prob, u)
-    J = __init_identity_jacobian(u, fu)
-    @bb du = similar(u)
 
-    if u isa Number
-        linsolve = FakeLinearSolveJLCache(J, fu)
+    if IJ === :true_jacobian
+        alg = get_concrete_algorithm(alg_, prob)
+        uf, _, J, fu_cache, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip);
+            lininit = Val(false))
+    elseif IJ === :identity
         alg = alg_
+        @bb du = similar(u)
+        uf, fu_cache, jac_cache = nothing, nothing, nothing
+        J = one.(u) # Identity Init Jacobian for Klement maintains a Diagonal Structure
+    else
+        error("Invalid `init_jacobian` value")
+    end
+
+    if IJ === :true_jacobian
+        linsolve = linsolve_caches(J, _vec(fu), _vec(du), p, alg_; linsolve_kwargs)
     else
-        # For General Julia Arrays default to LU Factorization
-        linsolve_alg = (alg_.linsolve === nothing && (u isa Array || u isa StaticArray)) ?
-                       LUFactorization() : nothing
-        alg = set_linsolve(alg_, linsolve_alg)
-        linsolve = linsolve_caches(J, _vec(fu), _vec(du), p, alg; linsolve_kwargs)
+        linsolve = nothing
     end
 
     abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu, u,
         termination_condition)
     trace = init_nonlinearsolve_trace(alg, u, fu, J, du; kwargs...)
 
     @bb u_cache = copy(u)
-    @bb fu_cache = copy(fu)
-    @bb J_cache = similar(J)
-    @bb J_cache_2 = similar(J)
+    @bb fu_cache_2 = copy(fu)
     @bb Jdu = similar(fu)
-    @bb Jdu_cache = similar(fu)
+    if IJ === :true_jacobian
+        @bb J_cache = similar(J)
+        @bb J_cache_2 = similar(J)
+        @bb Jdu_cache = similar(fu)
+    else
+        J_cache, J_cache_2, Jdu_cache = nothing, nothing, nothing
+    end
 
-    return GeneralKlementCache{iip}(f, alg, u, u_cache, fu, fu_cache, du, p, linsolve,
-        J, J_cache, J_cache_2, Jdu, Jdu_cache, 0, false, maxiters, internalnorm,
-        ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0),
+    return GeneralKlementCache{iip, IJ}(f, alg, u, u_cache, fu, fu_cache, fu_cache_2, du, p,
+        uf, linsolve, J, J_cache, J_cache_2, Jdu, Jdu_cache, 0, false, maxiters,
+        internalnorm,
+        ReturnCode.Default, abstol, reltol, prob, jac_cache, NLStats(1, 0, 0, 0, 0),
         init_linesearch_cache(alg.linesearch, f, u, p, fu, Val(iip)), tc_cache, trace)
 end
 
-function perform_step!(cache::GeneralKlementCache{iip}) where {iip}
+function perform_step!(cache::GeneralKlementCache{iip, IJ}) where {iip, IJ}
     @unpack linsolve, alg = cache
     T = eltype(cache.J)
-    singular, fact_done = __try_factorize_and_check_singular!(linsolve, cache.J)
 
-    if singular
+    if IJ === :true_jacobian
+        cache.stats.nsteps == 0 && (cache.J = jacobian!!(cache.J, cache))
+        ill_conditioned = __is_ill_conditioned(cache.J)
+    elseif IJ === :identity
+        ill_conditioned = __is_ill_conditioned(cache.J)
+    end
+
+    if ill_conditioned
         if cache.resets == alg.max_resets
             cache.force_stop = true
             cache.retcode = ReturnCode.ConvergenceFailure
             return nothing
         end
-        fact_done = false
-        cache.J = __reinit_identity_jacobian!!(cache.J)
+        if IJ === :true_jacobian && cache.stats.nsteps != 0
+            cache.J = jacobian!!(cache.J, cache)
+        else
+            cache.J = __reinit_identity_jacobian!!(cache.J)
+        end
         cache.resets += 1
     end
 
-    A = ifelse(cache.J isa SMatrix || cache.J isa Number || !fact_done, cache.J, nothing)
-
-    # u = u - J \ fu
-    linres = dolinsolve(cache, alg.precs, cache.linsolve; A,
-        b = _vec(cache.fu), linu = _vec(cache.du), cache.p, reltol = cache.abstol)
-    cache.linsolve = linres.cache
-    cache.du = _restructure(cache.du, linres.u)
+    if IJ === :identity
+        @bb @. cache.du = cache.fu / cache.J
+    else
+        # u = u - J \ fu
+        linres = dolinsolve(cache, alg.precs, cache.linsolve; A = cache.J,
+            b = _vec(cache.fu), linu = _vec(cache.du), cache.p, reltol = cache.abstol)
+        cache.linsolve = linres.cache
+        cache.du = _restructure(cache.du, linres.u)
+    end
 
     # Line Search
     α = perform_linesearch!(cache.ls_cache, cache.u, cache.du)
@@ -143,18 +187,25 @@ function perform_step!(cache::GeneralKlementCache{iip}) where {iip}
 
     # Update the Jacobian
     @bb cache.du .*= -1
-    @bb cache.J_cache .= cache.J' .^ 2
-    @bb @. cache.Jdu = cache.du^2
-    @bb cache.Jdu_cache = cache.J_cache × vec(cache.Jdu)
-    @bb cache.Jdu = cache.J × vec(cache.du)
-    @bb @. cache.fu_cache = (cache.fu - cache.fu_cache - cache.Jdu) /
-                            ifelse(iszero(cache.Jdu_cache), T(1e-5), cache.Jdu_cache)
-    @bb cache.J_cache = vec(cache.fu_cache) × transpose(_vec(cache.du))
-    @bb @. cache.J_cache *= cache.J
-    @bb cache.J_cache_2 = cache.J_cache × cache.J
-    @bb cache.J .+= cache.J_cache_2
-
-    @bb copyto!(cache.fu_cache, cache.fu)
+    if IJ === :identity
+        @bb @. cache.Jdu = (cache.J^2) * (cache.du^2)
+        @bb @. cache.J += ((cache.fu - cache.fu_cache_2 - cache.J * cache.du) /
+                           ifelse(iszero(cache.Jdu), T(1e-5), cache.Jdu)) * cache.du *
+                          (cache.J^2)
+    else
+        @bb cache.J_cache .= cache.J' .^ 2
+        @bb @. cache.Jdu = cache.du^2
+        @bb cache.Jdu_cache = cache.J_cache × vec(cache.Jdu)
+        @bb cache.Jdu = cache.J × vec(cache.du)
+        @bb @. cache.fu_cache_2 = (cache.fu - cache.fu_cache_2 - cache.Jdu) /
+                                  ifelse(iszero(cache.Jdu_cache), T(1e-5), cache.Jdu_cache)
+        @bb cache.J_cache = vec(cache.fu_cache_2) × transpose(_vec(cache.du))
+        @bb @. cache.J_cache *= cache.J
+        @bb cache.J_cache_2 = cache.J_cache × cache.J
+        @bb cache.J .+= cache.J_cache_2
+    end
+
+    @bb copyto!(cache.fu_cache_2, cache.fu)
 
     return nothing
 end