Update Klement to exploit the diagonal structure

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.0.0"
+version = "1.0.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/nlsolve/klement.jl

@@ -13,8 +13,6 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleKlement, args...;
     T = eltype(x)
     fx = _get_fx(prob, x)
 
-    singular_tol = eps(T)^(2 // 3)
-
     abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fx, x,
         termination_condition)
 
@@ -23,42 +21,14 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleKlement, args...;
     @bb xo = copy(x)
     @bb d = copy(x)
 
-    J = __init_identity_jacobian(fx, x)
-    @bb J_cache = similar(J)
+    J = one.(x)
     @bb δx² = similar(x)
-    @bb J_cache2 = similar(J)
-    @bb F = copy(J)
 
     for _ in 1:maxiters
-        if x isa Number
-            J < singular_tol && (J = __init_identity_jacobian!!(J))
-            F_ = J
-        else
-            @bb copyto!(F, J)
-            if setindex_trait(F) === CanSetindex()
-                F_ = lu!(F; check = false)
-            else
-                F_ = lu(F; check = false)
-            end
+        any(iszero, J) && (J = __init_identity_jacobian!!(J))
 
-            # Singularity test
-            if !issuccess(F_)
-                J = __init_identity_jacobian!!(J)
-                @bb copyto!(F, J)
-                if setindex_trait(J) === CanSetindex()
-                    F_ = lu!(F; check = false)
-                else
-                    F_ = lu(F; check = false)
-                end
-            end
-        end
+        @bb @. δx = fprev / J
 
-        @bb copyto!(δx, fprev)
-        if setindex_trait(δx) === CanSetindex()
-            ldiv!(F_, _vec(δx))
-        else
-            δx = _restructure(δx, F_ \ _vec(δx))
-        end
         @bb @. x = xo - δx
         fx = __eval_f(prob, fx, x)
 
@@ -67,15 +37,8 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleKlement, args...;
         tc_sol !== nothing && return tc_sol
 
         @bb δx .*= -1
-        @bb J_cache .= transpose(J) .^ 2
-        @bb @. δx² = δx^2
-        @bb d = J_cache × vec(δx²)
-        @bb δx² = J × vec(δx)
-        @bb @. fprev = (fx - fprev - δx²) / ifelse(iszero(d), singular_tol, d)
-        @bb J_cache = vec(fprev) × transpose(_vec(δx))
-        @bb @. J_cache *= J
-        @bb J_cache2 = J_cache × J
-        @bb @. J += J_cache2
+        @bb @. δx² = δx^2 * J^2
+        @bb @. J += (fx - fprev - J * δx) / ifelse(iszero(δx²), T(1e-5), δx²) * δx * (J^2)
 
         @bb copyto!(fprev, fx)
         @bb copyto!(xo, x)

src/utils.jl

@@ -212,6 +212,10 @@ function __init_identity_jacobian!!(J)
     J[diagind(J)] .= one(eltype(J))
     return J
 end
+function __init_identity_jacobian!!(J::AbstractVector)
+    fill!(J, one(eltype(J)))
+    return J
+end
 function __init_identity_jacobian(u::StaticArray, fu)
     S1, S2 = length(fu), length(u)
     J = SMatrix{S1, S2, eltype(u)}(I)
@@ -220,6 +224,9 @@ end
 function __init_identity_jacobian!!(J::SMatrix{S1, S2}) where {S1, S2}
     return SMatrix{S1, S2, eltype(J)}(I)
 end
+function __init_identity_jacobian!!(J::SVector{S1}) where {S1}
+    return ones(SVector{S1, eltype(J)})
+end
 
 function __init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2},
         ::Val{threshold}) where {S1, S2, T1, T2, threshold}

test/23_test_problems.jl

@@ -82,7 +82,7 @@ end
     alg_ops = (SimpleKlement(),)
 
     broken_tests = Dict(alg => Int[] for alg in alg_ops)
-    broken_tests[alg_ops[1]] = [1, 2, 4, 5, 6, 7, 11, 13, 22]
+    broken_tests[alg_ops[1]] = [1, 2, 4, 5, 6, 11, 12, 22]
 
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end