feat: use DifferentiationInterface for sparse AD

Project.toml

@@ -25,9 +25,13 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
@@ -42,7 +46,6 @@ NLSolvers = "337daf1e-9722-11e9-073e-8b9effe078ba"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
-Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
@@ -55,7 +58,6 @@ NonlinearSolveNLSolversExt = "NLSolvers"
 NonlinearSolveNLsolveExt = "NLsolve"
 NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
 NonlinearSolveSpeedMappingExt = "SpeedMapping"
-NonlinearSolveSymbolicsExt = "Symbolics"
 NonlinearSolveZygoteExt = "Zygote"
 
 [compat]
@@ -102,9 +104,13 @@ Reexport = "1.2"
 SIAMFANLEquations = "1.0.1"
 SciMLBase = "2.54.0"
 SciMLJacobianOperators = "0.1"
+SciMLOperators = "0.3.10"
+Setfield = "1.1.1"
 SimpleNonlinearSolve = "1.12.3"
 SparseArrays = "1.10"
+SparseConnectivityTracer = "0.6.5"
 SparseDiffTools = "2.22"
+SparseMatrixColorings = "0.4.2"
 SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.9"

docs/Project.toml

@@ -3,6 +3,7 @@ AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
@@ -16,6 +17,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
@@ -27,6 +29,7 @@ AlgebraicMultigrid = "0.5, 0.6"
 ArrayInterface = "6, 7"
 BenchmarkTools = "1"
 DiffEqBase = "6.136"
+DifferentiationInterface = "0.6"
 Documenter = "1"
 DocumenterCitations = "1"
 IncompleteLU = "0.2"
@@ -40,6 +43,7 @@ Random = "<0.0.1, 1"
 SciMLBase = "2.4"
 SciMLJacobianOperators = "0.1"
 SimpleNonlinearSolve = "1"
+SparseConnectivityTracer = "0.6.5"
 SparseDiffTools = "2.14"
 StaticArrays = "1"
 SteadyStateDiffEq = "2"

docs/src/basics/sparsity_detection.md

@@ -12,8 +12,6 @@ Let's say you have a Sparse Jacobian Prototype `jac_prototype`, in this case you
 create your problem as follows:
 
 ```julia
-prob = NonlinearProblem(NonlinearFunction(nlfunc; sparsity = jac_prototype), x0)
-# OR
 prob = NonlinearProblem(NonlinearFunction(nlfunc; jac_prototype = jac_prototype), x0)
 ```
 
@@ -22,9 +20,6 @@ NonlinearSolve will automatically perform matrix coloring and use sparse differe
 Now you can help the solver further by providing the color vector. This can be done by
 
 ```julia
-prob = NonlinearProblem(
-    NonlinearFunction(nlfunc; sparsity = jac_prototype, colorvec = colorvec), x0)
-# OR
 prob = NonlinearProblem(
     NonlinearFunction(nlfunc; jac_prototype = jac_prototype, colorvec = colorvec), x0)
 ```
@@ -34,10 +29,16 @@ If the `colorvec` is not provided, then it is computed on demand.
 !!! note
 
     One thing to be careful about in this case is that `colorvec` is dependent on the
-    autodiff backend used. Forward Mode and Finite Differencing will assume that the
-    colorvec is the column colorvec, while Reverse Mode will assume that the colorvec is the
+    autodiff backend used. `ADTypes.mode(ad) isa ADTypes.ForwardMode` will assume that the
+    colorvec is the column colorvec, otherwise we will assume that the colorvec is the
     row colorvec.
 
+!!! warning
+
+    Previously you could provide a `sparsity` argument to `NonlinearFunction` to specify
+    the jacobian prototype. However, to avoid confusion, this is now deprecated. Instead,
+    use the `jac_prototype` argument.
+
 ## Case II: Sparsity Detection algorithm is provided
 
 If you don't have a Sparse Jacobian Prototype, but you know the which sparsity detection
