Fix some correctness issues

avik-pal · avik-pal · commit eadf16ff89c5 · 2023-11-30T22:38:51.000-05:00
diff --git a/src/broyden.jl b/src/broyden.jl
@@ -96,7 +96,6 @@ function perform_step!(cache::GeneralBroydenCache{iip}) where {iip}
 
     update_trace!(cache, α)
     check_and_update!(cache, cache.fu, cache.u, cache.u_cache)
-    cache.stats.nf += 1
 
     cache.force_stop && return nothing
 
@@ -114,7 +113,7 @@ function perform_step!(cache::GeneralBroydenCache{iip}) where {iip}
     else
         @bb cache.du .*= -1
         @bb cache.J⁻¹dfu = cache.J⁻¹ × vec(cache.dfu)
-        @bb cache.u_cache = cache.J⁻¹ × vec(cache.du)
+        @bb cache.u_cache = transpose(cache.J⁻¹) × vec(cache.du)
         denom = dot(cache.du, cache.J⁻¹dfu)
         @bb @. cache.du = (cache.du - cache.J⁻¹dfu) / ifelse(iszero(denom), T(1e-5), denom)
         @bb cache.J⁻¹ += vec(cache.du) × transpose(cache.u_cache)
diff --git a/src/dfsane.jl b/src/dfsane.jl
@@ -58,6 +58,7 @@ end
     alg
     u
     u_cache
+    u_cache_2
     fu
     fu_cache
     du
@@ -95,6 +96,7 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::DFSane, args.
 
     @bb du = similar(u)
     @bb u_cache = copy(u)
+    @bb u_cache_2 = similar(u)
 
     fu = evaluate_f(prob, u)
     @bb fu_cache = copy(fu)
@@ -108,48 +110,43 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::DFSane, args.
         termination_condition)
     trace = init_nonlinearsolve_trace(alg, u, fu, nothing, du; kwargs...)
 
-    return DFSaneCache{iip}(alg, u, u_cache, fu, fu_cache, du, history, f_norm, f_norm_0,
-        alg.M, T(alg.σ_1), T(alg.σ_min), T(alg.σ_max), one(T), T(alg.γ), T(alg.τ_min),
-        T(alg.τ_max), alg.n_exp, prob.p, false, maxiters, internalnorm, ReturnCode.Default,
-        abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache, trace)
+    return DFSaneCache{iip}(alg, u, u_cache, u_cache_2, fu, fu_cache, du, history, f_norm,
+        f_norm_0, alg.M, T(alg.σ_1), T(alg.σ_min), T(alg.σ_max), one(T), T(alg.γ),
+        T(alg.τ_min), T(alg.τ_max), alg.n_exp, prob.p, false, maxiters, internalnorm,
+        ReturnCode.Default, abstol, reltol, prob, NLStats(1, 0, 0, 0, 0), tc_cache, trace)
 end
 
 function perform_step!(cache::DFSaneCache{iip}) where {iip}
     @unpack alg, f_norm, σ_n, σ_min, σ_max, α_1, γ, τ_min, τ_max, n_exp, M, prob = cache
     T = eltype(cache.u)
     f_norm_old = f_norm
 
-    # Spectral parameter range check
-    σ_n = sign(σ_n) * clamp(abs(σ_n), σ_min, σ_max)
-
     # Line search direction
     @bb @. cache.du = -σ_n * cache.fu
 
-    η = alg.η_strategy(cache.f_norm_0, cache.stats.nsteps, cache.u, cache.fu)
+    η = alg.η_strategy(cache.f_norm_0, cache.stats.nsteps + 1, cache.u, cache.fu)
 
     f_bar = maximum(cache.history)
     α₊ = α_1
     α₋ = α_1
 
-    @bb axpy!(α₊, cache.du, cache.u)
-
-    evaluate_f(cache, cache.u, cache.p)
+    @bb @. cache.u_cache_2 = cache.u + α₊ * cache.du
+    evaluate_f(cache, cache.u_cache_2, cache.p)
     f_norm = cache.internalnorm(cache.fu)^n_exp
-    α = α₊
+    α = -α₊
 
     inner_converged = false
     for k in 1:(cache.alg.max_inner_iterations)
         if f_norm ≤ f_bar + η - γ * α₊^2 * f_norm_old
-            α = α₊
+            α = -α₊
             inner_converged = true
             break
         end
 
         α₊ = α₊ * clamp(α₊ * f_norm_old / (f_norm + (T(2) * α₊ - T(1)) * f_norm_old),
             τ_min, τ_max)
-        @bb axpy!(-α₋, cache.du, cache.u)
-
-        evaluate_f(cache, cache.u, cache.p)
+        @bb @. cache.u_cache_2 = cache.u - α₋ * cache.du
+        evaluate_f(cache, cache.u_cache_2, cache.p)
         f_norm = cache.internalnorm(cache.fu)^n_exp
 
         if f_norm ≤ f_bar + η - γ * α₋^2 * f_norm_old
