Ariadne.jl

newton_krylov!(F!, u₀::AbstractArray, p = nothing, M::Int = length(u₀); kwargs...)

Takes an in-place residual function F!(res, u, p).

Arguments

• F!: F!(res, u, p) solves res = F(u) = 0
• u₀: Initial guess
• p: Parameters
• M: Length of the output of F!. Defaults to length(u₀)

Keyword Arguments

• tol_rel: Relative tolerance. Newton stops when ‖F(u)‖ <= tol_rel * ‖F(u₀)‖ + tol_abs.
• tol_abs: Absolute tolerance
• max_niter: Maximum number of Newton iterations
• forcing: Forcing-term strategy for the inexact Newton method, a Forcing such as EisenstatWalker() (default) or Fixed(η). It sets atol = 0 and rtol = η for each Krylov solve; krylov_kwargs can override both. If nothing, Ariadne does not set atol or rtol, so the Krylov solver uses its own default tolerances, or those given in krylov_kwargs.
• linesearch!: Line search strategy, an instance of a subtype of AbstractLineSearch.
• verbose::Int: Verbosity level. If verbose > 0, progress is logged with @info.
• M::Union{Nothing, Function}: If provided, M(ws.J) is passed as a keyword argument to the Krylov solver.
• N::Union{Nothing, Function}: If provided, N(ws.J) is passed as a keyword argument to the Krylov solver.- krylov_kwargs: Keyword arguments passed to the Krylov solver, e.g., (; restart = true, itmax = 100). The keys that configure the Krylov workspace (memory and window, e.g., memory = 50 for the restart length of GMRES(k)) only take effect when a NewtonKrylovWorkspace is created, i.e., in the methods that take a residual function, and are ignored by newton_krylov!(ws).
• algo, assume_p_const: Only for the methods that take a residual function; see NewtonKrylovWorkspace.
• on_krylov_failure: What to do if the Krylov solver does not reach its tolerance: :stop (default) returns with status :krylov_failed without updating u, :continue takes the Newton step anyway (and counts the failure in stats.krylov_failures).
• callback: A function called once for the initial guess and then after each Newton iteration, with signature callback(u, res, norm_res).

Return value

(u, (; solved, status, stats, t)), where status is

• :converged: the residual norm satisfies the tolerance (solved == true)
• :max_iterations: max_niter Newton iterations were taken without convergence
• :krylov_failed: the Krylov solver failed and on_krylov_failure == :stop
• :nonfinite: the residual norm became NaN or Inf

and stats are the Ariadne.Stats.

newton_krylov!(F!, u, p, res; kwargs...)

Takes an in-place residual function F!(res, u, p).

Arguments

• F!: F!(res, u, p) solves res = F(u) = 0
• u: Initial guess (modified in-place)
• p: Parameters
• res: Temporary for residual

Keyword Arguments

• tol_rel: Relative tolerance. Newton stops when ‖F(u)‖ <= tol_rel * ‖F(u₀)‖ + tol_abs.
• tol_abs: Absolute tolerance
• max_niter: Maximum number of Newton iterations
• forcing: Forcing-term strategy for the inexact Newton method, a Forcing such as EisenstatWalker() (default) or Fixed(η). It sets atol = 0 and rtol = η for each Krylov solve; krylov_kwargs can override both. If nothing, Ariadne does not set atol or rtol, so the Krylov solver uses its own default tolerances, or those given in krylov_kwargs.
• linesearch!: Line search strategy, an instance of a subtype of AbstractLineSearch.
• verbose::Int: Verbosity level. If verbose > 0, progress is logged with @info.
• M::Union{Nothing, Function}: If provided, M(ws.J) is passed as a keyword argument to the Krylov solver.
• N::Union{Nothing, Function}: If provided, N(ws.J) is passed as a keyword argument to the Krylov solver.- krylov_kwargs: Keyword arguments passed to the Krylov solver, e.g., (; restart = true, itmax = 100). The keys that configure the Krylov workspace (memory and window, e.g., memory = 50 for the restart length of GMRES(k)) only take effect when a NewtonKrylovWorkspace is created, i.e., in the methods that take a residual function, and are ignored by newton_krylov!(ws).
• algo, assume_p_const: Only for the methods that take a residual function; see NewtonKrylovWorkspace.
• on_krylov_failure: What to do if the Krylov solver does not reach its tolerance: :stop (default) returns with status :krylov_failed without updating u, :continue takes the Newton step anyway (and counts the failure in stats.krylov_failures).
• callback: A function called once for the initial guess and then after each Newton iteration, with signature callback(u, res, norm_res).