@@ -48,14 +49,18 @@ prob = NonlinearProblem(
     NonlinearFunction(nlfunc; sparsity = SymbolicsSparsityDetection()), x0)  # Remember to have Symbolics.jl loaded
 # OR
 prob = NonlinearProblem(
-    NonlinearFunction(nlfunc; sparsity = ApproximateJacobianSparsity()), x0)
+    NonlinearFunction(nlfunc; sparsity = TracerSparsityDetector()), x0) # From SparseConnectivityTracer.jl
 ```
 
-These Detection Algorithms are from [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl),
-refer to the documentation there for more details.
+Refer to the documentation of DifferentiationInterface.jl for more information on
+sparsity detection algorithms.
 
 ## Case III: Sparse AD Type is being Used
 
+!!! warning
+
+    This is now deprecated. Please use the previous two cases instead.
+
 If you constructed a Nonlinear Solver with a sparse AD type, for example
 
 ```julia
@@ -65,15 +70,6 @@ TrustRegion(; autodiff = AutoSparse(AutoZygote()))
 ```
 
 then NonlinearSolve will automatically perform matrix coloring and use sparse
-differentiation if none of `sparsity` or `jac_prototype` is provided. If `Symbolics.jl` is
-loaded, we default to using `SymbolicsSparsityDetection()`, else we default to using
-`ApproximateJacobianSparsity()`.
-
-`Case I/II` take precedence for sparsity detection and we perform sparse AD based on those
-options if those are provided.
-
-!!! warning
-
-    If you provide a non-sparse AD, and provide a `sparsity` or `jac_prototype` then
-    we will use dense AD. This is because, if you provide a specific AD type, we assume
-    that you know what you are doing and want to override the default choice of `nothing`.
+differentiation if none of `sparsity` or `jac_prototype` is provided. We default to using
+`TracerSparsityDetector()`. `Case I/II` take precedence for sparsity detection and we
+perform sparse AD based on those options if those are provided.

docs/src/tutorials/large_systems.md

@@ -60,7 +60,7 @@ broadcast). Use `dx=dy=1/32`.
 The resulting `NonlinearProblem` definition is:
 
 ```@example ill_conditioned_nlprob
-using NonlinearSolve, LinearAlgebra, SparseArrays, LinearSolve, SparseDiffTools
+using NonlinearSolve, LinearAlgebra, SparseArrays, LinearSolve
 
 const N = 32
 const xyd_brusselator = range(0, stop = 1, length = N)
@@ -125,23 +125,29 @@ include:
 
 ## Approximate Sparsity Detection & Sparse Jacobians
 
-In the next section, we will discuss how to declare a sparse Jacobian and how to use
-[Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl), to compute exact sparse
-jacobians. This is triggered if you pass in a sparse autodiff type such as
-`AutoSparse(AutoForwardDiff())`. If `Symbolics.jl` is loaded, then the default changes to
-Symbolic Sparsity Detection. See the manual entry on
-[Sparsity Detection](@ref sparsity-detection) for more details on the default.
+In the next section, we will show how to specify `sparsity` to trigger automatic sparsity
+detection.
 
 ```@example ill_conditioned_nlprob
 using BenchmarkTools # for @btime
 
 @btime solve(prob_brusselator_2d, NewtonRaphson());
-@btime solve(prob_brusselator_2d,
-    NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32))));
-@btime solve(prob_brusselator_2d,
-    NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32)),
+nothing # hide
+```
+
+```@example ill_conditioned_nlprob
+using SparseConnectivityTracer
+
+prob_brusselator_2d_autosparse = NonlinearProblem(
+    NonlinearFunction(brusselator_2d_loop; sparsity = TracerSparsityDetector()),
+    u0, p; abstol = 1e-10, reltol = 1e-10)
+
+@btime solve(prob_brusselator_2d_autosparse,
+    NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 32)));
+@btime solve(prob_brusselator_2d_autosparse,
+    NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 32),
         linsolve = KLUFactorization()));
-@btime solve(prob_brusselator_2d,
+@btime solve(prob_brusselator_2d_autosparse,
     NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 32),
         linsolve = KrylovJL_GMRES()));
 nothing # hide
@@ -160,8 +166,14 @@ declaration of Jacobian sparsity types. To see this in action, we can give an ex
 and `u` and call `jacobian_sparsity` on our function with the example arguments, and it will
 kick out a sparse matrix with our pattern, that we can turn into our `jac_prototype`.
 
+!!! tip
+
+    Alternatively you can use the `SparseConnectivityTracer.jl` package to automatically
+    generate a sparse Jacobian.
+
 ```@example ill_conditioned_nlprob
 using Symbolics
+
 du0 = copy(u0)
 jac_sparsity = Symbolics.jacobian_sparsity(
     (du, u) -> brusselator_2d_loop(du, u, p), du0, u0)
@@ -171,7 +183,7 @@ Notice that Julia gives a nice print out of the sparsity pattern. That's neat, a
 tedious to build by hand! Now we just pass it to the `NonlinearFunction` like as before:
 
 ```@example ill_conditioned_nlprob
-ff = NonlinearFunction(brusselator_2d_loop; sparsity = jac_sparsity)
+ff = NonlinearFunction(brusselator_2d_loop; jac_prototype = jac_sparsity)
 ```
 
 Build the `NonlinearProblem`:
@@ -275,8 +287,8 @@ function algebraicmultigrid(W, du, u, p, t, newW, Plprev, Prprev, solverdata)
 end
 
 @btime solve(prob_brusselator_2d_sparse,
-    NewtonRaphson(
-        linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid, concrete_jac = true));
+    NewtonRaphson(linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid,
+        concrete_jac = true));
 nothing # hide
 ```
 
@@ -296,8 +308,8 @@ function algebraicmultigrid2(W, du, u, p, t, newW, Plprev, Prprev, solverdata)
 end
 
 @btime solve(prob_brusselator_2d_sparse,
-    NewtonRaphson(
-        linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid2, concrete_jac = true));
+    NewtonRaphson(linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid2,
+        concrete_jac = true));
 nothing # hide
 ```
 
@@ -308,15 +320,22 @@ for the exact sparsity detection case, we left out the time it takes to perform
 sparsity detection. Let's compare the two by setting the sparsity detection algorithms.
 
 ```@example ill_conditioned_nlprob
-prob_brusselator_2d_exact = NonlinearProblem(
-    NonlinearFunction(brusselator_2d_loop; sparsity = SymbolicsSparsityDetection()),
+using DifferentiationInterface, SparseConnectivityTracer
+
+prob_brusselator_2d_exact_symbolics = NonlinearProblem(
+    NonlinearFunction(brusselator_2d_loop; sparsity = SymbolicsSparsityDetector()),
+    u0, p; abstol = 1e-10, reltol = 1e-10)
+prob_brusselator_2d_exact_tracer = NonlinearProblem(
+    NonlinearFunction(brusselator_2d_loop; sparsity = TracerSparsityDetector()),
     u0, p; abstol = 1e-10, reltol = 1e-10)
-prob_brusselator_2d_approx = NonlinearProblem(
-    NonlinearFunction(brusselator_2d_loop; sparsity = ApproximateJacobianSparsity()),
+prob_brusselator_2d_approx_di = NonlinearProblem(
+    NonlinearFunction(brusselator_2d_loop;
+        sparsity = DenseSparsityDetector(AutoForwardDiff(); atol = 1e-4)),
     u0, p; abstol = 1e-10, reltol = 1e-10)
 
-@btime solve(prob_brusselator_2d_exact, NewtonRaphson());
-@btime solve(prob_brusselator_2d_approx, NewtonRaphson());
+@btime solve(prob_brusselator_2d_exact_symbolics, NewtonRaphson());
+@btime solve(prob_brusselator_2d_exact_tracer, NewtonRaphson());
+@btime solve(prob_brusselator_2d_approx_di, NewtonRaphson());
 nothing # hide
 ```
 

