Add tests

Author: avik-pal

ext/SimpleNonlinearSolvePolyesterForwardDiffExt.jl

@@ -10,4 +10,10 @@ using SimpleNonlinearSolve, PolyesterForwardDiff
     return J
 end
 
-end
+@inline function SimpleNonlinearSolve.__polyester_forwarddiff_jacobian!(f::F, J, x,
+        chunksize) where {F}
+    PolyesterForwardDiff.threaded_jacobian!(f, J, x, chunksize)
+    return J
+end
+
+end

src/nlsolve/halley.jl

@@ -12,15 +12,15 @@ A low-overhead implementation of Halley's Method.
 
 ### Keyword Arguments
 
-  - `autodiff`: determines the backend used for the Hessian. Defaults to
-    `AutoForwardDiff()`. Valid choices are `AutoForwardDiff()` or `AutoFiniteDiff()`.
+  - `autodiff`: determines the backend used for the Hessian. Defaults to `nothing`. Valid
+    choices are `AutoForwardDiff()` or `AutoFiniteDiff()`.
 
 !!! warning
 
     Inplace Problems are currently not supported by this method.
 """
 @kwdef @concrete struct SimpleHalley <: AbstractNewtonAlgorithm
-    autodiff = AutoForwardDiff()
+    autodiff = nothing
 end
 
 function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...;
@@ -33,6 +33,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...;
     fx = _get_fx(prob, x)
     T = eltype(x)
 
+    autodiff = __get_concrete_autodiff(prob, alg.autodiff; polyester = Val(false))
     abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fx, x,
         termination_condition)
 
@@ -50,7 +51,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...;
 
     for i in 1:maxiters
         # Hessian Computation is unfortunately type unstable
-        fx, dfx, d2fx = compute_jacobian_and_hessian(alg.autodiff, prob, fx, x)
+        fx, dfx, d2fx = compute_jacobian_and_hessian(autodiff, prob, fx, x)
         setindex_trait(x) === CannotSetindex() && (A = dfx)
 
         # Factorize Once and Reuse

src/utils.jl

@@ -119,6 +119,11 @@ function value_and_jacobian(ad, f::F, y, x::Number, p, cache; J = nothing) where
         T = typeof(__standard_tag(ad.tag, x))
         out = f(ForwardDiff.Dual{T}(x, one(x)), p)
         return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
+    elseif ad isa AutoPolyesterForwardDiff
+        # Just use ForwardDiff
+        T = typeof(__standard_tag(nothing, x))
+        out = f(ForwardDiff.Dual{T}(x, one(x)), p)
+        return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
     elseif ad isa AutoFiniteDiff
         _f = Base.Fix2(f, p)
         return _f(x), FiniteDiff.finite_difference_derivative(_f, x, ad.fdtype)
@@ -153,7 +158,7 @@ function jacobian_cache(ad, f::F, y, x::X, p) where {F, X <: AbstractArray}
             J = ArrayInterface.can_setindex(x) ? similar(y, length(y), length(x)) : nothing
             return J, __get_jacobian_config(ad, _f, x)
         elseif ad isa AutoPolyesterForwardDiff
-            @assert ArrayInterface.can_setindex(x) "PolyesterForwardDiff requires mutable inputs."
+            @assert ArrayInterface.can_setindex(x) "PolyesterForwardDiff requires mutable inputs. Use AutoForwardDiff instead."
             J = similar(y, length(y), length(x))
             return J, __get_jacobian_config(ad, _f, x)
         elseif ad isa AutoFiniteDiff
@@ -362,16 +367,17 @@ end
 end
 
 # Decide which AD backend to use
-@inline __get_concrete_autodiff(prob, ad::ADTypes.AbstractADType) = ad
-@inline function __get_concrete_autodiff(prob, ::Nothing)
+@inline __get_concrete_autodiff(prob, ad::ADTypes.AbstractADType; kwargs...) = ad
+@inline function __get_concrete_autodiff(prob, ::Nothing; polyester::Val{P} = Val(true),
+        kwargs...) where {P}
     if ForwardDiff.can_dual(eltype(prob.u0))
-        if __is_extension_loaded(Val(:PolyesterForwardDiff)) && !(prob.u0 isa Number) &&
-           ArrayInterface.can_setindex(prob.u0)
+        if P && __is_extension_loaded(Val(:PolyesterForwardDiff)) &&
+           !(prob.u0 isa Number) && ArrayInterface.can_setindex(prob.u0)
             return AutoPolyesterForwardDiff()
         else
             return AutoForwardDiff()
         end
     else
         return AutoFiniteDiff()
     end
-end
+end

test/Project.toml

@@ -7,6 +7,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

test/basictests.jl

@@ -1,5 +1,6 @@
 using AllocCheck, BenchmarkTools, LinearSolve, SimpleNonlinearSolve, StaticArrays, Random,
     LinearAlgebra, Test, ForwardDiff, DiffEqBase
+import PolyesterForwardDiff
 
 _nameof(x) = applicable(nameof, x) ? nameof(x) : _nameof(typeof(x))
 
@@ -29,20 +30,21 @@ const TERMINATION_CONDITIONS = [
 @testset "$(alg)" for alg in (SimpleNewtonRaphson, SimpleTrustRegion)
     # Eval else the alg is type unstable
     @eval begin
-        function benchmark_nlsolve_oop(f, u0, p = 2.0; autodiff = AutoForwardDiff())
+        function benchmark_nlsolve_oop(f, u0, p = 2.0; autodiff = nothing)
             prob = NonlinearProblem{false}(f, u0, p)
             return solve(prob, $(alg)(; autodiff), abstol = 1e-9)
         end
 
-        function benchmark_nlsolve_iip(f, u0, p = 2.0; autodiff = AutoForwardDiff())
+        function benchmark_nlsolve_iip(f, u0, p = 2.0; autodiff = nothing)
             prob = NonlinearProblem{true}(f, u0, p)
             return solve(prob, $(alg)(; autodiff), abstol = 1e-9)
         end
     end
 
     @testset "AutoDiff: $(_nameof(autodiff))" for autodiff in (AutoFiniteDiff(),
-        AutoForwardDiff())
+        AutoForwardDiff(), AutoPolyesterForwardDiff())
         @testset "[OOP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0], @SVector[1.0, 1.0], 1.0)
+            u0 isa SVector && autodiff isa AutoPolyesterForwardDiff && continue
             sol = benchmark_nlsolve_oop(quadratic_f, u0; autodiff)
             @test SciMLBase.successful_retcode(sol)
             @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9)
@@ -103,7 +105,7 @@ end
 # --- SimpleHalley tests ---
 
 @testset "SimpleHalley" begin
-    function benchmark_nlsolve_oop(f, u0, p = 2.0; autodiff = AutoForwardDiff())
+    function benchmark_nlsolve_oop(f, u0, p = 2.0; autodiff = nothing)
         prob = NonlinearProblem{false}(f, u0, p)
         return solve(prob, SimpleHalley(; autodiff), abstol = 1e-9)
     end