Fix the tracing for NLLS

Author: avik-pal

src/core/approximate_jacobian.jl

@@ -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

@@ -191,7 +191,7 @@ 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

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

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),
+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
 

test/core/rootfind_tests.jl

@@ -120,7 +120,7 @@ end
         @test all(solve(probN, NewtonRaphson(; autodiff)).u .≈ sqrt(2.0))
     end
 
-    @testset "Termination condition: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
+    @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
         u0 in (1.0, [1.0, 1.0])
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
@@ -238,7 +238,7 @@ end
         end
     end
 
-    @testset "Termination condition: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
+    @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
         u0 in (1.0, [1.0, 1.0])
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
@@ -324,7 +324,7 @@ end
         end
     end
 
-    @testset "Termination condition: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
+    @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
         u0 in (1.0, [1.0, 1.0])
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
@@ -395,7 +395,7 @@ end
         end
     end
 
-    @testset "Termination condition: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
+    @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
         u0 in (1.0, [1.0, 1.0])
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
@@ -462,7 +462,7 @@ end
                   sqrt(2.0))
     end
 
-    @testset "Termination condition: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
+    @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
         u0 in (1.0, [1.0, 1.0])
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
@@ -514,7 +514,7 @@ end
     @test nlprob_iterator_interface(quadratic_f, p, Val(false), Broyden()) ≈ sqrt.(p)
     @test nlprob_iterator_interface(quadratic_f!, p, Val(true), Broyden()) ≈ sqrt.(p)
 
-    @testset "Termination condition: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
+    @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
         u0 in (1.0, [1.0, 1.0])
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
@@ -563,7 +563,7 @@ end
     @test nlprob_iterator_interface(quadratic_f, p, Val(false), Klement()) ≈ sqrt.(p)
     @test nlprob_iterator_interface(quadratic_f!, p, Val(true), Klement()) ≈ sqrt.(p)
 
-    @testset "Termination condition: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
+    @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
         u0 in (1.0, [1.0, 1.0])
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
@@ -613,7 +613,7 @@ end
     @test nlprob_iterator_interface(
         quadratic_f!, p, Val(true), LimitedMemoryBroyden())≈sqrt.(p) atol=1e-2
 
-    @testset "Termination condition: $(termination_condition) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
+    @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS,
         u0 in (1.0, [1.0, 1.0])
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)