Merge pull request #390 from SciML/ap/fixes

Misc. Improvements
SciML · Mar 6, 2024 · e1c0528 · e1c0528 · avik-pal · Mar 6, 2024
2 parents 52b3832 + 81edf48
commit e1c0528
Show file tree

Hide file tree

Showing 15 changed files with 121 additions and 47 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.7.3"
+version = "3.8.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -9,8 +9,8 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
 
 @recompile_invalidations begin
     using ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures, LazyArrays,
-          LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf,
-          SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools
+          LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf, SciMLBase,
+          SimpleNonlinearSolve, SparseArrays, SparseDiffTools
 
     import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing
     import DiffEqBase: AbstractNonlinearTerminationMode,
@@ -20,6 +20,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
     import FiniteDiff
     import ForwardDiff
     import ForwardDiff: Dual
+    import LineSearches
     import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A
     import RecursiveArrayTools: recursivecopy!, recursivefill!
 
@@ -29,7 +30,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
     import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix
 end
 
-@reexport using ADTypes, LineSearches, SciMLBase, SimpleNonlinearSolve
+@reexport using ADTypes, SciMLBase, SimpleNonlinearSolve
 
 const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences,
     ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode}
@@ -157,6 +158,7 @@ export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, GeodesicAcce
 # Globalization
 ## Line Search Algorithms
 export LineSearchesJL, NoLineSearch, RobustNonMonotoneLineSearch, LiFukushimaLineSearch
+export Static, HagerZhang, MoreThuente, StrongWolfe, BackTracking
 ## Trust Region Algorithms
 export RadiusUpdateSchemes
 

diff --git a/src/algorithms/klement.jl b/src/algorithms/klement.jl
@@ -25,8 +25,8 @@ over this.
         differentiable problems.
 """
 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}
+        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. \

diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl
@@ -159,8 +159,8 @@ function LinearAlgebra.mul!(y::AbstractVector, x::AbstractVector, J::BroydenLowR
     return y
 end
 
-function LinearAlgebra.mul!(
-        J::BroydenLowRankJacobian, u, vᵀ::LinearAlgebra.AdjOrTransAbsVec, α::Bool, β::Bool)
+function LinearAlgebra.mul!(J::BroydenLowRankJacobian, u::AbstractArray,
+        vᵀ::LinearAlgebra.AdjOrTransAbsVec, α::Bool, β::Bool)
     @assert α & β
     idx_update = mod1(J.idx + 1, size(J.U, 2))
     copyto!(@view(J.U[:, idx_update]), _vec(u))

diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl
@@ -66,8 +66,8 @@ function ApproximateJacobianSolveAlgorithm{concrete_jac, name}(;
         linesearch = LineSearchesJL(; method = linesearch)
     end
     return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(
-        linesearch, trustregion, descent, update_rule, reinit_rule,
-        max_resets, max_shrink_times, initialization)
+        linesearch, trustregion, descent, update_rule,
+        reinit_rule, max_resets, max_shrink_times, initialization)
 end
 
 @inline concrete_jac(::ApproximateJacobianSolveAlgorithm{CJ}) where {CJ} = CJ

diff --git a/src/default.jl b/src/default.jl
@@ -1,6 +1,7 @@
 # Poly Algorithms
 """
-    NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType}
+    NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS);
+        start_index = 1) where {pType}
 
 A general way to define PolyAlgorithms for `NonlinearProblem` and
 `NonlinearLeastSquaresProblem`. This is a container for a tuple of algorithms that will be
@@ -15,6 +16,10 @@ residual is returned.
     `NonlinearLeastSquaresProblem`. This is used to determine the correct problem type to
     dispatch on.
 
