Test the new NonlinearSolveBase.jl

avik-pal · avik-pal · commit 5afc4ba60f82 · 2024-02-26T22:23:29.000-05:00
diff --git a/Manifest.toml b/Manifest.toml
diff --git a/Project.toml b/Project.toml
@@ -1,41 +1,42 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.5.0"
+version = "1.6.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
-DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
+NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
-SimpleNonlinearSolveChainRulesCoreExt = "ChainRulesCore"
+SimpleNonlinearSolveChainRulesCoreDiffEqBaseExt = ["ChainRulesCore", "DiffEqBase"]
 SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
 SimpleNonlinearSolveStaticArraysExt = "StaticArrays"
 SimpleNonlinearSolveZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "0.2.6"
 AllocCheck = "0.1.1"
-ArrayInterface = "7.7"
 Aqua = "0.8"
+ArrayInterface = "7.7"
 CUDA = "5.2"
 ChainRulesCore = "1.22"
 ConcreteStructs = "0.2.3"
@@ -48,6 +49,7 @@ LinearAlgebra = "1.10"
 LinearSolve = "2.25"
 MaybeInplace = "0.1.1"
 NonlinearProblemLibrary = "0.1.2"
+NonlinearSolveBase = "1"
 Pkg = "1.10"
 PolyesterForwardDiff = "0.1.1"
 PrecompileTools = "1.2"
@@ -66,12 +68,12 @@ julia = "1.10"
 AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
+NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -83,4 +85,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "AllocCheck", "DiffEqBase", "ForwardDiff", "LinearAlgebra", "LinearSolve", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SciMLSensitivity", "StaticArrays", "Zygote", "CUDA", "PolyesterForwardDiff", "Reexport", "Test", "FiniteDiff"]
+test = ["Aqua", "AllocCheck", "NonlinearSolveBase", "ForwardDiff", "LinearAlgebra", "LinearSolve", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SciMLSensitivity", "StaticArrays", "Zygote", "CUDA", "PolyesterForwardDiff", "Reexport", "Test", "FiniteDiff"]
diff --git a/ext/SimpleNonlinearSolveChainRulesCoreDiffEqBaseExt.jl b/ext/SimpleNonlinearSolveChainRulesCoreDiffEqBaseExt.jl
@@ -1,4 +1,4 @@
-module SimpleNonlinearSolveChainRulesCoreExt
+module SimpleNonlinearSolveChainRulesCoreDiffEqBaseExt
 
 using ChainRulesCore, DiffEqBase, SciMLBase, SimpleNonlinearSolve
 
diff --git a/src/SimpleNonlinearSolve.jl b/src/SimpleNonlinearSolve.jl
@@ -3,22 +3,23 @@ module SimpleNonlinearSolve
 import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidations
 
 @recompile_invalidations begin
-    using ADTypes, ArrayInterface, ConcreteStructs, DiffEqBase, FastClosures, FiniteDiff,
-          ForwardDiff, Reexport, LinearAlgebra, SciMLBase
-
-    import DiffEqBase: AbstractNonlinearTerminationMode,
-                       AbstractSafeNonlinearTerminationMode,
-                       AbstractSafeBestNonlinearTerminationMode,
-                       NonlinearSafeTerminationReturnCode, get_termination_mode,
-                       NONLINEARSOLVE_DEFAULT_NORM
+    using ADTypes, ArrayInterface, FiniteDiff, ForwardDiff, NonlinearSolveBase, Reexport,
+          LinearAlgebra, SciMLBase
+
+    import ConcreteStructs: @concrete
     import DiffResults
+    import FastClosures: @closure
     import ForwardDiff: Dual
     import MaybeInplace: @bb, setindex_trait, CanSetindex, CannotSetindex
+    import NonlinearSolveBase: AbstractNonlinearTerminationMode,
+                               AbstractSafeNonlinearTerminationMode,
+                               AbstractSafeBestNonlinearTerminationMode,
+                               get_termination_mode, NONLINEARSOLVE_DEFAULT_NORM
     import SciMLBase: AbstractNonlinearAlgorithm, build_solution, isinplace, _unwrap_val
     import StaticArraysCore: StaticArray, SVector, SMatrix, SArray, MArray, MMatrix, Size
 end
 
-@reexport using ADTypes, SciMLBase
+@reexport using ADTypes, SciMLBase  # TODO: Reexport NonlinearSolveBase after the situation with NonlinearSolve.jl is resolved
 
 abstract type AbstractSimpleNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end
 abstract type AbstractBracketingAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm end
@@ -58,23 +59,28 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Nothing, args...;
 end
 
 # By Pass the highlevel checks for NonlinearProblem for Simple Algorithms
-function SciMLBase.solve(
-        prob::NonlinearProblem, alg::AbstractSimpleNonlinearSolveAlgorithm,
-        args...; sensealg = nothing, u0 = nothing, p = nothing, kwargs...)
-    if sensealg === nothing && haskey(prob.kwargs, :sensealg)
-        sensealg = prob.kwargs[:sensealg]
-    end
-    new_u0 = u0 !== nothing ? u0 : prob.u0
-    new_p = p !== nothing ? p : prob.p
-    return __internal_solve_up(
-        prob, sensealg, new_u0, u0 === nothing, new_p, p === nothing,
-        alg, args...; prob.kwargs..., kwargs...)
-end
+# Using eval to prevent ambiguity
+for pType in (NonlinearProblem, NonlinearLeastSquaresProblem)
+    @eval begin
+        function SciMLBase.solve(
+                prob::$(pType), alg::AbstractSimpleNonlinearSolveAlgorithm, args...;
+                sensealg = nothing, u0 = nothing, p = nothing, kwargs...)
+            if sensealg === nothing && haskey(prob.kwargs, :sensealg)
+                sensealg = prob.kwargs[:sensealg]
+            end
+            new_u0 = u0 !== nothing ? u0 : prob.u0
+            new_p = p !== nothing ? p : prob.p
+            return __internal_solve_up(
+                prob, sensealg, new_u0, u0 === nothing, new_p, p === nothing,
+                alg, args...; prob.kwargs..., kwargs...)
+        end
 
-function __internal_solve_up(_prob::NonlinearProblem, sensealg, u0, u0_changed, p,
-        p_changed, alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...)
-    prob = u0_changed || p_changed ? remake(_prob; u0, p) : _prob
-    return SciMLBase.__solve(prob, alg, args...; kwargs...)
+        function __internal_solve_up(_prob::$(pType), sensealg, u0, u0_changed, p,
+                p_changed, alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...)
+            prob = u0_changed || p_changed ? remake(_prob; u0, p) : _prob
+            return SciMLBase.__solve(prob, alg, args...; kwargs...)
+        end
+    end
 end
 
 @setup_workload begin
diff --git a/src/bracketing/bisection.jl b/src/bracketing/bisection.jl
@@ -26,7 +26,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Bisection, args...
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = __get_tolerance(nothing, abstol,
+    abstol = NonlinearSolveBase.get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)
diff --git a/src/bracketing/brent.jl b/src/bracketing/brent.jl
@@ -13,7 +13,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Brent, args...;
     fl, fr = f(left), f(right)
     ϵ = eps(convert(typeof(fl), 1))
 
-    abstol = __get_tolerance(nothing, abstol,
+    abstol = NonlinearSolveBase.get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)
diff --git a/src/bracketing/falsi.jl b/src/bracketing/falsi.jl
@@ -12,7 +12,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Falsi, args...;
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = __get_tolerance(nothing, abstol,
+    abstol = NonlinearSolveBase.get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)
diff --git a/src/bracketing/itp.jl b/src/bracketing/itp.jl
@@ -58,7 +58,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::ITP, args...;
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = __get_tolerance(nothing, abstol,
+    abstol = NonlinearSolveBase.get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)
diff --git a/src/bracketing/ridder.jl b/src/bracketing/ridder.jl
@@ -12,7 +12,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Ridder, args...;
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = __get_tolerance(nothing, abstol,
+    abstol = NonlinearSolveBase.get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)
diff --git a/src/linesearch.jl b/src/linesearch.jl
@@ -73,7 +73,7 @@ function (cache::LiFukushimaLineSearchCache)(u, δu)
     fx_norm = ϕ(T(0))
 
     # Non-Blocking exit if the norm is NaN or Inf
-    DiffEqBase.NAN_CHECK(fx_norm) && return cache.α
+    NonlinearSolveBase.NAN_CHECK(fx_norm) && return cache.α
 
     # Early Terminate based on Eq. 2.7
     du_norm = NONLINEARSOLVE_DEFAULT_NORM(δu)
@@ -84,12 +84,12 @@ function (cache::LiFukushimaLineSearchCache)(u, δu)
     fxλp_norm = ϕ(λ₂)
 
     if cache.nan_maxiters !== nothing
-        if DiffEqBase.NAN_CHECK(fxλp_norm)
+        if NonlinearSolveBase.NAN_CHECK(fxλp_norm)
             nan_converged = false
             for _ in 1:(cache.nan_maxiters)
                 λ₁, λ₂ = λ₂, cache.β * λ₂
                 fxλp_norm = ϕ(λ₂)
-                nan_converged = DiffEqBase.NAN_CHECK(fxλp_norm)::Bool
+                nan_converged = NonlinearSolveBase.NAN_CHECK(fxλp_norm)::Bool
                 nan_converged && break
             end
             nan_converged || return cache.α
diff --git a/src/nlsolve/lbroyden.jl b/src/nlsolve/lbroyden.jl
@@ -121,7 +121,7 @@ function __static_solve(prob::NonlinearProblem{<:SArray}, alg::SimpleLimitedMemo
 
     U, Vᵀ = __init_low_rank_jacobian(vec(x), vec(fx), threshold)
 
-    abstol = __get_tolerance(x, abstol, eltype(x))
+    abstol = NonlinearSolveBase.get_tolerance(x, abstol, eltype(x))
 
     xo, δx, fo, δf = x, -fx, fx, fx
 
diff --git a/src/utils.jl b/src/utils.jl
@@ -77,7 +77,7 @@ except `cache` (& `J` if not nothing) are mutated.
 function value_and_jacobian(ad, f::F, y, x::X, p, cache; J = nothing) where {F, X}
     if isinplace(f)
         _f = (du, u) -> f(du, u, p)
-        if DiffEqBase.has_jac(f)
+        if SciMLBase.has_jac(f)
             f.jac(J, x, p)
             _f(y, x)
             return y, J
@@ -97,7 +97,7 @@ function value_and_jacobian(ad, f::F, y, x::X, p, cache; J = nothing) where {F,
         end
     else
         _f = Base.Fix2(f, p)
-        if DiffEqBase.has_jac(f)
+        if SciMLBase.has_jac(f)
             return _f(x), f.jac(x, p)
         elseif ad isa AutoForwardDiff
             if ArrayInterface.can_setindex(x)
@@ -124,7 +124,7 @@ end
 function __polyester_forwarddiff_jacobian! end
 
 function value_and_jacobian(ad, f::F, y, x::Number, p, cache; J = nothing) where {F}
-    if DiffEqBase.has_jac(f)
+    if SciMLBase.has_jac(f)
         return f(x, p), f.jac(x, p)
     elseif ad isa AutoForwardDiff
         T = typeof(__standard_tag(ad.tag, x))
@@ -152,7 +152,7 @@ function jacobian_cache(ad, f::F, y, x::X, p) where {F, X <: AbstractArray}
     if isinplace(f)
         _f = (du, u) -> f(du, u, p)
         J = similar(y, length(y), length(x))
-        if DiffEqBase.has_jac(f)
+        if SciMLBase.has_jac(f)
             return J, nothing
         elseif ad isa AutoForwardDiff || ad isa AutoPolyesterForwardDiff
             return J, __get_jacobian_config(ad, _f, y, x)
@@ -163,7 +163,7 @@ function jacobian_cache(ad, f::F, y, x::X, p) where {F, X <: AbstractArray}
         end
     else
         _f = Base.Fix2(f, p)
-        if DiffEqBase.has_jac(f)
+        if SciMLBase.has_jac(f)
             return nothing, nothing
         elseif ad isa AutoForwardDiff
             J = ArrayInterface.can_setindex(x) ? similar(y, length(y), length(x)) : nothing
@@ -292,58 +292,27 @@ function init_termination_cache(abstol, reltol, du, u, ::Nothing)
     return init_termination_cache(abstol, reltol, du, u, AbsNormTerminationMode())
 end
 function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode)
-    T = promote_type(eltype(du), eltype(u))
-    abstol = __get_tolerance(u, abstol, T)
-    reltol = __get_tolerance(u, reltol, T)
     tc_cache = init(du, u, tc; abstol, reltol)
-    return DiffEqBase.get_abstol(tc_cache), DiffEqBase.get_reltol(tc_cache), tc_cache
+    return (NonlinearSolveBase.get_abstol(tc_cache),
+        NonlinearSolveBase.get_reltol(tc_cache), tc_cache)
 end
 
 function check_termination(tc_cache, fx, x, xo, prob, alg)
     return check_termination(tc_cache, fx, x, xo, prob, alg,
-        DiffEqBase.get_termination_mode(tc_cache))
-end
-function check_termination(tc_cache, fx, x, xo, prob, alg,
-        ::AbstractNonlinearTerminationMode)
-    if Bool(tc_cache(fx, x, xo))
-        return build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
-    end
-    return nothing
-end
-function check_termination(tc_cache, fx, x, xo, prob, alg,
-        ::AbstractSafeNonlinearTerminationMode)
-    if Bool(tc_cache(fx, x, xo))
-        if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success
-            retcode = ReturnCode.Success
-        elseif tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination
-            retcode = ReturnCode.ConvergenceFailure
-        elseif tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination
-            retcode = ReturnCode.Unstable
-        else
-            error("Unknown termination code: $(tc_cache.retcode)")
-        end
-        return build_solution(prob, alg, x, fx; retcode)
-    end
-    return nothing
+        NonlinearSolveBase.get_termination_mode(tc_cache))
 end
+
 function check_termination(tc_cache, fx, x, xo, prob, alg,
-        ::AbstractSafeBestNonlinearTerminationMode)
+        mode::AbstractNonlinearTerminationMode)
     if Bool(tc_cache(fx, x, xo))
