Handle Approximate Sparsity detection more gracefully

Author: avik-pal

docs/pages.jl

@@ -14,6 +14,7 @@ pages = ["index.md",
         "basics/NonlinearSolution.md",
         "basics/TerminationCondition.md",
         "basics/Logging.md",
+        "basics/SparsityDetection.md",
         "basics/FAQ.md"],
     "Solver Summaries and Recommendations" => Any["solvers/NonlinearSystemSolvers.md",
         "solvers/BracketingSolvers.md",

docs/src/basics/SparsityDetection.md

@@ -0,0 +1,77 @@
+# [(Semi-)Automatic Sparsity Detection](@id sparsity-detection)
+
+This section describes how to enable Sparsity Detection. For a detailed tutorial on how
+to use this in an actual problem, see
+[this tutorial on Efficiently Solving Large Sparse Ill-Conditioned Nonlinear Systems](@ref large_systems).
+
+Notation wise we are trying to solve for `x` such that `nlfunc(x) = 0`.
+
+## Case I: Sparse Jacobian Prototype is Provided
+
+Let's say you have a Sparse Jacobian Prototype `jac_prototype`, in this case you can
+create your problem as follows:
+
+```julia
+prob = NonlinearProblem(NonlinearFunction(nlfunc; sparsity = jac_prototype), x0)
+# OR
+prob = NonlinearProblem(NonlinearFunction(nlfunc; jac_prototype = jac_prototype), x0)
+```
+
+NonlinearSolve will automatically perform matrix coloring and use sparse differentiation.
+
+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)
+```
+
+!!! 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
+    row colorvec.
+
+## Case II: Sparsity Detection algorithm is provided
+
+If you don't have a Sparse Jacobian Prototype, but you know the which sparsity detection
+algorithm you want to use, then you can create your problem as follows:
+
+```julia
+prob = NonlinearProblem(NonlinearFunction(nlfunc; sparsity = SymbolicsSparsityDetection()),
+    x0)  # Remember to have Symbolics.jl loaded
+# OR
+prob = NonlinearProblem(NonlinearFunction(nlfunc; sparsity = ApproximateJacobianSparsity()),
+    x0)
+```
+
+These Detection Algorithms are from [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl),
+refer to the documentation there for more details.
+
+## Case III: Sparse AD Type is being Used
+
+If you constructed a Nonlinear Solver with a sparse AD type, for example
+
+```julia
+NewtonRaphson(; autodiff = AutoSparseForwardDiff())
+# OR
+TrustRegion(; autodiff = AutoSparseZygote())
+```
+
+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`.

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
+using NonlinearSolve, LinearAlgebra, SparseArrays, LinearSolve, SparseDiffTools
 
 const N = 32
 const xyd_brusselator = range(0, stop = 1, length = N)
@@ -100,7 +100,8 @@ function init_brusselator_2d(xyd)
 end
 
 u0 = init_brusselator_2d(xyd_brusselator)
-prob_brusselator_2d = NonlinearProblem(brusselator_2d_loop, u0, p)
+prob_brusselator_2d = NonlinearProblem(brusselator_2d_loop, u0, p; abstol = 1e-10,
+    reltol = 1e-10)
 ```
 
 ## Choosing Jacobian Types
@@ -127,8 +128,10 @@ include:
 
 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. For this to work it is important to not load `Symbolics.jl`. This is triggered
-if you pass in a sparse autodiff type such as `AutoSparseForwardDiff()`.
+jacobians. This is triggered if you pass in a sparse autodiff type such as
+`AutoSparseForwardDiff()`. 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.
 
 ```@example ill_conditioned_nlprob
 using BenchmarkTools # for @btime
@@ -223,6 +226,7 @@ used in the solution of the ODE. An example of this with using
 
 ```@example ill_conditioned_nlprob
 using IncompleteLU
+
 function incompletelu(W, du, u, p, t, newW, Plprev, Prprev, solverdata)
     if newW === nothing || newW
         Pl = ilu(W, τ = 50.0)
@@ -257,6 +261,7 @@ which is more automatic. The setup is very similar to before:
 
 ```@example ill_conditioned_nlprob
 using AlgebraicMultigrid
+
 function algebraicmultigrid(W, du, u, p, t, newW, Plprev, Prprev, solverdata)
     if newW === nothing || newW
         Pl = aspreconditioner(ruge_stuben(convert(AbstractMatrix, W)))
@@ -293,5 +298,22 @@ end
 nothing # hide
 ```
 
+## Let's compare the Sparsity Detection Methods
+
+We benchmarked the solvers before with approximate and exact sparsity detection. However,
+for the exact sparsity detection case, we left out the time it takes to perform exact
+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()), u0, p; abstol = 1e-10, reltol = 1e-10)
+prob_brusselator_2d_approx = NonlinearProblem(NonlinearFunction(brusselator_2d_loop;
+        sparsity = ApproximateJacobianSparsity()), u0, p; abstol = 1e-10, reltol = 1e-10)
+
+@btime solve(prob_brusselator_2d_exact, NewtonRaphson());
+@btime solve(prob_brusselator_2d_approx, NewtonRaphson());
+nothing # hide
+```
+
 For more information on the preconditioner interface, see the
 [linear solver documentation](https://docs.sciml.ai/LinearSolve/stable/basics/Preconditioners/).

src/jacobian.jl

@@ -15,17 +15,27 @@ function sparsity_detection_alg(f, ad::AbstractSparseADType)
             if is_extension_loaded(Val(:Symbolics))
                 return SymbolicsSparsityDetection()
             else
-                @warn "Symbolics.jl is not loaded and sparse AD mode $(ad) is being used. \
-                       Using approximate sparsity detection using ForwardDiff. This can \
-                       potentially fail or generate incorrect sparsity pattern for \
-                       complicated problems. Use with caution." maxlog=1
                 return ApproximateJacobianSparsity()
             end
         else
             jac_prototype = f.jac_prototype
         end
-    else
+    elseif f.sparsity isa SparseDiffTools.AbstractSparsityDetection
+        if f.jac_prototype === nothing
+            return f.sparsity
+        else
+            jac_prototype = f.jac_prototype
+        end
+    elseif f.sparsity isa AbstractMatrix
         jac_prototype = f.sparsity
+    elseif f.jac_prototype isa AbstractMatrix
+        jac_prototype = f.jac_prototype
+    else
+        error("`sparsity::typeof($(typeof(f.sparsity)))` & \
+               `jac_prototype::typeof($(typeof(f.jac_prototype)))` is not supported. \
+               Use `sparsity::AbstractMatrix` or `sparsity::AbstractSparsityDetection` or \
+               set to `nothing`. `jac_prototype` can be set to `nothing` or an \
+               `AbstractMatrix`.")
     end
 
     if SciMLBase.has_colorvec(f)

src/utils.jl

@@ -384,6 +384,7 @@ LazyArrays.applied_axes(::typeof(__zero), x) = axes(x)
     end
     return pinv(A)
 end
+@inline __safe_inv(A::SparseMatrixCSC) = __safe_inv(Matrix(A))
 
 LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x)
 LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x)