Special Case for Banded Matrices

Project.toml

@@ -25,14 +25,17 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [weakdeps]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 
 [extensions]
+NonlinearSolveBandedMatricesExt = "BandedMatrices"
 NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
 
 [compat]
+BandedMatrices = "1"
 ADTypes = "0.2"
 ArrayInterface = "6.0.24, 7"
 ConcreteStructs = "0.2"
@@ -58,6 +61,7 @@ Zygote = "0.6"
 julia = "1.9"
 
 [extras]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
@@ -77,4 +81,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath"]
+test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices"]

ext/NonlinearSolveBandedMatricesExt.jl

@@ -0,0 +1,8 @@
+module NonlinearSolveBandedMatricesExt
+
+using BandedMatrices, LinearAlgebra, NonlinearSolve, SparseArrays
+
+# This is used if we vcat a Banded Jacobian with a Diagonal Matrix in Levenberg
+@inline NonlinearSolve._vcat(B::BandedMatrix, D::Diagonal) = vcat(sparse(B), D)
+
+end

src/NonlinearSolve.jl

@@ -4,22 +4,29 @@ if isdefined(Base, :Experimental) && isdefined(Base.Experimental, Symbol("@max_m
     @eval Base.Experimental.@max_methods 1
 end
 
-using DiffEqBase, LinearAlgebra, LinearSolve, SparseArrays, SparseDiffTools
-import ArrayInterface: restructure
-import ForwardDiff
-
-import ADTypes: AbstractFiniteDifferencesMode
-import ArrayInterface: undefmatrix, matrix_colors, parameterless_type, ismutable, issingular
-import ConcreteStructs: @concrete
-import EnumX: @enumx
-import ForwardDiff: Dual
-import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A
-import RecursiveArrayTools: ArrayPartition,
-    AbstractVectorOfArray, recursivecopy!, recursivefill!
 import Reexport: @reexport
-import SciMLBase: AbstractNonlinearAlgorithm, NLStats, _unwrap_val, has_jac, isinplace
-import StaticArraysCore: StaticArray, SVector, SArray, MArray
-import UnPack: @unpack
+import PrecompileTools
+
+PrecompileTools.@recompile_invalidations begin
+    using DiffEqBase, LinearAlgebra, LinearSolve, SparseArrays, SparseDiffTools
+    import ArrayInterface: restructure
+
+    import ADTypes: AbstractFiniteDifferencesMode
+    import ArrayInterface: undefmatrix,
+        matrix_colors, parameterless_type, ismutable, issingular
+    import ConcreteStructs: @concrete
+    import EnumX: @enumx
+    import ForwardDiff
+    import ForwardDiff: Dual
+    import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A
+    import RecursiveArrayTools: ArrayPartition,
+        AbstractVectorOfArray, recursivecopy!, recursivefill!
+    import SciMLBase: AbstractNonlinearAlgorithm, NLStats, _unwrap_val, has_jac, isinplace
+    import StaticArraysCore: StaticArray, SVector, SArray, MArray
+    import UnPack: @unpack
+
+    using ADTypes, LineSearches, SciMLBase, SimpleNonlinearSolve
+end
 
 @reexport using ADTypes, LineSearches, SciMLBase, SimpleNonlinearSolve
 
@@ -81,25 +88,17 @@ include("jacobian.jl")
 include("ad.jl")
 include("default.jl")
 
-import PrecompileTools
-
-@static if VERSION ≥ v"1.10-"
-    PrecompileTools.@compile_workload begin
-        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)
+PrecompileTools.@compile_workload begin
+    for T in (Float32, Float64)
+        probs = (NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2)),
+            NonlinearProblem{false}((u, p) -> u .* u .- p, T[0.1], T[2]),
+            NonlinearProblem{true}((du, u, p) -> du .= u .* u .- p, T[0.1], T[2]))
 
-            for alg in precompile_algs
-                solve(prob, alg, abstol = T(1e-2))
-            end
+        precompile_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(),
+            PseudoTransient(), GeneralBroyden(), GeneralKlement(), nothing)
 
-            prob = NonlinearProblem{true}((du, u, p) -> du[1] = u[1] * u[1] - p[1], T[0.1],
-                T[2])
-            for alg in precompile_algs
-                solve(prob, alg, abstol = T(1e-2))
-            end
+        for prob in probs, alg in precompile_algs
+            solve(prob, alg, abstol = T(1e-2))
         end
     end
 end

src/levenberg.jl

@@ -208,7 +208,7 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip},
         rhs_tmp = nothing
     else
         # Preserve Types
-        mat_tmp = vcat(J, DᵀD)
+        mat_tmp = _vcat(J, DᵀD)
         fill!(mat_tmp, zero(eltype(u)))
         rhs_tmp = vcat(_vec(fu1), _vec(u))
         fill!(rhs_tmp, zero(eltype(u)))

src/utils.jl

@@ -257,14 +257,17 @@ function _try_factorize_and_check_singular!(linsolve, X)
 end
 _try_factorize_and_check_singular!(::Nothing, x) = _issingular(x), false
 
-_reshape(x, args...) = reshape(x, args...)
-_reshape(x::Number, args...) = x
+@inline _reshape(x, args...) = reshape(x, args...)
+@inline _reshape(x::Number, args...) = x
 
 @generated function _axpy!(α, x, y)
     hasmethod(axpy!, Tuple{α, x, y}) && return :(axpy!(α, x, y))
     return :(@. y += α * x)
 end
 
-_needs_square_A(_, ::Number) = true
-_needs_square_A(_, ::StaticArray) = true
-_needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve)
+@inline _needs_square_A(_, ::Number) = true
+@inline _needs_square_A(_, ::StaticArray) = true
+@inline _needs_square_A(alg, _) = LinearSolve.needs_square_A(alg.linsolve)
+
+# Define special concatenation for certain Array combinations
+@inline _vcat(x, y) = vcat(x, y)

test/GPU/Project.toml

@@ -1,7 +1,9 @@
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 
 [compat]
 CUDA = "5"
+LinearSolve = "2"
 NonlinearSolve = "2"

test/gpu.jl

@@ -1,4 +1,4 @@
-using CUDA, NonlinearSolve
+using CUDA, NonlinearSolve, LinearSolve
 
 CUDA.allowscalar(false)
 

test/misc.jl

@@ -0,0 +1,7 @@
+# Miscellaneous Tests
+using BandedMatrices, LinearAlgebra, NonlinearSolve, SparseArrays, Test
+
+b = BandedMatrix(Ones(5, 5), (1, 1))
+d = Diagonal(ones(5, 5))
+
+@test NonlinearSolve._vcat(b, d) == vcat(sparse(b), d)

test/runtests.jl

@@ -17,10 +17,7 @@ end
         @time @safetestset "Sparsity Tests" include("sparse.jl")
         @time @safetestset "Polyalgs" include("polyalgs.jl")
         @time @safetestset "Matrix Resizing" include("matrix_resizing.jl")
-        if VERSION ≥ v"1.10-"
-            # Takes too long to compile on older versions
-            @time @safetestset "Nonlinear Least Squares" include("nonlinear_least_squares.jl")
-        end
+        @time @safetestset "Nonlinear Least Squares" include("nonlinear_least_squares.jl")
     end
 
     if GROUP == "All" || GROUP == "23TestProblems"