Refactor to use function-based dispatch for sparse array checks

Per review feedback, replaced all `isa SparseMatrixCSC` and `isa AbstractSparseMatrix`
checks with function calls (`is_sparse_csc()` and `is_sparse()`) that default to false.

This allows the sparse code to completely compile out when SparseArrays is not loaded,
improving performance for non-sparse cases.

Changes:
- Replaced type checks with is_sparse() and is_sparse_csc() function calls
- Functions return false by default for any input
- Extension overloads them to return true for sparse types
- Removed unnecessary type aliases
- All sparse-specific code paths are now behind function checks that compile out

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

Author: claude

lib/OrdinaryDiffEqDifferentiation/ext/OrdinaryDiffEqDifferentiationSparseArraysExt.jl

@@ -2,26 +2,18 @@ module OrdinaryDiffEqDifferentiationSparseArraysExt
 
 using OrdinaryDiffEqDifferentiation
 import SparseArrays
-import SparseArrays: nonzeros, spzeros
+import SparseArrays: nonzeros, spzeros, SparseMatrixCSC, AbstractSparseMatrix
 
-# Set the type aliases when extension loads
-function __init__()
-    OrdinaryDiffEqDifferentiation.SparseMatrixCSC = SparseArrays.SparseMatrixCSC
-    OrdinaryDiffEqDifferentiation.AbstractSparseMatrix = SparseArrays.AbstractSparseMatrix
-end
+# Override the sparse checking functions
+OrdinaryDiffEqDifferentiation.is_sparse(::AbstractSparseMatrix) = true
+OrdinaryDiffEqDifferentiation.is_sparse_csc(::SparseMatrixCSC) = true
 
-# Define functions that were previously imported directly
-OrdinaryDiffEqDifferentiation.nonzeros(A::SparseArrays.AbstractSparseMatrix) = nonzeros(A)
+# 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 function to check if a type is sparse
-OrdinaryDiffEqDifferentiation.is_sparse_type(::Type{<:SparseArrays.SparseMatrixCSC}) = true
-OrdinaryDiffEqDifferentiation.is_sparse_type(::SparseArrays.SparseMatrixCSC) = true
-OrdinaryDiffEqDifferentiation.is_sparse_type(::Type{<:SparseArrays.AbstractSparseMatrix}) = true
-OrdinaryDiffEqDifferentiation.is_sparse_type(::SparseArrays.AbstractSparseMatrix) = true
-
 # Helper functions for accessing sparse matrix internals
-OrdinaryDiffEqDifferentiation.get_nzval(A::SparseArrays.SparseMatrixCSC) = A.nzval
-OrdinaryDiffEqDifferentiation.set_all_nzval!(A::SparseArrays.SparseMatrixCSC, val) = (A.nzval .= val; A)
+OrdinaryDiffEqDifferentiation.get_nzval(A::SparseMatrixCSC) = A.nzval
+OrdinaryDiffEqDifferentiation.set_all_nzval!(A::SparseMatrixCSC, val) = (A.nzval .= val; A)
 
 end

lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl

@@ -66,22 +66,21 @@ else
     struct OrdinaryDiffEqTag end
 end
 
-# Stub types and functions that will be defined by the SparseArrays extension
-# These are needed for type stability and method dispatch
-# Use a mutable placeholder that can be set by the extension
-SparseMatrixCSC = Any
-AbstractSparseMatrix = Any
+# 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
 
-# Stub functions that will be overridden by the extension
+# 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 is_sparse_type end
 function get_nzval end
 function set_all_nzval! end
 
-# Default implementations for non-sparse types
-is_sparse_type(::Type) = false
-is_sparse_type(::Any) = false
+# 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.")
 

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,7 +179,7 @@ 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
+               is_sparse_csc(integrator.f.jac_prototype)
                 set_all_nzval!(integrator.f.jac_prototype, true)
                 J .= true .* integrator.f.jac_prototype
                 set_all_nzval!(J, false)
@@ -203,7 +203,7 @@ 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
+               is_sparse_csc(integrator.f.jac_prototype)
                 set_all_nzval!(integrator.f.jac_prototype, true)
                 J .= true .* integrator.f.jac_prototype
                 set_all_nzval!(J, false)
@@ -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