diff --git a/Project.toml b/Project.toml index 5bd304bb3..69dd08304 100644 --- a/Project.toml +++ b/Project.toml @@ -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" @@ -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" @@ -55,8 +56,6 @@ NonlinearSolveNLSolversExt = "NLSolvers" NonlinearSolveNLsolveExt = "NLsolve" NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations" NonlinearSolveSpeedMappingExt = "SpeedMapping" -NonlinearSolveSymbolicsExt = "Symbolics" -NonlinearSolveZygoteExt = "Zygote" [compat] ADTypes = "1.9" @@ -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" @@ -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"] diff --git a/docs/Project.toml b/docs/Project.toml index 9fbee9e39..ab35d42ae 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -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" @@ -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" 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" @@ -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" diff --git a/docs/make.jl b/docs/make.jl index e11a04149..3e931a227 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,4 +1,4 @@ -using Documenter, DocumenterCitations +using Documenter, DocumenterCitations, DocumenterInterLinks using NonlinearSolve, SimpleNonlinearSolve, Sundials, SteadyStateDiffEq, SciMLBase, DiffEqBase using SciMLJacobianOperators @@ -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, @@ -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) diff --git a/docs/src/basics/autodiff.md b/docs/src/basics/autodiff.md index baa605363..d2d01e00b 100644 --- a/docs/src/basics/autodiff.md +++ b/docs/src/basics/autodiff.md @@ -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). diff --git a/docs/src/basics/sparsity_detection.md b/docs/src/basics/sparsity_detection.md index 97891be61..de924c398 100644 --- a/docs/src/basics/sparsity_detection.md +++ b/docs/src/basics/sparsity_detection.md @@ -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) ``` @@ -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) ``` @@ -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 + colorvec is the column colorvec, otherwise we will assume that the colorvec is the row colorvec. ## Case II: Sparsity Detection algorithm is provided @@ -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. + If you constructed a Nonlinear Solver with a sparse AD type, for example ```julia @@ -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. diff --git a/docs/src/tutorials/code_optimization.md b/docs/src/tutorials/code_optimization.md index ce114961b..b7eb9c174 100644 --- a/docs/src/tutorials/code_optimization.md +++ b/docs/src/tutorials/code_optimization.md @@ -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 diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index 368535287..17b768288 100644 --- a/docs/src/tutorials/large_systems.md +++ b/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 @@ -155,23 +161,31 @@ is non-zeros, otherwise the overhead of sparse matrices can be higher than the g sparse differentiation! One of the useful companion tools for NonlinearSolve.jl is -[Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl). This allows for automatic -declaration of Jacobian sparsity types. To see this in action, we can give an example `du` -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`. +[ADTypes.jl](https://github.com/SciML/ADTypes.jl) that specifies the interface for sparsity +detection via [`jacobian_sparsity`](@extref ADTypes.jacobian_sparsity). This allows for +automatic declaration of Jacobian sparsity types. To see this in action, we can give an +example `du` 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 + + External packages like `SparseConnectivityTracer.jl` and `Symbolics.jl` provide the + actual implementation of sparsity detection. ```@example ill_conditioned_nlprob -using Symbolics -du0 = copy(u0) -jac_sparsity = Symbolics.jacobian_sparsity( - (du, u) -> brusselator_2d_loop(du, u, p), du0, u0) +using SparseConnectivityTracer, ADTypes + +f! = (du, u) -> brusselator_2d_loop(du, u, p) +du0 = similar(u0) +jac_sparsity = ADTypes.jacobian_sparsity(f!, du0, u0, TracerSparsityDetector()) ``` Notice that Julia gives a nice print out of the sparsity pattern. That's neat, and would be 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 +289,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 +310,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 +322,18 @@ 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_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_tracer, NewtonRaphson()); +@btime solve(prob_brusselator_2d_approx_di, NewtonRaphson()); nothing # hide ``` diff --git a/ext/NonlinearSolveSymbolicsExt.jl b/ext/NonlinearSolveSymbolicsExt.jl deleted file mode 100644 index 8e5354cda..000000000 --- a/ext/NonlinearSolveSymbolicsExt.jl +++ /dev/null @@ -1,7 +0,0 @@ -module NonlinearSolveSymbolicsExt - -using NonlinearSolve: NonlinearSolve - -NonlinearSolve.is_extension_loaded(::Val{:Symbolics}) = true - -end diff --git a/ext/NonlinearSolveZygoteExt.jl b/ext/NonlinearSolveZygoteExt.jl deleted file mode 100644 index 8ed2e1853..000000000 --- a/ext/NonlinearSolveZygoteExt.jl +++ /dev/null @@ -1,7 +0,0 @@ -module NonlinearSolveZygoteExt - -using NonlinearSolve: NonlinearSolve - -NonlinearSolve.is_extension_loaded(::Val{:Zygote}) = true - -end diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 9f8048cfb..50051fcb0 100644 --- a/src/NonlinearSolve.jl +++ b/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, @@ -39,33 +30,41 @@ using MaybeInplace: @bb 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 +using SciMLBase: AbstractNonlinearAlgorithm, AbstractNonlinearProblem, _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 SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm, + LargestFirst + @reexport using SciMLBase, SimpleNonlinearSolve const DI = DifferentiationInterface -# Type-Inference Friendly Check for Extension Loading -is_extension_loaded(::Val) = false - const True = Val(true) const False = Val(false) include("abstract_types.jl") -include("adtypes.jl") include("timer_outputs.jl") include("internal/helpers.jl") diff --git a/src/adtypes.jl b/src/adtypes.jl deleted file mode 100644 index 0ee20effb..000000000 --- a/src/adtypes.jl +++ /dev/null @@ -1,140 +0,0 @@ -# This just documents the AD types from ADTypes.jl - -""" - AutoFiniteDiff(; fdtype = Val(:forward), fdjtype = fdtype, fdhtype = Val(:hcentral)) - -This uses [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl). While not necessarily -the most efficient, this is the only choice that doesn't require the `f` function to be -automatically differentiable, which means it applies to any choice. However, because it's -using finite differencing, one needs to be careful as this procedure introduces numerical -error into the derivative estimates. - - - Compatible with GPUs - - Can be used for Jacobian-Vector Products (JVPs) - - Can be used for Vector-Jacobian Products (VJPs) - - Supports both inplace and out-of-place functions - -### Keyword Arguments - - - `fdtype`: the method used for defining the gradient - - `fdjtype`: the method used for defining the Jacobian of constraints. - - `fdhtype`: the method used for defining the Hessian -""" -AutoFiniteDiff - -""" - AutoForwardDiff(; chunksize = nothing, tag = nothing) - AutoForwardDiff{chunksize, tagType}(tag::tagType) - -This uses the [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) package. It is -the fastest choice for square or wide systems. It is easy to use and compatible with most -Julia functions which have loose type restrictions. - - - Compatible with GPUs - - Can be used for Jacobian-Vector Products (JVPs) - - Supports both inplace and out-of-place functions - -For type-stability of internal operations, a positive `chunksize` must be provided. - -### Keyword Arguments - - - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this - number to a high value will lead to slowdowns. Use - [`NonlinearSolve.pickchunksize`](@ref) to get a proper value. - - `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag. -""" -AutoForwardDiff - -""" - AutoPolyesterForwardDiff(; chunksize = nothing) - -Uses [`PolyesterForwardDiff.jl`](https://github.com/JuliaDiff/PolyesterForwardDiff.jl) -to compute the jacobian. This is essentially parallelized `ForwardDiff.jl`. - - - Supports both inplace and out-of-place functions - -### Keyword Arguments - - - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting - this number to a high value will lead to slowdowns. Use - [`NonlinearSolve.pickchunksize`](@ref) to get a proper value. -""" -AutoPolyesterForwardDiff - -""" - AutoZygote() - -Uses [`Zygote.jl`](https://github.com/FluxML/Zygote.jl) package. This is the staple -reverse-mode AD that handles a large portion of Julia with good efficiency. - - - Compatible with GPUs - - Can be used for Vector-Jacobian Products (VJPs) - - Supports only out-of-place functions - -For VJPs this is the current best choice. This is the most efficient method for long -jacobians. -""" -AutoZygote - -""" - AutoEnzyme() - -Uses reverse mode [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). This is currently -experimental, and not extensively tested on our end. We only support Jacobian construction -and VJP support is currently not implemented. - - - Supports both inplace and out-of-place functions -""" -AutoEnzyme - -""" - AutoSparse(AutoEnzyme()) - -Sparse version of [`AutoEnzyme`](@ref) that uses -[Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) and the row color vector of -the Jacobian Matrix to efficiently compute the Sparse Jacobian. - - - Supports both inplace and out-of-place functions - -This is efficient only for long jacobians or if the maximum value of the row color vector is -significantly lower than the maximum value of the column color vector. - - AutoSparse(AutoFiniteDiff()) - -Sparse Version of [`AutoFiniteDiff`](@ref) that uses -[FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) and the column color vector of -the Jacobian Matrix to efficiently compute the Sparse Jacobian. - - - Supports both inplace and out-of-place functions - - AutoSparse(AutoForwardDiff(; chunksize = nothing, tag = nothing)) - AutoSparse(AutoForwardDiff{chunksize, tagType}(tag::tagType)) - -Sparse Version of [`AutoForwardDiff`](@ref) that uses -[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) and the column color vector of -the Jacobian Matrix to efficiently compute the Sparse Jacobian. - - - Supports both inplace and out-of-place functions - -For type-stability of internal operations, a positive `chunksize` must be provided. - -### Keyword Arguments - - - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this - number to a high value will lead to slowdowns. Use - [`NonlinearSolve.pickchunksize`](@ref) to get a proper value. - - - `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag. - - AutoSparse(AutoZygote()) - -Sparse version of [`AutoZygote`](@ref) that uses -[`Zygote.jl`](https://github.com/FluxML/Zygote.jl) and the row color vector of -the Jacobian Matrix to efficiently compute the Sparse Jacobian. - - - Supports only out-of-place functions - -This is efficient only for long jacobians or if the maximum value of the row color vector is -significantly lower than the maximum value of the column color vector. -""" -AutoSparse diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index a3555a17b..46ebc8dae 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -317,7 +317,7 @@ NLSolversJL(; method, autodiff = nothing) = NLSolversJL(method, autodiff) SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) -Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl) for solving +Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl/) for solving Fixed Point Problems. We allow using this algorithm to solve root finding problems as well. ### Keyword Arguments diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl index abf6f1099..afc24ce5b 100644 --- a/src/internal/helpers.jl +++ b/src/internal/helpers.jl @@ -80,7 +80,7 @@ function get_concrete_reverse_ad( else use_sparse_ad = false end - ad = if isinplace(prob) || !is_extension_loaded(Val(:Zygote)) # Use Finite Differencing + ad = if isinplace(prob) || !DI.check_available(AutoZygote()) # Use Finite Differencing use_sparse_ad ? AutoSparse(AutoFiniteDiff()) : AutoFiniteDiff() else use_sparse_ad ? AutoSparse(AutoZygote()) : AutoZygote() diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 67ab839a2..93d10f98b 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -37,7 +37,6 @@ Construct a cache for the Jacobian of `f` w.r.t. `u`. stats::NLStats autodiff di_extras - sdifft_extras end function reinit_cache!(cache::JacobianCache{iip}, args...; p = cache.p, @@ -58,33 +57,18 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, @bb fu = similar(fu_) - autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false) + autodiff = get_concrete_forward_ad( + autodiff, prob, Val(false); check_forward_mode = false) if !has_analytic_jac && needs_jac - sd = sparsity_detection_alg(f, autodiff) - sparse_jac = !(sd isa NoSparsityDetection) - # Eventually we want to do everything via DI. But for now, we just do the dense via DI - if sparse_jac - di_extras = nothing - uf = JacobianWrapper{iip}(f, p) - sdifft_extras = if iip - sparse_jacobian_cache(autodiff, sd, uf, fu, u) - else - sparse_jacobian_cache( - autodiff, sd, uf, __maybe_mutable(u, autodiff); fx = fu) - end + autodiff = construct_concrete_adtype(f, autodiff) + di_extras = if iip + DI.prepare_jacobian(f, fu, autodiff, u, Constant(prob.p)) else - sdifft_extras = nothing - di_extras = if iip - DI.prepare_jacobian(f, fu, autodiff, u, Constant(p)) - else - DI.prepare_jacobian(f, autodiff, u, Constant(p)) - end + DI.prepare_jacobian(f, autodiff, u, Constant(prob.p)) end else - sparse_jac = false di_extras = nothing - sdifft_extras = nothing end J = if !needs_jac @@ -95,42 +79,43 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, JacobianOperator(prob, fu, u; jvp_autodiff, vjp_autodiff) else if f.jac_prototype === nothing - if !sparse_jac - # While this is technically wasteful, it gives out the type of the Jacobian - # which is needed to create the linear solver cache - stats.njacs += 1 - if has_analytic_jac - __similar( - fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) + # While this is technically wasteful, it gives out the type of the Jacobian + # which is needed to create the linear solver cache + stats.njacs += 1 + if has_analytic_jac + __similar( + fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) + else + if iip + DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p)) else - if iip - DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p)) - else - DI.jacobian(f, autodiff, u, Constant(p)) - end + DI.jacobian(f, autodiff, u, Constant(p)) end - else - zero(init_jacobian(sdifft_extras; preserve_immutable = Val(true))) end else - similar(f.jac_prototype) + jac_proto = if eltype(f.jac_prototype) <: Bool + similar(f.jac_prototype, promote_type(eltype(fu), eltype(u))) + else + similar(f.jac_prototype) + end + fill!(jac_proto, false) + jac_proto end end - return JacobianCache{iip}( - J, f, fu, u, p, stats, autodiff, di_extras, sdifft_extras) + return JacobianCache{iip}(J, f, fu, u, p, stats, autodiff, di_extras) end function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats, autodiff = nothing, kwargs...) where {F} fu = f(u, p) if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f) - return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing, nothing) + return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing) end autodiff = get_dense_ad(get_concrete_forward_ad( autodiff, prob; check_forward_mode = false)) di_extras = DI.prepare_derivative(f, get_dense_ad(autodiff), u, Constant(prob.p)) - return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras, nothing) + return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras) end (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p) @@ -163,67 +148,121 @@ function (cache::JacobianCache{iip})( if iip if SciMLBase.has_jac(cache.f) cache.f.jac(J, u, p) - elseif cache.di_extras !== nothing + else DI.jacobian!( cache.f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p)) - else - uf = JacobianWrapper{iip}(cache.f, p) - sparse_jacobian!(J, cache.autodiff, cache.sdifft_extras, uf, cache.fu, u) end return J else if SciMLBase.has_jac(cache.f) return cache.f.jac(u, p) - elseif cache.di_extras !== nothing - return DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p)) else - uf = JacobianWrapper{iip}(cache.f, p) - if __can_setindex(typeof(J)) - sparse_jacobian!(J, cache.autodiff, cache.sdifft_extras, uf, u) - return J - else - return sparse_jacobian(cache.autodiff, cache.sdifft_extras, uf, u) - end + return DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p)) end end end -function sparsity_detection_alg(f::NonlinearFunction, ad::AbstractADType) - # TODO: Also handle case where colorvec is provided - f.sparsity === nothing && return NoSparsityDetection() - return sparsity_detection_alg(f, AutoSparse(ad; sparsity_detector = f.sparsity)) -end - -function sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse) +function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) + @assert !(ad isa AutoSparse) "This shouldn't happen. Open an issue." if f.sparsity === nothing if f.jac_prototype === nothing - is_extension_loaded(Val(:Symbolics)) && return SymbolicsSparsityDetection() - return ApproximateJacobianSparsity() + if SciMLBase.has_colorvec(f) + @warn "`colorvec` is provided but `sparsity` and `jac_prototype` is not \ + specified. `colorvec` will be ignored." + end + return ad # No sparse AD else - jac_prototype = f.jac_prototype + if !sparse_or_structured_prototype(f.jac_prototype) + if SciMLBase.has_colorvec(f) + @warn "`colorvec` is provided but `jac_prototype` is not a sparse \ + or structured matrix. `colorvec` will be ignored." + end + return ad + end + return AutoSparse( + ad; + sparsity_detector = KnownJacobianSparsityDetector(f.jac_prototype), + coloring_algorithm = select_fastest_coloring_algorithm( + f.jac_prototype, f, ad) + ) end - elseif f.sparsity isa AbstractSparsityDetection - f.jac_prototype === nothing && return f.sparsity - jac_prototype = f.jac_prototype - 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`.") + if f.sparsity isa AbstractMatrix + if f.jac_prototype !== f.sparsity + if f.jac_prototype !== nothing && + sparse_or_structured_prototype(f.jac_prototype) + throw(ArgumentError("`sparsity::AbstractMatrix` and a sparse or \ + structured `jac_prototype` cannot be both \ + provided. Pass only `jac_prototype`.")) + end + end + return AutoSparse( + ad; + sparsity_detector = KnownJacobianSparsityDetector(f.sparsity), + coloring_algorithm = select_fastest_coloring_algorithm( + f.sparsity, f, ad) + ) + end + + @assert f.sparsity isa ADTypes.AbstractSparsityDetector + sparsity_detector = f.sparsity + if f.jac_prototype === nothing + if SciMLBase.has_colorvec(f) + @warn "`colorvec` is provided but `jac_prototype` is not specified. \ + `colorvec` will be ignored." + end + return AutoSparse( + ad; + sparsity_detector, + coloring_algorithm = GreedyColoringAlgorithm(LargestFirst()) + ) + else + if sparse_or_structured_prototype(f.jac_prototype) + if !(sparsity_detector isa NoSparsityDetector) + @warn "`jac_prototype` is a sparse matrix but sparsity = $(f.sparsity) \ + has also been specified. Ignoring sparsity field and using \ + `jac_prototype` sparsity." + end + sparsity_detector = KnownJacobianSparsityDetector(f.jac_prototype) + end + + return AutoSparse( + ad; + sparsity_detector, + coloring_algorithm = select_fastest_coloring_algorithm( + f.jac_prototype, f, ad) + ) + end end +end +function select_fastest_coloring_algorithm( + prototype, f::NonlinearFunction, ad::AbstractADType) if SciMLBase.has_colorvec(f) - return PrecomputedJacobianColorvec(; jac_prototype, f.colorvec, - partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode) - else - return JacPrototypeSparsityDetection(; jac_prototype) + return ConstantColoringAlgorithm{ifelse( + ADTypes.mode(ad) isa ADTypes.ReverseMode, :row, :column)}( + prototype, f.colorvec) end + return GreedyColoringAlgorithm(LargestFirst()) +end + +function construct_concrete_adtype(f::NonlinearFunction, ad::AutoSparse) + Base.depwarn( + "Specifying a sparse AD type for Nonlinear Problems is deprecated. \ + Instead use the `sparsity`, `jac_prototype`, and `colorvec` to specify \ + the right sparsity pattern and coloring algorithm. Ignoring the sparsity \ + detection algorithm and coloring algorithm present in $(ad).", + :NonlinearSolve) + if f.sparsity === nothing && f.jac_prototype === nothing + @set! f.sparsity = TracerSparsityDetector() + end + return construct_concrete_adtype(f, get_dense_ad(ad)) end get_dense_ad(ad) = ad get_dense_ad(ad::AutoSparse) = ADTypes.dense_ad(ad) + +sparse_or_structured_prototype(::AbstractSparseMatrix) = true +function sparse_or_structured_prototype(prototype::AbstractMatrix) + return ArrayInterface.isstructured(prototype) +end diff --git a/src/utils.jl b/src/utils.jl index e13f9ca60..6ceb4c9d8 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -157,5 +157,6 @@ end function __similar(x, args...; kwargs...) y = similar(x, args...; kwargs...) - return zero(y) + fill!(y, false) + return y end diff --git a/test/core/rootfind_tests.jl b/test/core/rootfind_tests.jl index bca02a44b..6a56c6ac6 100644 --- a/test/core/rootfind_tests.jl +++ b/test/core/rootfind_tests.jl @@ -1,7 +1,7 @@ @testsetup module CoreRootfindTesting using Reexport @reexport using BenchmarkTools, LinearSolve, NonlinearSolve, StaticArrays, Random, - LinearAlgebra, ForwardDiff, Zygote, Enzyme, SparseDiffTools, DiffEqBase + LinearAlgebra, ForwardDiff, Zygote, Enzyme, DiffEqBase _nameof(x) = applicable(nameof, x) ? nameof(x) : _nameof(typeof(x)) diff --git a/test/misc/bruss_tests.jl b/test/misc/bruss_tests.jl index 8ea1fdb18..80b1bd540 100644 --- a/test/misc/bruss_tests.jl +++ b/test/misc/bruss_tests.jl @@ -1,5 +1,5 @@ @testitem "Brusselator 2D" tags=[:misc] begin - using LinearAlgebra, SparseArrays, Symbolics + using LinearAlgebra, SparseArrays, SparseConnectivityTracer, ADTypes const N = 32 const xyd_brusselator = range(0, stop = 1, length = N) @@ -46,20 +46,20 @@ sol = solve(prob_brusselator_2d, NewtonRaphson(); abstol = 1e-8) @test norm(sol.resid, Inf) < 1e-8 - sol = solve(prob_brusselator_2d, - NewtonRaphson(autodiff = AutoSparse(AutoForwardDiff())); abstol = 1e-8) + prob_brusselator_2d_sparse = NonlinearProblem( + NonlinearFunction(brusselator_2d_loop; sparsity = TracerSparsityDetector()), + u0, p) + sol = solve(prob_brusselator_2d_sparse, NewtonRaphson(); abstol = 1e-8) @test norm(sol.resid, Inf) < 1e-8 + # Deprecated sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = AutoSparse(AutoFiniteDiff())); abstol = 1e-8) @test norm(sol.resid, Inf) < 1e-8 - du0 = copy(u0) - jac_sparsity = Symbolics.jacobian_sparsity( - (du, u) -> brusselator_2d_loop(du, u, p), du0, u0) - jac_prototype = float.(jac_sparsity) - fill!(jac_prototype, 0) - @test all(iszero, jac_prototype) + f! = (du, u) -> brusselator_2d_loop(du, u, p) + du0 = similar(u0) + jac_prototype = ADTypes.jacobian_sparsity(f!, du0, u0, TracerSparsityDetector()) ff_iip = NonlinearFunction(brusselator_2d_loop; jac_prototype) prob_brusselator_2d = NonlinearProblem(ff_iip, u0, p) @@ -68,11 +68,6 @@ @test norm(sol.resid, Inf) < 1e-8 sol = solve(prob_brusselator_2d, - NewtonRaphson(autodiff = AutoSparse(AutoFiniteDiff())); abstol = 1e-8) + NewtonRaphson(autodiff = AutoFiniteDiff()); abstol = 1e-8) @test norm(sol.resid, Inf) < 1e-8 - - cache = init( - prob_brusselator_2d, NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff()))) - @test maximum(cache.jac_cache.sdifft_extras.coloring.colorvec) == 12 - @test cache.jac_cache.autodiff isa AutoSparse{<:AutoForwardDiff} end diff --git a/test/misc/qa_tests.jl b/test/misc/qa_tests.jl index 5c60e5339..6aafd4024 100644 --- a/test/misc/qa_tests.jl +++ b/test/misc/qa_tests.jl @@ -17,8 +17,7 @@ end @testitem "Explicit Imports" tags=[:misc] begin using NonlinearSolve, ADTypes, SimpleNonlinearSolve, SciMLBase import BandedMatrices, FastLevenbergMarquardt, FixedPointAcceleration, - LeastSquaresOptim, MINPACK, NLsolve, NLSolvers, SIAMFANLEquations, SpeedMapping, - Symbolics, Zygote + LeastSquaresOptim, MINPACK, NLsolve, NLSolvers, SIAMFANLEquations, SpeedMapping using ExplicitImports diff --git a/test/misc/structured_jacobian_tests.jl b/test/misc/structured_jacobian_tests.jl new file mode 100644 index 000000000..41b316911 --- /dev/null +++ b/test/misc/structured_jacobian_tests.jl @@ -0,0 +1,67 @@ +@testitem "Structured Jacobians" tags=[:misc] begin + using NonlinearSolve, SparseConnectivityTracer, BandedMatrices, LinearAlgebra, + SparseArrays + + N = 16 + p = rand(N) + u0 = rand(N) + + function f!(du, u, p) + for i in 2:(length(u) - 1) + du[i] = u[i - 1] - 2u[i] + u[i + 1] + p[i] + end + du[1] = -2u[1] + u[2] + p[1] + du[end] = u[end - 1] - 2u[end] + p[end] + return nothing + end + + function f(u, p) + du = similar(u, promote_type(eltype(u), eltype(p))) + f!(du, u, p) + return du + end + + for nlf in (f, f!) + @testset "Dense AD" begin + nlprob = NonlinearProblem(NonlinearFunction(nlf), u0, p) + + cache = init(nlprob, NewtonRaphson(); abstol = 1e-9) + @test cache.jac_cache.J isa Matrix + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + end + + @testset "Unstructured Sparse AD" begin + nlprob_autosparse = NonlinearProblem( + NonlinearFunction(nlf; sparsity = TracerSparsityDetector()), + u0, p) + + cache = init(nlprob_autosparse, NewtonRaphson(); abstol = 1e-9) + @test cache.jac_cache.J isa SparseMatrixCSC + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + end + + @testset "Structured Sparse AD: Banded Jacobian" begin + jac_prototype = BandedMatrix(-1 => ones(N - 1), 0 => ones(N), 1 => ones(N - 1)) + nlprob_sparse_structured = NonlinearProblem( + NonlinearFunction(nlf; jac_prototype), u0, p) + + cache = init(nlprob_sparse_structured, NewtonRaphson(); abstol = 1e-9) + @test cache.jac_cache.J isa BandedMatrix + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + end + + @testset "Structured Sparse AD: Tridiagonal Jacobian" begin + jac_prototype = Tridiagonal(ones(N - 1), ones(N), ones(N - 1)) + nlprob_sparse_structured = NonlinearProblem( + NonlinearFunction(nlf; jac_prototype), u0, p) + + cache = init(nlprob_sparse_structured, NewtonRaphson(); abstol = 1e-9) + @test cache.jac_cache.J isa Tridiagonal + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + end + end +end