refactor: create Quasi Newton Algorithm subpackage

Author: avik-pal

Project.toml

@@ -19,6 +19,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
 NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
+NonlinearSolveQuasiNewton = "9a2c21bd-3a47-402d-9113-8faf9a0ee114"
 NonlinearSolveSpectralMethods = "26075421-4e9a-44e1-8bd1-420ed7ad02b2"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"

common/nlls_problem_workloads.jl

@@ -0,0 +1 @@
+

docs/src/release_notes.md

@@ -4,6 +4,7 @@
 
 ### Breaking Changes in `NonlinearSolve.jl` v4
 
+  - `ApproximateJacobianSolveAlgorithm` has been renamed to `QuasiNewtonAlgorithm`.
   - See [common breaking changes](@ref common-breaking-changes-v4v2) below.
 
 ### Breaking Changes in `SimpleNonlinearSolve.jl` v2

lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl

@@ -15,6 +15,9 @@ Utils.make_sparse(x) = sparse(x)
 
 Utils.condition_number(J::AbstractSparseMatrix) = Utils.condition_number(Matrix(J))
 
-Utils.maybe_pinv!!_workspace(A::AbstractSparseMatrix) = Matrix(A)
+function Utils.maybe_pinv!!_workspace(A::AbstractSparseMatrix)
+    dense_A = Matrix(A)
+    return dense_A, copy(dense_A)
+end
 
 end

lib/NonlinearSolveBase/src/NonlinearSolveBase.jl

@@ -11,7 +11,7 @@ using DifferentiationInterface: DifferentiationInterface, Constant
 using EnzymeCore: EnzymeCore
 using FastClosures: @closure
 using FunctionProperties: hasbranching
-using LinearAlgebra: LinearAlgebra, Diagonal, norm, ldiv!, diagind
+using LinearAlgebra: LinearAlgebra, Diagonal, norm, ldiv!, diagind, pinv
 using Markdown: @doc_str
 using MaybeInplace: @bb
 using Preferences: @load_preference

lib/NonlinearSolveBase/src/abstract_types.jl

@@ -226,11 +226,11 @@ Abstract Type for all NonlinearSolveBase Caches.
 ### Interface Functions
 
   - `get_fu(cache)`: get the residual.
+
   - `get_u(cache)`: get the current state.
   - `set_fu!(cache, fu)`: set the residual.
   - `has_time_limit(cache)`: whether or not the solver has a maximum time limit.
   - `not_terminated(cache)`: whether or not the solver has terminated.
-
   - `SciMLBase.set_u!(cache, u)`: set the current state.
   - `SciMLBase.reinit!(cache, u0; kwargs...)`: reinitialize the cache with the initial state
     `u0` and any additional keyword arguments.
