From ed4751fab66596c3629389144d751893a1f1beef Mon Sep 17 00:00:00 2001 From: Songchen Tan Date: Wed, 2 Oct 2024 15:41:34 -0400 Subject: [PATCH 1/3] Add Householder's method --- Project.toml | 2 + ext/SimpleNonlinearSolveTaylorDiffExt.jl | 60 ++++++++++++++++++++++++ src/SimpleNonlinearSolve.jl | 2 + src/nlsolve/householder.jl | 13 +++++ 4 files changed, 77 insertions(+) create mode 100644 ext/SimpleNonlinearSolveTaylorDiffExt.jl create mode 100644 src/nlsolve/householder.jl diff --git a/Project.toml b/Project.toml index 33e3620..881ca66 100644 --- a/Project.toml +++ b/Project.toml @@ -26,12 +26,14 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c" [extensions] SimpleNonlinearSolveChainRulesCoreExt = "ChainRulesCore" SimpleNonlinearSolveReverseDiffExt = "ReverseDiff" SimpleNonlinearSolveTrackerExt = "Tracker" SimpleNonlinearSolveZygoteExt = "Zygote" +SimpleNonlinearSolveTaylorDiffExt = "TaylorDiff" [compat] ADTypes = "1.9" diff --git a/ext/SimpleNonlinearSolveTaylorDiffExt.jl b/ext/SimpleNonlinearSolveTaylorDiffExt.jl new file mode 100644 index 0000000..99b2ca7 --- /dev/null +++ b/ext/SimpleNonlinearSolveTaylorDiffExt.jl @@ -0,0 +1,60 @@ +module SimpleNonlinearSolveTaylorDiffExt +using SimpleNonlinearSolve +using SimpleNonlinearSolve: ImmutableNonlinearProblem, ReturnCode, build_solution, check_termination, init_termination_cache +using SimpleNonlinearSolve: __maybe_unaliased, _get_fx, __fixed_parameter_function +using MaybeInplace: @bb +using SciMLBase: isinplace + +import TaylorDiff + +@inline function __get_higher_order_derivatives(::SimpleHouseholder{N}, prob, f, x) where N + vN = Val(N) + l = map(one, x) + + if isinplace(prob) + fx = f(x) + invf = x -> inv.(f(x)) + bundle = TaylorDiff.derivatives(invf, fx, x, l, vN) + else + t = TaylorDiff.make_seed(x, l, vN) + ft = f(t) + fx = map(TaylorDiff.primal, ft) + bundle = inv.(ft) + end + num = TaylorDiff.extract_derivative(bundle, N - 1) + den = TaylorDiff.extract_derivative(bundle, N) + return num, den, fx +end + +function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleHouseholder{N}, + args...; abstol = nothing, reltol = nothing, maxiters = 1000, + termination_condition = nothing, alias_u0 = false, kwargs...) where N + x = __maybe_unaliased(prob.u0, alias_u0) + length(x) == 1 || + throw(ArgumentError("SimpleHouseholder only supports scalar problems")) + fx = _get_fx(prob, x) + @bb xo = copy(x) + f = __fixed_parameter_function(prob) + + abstol, reltol, tc_cache = init_termination_cache( + prob, abstol, reltol, fx, x, termination_condition) + + for i in 1:maxiters + num, den, fx = __get_higher_order_derivatives(alg, prob, f, x) + + if i == 1 + iszero(fx) && build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) + else + # Termination Checks + tc_sol = check_termination(tc_cache, fx, x, xo, prob, alg) + tc_sol !== nothing && return tc_sol + end + + @bb copyto!(xo, x) + @bb x .+= (N - 1) .* num ./ den + end + + return build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters) +end + +end diff --git a/src/SimpleNonlinearSolve.jl b/src/SimpleNonlinearSolve.jl index ba0b243..e6c939b 100644 --- a/src/SimpleNonlinearSolve.jl +++ b/src/SimpleNonlinearSolve.jl @@ -47,6 +47,7 @@ include("nlsolve/lbroyden.jl") include("nlsolve/klement.jl") include("nlsolve/trustRegion.jl") include("nlsolve/halley.jl") +include("nlsolve/householder.jl") include("nlsolve/dfsane.jl") ## Interval Nonlinear Solvers @@ -139,6 +140,7 @@ end export AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff export SimpleBroyden, SimpleDFSane, SimpleGaussNewton, SimpleHalley, SimpleKlement, SimpleLimitedMemoryBroyden, SimpleNewtonRaphson, SimpleTrustRegion +export SimpleHouseholder export Alefeld, Bisection, Brent, Falsi, ITP, Ridder end # module diff --git a/src/nlsolve/householder.jl b/src/nlsolve/householder.jl new file mode 100644 index 0000000..f012fff --- /dev/null +++ b/src/nlsolve/householder.jl @@ -0,0 +1,13 @@ +""" + SimpleHouseholder{order}() + +A low-overhead implementation of Householder's method. This method is non-allocating on scalar +and static array problems. + +Internally, this uses TaylorDiff.jl for the automatic differentiation. + +### Type Parameters + + - `order`: the convergence order of the Householder method. `order = 2` is the same as Newton's method, `order = 3` is the same as Halley's method, etc. +""" +struct SimpleHouseholder{order} <: AbstractNewtonAlgorithm end From 88bd4725ef3704f9043e7eba3b3424d80dbe1848 Mon Sep 17 00:00:00 2001 From: Songchen Tan Date: Tue, 8 Oct 2024 10:52:00 -0400 Subject: [PATCH 2/3] Add tests and doc --- ext/SimpleNonlinearSolveTaylorDiffExt.jl | 18 ++++++++--------- src/nlsolve/householder.jl | 9 ++++++--- test/core/rootfind_tests.jl | 25 ++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 12 deletions(-) diff --git a/ext/SimpleNonlinearSolveTaylorDiffExt.jl b/ext/SimpleNonlinearSolveTaylorDiffExt.jl index 99b2ca7..3257c68 100644 --- a/ext/SimpleNonlinearSolveTaylorDiffExt.jl +++ b/ext/SimpleNonlinearSolveTaylorDiffExt.jl @@ -7,20 +7,20 @@ using SciMLBase: isinplace import TaylorDiff -@inline function __get_higher_order_derivatives(::SimpleHouseholder{N}, prob, f, x) where N +@inline function __get_higher_order_derivatives(::SimpleHouseholder{N}, prob, f, x, fx) where N vN = Val(N) l = map(one, x) + t = TaylorDiff.make_seed(x, l, vN) if isinplace(prob) - fx = f(x) - invf = x -> inv.(f(x)) - bundle = TaylorDiff.derivatives(invf, fx, x, l, vN) + bundle = similar(fx, TaylorDiff.TaylorScalar{eltype(fx), N}) + f(bundle, t) + map!(TaylorDiff.primal, fx, bundle) else - t = TaylorDiff.make_seed(x, l, vN) - ft = f(t) - fx = map(TaylorDiff.primal, ft) - bundle = inv.(ft) + bundle = f(t) + fx = map(TaylorDiff.primal, bundle) end + bundle = inv.(bundle) num = TaylorDiff.extract_derivative(bundle, N - 1) den = TaylorDiff.extract_derivative(bundle, N) return num, den, fx @@ -40,7 +40,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleHousehold prob, abstol, reltol, fx, x, termination_condition) for i in 1:maxiters - num, den, fx = __get_higher_order_derivatives(alg, prob, f, x) + num, den, fx = __get_higher_order_derivatives(alg, prob, f, x, fx) if i == 1 iszero(fx) && build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) diff --git a/src/nlsolve/householder.jl b/src/nlsolve/householder.jl index f012fff..ce06c5f 100644 --- a/src/nlsolve/householder.jl +++ b/src/nlsolve/householder.jl @@ -1,10 +1,13 @@ """ SimpleHouseholder{order}() -A low-overhead implementation of Householder's method. This method is non-allocating on scalar -and static array problems. +A low-overhead implementation of Householder's method to arbitrary order. +This method is non-allocating on scalar and static array problems. -Internally, this uses TaylorDiff.jl for the automatic differentiation. +!!! warning + + Needs `TaylorDiff.jl` to be explicitly loaded before using this functionality. + Internally, this uses TaylorDiff.jl for automatic differentiation. ### Type Parameters diff --git a/test/core/rootfind_tests.jl b/test/core/rootfind_tests.jl index 272fa3c..76ba85d 100644 --- a/test/core/rootfind_tests.jl +++ b/test/core/rootfind_tests.jl @@ -91,6 +91,31 @@ end end end +@testitem "SimpleHouseholder" setup=[RootfindingTesting] tags=[:core] begin + using TaylorDiff + @testset "AutoDiff: TaylorDiff.jl" for order in (2, 3, 4) + @testset "[OOP] u0: $(nameof(typeof(u0)))" for u0 in ( + [1.0], @SVector[1.0], 1.0) + sol = benchmark_nlsolve_oop(quadratic_f, u0; solver = SimpleHouseholder{order}()) + @test SciMLBase.successful_retcode(sol) + @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) + end + + @testset "[IIP] u0: $(nameof(typeof(u0)))" for u0 in ([1.0],) + sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver = SimpleHouseholder{order}()) + @test SciMLBase.successful_retcode(sol) + @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) + end + end + + @testset "Termination condition: $(nameof(typeof(termination_condition))) u0: $(nameof(typeof(u0)))" for termination_condition in TERMINATION_CONDITIONS, + u0 in (1.0, [1.0], @SVector[1.0]) + + probN = NonlinearProblem(quadratic_f, u0, 2.0) + @test all(solve(probN, SimpleHouseholder{2}(); termination_condition).u .≈ sqrt(2.0)) + end +end + @testitem "Derivative Free Metods" setup=[RootfindingTesting] tags=[:core] begin @testset "$(nameof(typeof(alg)))" for alg in [ SimpleBroyden(), SimpleKlement(), SimpleDFSane(), From d382ad4ec949b546c9899e82df1ec198040b6d50 Mon Sep 17 00:00:00 2001 From: Songchen Tan Date: Thu, 24 Oct 2024 15:27:48 -0400 Subject: [PATCH 3/3] Fix format and add compat --- Project.toml | 4 +++- ext/SimpleNonlinearSolveTaylorDiffExt.jl | 8 +++++--- test/core/rootfind_tests.jl | 16 +++++++++------- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/Project.toml b/Project.toml index 881ca66..e7f4f62 100644 --- a/Project.toml +++ b/Project.toml @@ -65,6 +65,7 @@ SciMLBase = "2.37.0" Setfield = "1.1.1" StaticArrays = "1.9" StaticArraysCore = "1.4.2" +TaylorDiff = "0.2.5" Test = "1.10" Tracker = "0.2.33" Zygote = "0.6.69" @@ -88,9 +89,10 @@ ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["AllocCheck", "Aqua", "DiffEqBase", "ExplicitImports", "FiniteDiff", "ForwardDiff", "Hwloc", "InteractiveUtils", "LinearAlgebra", "NonlinearProblemLibrary", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "Reexport", "ReverseDiff", "StaticArrays", "Test", "Tracker", "Zygote"] +test = ["AllocCheck", "Aqua", "DiffEqBase", "ExplicitImports", "FiniteDiff", "ForwardDiff", "Hwloc", "InteractiveUtils", "LinearAlgebra", "NonlinearProblemLibrary", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "Reexport", "ReverseDiff", "StaticArrays", "TaylorDiff", "Test", "Tracker", "Zygote"] diff --git a/ext/SimpleNonlinearSolveTaylorDiffExt.jl b/ext/SimpleNonlinearSolveTaylorDiffExt.jl index 3257c68..da096a7 100644 --- a/ext/SimpleNonlinearSolveTaylorDiffExt.jl +++ b/ext/SimpleNonlinearSolveTaylorDiffExt.jl @@ -1,13 +1,15 @@ module SimpleNonlinearSolveTaylorDiffExt using SimpleNonlinearSolve -using SimpleNonlinearSolve: ImmutableNonlinearProblem, ReturnCode, build_solution, check_termination, init_termination_cache +using SimpleNonlinearSolve: ImmutableNonlinearProblem, ReturnCode, build_solution, + check_termination, init_termination_cache using SimpleNonlinearSolve: __maybe_unaliased, _get_fx, __fixed_parameter_function using MaybeInplace: @bb using SciMLBase: isinplace import TaylorDiff -@inline function __get_higher_order_derivatives(::SimpleHouseholder{N}, prob, f, x, fx) where N +@inline function __get_higher_order_derivatives( + ::SimpleHouseholder{N}, prob, f, x, fx) where {N} vN = Val(N) l = map(one, x) t = TaylorDiff.make_seed(x, l, vN) @@ -28,7 +30,7 @@ end function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleHouseholder{N}, args...; abstol = nothing, reltol = nothing, maxiters = 1000, - termination_condition = nothing, alias_u0 = false, kwargs...) where N + termination_condition = nothing, alias_u0 = false, kwargs...) where {N} x = __maybe_unaliased(prob.u0, alias_u0) length(x) == 1 || throw(ArgumentError("SimpleHouseholder only supports scalar problems")) diff --git a/test/core/rootfind_tests.jl b/test/core/rootfind_tests.jl index 76ba85d..721d7fa 100644 --- a/test/core/rootfind_tests.jl +++ b/test/core/rootfind_tests.jl @@ -1,6 +1,7 @@ @testsetup module RootfindingTesting using Reexport -@reexport using AllocCheck, StaticArrays, Random, LinearAlgebra, ForwardDiff, DiffEqBase +@reexport using AllocCheck, StaticArrays, Random, LinearAlgebra, ForwardDiff, DiffEqBase, + TaylorDiff import PolyesterForwardDiff quadratic_f(u, p) = u .* u .- p @@ -92,17 +93,17 @@ end end @testitem "SimpleHouseholder" setup=[RootfindingTesting] tags=[:core] begin - using TaylorDiff @testset "AutoDiff: TaylorDiff.jl" for order in (2, 3, 4) - @testset "[OOP] u0: $(nameof(typeof(u0)))" for u0 in ( - [1.0], @SVector[1.0], 1.0) - sol = benchmark_nlsolve_oop(quadratic_f, u0; solver = SimpleHouseholder{order}()) + @testset "[OOP] u0: $(nameof(typeof(u0)))" for u0 in ([1.0], @SVector[1.0], 1.0) + sol = benchmark_nlsolve_oop( + quadratic_f, u0; solver = SimpleHouseholder{order}()) @test SciMLBase.successful_retcode(sol) @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) end @testset "[IIP] u0: $(nameof(typeof(u0)))" for u0 in ([1.0],) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver = SimpleHouseholder{order}()) + sol = benchmark_nlsolve_iip( + quadratic_f!, u0; solver = SimpleHouseholder{order}()) @test SciMLBase.successful_retcode(sol) @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) end @@ -112,7 +113,8 @@ end u0 in (1.0, [1.0], @SVector[1.0]) probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, SimpleHouseholder{2}(); termination_condition).u .≈ sqrt(2.0)) + @test all(solve(probN, SimpleHouseholder{2}(); termination_condition).u .≈ + sqrt(2.0)) end end