diff --git a/lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl b/lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl index 67721a69bf..cae69826f9 100644 --- a/lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl +++ b/lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl @@ -23,7 +23,7 @@ import StaticArraysCore: StaticArray, StaticMatrix using SciMLBase: UJacobianWrapper, UDerivativeWrapper, _vec, _unwrap_val import SciMLBase: SciMLBase, @set, DEIntegrator, ODEFunction, SplitFunction, DAEFunction, remake, solve! import SciMLOperators: SciMLOperators, update_coefficients, update_coefficients!, MatrixOperator, AbstractSciMLOperator, - islinear, isconstant + islinear, isconstant, ScalarOperator import SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm, ColoringProblem, ncolors, column_colors, coloring, sparsity_pattern import OrdinaryDiffEqCore diff --git a/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl b/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl index ce072d3f8d..d8de1a11d3 100644 --- a/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl +++ b/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl @@ -578,6 +578,13 @@ function do_newJW(integrator, alg, nlsolver, repeat_step)::NTuple{2, Bool} end end +# `ScalarOperator` (λ·I) reports `axes(mm) == ()` like `UniformScaling`, but unlike +# `UniformScaling` it isn't matched by `isa UniformScaling` -- treat both as the same +# scalar-times-identity case rather than requiring `axes(mm) == axes(W)`. +_is_scalar_massmatrix(mm) = mm isa UniformScaling || mm isa ScalarOperator +_scalar_massmatrix_λ(mm::UniformScaling) = mm.λ +_scalar_massmatrix_λ(mm::ScalarOperator) = mm.val + @noinline _throwWJerror(W, J) = throw(DimensionMismatch("W: $(axes(W)), J: $(axes(J))")) @noinline function _throwWMerror(W, mass_matrix) throw(DimensionMismatch("W: $(axes(W)), mass matrix: $(axes(mass_matrix))")) @@ -608,14 +615,14 @@ function jacobian2W!( # check size and dimension iijj = axes(W) @boundscheck (iijj == axes(J) && length(iijj) == 2) || _throwWJerror(W, J) - mass_matrix isa UniformScaling || + _is_scalar_massmatrix(mass_matrix) || @boundscheck axes(mass_matrix) == axes(W) || _throwWMerror(W, mass_matrix) @inbounds begin invdtgamma = inv(dtgamma) - if mass_matrix isa UniformScaling + if _is_scalar_massmatrix(mass_matrix) copyto!(W, J) idxs = diagind(W) - λ = -mass_matrix.λ + λ = -_scalar_massmatrix_λ(mass_matrix) if ArrayInterface.fast_scalar_indexing(J) && ArrayInterface.fast_scalar_indexing(W) @inbounds for i in 1:size(J, 1) @@ -639,14 +646,14 @@ function jacobian2W!(W::Matrix, mass_matrix, dtgamma::Number, J::Matrix)::Nothin # check size and dimension iijj = axes(W) @boundscheck (iijj == axes(J) && length(iijj) == 2) || _throwWJerror(W, J) - mass_matrix isa UniformScaling || + _is_scalar_massmatrix(mass_matrix) || @boundscheck axes(mass_matrix) == axes(W) || _throwWMerror(W, mass_matrix) @inbounds begin invdtgamma = inv(dtgamma) - if mass_matrix isa UniformScaling + if _is_scalar_massmatrix(mass_matrix) copyto!(W, J) idxs = diagind(W) - λ = -mass_matrix.λ + λ = -_scalar_massmatrix_λ(mass_matrix) @inbounds for i in 1:size(J, 1) W[i, i] = muladd(λ, invdtgamma, J[i, i]) end @@ -661,12 +668,12 @@ end function jacobian2W(mass_matrix, dtgamma::Number, J::AbstractMatrix) # check size and dimension - mass_matrix isa UniformScaling || + _is_scalar_massmatrix(mass_matrix) || @boundscheck axes(mass_matrix) == axes(J) || _throwJMerror(J, mass_matrix) @inbounds begin invdtgamma = inv(dtgamma) - if mass_matrix isa UniformScaling - λ = -mass_matrix.λ + if _is_scalar_massmatrix(mass_matrix) + λ = -_scalar_massmatrix_λ(mass_matrix) W = J + (λ * invdtgamma) * I else W = muladd(-mass_matrix, invdtgamma, J) diff --git a/lib/OrdinaryDiffEqDifferentiation/test/runtests.jl b/lib/OrdinaryDiffEqDifferentiation/test/runtests.jl index 244320ddfc..e0c4506a09 100644 --- a/lib/OrdinaryDiffEqDifferentiation/test/runtests.jl +++ b/lib/OrdinaryDiffEqDifferentiation/test/runtests.jl @@ -30,6 +30,7 @@ end # Run functional tests if TEST_GROUP ∉ ("QA", "Sparse", "ModelingToolkit") @time @safetestset "DAE jacobian2W sparse" include("dae_jacobian2w_sparse_tests.jl") + @time @safetestset "ScalarOperator mass matrix" include("scalar_operator_massmatrix_tests.jl") @time @safetestset "nzval helpers" include("nzval_helpers_tests.jl") @time @safetestset "prepare_sparse_jac!" include("prepare_sparse_jac_tests.jl") @time @safetestset "OOP J_t Tracking" include("oop_jt_tracking_test.jl") diff --git a/lib/OrdinaryDiffEqDifferentiation/test/scalar_operator_massmatrix_tests.jl b/lib/OrdinaryDiffEqDifferentiation/test/scalar_operator_massmatrix_tests.jl new file mode 100644 index 0000000000..272d7744ed --- /dev/null +++ b/lib/OrdinaryDiffEqDifferentiation/test/scalar_operator_massmatrix_tests.jl @@ -0,0 +1,25 @@ +using OrdinaryDiffEqDifferentiation +using SciMLOperators +using LinearAlgebra +using Test + +# ScalarOperator (λ·I) reports axes(mm) == (), unlike UniformScaling, so it fell +# through the `mass_matrix isa UniformScaling` special case and hit the +# `axes(mass_matrix) == axes(W)` boundscheck meant for full mass matrices. +J = [1.0 2.0; 3.0 4.0] +λ = 2.0 +W_expected = J - λ * inv(0.5) * I + +W = similar(J) +OrdinaryDiffEqDifferentiation.jacobian2W!(W, ScalarOperator(λ), 0.5, J) +@test W ≈ W_expected + +W_uniform = similar(J) +OrdinaryDiffEqDifferentiation.jacobian2W!(W_uniform, λ * I, 0.5, J) +@test W ≈ W_uniform + +@test OrdinaryDiffEqDifferentiation.jacobian2W(ScalarOperator(λ), 0.5, J) ≈ W_expected + +W_dense = Matrix(J) +OrdinaryDiffEqDifferentiation.jacobian2W!(W_dense, ScalarOperator(λ), 0.5, Matrix(J)) +@test W_dense ≈ W_expected