From fec96d98585aa45d0a407040e7959518dc145a9b Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 3 Apr 2026 16:17:03 +0530 Subject: [PATCH 1/5] Add native StrongWolfeLineSearch (Nocedal & Wright Alg. 3.5/3.6) Signed-off-by: AdityaPandeyCN --- src/LineSearch.jl | 3 +- src/strong_wolfe.jl | 346 +++++++++++++++++++++++++++++++++ test/custom_optimizer_tests.jl | 4 +- test/root_finding_tests.jl | 8 +- 4 files changed, 356 insertions(+), 5 deletions(-) create mode 100644 src/strong_wolfe.jl diff --git a/src/LineSearch.jl b/src/LineSearch.jl index aeb982b..bf5446b 100644 --- a/src/LineSearch.jl +++ b/src/LineSearch.jl @@ -26,6 +26,7 @@ include("golden_section.jl") include("li_fukushima.jl") include("no_search.jl") include("robust_non_monotone.jl") +include("strong_wolfe.jl") include("line_searches_ext.jl") @@ -38,7 +39,7 @@ export LineSearchSolution export BackTracking export GoldenSection -export NoLineSearch, LiFukushimaLineSearch, RobustNonMonotoneLineSearch +export NoLineSearch, LiFukushimaLineSearch, RobustNonMonotoneLineSearch, StrongWolfeLineSearch export LineSearchesJL include("precompilation.jl") diff --git a/src/strong_wolfe.jl b/src/strong_wolfe.jl new file mode 100644 index 0000000..1faf7a9 --- /dev/null +++ b/src/strong_wolfe.jl @@ -0,0 +1,346 @@ +""" + StrongWolfeLineSearch(; autodiff = nothing, c1 = 1e-4, c2 = 0.9, + α_init = 1.0, maxiters::Int = 10) + +Strong Wolfe line search satisfying both Armijo (sufficient decrease) and +curvature conditions. Based on Nocedal & Wright, "Numerical Optimization" (2006), +Algorithms 3.5 and 3.6. + +`autodiff` is the automatic differentiation backend to use for computing the +directional derivative. Must be specified if analytic jacobian/jvp/vjp is not +available. + +""" +@kwdef @concrete struct StrongWolfeLineSearch <: AbstractLineSearchAlgorithm + autodiff = nothing + c1 = 1e-4 + c2 = 0.9 + α_init = 1.0 + maxiters::Int = 10 +end + +@concrete mutable struct StrongWolfeLineSearchCache <: AbstractLineSearchCache + ϕ + dϕ + f + p + deriv_op + u_cache + fu_cache + c1 + c2 + α + maxiters::Int + stats <: Union{SciMLBase.NLStats, Nothing} + alg <: StrongWolfeLineSearch +end + +@concrete struct StaticStrongWolfeLineSearchCache <: AbstractLineSearchCache + f + grad_f + p + c1 + c2 + α_init + maxiters::Int +end + +function CommonSolve.init( + prob::AbstractNonlinearProblem, alg::StrongWolfeLineSearch, + fu::Union{SArray, Number}, u::Union{SArray, Number}; + stats::Union{SciMLBase.NLStats, Nothing} = nothing, + grad_f = nothing, kwargs... + ) + if stats === nothing + T = promote_type(eltype(fu), eltype(u)) + return StaticStrongWolfeLineSearchCache( + prob.f, grad_f, prob.p, T(alg.c1), T(alg.c2), T(alg.α_init), alg.maxiters + ) + end + return generic_strongwolfe_init(prob, alg, fu, u; stats, kwargs...) +end + +function CommonSolve.init( + prob::AbstractNonlinearProblem, alg::StrongWolfeLineSearch, fu, u; + autodiff = nothing, kwargs... + ) + return generic_strongwolfe_init(prob, alg, fu, u; autodiff, kwargs...) +end + +function generic_strongwolfe_init( + prob::AbstractNonlinearProblem, alg::StrongWolfeLineSearch, + fu, u; stats::Union{SciMLBase.NLStats, Nothing} = nothing, + autodiff = nothing, kwargs... + ) + autodiff = autodiff !== nothing ? autodiff : alg.autodiff + + _, _, deriv_op = construct_jvp_or_vjp_operator(prob, fu, u; autodiff) + + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + T = promote_type(eltype(fu), eltype(u)) + + ϕ = @closure (f, p, u, du, α, u_cache, fu_cache) -> begin + @bb @. u_cache = u + α * du + fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) + add_nf!(stats) + return @fastmath norm(fu_cache)^2 / 2 + end + + dϕ = @closure (f, p, u, du, α, u_cache, fu_cache, deriv_op) -> begin + @bb @. u_cache = u + α * du + fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) + add_nf!(stats) + return deriv_op(du, u_cache, fu_cache, p) + end + + return StrongWolfeLineSearchCache( + ϕ, dϕ, prob.f, prob.p, deriv_op, u_cache, fu_cache, + T(alg.c1), T(alg.c2), T(alg.α_init), alg.maxiters, stats, alg + ) +end + +@inline function _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) + d1 = dϕ_lo + dϕ_hi - 3 * (ϕ_lo - ϕ_hi) / (a_lo - a_hi) + desc = d1 * d1 - dϕ_lo * dϕ_hi + d2 = sqrt(max(desc, zero(desc))) + ifelse( + desc < 0, + (a_lo + a_hi) / 2, + a_hi - (a_hi - a_lo) * ((dϕ_hi + d2 - d1) / (dϕ_hi - dϕ_lo + 2 * d2)) + ) +end + +@inline function _sw_phi(f, p, u, du, α) + u_new = u .+ α .* du + fu = f(u_new, p) + @fastmath norm(fu)^2 / 2 +end + +@inline function _sw_dphi(grad_f, p, u, du, α) + u_new = u .+ α .* du + g = grad_f(u_new, p) + dot(g, du) +end + +@inline function _sw_zoom_static(f, grad_f, p, u, du, + a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, c1, c2, maxiters) + T = typeof(a_lo) + α_out = a_lo + ok = false + done = false + + for _ in 1:maxiters + if !done + ϕ_hi = _sw_phi(f, p, u, du, a_hi) + dϕ_lo_val = _sw_dphi(grad_f, p, u, du, a_lo) + dϕ_hi_val = _sw_dphi(grad_f, p, u, du, a_hi) + + α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo_val, dϕ_hi_val) + bracket = T(0.01) * abs(a_hi - a_lo) + α_j = clamp(α_j, + min(a_lo, a_hi) + bracket, + max(a_lo, a_hi) - bracket) + + ϕ_j = _sw_phi(f, p, u, du, α_j) + + if (ϕ_j > ϕ_0 + c1 * α_j * dϕ_0) || (ϕ_j >= ϕ_lo) + a_hi = α_j + else + dϕ_j = _sw_dphi(grad_f, p, u, du, α_j) + if abs(dϕ_j) <= -c2 * dϕ_0 + α_out, ok, done = α_j, true, true + else + if dϕ_j * (a_hi - a_lo) >= zero(T) + a_hi = a_lo + end + a_lo = α_j + ϕ_lo = ϕ_j + end + end + end + end + + if !done + α_out = a_lo + end + (α_out, ok) +end + +@inline function _sw_search_static(f, grad_f, p, u, du, ϕ_0, dϕ_0, + c1, c2, α_init, maxiters) + T = typeof(α_init) + α_out = zero(T) + ok = false + + dϕ_0 >= zero(T) && return (zero(T), false) + + α_prev = zero(T) + α_i = α_init + ϕ_prev = ϕ_0 + done = false + + for i in 1:maxiters + if !done + ϕ_i = _sw_phi(f, p, u, du, α_i) + + if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) + α_z, ok_z = _sw_zoom_static(f, grad_f, p, u, du, + α_prev, α_i, ϕ_0, dϕ_0, ϕ_prev, c1, c2, maxiters) + α_out, ok, done = α_z, ok_z, true + else + dϕ_i = _sw_dphi(grad_f, p, u, du, α_i) + + if abs(dϕ_i) <= -c2 * dϕ_0 + α_out, ok, done = α_i, true, true + elseif dϕ_i >= zero(T) + α_z, ok_z = _sw_zoom_static(f, grad_f, p, u, du, + α_i, α_prev, ϕ_0, dϕ_0, ϕ_i, c1, c2, maxiters) + α_out, ok, done = α_z, ok_z, true + else + α_prev = α_i + ϕ_prev = ϕ_i + α_i *= T(2) + end + end + end + end + + if !done + α_out, ok = α_prev, true + end + (α_out, ok) +end + +@inline function _sw_zoom(ϕ, dϕ, a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, c1, c2, maxiters) + T = typeof(a_lo) + α_out = a_lo + ok = false + done = false + + for _ in 1:maxiters + if !done + ϕ_hi = ϕ(a_hi) + dϕ_lo_val = dϕ(a_lo) + dϕ_hi_val = dϕ(a_hi) + + α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo_val, dϕ_hi_val) + bracket = T(0.01) * abs(a_hi - a_lo) + α_j = clamp(α_j, + min(a_lo, a_hi) + bracket, + max(a_lo, a_hi) - bracket) + + ϕ_j = ϕ(α_j) + + if (ϕ_j > ϕ_0 + c1 * α_j * dϕ_0) || (ϕ_j >= ϕ_lo) + a_hi = α_j + else + dϕ_j = dϕ(α_j) + if abs(dϕ_j) <= -c2 * dϕ_0 + α_out, ok, done = α_j, true, true + else + if dϕ_j * (a_hi - a_lo) >= zero(T) + a_hi = a_lo + end + a_lo = α_j + ϕ_lo = ϕ_j + end + end + end + end + + if !done + α_out = a_lo + end + (α_out, ok) +end + +@inline function _sw_search(ϕ, dϕ, ϕ_0, dϕ_0, c1, c2, α_init, maxiters) + T = typeof(α_init) + α_out = zero(T) + ok = false + + dϕ_0 >= zero(T) && return (zero(T), false) + + α_prev = zero(T) + α_i = α_init + ϕ_prev = ϕ_0 + done = false + + for i in 1:maxiters + if !done + ϕ_i = ϕ(α_i) + + if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) + α_z, ok_z = _sw_zoom(ϕ, dϕ, α_prev, α_i, + ϕ_0, dϕ_0, ϕ_prev, c1, c2, maxiters) + α_out, ok, done = α_z, ok_z, true + else + dϕ_i = dϕ(α_i) + + if abs(dϕ_i) <= -c2 * dϕ_0 + α_out, ok, done = α_i, true, true + elseif dϕ_i >= zero(T) + α_z, ok_z = _sw_zoom(ϕ, dϕ, α_i, α_prev, + ϕ_0, dϕ_0, ϕ_i, c1, c2, maxiters) + α_out, ok, done = α_z, ok_z, true + else + α_prev = α_i + ϕ_prev = ϕ_i + α_i *= T(2) + end + end + end + end + + if !done + α_out, ok = α_prev, true + end + (α_out, ok) +end + +function CommonSolve.solve!(cache::StrongWolfeLineSearchCache, u, du) + T = promote_type(eltype(du), eltype(u)) + + ϕ = @closure α -> cache.ϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache) + dϕ = @closure α -> cache.dϕ(cache.f, cache.p, u, du, α, + cache.u_cache, cache.fu_cache, cache.deriv_op) + + ϕ_0 = ϕ(zero(T)) + dϕ_0 = dϕ(zero(T)) + + isfinite(ϕ_0) || return LineSearchSolution(cache.α, ReturnCode.Failure) + dϕ_0 >= zero(T) && return LineSearchSolution(cache.α, ReturnCode.Failure) + + α, ok = _sw_search(ϕ, dϕ, ϕ_0, dϕ_0, cache.c1, cache.c2, cache.α, cache.maxiters) + + ok && return LineSearchSolution(α, ReturnCode.Success) + return LineSearchSolution(cache.α, ReturnCode.Failure) +end + +function CommonSolve.solve!(cache::StaticStrongWolfeLineSearchCache, u, du) + T = promote_type(eltype(du), eltype(u)) + + ϕ_0 = _sw_phi(cache.f, cache.p, u, du, zero(T)) + dϕ_0 = _sw_dphi(cache.grad_f, cache.p, u, du, zero(T)) + + isfinite(ϕ_0) || return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) + dϕ_0 >= zero(T) && return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) + + α, ok = _sw_search_static(cache.f, cache.grad_f, cache.p, u, du, + ϕ_0, dϕ_0, cache.c1, cache.c2, T(cache.α_init), cache.maxiters) + + ok && return LineSearchSolution(α, ReturnCode.Success) + return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) +end + +function SciMLBase.reinit!( + cache::StrongWolfeLineSearchCache; p = missing, stats = missing, kwargs... + ) + p !== missing && (cache.p = p) + stats !== missing && (cache.stats = stats) + cache.α = oftype(cache.α, cache.alg.α_init) + cache.c1 = oftype(cache.c1, cache.alg.c1) + cache.c2 = oftype(cache.c2, cache.alg.c2) + return cache +end \ No newline at end of file diff --git a/test/custom_optimizer_tests.jl b/test/custom_optimizer_tests.jl index d93f25c..b831b33 100644 --- a/test/custom_optimizer_tests.jl +++ b/test/custom_optimizer_tests.jl @@ -120,6 +120,7 @@ end GoldenSection(; tol = 1.0e-4), BackTracking(; order = Val(3), autodiff), BackTracking(; order = Val(2), autodiff), + StrongWolfeLineSearch(; autodiff), ) fu, u, iter, alphas = gradient_descent(nlp, method; autodiff) @@ -143,6 +144,7 @@ end GoldenSection(; tol = 1.0e-4), BackTracking(; order = Val(3), autodiff), BackTracking(; order = Val(2), autodiff), + StrongWolfeLineSearch(; autodiff), ) fu, u, iter, alphas = gradient_descent(nlp, method; autodiff) @@ -152,4 +154,4 @@ end end end end -end +end \ No newline at end of file diff --git a/test/root_finding_tests.jl b/test/root_finding_tests.jl index 302ce61..1918a1b 100644 --- a/test/root_finding_tests.jl +++ b/test/root_finding_tests.jl @@ -140,7 +140,7 @@ end NoLineSearch(0.5), GoldenSection(; tol = 1.0e-4), RobustNonMonotoneLineSearch(), - RobustNonMonotoneLineSearch(; M = 1), #strictly monotonous case + RobustNonMonotoneLineSearch(; M = 1), RobustNonMonotoneLineSearch(; M = 15), ) converged, fu, u, iter, alphas = newton_raphson(nlp, method) @@ -156,6 +156,7 @@ end @testset "method: $(nameof(typeof(method)))" for method in ( BackTracking(; order = Val(3), autodiff), BackTracking(; order = Val(2), autodiff), + StrongWolfeLineSearch(; autodiff), ) converged, fu, u, iter, alphas = newton_raphson(nlp, method) @@ -174,7 +175,7 @@ end NoLineSearch(0.5), GoldenSection(; tol = 1.0e-4), RobustNonMonotoneLineSearch(), - RobustNonMonotoneLineSearch(; M = 1), #strictly monotonous case + RobustNonMonotoneLineSearch(; M = 1), RobustNonMonotoneLineSearch(; M = 15), ) converged, fu, u, iter, alphas = newton_raphson(nlp, method) @@ -189,6 +190,7 @@ end @testset "method: $(nameof(typeof(method)))" for method in ( BackTracking(; order = Val(3), autodiff), BackTracking(; order = Val(2), autodiff), + StrongWolfeLineSearch(; autodiff), ) converged, fu, u, iter, alphas = newton_raphson(nlp, method) @@ -197,4 +199,4 @@ end end end end -end +end \ No newline at end of file From c3feeff309f17786a694256b1b839e77107b414f Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sun, 5 Apr 2026 19:12:09 +0530 Subject: [PATCH 2/5] =?UTF-8?q?combine=20=CF=95/d=CF=95=20into=20single=20?= =?UTF-8?q?=CF=95d=CF=95=20eval,=20cache=20bracket=20values=20in=20zoom,?= =?UTF-8?q?=20fix=20stats=20closure=20bug,=20add=20grad=5Ff=20guard=20and?= =?UTF-8?q?=20=CE=B1=5Fmax=20cap?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: AdityaPandeyCN --- src/strong_wolfe.jl | 187 ++++++++++++++++++++++---------------------- 1 file changed, 95 insertions(+), 92 deletions(-) diff --git a/src/strong_wolfe.jl b/src/strong_wolfe.jl index 1faf7a9..d44b08c 100644 --- a/src/strong_wolfe.jl +++ b/src/strong_wolfe.jl @@ -9,7 +9,6 @@ Algorithms 3.5 and 3.6. `autodiff` is the automatic differentiation backend to use for computing the directional derivative. Must be specified if analytic jacobian/jvp/vjp is not available. - """ @kwdef @concrete struct StrongWolfeLineSearch <: AbstractLineSearchAlgorithm autodiff = nothing @@ -19,9 +18,8 @@ available. maxiters::Int = 10 end +# CPU path: stores raw ingredients, closure built fresh in solve! @concrete mutable struct StrongWolfeLineSearchCache <: AbstractLineSearchCache - ϕ - dϕ f p deriv_op @@ -35,6 +33,7 @@ end alg <: StrongWolfeLineSearch end +# GPU path: static cache for SArray/Number, no allocations @concrete struct StaticStrongWolfeLineSearchCache <: AbstractLineSearchCache f grad_f @@ -52,6 +51,9 @@ function CommonSolve.init( grad_f = nothing, kwargs... ) if stats === nothing + grad_f === nothing && error( + "StrongWolfeLineSearch requires `grad_f` for static (GPU) dispatch" + ) T = promote_type(eltype(fu), eltype(u)) return StaticStrongWolfeLineSearchCache( prob.f, grad_f, prob.p, T(alg.c1), T(alg.c2), T(alg.α_init), alg.maxiters @@ -73,33 +75,19 @@ function generic_strongwolfe_init( autodiff = nothing, kwargs... ) autodiff = autodiff !== nothing ? autodiff : alg.autodiff - _, _, deriv_op = construct_jvp_or_vjp_operator(prob, fu, u; autodiff) @bb u_cache = similar(u) @bb fu_cache = similar(fu) T = promote_type(eltype(fu), eltype(u)) - ϕ = @closure (f, p, u, du, α, u_cache, fu_cache) -> begin - @bb @. u_cache = u + α * du - fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) - add_nf!(stats) - return @fastmath norm(fu_cache)^2 / 2 - end - - dϕ = @closure (f, p, u, du, α, u_cache, fu_cache, deriv_op) -> begin - @bb @. u_cache = u + α * du - fu_cache = evaluate_f!!(f, fu_cache, u_cache, p) - add_nf!(stats) - return deriv_op(du, u_cache, fu_cache, p) - end - return StrongWolfeLineSearchCache( - ϕ, dϕ, prob.f, prob.p, deriv_op, u_cache, fu_cache, + prob.f, prob.p, deriv_op, u_cache, fu_cache, T(alg.c1), T(alg.c2), T(alg.α_init), alg.maxiters, stats, alg ) end +# cubic interpolation with bisection fallback when discriminant is negative @inline function _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) d1 = dϕ_lo + dϕ_hi - 3 * (ϕ_lo - ϕ_hi) / (a_lo - a_hi) desc = d1 * d1 - dϕ_lo * dϕ_hi @@ -111,51 +99,54 @@ end ) end -@inline function _sw_phi(f, p, u, du, α) +# combined f + grad_f eval for static/GPU path +@inline function _sw_phi_dphi(f, grad_f, p, u, du, α) u_new = u .+ α .* du fu = f(u_new, p) - @fastmath norm(fu)^2 / 2 -end - -@inline function _sw_dphi(grad_f, p, u, du, α) - u_new = u .+ α .* du + ϕ = @fastmath norm(fu)^2 / 2 g = grad_f(u_new, p) - dot(g, du) + dϕ = dot(g, du) + return ϕ, dϕ end +# N&W Algorithm 3.6 (zoom) for static path +# uses `done` flag instead of early return for GPU kernel compatibility @inline function _sw_zoom_static(f, grad_f, p, u, du, - a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, c1, c2, maxiters) + a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, dϕ_lo, c1, c2, maxiters) T = typeof(a_lo) α_out = a_lo ok = false done = false + # evaluate a_hi once, then maintain via updates below + ϕ_hi, dϕ_hi = _sw_phi_dphi(f, grad_f, p, u, du, a_hi) + for _ in 1:maxiters if !done - ϕ_hi = _sw_phi(f, p, u, du, a_hi) - dϕ_lo_val = _sw_dphi(grad_f, p, u, du, a_lo) - dϕ_hi_val = _sw_dphi(grad_f, p, u, du, a_hi) - - α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo_val, dϕ_hi_val) + α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) bracket = T(0.01) * abs(a_hi - a_lo) α_j = clamp(α_j, min(a_lo, a_hi) + bracket, max(a_lo, a_hi) - bracket) - ϕ_j = _sw_phi(f, p, u, du, α_j) + ϕ_j, dϕ_j = _sw_phi_dphi(f, grad_f, p, u, du, α_j) if (ϕ_j > ϕ_0 + c1 * α_j * dϕ_0) || (ϕ_j >= ϕ_lo) a_hi = α_j + ϕ_hi = ϕ_j + dϕ_hi = dϕ_j else - dϕ_j = _sw_dphi(grad_f, p, u, du, α_j) if abs(dϕ_j) <= -c2 * dϕ_0 α_out, ok, done = α_j, true, true else if dϕ_j * (a_hi - a_lo) >= zero(T) a_hi = a_lo + ϕ_hi = ϕ_lo + dϕ_hi = dϕ_lo end a_lo = α_j ϕ_lo = ϕ_j + dϕ_lo = dϕ_j end end end @@ -164,86 +155,93 @@ end if !done α_out = a_lo end - (α_out, ok) + return (α_out, ok) end +# N&W Algorithm 3.5 (outer search) for static path @inline function _sw_search_static(f, grad_f, p, u, du, ϕ_0, dϕ_0, c1, c2, α_init, maxiters) T = typeof(α_init) - α_out = zero(T) - ok = false + α_max = T(65536) dϕ_0 >= zero(T) && return (zero(T), false) α_prev = zero(T) α_i = α_init ϕ_prev = ϕ_0 + dϕ_prev = dϕ_0 done = false + α_out = zero(T) + ok = false for i in 1:maxiters if !done - ϕ_i = _sw_phi(f, p, u, du, α_i) + ϕ_i, dϕ_i = _sw_phi_dphi(f, grad_f, p, u, du, α_i) if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) α_z, ok_z = _sw_zoom_static(f, grad_f, p, u, du, - α_prev, α_i, ϕ_0, dϕ_0, ϕ_prev, c1, c2, maxiters) + α_prev, α_i, ϕ_0, dϕ_0, ϕ_prev, dϕ_prev, + c1, c2, maxiters) + α_out, ok, done = α_z, ok_z, true + elseif abs(dϕ_i) <= -c2 * dϕ_0 + α_out, ok, done = α_i, true, true + elseif dϕ_i >= zero(T) + α_z, ok_z = _sw_zoom_static(f, grad_f, p, u, du, + α_i, α_prev, ϕ_0, dϕ_0, ϕ_i, dϕ_i, + c1, c2, maxiters) α_out, ok, done = α_z, ok_z, true else - dϕ_i = _sw_dphi(grad_f, p, u, du, α_i) - - if abs(dϕ_i) <= -c2 * dϕ_0 - α_out, ok, done = α_i, true, true - elseif dϕ_i >= zero(T) - α_z, ok_z = _sw_zoom_static(f, grad_f, p, u, du, - α_i, α_prev, ϕ_0, dϕ_0, ϕ_i, c1, c2, maxiters) - α_out, ok, done = α_z, ok_z, true - else - α_prev = α_i - ϕ_prev = ϕ_i - α_i *= T(2) - end + α_prev = α_i + ϕ_prev = ϕ_i + dϕ_prev = dϕ_i + α_i = min(α_i * T(2), α_max) end end end if !done - α_out, ok = α_prev, true + α_out = α_prev + ok = false end - (α_out, ok) + return (α_out, ok) end -@inline function _sw_zoom(ϕ, dϕ, a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, c1, c2, maxiters) +# N&W Algorithm 3.6 (zoom) for mutable path +@inline function _sw_zoom(ϕdϕ, a_lo, a_hi, ϕ_0, dϕ_0, + ϕ_lo, dϕ_lo, c1, c2, maxiters) T = typeof(a_lo) α_out = a_lo ok = false done = false + ϕ_hi, dϕ_hi = ϕdϕ(a_hi) + for _ in 1:maxiters if !done - ϕ_hi = ϕ(a_hi) - dϕ_lo_val = dϕ(a_lo) - dϕ_hi_val = dϕ(a_hi) - - α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo_val, dϕ_hi_val) + α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) bracket = T(0.01) * abs(a_hi - a_lo) α_j = clamp(α_j, min(a_lo, a_hi) + bracket, max(a_lo, a_hi) - bracket) - ϕ_j = ϕ(α_j) + ϕ_j, dϕ_j = ϕdϕ(α_j) if (ϕ_j > ϕ_0 + c1 * α_j * dϕ_0) || (ϕ_j >= ϕ_lo) a_hi = α_j + ϕ_hi = ϕ_j + dϕ_hi = dϕ_j else - dϕ_j = dϕ(α_j) if abs(dϕ_j) <= -c2 * dϕ_0 α_out, ok, done = α_j, true, true else if dϕ_j * (a_hi - a_lo) >= zero(T) a_hi = a_lo + ϕ_hi = ϕ_lo + dϕ_hi = dϕ_lo end a_lo = α_j ϕ_lo = ϕ_j + dϕ_lo = dϕ_j end end end @@ -252,67 +250,73 @@ end if !done α_out = a_lo end - (α_out, ok) + return (α_out, ok) end -@inline function _sw_search(ϕ, dϕ, ϕ_0, dϕ_0, c1, c2, α_init, maxiters) +# N&W Algorithm 3.5 (outer search) for mutable path +@inline function _sw_search(ϕdϕ, ϕ_0, dϕ_0, c1, c2, α_init, maxiters) T = typeof(α_init) - α_out = zero(T) - ok = false + α_max = T(65536) dϕ_0 >= zero(T) && return (zero(T), false) α_prev = zero(T) α_i = α_init ϕ_prev = ϕ_0 + dϕ_prev = dϕ_0 done = false + α_out = zero(T) + ok = false for i in 1:maxiters if !done - ϕ_i = ϕ(α_i) + ϕ_i, dϕ_i = ϕdϕ(α_i) if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) - α_z, ok_z = _sw_zoom(ϕ, dϕ, α_prev, α_i, - ϕ_0, dϕ_0, ϕ_prev, c1, c2, maxiters) + α_z, ok_z = _sw_zoom(ϕdϕ, α_prev, α_i, + ϕ_0, dϕ_0, ϕ_prev, dϕ_prev, c1, c2, maxiters) + α_out, ok, done = α_z, ok_z, true + elseif abs(dϕ_i) <= -c2 * dϕ_0 + α_out, ok, done = α_i, true, true + elseif dϕ_i >= zero(T) + α_z, ok_z = _sw_zoom(ϕdϕ, α_i, α_prev, + ϕ_0, dϕ_0, ϕ_i, dϕ_i, c1, c2, maxiters) α_out, ok, done = α_z, ok_z, true else - dϕ_i = dϕ(α_i) - - if abs(dϕ_i) <= -c2 * dϕ_0 - α_out, ok, done = α_i, true, true - elseif dϕ_i >= zero(T) - α_z, ok_z = _sw_zoom(ϕ, dϕ, α_i, α_prev, - ϕ_0, dϕ_0, ϕ_i, c1, c2, maxiters) - α_out, ok, done = α_z, ok_z, true - else - α_prev = α_i - ϕ_prev = ϕ_i - α_i *= T(2) - end + α_prev = α_i + ϕ_prev = ϕ_i + dϕ_prev = dϕ_i + α_i = min(α_i * T(2), α_max) end end end if !done - α_out, ok = α_prev, true + α_out = α_prev + ok = false end - (α_out, ok) + return (α_out, ok) end +# closure built here, not in init, so it always reads cache.stats function CommonSolve.solve!(cache::StrongWolfeLineSearchCache, u, du) T = promote_type(eltype(du), eltype(u)) - ϕ = @closure α -> cache.ϕ(cache.f, cache.p, u, du, α, cache.u_cache, cache.fu_cache) - dϕ = @closure α -> cache.dϕ(cache.f, cache.p, u, du, α, - cache.u_cache, cache.fu_cache, cache.deriv_op) + ϕdϕ = @closure α -> begin + @bb @. cache.u_cache = u + α * du + cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + add_nf!(cache.stats) + obj = @fastmath norm(cache.fu_cache)^2 / 2 + deriv = cache.deriv_op(du, cache.u_cache, cache.fu_cache, cache.p) + return obj, deriv + end - ϕ_0 = ϕ(zero(T)) - dϕ_0 = dϕ(zero(T)) + ϕ_0, dϕ_0 = ϕdϕ(zero(T)) isfinite(ϕ_0) || return LineSearchSolution(cache.α, ReturnCode.Failure) dϕ_0 >= zero(T) && return LineSearchSolution(cache.α, ReturnCode.Failure) - α, ok = _sw_search(ϕ, dϕ, ϕ_0, dϕ_0, cache.c1, cache.c2, cache.α, cache.maxiters) + α, ok = _sw_search(ϕdϕ, ϕ_0, dϕ_0, cache.c1, cache.c2, cache.α, cache.maxiters) ok && return LineSearchSolution(α, ReturnCode.Success) return LineSearchSolution(cache.α, ReturnCode.Failure) @@ -321,8 +325,7 @@ end function CommonSolve.solve!(cache::StaticStrongWolfeLineSearchCache, u, du) T = promote_type(eltype(du), eltype(u)) - ϕ_0 = _sw_phi(cache.f, cache.p, u, du, zero(T)) - dϕ_0 = _sw_dphi(cache.grad_f, cache.p, u, du, zero(T)) + ϕ_0, dϕ_0 = _sw_phi_dphi(cache.f, cache.grad_f, cache.p, u, du, zero(T)) isfinite(ϕ_0) || return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) dϕ_0 >= zero(T) && return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) From d0e1be0301d19d6182d80dea98008618bf41e535 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sun, 5 Apr 2026 20:34:57 +0530 Subject: [PATCH 3/5] runic formatting Signed-off-by: AdityaPandeyCN --- src/strong_wolfe.jl | 69 ++++++++++++++++++++++------------ test/custom_optimizer_tests.jl | 2 +- test/root_finding_tests.jl | 6 +-- 3 files changed, 48 insertions(+), 29 deletions(-) diff --git a/src/strong_wolfe.jl b/src/strong_wolfe.jl index d44b08c..64ee2b2 100644 --- a/src/strong_wolfe.jl +++ b/src/strong_wolfe.jl @@ -12,7 +12,7 @@ available. """ @kwdef @concrete struct StrongWolfeLineSearch <: AbstractLineSearchAlgorithm autodiff = nothing - c1 = 1e-4 + c1 = 1.0e-4 c2 = 0.9 α_init = 1.0 maxiters::Int = 10 @@ -33,7 +33,7 @@ end alg <: StrongWolfeLineSearch end -# GPU path: static cache for SArray/Number, no allocations +# GPU path: static cache for SArray/Number @concrete struct StaticStrongWolfeLineSearchCache <: AbstractLineSearchCache f grad_f @@ -92,7 +92,7 @@ end d1 = dϕ_lo + dϕ_hi - 3 * (ϕ_lo - ϕ_hi) / (a_lo - a_hi) desc = d1 * d1 - dϕ_lo * dϕ_hi d2 = sqrt(max(desc, zero(desc))) - ifelse( + return ifelse( desc < 0, (a_lo + a_hi) / 2, a_hi - (a_hi - a_lo) * ((dϕ_hi + d2 - d1) / (dϕ_hi - dϕ_lo + 2 * d2)) @@ -110,9 +110,10 @@ end end # N&W Algorithm 3.6 (zoom) for static path -# uses `done` flag instead of early return for GPU kernel compatibility -@inline function _sw_zoom_static(f, grad_f, p, u, du, - a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, dϕ_lo, c1, c2, maxiters) +@inline function _sw_zoom_static( + f, grad_f, p, u, du, + a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, dϕ_lo, c1, c2, maxiters + ) T = typeof(a_lo) α_out = a_lo ok = false @@ -125,9 +126,11 @@ end if !done α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) bracket = T(0.01) * abs(a_hi - a_lo) - α_j = clamp(α_j, + α_j = clamp( + α_j, min(a_lo, a_hi) + bracket, - max(a_lo, a_hi) - bracket) + max(a_lo, a_hi) - bracket + ) ϕ_j, dϕ_j = _sw_phi_dphi(f, grad_f, p, u, du, α_j) @@ -159,8 +162,10 @@ end end # N&W Algorithm 3.5 (outer search) for static path -@inline function _sw_search_static(f, grad_f, p, u, du, ϕ_0, dϕ_0, - c1, c2, α_init, maxiters) +@inline function _sw_search_static( + f, grad_f, p, u, du, ϕ_0, dϕ_0, + c1, c2, α_init, maxiters + ) T = typeof(α_init) α_max = T(65536) @@ -179,16 +184,20 @@ end ϕ_i, dϕ_i = _sw_phi_dphi(f, grad_f, p, u, du, α_i) if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) - α_z, ok_z = _sw_zoom_static(f, grad_f, p, u, du, + α_z, ok_z = _sw_zoom_static( + f, grad_f, p, u, du, α_prev, α_i, ϕ_0, dϕ_0, ϕ_prev, dϕ_prev, - c1, c2, maxiters) + c1, c2, maxiters + ) α_out, ok, done = α_z, ok_z, true elseif abs(dϕ_i) <= -c2 * dϕ_0 α_out, ok, done = α_i, true, true elseif dϕ_i >= zero(T) - α_z, ok_z = _sw_zoom_static(f, grad_f, p, u, du, + α_z, ok_z = _sw_zoom_static( + f, grad_f, p, u, du, α_i, α_prev, ϕ_0, dϕ_0, ϕ_i, dϕ_i, - c1, c2, maxiters) + c1, c2, maxiters + ) α_out, ok, done = α_z, ok_z, true else α_prev = α_i @@ -207,8 +216,10 @@ end end # N&W Algorithm 3.6 (zoom) for mutable path -@inline function _sw_zoom(ϕdϕ, a_lo, a_hi, ϕ_0, dϕ_0, - ϕ_lo, dϕ_lo, c1, c2, maxiters) +@inline function _sw_zoom( + ϕdϕ, a_lo, a_hi, ϕ_0, dϕ_0, + ϕ_lo, dϕ_lo, c1, c2, maxiters + ) T = typeof(a_lo) α_out = a_lo ok = false @@ -220,9 +231,11 @@ end if !done α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) bracket = T(0.01) * abs(a_hi - a_lo) - α_j = clamp(α_j, + α_j = clamp( + α_j, min(a_lo, a_hi) + bracket, - max(a_lo, a_hi) - bracket) + max(a_lo, a_hi) - bracket + ) ϕ_j, dϕ_j = ϕdϕ(α_j) @@ -273,14 +286,18 @@ end ϕ_i, dϕ_i = ϕdϕ(α_i) if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) - α_z, ok_z = _sw_zoom(ϕdϕ, α_prev, α_i, - ϕ_0, dϕ_0, ϕ_prev, dϕ_prev, c1, c2, maxiters) + α_z, ok_z = _sw_zoom( + ϕdϕ, α_prev, α_i, + ϕ_0, dϕ_0, ϕ_prev, dϕ_prev, c1, c2, maxiters + ) α_out, ok, done = α_z, ok_z, true elseif abs(dϕ_i) <= -c2 * dϕ_0 α_out, ok, done = α_i, true, true elseif dϕ_i >= zero(T) - α_z, ok_z = _sw_zoom(ϕdϕ, α_i, α_prev, - ϕ_0, dϕ_0, ϕ_i, dϕ_i, c1, c2, maxiters) + α_z, ok_z = _sw_zoom( + ϕdϕ, α_i, α_prev, + ϕ_0, dϕ_0, ϕ_i, dϕ_i, c1, c2, maxiters + ) α_out, ok, done = α_z, ok_z, true else α_prev = α_i @@ -330,8 +347,10 @@ function CommonSolve.solve!(cache::StaticStrongWolfeLineSearchCache, u, du) isfinite(ϕ_0) || return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) dϕ_0 >= zero(T) && return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) - α, ok = _sw_search_static(cache.f, cache.grad_f, cache.p, u, du, - ϕ_0, dϕ_0, cache.c1, cache.c2, T(cache.α_init), cache.maxiters) + α, ok = _sw_search_static( + cache.f, cache.grad_f, cache.p, u, du, + ϕ_0, dϕ_0, cache.c1, cache.c2, T(cache.α_init), cache.maxiters + ) ok && return LineSearchSolution(α, ReturnCode.Success) return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) @@ -346,4 +365,4 @@ function SciMLBase.reinit!( cache.c1 = oftype(cache.c1, cache.alg.c1) cache.c2 = oftype(cache.c2, cache.alg.c2) return cache -end \ No newline at end of file +end diff --git a/test/custom_optimizer_tests.jl b/test/custom_optimizer_tests.jl index b831b33..c9e2f2a 100644 --- a/test/custom_optimizer_tests.jl +++ b/test/custom_optimizer_tests.jl @@ -154,4 +154,4 @@ end end end end -end \ No newline at end of file +end diff --git a/test/root_finding_tests.jl b/test/root_finding_tests.jl index 1918a1b..d1d195c 100644 --- a/test/root_finding_tests.jl +++ b/test/root_finding_tests.jl @@ -140,7 +140,7 @@ end NoLineSearch(0.5), GoldenSection(; tol = 1.0e-4), RobustNonMonotoneLineSearch(), - RobustNonMonotoneLineSearch(; M = 1), + RobustNonMonotoneLineSearch(; M = 1), #strictly monotonous case RobustNonMonotoneLineSearch(; M = 15), ) converged, fu, u, iter, alphas = newton_raphson(nlp, method) @@ -175,7 +175,7 @@ end NoLineSearch(0.5), GoldenSection(; tol = 1.0e-4), RobustNonMonotoneLineSearch(), - RobustNonMonotoneLineSearch(; M = 1), + RobustNonMonotoneLineSearch(; M = 1), #strictly monotonous case RobustNonMonotoneLineSearch(; M = 15), ) converged, fu, u, iter, alphas = newton_raphson(nlp, method) @@ -199,4 +199,4 @@ end end end end -end \ No newline at end of file +end From 038d7b99e50ee3ef6d2b285e4ec657c052626700 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 17 Apr 2026 21:43:36 +0530 Subject: [PATCH 4/5] refactor(StrongWolfe): unify CPU/GPU paths, fix redundant zoom eval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Thread (ϕ_hi, dϕ_hi) through _sw_zoom to skip duplicate eval at entry; separate zoom_maxiters from maxiters; promote α_max to alg field. Signed-off-by: AdityaPandeyCN --- src/strong_wolfe.jl | 224 ++++++++++++-------------------------------- 1 file changed, 61 insertions(+), 163 deletions(-) diff --git a/src/strong_wolfe.jl b/src/strong_wolfe.jl index 64ee2b2..afa1d83 100644 --- a/src/strong_wolfe.jl +++ b/src/strong_wolfe.jl @@ -1,6 +1,7 @@ """ StrongWolfeLineSearch(; autodiff = nothing, c1 = 1e-4, c2 = 0.9, - α_init = 1.0, maxiters::Int = 10) + α_init = 1.0, α_max = 65536.0, maxiters::Int = 10, + zoom_maxiters::Int = 10) Strong Wolfe line search satisfying both Armijo (sufficient decrease) and curvature conditions. Based on Nocedal & Wright, "Numerical Optimization" (2006), @@ -9,16 +10,20 @@ Algorithms 3.5 and 3.6. `autodiff` is the automatic differentiation backend to use for computing the directional derivative. Must be specified if analytic jacobian/jvp/vjp is not available. + +`maxiters` bounds the outer bracketing loop (Alg. 3.5). `zoom_maxiters` bounds +the inner zoom loop (Alg. 3.6) independently. """ @kwdef @concrete struct StrongWolfeLineSearch <: AbstractLineSearchAlgorithm autodiff = nothing c1 = 1.0e-4 c2 = 0.9 α_init = 1.0 + α_max = 65536.0 maxiters::Int = 10 + zoom_maxiters::Int = 10 end -# CPU path: stores raw ingredients, closure built fresh in solve! @concrete mutable struct StrongWolfeLineSearchCache <: AbstractLineSearchCache f p @@ -28,12 +33,13 @@ end c1 c2 α + α_max maxiters::Int + zoom_maxiters::Int stats <: Union{SciMLBase.NLStats, Nothing} alg <: StrongWolfeLineSearch end -# GPU path: static cache for SArray/Number @concrete struct StaticStrongWolfeLineSearchCache <: AbstractLineSearchCache f grad_f @@ -41,25 +47,25 @@ end c1 c2 α_init + α_max maxiters::Int + zoom_maxiters::Int end function CommonSolve.init( prob::AbstractNonlinearProblem, alg::StrongWolfeLineSearch, fu::Union{SArray, Number}, u::Union{SArray, Number}; - stats::Union{SciMLBase.NLStats, Nothing} = nothing, grad_f = nothing, kwargs... ) - if stats === nothing - grad_f === nothing && error( - "StrongWolfeLineSearch requires `grad_f` for static (GPU) dispatch" - ) - T = promote_type(eltype(fu), eltype(u)) - return StaticStrongWolfeLineSearchCache( - prob.f, grad_f, prob.p, T(alg.c1), T(alg.c2), T(alg.α_init), alg.maxiters - ) - end - return generic_strongwolfe_init(prob, alg, fu, u; stats, kwargs...) + grad_f === nothing && error( + "StrongWolfeLineSearch requires `grad_f` for static (GPU) dispatch" + ) + T = promote_type(eltype(fu), eltype(u)) + return StaticStrongWolfeLineSearchCache( + prob.f, grad_f, prob.p, + T(alg.c1), T(alg.c2), T(alg.α_init), T(alg.α_max), + alg.maxiters, alg.zoom_maxiters + ) end function CommonSolve.init( @@ -76,157 +82,52 @@ function generic_strongwolfe_init( ) autodiff = autodiff !== nothing ? autodiff : alg.autodiff _, _, deriv_op = construct_jvp_or_vjp_operator(prob, fu, u; autodiff) - @bb u_cache = similar(u) @bb fu_cache = similar(fu) T = promote_type(eltype(fu), eltype(u)) - return StrongWolfeLineSearchCache( prob.f, prob.p, deriv_op, u_cache, fu_cache, - T(alg.c1), T(alg.c2), T(alg.α_init), alg.maxiters, stats, alg + T(alg.c1), T(alg.c2), T(alg.α_init), T(alg.α_max), + alg.maxiters, alg.zoom_maxiters, stats, alg ) end -# cubic interpolation with bisection fallback when discriminant is negative @inline function _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) d1 = dϕ_lo + dϕ_hi - 3 * (ϕ_lo - ϕ_hi) / (a_lo - a_hi) desc = d1 * d1 - dϕ_lo * dϕ_hi d2 = sqrt(max(desc, zero(desc))) - return ifelse( - desc < 0, - (a_lo + a_hi) / 2, - a_hi - (a_hi - a_lo) * ((dϕ_hi + d2 - d1) / (dϕ_hi - dϕ_lo + 2 * d2)) - ) -end - -# combined f + grad_f eval for static/GPU path -@inline function _sw_phi_dphi(f, grad_f, p, u, du, α) - u_new = u .+ α .* du - fu = f(u_new, p) - ϕ = @fastmath norm(fu)^2 / 2 - g = grad_f(u_new, p) - dϕ = dot(g, du) - return ϕ, dϕ + candidate = a_hi - (a_hi - a_lo) * ((dϕ_hi + d2 - d1) / (dϕ_hi - dϕ_lo + 2 * d2)) + candidate = ifelse(isfinite(candidate), candidate, (a_lo + a_hi) / 2) + return ifelse(desc < 0, (a_lo + a_hi) / 2, candidate) end -# N&W Algorithm 3.6 (zoom) for static path -@inline function _sw_zoom_static( - f, grad_f, p, u, du, - a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, dϕ_lo, c1, c2, maxiters - ) - T = typeof(a_lo) - α_out = a_lo - ok = false - done = false - - # evaluate a_hi once, then maintain via updates below - ϕ_hi, dϕ_hi = _sw_phi_dphi(f, grad_f, p, u, du, a_hi) - - for _ in 1:maxiters - if !done - α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) - bracket = T(0.01) * abs(a_hi - a_lo) - α_j = clamp( - α_j, - min(a_lo, a_hi) + bracket, - max(a_lo, a_hi) - bracket - ) - - ϕ_j, dϕ_j = _sw_phi_dphi(f, grad_f, p, u, du, α_j) - - if (ϕ_j > ϕ_0 + c1 * α_j * dϕ_0) || (ϕ_j >= ϕ_lo) - a_hi = α_j - ϕ_hi = ϕ_j - dϕ_hi = dϕ_j - else - if abs(dϕ_j) <= -c2 * dϕ_0 - α_out, ok, done = α_j, true, true - else - if dϕ_j * (a_hi - a_lo) >= zero(T) - a_hi = a_lo - ϕ_hi = ϕ_lo - dϕ_hi = dϕ_lo - end - a_lo = α_j - ϕ_lo = ϕ_j - dϕ_lo = dϕ_j - end - end - end - end - - if !done - α_out = a_lo - end - return (α_out, ok) +struct _SWNonlinearEval{F, G, P, U, D} + f::F + grad_f::G + p::P + u::U + du::D end -# N&W Algorithm 3.5 (outer search) for static path -@inline function _sw_search_static( - f, grad_f, p, u, du, ϕ_0, dϕ_0, - c1, c2, α_init, maxiters - ) - T = typeof(α_init) - α_max = T(65536) - - dϕ_0 >= zero(T) && return (zero(T), false) - - α_prev = zero(T) - α_i = α_init - ϕ_prev = ϕ_0 - dϕ_prev = dϕ_0 - done = false - α_out = zero(T) - ok = false - - for i in 1:maxiters - if !done - ϕ_i, dϕ_i = _sw_phi_dphi(f, grad_f, p, u, du, α_i) - - if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) - α_z, ok_z = _sw_zoom_static( - f, grad_f, p, u, du, - α_prev, α_i, ϕ_0, dϕ_0, ϕ_prev, dϕ_prev, - c1, c2, maxiters - ) - α_out, ok, done = α_z, ok_z, true - elseif abs(dϕ_i) <= -c2 * dϕ_0 - α_out, ok, done = α_i, true, true - elseif dϕ_i >= zero(T) - α_z, ok_z = _sw_zoom_static( - f, grad_f, p, u, du, - α_i, α_prev, ϕ_0, dϕ_0, ϕ_i, dϕ_i, - c1, c2, maxiters - ) - α_out, ok, done = α_z, ok_z, true - else - α_prev = α_i - ϕ_prev = ϕ_i - dϕ_prev = dϕ_i - α_i = min(α_i * T(2), α_max) - end - end - end - - if !done - α_out = α_prev - ok = false - end - return (α_out, ok) +@inline function (e::_SWNonlinearEval)(α) + u_new = e.u .+ α .* e.du + fu = e.f(u_new, e.p) + ϕ = sum(abs2, fu) / 2 + dϕ = dot(e.grad_f(u_new, e.p), e.du) + return (ϕ, dϕ) end -# N&W Algorithm 3.6 (zoom) for mutable path +# N&W Algorithm 3.6. Fixed iteration count with `!done` guard is intentional +# for GPU warp-uniform execution; do not rewrite as early `break`. @inline function _sw_zoom( - ϕdϕ, a_lo, a_hi, ϕ_0, dϕ_0, - ϕ_lo, dϕ_lo, c1, c2, maxiters + eval_fn, a_lo, a_hi, ϕ_0, dϕ_0, + ϕ_lo, dϕ_lo, ϕ_hi, dϕ_hi, c1, c2, maxiters ) T = typeof(a_lo) α_out = a_lo ok = false done = false - ϕ_hi, dϕ_hi = ϕdϕ(a_hi) - for _ in 1:maxiters if !done α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) @@ -236,8 +137,7 @@ end min(a_lo, a_hi) + bracket, max(a_lo, a_hi) - bracket ) - - ϕ_j, dϕ_j = ϕdϕ(α_j) + ϕ_j, dϕ_j = eval_fn(α_j) if (ϕ_j > ϕ_0 + c1 * α_j * dϕ_0) || (ϕ_j >= ϕ_lo) a_hi = α_j @@ -259,17 +159,15 @@ end end end end - if !done α_out = a_lo end return (α_out, ok) end -# N&W Algorithm 3.5 (outer search) for mutable path -@inline function _sw_search(ϕdϕ, ϕ_0, dϕ_0, c1, c2, α_init, maxiters) +# N&W Algorithm 3.5. +@inline function _sw_search(eval_fn, ϕ_0, dϕ_0, c1, c2, α_init, α_max, maxiters, zoom_maxiters) T = typeof(α_init) - α_max = T(65536) dϕ_0 >= zero(T) && return (zero(T), false) @@ -283,20 +181,20 @@ end for i in 1:maxiters if !done - ϕ_i, dϕ_i = ϕdϕ(α_i) + ϕ_i, dϕ_i = eval_fn(α_i) if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) α_z, ok_z = _sw_zoom( - ϕdϕ, α_prev, α_i, - ϕ_0, dϕ_0, ϕ_prev, dϕ_prev, c1, c2, maxiters + eval_fn, α_prev, α_i, ϕ_0, dϕ_0, + ϕ_prev, dϕ_prev, ϕ_i, dϕ_i, c1, c2, zoom_maxiters ) α_out, ok, done = α_z, ok_z, true elseif abs(dϕ_i) <= -c2 * dϕ_0 α_out, ok, done = α_i, true, true elseif dϕ_i >= zero(T) α_z, ok_z = _sw_zoom( - ϕdϕ, α_i, α_prev, - ϕ_0, dϕ_0, ϕ_i, dϕ_i, c1, c2, maxiters + eval_fn, α_i, α_prev, ϕ_0, dϕ_0, + ϕ_i, dϕ_i, ϕ_prev, dϕ_prev, c1, c2, zoom_maxiters ) α_out, ok, done = α_z, ok_z, true else @@ -307,15 +205,13 @@ end end end end - if !done - α_out = α_prev + α_out = α_i ok = false end return (α_out, ok) end -# closure built here, not in init, so it always reads cache.stats function CommonSolve.solve!(cache::StrongWolfeLineSearchCache, u, du) T = promote_type(eltype(du), eltype(u)) @@ -323,18 +219,19 @@ function CommonSolve.solve!(cache::StrongWolfeLineSearchCache, u, du) @bb @. cache.u_cache = u + α * du cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) add_nf!(cache.stats) - obj = @fastmath norm(cache.fu_cache)^2 / 2 + obj = sum(abs2, cache.fu_cache) / 2 deriv = cache.deriv_op(du, cache.u_cache, cache.fu_cache, cache.p) return obj, deriv end ϕ_0, dϕ_0 = ϕdϕ(zero(T)) - isfinite(ϕ_0) || return LineSearchSolution(cache.α, ReturnCode.Failure) dϕ_0 >= zero(T) && return LineSearchSolution(cache.α, ReturnCode.Failure) - α, ok = _sw_search(ϕdϕ, ϕ_0, dϕ_0, cache.c1, cache.c2, cache.α, cache.maxiters) - + α, ok = _sw_search( + ϕdϕ, ϕ_0, dϕ_0, cache.c1, cache.c2, + cache.α, cache.α_max, cache.maxiters, cache.zoom_maxiters + ) ok && return LineSearchSolution(α, ReturnCode.Success) return LineSearchSolution(cache.α, ReturnCode.Failure) end @@ -342,16 +239,16 @@ end function CommonSolve.solve!(cache::StaticStrongWolfeLineSearchCache, u, du) T = promote_type(eltype(du), eltype(u)) - ϕ_0, dϕ_0 = _sw_phi_dphi(cache.f, cache.grad_f, cache.p, u, du, zero(T)) + eval_fn = _SWNonlinearEval(cache.f, cache.grad_f, cache.p, u, du) + ϕ_0, dϕ_0 = eval_fn(zero(T)) isfinite(ϕ_0) || return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) dϕ_0 >= zero(T) && return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) - α, ok = _sw_search_static( - cache.f, cache.grad_f, cache.p, u, du, - ϕ_0, dϕ_0, cache.c1, cache.c2, T(cache.α_init), cache.maxiters + α, ok = _sw_search( + eval_fn, ϕ_0, dϕ_0, cache.c1, cache.c2, + T(cache.α_init), T(cache.α_max), cache.maxiters, cache.zoom_maxiters ) - ok && return LineSearchSolution(α, ReturnCode.Success) return LineSearchSolution(T(cache.α_init), ReturnCode.Failure) end @@ -364,5 +261,6 @@ function SciMLBase.reinit!( cache.α = oftype(cache.α, cache.alg.α_init) cache.c1 = oftype(cache.c1, cache.alg.c1) cache.c2 = oftype(cache.c2, cache.alg.c2) + cache.α_max = oftype(cache.α_max, cache.alg.α_max) return cache end From 07dfc04550faddde9a5bc01582f1917933cbe74b Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Tue, 21 Apr 2026 16:01:34 +0530 Subject: [PATCH 5/5] test(StrongWolfe): add edge-case coverage and clarify loop guards Document the fixed-iteration loop structure used for GPU-uniform execution and add focused edge-case tests so the review concerns are covered explicitly. Made-with: Cursor --- src/strong_wolfe.jl | 99 +++++++++++++++++++------------------- test/root_finding_tests.jl | 55 +++++++++++++++++++++ 2 files changed, 105 insertions(+), 49 deletions(-) diff --git a/src/strong_wolfe.jl b/src/strong_wolfe.jl index afa1d83..e877e13 100644 --- a/src/strong_wolfe.jl +++ b/src/strong_wolfe.jl @@ -117,8 +117,8 @@ end return (ϕ, dϕ) end -# N&W Algorithm 3.6. Fixed iteration count with `!done` guard is intentional -# for GPU warp-uniform execution; do not rewrite as early `break`. +# N&W Algorithm 3.6. Keep the fixed iteration count for GPU warp-uniform +# execution; once `done`, later iterations just skip the body. @inline function _sw_zoom( eval_fn, a_lo, a_hi, ϕ_0, dϕ_0, ϕ_lo, dϕ_lo, ϕ_hi, dϕ_hi, c1, c2, maxiters @@ -129,33 +129,33 @@ end done = false for _ in 1:maxiters - if !done - α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) - bracket = T(0.01) * abs(a_hi - a_lo) - α_j = clamp( - α_j, - min(a_lo, a_hi) + bracket, - max(a_lo, a_hi) - bracket - ) - ϕ_j, dϕ_j = eval_fn(α_j) + done && continue + + α_j = _sw_interpolate(a_lo, a_hi, ϕ_lo, ϕ_hi, dϕ_lo, dϕ_hi) + bracket = T(0.01) * abs(a_hi - a_lo) + α_j = clamp( + α_j, + min(a_lo, a_hi) + bracket, + max(a_lo, a_hi) - bracket + ) + ϕ_j, dϕ_j = eval_fn(α_j) - if (ϕ_j > ϕ_0 + c1 * α_j * dϕ_0) || (ϕ_j >= ϕ_lo) - a_hi = α_j - ϕ_hi = ϕ_j - dϕ_hi = dϕ_j + if (ϕ_j > ϕ_0 + c1 * α_j * dϕ_0) || (ϕ_j >= ϕ_lo) + a_hi = α_j + ϕ_hi = ϕ_j + dϕ_hi = dϕ_j + else + if abs(dϕ_j) <= -c2 * dϕ_0 + α_out, ok, done = α_j, true, true else - if abs(dϕ_j) <= -c2 * dϕ_0 - α_out, ok, done = α_j, true, true - else - if dϕ_j * (a_hi - a_lo) >= zero(T) - a_hi = a_lo - ϕ_hi = ϕ_lo - dϕ_hi = dϕ_lo - end - a_lo = α_j - ϕ_lo = ϕ_j - dϕ_lo = dϕ_j + if dϕ_j * (a_hi - a_lo) >= zero(T) + a_hi = a_lo + ϕ_hi = ϕ_lo + dϕ_hi = dϕ_lo end + a_lo = α_j + ϕ_lo = ϕ_j + dϕ_lo = dϕ_j end end end @@ -165,7 +165,8 @@ end return (α_out, ok) end -# N&W Algorithm 3.5. +# N&W Algorithm 3.5. Keep the fixed iteration count for the same GPU +# warp-uniform reason as `_sw_zoom`; once `done`, later iterations skip work. @inline function _sw_search(eval_fn, ϕ_0, dϕ_0, c1, c2, α_init, α_max, maxiters, zoom_maxiters) T = typeof(α_init) @@ -180,29 +181,29 @@ end ok = false for i in 1:maxiters - if !done - ϕ_i, dϕ_i = eval_fn(α_i) + done && continue - if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) - α_z, ok_z = _sw_zoom( - eval_fn, α_prev, α_i, ϕ_0, dϕ_0, - ϕ_prev, dϕ_prev, ϕ_i, dϕ_i, c1, c2, zoom_maxiters - ) - α_out, ok, done = α_z, ok_z, true - elseif abs(dϕ_i) <= -c2 * dϕ_0 - α_out, ok, done = α_i, true, true - elseif dϕ_i >= zero(T) - α_z, ok_z = _sw_zoom( - eval_fn, α_i, α_prev, ϕ_0, dϕ_0, - ϕ_i, dϕ_i, ϕ_prev, dϕ_prev, c1, c2, zoom_maxiters - ) - α_out, ok, done = α_z, ok_z, true - else - α_prev = α_i - ϕ_prev = ϕ_i - dϕ_prev = dϕ_i - α_i = min(α_i * T(2), α_max) - end + ϕ_i, dϕ_i = eval_fn(α_i) + + if (ϕ_i > ϕ_0 + c1 * α_i * dϕ_0) || (ϕ_i >= ϕ_prev && i > 1) + α_z, ok_z = _sw_zoom( + eval_fn, α_prev, α_i, ϕ_0, dϕ_0, + ϕ_prev, dϕ_prev, ϕ_i, dϕ_i, c1, c2, zoom_maxiters + ) + α_out, ok, done = α_z, ok_z, true + elseif abs(dϕ_i) <= -c2 * dϕ_0 + α_out, ok, done = α_i, true, true + elseif dϕ_i >= zero(T) + α_z, ok_z = _sw_zoom( + eval_fn, α_i, α_prev, ϕ_0, dϕ_0, + ϕ_i, dϕ_i, ϕ_prev, dϕ_prev, c1, c2, zoom_maxiters + ) + α_out, ok, done = α_z, ok_z, true + else + α_prev = α_i + ϕ_prev = ϕ_i + dϕ_prev = dϕ_i + α_i = min(α_i * T(2), α_max) end end if !done diff --git a/test/root_finding_tests.jl b/test/root_finding_tests.jl index d1d195c..692bb21 100644 --- a/test/root_finding_tests.jl +++ b/test/root_finding_tests.jl @@ -200,3 +200,58 @@ end end end end + +@testitem "Native Strong Wolfe edge cases" tags = [:core] begin + quadratic_eval(α) = (0.5 * (α - 1.0)^2, α - 1.0) + + @testset "initial convergence" begin + ϕ_0, dϕ_0 = quadratic_eval(0.0) + α, ok = LineSearch._sw_search(quadratic_eval, ϕ_0, dϕ_0, 1.0e-4, 0.9, 1.0, 4.0, 10, 10) + + @test ok + @test α ≈ 1.0 + end + + @testset "initial need for bracketing" begin + ϕ_0, dϕ_0 = quadratic_eval(0.0) + α, ok = LineSearch._sw_search( + quadratic_eval, ϕ_0, dϕ_0, 1.0e-4, 0.1, 0.25, 4.0, 10, 10 + ) + + @test ok + @test α ≈ 1.0 + @test α > 0.25 + end + + @testset "initial point is on the upward slope" begin + uphill_eval(α) = (0.5 * (α + 1.0)^2, α + 1.0) + ϕ_0, dϕ_0 = uphill_eval(0.0) + α, ok = LineSearch._sw_search(uphill_eval, ϕ_0, dϕ_0, 1.0e-4, 0.9, 1.0, 4.0, 10, 10) + + @test !ok + @test α == 0.0 + end + + @testset "initial trial has already passed the minimum" begin + ϕ_0, dϕ_0 = quadratic_eval(0.0) + α, ok = LineSearch._sw_search( + quadratic_eval, ϕ_0, dϕ_0, 1.0e-4, 0.1, 3.0, 4.0, 10, 10 + ) + + @test ok + @test α ≈ 1.0 + @test α < 3.0 + end + + @testset "nonfinite trial values" begin + nonfinite_eval(α) = α > 2.0 ? (Inf, Inf) : quadratic_eval(α) + ϕ_0, dϕ_0 = nonfinite_eval(0.0) + α, ok = LineSearch._sw_search( + nonfinite_eval, ϕ_0, dϕ_0, 1.0e-4, 0.1, 3.0, 4.0, 10, 10 + ) + + @test ok + @test isfinite(α) + @test 0.0 < α <= 2.0 + end +end