prepare_user_sparsity: unwrap operator mass matrices before findall (#2929)

lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl

@@ -123,14 +123,20 @@ function prepare_user_sparsity(ad_alg, prob)
                     @. @view(jac_prototype[idxs]) = 1
                 end
             else
-                idxs = findall(!iszero, prob.f.mass_matrix)
+                # `findall`/`getindex` need a concrete matrix; a SciMLOperator mass
+                # matrix (e.g. `MatrixOperator`) doesn't support either directly.
+                # `convert` here only reads the operator's currently-cached array,
+                # not a live re-evaluation (no `update_coefficients!` call).
+                mm = prob.f.mass_matrix
+                mm = mm isa AbstractSciMLOperator ? convert(AbstractMatrix, mm) : mm
+                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

lib/OrdinaryDiffEqDifferentiation/test/prepare_user_sparsity_tests.jl

@@ -0,0 +1,48 @@
+using OrdinaryDiffEqDifferentiation
+using SparseArrays
+using SciMLOperators
+using SciMLBase
+using ADTypes
+using LinearAlgebra
+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

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 "OOP J_t Tracking" include("oop_jt_tracking_test.jl")
     @time @safetestset "Differentiation Trait Tests" include("differentiation_traits_tests.jl")