diff --git a/examples/rosenbrock.jl b/examples/rosenbrock.jl new file mode 100644 index 00000000..5a85c5e2 --- /dev/null +++ b/examples/rosenbrock.jl @@ -0,0 +1,58 @@ +# Using the example from: +# A. Pal et al., “NonlinearSolve.Jl: High-performance and robust solvers for systems of nonlinear equations in Julia,” arXiv [math.NA], 24-Mar-2024. +# https://arxiv.org/abs/2403.16341 +# Fig 1 + +function generalized_rosenbrock(x, _) + vcat( + 1 - x[1], + 10 .* (x[2:end] .- x[1:(end - 1)] .* x[1:(end - 1)]) + ) +end + + +using Ariadne + +N = 12 +x_start = vcat(-1.2, ones(N-1)) +# for N=6 we require 21 iterations in 7.1211e-5 seconds +# for N=7 we require 66 iterations in 0.000129203 seconds +# for N=8 we require 56 iterations in 0.000106193 seconds +# for N=9 we do not find a solution within 100_000 iterations +# for N=10 we do not find a solution within 100_000 iterations +# for N=11 we do not find a solution within 100_000 iterations + +_, stats = newton_krylov( + generalized_rosenbrock, + copy(x_start); + algo = :gmres, + max_niter = 100_000 +) + +# Using simple line search: BacktrackingLineSearch + +# for N=6 we require 115 iterations in 0.000415178 seconds +# for N=7 we require 182 iterations in 0.003204213 seconds +# for N=8 we require 282 iterations in 0.001153047 seconds +# for N=9 we require 465 iterations in 0.00423094 seconds +# for N=10 we require 884 iterations in 0.006603219 seconds +# NOTE: Pal et.al. report that for N=10 their backtracking implementation does not converge. +# They use abstol = 1e-8 we use 1e-12 +# for N=11 we require 1568 iterations in 0.01126607 seconds +# for N=12 we require 2346 iterations in 0.024341541 seconds +_, stats = newton_krylov( + generalized_rosenbrock, + copy(x_start); + algo = :gmres, + linesearch! = Ariadne.LineSearches.BacktrackingLineSearch(), + max_niter = 100_000 +) + +_, stats = newton_krylov( + generalized_rosenbrock, + copy(x_start); + algo = :gmres, + linesearch! = Ariadne.LineSearches.BacktrackingLineSearch(; parabolic = true), + max_niter = 100_000 +) + diff --git a/src/Ariadne.jl b/src/Ariadne.jl index 65e77d1a..38e99d0a 100644 --- a/src/Ariadne.jl +++ b/src/Ariadne.jl @@ -227,6 +227,14 @@ function Base.collect(JOp::Union{Adjoint{<:Any, <:AbstractJacobianOperator}, Tra return J end +## +# LineSearches +## + +include("linesearches.jl") +import .LineSearches: AbstractLineSearch, NoLineSearch +export NoLineSearch + ## # Newton-Krylov ## @@ -357,6 +365,7 @@ function newton_krylov!( tol_abs = 1.0e-12, # Scipy uses 6e-6 max_niter = 50, forcing::Union{Forcing, Nothing} = EisenstatWalker(), + linesearch!::AbstractLineSearch = NoLineSearch(), verbose = 0, algo = :gmres, M = nothing, @@ -405,26 +414,21 @@ function newton_krylov!( kwargs = (; atol = zero(η), rtol = η, kwargs...) end - # Solve: J d = res = F(u) - # Typically, the Newton method is formulated as J d = -F(u) - # with update u = u + d. - # To simplify the implementation, we solve J d = F(u) - # and update u = u - d instead. - # `res` is modified by J, so we create a copy `res` - # TODO: provide a temporary storage for `res` - krylov_solve!(workspace, J, copy(res); kwargs...) - - d = workspace.x # (negative) Newton direction - s = 1 # Scaling of the Newton step TODO: LineSearch + # Solve: J d = -res = -F(u) + # The Newton method is formulated as J d = -F(u) + # `res` is modified by J, so we create a `neg_res` copy here. + # TODO: provide cache for `neg_res` to avoid this allocation. + neg_res = map(-, res) + krylov_solve!(workspace, J, neg_res; kwargs...) - # Update u - u .= muladd.(-s, d, u) # u = u - s * d + d₀ = workspace.x # (negative) Newton direction - # Update residual and norm + # Perform line search + # Must update `res` and `u` in-place + # by calling: F!(res, u, p) # res = F(u) n_res_prior = n_res + n_res = linesearch!(F!, res, n_res_prior, u, p, d₀) - F!(res, u, p) # res = F(u) - n_res = norm(res) callback(u, res, n_res) if isinf(n_res) || isnan(n_res) diff --git a/src/linesearches.jl b/src/linesearches.jl new file mode 100644 index 00000000..66210a77 --- /dev/null +++ b/src/linesearches.jl @@ -0,0 +1,163 @@ +module LineSearches + +using LinearAlgebra + +""" + AbstractLineSearch + +Line search may update the solution `u` and the residual `res` in-place, +given the function `F!`, parameters `p`, and the Newton direction `d`. + +They must call `F!(res, u, p)` to update the residual after updating `u`. +""" +abstract type AbstractLineSearch end + +""" + NoLineSearch() + +A line search that does not perform any line search: it simply takes the full Newton step. +""" +struct NoLineSearch <: AbstractLineSearch end + +function (::NoLineSearch)(F!, res, _, u, p, d) + # No line search: take the full Newton step + u .+= d + F!(res, u, p) + return norm(res) +end + +""" + BacktrackingLineSearch(; n_iter_max = 10, parabolic = false) + +## References + +- Kelley, C. T. (2022). + Solving nonlinear equations with iterative methods: + Solvers and examples in Julia. + Society for Industrial and Applied Mathematics. +- +""" +Base.@kwdef struct BacktrackingLineSearch <: AbstractLineSearch + n_iter_max::Int = 10 + parabolic::Bool = false +end + +function (ls::BacktrackingLineSearch)(F!, res, n_res_prior, u, p, d) + alpha = 1.0e-4 + lambda = 1.0 + + @assert ls.n_iter_max > 0 "n_iter_max must be positive" + @assert alpha > 0 "alpha must be positive" + + n_res = Inf + u_trial = similar(u) + + ff0 = n_res_prior^2 + ffc = ff0 + ffm = ffc + lamc = lambda + + for i in 1:ls.n_iter_max + # Take a step of size s + u_trial .= muladd.(lambda, d, u) # u = u + lambda * d + F!(res, u_trial, p) + n_res = norm(res) + @info "Line search iteration $i: lambda = $lambda, residual norm = $n_res, previous residual norm = $n_res_prior" + + # Armijo condition + if n_res <= (1 - alpha * lambda) * n_res_prior + u .= u_trial + return n_res + end + + ffm = ffc + ffc = n_res^2 + lambda = update_lambda(i, ls.parabolic, lambda, lamc, ff0, ffc, ffm) + end + u .= u_trial + return n_res +end + +function update_lambda(i, parabolic, lambda, lamc, ff0, ffc, ffm) + if !parabolic + return lambda * 0.5 + end + if i == 1 + return lambda * 0.5 + else + return parab3p(lambda, lamc, ff0, ffc, ffm) + end +end + +# From https://github.com/ctkelley/SIAMFANLEquations.jl/blob/e5603e177dd007b065265641fb232d54020c4282/src/Tools/armijo.jl#L57 +""" +parab3p(lambdac, lambdam, ff0, ffc, ffm) + +Three point parabolic line search. + +input:\n + lambdac = current steplength + lambdam = previous steplength + ff0 = value of || F(x_c) ||^2 + ffc = value of || F(x_c + lambdac d) ||^2 + ffm = value of || F(x_c + lambdam d) ||^2 + +output:\n + lambdap = new value of lambda + +internal parameters:\n + sigma0 = .1, sigma1=.5, safeguarding bounds for the linesearch + +You get here if cutting the steplength in half doesn't get you +sufficient decrease. Now you have three points and can build a parabolic +model. I do not like cubic models because they either need four points +or a derivative. + +So let's think about how this works. I cheat a bit and check the model +for negative curvature, which I don't want to see. + + The polynomial is + + p(lambda) = ff0 + (c1 lambda + c2 lambda^2)/d1 + + d1 = (lambdac - lambdam)*lambdac*lambdam < 0 + So if c2 > 0 we have negative curvature and default to + lambdap = sigma0 * lambda + The logic is that negative curvature is telling us that + the polynomial model is not helping much, so it looks better + to take the smallest possible step. This is not what I did in the + matlab code because I did it wrong. I have sinced fixed it. + + So (Students, listen up!) if c2 < 0 then all we gotta do is minimize + (c1 lambda + c2 lambda^2)/d1 over [.1* lambdac, .5*lambdac] + This means to MAXIMIZE c1 lambda + c2 lambda^2 becase d1 < 0. + So I find the zero of the derivative and check the endpoints. + +""" +function parab3p(lambdac, lambdam, ff0, ffc, ffm) + # + # internal parameters + # + sigma0 = 0.1 + sigma1 = 0.5 + # + c2 = lambdam * (ffc - ff0) - lambdac * (ffm - ff0) + if c2 >= 0 + # + # Sanity check for negative curvature + # + lambdap = sigma0 * lambdac + else + # + # It's a convex parabola, so use calculus! + # + c1 = lambdac * lambdac * (ffm - ff0) - lambdam * lambdam * (ffc - ff0) + lambdap = -c1 * 0.5 / c2 + # + lambdaup = sigma1 * lambdac + lambdadown = sigma0 * lambdac + lambdap = max(lambdadown, min(lambdaup, lambdap)) + end +end + +end # module LineSearches