Merge pull request #399 from SciML/ap/nlls_term

Use a different termination norm for NLLS

Author: ChrisRackauckas

.github/workflows/Downstream.yml

@@ -55,6 +55,9 @@ jobs:
             @info "Not compatible with this release. No problem." exception=err
             exit(0)  # Exit immediately, as a success
           end
+        env:
+          RETESTITEMS_NWORKERS: 4
+          RETESTITEMS_NWORKER_THREADS: 2
       - uses: julia-actions/julia-processcoverage@v1
         with:
           directories: src,ext

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.8.4"
+version = "3.9.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -58,15 +58,15 @@ NonlinearSolveZygoteExt = "Zygote"
 [compat]
 ADTypes = "0.2.6"
 Aqua = "0.8"
-ArrayInterface = "7.7"
+ArrayInterface = "7.9"
 BandedMatrices = "1.4"
 BenchmarkTools = "1.4"
 ConcreteStructs = "0.2.3"
-CUDA = "5.1"
-DiffEqBase = "6.146.0"
+CUDA = "5.2"
+DiffEqBase = "6.149.0"
 Enzyme = "0.11.15"
 FastBroadcast = "0.2.8"
-FastClosures = "0.3"
+FastClosures = "0.3.2"
 FastLevenbergMarquardt = "0.1"
 FiniteDiff = "2.21"
 FixedPointAcceleration = "0.3"
@@ -82,21 +82,21 @@ NLsolve = "4.5"
 NLSolvers = "0.5"
 NaNMath = "1"
 NonlinearProblemLibrary = "0.1.2"
-OrdinaryDiffEq = "6.63"
+OrdinaryDiffEq = "6.74"
 Pkg = "1.10"
 PrecompileTools = "1.2"
 Preferences = "1.4"
 Printf = "1.10"
 Random = "1.91"
 ReTestItems = "1"
-RecursiveArrayTools = "3.4"
+RecursiveArrayTools = "3.8"
 Reexport = "1.2"
 SIAMFANLEquations = "1.0.1"
 SafeTestsets = "0.1"
-SciMLBase = "2.19.0"
+SciMLBase = "2.28.0"
 SimpleNonlinearSolve = "1.2"
 SparseArrays = "1.10"
-SparseDiffTools = "2.14"
+SparseDiffTools = "2.17"
 SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.7"
@@ -105,7 +105,7 @@ Sundials = "4.23.1"
 Symbolics = "5.13"
 Test = "1.10"
 TimerOutputs = "0.5.23"
-Zygote = "0.6.67"
+Zygote = "0.6.69"
 julia = "1.10"
 
 [extras]

src/core/approximate_jacobian.jl

@@ -167,7 +167,7 @@ function SciMLBase.__init(
             prob, alg.initialization, alg, f, fu, u, p; linsolve, maxiters, internalnorm)
 
         abstol, reltol, termination_cache = init_termination_cache(
-            abstol, reltol, fu, u, termination_condition)
+            prob, abstol, reltol, fu, u, termination_condition)
         linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs)
 
         J = initialization_cache(nothing)