Return value

(u, (; solved, status, stats, t)), where status is

• :converged: the residual norm satisfies the tolerance (solved == true)
• :max_iterations: max_niter Newton iterations were taken without convergence
• :krylov_failed: the Krylov solver failed and on_krylov_failure == :stop
• :nonfinite: the residual norm became NaN or Inf

and stats are the Ariadne.Stats.

newton_krylov!(ws::NewtonKrylovWorkspace, u; kwargs...)

Updates ws.u with the initial guess u and then calls newton_krylov!(ws; kwargs...).

Arguments

• ws: Pre-allocated NewtonKrylovWorkspace
• u: Initial guess (must have the same shape as ws.u)

Keyword Arguments

• tol_rel: Relative tolerance. Newton stops when ‖F(u)‖ <= tol_rel * ‖F(u₀)‖ + tol_abs.
• tol_abs: Absolute tolerance
• max_niter: Maximum number of Newton iterations
• forcing: Forcing-term strategy for the inexact Newton method, a Forcing such as EisenstatWalker() (default) or Fixed(η). It sets atol = 0 and rtol = η for each Krylov solve; krylov_kwargs can override both. If nothing, Ariadne does not set atol or rtol, so the Krylov solver uses its own default tolerances, or those given in krylov_kwargs.
• linesearch!: Line search strategy, an instance of a subtype of AbstractLineSearch.
• verbose::Int: Verbosity level. If verbose > 0, progress is logged with @info.
• M::Union{Nothing, Function}: If provided, M(ws.J) is passed as a keyword argument to the Krylov solver.
• N::Union{Nothing, Function}: If provided, N(ws.J) is passed as a keyword argument to the Krylov solver.- krylov_kwargs: Keyword arguments passed to the Krylov solver, e.g., (; restart = true, itmax = 100). The keys that configure the Krylov workspace (memory and window, e.g., memory = 50 for the restart length of GMRES(k)) only take effect when a NewtonKrylovWorkspace is created, i.e., in the methods that take a residual function, and are ignored by newton_krylov!(ws).
• algo, assume_p_const: Only for the methods that take a residual function; see NewtonKrylovWorkspace.
• on_krylov_failure: What to do if the Krylov solver does not reach its tolerance: :stop (default) returns with status :krylov_failed without updating u, :continue takes the Newton step anyway (and counts the failure in stats.krylov_failures).
• callback: A function called once for the initial guess and then after each Newton iteration, with signature callback(u, res, norm_res).

Return value

(u, (; solved, status, stats, t)), where status is

• :converged: the residual norm satisfies the tolerance (solved == true)
• :max_iterations: max_niter Newton iterations were taken without convergence
• :krylov_failed: the Krylov solver failed and on_krylov_failure == :stop
• :nonfinite: the residual norm became NaN or Inf

and stats are the Ariadne.Stats.

newton_krylov!(ws; kwargs...)

Arguments

• ws: Pre-allocated NewtonKrylovWorkspace

Keyword Arguments

• tol_rel: Relative tolerance. Newton stops when ‖F(u)‖ <= tol_rel * ‖F(u₀)‖ + tol_abs.
• tol_abs: Absolute tolerance
• max_niter: Maximum number of Newton iterations
• forcing: Forcing-term strategy for the inexact Newton method, a Forcing such as EisenstatWalker() (default) or Fixed(η). It sets atol = 0 and rtol = η for each Krylov solve; krylov_kwargs can override both. If nothing, Ariadne does not set atol or rtol, so the Krylov solver uses its own default tolerances, or those given in krylov_kwargs.
• linesearch!: Line search strategy, an instance of a subtype of AbstractLineSearch.
• verbose::Int: Verbosity level. If verbose > 0, progress is logged with @info.
• M::Union{Nothing, Function}: If provided, M(ws.J) is passed as a keyword argument to the Krylov solver.
• N::Union{Nothing, Function}: If provided, N(ws.J) is passed as a keyword argument to the Krylov solver.- krylov_kwargs: Keyword arguments passed to the Krylov solver, e.g., (; restart = true, itmax = 100). The keys that configure the Krylov workspace (memory and window, e.g., memory = 50 for the restart length of GMRES(k)) only take effect when a NewtonKrylovWorkspace is created, i.e., in the methods that take a residual function, and are ignored by newton_krylov!(ws).
• algo, assume_p_const: Only for the methods that take a residual function; see NewtonKrylovWorkspace.
• on_krylov_failure: What to do if the Krylov solver does not reach its tolerance: :stop (default) returns with status :krylov_failed without updating u, :continue takes the Newton step anyway (and counts the failure in stats.krylov_failures).
• callback: A function called once for the initial guess and then after each Newton iteration, with signature callback(u, res, norm_res).

