Merge branch 'main' into ap/termination_broyden

Author: ChrisRackauckas

Project.toml

@@ -4,7 +4,7 @@ authors = ["SciML"]
 version = "0.1.12"
 
 [deps]
-ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -22,8 +22,8 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 SimpleBatchedNonlinearSolveExt = "NNlib"
 
 [compat]
+ArrayInterface = "6, 7"
 ArrayInterfaceCore = "0.1.1"
-DiffEqBase = "6.118.1"
 FiniteDiff = "2"
 ForwardDiff = "0.10.3"
 NNlib = "0.8"

ext/SimpleBatchedNonlinearSolveExt.jl

@@ -1,6 +1,7 @@
 module SimpleBatchedNonlinearSolveExt
 
-using ArrayInterfaceCore, DiffEqBase, LinearAlgebra, SimpleNonlinearSolve, SciMLBase
+using ArrayInterface, DiffEqBase, LinearAlgebra, SimpleNonlinearSolve, SciMLBase
+
 isdefined(Base, :get_extension) ? (using NNlib) : (using ..NNlib)
 
 _batch_transpose(x) = reshape(x, 1, size(x)...)
@@ -20,7 +21,7 @@ function _batched_mul(x::AbstractArray{T1, 3}, y::AbstractArray{T2, 3}) where {T
 end
 
 function _init_J_batched(x::AbstractMatrix{T}) where {T}
-    J = ArrayInterfaceCore.zeromatrix(x[:, 1])
+    J = ArrayInterface.zeromatrix(x[:, 1])
     if ismutable(x)
         J[diagind(J)] .= one(eltype(x))
     else

src/SimpleNonlinearSolve.jl

@@ -5,7 +5,7 @@ using FiniteDiff, ForwardDiff
 using ForwardDiff: Dual
 using StaticArraysCore
 using LinearAlgebra
-import ArrayInterfaceCore
+import ArrayInterface
 using DiffEqBase
 
 @reexport using SciMLBase
@@ -37,12 +37,14 @@ include("ridder.jl")
 include("brent.jl")
 include("dfsane.jl")
 include("ad.jl")
+include("halley.jl")
 
 import SnoopPrecompile
 
 SnoopPrecompile.@precompile_all_calls begin for T in (Float32, Float64)
     prob_no_brack = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))
-    for alg in (SimpleNewtonRaphson, Broyden, Klement, SimpleTrustRegion, SimpleDFSane)
+    for alg in (SimpleNewtonRaphson, Halley, Broyden, Klement, SimpleTrustRegion,
+                SimpleDFSane)
         solve(prob_no_brack, alg(), abstol = T(1e-2))
     end
 
@@ -63,7 +65,7 @@ SnoopPrecompile.@precompile_all_calls begin for T in (Float32, Float64)
 end end
 
 # DiffEq styled algorithms
-export Bisection, Brent, Broyden, LBroyden, SimpleDFSane, Falsi, Klement,
+export Bisection, Brent, Broyden, LBroyden, SimpleDFSane, Falsi, Halley, Klement,
        Ridder, SimpleNewtonRaphson, SimpleTrustRegion
 
 end # module

src/halley.jl

