refactor: centralize mass matrix and W sparsity handling

Author: AayushSabharwal

src/problems/odeproblem.jl

@@ -35,24 +35,13 @@
     end
 
     M = calculate_massmatrix(sys)
-
-    _M = if sparse && !(u0 === nothing || M === I)
-        SparseArrays.sparse(M)
-    elseif u0 === nothing || M === I
-        M
-    else
-        ArrayInterface.restructure(u0 .* u0', M)
-    end
+    _M = concrete_massmatrix(M; sparse, u0)
 
     observedfun = ObservedFunctionCache(
         sys; steady_state, eval_expression, eval_module, checkbounds, cse)
 
-    if sparse
-        uElType = u0 === nothing ? Float64 : eltype(u0)
-        W_prototype = similar(W_sparsity(sys), uElType)
-    else
-        W_prototype = nothing
-    end
+    _W_sparsity = W_sparsity(sys)
+    W_prototype = calculate_W_prototype(_W_sparsity; u0, sparse)
 
     ODEFunction{iip, spec}(f;
         sys = sys,
@@ -61,7 +50,7 @@
         mass_matrix = _M,
         jac_prototype = W_prototype,
         observed = observedfun,
-        sparsity = sparsity ? W_sparsity(sys) : nothing,
+        sparsity = sparsity ? _W_sparsity : nothing,
         analytic = analytic,
         initialization_data)
 end

src/systems/codegen.jl

@@ -240,6 +240,16 @@ function calculate_massmatrix(sys::System; simplify = false)
     M == I ? I : M
 end
 
+function concrete_massmatrix(M; sparse = false, u0 = nothing)
+    if sparse && !(u0 === nothing || M === I)
+        SparseArrays.sparse(M)
+    elseif u0 === nothing || M === I
+        M
+    else
+        ArrayInterface.restructure(u0 .* u0', M)
+    end
+end
+
 function jacobian_sparsity(sys::System)
     sparsity = torn_system_jacobian_sparsity(sys)
     sparsity === nothing || return sparsity
@@ -266,6 +276,12 @@ function W_sparsity(sys::System)
     jac_sparsity .| M_sparsity
 end
 
+function calculate_W_prototype(W_sparsity; u0 = nothing, sparse = false)
+    sparse || return nothing
+    uElType = u0 === nothing ? Float64 : eltype(u0)
+    return similar(W_sparsity, uElType)
+end
+
 function isautonomous(sys::System)
     tgrad = calculate_tgrad(sys; simplify = true)
     all(iszero, tgrad)