Return value

(u, (; solved, status, stats, t)), where status is

• :converged: the residual norm satisfies the tolerance (solved == true)
• :max_iterations: max_niter Newton iterations were taken without convergence
• :krylov_failed: the Krylov solver failed and on_krylov_failure == :stop
• :nonfinite: the residual norm became NaN or Inf

and stats are the Ariadne.Stats.

newton_krylov(F, u₀::AbstractArray, p = nothing, M::Int = length(u₀); kwargs...)

Takes a out-of-place residual function F(u, p).

Arguments

• F: res = F(u₀, p) solves res = F(u₀) = 0
• u₀: Initial guess
• p: Parameters
• M: Length of the output of F. Defaults to length(u₀).

Keyword Arguments

• tol_rel: Relative tolerance. Newton stops when ‖F(u)‖ <= tol_rel * ‖F(u₀)‖ + tol_abs.
• tol_abs: Absolute tolerance
• max_niter: Maximum number of Newton iterations
• forcing: Forcing-term strategy for the inexact Newton method, a Forcing such as EisenstatWalker() (default) or Fixed(η). It sets atol = 0 and rtol = η for each Krylov solve; krylov_kwargs can override both. If nothing, Ariadne does not set atol or rtol, so the Krylov solver uses its own default tolerances, or those given in krylov_kwargs.
• linesearch!: Line search strategy, an instance of a subtype of AbstractLineSearch.
• verbose::Int: Verbosity level. If verbose > 0, progress is logged with @info.
• M::Union{Nothing, Function}: If provided, M(ws.J) is passed as a keyword argument to the Krylov solver.
• N::Union{Nothing, Function}: If provided, N(ws.J) is passed as a keyword argument to the Krylov solver.- krylov_kwargs: Keyword arguments passed to the Krylov solver, e.g., (; restart = true, itmax = 100). The keys that configure the Krylov workspace (memory and window, e.g., memory = 50 for the restart length of GMRES(k)) only take effect when a NewtonKrylovWorkspace is created, i.e., in the methods that take a residual function, and are ignored by newton_krylov!(ws).
• algo, assume_p_const: Only for the methods that take a residual function; see NewtonKrylovWorkspace.
• on_krylov_failure: What to do if the Krylov solver does not reach its tolerance: :stop (default) returns with status :krylov_failed without updating u, :continue takes the Newton step anyway (and counts the failure in stats.krylov_failures).
• callback: A function called once for the initial guess and then after each Newton iteration, with signature callback(u, res, norm_res).

Return value

(u, (; solved, status, stats, t)), where status is

• :converged: the residual norm satisfies the tolerance (solved == true)
• :max_iterations: max_niter Newton iterations were taken without convergence
• :krylov_failed: the Krylov solver failed and on_krylov_failure == :stop
• :nonfinite: the residual norm became NaN or Inf

and stats are the Ariadne.Stats.

NewtonKrylovWorkspace

Pre-allocated workspace for newton_krylov!. Holds the residual buffer, negated-residual buffer, Jacobian operator (with its Enzyme caches), and the Krylov solver workspace so that no intermediate arrays are allocated during the Newton iteration.

Constructor

NewtonKrylovWorkspace(F!, u, p, res, alg=Val(:gmres); assume_p_const = false,
                      lazy_zero_shadows = false, krylov_kwargs = (;))

• F!: in-place residual function F!(res, u, p)
• u: initial-guess array (used as template; the workspace holds a reference to it)
• p: parameters
• res: pre-allocated residual buffer
• algo: Krylov algorithm symbol (e.g. :gmres, :fgmres) passed as a Val.
• assume_p_const, lazy_zero_shadows: passed through to JacobianOperator
• krylov_kwargs: the keyword arguments of the Krylov solver. The keys that configure the Krylov workspace (memory and window, e.g., memory = 50 for the size of the Krylov basis of GMRES and FGMRES) are passed to Krylov.krylov_workspace, all other keys are ignored. The same krylov_kwargs can therefore also be passed to newton_krylov!(ws), which ignores the workspace keys.

Example

    ws = NewtonKrylovWorkspace(F!, u, p, res, Val(:gmres))
    newton_krylov!(ws)

    # To change the initial guess, pass it to newton_krylov!:
    newton_krylov!(ws, u_new)    

AbstractLineSearch

