check diagonal correctly

Author: avik-pal

src/broyden.jl

@@ -164,15 +164,15 @@ function perform_step!(cache::GeneralBroydenCache{iip, IJ, UR}) where {iip, IJ,
     T = eltype(cache.u)
 
     if IJ === :true_jacobian && cache.stats.nsteps == 0
-        if __isdiag(cache.J⁻¹) && cache.J⁻¹_cache !== nothing
+        if UR === :diagonal
             cache.J⁻¹_cache = __safe_inv(jacobian!!(cache.J⁻¹_cache, cache))
             cache.J⁻¹ = __get_diagonal!!(cache.J⁻¹, cache.J⁻¹_cache)
         else
             cache.J⁻¹ = __safe_inv(jacobian!!(cache.J⁻¹, cache))
         end
     end
 
-    if __isdiag(cache.J⁻¹)
+    if UR === :diagonal
         @bb @. cache.du = cache.J⁻¹ * cache.fu
     else
         @bb cache.du = cache.J⁻¹ × vec(cache.fu)
@@ -197,7 +197,7 @@ function perform_step!(cache::GeneralBroydenCache{iip, IJ, UR}) where {iip, IJ,
             return nothing
         end
         if IJ === :true_jacobian
-            if __isdiag(cache.J⁻¹) && cache.J⁻¹_cache !== nothing
+            if UR === :diagonal
                 cache.J⁻¹_cache = __safe_inv(jacobian!!(cache.J⁻¹_cache, cache))
                 cache.J⁻¹ = __get_diagonal!!(cache.J⁻¹, cache.J⁻¹_cache)
             else

src/klement.jl

@@ -198,7 +198,7 @@ function perform_step!(cache::GeneralKlementCache{iip, IJ}) where {iip, IJ}
         cache.resets += 1
     end
 
-    if __isdiag(cache.J)
+    if IJ === :true_jacobian_diagonal || IJ === :identity
         @bb @. cache.du = cache.fu / cache.J
     else
         # u = u - J \ fu
@@ -223,7 +223,7 @@ function perform_step!(cache::GeneralKlementCache{iip, IJ}) where {iip, IJ}
 
     # Update the Jacobian
     @bb cache.du .*= -1
-    if __isdiag(cache.J)
+    if IJ === :true_jacobian_diagonal || IJ === :identity
         @bb @. cache.Jdu = (cache.J^2) * (cache.du^2)
         @bb @. cache.J += ((cache.fu - cache.fu_cache_2 - cache.J * cache.du) /
                            ifelse(iszero(cache.Jdu), T(1e-5), cache.Jdu)) * cache.du *

src/utils.jl

@@ -556,7 +556,3 @@ end
 @inline __diag(x::AbstractMatrix) = diag(x)
 @inline __diag(x::AbstractVector) = x
 @inline __diag(x::Number) = x
-
-@inline __isdiag(::AbstractVector) = true
-@inline __isdiag(::Number) = true
-@inline __isdiag(::AbstractMatrix) = false