Unwrap operator mass matrices before findall in Differentiation

Seeding a sparsity pattern from the mass matrix uses findall and getindex,
which a SciMLOperator does not support, so a MatrixOperator mass matrix
combined with a jac_prototype threw MethodError: keys(::MatrixOperator).

The same pattern appeared at four places: prepare_user_sparsity and its
sparsity unwrap in alg_utils.jl, and build_jac_config and sparsity_colorvec
in derivative_wrappers.jl. Route all four through one concrete_mass_matrix
helper, which reads the operator's cached array rather than re-evaluating it,
since only the structural pattern matters and this runs before init.

Author: singhharsh1708

lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl

@@ -58,6 +58,13 @@ import DiffEqBase: OrdinaryDiffEqTag
 is_sparse(::Any) = false
 is_sparse_csc(::Any) = false
 
+# Seeding a sparsity pattern from the mass matrix uses `findall`/`getindex`, which a
+# SciMLOperator does not support. `convert` reads the operator's currently-cached array
+# rather than re-evaluating it, which is what is wanted here: only the structural
+# pattern matters, and this runs before `init`.
+concrete_mass_matrix(mm) = mm
+concrete_mass_matrix(mm::AbstractSciMLOperator) = convert(AbstractMatrix, mm)
+
 # These will error if called without the extension, but should never be called
 # on non-sparse types due to the is_sparse checks
 function nonzeros end

lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl

@@ -123,21 +123,22 @@ function prepare_user_sparsity(ad_alg, prob)
                     @. @view(jac_prototype[idxs]) = 1
                 end
             else
-                idxs = findall(!iszero, prob.f.mass_matrix)
+                mm = concrete_mass_matrix(prob.f.mass_matrix)
+                idxs = findall(!iszero, mm)
                 for idx in idxs
-                    sparsity[idx] = prob.f.mass_matrix[idx]
+                    sparsity[idx] = mm[idx]
                 end
 
                 if !isnothing(jac_prototype)
                     for idx in idxs
-                        jac_prototype[idx] = prob.f.mass_matrix[idx]
+                        jac_prototype[idx] = mm[idx]
                     end
                 end
             end
         end
 
         # KnownJacobianSparsityDetector needs an AbstractMatrix