src/NonlinearSolve.jl

@@ -7,12 +7,6 @@ end
 using Reexport: @reexport
 using PrecompileTools: @compile_workload, @setup_workload
 
-using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff,
-               AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, AutoSparse
-# FIXME: deprecated, remove in future
-using ADTypes: AutoSparseFiniteDiff, AutoSparseForwardDiff, AutoSparsePolyesterForwardDiff,
-               AutoSparseZygote
-
 using ArrayInterface: ArrayInterface, can_setindex, restructure, fast_scalar_indexing,
                       ismutable
 using ConcreteStructs: @concrete
@@ -22,11 +16,8 @@ using DiffEqBase: DiffEqBase, AbstractNonlinearTerminationMode,
                   NormTerminationMode, RelNormTerminationMode, RelSafeBestTerminationMode,
                   RelSafeTerminationMode, RelTerminationMode,
                   SimpleNonlinearSolveTerminationMode, SteadyStateDiffEqTerminationMode
-using DifferentiationInterface: DifferentiationInterface, Constant
 using FastBroadcast: @..
 using FastClosures: @closure
-using FiniteDiff: FiniteDiff
-using ForwardDiff: ForwardDiff, Dual
 using LazyArrays: LazyArrays, ApplyArray, cache
 using LinearAlgebra: LinearAlgebra, ColumnNorm, Diagonal, I, LowerTriangular, Symmetric,
                      UpperTriangular, axpy!, cond, diag, diagind, dot, issuccess, istril,
@@ -40,20 +31,34 @@ using Printf: @printf
 using Preferences: Preferences, @load_preference, @set_preferences!
 using RecursiveArrayTools: recursivecopy!
 using SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
-                 AbstractSciMLOperator, _unwrap_val, isinplace, NLStats
-using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJacOperator,
-                              JacVecOperator, StatefulJacobianOperator
-using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC
-using SparseDiffTools: SparseDiffTools, AbstractSparsityDetection,
-                       ApproximateJacobianSparsity, JacPrototypeSparsityDetection,
-                       NoSparsityDetection, PrecomputedJacobianColorvec,
-                       SymbolicsSparsityDetection, init_jacobian, sparse_jacobian,
-                       sparse_jacobian!, sparse_jacobian_cache
+                 _unwrap_val, isinplace, NLStats
+using SciMLOperators: AbstractSciMLOperator
+using Setfield: @set!
 using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix
 using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy,
                                  symbolic_container, parameter_values, state_values, getu,
                                  setu
 
+# AD Support
+using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff,
+               AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, AutoSparse,
+               NoSparsityDetector, KnownJacobianSparsityDetector
+using ADTypes: AutoSparseFiniteDiff, AutoSparseForwardDiff, AutoSparsePolyesterForwardDiff,
+               AutoSparseZygote # FIXME: deprecated, remove in future
+using DifferentiationInterface: DifferentiationInterface, Constant
+using FiniteDiff: FiniteDiff
+using ForwardDiff: ForwardDiff, Dual
+using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJacOperator,
+                              JacVecOperator, StatefulJacobianOperator
+
+## Sparse AD Support
+using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC
+using SparseConnectivityTracer: TracerSparsityDetector # This can be dropped in the next release
+using SparseDiffTools: SparseDiffTools, JacPrototypeSparsityDetection,
+                       PrecomputedJacobianColorvec, init_jacobian, sparse_jacobian,
+                       sparse_jacobian!, sparse_jacobian_cache
+using SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm
+
 @reexport using SciMLBase, SimpleNonlinearSolve
 
 const DI = DifferentiationInterface