+### Keyword Arguments
+
+  - `start_index`: the index to start at. Defaults to `1`.
+
 ### Example
 
 ```julia
@@ -25,11 +30,14 @@ alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden()))
 """
 struct NonlinearSolvePolyAlgorithm{pType, N, A} <: AbstractNonlinearSolveAlgorithm{:PolyAlg}
     algs::A
+    start_index::Int
 
-    function NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS)) where {pType}
+    function NonlinearSolvePolyAlgorithm(
+            algs, ::Val{pType} = Val(:NLS); start_index::Int = 1) where {pType}
         @assert pType ∈ (:NLS, :NLLS)
+        @assert 0 < start_index ≤ length(algs)
         algs = Tuple(algs)
-        return new{pType, length(algs), typeof(algs)}(algs)
+        return new{pType, length(algs), typeof(algs)}(algs, start_index)
     end
 end
 
@@ -73,7 +81,7 @@ end
 
 function reinit_cache!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...)
     foreach(c -> reinit_cache!(c, args...; kwargs...), cache.caches)
-    cache.current = 1
+    cache.current = cache.alg.start_index
     cache.nsteps = 0
     cache.total_time = 0.0
 end
@@ -91,7 +99,7 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb
                     alg.algs),
                 alg,
                 -1,
-                1,
+                alg.start_index,
                 0,
                 0.0,
                 maxtime,
@@ -134,7 +142,7 @@ end
 
     resids = map(x -> Symbol("$(x)_resid"), cache_syms)
     for (sym, resid) in zip(cache_syms, resids)
-        push!(calls, :($(resid) = get_fu($(sym))))
+        push!(calls, :($(resid) = @isdefined($(sym)) ? get_fu($(sym)) : nothing))
     end
     push!(calls,
         quote
@@ -194,25 +202,29 @@ for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProb
     @eval begin
         @generated function SciMLBase.__solve(
                 prob::$probType, alg::$algType{N}, args...; kwargs...) where {N}
-            calls = []
+            calls = [:(current = alg.start_index)]
             sol_syms = [gensym("sol") for _ in 1:N]
             for i in 1:N
                 cur_sol = sol_syms[i]
                 push!(calls,
                     quote
-                        $(cur_sol) = SciMLBase.__solve(prob, alg.algs[$(i)], args...; kwargs...)
-                        if SciMLBase.successful_retcode($(cur_sol))
-                            return SciMLBase.build_solution(
-                                prob, alg, $(cur_sol).u, $(cur_sol).resid;
-                                $(cur_sol).retcode, $(cur_sol).stats,
-                                original = $(cur_sol), trace = $(cur_sol).trace)
+                        if current == $i
+                            $(cur_sol) = SciMLBase.__solve(
+                                prob, alg.algs[$(i)], args...; kwargs...)
+                            if SciMLBase.successful_retcode($(cur_sol))
+                                return SciMLBase.build_solution(
+                                    prob, alg, $(cur_sol).u, $(cur_sol).resid;
+                                    $(cur_sol).retcode, $(cur_sol).stats,
+                                    original = $(cur_sol), trace = $(cur_sol).trace)
+                            end
+                            current = $(i + 1)
                         end
                     end)
             end
 
             resids = map(x -> Symbol("$(x)_resid"), sol_syms)
             for (sym, resid) in zip(sol_syms, resids)
-                push!(calls, :($(resid) = $(sym).resid))
+                push!(calls, :($(resid) = @isdefined($(sym)) ? $(sym).resid : nothing))
             end
 
             push!(calls, quote
@@ -263,6 +275,7 @@ function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve
         algs = (TrustRegion(; concrete_jac, linsolve, precs, autodiff),
             TrustRegion(; concrete_jac, linsolve, precs, autodiff,
                 radius_update_scheme = RadiusUpdateSchemes.Bastin),
+            NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),
             NewtonRaphson(; concrete_jac, linsolve, precs,
                 linesearch = LineSearchesJL(; method = BackTracking()), autodiff),
             TrustRegion(; concrete_jac, linsolve, precs,
@@ -276,7 +289,8 @@ end
 """
     FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing,
         linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false),
-        prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing) where {T}
+        prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing,
+        u0_len::Union{Int, Nothing} = nothing) where {T}
 
 A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods
 for more performance and then tries more robust techniques if the faster ones fail.
@@ -285,12 +299,19 @@ for more performance and then tries more robust techniques if the faster ones fa
 
   - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms
     are compatible with the problem type. Defaults to `Float64`.
+
+### Keyword Arguments
+
+  - `u0_len`: The length of the initial guess. If this is `nothing`, then the length of the
+    initial guess is not checked. If this is an integer and it is less than `25`, we use
+    jacobian based methods.
 """
 function FastShortcutNonlinearPolyalg(
         ::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing,
         precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false),
         prefer_simplenonlinearsolve::Val{SA} = Val(false),
-        autodiff = nothing) where {T, JAC, SA}
+        u0_len::Union{Int, Nothing} = nothing, autodiff = nothing) where {T, JAC, SA}
+    start_index = 1
     if JAC
         if __is_complex(T)
             algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),)
@@ -312,6 +333,7 @@ function FastShortcutNonlinearPolyalg(
                     SimpleKlement(),
                     NewtonRaphson(; concrete_jac, linsolve, precs, autodiff))
             else
+                start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1
                 algs = (SimpleBroyden(),
                     Broyden(; init_jacobian = Val(:true_jacobian), autodiff),
                     SimpleKlement(),
@@ -327,6 +349,8 @@ function FastShortcutNonlinearPolyalg(
                     Klement(; linsolve, precs, autodiff),
                     NewtonRaphson(; concrete_jac, linsolve, precs, autodiff))
             else
+                # TODO: This number requires a bit rigorous testing
+                start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1
                 algs = (Broyden(; autodiff),
                     Broyden(; init_jacobian = Val(:true_jacobian), autodiff),
                     Klement(; linsolve, precs, autodiff),
@@ -339,7 +363,7 @@ function FastShortcutNonlinearPolyalg(
             end
         end
     end
-    return NonlinearSolvePolyAlgorithm(algs, Val(:NLS))
+    return NonlinearSolvePolyAlgorithm(algs, Val(:NLS); start_index)
 end
 
 """
@@ -392,17 +416,19 @@ end
 ## can use that!
 function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
     must_use_jacobian = Val(prob.f.jac !== nothing)
-    return SciMLBase.__init(
-        prob, FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian),
-        args...; kwargs...)
+    return SciMLBase.__init(prob,
+        FastShortcutNonlinearPolyalg(
+            eltype(prob.u0); must_use_jacobian, u0_len = length(prob.u0)),
+        args...;
+        kwargs...)
 end
 
 function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...)
     must_use_jacobian = Val(prob.f.jac !== nothing)
     prefer_simplenonlinearsolve = Val(prob.u0 isa SArray)
     return SciMLBase.__solve(prob,
-        FastShortcutNonlinearPolyalg(
-            eltype(prob.u0); must_use_jacobian, prefer_simplenonlinearsolve),
+        FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian,
+            prefer_simplenonlinearsolve, u0_len = length(prob.u0)),
         args...;
         kwargs...)
 end

diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl
@@ -53,6 +53,10 @@ end
 
 LineSearchesJL(method; kwargs...) = LineSearchesJL(; method, kwargs...)
 function LineSearchesJL(; method = LineSearches.Static(), autodiff = nothing, α = true)
+    if method isa LineSearchesJL  # Prevent breaking old code
+        return LineSearchesJL(method.method, α, autodiff)
+    end
+
     if method isa AbstractNonlinearSolveLineSearchAlgorithm
         Base.depwarn("Passing a native NonlinearSolve line search algorithm to \
                       `LineSearchesJL` or `LineSearch` is deprecated. Pass the method \
@@ -65,6 +69,18 @@ end
 
 Base.@deprecate_binding LineSearch LineSearchesJL true
 
+Static(args...; kwargs...) = LineSearchesJL(LineSearches.Static(args...; kwargs...))
+HagerZhang(args...; kwargs...) = LineSearchesJL(LineSearches.HagerZhang(args...; kwargs...))
+function MoreThuente(args...; kwargs...)
+    return LineSearchesJL(LineSearches.MoreThuente(args...; kwargs...))
+end
+function BackTracking(args...; kwargs...)
+    return LineSearchesJL(LineSearches.BackTracking(args...; kwargs...))
+end
+function StrongWolfe(args...; kwargs...)
+    return LineSearchesJL(LineSearches.StrongWolfe(args...; kwargs...))
+end
+
 # Wrapper over LineSearches.jl algorithms
 @concrete mutable struct LineSearchesJLCache <: AbstractNonlinearSolveLineSearchCache
     f

diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl
@@ -83,7 +83,9 @@ function JacobianCache(
         JacobianOperator(prob, fu, u; jvp_autodiff, vjp_autodiff)
     else
         if has_analytic_jac
-            f.jac_prototype === nothing ? undefmatrix(u) : f.jac_prototype
+            f.jac_prototype === nothing ?
+            similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) :
+            copy(f.jac_prototype)
         elseif f.jac_prototype === nothing
             init_jacobian(jac_cache; preserve_immutable = Val(true))
         else

diff --git a/src/utils.jl b/src/utils.jl
@@ -99,6 +99,7 @@ function __findmin_caches(f, caches)
 end
 function __findmin(f, x)
     return findmin(x) do xᵢ
+        xᵢ === nothing && return Inf
         fx = f(xᵢ)
         return isnan(fx) ? Inf : fx
     end

diff --git a/test/core/23_test_problems_tests.jl b/test/core/23_test_problems_tests.jl
@@ -40,6 +40,16 @@ end
 export test_on_library, problems, dicts
 end
 
+@testitem "PolyAlgorithms" setup=[RobustnessTesting] begin
+    alg_ops = (RobustMultiNewton(), FastShortcutNonlinearPolyalg())
+
+    broken_tests = Dict(alg => Int[] for alg in alg_ops)
+    broken_tests[alg_ops[1]] = []
+    broken_tests[alg_ops[2]] = []
+
+    test_on_library(problems, dicts, alg_ops, broken_tests)
+end
+
 @testitem "NewtonRaphson" setup=[RobustnessTesting] begin
     alg_ops = (NewtonRaphson(),)
 
@@ -91,7 +101,7 @@ end
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end
 
-@testitem "Broyden" retries=5 setup=[RobustnessTesting] begin
+@testitem "Broyden" setup=[RobustnessTesting] begin
     alg_ops = (Broyden(), Broyden(; init_jacobian = Val(:true_jacobian)),
         Broyden(; update_rule = Val(:bad_broyden)),
         Broyden(; init_jacobian = Val(:true_jacobian), update_rule = Val(:bad_broyden)))

diff --git a/test/core/forward_ad_tests.jl b/test/core/forward_ad_tests.jl
@@ -64,8 +64,8 @@ end
 
 @testitem "ForwardDiff.jl Integration" setup=[ForwardADTesting] begin
     for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(),
-        PseudoTransient(; alpha_initial = 10.0), Broyden(), Klement(), DFSane(), nothing,
-        NLsolveJL(), CMINPACK(), KINSOL(; globalization_strategy = :LineSearch))
+        PseudoTransient(; alpha_initial = 10.0), Broyden(), Klement(), DFSane(),
+        nothing, NLsolveJL(), CMINPACK(), KINSOL(; globalization_strategy = :LineSearch))
         us = (2.0, @SVector[1.0, 1.0], [1.0, 1.0], ones(2, 2), @SArray ones(2, 2))
 
         @testset "Scalar AD" begin

diff --git a/test/core/nlls_tests.jl b/test/core/nlls_tests.jl
@@ -94,3 +94,22 @@ end
         @test maximum(abs, sol.resid) < 1e-6
     end
 end
+
+@testitem "NLLS Analytic Jacobian" begin
+    dataIn = 1:10
+    f(x, p) = x[1] * dataIn .^ 2 .+ x[2] * dataIn .+ x[3]
+    dataOut = f([1, 2, 3], nothing) + 0.1 * randn(10, 1)
+
+    resid(x, p) = f(x, p) - dataOut
+    jac(x, p) = [dataIn .^ 2 dataIn ones(10, 1)]
+    x0 = [1, 1, 1]
+
+    prob = NonlinearLeastSquaresProblem(resid, x0)
+    sol1 = solve(prob)
+
+    nlfunc = NonlinearFunction(resid; jac)
+    prob = NonlinearLeastSquaresProblem(nlfunc, x0)
+    sol2 = solve(prob)
+
+    @test sol1.u ≈ sol2.u
+end