fix: move the BandedArrays extension

SciML · Oct 26, 2024 · 373b98b · 373b98b
1 parent 798c966
commit 373b98b
Show file tree

Hide file tree

Showing 7 changed files with 24 additions and 16 deletions.
diff --git a/Project.toml b/Project.toml
@@ -35,7 +35,6 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
-BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
 FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
@@ -48,7 +47,6 @@ SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 
 [extensions]
-NonlinearSolveBandedMatricesExt = "BandedMatrices"
 NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
 NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"

diff --git a/ext/NonlinearSolveBandedMatricesExt.jl b/ext/NonlinearSolveBandedMatricesExt.jl
diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml
@@ -27,13 +27,15 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 
 [extensions]
+NonlinearSolveBaseBandedMatricesExt = "BandedMatrices"
 NonlinearSolveBaseDiffEqBaseExt = "DiffEqBase"
 NonlinearSolveBaseForwardDiffExt = "ForwardDiff"
 NonlinearSolveBaseLinearSolveExt = "LinearSolve"
@@ -45,6 +47,7 @@ ADTypes = "1.9"
 Adapt = "4.1.0"
 Aqua = "0.8.7"
 ArrayInterface = "7.9"
+BandedMatrices = "1.5"
 CommonSolve = "0.2.4"
 Compat = "4.15"
 ConcreteStructs = "0.2.3"

diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseBandedMatricesExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseBandedMatricesExt.jl
@@ -0,0 +1,16 @@
+module NonlinearSolveBaseBandedMatricesExt
+
+using BandedMatrices: BandedMatrix
+using LinearAlgebra: Diagonal
+using NonlinearSolveBase: NonlinearSolveBase, Utils
+
+# This is used if we vcat a Banded Jacobian with a Diagonal Matrix in Levenberg
+@inline function Utils.faster_vcat(B::BandedMatrix, D::Diagonal)
+    if Utils.is_extension_loaded(Val(:SparseArrays))
+        @warn "Load `SparseArrays` for an optimized vcat for BandedMatrices."
+        return vcat(B, D)
+    end
+    return vcat(Utils.make_sparse(B), D)
+end
+
+end
diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl
@@ -1,7 +1,7 @@
 module NonlinearSolveBaseSparseArraysExt
 
 using NonlinearSolveBase: NonlinearSolveBase, Utils
-using SparseArrays: AbstractSparseMatrix, AbstractSparseMatrixCSC, nonzeros
+using SparseArrays: AbstractSparseMatrix, AbstractSparseMatrixCSC, nonzeros, sparse
 
 function NonlinearSolveBase.NAN_CHECK(x::AbstractSparseMatrixCSC)
     return any(NonlinearSolveBase.NAN_CHECK, nonzeros(x))
@@ -11,4 +11,6 @@ NonlinearSolveBase.sparse_or_structured_prototype(::AbstractSparseMatrix) = true
 
 Utils.maybe_symmetric(x::AbstractSparseMatrix) = x
 
+Utils.make_sparse(x) = sparse(x)
+
 end
diff --git a/lib/NonlinearSolveBase/src/abstract_types.jl b/lib/NonlinearSolveBase/src/abstract_types.jl
@@ -135,8 +135,6 @@ for fname in (:preinverted_jacobian, :normal_form)
     end
 end
 
-
-
 """
     AbstractDampingFunction
 

diff --git a/lib/NonlinearSolveBase/src/utils.jl b/lib/NonlinearSolveBase/src/utils.jl
@@ -144,4 +144,6 @@ function evaluate_f!!(f::NonlinearFunction, fu, u, p)
     return f(u, p)
 end
 
+function make_sparse end
+
 end