diff --git a/lib/SimpleNonlinearSolve/src/halley.jl b/lib/SimpleNonlinearSolve/src/halley.jl index 77ea3a4a1..0e5417642 100644 --- 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 index 26e92dfe8..8fbdfffaf 100644 --- 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,7 +225,7 @@ 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) @@ -229,11 +233,6 @@ 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