Add a function to check if square A is needed

SciML · Oct 23, 2023 · 2fdc122 · 2fdc122
1 parent b4ea93f
commit 2fdc122
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 10 deletions.
diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -88,10 +88,8 @@ import PrecompileTools
         for T in (Float32, Float64)
             prob = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))
 
-            # precompile_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(),
-            #     PseudoTransient(), GeneralBroyden(), GeneralKlement(), nothing)
-            # DON'T MERGE
-            precompile_algs = ()
+            precompile_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(),
+                PseudoTransient(), GeneralBroyden(), GeneralKlement(), nothing)
 
             for alg in precompile_algs
                 solve(prob, alg, abstol = T(1e-2))

diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl
@@ -82,9 +82,7 @@ function SciMLBase.__init(prob::NonlinearLeastSquaresProblem{uType, iip}, alg_::
     alg = get_concrete_algorithm(alg_, prob)
     @unpack f, u0, p = prob
 
-    # Use QR if the user did not specify a linear solver
-    if alg.linsolve === nothing || alg.linsolve isa QRFactorization ||
-       alg.linsolve isa FastQRFactorization
+    if !needs_square_A(alg.linsolve) && !(u isa Number)
         linsolve_with_JᵀJ = Val(false)
     else
         linsolve_with_JᵀJ = Val(true)

diff --git a/src/levenberg.jl b/src/levenberg.jl
@@ -164,9 +164,7 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip},
     u = alias_u0 ? u0 : deepcopy(u0)
     fu1 = evaluate_f(prob, u)
 
-    # Use QR if the user did not specify a linear solver
-    if (alg.linsolve === nothing || alg.linsolve isa QRFactorization ||
-        alg.linsolve isa FastQRFactorization) && !(u isa Number)
+    if !needs_square_A(alg.linsolve) && !(u isa Number)
         linsolve_with_JᵀJ = Val(false)
     else
         linsolve_with_JᵀJ = Val(true)

diff --git a/src/utils.jl b/src/utils.jl
@@ -256,3 +256,24 @@ function _try_factorize_and_check_singular!(linsolve, X)
     return _issingular(X), false
 end
 _try_factorize_and_check_singular!(::Nothing, x) = _issingular(x), false
+
+# Needs Square Matrix
+"""
+    needs_square_A(alg)
+
+Returns `true` if the algorithm requires a square matrix.
+"""
+needs_square_A(::Nothing) = false
+function needs_square_A(alg)
+    try
+        A = ones(Float64, 3, 2)
+        b = ones(Float64, 3)
+        solve(LinearProblem(A, b), alg)
+        return false
+    catch err
+        return true
+    end
+end
+for alg in (:QRFactorization, :FastQRFactorization)
+    @eval needs_square_A(::$(alg)) = false
+end