From 53d4382e25e794cb620c3bb8640ce18b4ea87f23 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 10:14:31 +0000 Subject: [PATCH 1/3] Fix type instability in LinearSolve QR factorization fallback Requires https://github.com/SciML/LinearSolve.jl/pull/617 This now uses the DefaultLinearSolver sum type in order to do the QR factorization fallback. This removes the `additional_lincache` which was `::Any`. This fixes a core type instability on the `LinearSolveResult` that was causing a pretty decently large slow down to the Newton methods. This also makes it so if the user specifically says they want LU, it don't do anything different. --- .../ext/NonlinearSolveBaseLinearSolveExt.jl | 46 +------------------ lib/NonlinearSolveBase/src/linear_solve.jl | 1 - 2 files changed, 2 insertions(+), 45 deletions(-) diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl index 104c3969a..c6e88e3ab 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl @@ -21,53 +21,11 @@ function (cache::LinearSolveJLCache)(; linu !== nothing && NonlinearSolveBase.set_lincache_u!(cache, linu) linres = solve!(cache.lincache) - cache.lincache = linres.cache - # Unfortunately LinearSolve.jl doesn't have the most uniform ReturnCode handling if linres.retcode === ReturnCode.Failure - structured_mat = ArrayInterface.isstructured(cache.lincache.A) - is_gpuarray = ArrayInterface.device(cache.lincache.A) isa ArrayInterface.GPU - - if !(cache.linsolve isa QRFactorization{ColumnNorm}) && !is_gpuarray && - !structured_mat - if verbose - @warn "Potential Rank Deficient Matrix Detected. Attempting to solve using \ - Pivoted QR Factorization." - end - @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 = LinearAliasSpecifier(alias_A = false, alias_b = false) - ) - 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 - linres.retcode === ReturnCode.Failure && - return LinearSolveResult(; linres.u, success = false) - return LinearSolveResult(; linres.u) - elseif !(cache.linsolve isa QRFactorization{ColumnNorm}) - if verbose - if structured_mat || is_gpuarray - mat_desc = structured_mat ? "Structured" : "GPU" - @warn "Potential Rank Deficient Matrix Detected. But Matrix is \ - $(mat_desc). Currently, we don't attempt to solve Rank Deficient \ - $(mat_desc) Matrices. Please open an issue at \ - https://github.com/SciML/NonlinearSolve.jl" - end - end - end return LinearSolveResult(; linres.u, success = false) + else + return LinearSolveResult(; linres.u) end - - return LinearSolveResult(; linres.u) end function NonlinearSolveBase.needs_square_A(linsolve::SciMLLinearSolveAlgorithm, ::Any) diff --git a/lib/NonlinearSolveBase/src/linear_solve.jl b/lib/NonlinearSolveBase/src/linear_solve.jl index a6f0f34d9..d4469d0d0 100644 --- a/lib/NonlinearSolveBase/src/linear_solve.jl +++ b/lib/NonlinearSolveBase/src/linear_solve.jl @@ -6,7 +6,6 @@ end @concrete mutable struct LinearSolveJLCache <: AbstractLinearSolverCache lincache linsolve - additional_lincache::Any stats::NLStats end From b2d42dd38bd76e42c9d09e0a047a5db165d1181b Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 12:57:00 +0000 Subject: [PATCH 2/3] fix dispatch --- lib/NonlinearSolveBase/src/linear_solve.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/NonlinearSolveBase/src/linear_solve.jl b/lib/NonlinearSolveBase/src/linear_solve.jl index d4469d0d0..592524ec3 100644 --- a/lib/NonlinearSolveBase/src/linear_solve.jl +++ b/lib/NonlinearSolveBase/src/linear_solve.jl @@ -75,7 +75,7 @@ function construct_linear_solver(alg, linsolve, A, b, u; stats, kwargs...) # unlias here, we will later use these as caches lincache = init( linprob, linsolve; alias = LinearAliasSpecifier(alias_A = false, alias_b = false)) - return LinearSolveJLCache(lincache, linsolve, nothing, stats) + return LinearSolveJLCache(lincache, linsolve, stats) end function (cache::NativeJLLinearSolveCache)(; From ec605facdd428d269a5d234e569bc44070ec7dda Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 17:37:39 +0000 Subject: [PATCH 3/3] bump LinearSolve compat --- lib/NonlinearSolveBase/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index 44fd50228..563155a87 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -62,7 +62,7 @@ ForwardDiff = "0.10.36, 1" InteractiveUtils = "<0.0.1, 1" LineSearch = "0.1.4" LinearAlgebra = "1.10" -LinearSolve = "2.36.1, 3" +LinearSolve = "3.15" Markdown = "1.10" MaybeInplace = "0.1.4" Preferences = "1.4"