Trying to improve the Klement-solver

src/broyden.jl

@@ -16,7 +16,7 @@ function SciMLBase.solve(prob::NonlinearProblem,
     x = float(prob.u0)
     fₙ = f(x)
     T = eltype(x)
-    J⁻¹ = ArrayInterfaceCore.zeromatrix(x) + I
+    J⁻¹ = init_J(x)
 
     if SciMLBase.isinplace(prob)
         error("Broyden currently only supports out-of-place nonlinear problems")

src/klement.jl

@@ -4,7 +4,7 @@ Klement()
 ```
 
 A low-overhead implementation of [Klement](https://jatm.com.br/jatm/article/view/373).
-This method is non-allocating on scalar and static array problems.
+This method is non-allocating on scalar problems.
 """
 struct Klement <: AbstractSimpleNonlinearSolveAlgorithm end
 
@@ -16,7 +16,7 @@ function SciMLBase.solve(prob::NonlinearProblem,
     x = float(prob.u0)
     fₙ = f(x)
     T = eltype(x)
-    J = ArrayInterfaceCore.zeromatrix(x) + I
+    singular_tol = 1e-9
 
     if SciMLBase.isinplace(prob)
         error("Klement currently only supports out-of-place nonlinear problems")
@@ -29,33 +29,77 @@ function SciMLBase.solve(prob::NonlinearProblem,
     xₙ = x
     xₙ₋₁ = x
     fₙ₋₁ = fₙ
-    for _ in 1:maxiters
-        xₙ = xₙ₋₁ - inv(J) * fₙ₋₁
-        fₙ = f(xₙ)
-        Δxₙ = xₙ - xₙ₋₁
-        Δfₙ = fₙ - fₙ₋₁
 
-        # Prevent division by 0
-        denominator = max.(J' .^ 2 * Δxₙ .^ 2, 1e-9)
+    # x is scalar
+    if x isa Number
+        J = 1.0
+        for _ in 1:maxiters
+            xₙ = xₙ₋₁ - fₙ₋₁ / J
+            fₙ = f(xₙ)
 
-        k = (Δfₙ - J * Δxₙ) ./ denominator
-        J += (k * Δxₙ' .* J) * J
+            iszero(fₙ) &&
+                return SciMLBase.build_solution(prob, alg, xₙ, fₙ;
+                                                retcode = ReturnCode.Success)
 
-        # Prevent inverting singular matrix
-        if det(J) ≈ 0
-            J = ArrayInterfaceCore.zeromatrix(x) + I
+            if isapprox(xₙ, xₙ₋₁, atol = atol, rtol = rtol)
+                return SciMLBase.build_solution(prob, alg, xₙ, fₙ;
+                                                retcode = ReturnCode.Success)
+            end
+
+            Δxₙ = xₙ - xₙ₋₁
+            Δfₙ = fₙ - fₙ₋₁
+
+            # Prevent division by 0
+            denominator = max(J^2 * Δxₙ^2, 1e-9)
+
+            k = (Δfₙ - J * Δxₙ) / denominator
+            J += (k * Δxₙ * J) * J
+
+            # Singularity test
+            if J < singular_tol
+                J = 1.0
+            end
+
+            xₙ₋₁ = xₙ
+            fₙ₋₁ = fₙ
         end
+        # x is a vector
+    else
+        J = init_J(x)
+        for _ in 1:maxiters
+            F = lu(J, check = false)
+
+            # Singularity test
+            if any(abs.(F.U[diagind(F.U)]) .< singular_tol)
+                J = init_J(xₙ)
+                F = lu(J, check = false)
+            end
+
+            tmp = F \ fₙ₋₁
+            xₙ = xₙ₋₁ - tmp
+            fₙ = f(xₙ)
+
+            iszero(fₙ) &&
+                return SciMLBase.build_solution(prob, alg, xₙ, fₙ;
+                                                retcode = ReturnCode.Success)
+
+            if isapprox(xₙ, xₙ₋₁, atol = atol, rtol = rtol)
+                return SciMLBase.build_solution(prob, alg, xₙ, fₙ;
+                                                retcode = ReturnCode.Success)
+            end
+
+            Δxₙ = xₙ - xₙ₋₁
+            Δfₙ = fₙ - fₙ₋₁
+
+            # Prevent division by 0
+            denominator = max.(J' .^ 2 * Δxₙ .^ 2, 1e-9)
 
-        iszero(fₙ) &&
-            return SciMLBase.build_solution(prob, alg, xₙ, fₙ;
-                                            retcode = ReturnCode.Success)
+            k = (Δfₙ - J * Δxₙ) ./ denominator
+            J += (k * Δxₙ' .* J) * J
 
-        if isapprox(xₙ, xₙ₋₁, atol = atol, rtol = rtol)
-            return SciMLBase.build_solution(prob, alg, xₙ, fₙ;
-                                            retcode = ReturnCode.Success)
+            xₙ₋₁ = xₙ
+            fₙ₋₁ = fₙ
         end
-        xₙ₋₁ = xₙ
-        fₙ₋₁ = fₙ
     end
 
     return SciMLBase.build_solution(prob, alg, xₙ, fₙ; retcode = ReturnCode.MaxIters)

src/utils.jl

@@ -33,3 +33,13 @@ value_derivative(f::F, x::AbstractArray) where {F} = f(x), ForwardDiff.jacobian(
 value(x) = x
 value(x::Dual) = ForwardDiff.value(x)
 value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)
+
+function init_J(x)
+    J = ArrayInterfaceCore.zeromatrix(x)
+    if ismutable(x)
+        J[diagind(J)] .= one(eltype(x))
+    else
+        J += I
+    end
+    return J
+end