diff --git a/lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl b/lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl index c901b9d929..666fbe4c5c 100644 --- a/lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqDifferentiation/test/prepare_user_sparsity_tests.jl b/lib/OrdinaryDiffEqDifferentiation/test/prepare_user_sparsity_tests.jl new file mode 100644 index 0000000000..353cd265b7 --- /dev/null +++ b/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 diff --git a/lib/OrdinaryDiffEqDifferentiation/test/runtests.jl b/lib/OrdinaryDiffEqDifferentiation/test/runtests.jl index d5acc2886a..0e8bb90ee9 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 "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")