Line search updates ws.u in-place along the Newton direction d and calls evaluate!(ws) to refresh ws.res and obtain the new residual norm. It returns (norm_res, status): the residual norm of the new state and status = :success, or :failed if it did not find a step with sufficient decrease (and took its last trial step), which newton_krylov! counts in stats.linesearch_failures.

Implemented variants

• NoLineSearch
• BacktrackingLineSearch

Custom line searches

struct CustomLineSearch <: AbstractLineSearch
    # parameters for the line search
end

function (ls::CustomLineSearch)(ws, norm_res_prior, d; verbose = 0)
    # update ws.u
    ws.u .+= d # for example, take the full Newton step
    return evaluate!(ws), :success
end

A line search is called as ls(ws, norm_res_prior, d; verbose) and must accept the keyword argument verbose, the verbosity level of newton_krylov!.

NoLineSearch()

A line search that does not perform any line search: it simply takes the full Newton step.

BacktrackingLineSearch(; n_iter_max = 10, alpha = 1.0e-4,
                       reject_exceptions = (DomainError,), parabolic = false)

Armijo backtracking: the step length λ, starting from 1, is reduced until ‖F(u + λ d)‖ <= (1 - alpha λ) ‖F(u)‖, for at most n_iter_max trials. If no trial satisfies this condition, the last trial step is taken and the line search counts as failed (stats.linesearch_failures of newton_krylov!).

By default, λ is halved in each reduction. With parabolic = true, the first reduction halves λ and later reductions use the safeguarded three-point parabolic model of [1] (Ariadne.LineSearches.parabolic_step), which reduces λ by a factor in [0.1, 0.5]. The parabolic model fails the Armijo condition less often, but does not reduce the number of Newton iterations on, e.g., the generalized Rosenbrock problem.

Trial states whose residual evaluation throws an exception of one of the types reject_exceptions (e.g., a DomainError from sqrt or log of a quantity that became negative in a too long step) count as trial states with infinite residual norm, so the step length is reduced further. This also holds for such exceptions thrown in tasks, e.g., in a Threads.@threads loop of the residual. The caught exceptions are logged with @info if newton_krylov! is called with verbose > 0, and with @debug otherwise. Other exceptions are rethrown. If all trials throw, the line search returns Inf and newton_krylov! stops with status :nonfinite. Use reject_exceptions = () to rethrow all exceptions.

References

• Kelley, C. T. (2022). Solving nonlinear equations with iterative methods: Solvers and examples in Julia. Society for Industrial and Applied Mathematics.
• https://github.com/ctkelley/SIAMFANLEquations.jl

parabolic_step(λc, λm, ff0, ffc, ffm; σ₀ = 0.1, σ₁ = 0.5)

Safeguarded three-point parabolic model for the step length of a line search, as in [1] (parab3p of SIAMFANLEquations.jl): minimize the parabola through the squared residual norms ff0 at λ = 0, ffc at the current step length λc, and ffm at the previous step length λm, and clamp the result to [σ₀ λc, σ₁ λc].

The model is p(λ) = ff0 + (c₁ λ + c₂ λ²) / d₁ with d₁ = (λc - λm) λc λm < 0, so it is convex if c₂ < 0, and its minimum is at λ = -c₁ / (2 c₂). If the parabola has negative curvature, the model is not helpful and the smallest step length σ₀ λc is taken (the corrected behavior of parab3p, see the comments there). If a residual norm is not finite, e.g., because the trial state threw an exception, λc is halved (σ₁ λc).

Adapted from https://github.com/ctkelley/SIAMFANLEquations.jl/blob/e5603e177dd007b065265641fb232d54020c4282/src/Tools/armijo.jl#L57 (MIT license).

Stats

Statistics of a newton_krylov! solve.

• outer_iterations: number of Newton iterations
• inner_iterations: total number of Krylov iterations
• norm_res: norm of the final residual
• krylov_failures: number of Newton iterations in which the Krylov solver did not reach its tolerance
• linesearch_failures: number of Newton iterations in which the line search did not find a step with sufficient decrease (and took its last trial step)

Forcing

Implements forcing for inexact Newton-Krylov. The equation $‖F′(u)d + F(u)‖ <= η * ‖F(u)‖$ gives the inexact Newton termination criterion.

Implemented variants

• Fixed
• EisenstatWalker

Fixed(η = 0.1)

Constant forcing term: every inner Krylov solve uses the relative tolerance η.

EisenstatWalker(η₀ = 0.999, η_max = 0.999, γ = 0.9, α = 2.0)

Adaptive forcing term for inexact Newton-Krylov methods. The forcing term η is updated each iteration as in [1], Eq. (3.6).

Parameters

