From 8762dccb3934fc1437e8bd3a17848e5a227411a4 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 31 Oct 2024 13:21:43 -0400 Subject: [PATCH 1/3] refactor!: move preconditioners inside linear solvers --- docs/src/native/solvers.md | 4 -- docs/src/release_notes.md | 2 + docs/src/tutorials/large_systems.md | 67 ++++++++----------- .../ext/NonlinearSolveBaseLinearSolveExt.jl | 25 ++----- .../src/descent/damped_newton.jl | 5 +- lib/NonlinearSolveBase/src/descent/dogleg.jl | 10 +-- lib/NonlinearSolveBase/src/descent/newton.jl | 7 +- .../src/descent/steepest.jl | 4 +- lib/NonlinearSolveBase/src/linear_solve.jl | 23 ++----- .../src/gauss_newton.jl | 6 +- .../src/levenberg_marquardt.jl | 5 +- .../src/pseudo_transient.jl | 6 +- lib/NonlinearSolveFirstOrder/src/raphson.jl | 6 +- .../src/trust_region.jl | 6 +- .../test/rootfind_tests.jl | 42 ++++++------ lib/NonlinearSolveQuasiNewton/src/klement.jl | 6 +- src/polyalg.jl | 26 +++---- 17 files changed, 103 insertions(+), 147 deletions(-) diff --git a/docs/src/native/solvers.md b/docs/src/native/solvers.md index a3aa900e0..01583dda2 100644 --- a/docs/src/native/solvers.md +++ b/docs/src/native/solvers.md @@ -18,10 +18,6 @@ documentation. uses the LinearSolve.jl default algorithm choice. For more information on available algorithm choices, see the [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - - `precs`: the choice of preconditioners for the linear solver. Defaults to using no - preconditioners. For more information on specifying preconditioners for LinearSolve - algorithms, consult the - [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/). - `linesearch`: the line search algorithm to use. Defaults to [`NoLineSearch()`](@extref LineSearch.NoLineSearch), which means that no line search is performed. diff --git a/docs/src/release_notes.md b/docs/src/release_notes.md index bb0b5dec9..0d4ce25e8 100644 --- a/docs/src/release_notes.md +++ b/docs/src/release_notes.md @@ -5,6 +5,8 @@ ### Breaking Changes in `NonlinearSolve.jl` v4 - `ApproximateJacobianSolveAlgorithm` has been renamed to `QuasiNewtonAlgorithm`. + - Preconditioners for the linear solver needs to be specified with the linear solver + instead of `precs` keyword argument. - See [common breaking changes](@ref common-breaking-changes-v4v2) below. ### Breaking Changes in `SimpleNonlinearSolve.jl` v2 diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index 0a0be59e7..ae8ed7acd 100644 --- a/docs/src/tutorials/large_systems.md +++ b/docs/src/tutorials/large_systems.md @@ -100,7 +100,8 @@ end u0 = init_brusselator_2d(xyd_brusselator) prob_brusselator_2d = NonlinearProblem( - brusselator_2d_loop, u0, p; abstol = 1e-10, reltol = 1e-10) + brusselator_2d_loop, u0, p; abstol = 1e-10, reltol = 1e-10 +) ``` ## Choosing Jacobian Types @@ -140,7 +141,8 @@ using SparseConnectivityTracer prob_brusselator_2d_autosparse = NonlinearProblem( NonlinearFunction(brusselator_2d_loop; sparsity = TracerSparsityDetector()), - u0, p; abstol = 1e-10, reltol = 1e-10) + u0, p; abstol = 1e-10, reltol = 1e-10 +) @btime solve(prob_brusselator_2d_autosparse, NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 12))); @@ -235,7 +237,7 @@ choices, see the Any [LinearSolve.jl-compatible preconditioner](https://docs.sciml.ai/LinearSolve/stable/basics/Preconditioners/) can be used as a preconditioner in the linear solver interface. To define preconditioners, -one must define a `precs` function in compatible with nonlinear solvers which returns the +one must define a `precs` function in compatible with linear solvers which returns the left and right preconditioners, matrices which approximate the inverse of `W = I - gamma*J` used in the solution of the ODE. An example of this with using [IncompleteLU.jl](https://github.com/haampie/IncompleteLU.jl) is as follows: @@ -244,26 +246,18 @@ used in the solution of the ODE. An example of this with using # FIXME: On 1.10+ this is broken. Skipping this for now. using IncompleteLU -function incompletelu(W, du, u, p, t, newW, Plprev, Prprev, solverdata) - if newW === nothing || newW - Pl = ilu(W, τ = 50.0) - else - Pl = Plprev - end - Pl, nothing -end +incompletelu(W, p = nothing) = ilu(W, τ = 50.0), LinearAlgebra.I @btime solve(prob_brusselator_2d_sparse, - NewtonRaphson(linsolve = KrylovJL_GMRES(), precs = incompletelu, concrete_jac = true)); + NewtonRaphson(linsolve = KrylovJL_GMRES(precs = incompletelu), concrete_jac = true) +); nothing # hide ``` Notice a few things about this preconditioner. This preconditioner uses the sparse Jacobian, and thus we set `concrete_jac = true` to tell the algorithm to generate the Jacobian -(otherwise, a Jacobian-free algorithm is used with GMRES by default). Then `newW = true` -whenever a new `W` matrix is computed, and `newW = nothing` during the startup phase of the -solver. Thus, we do a check `newW === nothing || newW` and when true, it's only at these -points when we update the preconditioner, otherwise we just pass on the previous version. +(otherwise, a Jacobian-free algorithm is used with GMRES by default). + We use `convert(AbstractMatrix,W)` to get the concrete `W` matrix (matching `jac_prototype`, thus `SpraseMatrixCSC`) which we can use in the preconditioner's definition. Then we use `IncompleteLU.ilu` on that sparse matrix to generate the preconditioner. We return @@ -279,39 +273,36 @@ which is more automatic. The setup is very similar to before: ```@example ill_conditioned_nlprob using AlgebraicMultigrid -function algebraicmultigrid(W, du, u, p, t, newW, Plprev, Prprev, solverdata) - if newW === nothing || newW - Pl = aspreconditioner(ruge_stuben(convert(AbstractMatrix, W))) - else - Pl = Plprev - end - Pl, nothing +function algebraicmultigrid(W, p = nothing) + return aspreconditioner(ruge_stuben(convert(AbstractMatrix, W))), LinearAlgebra.I 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 ``` or with a Jacobi smoother: ```@example ill_conditioned_nlprob -function algebraicmultigrid2(W, du, u, p, t, newW, Plprev, Prprev, solverdata) - if newW === nothing || newW - A = convert(AbstractMatrix, W) - Pl = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben( - A, presmoother = AlgebraicMultigrid.Jacobi(rand(size(A, 1))), - postsmoother = AlgebraicMultigrid.Jacobi(rand(size(A, 1))))) - else - Pl = Plprev - end - Pl, nothing +function algebraicmultigrid2(W, p = nothing) + A = convert(AbstractMatrix, W) + Pl = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben( + A, presmoother = AlgebraicMultigrid.Jacobi(rand(size(A, 1))), + postsmoother = AlgebraicMultigrid.Jacobi(rand(size(A, 1))) + )) + return Pl, LinearAlgebra.I end -@btime solve(prob_brusselator_2d_sparse, - NewtonRaphson(linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid2, - concrete_jac = true)); +@btime solve( + prob_brusselator_2d_sparse, + NewtonRaphson( + linsolve = KrylovJL_GMRES(precs = algebraicmultigrid2), concrete_jac = true + ) +); nothing # hide ``` diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl index d1829804d..4a1234b25 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl @@ -11,8 +11,8 @@ using LinearAlgebra: ColumnNorm using NonlinearSolveBase: NonlinearSolveBase, LinearSolveJLCache, LinearSolveResult, Utils function (cache::LinearSolveJLCache)(; - A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, - cachedata = nothing, reuse_A_if_factorization = false, verbose = true, kwargs... + A = nothing, b = nothing, linu = nothing, + reuse_A_if_factorization = false, verbose = true, kwargs... ) cache.stats.nsolve += 1 @@ -20,24 +20,6 @@ function (cache::LinearSolveJLCache)(; b !== nothing && setproperty!(cache.lincache, :b, b) linu !== nothing && NonlinearSolveBase.set_lincache_u!(cache, linu) - Plprev = cache.lincache.Pl - Prprev = cache.lincache.Pr - - if cache.precs === nothing - Pl, Pr = nothing, nothing - else - Pl, Pr = cache.precs( - cache.lincache.A, du, linu, p, nothing, - A !== nothing, Plprev, Prprev, cachedata - ) - end - - if Pl !== nothing || Pr !== nothing - Pl, Pr = NonlinearSolveBase.wrap_preconditioners(Pl, Pr, linu) - cache.lincache.Pl = Pl - cache.lincache.Pr = Pr - end - linres = solve!(cache.lincache) cache.lincache = linres.cache # Unfortunately LinearSolve.jl doesn't have the most uniform ReturnCode handling @@ -58,7 +40,8 @@ function (cache::LinearSolveJLCache)(; linprob = LinearProblem(A, b; u0 = linres.u) cache.additional_lincache = init( linprob, QRFactorization(ColumnNorm()); alias_u0 = false, - alias_A = false, alias_b = false, cache.lincache.Pl, cache.lincache.Pr) + alias_A = false, alias_b = false + ) else cache.additional_lincache.A = A cache.additional_lincache.b = b diff --git a/lib/NonlinearSolveBase/src/descent/damped_newton.jl b/lib/NonlinearSolveBase/src/descent/damped_newton.jl index 47ed56207..f6523e399 100644 --- a/lib/NonlinearSolveBase/src/descent/damped_newton.jl +++ b/lib/NonlinearSolveBase/src/descent/damped_newton.jl @@ -1,7 +1,5 @@ """ - DampedNewtonDescent(; - linsolve = nothing, precs = nothing, initial_damping, damping_fn - ) + DampedNewtonDescent(; linsolve = nothing, initial_damping, damping_fn) A Newton descent algorithm with damping. The damping factor is computed using the `damping_fn` function. The descent direction is computed as ``(JᵀJ + λDᵀD) δu = -fu``. For @@ -20,7 +18,6 @@ The damping factor returned must be a non-negative number. """ @kwdef @concrete struct DampedNewtonDescent <: AbstractDescentDirection linsolve = nothing - precs = nothing initial_damping damping_fn <: AbstractDampingFunction end diff --git a/lib/NonlinearSolveBase/src/descent/dogleg.jl b/lib/NonlinearSolveBase/src/descent/dogleg.jl index f446e9cf6..0eafbff74 100644 --- a/lib/NonlinearSolveBase/src/descent/dogleg.jl +++ b/lib/NonlinearSolveBase/src/descent/dogleg.jl @@ -1,5 +1,5 @@ """ - Dogleg(; linsolve = nothing, precs = nothing) + Dogleg(; linsolve = nothing) Switch between Newton's method and the steepest descent method depending on the size of the trust region. The trust region is specified via keyword argument `trust_region` to @@ -15,18 +15,18 @@ end supports_trust_region(::Dogleg) = true get_linear_solver(alg::Dogleg) = get_linear_solver(alg.newton_descent) -function Dogleg(; linsolve = nothing, precs = nothing, damping = Val(false), +function Dogleg(; linsolve = nothing, damping = Val(false), damping_fn = missing, initial_damping = missing, kwargs...) if !Utils.unwrap_val(damping) - return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent(; linsolve, precs)) + return Dogleg(NewtonDescent(; linsolve), SteepestDescent(; linsolve)) end if damping_fn === missing || initial_damping === missing throw(ArgumentError("`damping_fn` and `initial_damping` must be supplied if \ `damping = Val(true)`.")) end return Dogleg( - DampedNewtonDescent(; linsolve, precs, damping_fn, initial_damping), - SteepestDescent(; linsolve, precs) + DampedNewtonDescent(; linsolve, damping_fn, initial_damping), + SteepestDescent(; linsolve) ) end diff --git a/lib/NonlinearSolveBase/src/descent/newton.jl b/lib/NonlinearSolveBase/src/descent/newton.jl index 810661507..a0aaa91d3 100644 --- a/lib/NonlinearSolveBase/src/descent/newton.jl +++ b/lib/NonlinearSolveBase/src/descent/newton.jl @@ -1,5 +1,5 @@ """ - NewtonDescent(; linsolve = nothing, precs = nothing) + NewtonDescent(; linsolve = nothing) Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is commonly referred to as the Gauss-Newton Descent. @@ -8,7 +8,6 @@ See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@r """ @kwdef @concrete struct NewtonDescent <: AbstractDescentDirection linsolve = nothing - precs = nothing end supports_line_search(::NewtonDescent) = true @@ -103,7 +102,7 @@ function InternalAPI.solve!( @static_timeit cache.timer "linear solve" begin linres = cache.lincache(; A = Utils.maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, - kwargs..., linu = Utils.safe_vec(δu), du = Utils.safe_vec(δu), + kwargs..., linu = Utils.safe_vec(δu), reuse_A_if_factorization = !new_jacobian || (idx !== Val(1)) ) end @@ -111,7 +110,7 @@ function InternalAPI.solve!( @static_timeit cache.timer "linear solve" begin linres = cache.lincache(; A = J, b = Utils.safe_vec(fu), - kwargs..., linu = Utils.safe_vec(δu), du = Utils.safe_vec(δu), + kwargs..., linu = Utils.safe_vec(δu), reuse_A_if_factorization = !new_jacobian || idx !== Val(1) ) end diff --git a/lib/NonlinearSolveBase/src/descent/steepest.jl b/lib/NonlinearSolveBase/src/descent/steepest.jl index deb1a9817..247490957 100644 --- a/lib/NonlinearSolveBase/src/descent/steepest.jl +++ b/lib/NonlinearSolveBase/src/descent/steepest.jl @@ -1,5 +1,5 @@ """ - SteepestDescent(; linsolve = nothing, precs = nothing) + SteepestDescent(; linsolve = nothing) Compute the descent direction as ``δu = -Jᵀfu``. The linear solver and preconditioner are only used if `J` is provided in the inverted form. @@ -8,7 +8,6 @@ See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref """ @kwdef @concrete struct SteepestDescent <: AbstractDescentDirection linsolve = nothing - precs = nothing end supports_line_search(::SteepestDescent) = true @@ -57,7 +56,6 @@ function InternalAPI.solve!( A = J === nothing ? nothing : transpose(J) linres = cache.lincache(; A, b = Utils.safe_vec(fu), kwargs..., linu = Utils.safe_vec(δu), - du = Utils.safe_vec(δu), reuse_A_if_factorization = !new_jacobian || idx !== Val(1) ) δu = Utils.restructure(SciMLBase.get_du(cache, idx), linres.u) diff --git a/lib/NonlinearSolveBase/src/linear_solve.jl b/lib/NonlinearSolveBase/src/linear_solve.jl index f3a2338c9..afe9c52d6 100644 --- a/lib/NonlinearSolveBase/src/linear_solve.jl +++ b/lib/NonlinearSolveBase/src/linear_solve.jl @@ -7,7 +7,6 @@ end lincache linsolve additional_lincache::Any - precs stats::NLStats end @@ -34,8 +33,8 @@ handled: ```julia (cache::LinearSolverCache)(; - A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, - weight = nothing, cachedata = nothing, reuse_A_if_factorization = false, kwargs...) + A = nothing, b = nothing, linu = nothing, reuse_A_if_factorization = false, kwargs... +) ``` Returns the solution of the system `u` and stores the updated cache in `cache.lincache`. @@ -60,15 +59,11 @@ aliasing arguments even after cache construction, i.e., if we passed in an `A` t not mutated, we do this by copying over `A` to a preconstructed cache. """ function construct_linear_solver(alg, linsolve, A, b, u; stats, kwargs...) - no_preconditioner = !hasfield(typeof(alg), :precs) || alg.precs === nothing - if (A isa Number && b isa Number) || (A isa Diagonal) return NativeJLLinearSolveCache(A, b, stats) elseif linsolve isa typeof(\) - !no_preconditioner && - error("Default Julia Backsolve Operator `\\` doesn't support Preconditioners") return NativeJLLinearSolveCache(A, b, stats) - elseif no_preconditioner && linsolve === nothing + elseif linsolve === nothing if (A isa SMatrix || A isa WrappedArray{<:Any, <:SMatrix}) return NativeJLLinearSolveCache(A, b, stats) end @@ -78,17 +73,9 @@ function construct_linear_solver(alg, linsolve, A, b, u; stats, kwargs...) @bb u_cache = copy(u_fixed) linprob = LinearProblem(A, b; u0 = u_cache, kwargs...) - if no_preconditioner - precs, Pl, Pr = nothing, nothing, nothing - else - precs = alg.precs - Pl, Pr = precs(A, nothing, u, ntuple(Returns(nothing), 6)...) - end - Pl, Pr = wrap_preconditioners(Pl, Pr, u) - # unlias here, we will later use these as caches - lincache = init(linprob, linsolve; alias_A = false, alias_b = false, Pl, Pr) - return LinearSolveJLCache(lincache, linsolve, nothing, precs, stats) + lincache = init(linprob, linsolve; alias_A = false, alias_b = false) + return LinearSolveJLCache(lincache, linsolve, nothing, stats) end function (cache::NativeJLLinearSolveCache)(; diff --git a/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl b/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl index 59b90b76e..081f1af27 100644 --- a/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl +++ b/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl @@ -1,6 +1,6 @@ """ GaussNewton(; - concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing, + concrete_jac = nothing, linsolve = nothing, linesearch = missing, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) @@ -9,12 +9,12 @@ matrices via colored automatic differentiation and preconditioned linear solvers for large-scale and numerically-difficult nonlinear systems. """ function GaussNewton(; - concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing, + concrete_jac = nothing, linsolve = nothing, linesearch = missing, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) return GeneralizedFirstOrderAlgorithm(; linesearch, - descent = NewtonDescent(; linsolve, precs), + descent = NewtonDescent(; linsolve), autodiff, vjp_autodiff, jvp_autodiff, concrete_jac, name = :GaussNewton diff --git a/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl b/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl index f9c11ae28..66906fe9e 100644 --- a/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl +++ b/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl @@ -1,6 +1,6 @@ """ LevenbergMarquardt(; - linsolve = nothing, precs = nothing, + linsolve = nothing, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, disable_geodesic = Val(false), damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, min_damping_D::Real = 1e-8, @@ -33,7 +33,7 @@ For the remaining arguments, see [`GeodesicAcceleration`](@ref) and [`NonlinearSolveFirstOrder.LevenbergMarquardtTrustRegion`](@ref) documentations. """ function LevenbergMarquardt(; - linsolve = nothing, precs = nothing, + linsolve = nothing, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, disable_geodesic = Val(false), damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, min_damping_D::Real = 1e-8, @@ -41,7 +41,6 @@ function LevenbergMarquardt(; ) descent = DampedNewtonDescent(; linsolve, - precs, initial_damping = damping_initial, damping_fn = LevenbergMarquardtDampingFunction( damping_increase_factor, damping_decrease_factor, min_damping_D diff --git a/lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl b/lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl index a985cc100..0c1984f85 100644 --- a/lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl +++ b/lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl @@ -1,7 +1,7 @@ """ PseudoTransient(; concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3, - linsolve = nothing, precs = nothing, + linsolve = nothing, autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing ) @@ -19,13 +19,13 @@ This implementation specifically uses "switched evolution relaxation" """ function PseudoTransient(; concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3, - linsolve = nothing, precs = nothing, + linsolve = nothing, autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing ) return GeneralizedFirstOrderAlgorithm(; linesearch, descent = DampedNewtonDescent(; - linsolve, precs, initial_damping = alpha_initial, + linsolve, initial_damping = alpha_initial, damping_fn = SwitchedEvolutionRelaxation() ), autodiff, diff --git a/lib/NonlinearSolveFirstOrder/src/raphson.jl b/lib/NonlinearSolveFirstOrder/src/raphson.jl index 5c14f10b7..d482f4846 100644 --- a/lib/NonlinearSolveFirstOrder/src/raphson.jl +++ b/lib/NonlinearSolveFirstOrder/src/raphson.jl @@ -1,6 +1,6 @@ """ NewtonRaphson(; - concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing, + concrete_jac = nothing, linsolve = nothing, linesearch = missing, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) @@ -9,12 +9,12 @@ matrices via colored automatic differentiation and preconditioned linear solvers for large-scale and numerically-difficult nonlinear systems. """ function NewtonRaphson(; - concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing, + concrete_jac = nothing, linsolve = nothing, linesearch = missing, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) return GeneralizedFirstOrderAlgorithm(; linesearch, - descent = NewtonDescent(; linsolve, precs), + descent = NewtonDescent(; linsolve), autodiff, vjp_autodiff, jvp_autodiff, concrete_jac, name = :NewtonRaphson diff --git a/lib/NonlinearSolveFirstOrder/src/trust_region.jl b/lib/NonlinearSolveFirstOrder/src/trust_region.jl index 92d2655ac..03a1b3f9a 100644 --- a/lib/NonlinearSolveFirstOrder/src/trust_region.jl +++ b/lib/NonlinearSolveFirstOrder/src/trust_region.jl @@ -1,6 +1,6 @@ """ TrustRegion(; - concrete_jac = nothing, linsolve = nothing, precs = nothing, + concrete_jac = nothing, linsolve = nothing, radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, @@ -23,7 +23,7 @@ For the remaining arguments, see [`NonlinearSolveFirstOrder.GenericTrustRegionSc documentation. """ function TrustRegion(; - concrete_jac = nothing, linsolve = nothing, precs = nothing, + concrete_jac = nothing, linsolve = nothing, radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, @@ -31,7 +31,7 @@ function TrustRegion(; max_shrink_times::Int = 32, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) - descent = Dogleg(; linsolve, precs) + descent = Dogleg(; linsolve) trustregion = GenericTrustRegionScheme(; method = radius_update_scheme, step_threshold, shrink_threshold, expand_threshold, shrink_factor, expand_factor, initial_trust_radius, max_trust_radius diff --git a/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl b/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl index f554e02d6..b352febaf 100644 --- a/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl +++ b/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl @@ -13,11 +13,6 @@ end u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) - preconditioners = [ - u0 -> nothing, - u0 -> ((args...) -> (Diagonal(rand!(similar(u0))), nothing)) - ] - @testset for ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff(), AutoEnzyme()) @testset "$(nameof(typeof(linesearch)))" for linesearch in ( LineSearchesJL(; method = LineSearches.Static(), autodiff = ad), @@ -43,13 +38,17 @@ end @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) ad isa AutoZygote && continue - @testset for preconditioner in preconditioners, - linsolve in (nothing, KrylovJL_GMRES(), \) - - precs = preconditioner(u0) - typeof(linsolve) <: typeof(\) && precs !== nothing && continue - - solver = NewtonRaphson(; linsolve, precs, linesearch, autodiff = ad) + @testset for linsolve in ( + nothing, + KrylovJL_GMRES(; precs = nothing), + KrylovJL_GMRES(; + precs = (A, p = nothing) -> ( + Diagonal(randn!(similar(A, size(A, 1)))), LinearAlgebra.I + ) + ), + \ + ) + solver = NewtonRaphson(; linsolve, linesearch, autodiff = ad) sol = solve_iip(quadratic_f!, u0; solver) @test SciMLBase.successful_retcode(sol) @@ -115,14 +114,17 @@ end @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) ad isa AutoZygote && continue - @testset for preconditioner in preconditioners, - linsolve in (nothing, KrylovJL_GMRES(), \) - - precs = preconditioner(u0) - typeof(linsolve) <: typeof(\) && precs !== nothing && continue - - solver = PseudoTransient(; - alpha_initial = 10.0, linsolve, precs, autodiff = ad) + @testset for linsolve in ( + nothing, + KrylovJL_GMRES(; precs = nothing), + KrylovJL_GMRES(; + precs = (A, p = nothing) -> ( + Diagonal(randn!(similar(A, size(A, 1)))), LinearAlgebra.I + ) + ), + \ + ) + solver = PseudoTransient(; alpha_initial = 10.0, linsolve, autodiff = ad) sol = solve_iip(quadratic_f!, u0; solver) @test SciMLBase.successful_retcode(sol) err = maximum(abs, quadratic_f(sol.u, 2.0)) diff --git a/lib/NonlinearSolveQuasiNewton/src/klement.jl b/lib/NonlinearSolveQuasiNewton/src/klement.jl index 4bd09f45e..34ad7f6d5 100644 --- a/lib/NonlinearSolveQuasiNewton/src/klement.jl +++ b/lib/NonlinearSolveQuasiNewton/src/klement.jl @@ -1,7 +1,7 @@ """ Klement(; max_resets = 100, linsolve = nothing, linesearch = nothing, - precs = nothing, alpha = nothing, init_jacobian::Val = Val(:identity), + alpha = nothing, init_jacobian::Val = Val(:identity), autodiff = nothing ) @@ -28,14 +28,14 @@ over this. """ function Klement(; max_resets = 100, linsolve = nothing, linesearch = nothing, - precs = nothing, alpha = nothing, init_jacobian::Val = Val(:identity), + alpha = nothing, init_jacobian::Val = Val(:identity), autodiff = nothing ) concrete_jac = Val(init_jacobian isa Val{:true_jacobian} || init_jacobian isa Val{:true_jacobian_diagonal}) return QuasiNewtonAlgorithm(; linesearch, - descent = NewtonDescent(; linsolve, precs), + descent = NewtonDescent(; linsolve), update_rule = KlementUpdateRule(), reinit_rule = IllConditionedJacobianReset(), max_resets, diff --git a/src/polyalg.jl b/src/polyalg.jl index 5cf21d6e1..f16dd7e0e 100644 --- a/src/polyalg.jl +++ b/src/polyalg.jl @@ -322,7 +322,7 @@ end RobustMultiNewton( ::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = nothing, + linsolve = nothing, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) @@ -342,12 +342,14 @@ or more precision / more stable linear solver choice is required). function RobustMultiNewton( ::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = nothing, + linsolve = nothing, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) where {T} - common_kwargs = (; concrete_jac, linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + common_kwargs = (; concrete_jac, linsolve, autodiff, vjp_autodiff, jvp_autodiff) if T <: Complex # Let's atleast have something here for complex numbers - algs = (NewtonRaphson(; common_kwargs...),) + algs = ( + NewtonRaphson(; common_kwargs...), + ) else algs = ( TrustRegion(; common_kwargs...), @@ -365,7 +367,7 @@ end FastShortcutNonlinearPolyalg( ::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = nothing, + linsolve = nothing, must_use_jacobian::Val = Val(false), prefer_simplenonlinearsolve::Val = Val(false), autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, @@ -389,14 +391,14 @@ for more performance and then tries more robust techniques if the faster ones fa function FastShortcutNonlinearPolyalg( ::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = nothing, + linsolve = nothing, must_use_jacobian::Val = Val(false), prefer_simplenonlinearsolve::Val = Val(false), autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, u0_len::Union{Int, Nothing} = nothing ) where {T} start_index = 1 - common_kwargs = (; concrete_jac, linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + common_kwargs = (; concrete_jac, linsolve, autodiff, vjp_autodiff, jvp_autodiff) if must_use_jacobian isa Val{true} if T <: Complex algs = (NewtonRaphson(; common_kwargs...),) @@ -436,7 +438,7 @@ function FastShortcutNonlinearPolyalg( algs = ( Broyden(; autodiff), Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - Klement(; linsolve, precs, autodiff), + Klement(; linsolve, autodiff), NewtonRaphson(; common_kwargs...) ) else @@ -445,7 +447,7 @@ function FastShortcutNonlinearPolyalg( algs = ( Broyden(; autodiff), Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - Klement(; linsolve, precs, autodiff), + Klement(; linsolve, autodiff), NewtonRaphson(; common_kwargs...), NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), TrustRegion(; common_kwargs...), @@ -461,7 +463,7 @@ end FastShortcutNLLSPolyalg( ::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = nothing, + linsolve = nothing, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) @@ -476,10 +478,10 @@ for more performance and then tries more robust techniques if the faster ones fa function FastShortcutNLLSPolyalg( ::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = nothing, + linsolve = nothing, autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing ) where {T} - common_kwargs = (; linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + common_kwargs = (; linsolve, autodiff, vjp_autodiff, jvp_autodiff) if T <: Complex algs = ( GaussNewton(; common_kwargs..., concrete_jac), From 71e9806ab027a3efb16b82628a34d54aef255157 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 31 Oct 2024 15:27:53 -0400 Subject: [PATCH 2/3] test: precs need concrete jacobian --- docs/src/tutorials/large_systems.md | 2 +- .../test/rootfind_tests.jl | 42 ++++++++++++------- 2 files changed, 27 insertions(+), 17 deletions(-) diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index ae8ed7acd..434acb185 100644 --- a/docs/src/tutorials/large_systems.md +++ b/docs/src/tutorials/large_systems.md @@ -274,7 +274,7 @@ which is more automatic. The setup is very similar to before: using AlgebraicMultigrid function algebraicmultigrid(W, p = nothing) - return aspreconditioner(ruge_stuben(convert(AbstractMatrix, W))), LinearAlgebra.I + return aspreconditioner(ruge_stuben(convert(AbstractMatrix, W))), LinearAlgebra.I end @btime solve(prob_brusselator_2d_sparse, diff --git a/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl b/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl index b352febaf..3aa61b256 100644 --- a/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl +++ b/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl @@ -38,17 +38,22 @@ end @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) ad isa AutoZygote && continue - @testset for linsolve in ( - nothing, - KrylovJL_GMRES(; precs = nothing), - KrylovJL_GMRES(; - precs = (A, p = nothing) -> ( - Diagonal(randn!(similar(A, size(A, 1)))), LinearAlgebra.I + @testset for (concrete_jac, linsolve) in ( + (Val(false), nothing), + (Val(false), KrylovJL_GMRES(; precs = nothing)), + ( + Val(true), + KrylovJL_GMRES(; + precs = (A, p = nothing) -> ( + Diagonal(randn!(similar(A, size(A, 1)))), LinearAlgebra.I + ) ) ), - \ + (Val(false), \) ) - solver = NewtonRaphson(; linsolve, linesearch, autodiff = ad) + solver = NewtonRaphson(; + linsolve, linesearch, autodiff = ad, concrete_jac + ) sol = solve_iip(quadratic_f!, u0; solver) @test SciMLBase.successful_retcode(sol) @@ -114,17 +119,22 @@ end @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) ad isa AutoZygote && continue - @testset for linsolve in ( - nothing, - KrylovJL_GMRES(; precs = nothing), - KrylovJL_GMRES(; - precs = (A, p = nothing) -> ( - Diagonal(randn!(similar(A, size(A, 1)))), LinearAlgebra.I + @testset for (concrete_jac, linsolve) in ( + (Val(false), nothing), + (Val(false), KrylovJL_GMRES(; precs = nothing)), + ( + Val(true), + KrylovJL_GMRES(; + precs = (A, p = nothing) -> ( + Diagonal(randn!(similar(A, size(A, 1)))), LinearAlgebra.I + ) ) ), - \ + (Val(false), \) ) - solver = PseudoTransient(; alpha_initial = 10.0, linsolve, autodiff = ad) + solver = PseudoTransient(; + alpha_initial = 10.0, linsolve, autodiff = ad, concrete_jac + ) sol = solve_iip(quadratic_f!, u0; solver) @test SciMLBase.successful_retcode(sol) err = maximum(abs, quadratic_f(sol.u, 2.0)) From 86495eb2a5809db2af4fd4a3e4154754057050e8 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 31 Oct 2024 15:53:39 -0400 Subject: [PATCH 3/3] docs: fix references --- lib/NonlinearSolveBase/src/jacobian.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/lib/NonlinearSolveBase/src/jacobian.jl b/lib/NonlinearSolveBase/src/jacobian.jl index 6ab31efd2..363b7b77a 100644 --- a/lib/NonlinearSolveBase/src/jacobian.jl +++ b/lib/NonlinearSolveBase/src/jacobian.jl @@ -29,8 +29,7 @@ Construct a cache for the Jacobian of `f` w.r.t. `u`. - `jvp_autodiff`: Automatic Differentiation or Finite Differencing backend for computing the Jacobian-vector product. - `linsolve`: Linear Solver Algorithm used to determine if we need a concrete jacobian - or if possible we can just use a [`SciMLJacobianOperators.JacobianOperator`](@ref) - instead. + or if possible we can just use a `JacobianOperator` instead. """ function construct_jacobian_cache( prob, alg, f::NonlinearFunction, fu, u = prob.u0, p = prob.p; stats,