Make it a ext

Author: tansongchen

Project.toml

@@ -44,6 +44,7 @@ PETSc = "ace2c81b-2b5f-4b1e-a30d-d662738edfe0"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
+TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c"
 
 [extensions]
 NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
@@ -56,6 +57,7 @@ NonlinearSolvePETScExt = ["PETSc", "MPI"]
 NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
 NonlinearSolveSpeedMappingExt = "SpeedMapping"
 NonlinearSolveSundialsExt = "Sundials"
+NonlinearSolveTaylorDiffExt = "TaylorDiff"
 
 [compat]
 ADTypes = "1.9"
@@ -113,7 +115,6 @@ StaticArrays = "1.9"
 StaticArraysCore = "1.4"
 Sundials = "4.23.1"
 SymbolicIndexingInterface = "0.3.31"
-Symbolics = "6"
 TaylorDiff = "0.3"
 Test = "1.10"
 Zygote = "0.6.69"
@@ -148,8 +149,9 @@ SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
+TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote"]
+test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "TaylorDiff", "Test", "Zygote"]

ext/NonlinearSolveTaylorDiffExt.jl

@@ -0,0 +1,19 @@
+module NonlinearSolveTaylorDiffExt
+using NonlinearSolve: HalleyDescentCache, NonlinearFunction
+import NonlinearSolve: evaluate_hvvp
+using TaylorDiff: derivative, derivative!
+using FastClosures: @closure
+
+function evaluate_hvvp(
+        hvvp, cache::HalleyDescentCache, f::NonlinearFunction{iip}, p, u, δu) where {iip}
+    if iip
+        binary_f = @closure (y, x) -> f(y, x, p)
+        derivative!(hvvp, binary_f, cache.fu, u, δu, Val(2))
+    else
+        unary_f = Base.Fix2(f, p)
+        hvvp = derivative(unary_f, u, δu, Val(2))
+    end
+    hvvp
+end
+
+end

lib/NonlinearSolveBase/src/descent/halley.jl

@@ -5,15 +5,15 @@ Improve the NewtonDescent with higher-order terms. First compute the descent dir
 Then compute the hessian-vector-vector product and solve for the second-order correction term as ``J b = H a a``.
 Finally, compute the descent direction as ``δu = a * a / (b / 2 - a)``.
 
+Note that `import TaylorDiff` is required to use this descent algorithm.
+
 See also [`NewtonDescent`](@ref).
 """
 @kwdef @concrete struct HalleyDescent <: AbstractDescentAlgorithm
     linsolve = nothing
     precs = DEFAULT_PRECS
 end
 
-using TaylorDiff: derivative
-
 function Base.show(io::IO, d::HalleyDescent)
     modifiers = String[]
     d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)")
@@ -30,6 +30,7 @@ supports_line_search(::HalleyDescent) = true
     δus
     b
     fu
+    hvvp
     lincache
     timer
 end
@@ -43,13 +44,14 @@ function __internal_init(prob::NonlinearProblem, alg::HalleyDescent, J, fu, u; s
     @bb δu = similar(u)
     @bb b = similar(u)
     @bb fu = similar(fu)
+    @bb hvvp = similar(fu)
     δus = N ≤ 1 ? nothing : map(2:N) do i
         @bb δu_ = similar(u)
     end
-    INV && return HalleyDescentCache{true}(prob.f, prob.p, δu, δus, b, nothing, timer)
-    lincache = LinearSolverCache(
+    lincache = INV ? nothing :
+               LinearSolverCache(
         alg, alg.linsolve, J, _vec(fu), _vec(u); stats, abstol, reltol, linsolve_kwargs...)
-    return HalleyDescentCache{false}(prob.f, prob.p, δu, δus, b, fu, lincache, timer)
+    return HalleyDescentCache{false}(prob.f, prob.p, δu, δus, b, fu, hvvp, lincache, timer)
 end
 
 function __internal_solve!(cache::HalleyDescentCache{INV}, J, fu, u, idx::Val = Val(1);
@@ -73,7 +75,7 @@ function __internal_solve!(cache::HalleyDescentCache{INV}, J, fu, u, idx::Val =
     end
     b = cache.b
     # compute the hessian-vector-vector product
-    hvvp = evaluate_hvvp(cache, cache.f, cache.p, u, δu)
+    hvvp = evaluate_hvvp(cache.hvvp, cache, cache.f, cache.p, u, δu)
     # second linear solve, reuse factorization if possible
     if INV
         @bb b = J × vec(hvvp)
@@ -94,13 +96,4 @@ function __internal_solve!(cache::HalleyDescentCache{INV}, J, fu, u, idx::Val =
     return DescentResult(; δu)
 end
 
-function evaluate_hvvp(
-        cache::HalleyDescentCache, f::NonlinearFunction{iip}, p, u, δu) where {iip}
-    if iip
-        binary_f = @closure (y, x) -> f(y, x, p)
-        derivative(binary_f, cache.fu, u, δu, Val{3}())
-    else
-        unary_f = Base.Fix2(f, p)
-        derivative(unary_f, u, δu, Val{3}())
-    end
-end
+evaluate_hvvp(hvvp, cache, f, p, u, δu) = error("not implemented. please import TaylorDiff")

lib/NonlinearSolveFirstOrder/src/halley.jl

@@ -1,5 +1,5 @@
 """
-    Halley(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(),
+    Halley(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing,
         precs = DEFAULT_PRECS, autodiff = nothing)
 
 An experimental Halley's method implementation. Improves the convergence rate of Newton's method by using second-order derivative information to correct the descent direction.

test/23_test_problems_tests.jl

@@ -1,5 +1,6 @@
 @testsetup module RobustnessTesting
 using NonlinearSolve, LinearAlgebra, LinearSolve, NonlinearProblemLibrary, Test
+import TaylorDiff
 
 problems = NonlinearProblemLibrary.problems
 dicts = NonlinearProblemLibrary.dicts