diff --git a/Project.toml b/Project.toml index 57698c7bf..726dedac7 100644 --- a/Project.toml +++ b/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"] diff --git a/ext/NonlinearSolveBandedMatricesExt.jl b/ext/NonlinearSolveBandedMatricesExt.jl new file mode 100644 index 000000000..7009c6273 --- /dev/null +++ b/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 diff --git a/src/levenberg.jl b/src/levenberg.jl index 6e9fd8582..b4132ec9e 100644 --- a/src/levenberg.jl +++ b/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))) diff --git a/src/utils.jl b/src/utils.jl index 6398b058b..a4231cf9a 100644 --- a/src/utils.jl +++ b/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) diff --git a/test/GPU/Project.toml b/test/GPU/Project.toml index 97cdce026..84648a7e4 100644 --- a/test/GPU/Project.toml +++ b/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" diff --git a/test/gpu.jl b/test/gpu.jl index d8ff3b8e7..43f12672a 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -1,4 +1,4 @@ -using CUDA, NonlinearSolve +using CUDA, NonlinearSolve, LinearSolve CUDA.allowscalar(false) diff --git a/test/misc.jl b/test/misc.jl new file mode 100644 index 000000000..e00a8fb01 --- /dev/null +++ b/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)