@@ -0,0 +1,91 @@
+"""
+```julia
+Halley(; chunk_size = Val{0}(), autodiff = Val{true}(),
+                                 diff_type = Val{:forward})
+```
+
+A low-overhead implementation of Halley's Method. This method is non-allocating on scalar
+and static array problems.
+
+!!! note
+
+    As part of the decreased overhead, this method omits some of the higher level error
+    catching of the other methods. Thus, to see better error messages, use one of the other
+    methods like `NewtonRaphson`
+
+### Keyword Arguments
+
+- `chunk_size`: the chunk size used by the internal ForwardDiff.jl automatic differentiation
+  system. This allows for multiple derivative columns to be computed simultaneously,
+  improving performance. Defaults to `0`, which is equivalent to using ForwardDiff.jl's
+  default chunk size mechanism. For more details, see the documentation for
+  [ForwardDiff.jl](https://juliadiff.org/ForwardDiff.jl/stable/).
+- `autodiff`: whether to use forward-mode automatic differentiation for the Jacobian.
+  Note that this argument is ignored if an analytical Jacobian is passed; as that will be
+  used instead. Defaults to `Val{true}`, which means ForwardDiff.jl is used by default.
+  If `Val{false}`, then FiniteDiff.jl is used for finite differencing.
+- `diff_type`: the type of finite differencing used if `autodiff = false`. Defaults to
+  `Val{:forward}` for forward finite differences. For more details on the choices, see the
+  [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) documentation.
+"""
+struct Halley{CS, AD, FDT} <: AbstractNewtonAlgorithm{CS, AD, FDT}
+    function Halley(; chunk_size = Val{0}(), autodiff = Val{true}(),
+                    diff_type = Val{:forward})
+        new{SciMLBase._unwrap_val(chunk_size), SciMLBase._unwrap_val(autodiff),
+            SciMLBase._unwrap_val(diff_type)}()
+    end
+end
+
+function SciMLBase.__solve(prob::NonlinearProblem,
+                           alg::Halley, args...; abstol = nothing,
+                           reltol = nothing,
+                           maxiters = 1000, kwargs...)
+    f = Base.Fix2(prob.f, prob.p)
+    x = float(prob.u0)
+    fx = f(x)
+    # fx = float(prob.u0)
+    if !isa(fx, Number) || !isa(x, Number)
+        error("Halley currently only supports scalar-valued single-variable functions")
+    end
+    T = typeof(x)
+
+    if SciMLBase.isinplace(prob)
+        error("Halley currently only supports out-of-place nonlinear problems")
+    end
+
+    atol = abstol !== nothing ? abstol :
+           real(oneunit(eltype(T))) * (eps(real(one(eltype(T)))))^(4 // 5)
+    rtol = reltol !== nothing ? reltol : eps(real(one(eltype(T))))^(4 // 5)
+
+    if typeof(x) <: Number
+        xo = oftype(one(eltype(x)), Inf)
+    else
+        xo = map(x -> oftype(one(eltype(x)), Inf), x)
+    end
+
+    for i in 1:maxiters
+        if alg_autodiff(alg)
+            fx = f(x)
+            dfdx(x) = ForwardDiff.derivative(f, x)
+            dfx = dfdx(x)
+            d2fx = ForwardDiff.derivative(dfdx, x)
+        else
+            fx = f(x)
+            dfx = FiniteDiff.finite_difference_derivative(f, x, diff_type(alg), eltype(x),
+                                                          fx)
+            d2fx = FiniteDiff.finite_difference_derivative(x -> FiniteDiff.finite_difference_derivative(f,
+                                                                                                        x),
+                                                           x, diff_type(alg), eltype(x), fx)
+        end
+        iszero(fx) &&
+            return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
+        Δx = (2 * dfx^2 - fx * d2fx) \ (2fx * dfx)
+        x -= Δx
+        if isapprox(x, xo, atol = atol, rtol = rtol)
+            return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
+        end
+        xo = x
+    end
+
+    return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters)
+end

src/utils.jl

@@ -35,7 +35,7 @@ value(x::Dual) = ForwardDiff.value(x)
 value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)
 
 function init_J(x)
-    J = ArrayInterfaceCore.zeromatrix(x)
+    J = ArrayInterface.zeromatrix(x)
     if ismutable(x)
         J[diagind(J)] .= one(eltype(x))
     else

test/basictests.jl

@@ -43,6 +43,29 @@ if VERSION >= v"1.7"
     @test (@ballocated benchmark_scalar(sf, csu0)) == 0
 end
 
+# Halley
+function benchmark_scalar(f, u0)
+    probN = NonlinearProblem{false}(f, u0)
+    sol = (solve(probN, Halley()))
+end
+
+# function ff(u, p)
+#     u .* u .- 2
+# end
+# const cu0 = @SVector[1.0, 1.0]
+function sf(u, p)
+    u * u - 2
+end
+const csu0 = 1.0
+
+sol = benchmark_scalar(sf, csu0)
+@test sol.retcode === ReturnCode.Success
+@test sol.u * sol.u - 2 < 1e-9
+
+if VERSION >= v"1.7"
+    @test (@ballocated benchmark_scalar(sf, csu0)) == 0
+end
+
 # Broyden
 function benchmark_scalar(f, u0, alg)
     probN = NonlinearProblem{false}(f, u0)
@@ -115,7 +138,7 @@ end
 # Scalar
 f, u0 = (u, p) -> u * u - p, 1.0
 for alg in (SimpleNewtonRaphson(), LBroyden(), Klement(), SimpleTrustRegion(),
-            SimpleDFSane(), BROYDEN_SOLVERS...)
+            SimpleDFSane(), Halley(), BROYDEN_SOLVERS...)
     g = function (p)
         probN = NonlinearProblem{false}(f, oftype(p, u0), p)
         sol = solve(probN, alg)
@@ -181,7 +204,7 @@ for alg in [Bisection(), Falsi(), Ridder(), Brent()]
 end
 
 for alg in (SimpleNewtonRaphson(), LBroyden(), Klement(), SimpleTrustRegion(),
-            SimpleDFSane(), BROYDEN_SOLVERS...)
+            SimpleDFSane(), Halley(), BROYDEN_SOLVERS...)
     global g, p
     g = function (p)
         probN = NonlinearProblem{false}(f, 0.5, p)
@@ -206,6 +229,12 @@ for alg in (SimpleNewtonRaphson(), SimpleNewtonRaphson(; autodiff = false),
     @test sol.u[end] ≈ sqrt(2.0)
 end
 
+# Separate Error check for Halley; will be included in above error checks for the improved Halley
+f, u0 = (u, p) -> u * u - 2.0, 1.0
+probN = NonlinearProblem(f, u0)
+
+@test solve(probN, Halley()).u ≈ sqrt(2.0)
+
 for u0 in [1.0, [1, 1.0]]
     local f, probN, sol
     f = (u, p) -> u .* u .- 2.0