• η₀: Initial forcing term. Default: 0.999.
• η_max: Maximum forcing term; must satisfy η_max ∈ (0, 1). Default: 0.999.
• γ: Contraction factor for the convergence rate estimate. Default: 0.9.
• α: Exponent for the convergence rate estimate; should be in (1, 2]. Default: 2.0.

References

• Kelley, C. T. (2022). Solving nonlinear equations with iterative methods: Solvers and examples in Julia. Society for Industrial and Applied Mathematics.
• Eisenstat, S. C., & Walker, H. F. (1996). Choosing the forcing terms in an inexact Newton method. SIAM Journal on Scientific Computing, 17(1), 16-32.

SparseJacobian(f!, res, u, p, pattern; coloring = GreedyColoringAlgorithm(),
               structure = :nonsymmetric, partition = :column, batchsize = 8,
               check_pattern = true)

Workspace for the assembly of the sparse Jacobian ∂f/∂u of the in-place function f!(res, u, p) at the state u with the sparsity pattern (any matrix whose structural nonzeros are the possible nonzeros of the Jacobian), by colored AD. res is a residual buffer. The Jacobian operator aliases u and p: assemble! assembles the Jacobian at their current values, so mutate them in place between assemblies.

The columns are grouped by a coloring of SparseMatrixColorings.jl for the ColoringProblem(; structure, partition), and batchsize colors are computed together by one batched Enzyme.jl pass (BatchedJacobianOperator, Julia ≥ 1.11; on older versions or for batchsize = 1 one JacobianOperator product per color is used):

• partition = :column: a distance-2 column coloring, one forward-mode Jacobian-vector product per color;
• partition = :row: a distance-2 row coloring, one reverse-mode vector-Jacobian product per color (fewer colors for Jacobians with dense rows, e.g., length(res) < length(u));
• structure = :symmetric (with partition = :column): a star coloring for Jacobians that are symmetric (the pattern must be symmetric), which needs fewer colors.

coloring is a coloring algorithm of SparseMatrixColorings.jl (e.g., GreedyColoringAlgorithm(LargestFirst())) or a vector with the color of each column (row). With a tuple of orders, e.g., GreedyColoringAlgorithm((NaturalOrder(), LargestFirst(), SmallestLast(), IncidenceDegree(), DynamicLargestFirst())), the coloring with the fewest colors is kept (for the 5-point Laplacian on a 60 × 60 grid, 5 instead of the 7 colors of the default natural order). The coloring is computed once and reused for all assemblies.

Given colors fit discretizations with a known block structure. For example, for a DG-like discretization with N nodes per cell and K cells in a periodic 1D mesh, where the residual of a cell depends on the cell and its two neighbors (K divisible by 3), the columns with the same element-local index i in every third cell can share a color:

pattern = kron(spdiagm(-1 => ones(K - 1), 0 => ones(K), 1 => ones(K - 1),
                       K - 1 => ones(1), 1 - K => ones(1)), ones(N, N))
colors = [i + N * mod(k - 1, 3) for i in 1:N, k in 1:K]  # 3N colors
A = SparseJacobian(f!, res, u, p, pattern; coloring = vec(colors))

The diagonal is always included in the pattern of square Jacobians, so that Diagonal(d) - J has the same pattern.

With check_pattern = true, each assembly also checks whether the products have nonzeros in rows (columns for partition = :row) that no column (row) of the color covers, i.e., whether pattern misses nonzeros of the Jacobian (e.g., a pattern detected at a different state for a function with branches). These entries are dropped from A.J; their number and largest magnitude in the last assembly are A.missed_entries and A.missed_max.

The Jacobian is stored in A.J.

assemble!(A::SparseJacobian) -> A.J

Assemble the sparse Jacobian ∂f/∂u of f!(res, u, p) at the current state u (and parameters p) of A into A.J by colored AD.

jacobian_sparsity(f!, u, p; kwargs...)

Detect the sparsity pattern of the Jacobian of f!(res, u, p) with respect to u. Requires SparseConnectivityTracer.jl to be loaded and f! to be generic in the element type of u and res (buffers in p must not have a fixed element type). The keyword argument detector defaults to TracerSparsityDetector() (global patterns); use TracerLocalSparsityDetector() for local patterns at u.

For discretizations with caches of fixed element type, compute the pattern yourself and pass it to SparseJacobian.

JacobianOperator

Efficient implementation of J(f,x,p) * v and v * J(f, x,p)'

BatchedJacobianOperator{N}

Ariadne.evaluate!(ws::NewtonKrylovWorkspace) -> norm_res

Evaluate F!(ws.res, ws.u, ws.p) in-place and return norm(ws.res).