refactor: migrate to LineSearch.jl

avik-pal · avik-pal · commit 40c4df9c8451 · 2024-09-26T05:03:16.000-04:00
diff --git a/Project.toml b/Project.toml
@@ -13,6 +13,7 @@ FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
+LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
@@ -78,6 +79,7 @@ Hwloc = "3"
 InteractiveUtils = "<0.0.1, 1"
 LazyArrays = "1.8.2, 2"
 LeastSquaresOptim = "0.8.5"
+LineSearch = "0.1"
 LineSearches = "7.2"
 LinearAlgebra = "1.10"
 LinearSolve = "2.30"
diff --git a/docs/src/devdocs/internal_interfaces.md b/docs/src/devdocs/internal_interfaces.md
@@ -38,13 +38,6 @@ NonlinearSolve.AbstractDampingFunction
 NonlinearSolve.AbstractDampingFunctionCache
 ```
 
-## Line Search
-
-```@docs
-NonlinearSolve.AbstractNonlinearSolveLineSearchAlgorithm
-NonlinearSolve.AbstractNonlinearSolveLineSearchCache
-```
-
 ## Trust Region
 
 ```@docs
diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -30,6 +30,8 @@ using LazyArrays: LazyArrays, ApplyArray, cache
 using LinearAlgebra: LinearAlgebra, ColumnNorm, Diagonal, I, LowerTriangular, Symmetric,
                      UpperTriangular, axpy!, cond, diag, diagind, dot, issuccess, istril,
                      istriu, lu, mul!, norm, pinv, tril!, triu!
+using LineSearch: LineSearch, AbstractLineSearchAlgorithm, AbstractLineSearchCache,
+                  NoLineSearch
 using LineSearches: LineSearches
 using LinearSolve: LinearSolve, LUFactorization, QRFactorization, ComposePreconditioner,
                    InvPreconditioner, needs_concrete_A, AbstractFactorization,
@@ -103,54 +105,54 @@ include("algorithms/extension_algs.jl")
 include("utils.jl")
 include("default.jl")
 
-@setup_workload begin
-    nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1),
-        (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1]))
-    probs_nls = NonlinearProblem[]
-    for (fn, u0) in nlfuncs
-        push!(probs_nls, NonlinearProblem(fn, u0, 2.0))
-    end
-
-    nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(),
-        PseudoTransient(), Broyden(), Klement(), DFSane(), nothing)
-
-    probs_nlls = NonlinearLeastSquaresProblem[]
-    nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]),
-        (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]),
-        (
-            NonlinearFunction{true}(
-                (du, u, p) -> du[1] = u[1] * u[1] - p, resid_prototype = zeros(1)),
-            [0.1, 0.0]),
-        (
-            NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p),
-                resid_prototype = zeros(4)),
-            [0.1, 0.1]))
-    for (fn, u0) in nlfuncs
-        push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0))
-    end
-
-    nlls_algs = (LevenbergMarquardt(), GaussNewton(), TrustRegion(),
-        LevenbergMarquardt(; linsolve = LUFactorization()),
-        GaussNewton(; linsolve = LUFactorization()),
-        TrustRegion(; linsolve = LUFactorization()), nothing)
-
-    @compile_workload begin
-        @sync begin
-            for T in (Float32, Float64), (fn, u0) in nlfuncs
-                Threads.@spawn NonlinearProblem(fn, T.(u0), T(2))
-            end
-            for (fn, u0) in nlfuncs
-                Threads.@spawn NonlinearLeastSquaresProblem(fn, u0, 2.0)
-            end
-            for prob in probs_nls, alg in nls_algs
-                Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
-            end
-            for prob in probs_nlls, alg in nlls_algs
-                Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
-            end
-        end
-    end
-end
+# @setup_workload begin
+#     nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1),
+#         (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1]))
+#     probs_nls = NonlinearProblem[]
+#     for (fn, u0) in nlfuncs
+#         push!(probs_nls, NonlinearProblem(fn, u0, 2.0))
+#     end
+
+#     nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(),
+#         PseudoTransient(), Broyden(), Klement(), DFSane(), nothing)
+
+#     probs_nlls = NonlinearLeastSquaresProblem[]
+#     nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]),
+#         (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]),
+#         (
+#             NonlinearFunction{true}(
+#                 (du, u, p) -> du[1] = u[1] * u[1] - p, resid_prototype = zeros(1)),
+#             [0.1, 0.0]),
+#         (
+#             NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p),
+#                 resid_prototype = zeros(4)),
+#             [0.1, 0.1]))
+#     for (fn, u0) in nlfuncs
+#         push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0))
+#     end
+
+#     nlls_algs = (LevenbergMarquardt(), GaussNewton(), TrustRegion(),
+#         LevenbergMarquardt(; linsolve = LUFactorization()),
+#         GaussNewton(; linsolve = LUFactorization()),
+#         TrustRegion(; linsolve = LUFactorization()), nothing)
+
+#     @compile_workload begin
+#         @sync begin
+#             for T in (Float32, Float64), (fn, u0) in nlfuncs
+#                 Threads.@spawn NonlinearProblem(fn, T.(u0), T(2))
+#             end
+#             for (fn, u0) in nlfuncs
+#                 Threads.@spawn NonlinearLeastSquaresProblem(fn, u0, 2.0)
+#             end
+#             for prob in probs_nls, alg in nls_algs
+#                 Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
+#             end
+#             for prob in probs_nlls, alg in nlls_algs
+#                 Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
+#             end
+#         end
+#     end
+# end
 
 # Core Algorithms
 export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane
