Merge pull request #2662 from jClugstor/sparse_fixes

Fix sparsity pattern mismatch and test

Author: ChrisRackauckas

Project.toml

@@ -150,6 +150,7 @@ Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
@@ -175,4 +176,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["Calculus", "ComponentArrays", "Symbolics", "AlgebraicMultigrid", "IncompleteLU", "DiffEqCallbacks", "DiffEqDevTools", "ODEProblemLibrary", "ElasticArrays", "InteractiveUtils", "ParameterizedFunctions", "PoissonRandom", "Printf", "Random", "ReverseDiff", "SafeTestsets", "SparseArrays", "Statistics", "StructArrays", "Test", "Unitful", "ModelingToolkit", "Pkg", "NLsolve", "RecursiveFactorization", "Enzyme", "SparseConnectivityTracer", "SparseMatrixColorings"]
+test = ["Calculus", "ComponentArrays", "Symbolics", "AlgebraicMultigrid", "IncompleteLU", "DiffEqCallbacks", "DifferentiationInterface", "DiffEqDevTools", "ODEProblemLibrary", "ElasticArrays", "InteractiveUtils", "ParameterizedFunctions", "PoissonRandom", "Printf", "Random", "ReverseDiff", "SafeTestsets", "SparseArrays", "Statistics", "StructArrays", "Test", "Unitful", "ModelingToolkit", "Pkg", "NLsolve", "RecursiveFactorization", "Enzyme", "SparseConnectivityTracer", "SparseMatrixColorings"]

lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl

@@ -106,7 +106,7 @@ function prepare_user_sparsity(ad_alg, prob)
     sparsity = prob.f.sparsity
 
     if !isnothing(sparsity) && !(ad_alg isa AutoSparse)
-        if sparsity isa SparseMatrixCSC
+        if sparsity isa SparseMatrixCSC && !SciMLBase.has_jac(prob.f)
             if prob.f.mass_matrix isa UniformScaling
                 idxs = diagind(sparsity)
                 @. @view(sparsity[idxs]) = 1

lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl

@@ -172,7 +172,21 @@ function calc_J!(J, integrator, cache, next_step::Bool = false)
         if DiffEqBase.has_jac(f)
             duprev = integrator.duprev
             uf = cache.uf
-            f.jac(J, duprev, uprev, p, uf.α * uf.invγdt, t)
+            # need to do some jank here to account for sparsity pattern of W
+            # https://github.com/SciML/OrdinaryDiffEq.jl/issues/2653
+
+            # we need to set all nzval to a non-zero number
+            # otherwise in the following line any zero gets interpreted as a structural zero
+            if !isnothing(integrator.f.jac_prototype) && 
+                integrator.f.jac_prototype isa SparseMatrixCSC
+
+                integrator.f.jac_prototype.nzval .= true
+                J .= true .* integrator.f.jac_prototype
+                J.nzval .= false
+                f.jac(J, duprev, uprev, p, uf.α * uf.invγdt, t)
+            else
+                f.jac(J, duprev, uprev, p, uf.α * uf.invγdt, t)
+            end
         else
             @unpack du1, uf, jac_config = cache
             # using `dz` as temporary array
@@ -183,7 +197,21 @@ function calc_J!(J, integrator, cache, next_step::Bool = false)
         end
     else
         if DiffEqBase.has_jac(f)
-            f.jac(J, uprev, p, t)
+            # need to do some jank here to account for sparsity pattern of W
+            # https://github.com/SciML/OrdinaryDiffEq.jl/issues/2653
+
+            # we need to set all nzval to a non-zero number
+            # otherwise in the following line any zero gets interpreted as a structural zero
+            if !isnothing(integrator.f.jac_prototype) &&
+               integrator.f.jac_prototype isa SparseMatrixCSC
+               
+                integrator.f.jac_prototype.nzval .= true
+                J .= true .* integrator.f.jac_prototype
+                J.nzval .= false
+                f.jac(J, uprev, p, t)
+            else 
+                f.jac(J, uprev, p, t)
+            end
         else
             @unpack du1, uf, jac_config = cache
             uf.f = nlsolve_f(f, alg)

test/interface/sparsediff_tests.jl

@@ -3,6 +3,9 @@ using OrdinaryDiffEq
 using SparseArrays
 using LinearAlgebra
 using LinearSolve
+import DifferentiationInterface as DI
+using SparseConnectivityTracer
+using SparseMatrixColorings
 using ADTypes
 using Enzyme
 
@@ -84,3 +87,26 @@ for f in [f_oop, f_ip]
     end
 end
 
+# test for https://github.com/SciML/OrdinaryDiffEq.jl/issues/2653#issuecomment-2778430025
+
+using LinearAlgebra, SparseArrays
+using OrdinaryDiffEq
+
+function f(du, u, p, t)
+    du[1] = u[1]
+    return du
+end
+
+function jac(J::SparseMatrixCSC, u, p, t)
+    @assert nnz(J) == 1  # mirrors the strict behavior of SparseMatrixColorings
+    nonzeros(J)[1] = 1
+    return J
+end
+
+u0 = ones(10)
+jac_prototype = sparse(Diagonal(vcat(1, zeros(9))))
+
+fun = ODEFunction(f; jac, jac_prototype)
+prob = ODEProblem(fun, u0, (0.0, 1.0))
+@test_nowarn sol = solve(prob, Rodas4(); reltol = 1e-8, abstol = 1e-8)
+