diff --git a/Project.toml b/Project.toml index 26d9cff86..3080b7178 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" @@ -78,4 +82,4 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" [targets] -test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "DiffEqBase"] +test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase"] 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/NonlinearSolve.jl b/src/NonlinearSolve.jl index e6caaf4da..1b18666e6 100644 --- a/src/NonlinearSolve.jl +++ b/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!, recursive_unitless_bottom_eltype 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 diff --git a/src/levenberg.jl b/src/levenberg.jl index 0aa42f5c4..1fe2feed6 100644 --- a/src/levenberg.jl +++ b/src/levenberg.jl @@ -225,7 +225,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 016d09a4b..7aa19f6ca 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -314,14 +314,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) diff --git a/test/runtests.jl b/test/runtests.jl index 248de16b9..d4f817d0a 100644 --- a/test/runtests.jl +++ b/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"