@@ -11,7 +11,7 @@ Fetch required packages
1111```julia
1212using NonlinearSolve, SparseDiffTools, LinearAlgebra, SparseArrays, DiffEqDevTools,
1313 CairoMakie, Symbolics, BenchmarkTools, PolyesterForwardDiff, LinearSolve, Sundials,
14- Enzyme, SparseConnectivityTracer, DifferentiationInterface
14+ Enzyme, SparseConnectivityTracer, DifferentiationInterface, SparseMatrixColorings
1515import NLsolve, MINPACK, PETSc, RecursiveFactorization
1616
1717const RUS = RadiusUpdateSchemes;
@@ -156,31 +156,38 @@ function cache_and_compute_10_jacobians(adtype, f!::F, y, x, p) where {F}
156156 return J
157157end
158158
159- Ns = [2^i for i in 1 :8]
159+ Ns = [2^i for i in 3 :8];
160160
161161adtypes = [
162162 (
163- AutoSparse(AutoFiniteDiff(); sparsity_detector=TracerSparsityDetector()),
163+ AutoSparse(
164+ AutoFiniteDiff();
165+ sparsity_detector=TracerSparsityDetector(),
166+ coloring_algorithm=GreedyColoringAlgorithm(LargestFirst())
167+ ),
164168 [:finitediff, :exact_sparse]
165169 ),
166170 (
167171 AutoSparse(
168172 AutoPolyesterForwardDiff(; chunksize=8);
169- sparsity_detector=TracerSparsityDetector()
173+ sparsity_detector=TracerSparsityDetector(),
174+ coloring_algorithm=GreedyColoringAlgorithm(LargestFirst())
170175 ),
171176 [:polyester, :exact_sparse]
172177 ),
173178 (
174179 AutoSparse(
175180 AutoEnzyme(; mode=Enzyme.Forward);
176- sparsity_detector=TracerSparsityDetector()
181+ sparsity_detector=TracerSparsityDetector(),
182+ coloring_algorithm=GreedyColoringAlgorithm(LargestFirst())
177183 ),
178184 [:enzyme, :exact_sparse]
179185 ),
180186 (
181187 AutoSparse(
182188 AutoFiniteDiff();
183- sparsity_detector=DenseSparsityDetector(AutoFiniteDiff(); atol=1e-5)
189+ sparsity_detector=DenseSparsityDetector(AutoFiniteDiff(); atol=1e-5),
190+ coloring_algorithm=GreedyColoringAlgorithm(LargestFirst())
184191 ),
185192 [:finitediff, :approx_sparse]
186193 ),
@@ -189,7 +196,8 @@ adtypes = [
189196 AutoPolyesterForwardDiff(; chunksize=8);
190197 sparsity_detector=DenseSparsityDetector(
191198 AutoPolyesterForwardDiff(; chunksize=8); atol=1e-5
192- )
199+ ),
200+ coloring_algorithm=GreedyColoringAlgorithm(LargestFirst())
193201 ),
194202 [:polyester, :approx_sparse]
195203 ),
@@ -198,17 +206,18 @@ adtypes = [
198206 AutoEnzyme(; mode=Enzyme.Forward);
199207 sparsity_detector=DenseSparsityDetector(
200208 AutoEnzyme(; mode=Enzyme.Forward); atol=1e-5
201- )
209+ ),
210+ coloring_algorithm=GreedyColoringAlgorithm(LargestFirst())
202211 ),
203212 [:enzyme, :approx_sparse]
204213 ),
205214 (
206215 AutoPolyesterForwardDiff(; chunksize=8),
207216 [:polyester, :none]
208217 ),
209- ]
218+ ];
210219
211- times = Matrix{Float64}(undef, length(Ns), length(adtypes))
220+ times = Matrix{Float64}(undef, length(Ns), length(adtypes));
212221
213222for (i, N) in enumerate(Ns)
214223 str = "$(lpad(N, 10)) "
@@ -227,6 +236,7 @@ for (i, N) in enumerate(Ns)
227236 end
228237 println(str)
229238end
239+ nothing
230240```
231241
232242Plotting the results.
@@ -255,10 +265,11 @@ fig = begin
255265 colormap=:tableau_20)
256266
257267 ax = Axis(fig[1, 2]; title="Scaling of Sparse Jacobian Computation",
258- titlesize=22, titlegap=10, xscale=log10 , yscale=log10 ,
268+ titlesize=22, titlegap=10, xscale=log2 , yscale=log2 ,
259269 xticksize=20, yticksize=20, xticklabelsize=20, yticklabelsize=20,
260270 xtickwidth=2.5, ytickwidth=2.5, spinewidth=2.5,
261- xlabel=L"Input Dimension ($\mathbf{N}$)", ylabel=L"Time $\mathbf{(s)}$", xlabelsize=22,
271+ xlabel=L"Input Dimension ($\mathbf{N}$)",
272+ ylabel=L"Time $\mathbf{(s)}$", xlabelsize=22,
262273 ylabelsize=22, yaxisposition=:right)
263274
264275 colors = cgrad(:tableau_20, length(adtypes); categorical=true)
@@ -334,7 +345,7 @@ fig = begin
334345 [symbol_to_adname[adtypes[idx][2][1]] for idx in local_sparse_idxs],
335346 [symbol_to_adname[adtypes[idx][2][1]] for idx in non_sparse_idxs],
336347 ],
337- ["Exact Sparsity", "Approx. Local Sparsity", "No Sparsity "];
348+ ["Exact Sparsity", "Approx. Local Sparsity", "Dense "];
338349 position=:rb, framevisible=true, framewidth=2.5, titlesize=18,
339350 labelsize=16, patchsize=(40.0f0, 20.0f0)
340351 )
@@ -447,9 +458,9 @@ fig = begin
447458 theme = Theme(Lines = (cycle = cycle,), Scatter = (cycle = cycle,))
448459 LINESTYLES = Dict(
449460 (:nonlinearsolve, :none) => :solid,
450- # (:nonlinearsolve, :approx) => :dash,
451461 (:nonlinearsolve, :exact) => :dashdot,
452462 # (:simplenonlinearsolve, :none) => :solid,
463+ (:wrapper, :exact) => :dash,
453464 (:wrapper, :none) => :dot,
454465 )
455466
@@ -553,20 +564,20 @@ Ns = 2 .^ (2:7)
553564solvers_scaling_jacobian_free = [
554565 (; pkg = :nonlinearsolve, name = "Newton Krylov", alg = NewtonRaphson(; linsolve = KrylovJL_GMRES())),
555566 (; pkg = :nonlinearsolve, name = "Newton Krylov (ILU)", alg = NewtonRaphson(; linsolve = KrylovJL_GMRES(; precs = incompletelu), concrete_jac = true)),
556- (; pkg = :nonlinearsolve, name = "Newton Krylov (AMG)", alg = NewtonRaphson(; linsolve = KrylovJL_CG (; precs = algebraicmultigrid), concrete_jac = true)),
557- (; pkg = :nonlinearsolve, name = "Newton Krylov (AMG Jacobi)", alg = NewtonRaphson(; linsolve = KrylovJL_CG (; precs = algebraicmultigrid_jacobi), concrete_jac = true)),
567+ (; pkg = :nonlinearsolve, name = "Newton Krylov (AMG)", alg = NewtonRaphson(; linsolve = KrylovJL_GMRES (; precs = algebraicmultigrid), concrete_jac = true)),
568+ (; pkg = :nonlinearsolve, name = "Newton Krylov (AMG Jacobi)", alg = NewtonRaphson(; linsolve = KrylovJL_GMRES (; precs = algebraicmultigrid_jacobi), concrete_jac = true)),
558569 (; pkg = :wrapper, name = "Newton Krylov [PETSc]", alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "gmres", snes_mf = true)),
559570 (; pkg = :wrapper, name = "Newton Krylov (ILU) [PETSc]", alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "gmres", pc_type = "ilu")),
560- (; pkg = :wrapper, name = "Newton Krylov (AMG) [PETSc]", alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "cg ", pc_type = "gamg")),
561- (; pkg = :wrapper, name = "Newton Krylov (AMG Jacobi) [PETSc]", alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "cg ", pc_type = "gamg", mg_levels_ksp_type = "richardson", mg_levels_pc_type = "jacobi")),
571+ (; pkg = :wrapper, name = "Newton Krylov (AMG) [PETSc]", alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "gmres ", pc_type = "gamg")),
572+ (; pkg = :wrapper, name = "Newton Krylov (AMG Jacobi) [PETSc]", alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "gmres ", pc_type = "gamg", mg_levels_ksp_type = "richardson", mg_levels_pc_type = "jacobi")),
562573 (; pkg = :nonlinearsolve, name = "TR Krylov", alg = TrustRegion(; linsolve = KrylovJL_GMRES())),
563574 (; pkg = :nonlinearsolve, name = "TR Krylov (ILU)", alg = TrustRegion(; linsolve = KrylovJL_GMRES(; precs = incompletelu), concrete_jac = true)),
564- (; pkg = :nonlinearsolve, name = "TR Krylov (AMG)", alg = TrustRegion(; linsolve = KrylovJL_CG (; precs = algebraicmultigrid), concrete_jac = true)),
565- (; pkg = :nonlinearsolve, name = "TR Krylov (AMG Jacobi)", alg = TrustRegion(; linsolve = KrylovJL_CG (; precs = algebraicmultigrid_jacobi), concrete_jac = true)),
575+ (; pkg = :nonlinearsolve, name = "TR Krylov (AMG)", alg = TrustRegion(; linsolve = KrylovJL_GMRES (; precs = algebraicmultigrid), concrete_jac = true)),
576+ (; pkg = :nonlinearsolve, name = "TR Krylov (AMG Jacobi)", alg = TrustRegion(; linsolve = KrylovJL_GMRES (; precs = algebraicmultigrid_jacobi), concrete_jac = true)),
566577 (; pkg = :wrapper, name = "TR Krylov [PETSc]", alg = PETScSNES(; snes_type = "newtontr", ksp_type = "gmres", snes_mf = true)),
567578 (; pkg = :wrapper, name = "TR Krylov (ILU) [PETSc]", alg = PETScSNES(; snes_type = "newtontr", ksp_type = "gmres", pc_type = "ilu")),
568- (; pkg = :wrapper, name = "TR Krylov (AMG) [PETSc]", alg = PETScSNES(; snes_type = "newtontr", ksp_type = "cg ", pc_type = "gamg")),
569- (; pkg = :wrapper, name = "TR Krylov (AMG Jacobi) [PETSc]", alg = PETScSNES(; snes_type = "newtontr", ksp_type = "cg ", pc_type = "gamg", mg_levels_ksp_type = "richardson", mg_levels_pc_type = "jacobi")),
579+ (; pkg = :wrapper, name = "TR Krylov (AMG) [PETSc]", alg = PETScSNES(; snes_type = "newtontr", ksp_type = "gmres ", pc_type = "gamg")),
580+ (; pkg = :wrapper, name = "TR Krylov (AMG Jacobi) [PETSc]", alg = PETScSNES(; snes_type = "newtontr", ksp_type = "gmres ", pc_type = "gamg", mg_levels_ksp_type = "richardson", mg_levels_pc_type = "jacobi")),
570581 (; pkg = :wrapper, name = "Newton Krylov [Sundials]", alg = KINSOL(; linear_solver = :GMRES)),
571582]
572583
@@ -787,7 +798,7 @@ fig = begin
787798 cycle = Cycle([:marker], covary = true)
788799 plot_theme = Theme(Lines = (; cycle), Scatter = (; cycle))
789800
790- with_theme(plot_theme) do
801+ with_theme(plot_theme) do
791802 fig = Figure(; size = (WIDTH, HEIGHT))
792803 # `textbf` doesn't work
793804 ax = Axis(fig[1, 1], ylabel = L"Time $\mathbf{(s)}$",
0 commit comments