fix: more options

Author: avik-pal

.gitignore

@@ -22,4 +22,5 @@ pdf/
 notebook/
 markdown/
 
-.vscode
+.vscode
+*.html

benchmarks/NonlinearProblem/Manifest.toml

@@ -2223,9 +2223,9 @@ version = "0.2.0"
 
 [[deps.SciMLBase]]
 deps = ["ADTypes", "Accessors", "ArrayInterface", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "Moshi", "PrecompileTools", "Preferences", "Printf", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "SciMLStructures", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface"]
-git-tree-sha1 = "c4ce89e19f2a220e34c0f377fb41516b642ec899"
+git-tree-sha1 = "f2db9ab9d33130df3be35be9438da51a3416d528"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
-version = "2.83.1"
+version = "2.84.0"
 
     [deps.SciMLBase.extensions]
     SciMLBaseChainRulesCoreExt = "ChainRulesCore"

benchmarks/NonlinearProblem/Project.toml

@@ -59,3 +59,6 @@ StableRNGs = "1"
 StaticArrays = "1"
 Sundials = "4.22"
 Symbolics = "5, 6"
+
+[sources]
+NonlinearSolve = {url = "https://github.com/SciML/NonlinearSolve.jl", rev = "ap/petsc_debug"}

benchmarks/NonlinearProblem/bruss.jmd

@@ -360,13 +360,13 @@ solvers_scaling = [
     (; pkg = :nonlinearsolve, sparsity = :exact,  name = "NR (Exact Sparsity)",                 alg = NewtonRaphson()),
     (; pkg = :wrapper,        sparsity = :none,   name = "NR [NLsolve.jl]",                     alg = NLsolveJL(; method = :newton, autodiff = :forward)),
     (; pkg = :wrapper,        sparsity = :none,   name = "Mod. NR [Sundials]",                  alg = KINSOL()),
-    (; pkg = :wrapper,        sparsity = :none,   name = "NR [PETSc] (No Sparsity)",            alg = PETScSNES(; snes_type = "newtonls")),
-    (; pkg = :wrapper,        sparsity = :exact,  name = "NR [PETSc] (Exact Sparsity)",         alg = PETScSNES(; snes_type = "newtonls")),
+    (; pkg = :wrapper,        sparsity = :none,   name = "NR [PETSc] (No Sparsity)",            alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", autodiff = missing)),
+    (; pkg = :wrapper,        sparsity = :exact,  name = "NR [PETSc] (Exact Sparsity)",         alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic")),
 
     (; pkg = :nonlinearsolve, sparsity = :none,   name = "TR (No Sparsity)",                    alg = TrustRegion(; radius_update_scheme = RUS.NLsolve)),
     (; pkg = :nonlinearsolve, sparsity = :exact,  name = "TR (Exact Sparsity)",                 alg = TrustRegion(; radius_update_scheme = RUS.NLsolve)),
     (; pkg = :wrapper,        sparsity = :none,   name = "TR [NLsolve.jl]",                     alg = NLsolveJL(; autodiff = :forward)),
-    (; pkg = :wrapper,        sparsity = :none,   name = "TR [PETSc] (No Sparsity)",            alg = PETScSNES(; snes_type = "newtontr")),
+    (; pkg = :wrapper,        sparsity = :none,   name = "TR [PETSc] (No Sparsity)",            alg = PETScSNES(; snes_type = "newtontr", autodiff = missing)),
     (; pkg = :wrapper,        sparsity = :exact,  name = "TR [PETSc] (Exact Sparsity)",         alg = PETScSNES(; snes_type = "newtontr")),
 
     (; pkg = :wrapper,        sparsity = :none,   name = "Mod. Powell [MINPACK]",               alg = CMINPACK()),
@@ -405,7 +405,7 @@ for (i, N) in enumerate(Ns)
             (alg isa GeneralizedFirstOrderAlgorithm && alg.name == :TrustRegion && N > 64) ||
             (alg isa NLsolveJL && N > 128 && alg.method == :newton) ||
             (alg isa GeneralizedFirstOrderAlgorithm && alg.name == :NewtonRaphson && N > 128 && ptype == :none) ||
-            (alg isa PETScSNES && N > 64 && ptype == :exact)
+            (alg isa PETScSNES && N > 64)
             # The last benchmark failed so skip this too
             runtimes_scaling[j, i] = NaN
             @warn "$(name): Would Have Timed out"
@@ -551,15 +551,23 @@ end
 Ns = 2 .^ (2:7)
 
 solvers_scaling_jacobian_free = [
-    (; pkg = :nonlinearsolve,  name = "Newton Krylov",              alg = NewtonRaphson(; linsolve = KrylovJL_GMRES())),
-    (; pkg = :nonlinearsolve,  name = "Newton Krylov (ILU)",        alg = NewtonRaphson(; linsolve = KrylovJL_GMRES(; precs = incompletelu), concrete_jac = true)),
-    (; pkg = :nonlinearsolve,  name = "Newton Krylov (AMG)",        alg = NewtonRaphson(; linsolve = KrylovJL_GMRES(; precs = algebraicmultigrid), concrete_jac = true)),
-    (; pkg = :nonlinearsolve,  name = "Newton Krylov (AMG Jacobi)", alg = NewtonRaphson(; linsolve = KrylovJL_GMRES(; precs = algebraicmultigrid_jacobi), concrete_jac = true)),
-    (; pkg = :nonlinearsolve,  name = "TR Krylov",                  alg = TrustRegion(; linsolve = KrylovJL_GMRES())),
-    (; pkg = :nonlinearsolve,  name = "TR Krylov (ILU)",            alg = TrustRegion(; linsolve = KrylovJL_GMRES(; precs = incompletelu), concrete_jac = true)),
-    (; pkg = :nonlinearsolve,  name = "TR Krylov (AMG)",            alg = TrustRegion(; linsolve = KrylovJL_GMRES(; precs = algebraicmultigrid), concrete_jac = true)),
-    (; pkg = :nonlinearsolve,  name = "TR Krylov (AMG Jacobi)",     alg = TrustRegion(; linsolve = KrylovJL_GMRES(; precs = algebraicmultigrid_jacobi), concrete_jac = true)),
-    (; pkg = :wrapper,         name = "Newton Krylov [Sundials]",   alg = KINSOL(; linear_solver = :GMRES)),
+    (; pkg = :nonlinearsolve,  name = "Newton Krylov",                      alg = NewtonRaphson(; linsolve = KrylovJL_GMRES())),
+    (; pkg = :nonlinearsolve,  name = "Newton Krylov (ILU)",                alg = NewtonRaphson(; linsolve = KrylovJL_GMRES(; precs = incompletelu), concrete_jac = true)),
+    (; pkg = :nonlinearsolve,  name = "Newton Krylov (AMG)",                alg = NewtonRaphson(; linsolve = KrylovJL_CG(; precs = algebraicmultigrid), concrete_jac = true)),
+    (; pkg = :nonlinearsolve,  name = "Newton Krylov (AMG Jacobi)",         alg = NewtonRaphson(; linsolve = KrylovJL_CG(; precs = algebraicmultigrid_jacobi), concrete_jac = true)),
+    (; pkg = :wrapper,         name = "Newton Krylov [PETSc]",              alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "gmres", snes_mf = true)),
+    (; pkg = :wrapper,         name = "Newton Krylov (ILU) [PETSc]",        alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "gmres", pc_type = "ilu")),
+    (; pkg = :wrapper,         name = "Newton Krylov (AMG) [PETSc]",        alg = PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "cg", pc_type = "gamg")),
+    (; 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")),
+    (; pkg = :nonlinearsolve,  name = "TR Krylov",                          alg = TrustRegion(; linsolve = KrylovJL_GMRES())),
+    (; pkg = :nonlinearsolve,  name = "TR Krylov (ILU)",                    alg = TrustRegion(; linsolve = KrylovJL_GMRES(; precs = incompletelu), concrete_jac = true)),
+    (; pkg = :nonlinearsolve,  name = "TR Krylov (AMG)",                    alg = TrustRegion(; linsolve = KrylovJL_CG(; precs = algebraicmultigrid), concrete_jac = true)),
+    (; pkg = :nonlinearsolve,  name = "TR Krylov (AMG Jacobi)",             alg = TrustRegion(; linsolve = KrylovJL_CG(; precs = algebraicmultigrid_jacobi), concrete_jac = true)),
+    (; pkg = :wrapper,         name = "TR Krylov [PETSc]",                  alg = PETScSNES(; snes_type = "newtontr", ksp_type = "gmres", snes_mf = true)),
+    (; pkg = :wrapper,         name = "TR Krylov (ILU) [PETSc]",            alg = PETScSNES(; snes_type = "newtontr", ksp_type = "gmres", pc_type = "ilu")),
+    (; pkg = :wrapper,         name = "TR Krylov (AMG) [PETSc]",            alg = PETScSNES(; snes_type = "newtontr", ksp_type = "cg", pc_type = "gamg")),
+    (; 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")),
+    (; pkg = :wrapper,         name = "Newton Krylov [Sundials]",           alg = KINSOL(; linear_solver = :GMRES)),
 ]
 
 runtimes_scaling = zeros(length(solvers_scaling_jacobian_free), length(Ns)) .- 1
@@ -695,16 +703,24 @@ exploit sparsity will automatically do so.
 solvers_all = [
     (; pkg = :nonlinearsolve, name = "Default PolyAlg",           solver = Dict(:alg => FastShortcutNonlinearPolyalg())),
     (; pkg = :nonlinearsolve, name = "RobustMultiNewton (GMRES)", solver = Dict(:alg => RobustMultiNewton(; linsolve = KrylovJL_GMRES()))),
+
     (; pkg = :nonlinearsolve, name = "Newton Raphson",            solver = Dict(:alg => NewtonRaphson(; linsolve = nothing))),
     (; pkg = :nonlinearsolve, name = "Newton Krylov",             solver = Dict(:alg => NewtonRaphson(; linsolve = KrylovJL_GMRES()))),
     (; pkg = :nonlinearsolve, name = "Trust Region",              solver = Dict(:alg => TrustRegion())),
     (; pkg = :nonlinearsolve, name = "TR Krylov",                 solver = Dict(:alg => TrustRegion(; linsolve = KrylovJL_GMRES()))),
+
     (; pkg = :wrapper,        name = "NR [NLsolve.jl]",           solver = Dict(:alg => NLsolveJL(; method = :newton, autodiff = :forward))),
     (; pkg = :wrapper,        name = "TR [NLsolve.jl]",           solver = Dict(:alg => NLsolveJL(; autodiff = :forward))),
+
     (; pkg = :wrapper,        name = "NR [Sundials]",             solver = Dict(:alg => KINSOL())),
     (; pkg = :wrapper,        name = "Newton Krylov [Sundials]",  solver = Dict(:alg => KINSOL(; linear_solver = :GMRES))),
 
     (; pkg = :wrapper,        name = "Mod. Powell [MINPACK]",     solver = Dict(:alg => CMINPACK())),
+
+    (; pkg = :wrapper,        name = "NR [PETSc]",  solver = Dict(:alg => PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", autodiff = missing))),
+    (; pkg = :wrapper,        name = "TR [PETSc]",  solver = Dict(:alg => PETScSNES(; snes_type = "newtontr", autodiff = missing))),
+    (; pkg = :wrapper,        name = "Newton Krylov [PETSc]",  solver = Dict(:alg => PETScSNES(; snes_type = "newtonls", snes_linesearch_type = "basic", ksp_type = "gmres", autodiff = missing, snes_mf = true))),
+    (; pkg = :wrapper,        name = "TR Krylov [PETSc]",  solver = Dict(:alg => PETScSNES(; snes_type = "newtontr", ksp_type = "gmres", autodiff = missing, snes_mf = true))),
 ];
 ```
 

benchmarks/NonlinearProblem/nonlinear_solver_23_tests.jmd

@@ -21,12 +21,12 @@ Furthermore, for NonlinearSolve.jl's Newton Raphson method we try the following
 - `BackTracking`
 
 and for NonlinearSolve.jl's Trust Region we try the following Radius Update schemes (in addition to the default):
-- `NLsolve` 
-- `NocedalWright` 
-- `Hei` 
-- `Yuan` 
-- `Bastin` 
-- `Fan` 
+- `NLsolve`
+- `NocedalWright`
+- `Hei`
+- `Yuan`
+- `Bastin`
+- `Fan`
 and finally for NonlinearSolve.jl's Levenberg-Marquardt method why try using both the default `α_geodesic` value (`0.75`) and a modified value (`0.5`), and also with and without setting the `CholeskyFactorization` linear solver.
 
 For each benchmarked problem, the second, third, and fourth plots compares the performance of NonlinearSolve's Newton Raphson, Trust Region, and Levenberg-Marquardt methods, respectively. The first plot compares the best methods from each of these categories to the various methods available from other packages. At the end of the benchmarks, we print a summary table of which solvers succeeded for which problems.
@@ -38,7 +38,7 @@ Fetch required packages.
 ```julia
 using NonlinearSolve, LinearSolve, StaticArrays, Sundials, Setfield,
     BenchmarkTools, LinearAlgebra, DiffEqDevTools, NonlinearProblemLibrary, CairoMakie,
-    RecursiveFactorization
+    RecursiveFactorization, Enzyme
 import PolyesterForwardDiff, MINPACK, NLsolve, LineSearches
 
 const RUS = RadiusUpdateSchemes;
@@ -92,13 +92,34 @@ reltols = 1.0 ./ 10.0 .^ (4:12);
 Prepares various helper functions for benchmarking a specific problem.
 
 ```julia
+function set_ad_chunksize(solvers, u0)
+    ck = NonlinearSolve.pickchunksize(u0)
+    for i in eachindex(solvers)
+        @set! solvers[i].solver[:alg] = __set_ad_chunksize(solvers[i].solver[:alg], ck)
+    end
+    return solvers
+end
+
+function __set_ad_chunksize(solver::GeneralizedFirstOrderAlgorithm, ck)
+    ad = AutoPolyesterForwardDiff(; chunksize = ck)
+    return GeneralizedFirstOrderAlgorithm(; solver.descent, solver.linesearch,
+        solver.trustregion, jvp_autodiff = ad, solver.max_shrink_times, solver.vjp_autodiff,
+        concrete_jac = solver.concrete_jac, name = solver.name)
+end
+function __set_ad_chunksize(solver::NonlinearSolvePolyAlgorithm, ck)
+    algs = [__set_ad_chunksize(alg, ck) for alg in solver.algs]
+    return NonlinearSolvePolyAlgorithm(algs; solver.start_index)
+end
+__set_ad_chunksize(solver, ck) = solver
+
 # Benchmarks a specific problem, checks which solvers can solve it and their performance
 function benchmark_problem!(prob_name; solver_tracker=solver_tracker)
     # Finds the problem and the true solution.
     prob = nlprob_23_testcases[prob_name]
 
     # Finds the solvers that can solve the problem
-    successful_solvers = filter(Base.Fix1(check_solver, prob), solvers_all);
+    solvers_concrete = set_ad_chunksize(solvers_all, prob.prob.u0);
+    successful_solvers = filter(solver -> check_solver(prob, solver), solvers_concrete);
     push!(solver_tracker, prob_name => successful_solvers);
 
     # Handles the non-general cases.
@@ -161,12 +182,11 @@ function check_solver(prob, solver)
     return true
 end
 
-# Adds an additional, selected, solver to the general solver set.
 # Adds an additional, selected, solver to the general solver set.
 function add_solver!(solvers_general, selected_solver_name, additional_solver_set, wp)
     if isnothing(selected_solver_name)
         isempty(wp.wps) && return
-        selected_idx = argmin(median.(getfield.(wp.wps, :times))) 
+        selected_idx = argmin(median.(getfield.(wp.wps, :times)))
     else
         selected_idx = findfirst(s -> s.name==selected_solver_name, additional_solver_set)
         isnothing(selected_solver) && error("The $(selected_solver_name) was designated to \