Skip to content

Commit

Permalink
Merge pull request #320 from SciML/ap/approx_sparsity
Browse files Browse the repository at this point in the history
Use approximate sparsity detection by default
  • Loading branch information
ChrisRackauckas authored Dec 12, 2023
2 parents d7ef4af + 540ef5b commit 89bd481
Show file tree
Hide file tree
Showing 11 changed files with 165 additions and 14 deletions.
4 changes: 3 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "NonlinearSolve"
uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
authors = ["SciML"]
version = "3.0.1"
version = "3.0.2"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
Expand Down Expand Up @@ -35,6 +35,7 @@ FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

[extensions]
Expand All @@ -43,6 +44,7 @@ NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
NonlinearSolveMINPACKExt = "MINPACK"
NonlinearSolveNLsolveExt = "NLsolve"
NonlinearSolveSymbolicsExt = "Symbolics"
NonlinearSolveZygoteExt = "Zygote"

[compat]
Expand Down
6 changes: 3 additions & 3 deletions docs/make.jl
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,11 @@ include("pages.jl")

makedocs(; sitename = "NonlinearSolve.jl",
authors = "Chris Rackauckas",
modules = [NonlinearSolve, SciMLBase, DiffEqBase, SimpleNonlinearSolve, Sundials,
SteadyStateDiffEq],
modules = [NonlinearSolve, SimpleNonlinearSolve, SteadyStateDiffEq, Sundials,
DiffEqBase, SciMLBase],
clean = true, doctest = false, linkcheck = true,
linkcheck_ignore = ["https://twitter.com/ChrisRackauckas/status/1544743542094020615"],
warnonly = [:missing_docs, :cross_references],
warnonly = [:cross_references], checkdocs = :export,
format = Documenter.HTML(assets = ["assets/favicon.ico"],
canonical = "https://docs.sciml.ai/NonlinearSolve/stable/"),
pages)
Expand Down
1 change: 1 addition & 0 deletions docs/pages.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
3 changes: 3 additions & 0 deletions docs/src/api/nonlinearsolve.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ PseudoTransient
DFSane
Broyden
Klement
LimitedMemoryBroyden
```

## Nonlinear Least Squares Solvers
Expand Down Expand Up @@ -50,4 +51,6 @@ RadiusUpdateSchemes.Hei
RadiusUpdateSchemes.Yuan
RadiusUpdateSchemes.Bastin
RadiusUpdateSchemes.Fan
RadiusUpdateSchemes.NLsolve
RadiusUpdateSchemes.NocedalWright
```
1 change: 1 addition & 0 deletions docs/src/api/simplenonlinearsolve.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ i.e. `IntervalNonlinearProblem`.

```@docs
ITP
Alefeld
Bisection
Falsi
Ridder
Expand Down
77 changes: 77 additions & 0 deletions docs/src/basics/SparsityDetection.md
Original file line number Diff line number Diff line change
@@ -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`.
55 changes: 48 additions & 7 deletions docs/src/tutorials/large_systems.md
Original file line number Diff line number Diff line change
Expand Up @@ -60,12 +60,14 @@ broadcast). Use `dx=dy=1/32`.
The resulting `NonlinearProblem` definition is:

```@example ill_conditioned_nlprob
using NonlinearSolve, LinearAlgebra, SparseArrays, LinearSolve, Symbolics
using NonlinearSolve, LinearAlgebra, SparseArrays, LinearSolve, SparseDiffTools
const N = 32
const xyd_brusselator = range(0, stop = 1, length = N)
brusselator_f(x, y) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * 5.0
limit(a, N) = a == N + 1 ? 1 : a == 0 ? N : a
function brusselator_2d_loop(du, u, p)
A, B, alpha, dx = p
alpha = alpha / dx^2
Expand Down Expand Up @@ -96,8 +98,10 @@ function init_brusselator_2d(xyd)
end
u
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
Expand All @@ -120,6 +124,27 @@ include:
- 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
`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
@btime solve(prob_brusselator_2d, NewtonRaphson());
@btime solve(prob_brusselator_2d, NewtonRaphson(; autodiff = AutoSparseForwardDiff()));
@btime solve(prob_brusselator_2d,
NewtonRaphson(; autodiff = AutoSparseForwardDiff(), linsolve = KLUFactorization()));
@btime solve(prob_brusselator_2d,
NewtonRaphson(; autodiff = AutoSparseForwardDiff(), linsolve = KrylovJL_GMRES()));
nothing # hide
```

## Declaring a Sparse Jacobian with Automatic Sparsity Detection

Jacobian sparsity is declared by the `jac_prototype` argument in the `NonlinearFunction`.
Expand Down Expand Up @@ -156,7 +181,6 @@ prob_brusselator_2d_sparse = NonlinearProblem(ff, u0, p)
Now let's see how the version with sparsity compares to the version without:

```@example ill_conditioned_nlprob
using BenchmarkTools # for @btime
@btime solve(prob_brusselator_2d, NewtonRaphson());
@btime solve(prob_brusselator_2d_sparse, NewtonRaphson());
@btime solve(prob_brusselator_2d_sparse, NewtonRaphson(linsolve = KLUFactorization()));
Expand Down Expand Up @@ -202,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)
Expand Down Expand Up @@ -236,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)))
Expand All @@ -258,10 +284,8 @@ function algebraicmultigrid2(W, du, u, p, t, newW, Plprev, Prprev, solverdata)
if newW === nothing || newW
A = convert(AbstractMatrix, W)
Pl = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben(A,
presmoother = AlgebraicMultigrid.Jacobi(rand(size(A,
1))),
postsmoother = AlgebraicMultigrid.Jacobi(rand(size(A,
1)))))
presmoother = AlgebraicMultigrid.Jacobi(rand(size(A, 1))),
postsmoother = AlgebraicMultigrid.Jacobi(rand(size(A, 1)))))
else
Pl = Plprev
end
Expand All @@ -274,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/).
7 changes: 7 additions & 0 deletions ext/NonlinearSolveSymbolicsExt.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
module NonlinearSolveSymbolicsExt

import NonlinearSolve, Symbolics

NonlinearSolve.is_extension_loaded(::Val{:Symbolics}) = true

end
22 changes: 20 additions & 2 deletions src/jacobian.jl
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,30 @@ sparsity_detection_alg(_, _) = NoSparsityDetection()
function sparsity_detection_alg(f, ad::AbstractSparseADType)
if f.sparsity === nothing
if f.jac_prototype === nothing
return SymbolicsSparsityDetection()
if is_extension_loaded(Val(:Symbolics))
return SymbolicsSparsityDetection()
else
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)
Expand Down
2 changes: 1 addition & 1 deletion src/lbroyden.jl
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ An implementation of `LimitedMemoryBroyden` with resetting and line search.
- `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref),
which means that no line search is performed. Algorithms from `LineSearches.jl` can be
used here directly, and they will be converted to the correct `LineSearch`. It is
recommended to use [`LiFukushimaLineSearchCache`](@ref) -- a derivative free linesearch
recommended to use [`LiFukushimaLineSearch`](@ref) -- a derivative free linesearch
specifically designed for Broyden's method.
"""
@concrete struct LimitedMemoryBroyden{threshold} <: AbstractNewtonAlgorithm{false, Nothing}
Expand Down
1 change: 1 addition & 0 deletions src/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 89bd481

Please sign in to comment.