Address some of the comments

SciML · Jan 16, 2024 · 3db75f0 · 3db75f0
1 parent 210e544
commit 3db75f0
Show file tree

Hide file tree

Showing 9 changed files with 30 additions and 18 deletions.
diff --git a/docs/src/basics/faq.md b/docs/src/basics/faq.md
@@ -138,14 +138,12 @@ computation and the type of this chunksize can't be statically inferred. To fix
 directly specify the chunksize:
 
 ```@example type_unstable
-@code_warntype solve(prob, NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 2)))
+@code_warntype solve(prob, NewtonRaphson(;
+    autodiff = AutoForwardDiff(; chunksize = NonlinearSolve.pickchunksize(prob.u0))))
 nothing # hide
 ```
 
-And boom! Type stable again. For selecting the chunksize the method is:
-
- 1. For small inputs `≤ 12` use `chunksize = <length of input>`
- 2. For larger inputs, use `chunksize = 12`
-
-In general, the chunksize should be `≤ length of input`. However, a very large chunksize
-can lead to excessive compilation times and slowdown.
+And boom! Type stable again. We always recommend picking the chunksize via
+[`NonlinearSolve.pickchunksize`](@ref), however, if you manually specify the chunksize, it
+must be `≤ length of input`. However, a very large chunksize can lead to excessive
+compilation times and slowdown.
diff --git a/docs/src/solvers/fixed_point_solvers.md b/docs/src/solvers/fixed_point_solvers.md
@@ -27,7 +27,7 @@ Using [native NonlinearSolve.jl methods](@ref nonlinearsystemsolvers) is the rec
 approach. For systems where constructing Jacobian Matrices are expensive, we recommend
 using a Krylov Method with one of those solvers.
 
-## Full List of Methods
+## [Full List of Methods](@id fixed_point_methods_full_list)
 
 We are only listing the methods that natively solve fixed point problems.
 

diff --git a/docs/src/solvers/nonlinear_least_squares_solvers.md b/docs/src/solvers/nonlinear_least_squares_solvers.md
@@ -36,7 +36,7 @@ arrays.
   - `SimpleGaussNewton()`: Simple Gauss Newton implementation using QR factorizations for
     numerical stability (aliased to [`SimpleNewtonRaphson`](@ref)).
 
-### FastLevenbergMarquardt.jl
+### [FastLevenbergMarquardt.jl](@id fastlm_wrapper_summary)
 
 A wrapper over
 [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl). Note that
@@ -46,7 +46,7 @@ benchmarks demonstrate [`LevenbergMarquardt()`](@ref) usually outperforms.
   - [`FastLevenbergMarquardtJL(linsolve = :cholesky)`](@ref), can also choose
     `linsolve = :qr`.
 
