From 5f763136e6f66c4f46edecfa246f4035f01bc9dc Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 1 Oct 2024 15:27:29 -0400 Subject: [PATCH 01/17] refactor: reorder imports --- Project.toml | 4 ++++ src/NonlinearSolve.jl | 22 ++++++++++++---------- src/utils.jl | 3 ++- 3 files changed, 18 insertions(+), 11 deletions(-) diff --git a/Project.toml b/Project.toml index 5bd304bb3..b8a21339f 100644 --- a/Project.toml +++ b/Project.toml @@ -27,7 +27,9 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" +SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" @@ -104,7 +106,9 @@ SciMLBase = "2.54.0" SciMLJacobianOperators = "0.1" SimpleNonlinearSolve = "1.12.3" SparseArrays = "1.10" +SparseConnectivityTracer = "0.6.5" SparseDiffTools = "2.22" +SparseMatrixColorings = "0.4.2" SpeedMapping = "0.3" StableRNGs = "1" StaticArrays = "1.9" diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 9f8048cfb..4e0c0f7cf 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, @@ -43,7 +34,6 @@ using SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearP AbstractSciMLOperator, _unwrap_val, isinplace, NLStats using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJacOperator, JacVecOperator, StatefulJacobianOperator -using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC using SparseDiffTools: SparseDiffTools, AbstractSparsityDetection, ApproximateJacobianSparsity, JacPrototypeSparsityDetection, NoSparsityDetection, PrecomputedJacobianColorvec, @@ -54,6 +44,18 @@ using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingPro symbolic_container, parameter_values, state_values, getu, setu +# AD Support +using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff, + AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, AutoSparse +using ADTypes: AutoSparseFiniteDiff, AutoSparseForwardDiff, AutoSparsePolyesterForwardDiff, + AutoSparseZygote # FIXME: deprecated, remove in future +using DifferentiationInterface: DifferentiationInterface, Constant +using FiniteDiff: FiniteDiff +using ForwardDiff: ForwardDiff, Dual + +## Sparse AD Support +using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC + @reexport using SciMLBase, SimpleNonlinearSolve const DI = DifferentiationInterface 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 From 78e19a978de78258526a5b9eb03f8ec424b46e08 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 1 Oct 2024 15:35:22 -0400 Subject: [PATCH 02/17] refactor: remove Symbolics --- Project.toml | 2 -- ext/NonlinearSolveSymbolicsExt.jl | 7 ------- src/NonlinearSolve.jl | 5 +++-- src/internal/jacobian.jl | 2 +- 4 files changed, 4 insertions(+), 12 deletions(-) delete mode 100644 ext/NonlinearSolveSymbolicsExt.jl diff --git a/Project.toml b/Project.toml index b8a21339f..596eb5656 100644 --- a/Project.toml +++ b/Project.toml @@ -44,7 +44,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] @@ -57,7 +56,6 @@ NonlinearSolveNLSolversExt = "NLSolvers" NonlinearSolveNLsolveExt = "NLsolve" NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations" NonlinearSolveSpeedMappingExt = "SpeedMapping" -NonlinearSolveSymbolicsExt = "Symbolics" NonlinearSolveZygoteExt = "Zygote" [compat] 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/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 4e0c0f7cf..bd424d37f 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -37,8 +37,8 @@ using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJac using SparseDiffTools: SparseDiffTools, AbstractSparsityDetection, ApproximateJacobianSparsity, JacPrototypeSparsityDetection, NoSparsityDetection, PrecomputedJacobianColorvec, - SymbolicsSparsityDetection, init_jacobian, sparse_jacobian, - sparse_jacobian!, sparse_jacobian_cache + init_jacobian, sparse_jacobian, sparse_jacobian!, + sparse_jacobian_cache using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy, symbolic_container, parameter_values, state_values, getu, @@ -55,6 +55,7 @@ using ForwardDiff: ForwardDiff, Dual ## Sparse AD Support using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC +using SparseConnectivityTracer: TracerSparsityDetector @reexport using SciMLBase, SimpleNonlinearSolve diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 67ab839a2..bf178eff1 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -197,7 +197,7 @@ end function sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse) if f.sparsity === nothing if f.jac_prototype === nothing - is_extension_loaded(Val(:Symbolics)) && return SymbolicsSparsityDetection() + # is_extension_loaded(Val(:Symbolics)) && return SymbolicsSparsityDetection() return ApproximateJacobianSparsity() else jac_prototype = f.jac_prototype From fc7fcdbd68222ab010eb35fc1fd5750e5dd7c047 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 1 Oct 2024 17:12:31 -0400 Subject: [PATCH 03/17] feat: use DI for sparse jacobians --- Project.toml | 4 + docs/Project.toml | 4 + docs/src/basics/sparsity_detection.md | 40 +++--- docs/src/tutorials/large_systems.md | 83 +++++++----- src/NonlinearSolve.jl | 22 ++-- src/internal/jacobian.jl | 183 ++++++++++++++++++++------ test/misc/bruss_tests.jl | 27 ++-- 7 files changed, 247 insertions(+), 116 deletions(-) diff --git a/Project.toml b/Project.toml index 596eb5656..7d8f1f907 100644 --- a/Project.toml +++ b/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" diff --git a/docs/Project.toml b/docs/Project.toml index 9fbee9e39..48d18ffa9 100644 --- a/docs/Project.toml +++ b/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" diff --git a/docs/src/basics/sparsity_detection.md b/docs/src/basics/sparsity_detection.md index 97891be61..5124a0009 100644 --- a/docs/src/basics/sparsity_detection.md +++ b/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. diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index 368535287..8ed6ca169 100644 --- a/docs/src/tutorials/large_systems.md +++ b/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 ``` diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index bd424d37f..2db06f9b0 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -31,14 +31,9 @@ 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, @@ -46,16 +41,23 @@ using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingPro # 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 diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index bf178eff1..0bf7e4380 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -58,31 +58,34 @@ 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 + autodiff = construct_concrete_adtype(f, autodiff) + using_sparsedifftools = autodiff isa StructuredMatrixAutodiff + # DI can't handle structured matrices + if using_sparsedifftools 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) + autodiff.autodiff, autodiff.sparsity_detection, uf, fu, u) + else + sparse_jacobian_cache(autodiff.autodiff, autodiff.sparsity_detection, + uf, __maybe_mutable(u, autodiff); fx = fu) end + autodiff = autodiff.autodiff # For saving we unwrap else sdifft_extras = nothing di_extras = if iip - DI.prepare_jacobian(f, fu, autodiff, u, Constant(p)) + DI.prepare_jacobian(f, fu, autodiff, u, Constant(prob.p)) else - DI.prepare_jacobian(f, autodiff, u, Constant(p)) + DI.prepare_jacobian(f, autodiff, u, Constant(prob.p)) end end else - sparse_jac = false + using_sparsedifftools = false di_extras = nothing sdifft_extras = nothing end @@ -95,7 +98,7 @@ 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 + if !using_sparsedifftools # 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 @@ -113,7 +116,13 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, 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 @@ -188,41 +197,137 @@ function (cache::JacobianCache{iip})( 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 + if ArrayInterface.isstructured(f.jac_prototype) + return select_fastest_structured_matrix_autodiff(f.jac_prototype, f, ad) + end + + return AutoSparse( + ad; + sparsity_detector = KnownJacobianSparsityDetector(f.jac_prototype), + coloring_algorithm = select_fastest_coloring_algorithm( + f.jac_prototype, f, ad) + ) + end + else + if f.sparsity isa AbstractMatrix + if f.jac_prototype !== nothing && f.jac_prototype !== f.sparsity + throw(ArgumentError("`sparsity::AbstractMatrix` and `jac_prototype` cannot \ + be both provided. Pass only `jac_prototype`.")) + end + Base.depwarn("`sparsity::typeof($(typeof(f.sparsity)))` is deprecated. \ + Pass it as `jac_prototype` instead.", + :NonlinearSolve) + if ArrayInterface.isstructured(f.sparsity) + return select_fastest_structured_matrix_autodiff(f.sparsity, f, ad) + 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() + ) else - jac_prototype = f.jac_prototype + if ArrayInterface.isstructured(f.jac_prototype) + return select_fastest_structured_matrix_autodiff(f.jac_prototype, f, ad) + end + + if f.jac_prototype isa AbstractSparseMatrix + 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 - 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 + end +end + +@concrete struct StructuredMatrixAutodiff <: AbstractADType + autodiff <: AbstractADType + sparsity_detection +end + +function select_fastest_structured_matrix_autodiff( + prototype::AbstractMatrix, f::NonlinearFunction, ad::AbstractADType) + sparsity_detection = if SciMLBase.has_colorvec(f) + PrecomputedJacobianColorvec(; + jac_prototype = prototype, + f.colorvec, + partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode + ) 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 ArrayInterface.fast_matrix_colors(prototype) + colorvec = if ADTypes.mode(ad) isa ADTypes.ForwardMode + ArrayInterface.matrix_colors(prototype) + else + ArrayInterface.matrix_colors(prototype') + end + PrecomputedJacobianColorvec(; + jac_prototype = prototype, + colorvec, + partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode + ) + else + JacPrototypeSparsityDetection(; jac_prototype = prototype) + end end + return StructuredMatrixAutodiff(ad, sparsity_detection) +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() +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 diff --git a/test/misc/bruss_tests.jl b/test/misc/bruss_tests.jl index 8ea1fdb18..7b63e421c 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, Symbolics const N = 32 const xyd_brusselator = range(0, stop = 1, length = N) @@ -46,20 +46,26 @@ 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 + + prob_brusselator_2d_sparse = NonlinearProblem( + NonlinearFunction(brusselator_2d_loop; sparsity = SymbolicsSparsityDetector()), + 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( + jac_prototype = 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) ff_iip = NonlinearFunction(brusselator_2d_loop; jac_prototype) prob_brusselator_2d = NonlinearProblem(ff_iip, u0, p) @@ -68,11 +74,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 From 607f5c46394956eaf2f5112ee8f554b565fc8d06 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 1 Oct 2024 17:17:36 -0400 Subject: [PATCH 04/17] chore: apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/basics/sparsity_detection.md | 4 ++-- docs/src/tutorials/large_systems.md | 22 +++++++++++----------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/src/basics/sparsity_detection.md b/docs/src/basics/sparsity_detection.md index 5124a0009..2744d74d3 100644 --- a/docs/src/basics/sparsity_detection.md +++ b/docs/src/basics/sparsity_detection.md @@ -34,7 +34,7 @@ If the `colorvec` is not provided, then it is computed on demand. 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. @@ -58,7 +58,7 @@ 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 diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index 8ed6ca169..4955f9084 100644 --- a/docs/src/tutorials/large_systems.md +++ b/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: @@ -117,11 +117,11 @@ 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 @@ -167,7 +167,7 @@ and `u` and call `jacobian_sparsity` on our function with the example arguments, 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. @@ -224,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. @@ -330,7 +330,7 @@ prob_brusselator_2d_exact_tracer = NonlinearProblem( u0, p; abstol = 1e-10, reltol = 1e-10) prob_brusselator_2d_approx_di = NonlinearProblem( NonlinearFunction(brusselator_2d_loop; - sparsity = DenseSparsityDetector(AutoForwardDiff(); atol=1e-4)), + sparsity = DenseSparsityDetector(AutoForwardDiff(); atol = 1e-4)), u0, p; abstol = 1e-10, reltol = 1e-10) @btime solve(prob_brusselator_2d_exact_symbolics, NewtonRaphson()); From 35bea3e04d535f37d72900549ece2583590f935f Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 1 Oct 2024 19:03:12 -0400 Subject: [PATCH 05/17] chore: apply suggestions from code review --- docs/src/basics/sparsity_detection.md | 9 +++++---- src/NonlinearSolve.jl | 3 ++- src/internal/jacobian.jl | 8 ++++---- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/docs/src/basics/sparsity_detection.md b/docs/src/basics/sparsity_detection.md index 2744d74d3..241399d99 100644 --- a/docs/src/basics/sparsity_detection.md +++ b/docs/src/basics/sparsity_detection.md @@ -37,7 +37,8 @@ If the `colorvec` is not provided, then it is computed on demand. 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. + use the `jac_prototype` argument. `sparsity` must be used to exclusively specify the + sparsity detection algorithm. ## Case II: Sparsity Detection algorithm is provided @@ -46,14 +47,14 @@ 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 = TracerSparsityDetector()), x0) # From SparseConnectivityTracer.jl ``` -Refer to the documentation of DifferentiationInterface.jl for more information on -sparsity detection algorithms. +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 diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 2db06f9b0..575c4cc1c 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -57,7 +57,8 @@ using SparseConnectivityTracer: TracerSparsityDetector # This can be dropped in using SparseDiffTools: SparseDiffTools, JacPrototypeSparsityDetection, PrecomputedJacobianColorvec, init_jacobian, sparse_jacobian, sparse_jacobian!, sparse_jacobian_cache -using SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm +using SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm, + LargestFirst @reexport using SciMLBase, SimpleNonlinearSolve diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index 0bf7e4380..f8264c79a 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -64,7 +64,7 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, if !has_analytic_jac && needs_jac autodiff = construct_concrete_adtype(f, autodiff) using_sparsedifftools = autodiff isa StructuredMatrixAutodiff - # DI can't handle structured matrices + # SparseMatrixColorings can't handle structured matrices if using_sparsedifftools di_extras = nothing uf = JacobianWrapper{iip}(f, p) @@ -249,7 +249,7 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) return AutoSparse( ad; sparsity_detector, - coloring_algorithm = GreedyColoringAlgorithm() + coloring_algorithm = GreedyColoringAlgorithm(LargestFirst()) ) else if ArrayInterface.isstructured(f.jac_prototype) @@ -304,7 +304,7 @@ function select_fastest_structured_matrix_autodiff( JacPrototypeSparsityDetection(; jac_prototype = prototype) end end - return StructuredMatrixAutodiff(ad, sparsity_detection) + return StructuredMatrixAutodiff(AutoSparse(ad), sparsity_detection) end function select_fastest_coloring_algorithm( @@ -314,7 +314,7 @@ function select_fastest_coloring_algorithm( ADTypes.mode(ad) isa ADTypes.ReverseMode, :row, :column)}( prototype, f.colorvec) end - return GreedyColoringAlgorithm() + return GreedyColoringAlgorithm(LargestFirst()) end function construct_concrete_adtype(f::NonlinearFunction, ad::AutoSparse) From fb8b1eee0d8bce6f4d9e6898fcccca47bec89301 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 2 Oct 2024 10:46:33 -0400 Subject: [PATCH 06/17] test: structured jacobians --- test/misc/structured_jacobian_tests.jl | 67 ++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 test/misc/structured_jacobian_tests.jl 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 From 294ec090b7a377143704886304cb03355547c265 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 2 Oct 2024 12:17:49 -0400 Subject: [PATCH 07/17] feat: using DI for structured Jacobians --- Project.toml | 4 +- src/NonlinearSolve.jl | 3 - src/internal/jacobian.jl | 118 +++++++-------------------------------- 3 files changed, 22 insertions(+), 103 deletions(-) diff --git a/Project.toml b/Project.toml index 7d8f1f907..0f77d7966 100644 --- a/Project.toml +++ b/Project.toml @@ -30,7 +30,6 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" -SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" @@ -146,6 +145,7 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4" +SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -155,4 +155,4 @@ 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", "SparseDiffTools", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"] diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 575c4cc1c..ab5e31afc 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -54,9 +54,6 @@ using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJac ## Sparse AD Support using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC 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, LargestFirst diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index f8264c79a..c152fc749 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, @@ -63,31 +62,13 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, if !has_analytic_jac && needs_jac autodiff = construct_concrete_adtype(f, autodiff) - using_sparsedifftools = autodiff isa StructuredMatrixAutodiff - # SparseMatrixColorings can't handle structured matrices - if using_sparsedifftools - di_extras = nothing - uf = JacobianWrapper{iip}(f, p) - sdifft_extras = if iip - sparse_jacobian_cache( - autodiff.autodiff, autodiff.sparsity_detection, uf, fu, u) - else - sparse_jacobian_cache(autodiff.autodiff, autodiff.sparsity_detection, - uf, __maybe_mutable(u, autodiff); fx = fu) - end - autodiff = autodiff.autodiff # For saving we unwrap + 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(prob.p)) - else - DI.prepare_jacobian(f, autodiff, u, Constant(prob.p)) - end + DI.prepare_jacobian(f, autodiff, u, Constant(prob.p)) end else - using_sparsedifftools = false di_extras = nothing - sdifft_extras = nothing end J = if !needs_jac @@ -98,22 +79,18 @@ 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 !using_sparsedifftools - # 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 jac_proto = if eltype(f.jac_prototype) <: Bool @@ -126,20 +103,19 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, 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) @@ -172,27 +148,16 @@ 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 @@ -207,10 +172,6 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) end return ad # No sparse AD else - if ArrayInterface.isstructured(f.jac_prototype) - return select_fastest_structured_matrix_autodiff(f.jac_prototype, f, ad) - end - return AutoSparse( ad; sparsity_detector = KnownJacobianSparsityDetector(f.jac_prototype), @@ -227,10 +188,6 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) Base.depwarn("`sparsity::typeof($(typeof(f.sparsity)))` is deprecated. \ Pass it as `jac_prototype` instead.", :NonlinearSolve) - if ArrayInterface.isstructured(f.sparsity) - return select_fastest_structured_matrix_autodiff(f.sparsity, f, ad) - end - return AutoSparse( ad; sparsity_detector = KnownJacobianSparsityDetector(f.sparsity), @@ -252,11 +209,8 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) coloring_algorithm = GreedyColoringAlgorithm(LargestFirst()) ) else - if ArrayInterface.isstructured(f.jac_prototype) - return select_fastest_structured_matrix_autodiff(f.jac_prototype, f, ad) - end - - if f.jac_prototype isa AbstractSparseMatrix + if f.jac_prototype isa AbstractSparseMatrix || + ArrayInterface.isstructured(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 \ @@ -275,38 +229,6 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) end end -@concrete struct StructuredMatrixAutodiff <: AbstractADType - autodiff <: AbstractADType - sparsity_detection -end - -function select_fastest_structured_matrix_autodiff( - prototype::AbstractMatrix, f::NonlinearFunction, ad::AbstractADType) - sparsity_detection = if SciMLBase.has_colorvec(f) - PrecomputedJacobianColorvec(; - jac_prototype = prototype, - f.colorvec, - partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode - ) - else - if ArrayInterface.fast_matrix_colors(prototype) - colorvec = if ADTypes.mode(ad) isa ADTypes.ForwardMode - ArrayInterface.matrix_colors(prototype) - else - ArrayInterface.matrix_colors(prototype') - end - PrecomputedJacobianColorvec(; - jac_prototype = prototype, - colorvec, - partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode - ) - else - JacPrototypeSparsityDetection(; jac_prototype = prototype) - end - end - return StructuredMatrixAutodiff(AutoSparse(ad), sparsity_detection) -end - function select_fastest_coloring_algorithm( prototype, f::NonlinearFunction, ad::AbstractADType) if SciMLBase.has_colorvec(f) From 7836f04a221d12e6c14c1ff770b8aa2b81bbaf1e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 2 Oct 2024 12:51:55 -0400 Subject: [PATCH 08/17] docs: add a table to guarantee selections --- docs/src/basics/sparsity_detection.md | 39 ++++++++++++++++++++------- src/internal/jacobian.jl | 28 +++++++++++++------ 2 files changed, 50 insertions(+), 17 deletions(-) diff --git a/docs/src/basics/sparsity_detection.md b/docs/src/basics/sparsity_detection.md index 241399d99..fc03528ac 100644 --- a/docs/src/basics/sparsity_detection.md +++ b/docs/src/basics/sparsity_detection.md @@ -6,6 +6,34 @@ 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 @@ -27,19 +55,12 @@ prob = NonlinearProblem( 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. `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. `sparsity` must be used to exclusively specify the - sparsity detection algorithm. - ## Case II: Sparsity Detection algorithm is provided If you don't have a Sparse Jacobian Prototype, but you know the which sparsity detection @@ -59,7 +80,7 @@ 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 diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index c152fc749..93d10f98b 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -172,6 +172,13 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) end return ad # No sparse AD else + 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), @@ -181,13 +188,14 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) end else if f.sparsity isa AbstractMatrix - if f.jac_prototype !== nothing && f.jac_prototype !== f.sparsity - throw(ArgumentError("`sparsity::AbstractMatrix` and `jac_prototype` cannot \ - be both provided. Pass only `jac_prototype`.")) + 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 - Base.depwarn("`sparsity::typeof($(typeof(f.sparsity)))` is deprecated. \ - Pass it as `jac_prototype` instead.", - :NonlinearSolve) return AutoSparse( ad; sparsity_detector = KnownJacobianSparsityDetector(f.sparsity), @@ -209,8 +217,7 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) coloring_algorithm = GreedyColoringAlgorithm(LargestFirst()) ) else - if f.jac_prototype isa AbstractSparseMatrix || - ArrayInterface.isstructured(f.jac_prototype) + 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 \ @@ -254,3 +261,8 @@ 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 From 89b679c270e31dd7343868533872c6dae5f03f74 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 2 Oct 2024 12:54:28 -0400 Subject: [PATCH 09/17] chore: apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/src/basics/sparsity_detection.md | 48 +++++++++++++-------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/docs/src/basics/sparsity_detection.md b/docs/src/basics/sparsity_detection.md index fc03528ac..de924c398 100644 --- a/docs/src/basics/sparsity_detection.md +++ b/docs/src/basics/sparsity_detection.md @@ -9,30 +9,30 @@ 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())` | +|:-------------------------- |:----------------- |:------------ |:------------------------------------------------ |:----------------------------------------- | +| ❌ | ❌ | `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` | 🔴 | 🔴 | +| `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. +| `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 @@ -55,7 +55,7 @@ prob = NonlinearProblem( 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. `ADTypes.mode(ad) isa ADTypes.ForwardMode` will assume that the colorvec is the column colorvec, otherwise we will assume that the colorvec is the @@ -80,7 +80,7 @@ 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 From 76fe551ca22970a5afc5bb32bba7ebd3a866ca09 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 2 Oct 2024 13:16:14 -0400 Subject: [PATCH 10/17] fix: remove stale load --- src/NonlinearSolve.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index ab5e31afc..d4f0898e1 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -30,8 +30,7 @@ using MaybeInplace: @bb using Printf: @printf using Preferences: Preferences, @load_preference, @set_preferences! using RecursiveArrayTools: recursivecopy! -using SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem, - _unwrap_val, isinplace, NLStats +using SciMLBase: AbstractNonlinearAlgorithm, AbstractNonlinearProblem, _unwrap_val, isinplace, NLStats using SciMLOperators: AbstractSciMLOperator using Setfield: @set! using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix From acf377f9c1481b3f719a37f5d97b8023a7e2ca9d Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 2 Oct 2024 13:19:33 -0400 Subject: [PATCH 11/17] chore: apply formatting suggestion Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/NonlinearSolve.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index d4f0898e1..4d78e2195 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -30,7 +30,8 @@ using MaybeInplace: @bb using Printf: @printf using Preferences: Preferences, @load_preference, @set_preferences! using RecursiveArrayTools: recursivecopy! -using SciMLBase: AbstractNonlinearAlgorithm, AbstractNonlinearProblem, _unwrap_val, isinplace, NLStats +using SciMLBase: AbstractNonlinearAlgorithm, AbstractNonlinearProblem, _unwrap_val, + isinplace, NLStats using SciMLOperators: AbstractSciMLOperator using Setfield: @set! using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix From 1356cfeb7b8bd45933aec6218b3cd123ceed4e00 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 3 Oct 2024 09:37:00 -0400 Subject: [PATCH 12/17] docs: remove Symbolics and SparseDiffTools --- docs/Project.toml | 6 +-- docs/make.jl | 6 ++- docs/src/basics/autodiff.md | 66 +++++++++-------------------- docs/src/tutorials/large_systems.md | 19 ++++----- test/misc/qa_tests.jl | 2 +- 5 files changed, 35 insertions(+), 64 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index 48d18ffa9..30a7f5035 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,4 +1,5 @@ [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" @@ -18,13 +19,12 @@ 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" 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" @@ -44,8 +44,6 @@ SciMLBase = "2.4" SciMLJacobianOperators = "0.1" SimpleNonlinearSolve = "1" SparseConnectivityTracer = "0.6.5" -SparseDiffTools = "2.14" StaticArrays = "1" SteadyStateDiffEq = "2" Sundials = "4.11" -Symbolics = "4, 5, 6" diff --git a/docs/make.jl b/docs/make.jl index e11a04149..88381a5b6 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -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..46bc0c576 100644 --- a/docs/src/basics/autodiff.md +++ b/docs/src/basics/autodiff.md @@ -1,60 +1,32 @@ # 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): Finite differencing using `FiniteDiff.jl`, not + optimal but always applicable. + - [`AutoFiniteDifferences`](@extref ADTypes): 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): The best choice for dense problems. + - [`AutoPolyesterForwardDiff`](@extref ADTypes): Might be faster than + [`AutoForwardDiff`](@extref ADTypes) 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. + - [`AutoEnzyme`](@ref): Uses `Enzyme.jl` Reverse Mode and works for both in-place and + out-of-place functions. -!!! note - - If `PolyesterForwardDiff.jl` is installed and loaded, then `SimpleNonlinearSolve.jl` - will automatically use `AutoPolyesterForwardDiff` as the default AD backend. +!!! tip -!!! note - - 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/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index 4955f9084..a27105227 100644 --- a/docs/src/tutorials/large_systems.md +++ b/docs/src/tutorials/large_systems.md @@ -161,22 +161,23 @@ 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 +[ADTypes.jl](https://github.com/SciML/ADTypes.jl) that specifies the interface for sparsity +detection via [`jacobian_sparsity`](@extref ADTypes). 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 - Alternatively you can use the `SparseConnectivityTracer.jl` package to automatically - generate a sparse Jacobian. + External packages like `SparseConnectivityTracer.jl` and `Symbolics.jl` provide the + actual implementation of sparsity detection. ```@example ill_conditioned_nlprob -using Symbolics +using SparseConnectivityTracer, ADTypes -du0 = copy(u0) -jac_sparsity = Symbolics.jacobian_sparsity( - (du, u) -> brusselator_2d_loop(du, u, p), du0, u0) +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 @@ -322,9 +323,6 @@ sparsity detection. Let's compare the two by setting the sparsity detection algo ```@example ill_conditioned_nlprob 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) @@ -333,7 +331,6 @@ prob_brusselator_2d_approx_di = NonlinearProblem( sparsity = DenseSparsityDetector(AutoForwardDiff(); atol = 1e-4)), u0, p; abstol = 1e-10, reltol = 1e-10) -@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 diff --git a/test/misc/qa_tests.jl b/test/misc/qa_tests.jl index 5c60e5339..f6c0d8cb4 100644 --- a/test/misc/qa_tests.jl +++ b/test/misc/qa_tests.jl @@ -18,7 +18,7 @@ end using NonlinearSolve, ADTypes, SimpleNonlinearSolve, SciMLBase import BandedMatrices, FastLevenbergMarquardt, FixedPointAcceleration, LeastSquaresOptim, MINPACK, NLsolve, NLSolvers, SIAMFANLEquations, SpeedMapping, - Symbolics, Zygote + Zygote using ExplicitImports From 8449406a07982db3a331d69799a708d32004fc60 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 3 Oct 2024 09:39:50 -0400 Subject: [PATCH 13/17] docs: remove unnecessary ADTypes docs --- docs/src/basics/autodiff.md | 7 +- src/NonlinearSolve.jl | 1 - src/adtypes.jl | 140 ------------------------------------ 3 files changed, 4 insertions(+), 144 deletions(-) delete mode 100644 src/adtypes.jl diff --git a/docs/src/basics/autodiff.md b/docs/src/basics/autodiff.md index 46bc0c576..df59ceafc 100644 --- a/docs/src/basics/autodiff.md +++ b/docs/src/basics/autodiff.md @@ -22,9 +22,10 @@ ## 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 works for both in-place and - out-of-place functions. + - [`AutoZygote`](@extref ADTypes): The fastest choice for non-mutating array-based (BLAS) + functions. + - [`AutoEnzyme`](@extref ADTypes): Uses `Enzyme.jl` Reverse Mode and works for both + in-place and out-of-place functions. !!! tip diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 4d78e2195..5043ed865 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -68,7 +68,6 @@ 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 From fa575c6f7cabfa2b75cd867275223d29938aac7a Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 3 Oct 2024 09:43:22 -0400 Subject: [PATCH 14/17] test: remove sparsedifftools and symbolics from tests --- Project.toml | 6 +----- docs/src/basics/autodiff.md | 4 ++-- test/core/rootfind_tests.jl | 2 +- test/misc/bruss_tests.jl | 8 ++++---- 4 files changed, 8 insertions(+), 12 deletions(-) diff --git a/Project.toml b/Project.toml index 0f77d7966..7e0bcb4c8 100644 --- a/Project.toml +++ b/Project.toml @@ -108,7 +108,6 @@ Setfield = "1.1.1" SimpleNonlinearSolve = "1.12.3" SparseArrays = "1.10" SparseConnectivityTracer = "0.6.5" -SparseDiffTools = "2.22" SparseMatrixColorings = "0.4.2" SpeedMapping = "0.3" StableRNGs = "1" @@ -116,7 +115,6 @@ 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" @@ -145,14 +143,12 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4" -SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" 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", "SparseDiffTools", "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/src/basics/autodiff.md b/docs/src/basics/autodiff.md index df59ceafc..dc3166c6e 100644 --- a/docs/src/basics/autodiff.md +++ b/docs/src/basics/autodiff.md @@ -1,7 +1,7 @@ # 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. @@ -28,6 +28,6 @@ in-place and out-of-place functions. !!! tip - + For sparsity detection and sparse AD take a look at [sparsity detection](@ref sparsity-detection). 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 7b63e421c..b7abd634b 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, SparseConnectivityTracer, Symbolics + using LinearAlgebra, SparseArrays, SparseConnectivityTracer, ADTypes const N = 32 const xyd_brusselator = range(0, stop = 1, length = N) @@ -63,9 +63,9 @@ NewtonRaphson(autodiff = AutoSparse(AutoFiniteDiff())); abstol = 1e-8) @test norm(sol.resid, Inf) < 1e-8 - du0 = copy(u0) - jac_prototype = Symbolics.jacobian_sparsity( - (du, u) -> brusselator_2d_loop(du, u, p), du0, u0) + 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) From d6b5536314798dd745a89805cbe2db67166debd6 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 3 Oct 2024 09:45:43 -0400 Subject: [PATCH 15/17] refactor: remove Zygote extension --- Project.toml | 2 -- ext/NonlinearSolveZygoteExt.jl | 7 ------- src/NonlinearSolve.jl | 3 --- src/internal/helpers.jl | 2 +- test/misc/bruss_tests.jl | 6 ------ test/misc/qa_tests.jl | 3 +-- 6 files changed, 2 insertions(+), 21 deletions(-) delete mode 100644 ext/NonlinearSolveZygoteExt.jl diff --git a/Project.toml b/Project.toml index 7e0bcb4c8..69dd08304 100644 --- a/Project.toml +++ b/Project.toml @@ -45,7 +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" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] NonlinearSolveBandedMatricesExt = "BandedMatrices" @@ -57,7 +56,6 @@ NonlinearSolveNLSolversExt = "NLSolvers" NonlinearSolveNLsolveExt = "NLsolve" NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations" NonlinearSolveSpeedMappingExt = "SpeedMapping" -NonlinearSolveZygoteExt = "Zygote" [compat] ADTypes = "1.9" 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 5043ed865..50051fcb0 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -61,9 +61,6 @@ using SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm, const DI = DifferentiationInterface -# Type-Inference Friendly Check for Extension Loading -is_extension_loaded(::Val) = false - const True = Val(true) const False = Val(false) 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/test/misc/bruss_tests.jl b/test/misc/bruss_tests.jl index b7abd634b..80b1bd540 100644 --- a/test/misc/bruss_tests.jl +++ b/test/misc/bruss_tests.jl @@ -52,12 +52,6 @@ sol = solve(prob_brusselator_2d_sparse, NewtonRaphson(); abstol = 1e-8) @test norm(sol.resid, Inf) < 1e-8 - prob_brusselator_2d_sparse = NonlinearProblem( - NonlinearFunction(brusselator_2d_loop; sparsity = SymbolicsSparsityDetector()), - 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) diff --git a/test/misc/qa_tests.jl b/test/misc/qa_tests.jl index f6c0d8cb4..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, - Zygote + LeastSquaresOptim, MINPACK, NLsolve, NLSolvers, SIAMFANLEquations, SpeedMapping using ExplicitImports From c562fa11079a0d030c9716b465d8831df6d664ec Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 3 Oct 2024 10:27:19 -0400 Subject: [PATCH 16/17] docs: add documenter interlinks as a dep --- docs/Project.toml | 2 ++ docs/make.jl | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/Project.toml b/docs/Project.toml index 30a7f5035..ab35d42ae 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -7,6 +7,7 @@ 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" @@ -32,6 +33,7 @@ DiffEqBase = "6.136" DifferentiationInterface = "0.6" Documenter = "1" DocumenterCitations = "1" +DocumenterInterLinks = "1.0.0" IncompleteLU = "0.2" InteractiveUtils = "<0.0.1, 1" LinearSolve = "2" diff --git a/docs/make.jl b/docs/make.jl index 88381a5b6..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 From 6dcb0da7aeda2082e1b9fbc0b599fa8f1abd077a Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 3 Oct 2024 11:13:31 -0400 Subject: [PATCH 17/17] docs: fix external references --- docs/src/basics/autodiff.md | 23 ++++++++++++----------- docs/src/tutorials/code_optimization.md | 2 +- docs/src/tutorials/large_systems.md | 9 +++++---- src/algorithms/extension_algs.jl | 2 +- 4 files changed, 19 insertions(+), 17 deletions(-) diff --git a/docs/src/basics/autodiff.md b/docs/src/basics/autodiff.md index dc3166c6e..d2d01e00b 100644 --- a/docs/src/basics/autodiff.md +++ b/docs/src/basics/autodiff.md @@ -8,24 +8,25 @@ ## Summary of Finite Differencing Backends - - [`AutoFiniteDiff`](@extref ADTypes): Finite differencing using `FiniteDiff.jl`, not - optimal but always applicable. - - [`AutoFiniteDifferences`](@extref ADTypes): Finite differencing using - `FiniteDifferences.jl`, 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`](@extref ADTypes): The best choice for dense problems. - - [`AutoPolyesterForwardDiff`](@extref ADTypes): Might be faster than - [`AutoForwardDiff`](@extref ADTypes) for large problems. Requires + - [`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`](@extref ADTypes): The fastest choice for non-mutating array-based (BLAS) - functions. - - [`AutoEnzyme`](@extref ADTypes): Uses `Enzyme.jl` Reverse Mode and works for both - in-place and out-of-place functions. + - [`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. !!! tip 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 a27105227..17b768288 100644 --- a/docs/src/tutorials/large_systems.md +++ b/docs/src/tutorials/large_systems.md @@ -162,10 +162,11 @@ sparse differentiation! One of the useful companion tools for NonlinearSolve.jl is [ADTypes.jl](https://github.com/SciML/ADTypes.jl) that specifies the interface for sparsity -detection via [`jacobian_sparsity`](@extref ADTypes). 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`. +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 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