Merge pull request #49 from yash2798/ys/halley

Updated halley's method for multivariate functions
SciML · Feb 23, 2023 · cda307b · cda307b
2 parents 9b04692 + 7f54266
commit cda307b
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 26 deletions.
diff --git a/lib/SimpleNonlinearSolve/src/halley.jl b/lib/SimpleNonlinearSolve/src/halley.jl
@@ -43,9 +43,8 @@ function SciMLBase.__solve(prob::NonlinearProblem,
     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")
+    if isa(x, AbstractArray)
+        n = length(x)
     end
     T = typeof(x)
 
@@ -65,22 +64,45 @@ function SciMLBase.__solve(prob::NonlinearProblem,
 
     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)
+            if isa(x, Number)
+                fx = f(x)
+                dfx = ForwardDiff.derivative(f, x)
+                d2fx = ForwardDiff.derivative(x -> ForwardDiff.derivative(f, x), x)
+            else
+                fx = f(x)
+                dfx = ForwardDiff.jacobian(f, x)
+                d2fx = ForwardDiff.jacobian(x -> ForwardDiff.jacobian(f, x), x)
+                ai = -(dfx \ fx)
+                A = reshape(d2fx * ai, (n, n))
+                bi = (dfx) \ (A * ai)
+                ci = (ai .* ai) ./ (ai .+ (0.5 .* bi))
+            end
         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)
+            if isa(x, Number)
+                fx = f(x)
+                dfx = FiniteDiff.finite_difference_derivative(f, x, diff_type(alg), eltype(x))
+                d2fx = FiniteDiff.finite_difference_derivative(x -> FiniteDiff.finite_difference_derivative(f, x), x,
+                                                                                            diff_type(alg), eltype(x))
+            else
+                fx = f(x)
+                dfx = FiniteDiff.finite_difference_jacobian(f, x, diff_type(alg), eltype(x))
+                d2fx = FiniteDiff.finite_difference_jacobian(x -> FiniteDiff.finite_difference_jacobian(f, x), x, 
+                                                                        diff_type(alg), eltype(x))
+                ai = -(dfx \ fx)
+                A = reshape(d2fx * ai, (n, n))
+                bi = (dfx) \ (A * ai)
+                ci = (ai .* ai) ./ (ai .+ (0.5 .* bi))
+            end
         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 isa(x, Number)
+            Δx = (2 * dfx^2 - fx * d2fx) \ (2fx * dfx)
+            x -= Δx
+        else
+            Δx = ci
+            x += Δx
+        end
         if isapprox(x, xo, atol = atol, rtol = rtol)
             return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
         end

diff --git a/lib/SimpleNonlinearSolve/test/basictests.jl b/lib/SimpleNonlinearSolve/test/basictests.jl
@@ -49,10 +49,10 @@ function benchmark_scalar(f, u0)
     sol = (solve(probN, Halley()))
 end
 
-# function ff(u, p)
-#     u .* u .- 2
-# end
-# const cu0 = @SVector[1.0, 1.0]
+function ff(u, p)
+    u .* u .- 2
+end
+const cu0 = @SVector[1.0, 1.0]
 function sf(u, p)
     u * u - 2
 end
@@ -62,6 +62,10 @@ sol = benchmark_scalar(sf, csu0)
 @test sol.retcode === ReturnCode.Success
 @test sol.u * sol.u - 2 < 1e-9
 
+sol = benchmark_scalar(ff, cu0)
+@test sol.retcode === ReturnCode.Success
+@test sol.u .* sol.u .- 2 < [1e-9, 1e-9]
+
 if VERSION >= v"1.7"
     @test (@ballocated benchmark_scalar(sf, csu0)) == 0
 end
@@ -122,7 +126,7 @@ using ForwardDiff
 f, u0 = (u, p) -> u .* u .- p, @SVector[1.0, 1.0]
 
 for alg in (SimpleNewtonRaphson(), LBroyden(), Klement(), SimpleTrustRegion(),
-            SimpleDFSane(), BROYDEN_SOLVERS...)
+            SimpleDFSane(), Halley(), BROYDEN_SOLVERS...)
     g = function (p)
         probN = NonlinearProblem{false}(f, csu0, p)
         sol = solve(probN, alg, abstol = 1e-9)
@@ -221,19 +225,14 @@ probN = NonlinearProblem(f, u0)
 
 for alg in (SimpleNewtonRaphson(), SimpleNewtonRaphson(; autodiff = false),
             SimpleTrustRegion(),
-            SimpleTrustRegion(; autodiff = false), LBroyden(), Klement(), SimpleDFSane(),
+            SimpleTrustRegion(; autodiff = false), Halley(), Halley(; autodiff = false), LBroyden(), Klement(), SimpleDFSane(),
             BROYDEN_SOLVERS...)
     sol = solve(probN, alg)
 
     @test sol.retcode == ReturnCode.Success
     @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