Fix jacobian2W! DimensionMismatch for ScalarOperator mass matrix (#3915)

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

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)

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")

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