@@ -206,7 +206,7 @@ function SciMLBase.__init(
         update_rule_cache = __internal_init(
             prob, alg.update_rule, J, fu, u, du; internalnorm)
 
-        trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du;
+        trace = init_nonlinearsolve_trace(prob, alg, u, fu, ApplyArray(__zero, J), du;
             uses_jacobian_inverse = Val(INV), kwargs...)
 
         return ApproximateJacobianSolveCache{INV, GB, iip, maxtime !== nothing}(

src/core/generalized_first_order.jl

@@ -156,7 +156,7 @@ function SciMLBase.__init(
         linsolve = get_linear_solver(alg.descent)
 
         abstol, reltol, termination_cache = init_termination_cache(
-            abstol, reltol, fu, u, termination_condition)
+            prob, abstol, reltol, fu, u, termination_condition)
         linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs)
 
         jac_cache = JacobianCache(
@@ -191,7 +191,8 @@ function SciMLBase.__init(
             GB = :LineSearch
         end
 
-        trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...)
+        trace = init_nonlinearsolve_trace(
+            prob, alg, u, fu, ApplyArray(__zero, J), du; kwargs...)
 
         return GeneralizedFirstOrderAlgorithmCache{iip, GB, maxtime !== nothing}(
             fu, u, u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache,

src/core/spectral_methods.jl

@@ -133,8 +133,8 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane
             prob, alg.linesearch, prob.f, fu, u, prob.p; maxiters, internalnorm, kwargs...)
 
         abstol, reltol, tc_cache = init_termination_cache(
-            abstol, reltol, fu, u_cache, termination_condition)
-        trace = init_nonlinearsolve_trace(alg, u, fu, nothing, du; kwargs...)
+            prob, abstol, reltol, fu, u_cache, termination_condition)
+        trace = init_nonlinearsolve_trace(prob, alg, u, fu, nothing, du; kwargs...)
 
         if alg.σ_1 === nothing
             σ_n = dot(u, u) / dot(u, fu)

src/internal/termination.jl

@@ -1,9 +1,23 @@
-function init_termination_cache(abstol, reltol, du, u, ::Nothing)
-    return init_termination_cache(
-        abstol, reltol, du, u, AbsSafeBestTerminationMode(; max_stalled_steps = 32))
+function init_termination_cache(prob::NonlinearProblem, abstol, reltol, du, u, ::Nothing)
+    return init_termination_cache(prob, abstol, reltol, du, u,
+        AbsSafeBestTerminationMode(Base.Fix1(maximum, abs); max_stalled_steps = 32))
 end
-function init_termination_cache(abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode)
-    tc_cache = init(du, u, tc; abstol, reltol, use_deprecated_retcodes = Val(false))
+function init_termination_cache(
+        prob::NonlinearLeastSquaresProblem, abstol, reltol, du, u, ::Nothing)
+    return init_termination_cache(prob, abstol, reltol, du, u,
+        AbsSafeBestTerminationMode(Base.Fix2(norm, 2); max_stalled_steps = 32))
+end
+
+function init_termination_cache(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem},
+        abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode)
+    tc_ = if hasfield(typeof(tc), :internalnorm) && tc.internalnorm === nothing
+        internalnorm = ifelse(
+            prob isa NonlinearProblem, Base.Fix1(maximum, abs), Base.Fix2(norm, 2))
+        DiffEqBase.set_termination_mode_internalnorm(tc, internalnorm)
+    else
+        tc
+    end
+    tc_cache = init(du, u, tc_; abstol, reltol, use_deprecated_retcodes = Val(false))
     return DiffEqBase.get_abstol(tc_cache), DiffEqBase.get_reltol(tc_cache), tc_cache
 end
 

src/internal/tracing.jl

@@ -52,7 +52,7 @@ for Tr in (:TraceMinimal, :TraceWithJacobianConditionNumber, :TraceAll)
 end
 
 # NonlinearSolve Tracing Utilities
-@concrete struct NonlinearSolveTraceEntry
+@concrete struct NonlinearSolveTraceEntry{nType}
     iteration::Int
     fnorm
     stepnorm
@@ -63,19 +63,27 @@ end
     δu
 end
 
-function __show_top_level(io::IO, entry::NonlinearSolveTraceEntry)
+function __show_top_level(io::IO, entry::NonlinearSolveTraceEntry{nType}) where {nType}
     if entry.condJ === nothing
         @printf io "%-8s %-20s %-20s\n" "----" "-------------" "-----------"
-        @printf io "%-8s %-20s %-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm"
+        if nType === :L2
+            @printf io "%-8s %-20s %-20s\n" "Iter" "f(u) 2-norm" "Step 2-norm"
+        else
+            @printf io "%-8s %-20s %-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm"
+        end
         @printf io "%-8s %-20s %-20s\n" "----" "-------------" "-----------"
     else
         @printf io "%-8s %-20s %-20s %-20s\n" "----" "-------------" "-----------" "-------"
-        @printf io "%-8s %-20s %-20s %-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm" "cond(J)"
+        if nType === :L2
+            @printf io "%-8s %-20s %-20s %-20s\n" "Iter" "f(u) 2-norm" "Step 2-norm" "cond(J)"
+        else
+            @printf io "%-8s %-20s %-20s %-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm" "cond(J)"
+        end
         @printf io "%-8s %-20s %-20s %-20s\n" "----" "-------------" "-----------" "-------"
     end
 end
 
-function Base.show(io::IO, entry::NonlinearSolveTraceEntry)
+function Base.show(io::IO, entry::NonlinearSolveTraceEntry{nType}) where {nType}
     entry.iteration == 0 && __show_top_level(io, entry)
     if entry.iteration < 0
         # Special case for final entry
@@ -89,25 +97,32 @@ function Base.show(io::IO, entry::NonlinearSolveTraceEntry)
     return nothing
 end
 
-function NonlinearSolveTraceEntry(iteration, fu, δu)
-    return NonlinearSolveTraceEntry(
-        iteration, norm(fu, Inf), norm(δu, 2), nothing, nothing, nothing, nothing, nothing)
+function NonlinearSolveTraceEntry(prob::AbstractNonlinearProblem, iteration, fu, δu)
+    nType = ifelse(prob isa NonlinearLeastSquaresProblem, :L2, :Inf)
+    fnorm = prob isa NonlinearLeastSquaresProblem ? norm(fu, 2) : norm(fu, Inf)
+    return NonlinearSolveTraceEntry{nType}(
+        iteration, fnorm, norm(δu, 2), nothing, nothing, nothing, nothing, nothing)
 end
 