-        if tc_cache.retcode == NonlinearSafeTerminationReturnCode.Success
-            retcode = ReturnCode.Success
-        elseif tc_cache.retcode == NonlinearSafeTerminationReturnCode.PatienceTermination
-            retcode = ReturnCode.ConvergenceFailure
-        elseif tc_cache.retcode == NonlinearSafeTerminationReturnCode.ProtectiveTermination
-            retcode = ReturnCode.Unstable
-        else
-            error("Unknown termination code: $(tc_cache.retcode)")
-        end
-        if isinplace(prob)
-            prob.f(fx, x, prob.p)
-        else
-            fx = prob.f(x, prob.p)
+        if mode isa AbstractSafeBestNonlinearTerminationMode
+            if isinplace(prob)
+                prob.f(fx, x, prob.p)
+            else
+                fx = prob.f(x, prob.p)
+            end
         end
-        return build_solution(prob, alg, tc_cache.u, fx; retcode)
+        return build_solution(prob, alg, x, fx; retcode = tc_cache.retcode)
     end
     return nothing
 end
@@ -382,12 +351,5 @@ end
 @inline __reshape(x::Number, args...) = x
 @inline __reshape(x::AbstractArray, args...) = reshape(x, args...)
 
-# Override cases which might be used in a kernel launch
-__get_tolerance(x, η, ::Type{T}) where {T} = DiffEqBase._get_tolerance(η, T)
-function __get_tolerance(x::Union{SArray, Number}, ::Nothing, ::Type{T}) where {T}
-    η = real(oneunit(T)) * (eps(real(one(T))))^(real(T)(0.8))
-    return T(η)
-end
-
 # Extension
 function __zygote_compute_nlls_vjp end
