change singularity handeling, and init of J

Author: Deltadahl

src/klement.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("Klement currently only supports out-of-place nonlinear problems")
@@ -30,8 +30,19 @@ function SciMLBase.solve(prob::NonlinearProblem,
     xₙ₋₁ = x
     fₙ₋₁ = fₙ
     for _ in 1:maxiters
-        xₙ = xₙ₋₁ - inv(J) * fₙ₋₁
+        tmp = J \ 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ₙ₋₁
 
@@ -41,19 +52,11 @@ function SciMLBase.solve(prob::NonlinearProblem,
         k = (Δfₙ - J * Δxₙ) ./ denominator
         J += (k * Δxₙ' .* J) * J
 
-        # Prevent inverting singular matrix
-        if det(J) ≈ 0
-            J = ArrayInterfaceCore.zeromatrix(x) + I
+        # Singularity test
+        if cond(J) > 1e9
+            J = init_J(xₙ)
         end
 
-        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ₙ
         fₙ₋₁ = fₙ
     end

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