Default to polyester forward diff

SciML · Jan 16, 2024 · d50f848 · d50f848
1 parent 87750c4
commit d50f848
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.5.0"
+version = "3.6.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -36,6 +36,7 @@ FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
@@ -48,6 +49,7 @@ NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
 NonlinearSolveMINPACKExt = "MINPACK"
 NonlinearSolveNLsolveExt = "NLsolve"
+NonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
 NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
 NonlinearSolveSpeedMappingExt = "SpeedMapping"
 NonlinearSolveSymbolicsExt = "Symbolics"
@@ -83,6 +85,7 @@ Pkg = "1.10"
 PrecompileTools = "1.2"
 Preferences = "1.4"
 Printf = "1.10"
+PolyesterForwardDiff = "0.1.1"
 Random = "1.91"
 RecursiveArrayTools = "3.4"
 Reexport = "1.2"
@@ -91,7 +94,7 @@ SafeTestsets = "0.1"
 SciMLBase = "2.19.0"
 SimpleNonlinearSolve = "1.2"
 SparseArrays = "1.10"
-SparseDiffTools = "2.14"
+SparseDiffTools = "2.16"
 SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.7"
@@ -121,6 +124,7 @@ NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -134,4 +138,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration", "SIAMFANLEquations", "Sundials"]
+test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration", "SIAMFANLEquations", "Sundials", "PolyesterForwardDiff"]
diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md
@@ -61,6 +61,7 @@ The resulting `NonlinearProblem` definition is:
 
 ```@example ill_conditioned_nlprob
 using NonlinearSolve, LinearAlgebra, SparseArrays, LinearSolve, SparseDiffTools
+import PolyesterForwardDiff
 
 const N = 32
 const xyd_brusselator = range(0, stop = 1, length = N)

diff --git a/ext/NonlinearSolvePolyesterForwardDiffExt.jl b/ext/NonlinearSolvePolyesterForwardDiffExt.jl
@@ -0,0 +1,7 @@
+module NonlinearSolvePolyesterForwardDiffExt
+
+using NonlinearSolve, PolyesterForwardDiff
+
+NonlinearSolve.is_extension_loaded(::Val{:PolyesterForwardDiff}) = true
+
+end
diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl
@@ -63,7 +63,14 @@ function get_concrete_forward_ad(autodiff, prob, sp::Val{test_sparse} = True, ar
         use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff()
     else
         tag = ForwardDiff.Tag(NonlinearSolveTag(), eltype(prob.u0))
-        (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)(; tag)
+        if use_sparse_ad
+            AutoSparseForwardDiff(; tag) # Polyester Sparse Mode is not implemented yet
+        elseif is_extension_loaded(Val(:PolyesterForwardDiff)) &&
+               !(prob.u0 isa Number || prob.u0 isa SArray)
+            AutoPolyesterForwardDiff()
+        else
+            AutoForwardDiff(; tag)
+        end
     end
     return ad
 end