diff --git a/test/core/23_test_problems_tests.jl b/test/core/23_test_problems_tests.jl
@@ -1,5 +1,5 @@
 @testsetup module RobustnessTesting
-using LinearAlgebra, NonlinearProblemLibrary, DiffEqBase, Test
+using LinearAlgebra, NonlinearProblemLibrary, NonlinearSolveBase, SciMLBase, Test
 
 problems = NonlinearProblemLibrary.problems
 dicts = NonlinearProblemLibrary.dicts
diff --git a/test/core/rootfind_tests.jl b/test/core/rootfind_tests.jl
@@ -1,7 +1,8 @@
 @testsetup module RootfindingTesting
 using Reexport
 @reexport using AllocCheck,
-                LinearSolve, StaticArrays, Random, LinearAlgebra, ForwardDiff, DiffEqBase
+                LinearSolve, StaticArrays, Random, LinearAlgebra, ForwardDiff,
+                NonlinearSolveBase
 import PolyesterForwardDiff
 
 quadratic_f(u, p) = u .* u .- p
@@ -89,7 +90,7 @@ end
     end
 end
 
-@testitem "Derivative Free Metods" setup=[RootfindingTesting] begin
+@testitem "Derivative Free Methods" setup=[RootfindingTesting] begin
     @testset "$(nameof(typeof(alg)))" for alg in [SimpleBroyden(), SimpleKlement(),
         SimpleDFSane(), SimpleLimitedMemoryBroyden(),
         SimpleBroyden(; linesearch = Val(true)),

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-module SimpleNonlinearSolveChainRulesCoreExt`
	`1`	`+module SimpleNonlinearSolveChainRulesCoreDiffEqBaseExt`
`2`	`2`
`3`	`3`	`using ChainRulesCore, DiffEqBase, SciMLBase, SimpleNonlinearSolve`
`4`	`4`