-function NonlinearSolveTraceEntry(iteration, fu, δu, J)
-    return NonlinearSolveTraceEntry(iteration, norm(fu, Inf), norm(δu, 2),
-        __cond(J), nothing, nothing, nothing, nothing)
+function NonlinearSolveTraceEntry(prob::AbstractNonlinearProblem, iteration, fu, δu, J)
+    nType = ifelse(prob isa NonlinearLeastSquaresProblem, :L2, :Inf)
+    fnorm = prob isa NonlinearLeastSquaresProblem ? norm(fu, 2) : norm(fu, Inf)
+    return NonlinearSolveTraceEntry{nType}(
+        iteration, fnorm, norm(δu, 2), __cond(J), nothing, nothing, nothing, nothing)
 end
 
-function NonlinearSolveTraceEntry(iteration, fu, δu, J, u)
-    return NonlinearSolveTraceEntry(iteration, norm(fu, Inf), norm(δu, 2), __cond(J),
+function NonlinearSolveTraceEntry(prob::AbstractNonlinearProblem, iteration, fu, δu, J, u)
+    nType = ifelse(prob isa NonlinearLeastSquaresProblem, :L2, :Inf)
+    fnorm = prob isa NonlinearLeastSquaresProblem ? norm(fu, 2) : norm(fu, Inf)
+    return NonlinearSolveTraceEntry{nType}(iteration, fnorm, norm(δu, 2), __cond(J),
         __copy(J), __copy(u), __copy(fu), __copy(δu))
 end
 
 @concrete struct NonlinearSolveTrace{
     show_trace, store_trace, Tr <: AbstractNonlinearSolveTraceLevel}
     history
     trace_level::Tr
+    prob
 end
 
 function reset!(trace::NonlinearSolveTrace)
@@ -123,61 +138,63 @@ function Base.show(io::IO, trace::NonlinearSolveTrace)
     return nothing
 end
 
-function init_nonlinearsolve_trace(alg, u, fu, J, δu; show_trace::Val = Val(false),
+function init_nonlinearsolve_trace(prob, alg, u, fu, J, δu; show_trace::Val = Val(false),
         trace_level::AbstractNonlinearSolveTraceLevel = TraceMinimal(),
         store_trace::Val = Val(false), uses_jac_inverse = Val(false), kwargs...)
     return init_nonlinearsolve_trace(
-        alg, show_trace, trace_level, store_trace, u, fu, J, δu, uses_jac_inverse)
+        prob, alg, show_trace, trace_level, store_trace, u, fu, J, δu, uses_jac_inverse)
 end
 
-function init_nonlinearsolve_trace(
-        alg, ::Val{show_trace}, trace_level::AbstractNonlinearSolveTraceLevel,
-        ::Val{store_trace}, u, fu, J, δu,
-        ::Val{uses_jac_inverse}) where {show_trace, store_trace, uses_jac_inverse}
+function init_nonlinearsolve_trace(prob::AbstractNonlinearProblem, alg, ::Val{show_trace},
+        trace_level::AbstractNonlinearSolveTraceLevel, ::Val{store_trace}, u, fu, J,
+        δu, ::Val{uses_jac_inverse}) where {show_trace, store_trace, uses_jac_inverse}
     if show_trace
         print("\nAlgorithm: ")
         Base.printstyled(alg, "\n\n"; color = :green, bold = true)
     end
     J_ = uses_jac_inverse ? (trace_level isa TraceMinimal ? J : __safe_inv(J)) : J
     history = __init_trace_history(
-        Val{show_trace}(), trace_level, Val{store_trace}(), u, fu, J_, δu)
-    return NonlinearSolveTrace{show_trace, store_trace}(history, trace_level)
+        prob, Val{show_trace}(), trace_level, Val{store_trace}(), u, fu, J_, δu)
+    return NonlinearSolveTrace{show_trace, store_trace}(history, trace_level, prob)
 end
 
-function __init_trace_history(::Val{show_trace}, trace_level, ::Val{store_trace},
-        u, fu, J, δu) where {show_trace, store_trace}
+function __init_trace_history(
+        prob::AbstractNonlinearProblem, ::Val{show_trace}, trace_level,
+        ::Val{store_trace}, u, fu, J, δu) where {show_trace, store_trace}
     !store_trace && !show_trace && return nothing
-    entry = __trace_entry(trace_level, 0, u, fu, J, δu)
+    entry = __trace_entry(prob, trace_level, 0, u, fu, J, δu)
     show_trace && show(entry)
     store_trace && return NonlinearSolveTraceEntry[entry]
     return nothing
 end
 