@@ -160,9 +157,8 @@ function perform_step!(cache::DFSaneCache{iip}) where {iip}
 
         α₋ = α₋ * clamp(α₋ * f_norm_old / (f_norm + (T(2) * α₋ - T(1)) * f_norm_old),
             τ_min, τ_max)
-        @bb axpy!(α₊, cache.du, cache.u)
-
-        evaluate_f(cache, cache.u, cache.p)
+        @bb @. cache.u_cache_2 = cache.u + α₊ * cache.du
+        evaluate_f(cache, cache.u_cache_2, cache.p)
         f_norm = cache.internalnorm(cache.fu)^n_exp
     end
 
@@ -171,21 +167,20 @@ function perform_step!(cache::DFSaneCache{iip}) where {iip}
         cache.force_stop = true
     end
 
+    @bb copyto!(cache.u, cache.u_cache_2)
+
     update_trace!(cache, α)
     check_and_update!(cache, cache.fu, cache.u, cache.u_cache)
 
     # Update spectral parameter
     @bb @. cache.u_cache = cache.u - cache.u_cache
     @bb @. cache.fu_cache = cache.fu - cache.fu_cache
 
-    α₊ = sum(abs2, cache.u_cache)
-    @bb @. cache.u_cache *= cache.fu_cache
-    α₋ = sum(cache.u_cache)
-    cache.σ_n = α₊ / α₋
+    cache.σ_n = dot(cache.u_cache, cache.u_cache) / dot(cache.fu_cache, cache.u_cache)
 
     # Spectral parameter bounds check
     if !(σ_min ≤ abs(cache.σ_n) ≤ σ_max)
-        test_norm = sqrt(sum(abs2, cache.fuprev))
+        test_norm = dot(cache.fu, cache.fu)
         cache.σ_n = clamp(inv(test_norm), T(1), T(1e5))
     end
 
@@ -196,7 +191,6 @@ function perform_step!(cache::DFSaneCache{iip}) where {iip}
 
     # Update history
     cache.history[cache.stats.nsteps % M + 1] = f_norm
-    cache.stats.nf += 1
     return nothing
 end
 
diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl
@@ -142,7 +142,6 @@ function perform_step!(cache::GaussNewtonCache{iip}) where {iip}
     @bb copyto!(cache.u_cache, cache.u)
     @bb copyto!(cache.dfu, cache.fu)
 
-    cache.stats.nf += 1
     cache.stats.njacs += 1
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
diff --git a/src/klement.jl b/src/klement.jl
@@ -1,6 +1,6 @@
 """
-    GeneralKlement(; max_resets = 5, linsolve = nothing,
-                     linesearch = nothing, precs = DEFAULT_PRECS)
+    GeneralKlement(; max_resets = 5, linsolve = nothing, linesearch = nothing,
+        precs = DEFAULT_PRECS)
 
 An implementation of `Klement` with line search, preconditioning and customizable linear
 solves.
@@ -91,8 +91,8 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg_::GeneralKleme
         termination_condition)
     trace = init_nonlinearsolve_trace(alg, u, fu, J, du; kwargs...)
 
-    @bb u_cache = similar(u)
-    @bb fu_cache = similar(fu)
+    @bb u_cache = copy(u)
+    @bb fu_cache = copy(fu)
     @bb J_cache = similar(J)
     @bb J_cache_2 = similar(J)
     @bb Jdu = similar(fu)
@@ -139,7 +139,6 @@ function perform_step!(cache::GeneralKlementCache{iip}) where {iip}
 
     @bb copyto!(cache.u_cache, cache.u)
 
-    cache.stats.nf += 1
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
 
@@ -152,8 +151,8 @@ function perform_step!(cache::GeneralKlementCache{iip}) where {iip}
     @bb cache.Jdu_cache = cache.J_cache × vec(cache.Jdu)
     @bb cache.Jdu = cache.J × vec(cache.du)
     @bb @. cache.fu_cache = (cache.fu - cache.fu_cache - cache.Jdu) /
-                            max(cache.Jdu_cache, eps(real(T)))
-    @bb cache.J_cache = vec(cache.fu) × transpose(_vec(cache.du))
+                            ifelse(iszero(cache.Jdu_cache), T(1e-5), cache.Jdu_cache)
+    @bb cache.J_cache = vec(cache.fu_cache) × transpose(_vec(cache.du))
     @bb @. cache.J_cache *= cache.J
     @bb cache.J_cache_2 = cache.J_cache × cache.J
     @bb cache.J .+= cache.J_cache_2
diff --git a/src/lbroyden.jl b/src/lbroyden.jl
@@ -125,7 +125,6 @@ function perform_step!(cache::LimitedMemoryBroydenCache{iip}) where {iip}
         ApplyArray(*, Vᵀ_part, U_part), cache.du, α)
 
     check_and_update!(cache, cache.fu, cache.u, cache.u_cache)
-    cache.stats.nf += 1
 
     cache.force_stop && return nothing
 
diff --git a/src/levenberg.jl b/src/levenberg.jl
@@ -312,7 +312,6 @@ function perform_step!(cache::LevenbergMarquardtCache{true, fastls}) where {fast
         _vec(cache.δ) .= _vec(v) .+ _vec(cache.a) ./ 2
         @unpack δ, loss_old, norm_v_old, v_old, b_uphill = cache
         f(cache.fu_tmp, u .+ δ, p)
-        cache.stats.nf += 1
         loss = cache.internalnorm(cache.fu_tmp)
 
         # Condition to accept uphill steps (evaluates to `loss ≤ loss_old` in iteration 1).
@@ -411,7 +410,6 @@ function perform_step!(cache::LevenbergMarquardtCache{false, fastls}) where {fas
         cache.δ = _restructure(cache.δ, _vec(v) .+ _vec(cache.a) ./ 2)
         @unpack δ, loss_old, norm_v_old, v_old, b_uphill = cache
         fu_new = f(u .+ δ, p)
-        cache.stats.nf += 1
         loss = cache.internalnorm(fu_new)
 
         # Condition to accept uphill steps (evaluates to `loss ≤ loss_old` in iteration 1).
diff --git a/src/linesearch.jl b/src/linesearch.jl
@@ -1,5 +1,5 @@
 """
