Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use LinearSolve precs setup #462

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions docs/src/native/solvers.md
Original file line number Diff line number Diff line change
Expand Up @@ -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()`](@ref),
which means that no line search is performed. Algorithms from
[`LineSearches.jl`](https://github.com/JuliaNLSolvers/LineSearches.jl/) must be
Expand Down
46 changes: 15 additions & 31 deletions docs/src/tutorials/large_systems.md
Original file line number Diff line number Diff line change
Expand Up @@ -221,35 +221,27 @@ 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 solver 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:

```julia
# 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
function incompletelu(W, p=nothing)
Pl = ilu(W, τ = 50.0)
Pl, nothing
end

@btime solve(prob_brusselator_2d_sparse,
NewtonRaphson(linsolve = KrylovJL_GMRES(), precs = incompletelu, concrete_jac = true));
nothing # hide
NewtonRaphson(linsolve = KrylovJL_GMRES(precs = incompletelu), concrete_jac = true));
nothing
```

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
Expand All @@ -265,39 +257,31 @@ 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
function algebraicmultigrid(W, p=nothing)
Pl = aspreconditioner(ruge_stuben(convert(AbstractMatrix, W)))
Pl, nothing
end

@btime solve(prob_brusselator_2d_sparse,
NewtonRaphson(
linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid, concrete_jac = true));
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
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)))))
Pl, nothing
end

@btime solve(prob_brusselator_2d_sparse,
NewtonRaphson(
linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid2, concrete_jac = true));
linsolve = KrylovJL_GMRES(precs = algebraicmultigrid2), concrete_jac = true));
nothing # hide
```

Expand Down
6 changes: 3 additions & 3 deletions src/algorithms/gauss_newton.jl
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
"""
GaussNewton(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(),
precs = DEFAULT_PRECS, adkwargs...)
adkwargs...)

An advanced GaussNewton implementation with support for efficient handling of sparse
matrices via colored automatic differentiation and preconditioned linear solvers. Designed
for large-scale and numerically-difficult nonlinear least squares problems.
"""
function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS,
function GaussNewton(; concrete_jac = nothing, linsolve = nothing,
linesearch = NoLineSearch(), vjp_autodiff = nothing, autodiff = nothing)
descent = NewtonDescent(; linsolve, precs)
descent = NewtonDescent(; linsolve)
return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = :GaussNewton, descent,
jacobian_ad = autodiff, reverse_ad = vjp_autodiff, linesearch)
end
6 changes: 3 additions & 3 deletions src/algorithms/klement.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""
Klement(; max_resets = 100, linsolve = NoLineSearch(), linesearch = nothing,
precs = DEFAULT_PRECS, alpha = nothing, init_jacobian::Val = Val(:identity),
alpha = nothing, init_jacobian::Val = Val(:identity),
autodiff = nothing)

An implementation of `Klement` [klement2014using](@citep) with line search, preconditioning
Expand All @@ -25,7 +25,7 @@ over this.
differentiable problems.
"""
function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = nothing,
linesearch = NoLineSearch(), precs = DEFAULT_PRECS,
linesearch = NoLineSearch(),
autodiff = nothing, init_jacobian::Val{IJ} = Val(:identity)) where {IJ}
if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
Base.depwarn(
Expand All @@ -48,7 +48,7 @@ function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = nothing,
CJ = IJ === :true_jacobian || IJ === :true_jacobian_diagonal

return ApproximateJacobianSolveAlgorithm{CJ, :Klement}(;
linesearch, descent = NewtonDescent(; linsolve, precs),
linesearch, descent = NewtonDescent(; linsolve),
update_rule = KlementUpdateRule(),
reinit_rule = IllConditionedJacobianReset(), max_resets, initialization)
end
Expand Down
5 changes: 2 additions & 3 deletions src/algorithms/levenberg_marquardt.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""
LevenbergMarquardt(; linsolve = nothing,
precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75,
damping_initial::Real = 1.0, α_geodesic::Real = 0.75,
damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0,
finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, autodiff = nothing,
min_damping_D::Real = 1e-8, disable_geodesic = Val(false))
Expand Down Expand Up @@ -31,7 +31,7 @@ For the remaining arguments, see [`GeodesicAcceleration`](@ref) and
[`NonlinearSolve.LevenbergMarquardtTrustRegion`](@ref) documentations.
"""
function LevenbergMarquardt(;
concrete_jac = missing, linsolve = nothing, precs = DEFAULT_PRECS,
concrete_jac = missing, linsolve = nothing,
damping_initial::Real = 1.0, α_geodesic::Real = 0.75,
damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0,
finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0,
Expand All @@ -45,7 +45,6 @@ function LevenbergMarquardt(;
end

descent = DampedNewtonDescent(; linsolve,
precs,
initial_damping = damping_initial,
damping_fn = LevenbergMarquardtDampingFunction(
damping_increase_factor, damping_decrease_factor, min_damping_D))
Expand Down
6 changes: 3 additions & 3 deletions src/algorithms/pseudo_transient.jl
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""
PseudoTransient(; concrete_jac = nothing, linsolve = nothing,
linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(),
precs = DEFAULT_PRECS, autodiff = nothing)
autodiff = nothing)

An implementation of PseudoTransient Method [coffey2003pseudotransient](@cite) that is used
to solve steady state problems in an accelerated manner. It uses an adaptive time-stepping
Expand All @@ -17,8 +17,8 @@ This implementation specifically uses "switched evolution relaxation"
"""
function PseudoTransient(; concrete_jac = nothing, linsolve = nothing,
linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(),
precs = DEFAULT_PRECS, autodiff = nothing, alpha_initial = 1e-3)
descent = DampedNewtonDescent(; linsolve, precs, initial_damping = alpha_initial,
autodiff = nothing, alpha_initial = 1e-3)
descent = DampedNewtonDescent(; linsolve, initial_damping = alpha_initial,
damping_fn = SwitchedEvolutionRelaxation())
return GeneralizedFirstOrderAlgorithm(;
concrete_jac, name = :PseudoTransient, linesearch, descent, jacobian_ad = autodiff)
Expand Down
6 changes: 3 additions & 3 deletions src/algorithms/raphson.jl
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
"""
NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = NoLineSearch(),
precs = DEFAULT_PRECS, autodiff = nothing)
autodiff = nothing)

An advanced NewtonRaphson implementation with support for efficient handling of sparse
matrices via colored automatic differentiation and preconditioned linear solvers. Designed
for large-scale and numerically-difficult nonlinear systems.
"""
function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing,
linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing)
descent = NewtonDescent(; linsolve, precs)
linesearch = NoLineSearch(), autodiff = nothing)
descent = NewtonDescent(; linsolve)
return GeneralizedFirstOrderAlgorithm(;
concrete_jac, name = :NewtonRaphson, linesearch, descent, jacobian_ad = autodiff)
end
6 changes: 3 additions & 3 deletions src/algorithms/trust_region.jl
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
"""
TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS,
TrustRegion(; 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,
Expand All @@ -19,13 +19,13 @@ for large-scale and numerically-difficult nonlinear systems.
For the remaining arguments, see [`NonlinearSolve.GenericTrustRegionScheme`](@ref)
documentation.
"""
function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS,
function TrustRegion(; 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,
shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1,
max_shrink_times::Int = 32, autodiff = nothing, vjp_autodiff = nothing)
descent = Dogleg(; linsolve, precs)
descent = Dogleg(; linsolve)
if autodiff !== nothing && ADTypes.mode(autodiff) isa ADTypes.ForwardMode
forward_ad = autodiff
else
Expand Down
Loading
Loading