Add Halley's method via descent API

Author: tansongchen

Project.toml

@@ -27,6 +27,7 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]

src/NonlinearSolve.jl

@@ -48,6 +48,7 @@ include("timer_outputs.jl")
 include("internal/helpers.jl")
 
 include("descent/newton.jl")
+include("descent/halley.jl")
 include("descent/steepest.jl")
 include("descent/dogleg.jl")
 include("descent/damped_newton.jl")
@@ -80,6 +81,7 @@ include("algorithms/gauss_newton.jl")
 include("algorithms/levenberg_marquardt.jl")
 include("algorithms/trust_region.jl")
 include("algorithms/extension_algs.jl")
+include("algorithms/halley.jl")
 
 include("utils.jl")
 include("default.jl")
@@ -141,7 +143,7 @@ include("default.jl")
 end
 
 # Core Algorithms
-export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane
+export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane, Halley
 export GaussNewton, LevenbergMarquardt, TrustRegion
 export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPolyalg,
        FastShortcutNLLSPolyalg
@@ -154,7 +156,7 @@ export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, NLSol
 export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, GeneralizedDFSane
 
 # Descent Algorithms
-export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, GeodesicAcceleration
+export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, GeodesicAcceleration, HalleyDescent
 
 # Globalization
 ## Line Search Algorithms

src/algorithms/halley.jl

@@ -0,0 +1,12 @@
+"""
+    Halley(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(),
+        precs = DEFAULT_PRECS, autodiff = nothing)
+
+An experimental Halley's method implementation.
+"""
+function Halley(; concrete_jac = nothing, linsolve = nothing,
+        linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing)
+    descent = HalleyDescent(; linsolve, precs)
+    return GeneralizedFirstOrderAlgorithm(;
+        concrete_jac, name = :Halley, linesearch, descent, jacobian_ad = autodiff)
+end

src/descent/halley.jl

@@ -0,0 +1,87 @@
+"""
+    HalleyDescent(; linsolve = nothing, precs = DEFAULT_PRECS)
+
+Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is
+commonly referred to as the Gauss-Newton Descent.
+
+See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@ref).
+"""
+@kwdef @concrete struct HalleyDescent <: AbstractDescentAlgorithm
+    linsolve = nothing
+    precs = DEFAULT_PRECS
+end
+
+using TaylorDiff: derivative
+
+function Base.show(io::IO, d::HalleyDescent)
+    modifiers = String[]
+    d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)")
+    d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)")
+    print(io, "HalleyDescent($(join(modifiers, ", ")))")
+end
+
+supports_line_search(::HalleyDescent) = true
+
+@concrete mutable struct HalleyDescentCache{pre_inverted} <:
+                         AbstractDescentCache
+    f
+    p
+    δu
+    δus
+    b
+    lincache
+    timer
+end
+
+@internal_caches HalleyDescentCache :lincache
+
+function __internal_init(
+        prob::NonlinearProblem, alg::HalleyDescent, J, fu, u; shared::Val{N} = Val(1),
+        pre_inverted::Val{INV} = False, linsolve_kwargs = (;), abstol = nothing,
+        reltol = nothing, timer = get_timer_output(), kwargs...) where {INV, N}
+    @bb δu = similar(u)
+    @bb b = similar(u)
+    δus = N ≤ 1 ? nothing : map(2:N) do i
+        @bb δu_ = similar(u)
+    end
+    INV && return HalleyDescentCache{true}(prob.f, prob.p, δu, δus, b, nothing, timer)
+    lincache = LinearSolverCache(
+        alg, alg.linsolve, J, _vec(fu), _vec(u); abstol, reltol, linsolve_kwargs...)
+    return HalleyDescentCache{false}(prob.f, prob.p, δu, δus, b, lincache, timer)
+end
+
+function __internal_solve!(
+        cache::HalleyDescentCache{INV}, J, fu, u, idx::Val = Val(1);
+        skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) where {INV}
+    δu = get_du(cache, idx)
+    skip_solve && return δu, true, (;)
+    if INV
+        @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`."
+        @bb δu = J × vec(fu)
+    else
+        @static_timeit cache.timer "linear solve 1" begin
+            δu = cache.lincache(;
+                A = J, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu),
+                reuse_A_if_factorization = !new_jacobian || (idx !== Val(1)))
+            δu = _restructure(get_du(cache, idx), δu)
+        end
+    end
+    b = cache.b
+    # compute the hessian-vector-vector product
+    hvvp = derivative(x -> cache.f(x, cache.p), u, δu, 2)
+    # second linear solve, reuse factorization if possible
+    if INV
+        @bb b = J × vec(hvvp)
+    else
+        @static_timeit cache.timer "linear solve 2" begin
+            b = cache.lincache(;
+                A = J, b = _vec(hvvp), kwargs..., linu = _vec(b), du = _vec(b),
+                reuse_A_if_factorization = true)
+            b = _restructure(cache.b, b)
+        end
+    end
+    @bb @. δu = δu * δu / (b / 2 - δu)
+    set_du!(cache, δu, idx)
+    cache.b = b
+    return δu, true, (;)
+end