-function __trace_entry(::TraceMinimal, iter, u, fu, J, δu, α = 1)
-    return NonlinearSolveTraceEntry(iter, fu, δu .* α)
+function __trace_entry(prob, ::TraceMinimal, iter, u, fu, J, δu, α = 1)
+    return NonlinearSolveTraceEntry(prob, iter, fu, δu .* α)
 end
-function __trace_entry(::TraceWithJacobianConditionNumber, iter, u, fu, J, δu, α = 1)
-    return NonlinearSolveTraceEntry(iter, fu, δu .* α, J)
+function __trace_entry(prob, ::TraceWithJacobianConditionNumber, iter, u, fu, J, δu, α = 1)
+    return NonlinearSolveTraceEntry(prob, iter, fu, δu .* α, J)
 end
-function __trace_entry(::TraceAll, iter, u, fu, J, δu, α = 1)
-    return NonlinearSolveTraceEntry(iter, fu, δu .* α, J, u)
+function __trace_entry(prob, ::TraceAll, iter, u, fu, J, δu, α = 1)
+    return NonlinearSolveTraceEntry(prob, iter, fu, δu .* α, J, u)
 end
 
 function update_trace!(trace::NonlinearSolveTrace{ShT, StT}, iter, u, fu, J, δu,
         α = 1; last::Val{L} = Val(false)) where {ShT, StT, L}
     !StT && !ShT && return nothing
 
     if L
-        entry = NonlinearSolveTraceEntry(
-            -1, norm(fu, Inf), NaN32, nothing, nothing, nothing, nothing, nothing)
+        nType = ifelse(trace.prob isa NonlinearLeastSquaresProblem, :L2, :Inf)
+        fnorm = trace.prob isa NonlinearLeastSquaresProblem ? norm(fu, 2) : norm(fu, Inf)
+        entry = NonlinearSolveTraceEntry{nType}(
+            -1, fnorm, NaN32, nothing, nothing, nothing, nothing, nothing)
         ShT && show(entry)
         return trace
     end
 
     show_now = ShT && (mod1(iter, trace.trace_level.print_frequency) == 1)
     store_now = StT && (mod1(iter, trace.trace_level.store_frequency) == 1)
     (show_now || store_now) &&
-        (entry = __trace_entry(trace.trace_level, iter, u, fu, J, δu, α))
+        (entry = __trace_entry(trace.prob, trace.trace_level, iter, u, fu, J, δu, α))
     store_now && push!(trace.history, entry)
     show_now && show(entry)
     return trace

test/core/forward_ad_tests.jl

@@ -79,7 +79,6 @@ end
                     gs = abs.(ForwardDiff.derivative(solve_with(Val{mode}(), u0, alg), p))
                     gs_true = abs.(jacobian_f(u0, p))
                     if !(isapprox(gs, gs_true, atol = 1e-5))
-                        @show sol.retcode, sol.u
                         @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_gradient=gs true_gradient=gs_true
                     else
                         @test abs.(gs)≈abs.(gs_true) atol=1e-5

test/core/nlls_tests.jl

@@ -6,11 +6,11 @@ using Reexport
 true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4])
 true_function(y, x, θ) = (@. y = θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]))
 
-θ_true = [1.0, 0.1, 2.0, 0.5]
+const θ_true = [1.0, 0.1, 2.0, 0.5]
 
-x = [-1.0, -0.5, 0.0, 0.5, 1.0]
+const x = [-1.0, -0.5, 0.0, 0.5, 1.0]
 
-y_target = true_function(x, θ_true)
+const y_target = true_function(x, θ_true)
 
 function loss_function(θ, p)
     ŷ = true_function(p, θ)
@@ -23,7 +23,7 @@ function loss_function(resid, θ, p)
     return resid
 end
 
-θ_init = θ_true .+ randn!(StableRNG(0), similar(θ_true)) * 0.1
+const θ_init = θ_true .+ randn!(StableRNG(0), similar(θ_true)) * 0.1
 
 solvers = []
 for linsolve in [nothing, LUFactorization(), KrylovJL_GMRES(), KrylovJL_LSMR()]
@@ -56,9 +56,9 @@ end
     nlls_problems = [prob_oop, prob_iip]
 
     for prob in nlls_problems, solver in solvers
-        sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8)
+        sol = solve(prob, solver; maxiters = 10000, abstol = 1e-6)
         @test SciMLBase.successful_retcode(sol)
-        @test maximum(abs, sol.resid) < 1e-6
+        @test norm(sol.resid, 2) < 1e-6
     end
 end
 
@@ -90,8 +90,9 @@ end
             x)]
 
     for prob in probs, solver in solvers
-        sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8)
-        @test maximum(abs, sol.resid) < 1e-6
+        sol = solve(prob, solver; maxiters = 10000, abstol = 1e-6)
+        @test SciMLBase.successful_retcode(sol)
+        @test norm(sol.resid, 2) < 1e-6
     end
 end