Switch to QR Pivoted on linear solve failure

Author: avik-pal

src/descent/dogleg.jl

@@ -81,8 +81,8 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} =
                                     want to use a Trust Region."
     δu = get_du(cache, idx)
     T = promote_type(eltype(u), eltype(fu))
-    δu_newton, _, _ = __internal_solve!(
-        cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...)
+    δu_newton = __internal_solve!(
+        cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...).δu
 
     # Newton's Step within the trust region
     if cache.internalnorm(δu_newton) ≤ trust_region
@@ -102,8 +102,8 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} =
         @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy)
         δuJᵀJδu = __dot(δu_cauchy, cache.δu_cache_mul)
     else
-        δu_cauchy, _, _ = __internal_solve!(
-            cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...)
+        δu_cauchy = __internal_solve!(
+            cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...).δu
         J_ = INV ? inv(J) : J
         l_grad = cache.internalnorm(δu_cauchy)
         @bb cache.JᵀJ_cache = J × vec(δu_cauchy)  # TODO: Rename

src/descent/geodesic_acceleration.jl

@@ -106,8 +106,8 @@ function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{
         skip_solve::Bool = false, kwargs...) where {N}
     a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx)
     skip_solve && return DescentResult(; δu, extras = (; a, v))
-    v, _, _ = __internal_solve!(
-        cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...)
+    v = __internal_solve!(
+        cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...).δu
 
     @bb @. cache.u_cache = u + cache.h * v
     cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p)
@@ -116,8 +116,8 @@ function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{
     Jv = _restructure(cache.fu_cache, cache.Jv)
     @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv)
 
-    a, _, _ = __internal_solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N);
-        skip_solve, kwargs..., reuse_A_if_factorization = true)
+    a = __internal_solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N);
+        skip_solve, kwargs..., reuse_A_if_factorization = true).δu
 
     norm_v = cache.internalnorm(v)
     norm_a = cache.internalnorm(a)

src/internal/linear_solve.jl

@@ -1,5 +1,8 @@
 import LinearSolve: AbstractFactorization, DefaultAlgorithmChoice, DefaultLinearSolver
 
+const LinearSolveFailureCode = isdefined(ReturnCode, :InternalLinearSolveFailure) ?
+                               ReturnCode.InternalLinearSolveFailure : ReturnCode.Failure
+
 """
     LinearSolverCache(alg, linsolve, A, b, u; kwargs...)
 
@@ -23,6 +26,15 @@ handled:
 
 Returns the solution of the system `u` and stores the updated cache in `cache.lincache`.
 
+#### Special Handling for Rank-deficient Matrix `A`
+
+If we detect a failure in the linear solve (mostly due to using an algorithm that doesn't
+support rank-deficient matrices), we emit a warning and attempt to solve the problem using
+Pivoted QR factorization. This is quite efficient if there are only a few rank-deficient
+that originate in the problem. However, if these are quite frequent for the main nonlinear
+system, then it is recommended to use a different linear solver that supports rank-deficient
+matrices.
+
 #### Keyword Arguments
 
   - `reuse_A_if_factorization`: If `true`, then the factorization of `A` is reused if
@@ -36,6 +48,7 @@ not mutated, we do this by copying over `A` to a preconstructed cache.
 @concrete mutable struct LinearSolverCache <: AbstractLinearSolverCache
     lincache
     linsolve
+    additional_lincache::Any
     A
     b
     precs
@@ -71,7 +84,7 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...)
        (linsolve === nothing && A isa SMatrix) ||
        (A isa Diagonal) ||
        (linsolve isa typeof(\))
-        return LinearSolverCache(nothing, nothing, A, b, nothing, 0, 0)
+        return LinearSolverCache(nothing, nothing, nothing, A, b, nothing, 0, 0)
     end
     @bb u_ = copy(u_fixed)
     linprob = LinearProblem(A, b; u0 = u_, kwargs...)
@@ -89,7 +102,7 @@ function LinearSolverCache(alg, linsolve, A, b, u; kwargs...)
     # Unalias here, we will later use these as caches
     lincache = init(linprob, linsolve; alias_A = false, alias_b = false, Pl, Pr)
 
-    return LinearSolverCache(lincache, linsolve, nothing, nothing, precs, 0, 0)
+    return LinearSolverCache(lincache, linsolve, nothing, nothing, nothing, precs, 0, 0)
 end
 
 # Direct Linear Solve Case without Caching
@@ -108,6 +121,7 @@ function (cache::LinearSolverCache{Nothing})(;
     end
     return res
 end
+
 # Use LinearSolve.jl
 function (cache::LinearSolverCache)(;
         A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing,
@@ -139,8 +153,38 @@ function (cache::LinearSolverCache)(;
         cache.lincache.Pr = Pr
     end
 
+    # display(A)
+
     linres = solve!(cache.lincache)
+    # @show cache.lincache.cacheval
+    # @show LinearAlgebra.issuccess(cache.lincache.cacheval)
     cache.lincache = linres.cache
+    # Unfortunately LinearSolve.jl doesn't have the most uniform ReturnCode handling
+    if linres.retcode === ReturnCode.Failure
+        # TODO: We need to guard this somehow because this will surely fail if A is on GPU
+        # TODO: or some fancy Matrix type
+        if !(cache.linsolve isa QRFactorization{ColumnNorm})
+            @warn "Potential Rank Deficient Matrix Detected. Attempting to solve using \
+                   Pivoted QR Factorization."
+            @assert (A !== nothing)&&(b !== nothing) "This case is not yet supported. \
+                                                      Please open an issue at \
+                                                      https://github.com/SciML/NonlinearSolve.jl"
+            if cache.additional_lincache === nothing # First time
+                linprob = LinearProblem(A, b; u0 = linres.u)
+                cache.additional_lincache = init(
+                    linprob, QRFactorization(ColumnNorm()); alias_u0 = false,
+                    alias_A = false, alias_b = false, cache.lincache.Pl, cache.lincache.Pr)
+            else
+                cache.additional_lincache.A = A
+                cache.additional_lincache.b = b
+                cache.additional_lincache.Pl = cache.lincache.Pl
+                cache.additional_lincache.Pr = cache.lincache.Pr
+            end
+            linres = solve!(cache.additional_lincache)
+            cache.additional_lincache = linres.cache
+            return linres.u
+        end
+    end
 
     return linres.u
 end