Merge pull request #2853 from ChrisRackauckas-Claude/sparsearrays-extension

Move SparseArrays to extension for OrdinaryDiffEqDifferentiation

Author: ChrisRackauckas

lib/OrdinaryDiffEqDifferentiation/Project.toml

@@ -15,7 +15,6 @@ ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
-SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -24,6 +23,12 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 
+[weakdeps]
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[extensions]
+OrdinaryDiffEqDifferentiationSparseArraysExt = "SparseArrays"
+
 [extras]
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
@@ -32,6 +37,7 @@ DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 ForwardDiff = "0.10.38, 1"
@@ -48,8 +54,8 @@ ConstructionBase = "1.5.8"
 LinearAlgebra = "1.10"
 SciMLBase = "2.99"
 OrdinaryDiffEqCore = "1.29.0"
-SparseArrays = "1.10"
 ConcreteStructs = "0.2"
+SparseArrays = "1.10"
 Aqua = "0.8.11"
 ArrayInterface = "7.19"
 StaticArrays = "1.9"
@@ -63,7 +69,7 @@ SafeTestsets = "0.1.0"
 SciMLOperators = "1.4"
 
 [targets]
-test = ["DiffEqDevTools", "Random", "SafeTestsets", "Test", "JET", "Aqua", "AllocCheck"]
+test = ["DiffEqDevTools", "Random", "SafeTestsets", "Test", "JET", "Aqua", "AllocCheck", "SparseArrays"]
 
 [sources.OrdinaryDiffEqCore]
 path = "../OrdinaryDiffEqCore"

lib/OrdinaryDiffEqDifferentiation/ext/OrdinaryDiffEqDifferentiationSparseArraysExt.jl

@@ -0,0 +1,19 @@
+module OrdinaryDiffEqDifferentiationSparseArraysExt
+
+using OrdinaryDiffEqDifferentiation
+import SparseArrays
+import SparseArrays: nonzeros, spzeros, SparseMatrixCSC, AbstractSparseMatrix
+
+# Override the sparse checking functions
+OrdinaryDiffEqDifferentiation.is_sparse(::AbstractSparseMatrix) = true
+OrdinaryDiffEqDifferentiation.is_sparse_csc(::SparseMatrixCSC) = true
+
+# Override the sparse array manipulation functions
+OrdinaryDiffEqDifferentiation.nonzeros(A::AbstractSparseMatrix) = nonzeros(A)
+OrdinaryDiffEqDifferentiation.spzeros(T::Type, m::Integer, n::Integer) = spzeros(T, m, n)
+
+# Helper functions for accessing sparse matrix internals
+OrdinaryDiffEqDifferentiation.get_nzval(A::SparseMatrixCSC) = A.nzval
+OrdinaryDiffEqDifferentiation.set_all_nzval!(A::SparseMatrixCSC, val) = (A.nzval .= val; A)
+
+end

lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl

@@ -13,7 +13,6 @@ import DiffEqBase
 import LinearAlgebra
 import LinearAlgebra: Diagonal, I, UniformScaling, diagind, mul!, lmul!, axpby!, opnorm, lu
 import LinearAlgebra: LowerTriangular, UpperTriangular
-import SparseArrays: SparseMatrixCSC, AbstractSparseMatrix, nonzeros, sparse, spzeros
 import ArrayInterface
 import ArrayInterface: fast_scalar_indexing, zeromatrix, lu_instance
 
@@ -67,6 +66,24 @@ else
     struct OrdinaryDiffEqTag end
 end
 
+# Functions for sparse array handling - will be overloaded by extension
+# Default implementations return false/error for non-sparse types
+is_sparse(::Any) = false
+is_sparse_csc(::Any) = false
+
+# 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
+function spzeros end
+function get_nzval end
+function set_all_nzval! end
+
+# Provide error messages if these are called without extension
+nonzeros(A) = error("SparseArrays extension not loaded. Please load SparseArrays to use sparse matrix functionality.")
+spzeros(args...) = error("SparseArrays extension not loaded. Please load SparseArrays to use sparse matrix functionality.")
+get_nzval(A) = error("SparseArrays extension not loaded. Please load SparseArrays to use sparse matrix functionality.")
+set_all_nzval!(A, val) = error("SparseArrays extension not loaded. Please load SparseArrays to use sparse matrix functionality.")
+
 include("alg_utils.jl")
 include("linsolve_utils.jl")
 include("derivative_utils.jl")

lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl

@@ -108,7 +108,7 @@ function prepare_user_sparsity(ad_alg, prob)
     sparsity = prob.f.sparsity
 
     if !isnothing(sparsity) && !(ad_alg isa AutoSparse)
-        if sparsity isa SparseMatrixCSC && !SciMLBase.has_jac(prob.f)
+        if is_sparse_csc(sparsity) && !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

@@ -179,10 +179,10 @@ function calc_J!(J, integrator, cache, next_step::Bool = false)
             # 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
+               is_sparse_csc(integrator.f.jac_prototype)
+                set_all_nzval!(integrator.f.jac_prototype, true)
                 J .= true .* integrator.f.jac_prototype
-                J.nzval .= false
+                set_all_nzval!(J, false)
                 f.jac(J, duprev, uprev, p, uf.α * uf.invγdt, t)
             else
                 f.jac(J, duprev, uprev, p, uf.α * uf.invγdt, t)
@@ -203,10 +203,10 @@ function calc_J!(J, integrator, cache, next_step::Bool = false)
             # 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
+               is_sparse_csc(integrator.f.jac_prototype)
+                set_all_nzval!(integrator.f.jac_prototype, true)
                 J .= true .* integrator.f.jac_prototype
-                J.nzval .= false
+                set_all_nzval!(J, false)
                 f.jac(J, uprev, p, t)
             else
                 f.jac(J, uprev, p, t)
@@ -278,7 +278,7 @@ mutable struct WOperator{IIP, T,
             if AJ isa AbstractMatrix
                 mm = mass_matrix isa MatrixOperator ?
                      convert(AbstractMatrix, mass_matrix) : mass_matrix
-                if AJ isa AbstractSparseMatrix
+                if is_sparse(AJ)
 
                     # If gamma is zero, then it's just an initialization and we want to make sure
                     # we get the right sparsity pattern. If gamma is not zero, then it's a case where

lib/OrdinaryDiffEqDifferentiation/src/derivative_wrappers.jl

@@ -249,7 +249,7 @@ function build_jac_config(alg, f::F1, uf::F2, du1, uprev,
         (concrete_jac(alg) !== nothing && concrete_jac(alg)))
         jac_prototype = f.jac_prototype
 
-        if jac_prototype isa SparseMatrixCSC
+        if is_sparse_csc(jac_prototype)
             if f.mass_matrix isa UniformScaling
                 idxs = diagind(jac_prototype)
                 @. @view(jac_prototype[idxs]) = 1
@@ -396,7 +396,7 @@ end
 function sparsity_colorvec(f, x)
     sparsity = f.sparsity
 
-    if sparsity isa SparseMatrixCSC
+    if is_sparse_csc(sparsity)
         if f.mass_matrix isa UniformScaling
             idxs = diagind(sparsity)
             @. @view(sparsity[idxs]) = 1