Use DifferentiationInterface

Author: avik-pal

Project.toml

@@ -1,14 +1,15 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.8.1"
+version = "1.9.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -21,15 +22,13 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 SimpleNonlinearSolveChainRulesCoreExt = "ChainRulesCore"
-SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
 SimpleNonlinearSolveReverseDiffExt = "ReverseDiff"
 SimpleNonlinearSolveStaticArraysExt = "StaticArrays"
 SimpleNonlinearSolveTrackerExt = "Tracker"
@@ -45,6 +44,7 @@ ChainRulesCore = "1.22"
 ConcreteStructs = "0.2.3"
 DiffEqBase = "6.149"
 DiffResults = "1.1"
+DifferentiationInterface = "0.4"
 ExplicitImports = "1.5.0"
 FastClosures = "0.3.2"
 FiniteDiff = "2.22"

ext/SimpleNonlinearSolvePolyesterForwardDiffExt.jl

@@ -3,13 +3,15 @@ module SimpleNonlinearSolve
 using PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidations
 
 @recompile_invalidations begin
-    using ADTypes: ADTypes, AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff
+    using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff,
+                   AutoPolyesterForwardDiff
     using ArrayInterface: ArrayInterface
     using ConcreteStructs: @concrete
     using DiffEqBase: DiffEqBase, AbstractNonlinearTerminationMode,
                       AbstractSafeNonlinearTerminationMode,
                       AbstractSafeBestNonlinearTerminationMode, AbsNormTerminationMode,
                       NONLINEARSOLVE_DEFAULT_NORM
+    using DifferentiationInterface: DifferentiationInterface
     using DiffResults: DiffResults
     using FastClosures: @closure
     using FiniteDiff: FiniteDiff
@@ -25,6 +27,8 @@ using PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidati
     using StaticArraysCore: StaticArray, SVector, SMatrix, SArray, MArray, Size
 end
 
+const DI = DifferentiationInterface
+
 @reexport using SciMLBase
 
 abstract type AbstractSimpleNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end

src/SimpleNonlinearSolve.jl

@@ -12,8 +12,9 @@ A low-overhead implementation of Halley's Method.
 
 ### Keyword Arguments
 
-  - `autodiff`: determines the backend used for the Hessian. Defaults to `nothing`. Valid
-    choices are `AutoForwardDiff()` or `AutoFiniteDiff()`.
+  - `autodiff`: determines the backend used for the Hessian. Defaults to `nothing` (i.e.
+    automatic backend selection). Valid choices include backends from
+    `DifferentiationInterface.jl`.
 
 !!! warning
 
@@ -26,13 +27,11 @@ end
 function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...;
         abstol = nothing, reltol = nothing, maxiters = 1000,
         alias_u0 = false, termination_condition = nothing, kwargs...)
-    isinplace(prob) &&
-        error("SimpleHalley currently only supports out-of-place nonlinear problems")
-
     x = __maybe_unaliased(prob.u0, alias_u0)
     fx = _get_fx(prob, x)
     T = eltype(x)
 
+    f = __fixed_parameter_function(prob)
     autodiff = __get_concrete_autodiff(prob, alg.autodiff)
     abstol, reltol, tc_cache = init_termination_cache(
         prob, abstol, reltol, fx, x, termination_condition)
@@ -51,7 +50,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(autodiff, prob, fx, x)
+        fx, dfx, d2fx = compute_jacobian_and_hessian(autodiff, prob, f, fx, x)
         setindex_trait(x) === CannotSetindex() && (A = dfx)
 
         # Factorize Once and Reuse
@@ -78,9 +77,8 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...;
         cᵢ = _restructure(cᵢ, cᵢ_)
 
         if i == 1
-            if iszero(fx)
+            iszero(fx) &&
                 return build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
-            end
         else
             # Termination Checks
             tc_sol = check_termination(tc_cache, fx, x, xo, prob, alg)

src/nlsolve/halley.jl

@@ -13,9 +13,9 @@ and static array problems.
 
 ### Keyword Arguments
 
-  - `autodiff`: determines the backend used for the Jacobian. Defaults to
-    `nothing`. Valid choices are `AutoPolyesterForwardDiff()`, `AutoForwardDiff()` or
-    `AutoFiniteDiff()`.
+  - `autodiff`: determines the backend used for the Jacobian. Defaults to  `nothing` (i.e.
+    automatic backend selection). Valid choices include jacobian backends from
+    `DifferentiationInterface.jl`.
 """
 @kwdef @concrete struct SimpleNewtonRaphson <: AbstractNewtonAlgorithm
     autodiff = nothing
@@ -30,13 +30,14 @@ function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresPr
     fx = _get_fx(prob, x)
     autodiff = __get_concrete_autodiff(prob, alg.autodiff)
     @bb xo = copy(x)
-    J, jac_cache = jacobian_cache(autodiff, prob.f, fx, x, prob.p)
+    f = __fixed_parameter_function(prob)
+    J, jac_cache = jacobian_cache(autodiff, prob, f, fx, x)
 
     abstol, reltol, tc_cache = init_termination_cache(
         prob, abstol, reltol, fx, x, termination_condition)
 
     for i in 1:maxiters
-        fx, dfx = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J)
+        fx, dfx = value_and_jacobian(autodiff, prob, f, fx, x, jac_cache; J)
 
         if i == 1
             iszero(fx) && build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)

src/nlsolve/raphson.jl

@@ -10,9 +10,9 @@ scalar and static array problems.
 
 ### Keyword Arguments
 
-  - `autodiff`: determines the backend used for the Jacobian. Defaults to
-    `nothing`. Valid choices are `AutoPolyesterForwardDiff()`, `AutoForwardDiff()` or
-    `AutoFiniteDiff()`.
+  - `autodiff`: determines the backend used for the Jacobian. Defaults to `nothing` (i.e.
+    automatic backend selection). Valid choices include jacobian backends from
+    `DifferentiationInterface.jl`.
   - `max_trust_radius`: the maximum radius of the trust region. Defaults to
     `max(norm(f(u0)), maximum(u0) - minimum(u0))`.
   - `initial_trust_radius`: the initial trust region radius. Defaults to
@@ -85,8 +85,9 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleTrustRegion, args.
     fx = _get_fx(prob, x)
     norm_fx = norm(fx)
     @bb xo = copy(x)
-    J, jac_cache = jacobian_cache(autodiff, prob.f, fx, x, prob.p)
-    fx, ∇f = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J)
+    f = __fixed_parameter_function(prob)
+    J, jac_cache = jacobian_cache(autodiff, prob, f, fx, x)
+    fx, ∇f = value_and_jacobian(autodiff, prob, f, fx, x, jac_cache; J)
 
     abstol, reltol, tc_cache = init_termination_cache(
         prob, abstol, reltol, fx, x, termination_condition)
@@ -144,7 +145,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleTrustRegion, args.
             # Take the step.
             @bb @. xo = x
 
-            fx, ∇f = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J)
+            fx, ∇f = value_and_jacobian(autodiff, prob, f, fx, x, jac_cache; J)
 
             # Update the trust region radius.
             if !_unwrap_val(alg.nlsolve_update_rule) && r > η₃