@@ -339,3 +339,170 @@ Abstract Type for all Jacobian Caches used in NonlinearSolveBase. Subtypes of th
 meant to be constructured via [`construct_jacobian_cache`](@ref).
 """
 abstract type AbstractJacobianCache <: AbstractNonlinearSolveBaseAPI end
+
+"""
+    AbstractApproximateJacobianStructure
+
+Abstract Type for all Approximate Jacobian Structures used in NonlinearSolve.jl.
+
+### Interface Functions
+
+  - `stores_full_jacobian(alg)`: whether or not the algorithm stores the full Jacobian.
+    Defaults to `false`.
+  - `get_full_jacobian(cache, alg, J)`: get the full Jacobian. Defaults to throwing an
+    error if `stores_full_jacobian(alg)` is `false`.
+"""
+abstract type AbstractApproximateJacobianStructure <: AbstractNonlinearSolveBaseAPI end
+
+stores_full_jacobian(::AbstractApproximateJacobianStructure) = false
+function get_full_jacobian(cache, alg::AbstractApproximateJacobianStructure, J)
+    stores_full_jacobian(alg) && return J
+    error("This algorithm does not store the full Jacobian. Define `get_full_jacobian` for \
+           this algorithm.")
+end
+
+"""
+    AbstractJacobianInitialization
+
+Abstract Type for all Jacobian Initialization Algorithms used in NonlinearSolveBase.
+
+### Interface Functions
+
+  - `jacobian_initialized_preinverted(alg)`: whether or not the Jacobian is initialized
+    preinverted. Defaults to `false`.
+
+### `InternalAPI.init` specification
+
+```julia
+InternalAPI.init(
+    prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, solver,
+    f, fu, u, p;
+    linsolve = missing, internalnorm::IN = L2_NORM, kwargs...
+)::AbstractJacobianCache
+```
+
+All subtypes need to define
+`(cache::AbstractJacobianCache)(alg::NewSubType, fu, u)` which reinitializes the Jacobian in
+`cache.J`.
+"""
+abstract type AbstractJacobianInitialization <: AbstractNonlinearSolveBaseAPI end
+
+jacobian_initialized_preinverted(::AbstractJacobianInitialization) = false
+
+"""
+    AbstractApproximateJacobianUpdateRule
+
+Abstract Type for all Approximate Jacobian Update Rules used in NonlinearSolveBase.
+
+### Interface Functions
+
+  - `store_inverse_jacobian(alg)`: Return `alg.store_inverse_jacobian`
+
+### `InternalAPI.init` specification
+
+```julia
+InternalAPI.init(
+    prob::AbstractNonlinearProblem, alg::AbstractApproximateJacobianUpdateRule, J, fu, u,
+    du, args...; internalnorm = L2_NORM, kwargs...
+)::AbstractApproximateJacobianUpdateRuleCache
+```
+"""
+abstract type AbstractApproximateJacobianUpdateRule <: AbstractNonlinearSolveBaseAPI end
+
+function store_inverse_jacobian(rule::AbstractApproximateJacobianUpdateRule)
+    return rule.store_inverse_jacobian
+end
+
+"""
+    AbstractApproximateJacobianUpdateRuleCache
+
+Abstract Type for all Approximate Jacobian Update Rule Caches used in NonlinearSolveBase.
+
+### Interface Functions
+
+  - `store_inverse_jacobian(cache)`: Return `store_inverse_jacobian(cache.rule)`
+
+### `InternalAPI.solve!` specification
+
+```julia
+InternalAPI.solve!(
+    cache::AbstractApproximateJacobianUpdateRuleCache, J, fu, u, du; kwargs...
+) --> J / J⁻¹
+```
+"""
+abstract type AbstractApproximateJacobianUpdateRuleCache <: AbstractNonlinearSolveBaseAPI end
+
+function store_inverse_jacobian(cache::AbstractApproximateJacobianUpdateRuleCache)
+    return store_inverse_jacobian(cache.rule)
+end
+
+"""
+    AbstractResetCondition
+
+Condition for resetting the Jacobian in Quasi-Newton's methods.
+
+### `InternalAPI.init` specification
+
+```julia
+InternalAPI.init(
+    alg::AbstractResetCondition, J, fu, u, du, args...; kwargs...
+)::AbstractResetConditionCache
+```
+"""
+abstract type AbstractResetCondition <: AbstractNonlinearSolveBaseAPI end
+
+"""
+    AbstractResetConditionCache
+
+Abstract Type for all Reset Condition Caches used in NonlinearSolveBase.
+
+### `InternalAPI.solve!` specification
+
+```julia
+InternalAPI.solve!(
+    cache::AbstractResetConditionCache, J, fu, u, du; kwargs...
+)::Bool
+```
+"""
+abstract type AbstractResetConditionCache <: AbstractNonlinearSolveBaseAPI end
+
+"""
+    AbstractTrustRegionMethod
+
+Abstract Type for all Trust Region Methods used in NonlinearSolveBase.
+
+### `InternalAPI.init` specification
+
+```julia
+InternalAPI.init(
+    prob::AbstractNonlinearProblem, alg::AbstractTrustRegionMethod, f, fu, u, p, args...;
+    internalnorm = L2_NORM, kwargs...
+)::AbstractTrustRegionMethodCache
+```
+"""
+abstract type AbstractTrustRegionMethod <: AbstractNonlinearSolveBaseAPI end
+
+"""
+    AbstractTrustRegionMethodCache
+
+Abstract Type for all Trust Region Method Caches used in NonlinearSolveBase.
+
+### Interface Functions
+
+  - `last_step_accepted(cache)`: whether or not the last step was accepted. Defaults to
+    `cache.last_step_accepted`. Should if overloaded if the field is not present.
+
+### `InternalAPI.solve!` specification
+
+```julia
+InternalAPI.solve!(
+    cache::AbstractTrustRegionMethodCache, J, fu, u, δu, descent_stats; kwargs...
+)
+```
+
+Returns `last_step_accepted`, updated `u_cache` and `fu_cache`. If the last step was
+accepted then these values should be copied into the toplevel cache.
+"""
+abstract type AbstractTrustRegionMethodCache <: AbstractNonlinearSolveBaseAPI end
+
+last_step_accepted(cache::AbstractTrustRegionMethodCache) = cache.last_step_accepted

lib/NonlinearSolveBase/src/jacobian.jl

@@ -122,8 +122,7 @@ end
     di_extras
 end
 
-function InternalAPI.reinit!(
-        cache::JacobianCache, args...; p = cache.p, u0 = cache.u, kwargs...)
+function InternalAPI.reinit!(cache::JacobianCache; p = cache.p, u0 = cache.u, kwargs...)
     cache.u = u0
     cache.p = p
 end

lib/NonlinearSolveBase/src/tracing.jl

@@ -205,7 +205,7 @@ function update_trace!(
     return trace
 end
 
-function update_trace!(cache, α = true)
+function update_trace!(cache, α = true; uses_jac_inverse = Val(false))
     trace = Utils.safe_getproperty(cache, Val(:trace))
     trace === missing && return nothing
 
@@ -214,11 +214,8 @@ function update_trace!(cache, α = true)
         update_trace!(
             trace, cache.nsteps + 1, get_u(cache), get_fu(cache), nothing, cache.du, α
         )
-        # XXX: Implement
-        # elseif cache isa ApproximateJacobianSolveCache && store_inverse_jacobian(cache)
-        #     update_trace!(trace, cache.nsteps + 1, get_u(cache), get_fu(cache),
-        #         ApplyArray(__safe_inv, J), cache.du, α)
     else
+        J = uses_jac_inverse isa Val{true} ? pinv(J) : J
         update_trace!(trace, cache.nsteps + 1, get_u(cache), get_fu(cache), J, cache.du, α)
     end
 end

lib/NonlinearSolveBase/src/utils.jl

@@ -2,12 +2,12 @@ module Utils
 
 using ArrayInterface: ArrayInterface
 using FastClosures: @closure
-using LinearAlgebra: LinearAlgebra, Diagonal, Symmetric, norm, dot, cond, diagind, pinv
+using LinearAlgebra: LinearAlgebra, Diagonal, Symmetric, norm, dot, cond, diagind, pinv, inv
 using MaybeInplace: @bb
 using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition
 using SciMLOperators: AbstractSciMLOperator
 using SciMLBase: SciMLBase, AbstractNonlinearProblem, NonlinearFunction
-using StaticArraysCore: StaticArray, SArray
+using StaticArraysCore: StaticArray, SArray, SMatrix
 
 using ..NonlinearSolveBase: NonlinearSolveBase, L2_NORM, Linf_NORM
 
@@ -180,7 +180,7 @@ condition_number(::Any) = -1
 
 # XXX: Move to NonlinearSolveQuasiNewton
 # compute `pinv` if `inv` won't work
-maybe_pinv!!_workspace(A) = nothing
+maybe_pinv!!_workspace(A) = nothing, A
 
 maybe_pinv!!(workspace, A::Union{Number, AbstractMatrix}) = pinv(A)
 function maybe_pinv!!(workspace, A::Diagonal)
@@ -193,12 +193,12 @@ function maybe_pinv!!(workspace, A::StridedMatrix)
     LinearAlgebra.checksquare(A)
     if LinearAlgebra.istriu(A)
         issingular = any(iszero, @view(A[diagind(A)]))
-        A_ = UpperTriangular(A)
-        !issingular && return triu!(parent(inv(A_)))
+        A_ = LinearAlgebra.UpperTriangular(A)
+        !issingular && return LinearAlgebra.triu!(parent(inv(A_)))
     elseif LinearAlgebra.istril(A)
-        A_ = LowerTriangular(A)
+        A_ = LinearAlgebra.LowerTriangular(A)
         issingular = any(iszero, @view(A_[diagind(A_)]))
-        !issingular && return tril!(parent(inv(A_)))
+        !issingular && return LinearAlgebra.tril!(parent(inv(A_)))
     else
         F = LinearAlgebra.lu(A; check = false)
         if issuccess(F)
@@ -209,4 +209,34 @@ function maybe_pinv!!(workspace, A::StridedMatrix)
     return pinv(A)
 end
 
+function initial_jacobian_scaling_alpha(α, u, fu, ::Any)
+    return convert(promote_type(eltype(u), eltype(fu)), α)
+end
+function initial_jacobian_scaling_alpha(::Nothing, u, fu, internalnorm::F) where {F}
+    fu_norm = internalnorm(fu)
+    fu_norm < 1e-5 && return initial_jacobian_scaling_alpha(true, u, fu, internalnorm)
+    return (2 * fu_norm) / max(L2_NORM(u), true)
+end
+
+make_identity!!(::T, α) where {T <: Number} = T(α)
+function make_identity!!(A::AbstractVector{T}, α) where {T}
+    @bb @. A = T(α)
+    return A
+end
+function make_identity!!(::SMatrix{S1, S2, T, L}, α) where {S1, S2, T, L}
+    return SMatrix{S1, S2, T, L}(LinearAlgebra.I * α)
+end
+function make_identity!!(A::AbstractMatrix{T}, α) where {T}
+    A = ArrayInterface.can_setindex(A) ? A : similar(A)
+    fill!(A, false)
+    if ArrayInterface.fast_scalar_indexing(A)
+        @simd ivdep for i in axes(A, 1)
+            @inbounds A[i, i] = α
+        end
+    else
+        A[diagind(A)] .= α
+    end
+    return A
+end
+
 end

lib/NonlinearSolveFirstOrder/test/runtests.jl

@@ -0,0 +1 @@
+