-### LeastSquaresOptim.jl
+### [LeastSquaresOptim.jl](@id lso_wrapper_summary)
 
 A wrapper over
 [LeastSquaresOptim.jl](https://github.com/matthieugomez/LeastSquaresOptim.jl). Has a core

diff --git a/docs/src/solvers/nonlinear_system_solvers.md b/docs/src/solvers/nonlinear_system_solvers.md
@@ -32,6 +32,8 @@ solving of very large systems. Meanwhile, [`SimpleNewtonRaphson`](@ref) and
 [`SimpleTrustRegion`](@ref) are implementations which are specialized for small equations.
 They are non-allocating on static arrays and thus really well-optimized for small systems,
 thus usually outperforming the other methods when such types are used for `u0`.
+Additionally, these solvers can be used inside GPU kernels. See
+[PSOGPU.jl](https://github.com/SciML/PSOGPU.jl) for an example of this.
 
 ## Full List of Methods
 
@@ -134,8 +136,9 @@ Submethod choices for this algorithm include:
 
 ### MINPACK.jl
 
-MINPACK.jl methods are good for medium-sized nonlinear solves. It does not scale due to
-the lack of sparse Jacobian support, though the methods are very robust and stable.
+MINPACK.jl is a wrapper package for bringing the Fortran solvers from MINPACK. However, our
+benchmarks reveal that these methods are rarely competitive with our native solvers. Thus,
+our recommendation is to use these only for benchmarking and debugging purposes.
 
   - [`CMINPACK()`](@ref): A wrapper for using the classic MINPACK method through
     [MINPACK.jl](https://github.com/sglyon/MINPACK.jl)
@@ -161,3 +164,7 @@ SIAMFANLEquations.jl is a wrapper for the methods in the SIAMFANLEquations.jl li
 
   - [`SIAMFANLEquationsJL()`](@ref): A wrapper for using the methods in
     [SIAMFANLEquations.jl](https://github.com/ctkelley/SIAMFANLEquations.jl)
+
+Other solvers listed in [Fixed Point Solvers](@ref fixed_point_methods_full_list),
+[FastLevenbergMarquardt.jl](@ref fastlm_wrapper_summary) and
+[LeastSquaresOptim.jl](@ref lso_wrapper_summary) can also solve nonlinear systems.
diff --git a/docs/src/solvers/steady_state_solvers.md b/docs/src/solvers/steady_state_solvers.md
@@ -23,7 +23,7 @@ the direct [`SteadyStateProblem`](@ref) approach can give different answers (i.e
 correct unique fixed point) on ODEs with non-autonomous dynamics.
 
 If you have an unstable equilibrium and you want to solve for the unstable equilibrium,
-then [`DynamicSS`](@ref) might converge to the equilibrium based on the initial condition.
+then [`DynamicSS`](@ref) will not converge to that equilibrium for any initial condition.
 However, Nonlinear Solvers don't suffer from this issue, and thus it's recommended to
 use a nonlinear solver if you want to solve for the unstable equilibrium.
 

diff --git a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl
@@ -17,7 +17,7 @@ import StaticArraysCore: SArray
 end
 @inline _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, ::SArray) where {linsolve} = linsolve
 
-function SciMLBase.__solve(prob::NonlinearLeastSquaresProblem,
+function SciMLBase.__solve(prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem},
         alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = nothing,
         reltol = nothing, maxiters = 1000, termination_condition = nothing, kwargs...)
     NonlinearSolve.__test_termination_condition(termination_condition,

diff --git a/src/default.jl b/src/default.jl
@@ -189,7 +189,7 @@ function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve
         # Let's atleast have something here for complex numbers
         algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),)
     else
-        algs = (TrustRegion(; concrete_jac, linsolve, precs),
+        algs = (TrustRegion(; concrete_jac, linsolve, precs, autodiff),
             TrustRegion(; concrete_jac, linsolve, precs, autodiff,
                 radius_update_scheme = RadiusUpdateSchemes.Bastin),
             NewtonRaphson(; concrete_jac, linsolve, precs,

diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl
@@ -53,7 +53,6 @@ function get_concrete_forward_ad(autodiff::ADTypes.AbstractADType, prob,
 end
 function get_concrete_forward_ad(autodiff, prob, sp::Val{test_sparse} = True, args...;
         kwargs...) where {test_sparse}
-    # TODO: Default to PolyesterForwardDiff for non sparse problems
     if test_sparse
         (; sparsity, jac_prototype) = prob.f
         use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing
@@ -96,7 +95,6 @@ function get_concrete_reverse_ad(autodiff::ADTypes.AbstractADType, prob,
 end
 function get_concrete_reverse_ad(autodiff, prob, sp::Val{test_sparse} = True, args...;
         kwargs...) where {test_sparse}
-    # TODO: Default to Enzyme / ReverseDiff for inplace problems?
     if test_sparse
         (; sparsity, jac_prototype) = prob.f
         use_sparse_ad = sparsity !== nothing || jac_prototype !== nothing

diff --git a/src/utils.jl b/src/utils.jl
@@ -148,3 +148,12 @@ function Base.merge(s1::ImmutableNLStats, s2::ImmutableNLStats)
     return ImmutableNLStats(s1.nf + s2.nf, s1.njacs + s2.njacs, s1.nfactors + s2.nfactors,
         s1.nsolve + s2.nsolve, s1.nsteps + s2.nsteps)
 end
+
+"""
+    pickchunksize(x) = pickchunksize(length(x))
+    pickchunksize(x::Int)
+
+Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the input length.
+"""
+@inline pickchunksize(x) = pickchunksize(length(x))
+@inline pickchunksize(x::Int) = ForwardDiff.pickchunksize(x)