Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: use DifferentiationInterface for sparse AD #468

Merged
merged 17 commits into from
Oct 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,12 @@ 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"
SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
Expand All @@ -42,8 +45,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]
NonlinearSolveBandedMatricesExt = "BandedMatrices"
Expand All @@ -55,8 +56,6 @@ NonlinearSolveNLSolversExt = "NLSolvers"
NonlinearSolveNLsolveExt = "NLsolve"
NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
NonlinearSolveSpeedMappingExt = "SpeedMapping"
NonlinearSolveSymbolicsExt = "Symbolics"
NonlinearSolveZygoteExt = "Zygote"

[compat]
ADTypes = "1.9"
Expand Down Expand Up @@ -102,16 +101,18 @@ 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"
SparseDiffTools = "2.22"
SparseConnectivityTracer = "0.6.5"
SparseMatrixColorings = "0.4.2"
SpeedMapping = "0.3"
StableRNGs = "1"
StaticArrays = "1.9"
StaticArraysCore = "1.4"
Sundials = "4.23.1"
SymbolicIndexingInterface = "0.3.31"
Symbolics = "6.12"
Test = "1.10"
TimerOutputs = "0.5.23"
Zygote = "0.6.69"
Expand Down Expand Up @@ -144,9 +145,8 @@ SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

[targets]
test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote"]
12 changes: 8 additions & 4 deletions docs/Project.toml
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
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"
DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"
IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
Expand All @@ -16,19 +19,21 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e"
SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
avik-pal marked this conversation as resolved.
Show resolved Hide resolved
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"

[compat]
ADTypes = "1.9.0"
AlgebraicMultigrid = "0.5, 0.6"
ArrayInterface = "6, 7"
BenchmarkTools = "1"
DiffEqBase = "6.136"
DifferentiationInterface = "0.6"
Documenter = "1"
DocumenterCitations = "1"
DocumenterInterLinks = "1.0.0"
IncompleteLU = "0.2"
InteractiveUtils = "<0.0.1, 1"
LinearSolve = "2"
Expand All @@ -40,8 +45,7 @@ Random = "<0.0.1, 1"
SciMLBase = "2.4"
SciMLJacobianOperators = "0.1"
SimpleNonlinearSolve = "1"
SparseDiffTools = "2.14"
SparseConnectivityTracer = "0.6.5"
StaticArrays = "1"
SteadyStateDiffEq = "2"
Sundials = "4.11"
Symbolics = "4, 5, 6"
8 changes: 6 additions & 2 deletions docs/make.jl
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
using Documenter, DocumenterCitations
using Documenter, DocumenterCitations, DocumenterInterLinks
using NonlinearSolve, SimpleNonlinearSolve, Sundials, SteadyStateDiffEq, SciMLBase,
DiffEqBase
using SciMLJacobianOperators
Expand All @@ -12,6 +12,10 @@ include("pages.jl")

bib = CitationBibliography(joinpath(@__DIR__, "src", "refs.bib"))

interlinks = InterLinks(
"ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/",
)