-        sparsity = sparsity isa MatrixOperator ? sparsity.A : sparsity
+        sparsity = concrete_mass_matrix(sparsity)
 
         color_alg = SciMLBase.has_colorvec(prob.f) ?
             ConstantColoringAlgorithm(

lib/OrdinaryDiffEqDifferentiation/src/derivative_wrappers.jl

@@ -295,8 +295,9 @@ function build_jac_config(
                 idxs = diagind(jac_prototype)
                 @. @view(jac_prototype[idxs]) = 1
             else
-                idxs = findall(!iszero, f.mass_matrix)
-                @. @view(jac_prototype[idxs]) = @view(f.mass_matrix[idxs])
+                mm = concrete_mass_matrix(f.mass_matrix)
+                idxs = findall(!iszero, mm)
+                @. @view(jac_prototype[idxs]) = @view(mm[idxs])
             end
         end
 
@@ -480,8 +481,9 @@ function sparsity_colorvec(f::F, x) where {F}
             idxs = diagind(sparsity)
             @. @view(sparsity[idxs]) = 1
         else
-            idxs = findall(!iszero, f.mass_matrix)
-            @. @view(sparsity[idxs]) = @view(f.mass_matrix[idxs])
+            mm = concrete_mass_matrix(f.mass_matrix)
+            idxs = findall(!iszero, mm)
+            @. @view(sparsity[idxs]) = @view(mm[idxs])
         end
     end
 

lib/OrdinaryDiffEqDifferentiation/test/prepare_user_sparsity_tests.jl

@@ -0,0 +1,72 @@
+using OrdinaryDiffEqDifferentiation
+using SparseArrays
+using SciMLOperators
+using SciMLBase
+using ADTypes
+using Test
+
+f!(du, u, p, t) = (du .= u; nothing)
+ad_alg = AutoForwardDiff()
+
+# `findall`/`getindex` on the mass matrix need a concrete matrix; a SciMLOperator
+# (e.g. `MatrixOperator`) supports neither directly, so it must be unwrapped first.
+@testset "MatrixOperator mass matrix" begin
+    M = sparse([1.0 2.0; 0.0 1.0])
+    mass_matrix = MatrixOperator(M; update_func = (A, u, p, t) -> M)
+    jac_prototype = sparse([1.0 0.0; 0.0 1.0])
+    sparsity = copy(jac_prototype)
+    odef = ODEFunction(f!; mass_matrix, jac_prototype, sparsity)
+    prob = ODEProblem(odef, ones(2), (0.0, 1.0))
+
+    OrdinaryDiffEqDifferentiation.prepare_user_sparsity(ad_alg, prob)
+    @test Matrix(prob.f.sparsity) == Matrix(M)
+    @test Matrix(prob.f.jac_prototype) == Matrix(M)
+end
+
+@testset "UniformScaling mass matrix (unchanged path)" begin
+    jac_prototype = sparse([1.0 0.0; 0.0 1.0])
+    sparsity = copy(jac_prototype)
+    odef = ODEFunction(f!; jac_prototype, sparsity)
+    prob = ODEProblem(odef, ones(2), (0.0, 1.0))
+
+    OrdinaryDiffEqDifferentiation.prepare_user_sparsity(ad_alg, prob)
+    @test Matrix(prob.f.sparsity) == Matrix(jac_prototype)
+    @test Matrix(prob.f.jac_prototype) == Matrix(jac_prototype)
+end
+
+@testset "Plain sparse mass matrix (unchanged path)" begin
+    M = sparse([1.0 2.0; 0.0 1.0])
+    jac_prototype = sparse([1.0 0.0; 0.0 1.0])
+    sparsity = copy(jac_prototype)
+    odef = ODEFunction(f!; mass_matrix = M, jac_prototype, sparsity)
+    prob = ODEProblem(odef, ones(2), (0.0, 1.0))
+
+    OrdinaryDiffEqDifferentiation.prepare_user_sparsity(ad_alg, prob)
+    @test Matrix(prob.f.sparsity) == Matrix(M)
+    @test Matrix(prob.f.jac_prototype) == Matrix(M)
+end
+
+@testset "DiagonalOperator mass matrix" begin
+    mass_matrix = DiagonalOperator([1.0, 2.0])
+    jac_prototype = sparse([1.0 0.0; 0.0 1.0])
+    sparsity = copy(jac_prototype)
+    odef = ODEFunction(f!; mass_matrix, jac_prototype, sparsity)
+    prob = ODEProblem(odef, ones(2), (0.0, 1.0))
+
+    OrdinaryDiffEqDifferentiation.prepare_user_sparsity(ad_alg, prob)
+    @test Matrix(prob.f.sparsity) == [1.0 0.0; 0.0 2.0]
+    @test Matrix(prob.f.jac_prototype) == [1.0 0.0; 0.0 2.0]
+end
+
+# `sparsity` defaults to (and aliases) `jac_prototype`, which is how #2929's report
+# constructs the problem.
+@testset "MatrixOperator mass matrix, sparsity defaulted" begin
+    M = sparse([1.0 2.0; 0.0 1.0])
+    mass_matrix = MatrixOperator(M; update_func = (A, u, p, t) -> M)
+    jac_prototype = sparse([1.0 0.0; 0.0 1.0])
+    odef = ODEFunction(f!; mass_matrix, jac_prototype)
+    prob = ODEProblem(odef, ones(2), (0.0, 1.0))
+
+    OrdinaryDiffEqDifferentiation.prepare_user_sparsity(ad_alg, prob)
+    @test Matrix(prob.f.jac_prototype) == Matrix(M)
+end

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 "prepare_user_sparsity mass matrix" include("prepare_user_sparsity_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")