refactor: move algorithms to First Order sub-package

lib/NonlinearSolveFirstOrder/Project.toml

@@ -3,15 +3,53 @@ uuid = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d"
 authors = ["Avik Pal <[email protected]> and contributors"]
 version = "1.0.0"
 
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
+NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+
 [compat]
+ADTypes = "1.9.0"
 Aqua = "0.8"
+ArrayInterface = "7.16.0"
+CommonSolve = "0.2.4"
+ConcreteStructs = "0.2.3"
+DiffEqBase = "6.155.3"
 ExplicitImports = "1.5"
+FiniteDiff = "2.26.0"
+ForwardDiff = "0.10.36"
 Hwloc = "3"
 InteractiveUtils = "<0.0.1, 1"
+LineSearch = "0.1.4"
+LinearAlgebra = "1.11.0"
+LinearSolve = "2.36.1"
+MaybeInplace = "0.1.4"
 NonlinearProblemLibrary = "0.1.2"
+NonlinearSolveBase = "1.1"
 Pkg = "1.10"
+PrecompileTools = "1.2"
 ReTestItems = "1.24"
+Reexport = "1"
+SciMLBase = "2.54"
+SciMLOperators = "0.3.11"
+Setfield = "1.1.1"
 StableRNGs = "1"
+StaticArraysCore = "1.4.3"
 Test = "1.10"
 julia = "1.10"
 

lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl

@@ -3,23 +3,29 @@ module NonlinearSolveFirstOrder
 using Reexport: @reexport
 using PrecompileTools: @compile_workload, @setup_workload
 
+using ADTypes: ADTypes
 using ArrayInterface: ArrayInterface
 using CommonSolve: CommonSolve
 using ConcreteStructs: @concrete
-using DiffEqBase: DiffEqBase   # Needed for `init` / `solve` dispatches
+using DiffEqBase: DiffEqBase    # Needed for `init` / `solve` dispatches
+using FiniteDiff: FiniteDiff    # Default Finite Difference Method
+using ForwardDiff: ForwardDiff  # Default Forward Mode AD
 using LinearAlgebra: LinearAlgebra, Diagonal, dot, inv, diag
-using LinearSolve: LinearSolve # Trigger Linear Solve extension in NonlinearSolveBase
+using LinearSolve: LinearSolve  # Trigger Linear Solve extension in NonlinearSolveBase
 using MaybeInplace: @bb
 using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm,
                           AbstractNonlinearSolveCache, AbstractResetCondition,
                           AbstractResetConditionCache, AbstractApproximateJacobianStructure,
                           AbstractJacobianCache, AbstractJacobianInitialization,
                           AbstractApproximateJacobianUpdateRule, AbstractDescentDirection,
                           AbstractApproximateJacobianUpdateRuleCache,
+                          AbstractDampingFunction, AbstractDampingFunctionCache,
                           Utils, InternalAPI, get_timer_output, @static_timeit,
-                          update_trace!, L2_NORM, NewtonDescent
+                          update_trace!, L2_NORM,
+                          NewtonDescent, DampedNewtonDescent
 using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode
 using SciMLOperators: AbstractSciMLOperator
+using Setfield: @set!
 using StaticArraysCore: StaticArray, Size, MArray
 
 include("raphson.jl")

lib/NonlinearSolveFirstOrder/src/gauss_newton.jl

@@ -1 +1,22 @@
+"""
+    GaussNewton(;
+        concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing,
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+    )
 
+An advanced GaussNewton implementation with support for efficient handling of sparse
+matrices via colored automatic differentiation and preconditioned linear solvers. Designed
+for large-scale and numerically-difficult nonlinear systems.
+"""
+function GaussNewton(;
+        concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing,
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+)
+    return GeneralizedFirstOrderAlgorithm(;
+        linesearch,
+        descent = NewtonDescent(; linsolve, precs),
+        autodiff, vjp_autodiff, jvp_autodiff,
+        concrete_jac,
+        name = :GaussNewton
+    )
+end

lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl

@@ -1 +1,90 @@
+"""
+    PseudoTransient(;
+        concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3,
+        linsolve = nothing, precs = nothing,
+        autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing
+    )
 
+An implementation of PseudoTransient Method [coffey2003pseudotransient](@cite) that is used
+to solve steady state problems in an accelerated manner. It uses an adaptive time-stepping
+to integrate an initial value of nonlinear problem until sufficient accuracy in the desired
+steady-state is achieved to switch over to Newton's method and gain a rapid convergence.
+This implementation specifically uses "switched evolution relaxation"
+[kelley1998convergence](@cite) SER method.
+
+### Keyword Arguments
+
+  - `alpha_initial` : the initial pseudo time step. It defaults to `1e-3`. If it is small,
+    you are going to need more iterations to converge but it can be more stable.
+"""
+function PseudoTransient(;
+        concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3,
+        linsolve = nothing, precs = nothing,
+        autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing
+)
+    return GeneralizedFirstOrderAlgorithm(;
+        linesearch,
+        descent = DampedNewtonDescent(;
+            linsolve, precs, initial_damping = alpha_initial,
+            damping_fn = SwitchedEvolutionRelaxation()
+        ),
+        autodiff,
+        jvp_autodiff,
+        vjp_autodiff,
+        concrete_jac,
+        name = :PseudoTransient
+    )
+end
+
+"""
+    SwitchedEvolutionRelaxation()
+
+Method for updating the damping parameter in the [`PseudoTransient`](@ref) method based on
+"switched evolution relaxation" [kelley1998convergence](@cite) SER method.
+"""
+struct SwitchedEvolutionRelaxation <: AbstractDampingFunction end
+
+"""
+    SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache
+
+Cache for the [`SwitchedEvolutionRelaxation`](@ref) method.
+"""
+@concrete mutable struct SwitchedEvolutionRelaxationCache <: AbstractDampingFunctionCache
+    res_norm
+    α⁻¹
+    internalnorm
+end
+
+function NonlinearSolveBase.requires_normal_form_jacobian(::Union{
+        SwitchedEvolutionRelaxation, SwitchedEvolutionRelaxationCache})
+    return false
+end
+
+function NonlinearSolveBase.requires_normal_form_rhs(::Union{
+        SwitchedEvolutionRelaxation, SwitchedEvolutionRelaxationCache})
+    return false
+end
+
+function InternalAPI.init(
+        prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation,
+        initial_damping, J, fu, u, args...;
+        internalnorm::F = L2_NORM, kwargs...
+) where {F}
+    T = promote_type(eltype(u), eltype(fu))
+    return SwitchedEvolutionRelaxationCache(
+        internalnorm(fu),
+        T(inv(initial_damping)),
+        internalnorm
+    )
+end
+
+(damping::SwitchedEvolutionRelaxationCache)(::Nothing) = damping.α⁻¹
+
+function InternalAPI.solve!(
+        damping::SwitchedEvolutionRelaxationCache, J, fu, args...; kwargs...
+)
+    res_norm = damping.internalnorm(fu)
+    damping.α⁻¹ *= res_norm / damping.res_norm
+    damping.res_norm = res_norm
+    return damping.α⁻¹
+end

lib/NonlinearSolveFirstOrder/src/raphson.jl

@@ -1 +1,22 @@
+"""
+    NewtonRaphson(;
+        concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing,
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+    )
 
+An advanced NewtonRaphson implementation with support for efficient handling of sparse
+matrices via colored automatic differentiation and preconditioned linear solvers. Designed
+for large-scale and numerically-difficult nonlinear systems.
+"""
+function NewtonRaphson(;
+        concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing,
+        autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing
+)
+    return GeneralizedFirstOrderAlgorithm(;
+        linesearch,
+        descent = NewtonDescent(; linsolve, precs),
+        autodiff, vjp_autodiff, jvp_autodiff,
+        concrete_jac,
+        name = :NewtonRaphson
+    )
+end