makedocs(; sitename = "NonlinearSolve.jl",
authors = "Chris Rackauckas",
modules = [NonlinearSolve, SimpleNonlinearSolve, SteadyStateDiffEq,
Expand All @@ -23,7 +27,7 @@ makedocs(; sitename = "NonlinearSolve.jl",
"https://link.springer.com/article/10.1007/s40096-020-00339-4"],
checkdocs = :exports,
warnonly = [:missing_docs],
plugins = [bib],
plugins = [bib, interlinks],
format = Documenter.HTML(assets = ["assets/favicon.ico", "assets/citations.css"],
canonical = "https://docs.sciml.ai/NonlinearSolve/stable/"),
pages)
Expand Down
70 changes: 22 additions & 48 deletions docs/src/basics/autodiff.md
Original file line number Diff line number Diff line change
@@ -1,60 +1,34 @@
# Automatic Differentiation Backends

!!! note

We support all backends supported by DifferentiationInterface.jl. Please refer to
the [backends page](https://gdalle.github.io/DifferentiationInterface.jl/DifferentiationInterface/stable/explanation/backends/)
for more information.

## Summary of Finite Differencing Backends

- [`AutoFiniteDiff`](@ref): Finite differencing, not optimal but always applicable.
- [`AutoFiniteDiff`](@extref ADTypes.AutoFiniteDiff): Finite differencing using
`FiniteDiff.jl`, not optimal but always applicable.
- [`AutoFiniteDifferences`](@extref ADTypes.AutoFiniteDifferences): Finite differencing
using `FiniteDifferences.jl`, not optimal but always applicable.

## Summary of Forward Mode AD Backends

- [`AutoForwardDiff`](@ref): The best choice for dense problems.
- [`AutoPolyesterForwardDiff`](@ref): Might be faster than [`AutoForwardDiff`](@ref) for
large problems. Requires `PolyesterForwardDiff.jl` to be installed and loaded.
- [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff): The best choice for dense
problems.
- [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff): Might be faster
than [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff) for large problems. Requires
`PolyesterForwardDiff.jl` to be installed and loaded.

## Summary of Reverse Mode AD Backends

- [`AutoZygote`](@ref): The fastest choice for non-mutating array-based (BLAS) functions.
- [`AutoEnzyme`](@ref): Uses `Enzyme.jl` Reverse Mode and should be considered
experimental.

!!! note

If `PolyesterForwardDiff.jl` is installed and loaded, then `SimpleNonlinearSolve.jl`
will automatically use `AutoPolyesterForwardDiff` as the default AD backend.
- [`AutoZygote`](@extref ADTypes.AutoZygote): The fastest choice for non-mutating
array-based (BLAS) functions.
- [`AutoEnzyme`](@extref ADTypes.AutoEnzyme): Uses `Enzyme.jl` Reverse Mode and works for
both in-place and out-of-place functions.

!!! note
!!! tip

The sparse versions of the methods refer to automated sparsity detection. These
methods automatically discover the sparse Jacobian form from the function `f`. Note that
all methods specialize the differentiation on a sparse Jacobian if the sparse Jacobian
is given as `prob.f.jac_prototype` in the `NonlinearFunction` definition, and the
`AutoSparse` here simply refers to whether this `jac_prototype` should be generated
automatically. For more details, see
[SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl) and
[Sparsity Detection Manual Entry](@ref sparsity-detection), as well as the
documentation of [ADTypes.jl](https://github.com/SciML/ADTypes.jl).

## API Reference

```@docs
AutoSparse
```

### Finite Differencing Backends

```@docs
AutoFiniteDiff
```

### Forward Mode AD Backends

```@docs
AutoForwardDiff
AutoPolyesterForwardDiff
```

### Reverse Mode AD Backends

```@docs
AutoZygote
AutoEnzyme
```
For sparsity detection and sparse AD take a look at
[sparsity detection](@ref sparsity-detection).
64 changes: 41 additions & 23 deletions docs/src/basics/sparsity_detection.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,40 @@ to use this in an actual problem, see

Notation wise we are trying to solve for `x` such that `nlfunc(x) = 0`.

## Big Table for Determining Sparsity Detection and Coloring Algorithms

| `f.sparsity` | `f.jac_prototype` | `f.colorvec` | Sparsity Detection | Coloring Algorithm |
|:-------------------------- |:----------------- |:------------ |:------------------------------------------------ |:----------------------------------------- |
| ❌ | ❌ | `Any` | `NoSparsityDetector()` | `NoColoringAlgorithm()` |
| ❌ | Not Structured | `Any` | `NoSparsityDetector()` | `NoColoringAlgorithm()` |
| ❌ | Structured | ✅ | `KnownJacobianSparsityDetector(f.jac_prototype)` | `GreedyColoringAlgorithm(LargestFirst())` |
| ❌ | Structured | ❌ | `KnownJacobianSparsityDetector(f.jac_prototype)` | `GreedyColoringAlgorithm(LargestFirst())` |
| - | - | - | - | - |
| `AbstractMatrix` | ❌ | ✅ | `KnownJacobianSparsityDetector(f.sparsity)` | `ConstantColoringAlgorithm(f.colorvec)` |
| `AbstractMatrix` | ❌ | ❌ | `KnownJacobianSparsityDetector(f.sparsity)` | `GreedyColoringAlgorithm(LargestFirst())` |
| `AbstractMatrix` | Not Structured | ✅ | `KnownJacobianSparsityDetector(f.sparsity)` | `ConstantColoringAlgorithm(f.colorvec)` |
| `AbstractMatrix` | Not Structured | ❌ | `KnownJacobianSparsityDetector(f.sparsity)` | `GreedyColoringAlgorithm(LargestFirst())` |
| `AbstractMatrix` | Structured | `Any` | 🔴 | 🔴 |
| - | - | - | - | - |
| `AbstractSparsityDetector` | ❌ | `Any` | `f.sparsity` | `GreedyColoringAlgorithm(LargestFirst())` |
| `AbstractSparsityDetector` | Not Structured | ✅ | `f.sparsity` | `ConstantColoringAlgorithm(f.colorvec)` |
| `AbstractSparsityDetector` | Not Structured | ❌ | `f.sparsity` | `GreedyColoringAlgorithm(LargestFirst())` |
| `AbstractSparsityDetector` | Structured | ✅ | `KnownJacobianSparsityDetector(f.jac_prototype)` | `ConstantColoringAlgorithm(f.colorvec)` |
| `AbstractSparsityDetector` | Structured | ❌ | `KnownJacobianSparsityDetector(f.jac_prototype)` | `GreedyColoringAlgorithm(LargestFirst())` |

1. `Structured` means either a `AbstractSparseMatrix` or `ArrayInterface.isstructured(x)` is true.
2. ❌ means not provided (default)
3. ✅ means provided
4. 🔴 means an error will be thrown
5. Providing a colorvec without specifying either sparsity or jac_prototype with a sparse or structured matrix will cause us to ignore the colorvec.
6. The function calls demonstrated above are simply pseudo-code to show the general idea.

## 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)
```

Expand All @@ -22,9 +48,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)
```
Expand All @@ -34,8 +57,8 @@ 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
avik-pal marked this conversation as resolved.
Show resolved Hide resolved
colorvec is the column colorvec, otherwise we will assume that the colorvec is the
row colorvec.

## Case II: Sparsity Detection algorithm is provided
Expand All @@ -45,17 +68,21 @@ 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
NonlinearFunction(nlfunc; sparsity = SymbolicsSparsityDetector()), 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 and SparseConnectivityTracer.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.

avik-pal marked this conversation as resolved.
Show resolved Hide resolved
If you constructed a Nonlinear Solver with a sparse AD type, for example

```julia
Expand All @@ -65,15 +92,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.
2 changes: 1 addition & 1 deletion docs/src/tutorials/code_optimization.md
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ end

Allocations are only expensive if they are “heap allocations”. For a more in-depth
definition of heap allocations,
[there are many sources online](http://net-informations.com/faq/net/stack-heap.htm).
[there are many sources online](https://net-informations.com/faq/net/stack-heap.htm).
But a good working definition is that heap allocations are variable-sized slabs of memory
which have to be pointed to, and this pointer indirection costs time. Additionally, the heap
has to be managed, and the garbage controllers has to actively keep track of what's on the
Expand Down
Loading
Loading