diff --git a/src/abstract_types.jl b/src/abstract_types.jl
@@ -106,22 +106,6 @@ function last_step_accepted(cache::AbstractDescentCache)
     return true
 end
 
-"""
-    AbstractNonlinearSolveLineSearchAlgorithm
-
-Abstract Type for all Line Search Algorithms used in NonlinearSolve.jl.
-
-### `__internal_init` specification
-
-```julia
-__internal_init(
-    prob::AbstractNonlinearProblem, alg::AbstractNonlinearSolveLineSearchAlgorithm, f::F,
-    fu, u, p, args...; internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN} -->
-AbstractNonlinearSolveLineSearchCache
-```
-"""
-abstract type AbstractNonlinearSolveLineSearchAlgorithm end
-
 """
     AbstractNonlinearSolveLineSearchCache
 
@@ -512,9 +496,9 @@ SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = ii
 abstract type AbstractNonlinearSolveTraceLevel end
 
 # Default Printing
-for aType in (AbstractTrustRegionMethod, AbstractNonlinearSolveLineSearchAlgorithm,
-    AbstractResetCondition, AbstractApproximateJacobianUpdateRule,
-    AbstractDampingFunction, AbstractNonlinearSolveExtensionAlgorithm)
+for aType in (AbstractTrustRegionMethod, AbstractResetCondition,
+    AbstractApproximateJacobianUpdateRule, AbstractDampingFunction,
+    AbstractNonlinearSolveExtensionAlgorithm)
     @eval function Base.show(io::IO, alg::$(aType))
         print(io, "$(nameof(typeof(alg)))()")
     end
diff --git a/src/algorithms/dfsane.jl b/src/algorithms/dfsane.jl
@@ -19,8 +19,9 @@ For other keyword arguments, see [`RobustNonMonotoneLineSearch`](@ref).
 function DFSane(; σ_min = 1 // 10^10, σ_max = 1e10, σ_1 = 1, M::Int = 10, γ = 1 // 10^4,
         τ_min = 1 // 10, τ_max = 1 // 2, n_exp::Int = 2, max_inner_iterations::Int = 100,
         η_strategy::ETA = (fn_1, n, x_n, f_n) -> fn_1 / n^2) where {ETA}
-    linesearch = RobustNonMonotoneLineSearch(;
-        gamma = γ, sigma_1 = σ_1, M, tau_min = τ_min, tau_max = τ_max,
-        n_exp, η_strategy, maxiters = max_inner_iterations)
+    # linesearch = RobustNonMonotoneLineSearch(;
+    #     gamma = γ, sigma_1 = σ_1, M, tau_min = τ_min, tau_max = τ_max,
+    #     n_exp, η_strategy, maxiters = max_inner_iterations)
+    linesearch = NoLineSearch()
     return GeneralizedDFSane{:DFSane}(linesearch, σ_min, σ_max, nothing)
 end
diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl
@@ -27,12 +27,12 @@ over this.
 function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = nothing,
         linesearch = NoLineSearch(), precs = DEFAULT_PRECS,
         autodiff = nothing, init_jacobian::Val{IJ} = Val(:identity)) where {IJ}
-    if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
-        Base.depwarn(
-            "Passing in a `LineSearches.jl` algorithm directly is deprecated. \
-             Please use `LineSearchesJL` instead.", :Klement)
-        linesearch = LineSearchesJL(; method = linesearch)
-    end
+    # if !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
+    #     Base.depwarn(
+    #         "Passing in a `LineSearches.jl` algorithm directly is deprecated. \
+    #          Please use `LineSearchesJL` instead.", :Klement)
+    #     linesearch = LineSearchesJL(; method = linesearch)
+    # end
 
     if IJ === :identity
         initialization = IdentityInitialization(alpha, DiagonalStructure())
diff --git a/src/algorithms/pseudo_transient.jl b/src/algorithms/pseudo_transient.jl
@@ -1,7 +1,6 @@
 """
     PseudoTransient(; concrete_jac = nothing, linsolve = nothing,
-        linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(),
-        precs = DEFAULT_PRECS, autodiff = nothing)
+        linesearch = NoLineSearch(), precs = DEFAULT_PRECS, 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
@@ -16,8 +15,8 @@ This implementation specifically uses "switched evolution relaxation"
     you are going to need more iterations to converge but it can be more stable.
 """
 function PseudoTransient(; concrete_jac = nothing, linsolve = nothing,
-        linesearch::AbstractNonlinearSolveLineSearchAlgorithm = NoLineSearch(),
-        precs = DEFAULT_PRECS, autodiff = nothing, alpha_initial = 1e-3)
+        linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing,
+        alpha_initial = 1e-3)
     descent = DampedNewtonDescent(; linsolve, precs, initial_damping = alpha_initial,
         damping_fn = SwitchedEvolutionRelaxation())
     return GeneralizedFirstOrderAlgorithm(;
diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl
@@ -59,12 +59,12 @@ function ApproximateJacobianSolveAlgorithm{concrete_jac, name}(;
         linesearch = missing, trustregion = missing, descent, update_rule,
         reinit_rule, initialization, max_resets::Int = typemax(Int),
         max_shrink_times::Int = typemax(Int)) where {concrete_jac, name}
-    if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
-        Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \
-                      Please use `LineSearchesJL` instead.",
-            :GeneralizedFirstOrderAlgorithm)
-        linesearch = LineSearchesJL(; method = linesearch)
-    end
+    # if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
+    #     Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \
+    #                   Please use `LineSearchesJL` instead.",
+    #         :GeneralizedFirstOrderAlgorithm)
+    #     linesearch = LineSearchesJL(; method = linesearch)
+    # end
     return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(
         linesearch, trustregion, descent, update_rule,
         reinit_rule, max_resets, max_shrink_times, initialization)
diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl
@@ -66,12 +66,12 @@ function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(;
             jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ReverseMode,
             jacobian_ad, nothing))
 
-    if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
-        Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \
-                      Please use `LineSearchesJL` instead.",
-            :GeneralizedFirstOrderAlgorithm)
-        linesearch = LineSearchesJL(; method = linesearch)
-    end
+    # if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
+    #     Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \
+    #                   Please use `LineSearchesJL` instead.",
+    #         :GeneralizedFirstOrderAlgorithm)
+    #     linesearch = LineSearchesJL(; method = linesearch)
+    # end
 
     return GeneralizedFirstOrderAlgorithm{concrete_jac, name}(
         linesearch, trustregion, descent, max_shrink_times,
@@ -199,8 +199,11 @@ function SciMLBase.__init(
         if alg.linesearch !== missing
             supports_line_search(alg.descent) || error("Line Search not supported by \
                                                         $(alg.descent).")
-            linesearch_cache = __internal_init(
-                prob, alg.linesearch, f, fu, u, p; stats, internalnorm, kwargs...)
+            linesearch_ad = alg.forward_ad === nothing ?
+                            (alg.reverse_ad === nothing ? alg.jacobian_ad :
+                             alg.reverse_ad) : alg.forward_ad
+            linesearch_cache = init(
+                prob, alg.linesearch, fu, u; stats, autodiff = linesearch_ad, kwargs...)
             GB = :LineSearch
         end
 
@@ -264,8 +267,9 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB};
         cache.make_new_jacobian = true
         if GB === :LineSearch
             @static_timeit cache.timer "linesearch" begin
-                linesearch_failed, α = __internal_solve!(
-                    cache.linesearch_cache, cache.u, δu)
+                linesearch_sol = solve!(cache.linesearch_cache, cache.u, δu)
+                linesearch_failed = !SciMLBase.successful_retcode(linesearch_sol.retcode)
+                α = linesearch_sol.step_size
             end
             if linesearch_failed
                 cache.retcode = ReturnCode.InternalLineSearchFailed
diff --git a/src/core/spectral_methods.jl b/src/core/spectral_methods.jl
@@ -9,9 +9,8 @@ Method.
 
 ### Arguments
 
-  - `linesearch`: Globalization using a Line Search Method. This needs to follow the
-    [`NonlinearSolve.AbstractNonlinearSolveLineSearchAlgorithm`](@ref) interface. This
-    is not optional currently, but that restriction might be lifted in the future.
+  - `linesearch`: Globalization using a Line Search Method. This is not optional currently,
+    but that restriction might be lifted in the future.
   - `σ_min`: The minimum spectral parameter allowed. This is used to ensure that the
     spectral parameter is not too small.
   - `σ_max`: The maximum spectral parameter allowed. This is used to ensure that the
diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl
