Fix all tests

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "2.9.0"
+version = "2.10.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -72,7 +72,7 @@ Reexport = "0.2, 1"
 SafeTestsets = "0.1"
 SciMLBase = "2.9"
 SciMLOperators = "0.3"
-SimpleNonlinearSolve = "1"
+SimpleNonlinearSolve = "0.1.23"
 SparseArrays = "<0.0.1, 1"
 SparseDiffTools = "2.14"
 StaticArrays = "1"

src/broyden.jl

@@ -116,7 +116,7 @@ function perform_step!(cache::GeneralBroydenCache{iip}) where {iip}
         @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)
+        @bb cache.J⁻¹ += vec(cache.du) × transpose(_vec(cache.u_cache))
     end
 
     @bb copyto!(cache.fu_cache, cache.fu)

src/jacobian.jl

@@ -262,6 +262,6 @@ function __lr_mul(cache::AbstractNonlinearSolveCache, JᵀJ::KrylovJᵀJ, Jᵀf)
     return dot(_vec(Jᵀf), _vec(cache.lr_mul_cache))
 end
 function __lr_mul(cache::AbstractNonlinearSolveCache, JᵀJ, Jᵀf)
-    @bb cache.lr_mul_cache = JᵀJ × Jᵀf
+    @bb cache.lr_mul_cache = JᵀJ × vec(Jᵀf)
     return dot(_vec(Jᵀf), _vec(cache.lr_mul_cache))
 end

src/lbroyden.jl

@@ -187,7 +187,7 @@ function _rmatvec!!(y, xᵀU, U, Vᵀ, x)
     x_ = vec(x)
     xᵀU_ = view(xᵀU, 1:η)
     @bb xᵀU_ = transpose(U) × x_
-    @bb y = transpose(Vᵀ) × xᵀU_
+    @bb y = transpose(Vᵀ) × vec(xᵀU_)
     @bb @. y -= x
     return y
 end
@@ -202,7 +202,7 @@ function _matvec!!(y, Vᵀx, U, Vᵀ, x)
     x_ = vec(x)
     Vᵀx_ = view(Vᵀx, 1:η)
     @bb Vᵀx_ = Vᵀ × x_
-    @bb y = U × Vᵀx_
+    @bb y = U × vec(Vᵀx_)
     @bb @. y -= x
     return y
 end

src/levenberg.jl

@@ -224,7 +224,7 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip},
     @bb u_cache = copy(u)
     @bb u_cache_2 = similar(u)
     @bb fu_cache_2 = similar(fu)
-    Jv = J * v
+    Jv = J * _vec(v)
     @bb v_cache = zero(v)
 
     return LevenbergMarquardtCache{iip, fastls}(f, alg, u, u_cache, u_cache_2, fu, fu_cache,
@@ -265,16 +265,15 @@ function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip,
         end
         linres = dolinsolve(alg.precs, linsolve; A = cache.mat_tmp,
             b = cache.rhs_tmp, linu = _vec(cache.v), cache.p, reltol = cache.abstol)
-        cache.linsolve = linres.cache
-        @bb @. cache.v = -linres.u
     else
         @bb cache.u_cache_2 = transpose(cache.J) × cache.fu
         @bb @. cache.mat_tmp = cache.JᵀJ + cache.λ * cache.DᵀD
         linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(cache.mat_tmp),
             b = _vec(cache.u_cache_2), linu = _vec(cache.v), cache.p, reltol = cache.abstol)
-        cache.linsolve = linres.cache
-        @bb @. cache.v = -linres.u
     end
+    cache.linsolve = linres.cache
+    linu = _restructure(cache.v, linres.u)
+    @bb @. cache.v = -linu
 
     update_trace!(cache.trace, cache.stats.nsteps + 1, get_u(cache), get_fu(cache), cache.J,
         cache.v)
@@ -285,9 +284,9 @@ function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip,
 
     # The following lines do: cache.a = -cache.mat_tmp \ cache.fu_tmp
     # NOTE: Don't pass `A`` in again, since we want to reuse the previous solve
-    @bb cache.Jv = cache.J × cache.v
-    @bb @. cache.fu_cache_2 = (2 / cache.h) *
-                              ((cache.fu_cache_2 - cache.fu) / cache.h - cache.Jv)
+    @bb cache.Jv = cache.J × vec(cache.v)
+    Jv = _restructure(cache.fu_cache_2, cache.Jv)
+    @bb @. cache.fu_cache_2 = (2 / cache.h) * ((cache.fu_cache_2 - cache.fu) / cache.h - Jv)
     if fastls
         if setindex_trait(cache.rhs_tmp) === CanSetindex()
             cache.rhs_tmp[1:length(cache.fu)] .= _vec(cache.fu_cache_2)
@@ -296,15 +295,14 @@ function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip,
         end
         linres = dolinsolve(alg.precs, linsolve; b = cache.rhs_tmp, linu = _vec(cache.a),
             cache.p, reltol = cache.abstol)
-        cache.linsolve = linres.cache
-        @bb @. cache.a = -linres.u
     else
         @bb cache.u_cache_2 = transpose(cache.J) × cache.fu_cache_2
         linres = dolinsolve(alg.precs, linsolve; b = _vec(cache.u_cache_2),
             linu = _vec(cache.a), cache.p, reltol = cache.abstol)
-        cache.linsolve = linres.cache
-        @bb @. cache.a = -linres.u
     end
+    cache.linsolve = linres.cache
+    linu = _restructure(cache.a, linres.u)
+    @bb @. cache.a = -linu
 
     cache.stats.nsolve += 2
     cache.stats.nfactors += 2

src/trustRegion.jl

@@ -504,15 +504,16 @@ function dogleg!(cache::TrustRegionCache{iip}) where {iip}
     # outside of trust region
     l_grad = cache.internalnorm(cache.Jᵀf) # length of the gradient
     d_cauchy = l_grad^3 / __lr_mul(cache)
+    g = _restructure(cache.du, cache.Jᵀf)
     if d_cauchy ≥ cache.trust_r
         # step to the end of the trust region
-        @bb @. cache.du = -(cache.trust_r / l_grad) * cache.Jᵀf
+        @bb @. cache.du = -(cache.trust_r / l_grad) * g
         return
     end
 
     # Take the intersection of dogleg with trust region if Cauchy point lies inside the
     # trust region
-    @bb @. cache.u_cauchy = -(d_cauchy / l_grad) * cache.Jᵀf # compute Cauchy point
+    @bb @. cache.u_cauchy = -(d_cauchy / l_grad) * g # compute Cauchy point
     @bb @. cache.u_cache_2 = cache.u_gauss_newton - cache.u_cauchy # calf of the dogleg
 
     a = dot(cache.u_cache_2, cache.u_cache_2)
@@ -547,13 +548,19 @@ function not_terminated(cache::TrustRegionCache)
     return true
 end
 
-# FIXME: Update the JacVec Operator for Yuan
+# FIXME: Reinit `JᵀJ` operator if `p` is changed
 function __reinit_internal!(cache::TrustRegionCache; kwargs...)
+    if cache.jvp_operator !== nothing
+        cache.jvp_operator = __jacvec(cache.uf, cache.u; cache.fu,
+            autodiff = __get_nonsparse_ad(cache.alg.ad))
+        @bb cache.Jᵀf = cache.jvp_operator × cache.fu
+    end
     cache.loss = __trust_region_loss(cache, cache.fu)
+    cache.loss_new = cache.loss
     cache.shrink_counter = 0
     cache.trust_r = convert(eltype(cache.u),
-        ifelse(cache.alg.initial_trust_radius == 0, cache.alg.initial_trust_radius,
-            cache.max_trust_r / 11))
+        ifelse(cache.alg.initial_trust_radius == 0, cache.max_trust_r / 11,
+            cache.alg.initial_trust_radius))
     cache.make_new_J = true
     return nothing
 end

test/23_test_problems.jl

@@ -95,10 +95,10 @@ end
     alg_ops = (GeneralBroyden(; max_resets = 10),)
 
     broken_tests = Dict(alg => Int[] for alg in alg_ops)
-    broken_tests[alg_ops[1]] = [1, 2, 4, 5, 6, 11, 12, 13, 14]
+    broken_tests[alg_ops[1]] = [1, 4, 5, 6, 11, 12, 13, 14]
 
     skip_tests = Dict(alg => Int[] for alg in alg_ops)
-    skip_tests[alg_ops[1]] = [22]
+    skip_tests[alg_ops[1]] = [2, 22]
 
     test_on_library(problems, dicts, alg_ops, broken_tests; skip_tests)
 end

test/gpu.jl

@@ -6,28 +6,22 @@ A = cu(rand(4, 4))
 u0 = cu(rand(4))
 b = cu(rand(4))
 
-function f(du, u, p)
-    du .= A * u .+ b
-end
+linear_f(du, u, p) = (du .= A * u .+ b)
 
-prob = NonlinearProblem(f, u0)
+prob = NonlinearProblem(linear_f, u0)
 
-# TrustRegion is broken
-# LimitedMemoryBroyden will diverge!
 for alg in (NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()),
     PseudoTransient(; alpha_initial = 1.0f0), GeneralKlement(), GeneralBroyden(),
-    LimitedMemoryBroyden())
+    LimitedMemoryBroyden(), TrustRegion())
     @test_nowarn sol = solve(prob, alg; abstol = 1.0f-8, reltol = 1.0f-8)
 end
 
-f(u, p) = A * u .+ b
+linear_f(u, p) = A * u .+ b
 
-prob = NonlinearProblem{false}(f, u0)
+prob = NonlinearProblem{false}(linear_f, u0)
 
-# TrustRegion is broken
-# LimitedMemoryBroyden will diverge!
 for alg in (NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()),
     PseudoTransient(; alpha_initial = 1.0f0), GeneralKlement(), GeneralBroyden(),
-    LimitedMemoryBroyden())
+    LimitedMemoryBroyden(), TrustRegion())
     @test_nowarn sol = solve(prob, alg; abstol = 1.0f-8, reltol = 1.0f-8)
 end

test/infeasible.jl

@@ -56,15 +56,10 @@ end
     @test all(!isnan, sol.u)
     @test !SciMLBase.successful_retcode(sol.retcode)
 
-    try
-        u0 = @SVector [0.0, 0.0, 0.0]
-        prob = NonlinearProblem(f1, u0)
-        sol = solve(prob)
+    u0 = @SVector [0.0, 0.0, 0.0]
+    prob = NonlinearProblem(f1, u0)
+    sol = solve(prob)
 
-        @test all(!isnan, sol.u)
-        @test !SciMLBase.successful_retcode(sol.retcode)
-    catch err
-        # Static Arrays has different default linearsolve which throws an error
-        @test err isa SingularException
-    end
+    @test all(!isnan, sol.u)
+    @test !SciMLBase.successful_retcode(sol.retcode)
 end

test/matrix_resizing.jl

@@ -8,7 +8,7 @@ prob = NonlinearProblem(ff, u0, p)
 
 for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(),
     RobustMultiNewton(), FastShortcutNonlinearPolyalg(), GeneralBroyden(), GeneralKlement(),
-    LimitedMemoryBroyden())
+    LimitedMemoryBroyden(; threshold = 2))
     @test vec(solve(prob, alg).u) == solve(vecprob, alg).u
 end
 
@@ -19,6 +19,7 @@ vecprob = NonlinearProblem(fiip, vec(u0), p)
 prob = NonlinearProblem(fiip, u0, p)
 
 for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(),
-    RobustMultiNewton(), FastShortcutNonlinearPolyalg(), GeneralBroyden(), GeneralKlement())
+    RobustMultiNewton(), FastShortcutNonlinearPolyalg(), GeneralBroyden(), GeneralKlement(),
+    LimitedMemoryBroyden(; threshold = 2))
     @test vec(solve(prob, alg).u) == solve(vecprob, alg).u
 end