From 4d2f37a86a85b586ef4fbec0aa1c28881f4abb17 Mon Sep 17 00:00:00 2001
From: Songchen Tan <i@tansongchen.com>
Date: Fri, 1 Nov 2024 14:30:47 -0400
Subject: [PATCH] Make it a ext

---
 Project.toml                                 |  6 +++--
 ext/NonlinearSolveTaylorDiffExt.jl           | 19 +++++++++++++++
 lib/NonlinearSolveBase/src/descent/halley.jl | 25 +++++++-------------
 lib/NonlinearSolveFirstOrder/src/halley.jl   |  2 +-
 test/23_test_problems_tests.jl               |  1 +
 5 files changed, 34 insertions(+), 19 deletions(-)
 create mode 100644 ext/NonlinearSolveTaylorDiffExt.jl

diff --git a/Project.toml b/Project.toml
index 76ccfd389..9b8de1d6c 100644
--- a/Project.toml
+++ b/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"]
diff --git a/ext/NonlinearSolveTaylorDiffExt.jl b/ext/NonlinearSolveTaylorDiffExt.jl
new file mode 100644
index 000000000..100a9b0ff
--- /dev/null
+++ b/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
diff --git a/lib/NonlinearSolveBase/src/descent/halley.jl b/lib/NonlinearSolveBase/src/descent/halley.jl
index 596d78f19..fe6b67d5a 100644
--- a/lib/NonlinearSolveBase/src/descent/halley.jl
+++ b/lib/NonlinearSolveBase/src/descent/halley.jl
@@ -5,6 +5,8 @@ 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
@@ -12,8 +14,6 @@ See also [`NewtonDescent`](@ref).
     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")
diff --git a/lib/NonlinearSolveFirstOrder/src/halley.jl b/lib/NonlinearSolveFirstOrder/src/halley.jl
index bf9714f7b..d4a7e522d 100644
--- a/lib/NonlinearSolveFirstOrder/src/halley.jl
+++ b/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.
diff --git a/test/23_test_problems_tests.jl b/test/23_test_problems_tests.jl
index 8fa4c47b6..2075a1fcd 100644
--- a/test/23_test_problems_tests.jl
+++ b/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