-    LineSearch(method = nothing, autodiff = nothing, alpha = true)
+    LineSearch(; method = nothing, autodiff = nothing, alpha = true)
 
 Wrapper over algorithms from
 [LineSearches.jl](https://github.com/JuliaNLSolvers/LineSearches.jl/). Allows automatic
diff --git a/src/pseudotransient.jl b/src/pseudotransient.jl
@@ -145,7 +145,6 @@ function perform_step!(cache::PseudoTransientCache{iip}) where {iip}
     check_and_update!(cache, cache.fu, cache.u, cache.u_cache)
 
     @bb copyto!(cache.u_cache, cache.u)
-    cache.stats.nf += 1
     cache.stats.njacs += 1
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
diff --git a/src/raphson.jl b/src/raphson.jl
@@ -119,7 +119,6 @@ function perform_step!(cache::NewtonRaphsonCache{iip}) where {iip}
     check_and_update!(cache, cache.fu, cache.u, cache.u_cache)
 
     @bb copyto!(cache.u_cache, cache.u)
-    cache.stats.nf += 1
     cache.stats.njacs += 1
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
diff --git a/src/trustRegion.jl b/src/trustRegion.jl
@@ -376,7 +376,6 @@ function perform_step!(cache::TrustRegionCache{iip}) where {iip}
     @bb @. cache.u_cache_2 = cache.u + cache.du
     evaluate_f(cache, cache.u_tmp, cache.p, Val{:fu_cache_2}())
     trust_region_step!(cache)
-    cache.stats.nf += 1
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
     return nothing
diff --git a/src/utils.jl b/src/utils.jl
@@ -197,6 +197,7 @@ end
 
 function evaluate_f(cache::AbstractNonlinearSolveCache, u, p,
         fu_sym::Val{FUSYM} = Val(nothing)) where {FUSYM}
+    cache.stats.nf += 1
     if FUSYM === nothing
         if isinplace(cache)
             cache.prob.f(get_fu(cache), u, p)
diff --git a/test/23_test_problems.jl b/test/23_test_problems.jl
@@ -12,7 +12,7 @@ function test_on_library(problems, dicts, alg_ops, broken_tests, ϵ = 1e-4;
         @testset "$idx: $(dict["title"])" begin
             for alg in alg_ops
                 try
-                    sol = solve(nlprob, alg;
+                    sol = solve(nlprob, alg; maxiters = 10000,
                         termination_condition = AbsNormTerminationMode())
                     problem(res, sol.u, nothing)
 
@@ -23,7 +23,8 @@ function test_on_library(problems, dicts, alg_ops, broken_tests, ϵ = 1e-4;
                     end
                     broken = idx in broken_tests[alg] ? true : false
                     @test norm(res)≤ϵ broken=broken
-                catch
+                catch err
+                    @error err
                     broken = idx in broken_tests[alg] ? true : false
                     if broken
                         @test false broken=true
@@ -83,7 +84,7 @@ end
     alg_ops = (DFSane(),)
 
     broken_tests = Dict(alg => Int[] for alg in alg_ops)
-    broken_tests[alg_ops[1]] = [1, 2, 3, 4, 5, 6, 11, 22]
+    broken_tests[alg_ops[1]] = [1, 2, 3, 5, 6, 21]
 
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end
