feat: use DI for sparse jacobians

Author: avik-pal

Project.toml

@@ -25,6 +25,8 @@ 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"
@@ -102,6 +104,8 @@ 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"

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

@@ -2,15 +2,15 @@
 
 This tutorial is for getting into the extra features of using NonlinearSolve.jl. Solving
 ill-conditioned nonlinear systems requires specializing the linear solver on properties of
-the Jacobian in order to cut down on the ``\mathcal{O}(n^3)`` linear solve and the
-``\mathcal{O}(n^2)`` back-solves. This tutorial is designed to explain the advanced usage of
+the Jacobian in order to cut down on the `\mathcal{O}(n^3)` linear solve and the
+`\mathcal{O}(n^2)` back-solves. This tutorial is designed to explain the advanced usage of
 NonlinearSolve.jl by solving the steady state stiff Brusselator partial differential
 equation (BRUSS) using NonlinearSolve.jl.
 
 ## Definition of the Brusselator Equation
 
 !!! note
-    
+
     Feel free to skip this section: it simply defines the example problem.
 
 The Brusselator PDE is defined as follows:
@@ -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)
@@ -117,31 +117,37 @@ However, if you know the sparsity of your problem, then you can pass a different
 type. For example, a `SparseMatrixCSC` will give a sparse matrix. Other sparse matrix types
 include:
 
-  - Bidiagonal
-  - Tridiagonal
-  - SymTridiagonal
-  - BandedMatrix ([BandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BandedMatrices.jl))
-  - BlockBandedMatrix ([BlockBandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BlockBandedMatrices.jl))
+- Bidiagonal
+- Tridiagonal
+- SymTridiagonal
+- BandedMatrix ([BandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BandedMatrices.jl))
+- BlockBandedMatrix ([BlockBandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BlockBandedMatrices.jl))
 
 ## 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`:
@@ -212,7 +224,7 @@ choices, see the
 `linsolve` choices are any valid [LinearSolve.jl](https://linearsolve.sciml.ai/dev/) solver.
 
 !!! note
-    
+
     Switching to a Krylov linear solver will automatically change the nonlinear problem
     solver into Jacobian-free mode, dramatically reducing the memory required. This can be
     overridden by adding `concrete_jac=true` to the algorithm.
@@ -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

@@ -31,31 +31,33 @@ 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 SparseDiffTools: SparseDiffTools, AbstractSparsityDetection,
-                       ApproximateJacobianSparsity, JacPrototypeSparsityDetection,
-                       NoSparsityDetection, PrecomputedJacobianColorvec,
-                       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
+               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
+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