fix: calculate_jacobian(sparse) should use W-sparsity)

Author: vyudu

src/systems/diffeqs/abstractodesystem.jl

@@ -73,6 +73,15 @@ function calculate_jacobian(sys::AbstractODESystem;
 
     if sparse
         jac = sparsejacobian(rhs, dvs, simplify = simplify)
+        W_s = W_sparsity(sys)
+        (Is, Js, Vs) = findnz(W_s)
+        # Add nonzeros of W as non-structural zeros of the Jacobian (to ensure equal results for oop and iip Jacobian.)
+        for (i, j) in zip(Is, Js)
+            iszero(jac[i, j]) && begin
+                jac[i, j] = 1
+                jac[i, j] = 0
+            end
+        end
     else
         jac = jacobian(rhs, dvs, simplify = simplify)
     end
@@ -294,12 +303,7 @@ function jacobian_dae_sparsity(sys::AbstractODESystem)
 end
 
 function W_sparsity(sys::AbstractODESystem) 
-    if has_tearing_state(sys)
-        jac_sparsity = jacobian_sparsity(sys)
-    else
-        jac = calculate_jacobian(sys; sparse = true)
-        jac_sparsity = SparseMatrixCSC{Bool, Int64}((!iszero).(jac))
-    end
+    jac_sparsity = jacobian_sparsity(sys)
     (n, n) = size(jac_sparsity)
     M = calculate_massmatrix(sys)
     M_sparsity = M isa UniformScaling ? sparse(I(n)) : SparseMatrixCSC{Bool, Int64}((!iszero).(M))

src/systems/diffeqs/sdesystem.jl

@@ -164,6 +164,7 @@ struct SDESystem <: AbstractODESystem
     """
     is_dde::Bool
     isscheduled::Bool
+    tearing_state::Any
 
     function SDESystem(tag, deqs, neqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed,
             tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults,
@@ -173,7 +174,8 @@ struct SDESystem <: AbstractODESystem
             metadata = nothing, gui_metadata = nothing, namespacing = true,
             complete = false, index_cache = nothing, parent = nothing, is_scalar_noise = false,
             is_dde = false,
-            isscheduled = false;
+            isscheduled = false,
+            tearing_state = nothing;
             checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
             check_independent_variables([iv])
@@ -198,7 +200,7 @@ struct SDESystem <: AbstractODESystem
             ctrl_jac, Wfact, Wfact_t, name, description, systems,
             defaults, guesses, initializesystem, initialization_eqs, connector_type, cevents,
             devents, parameter_dependencies, assertions, metadata, gui_metadata, namespacing,
-            complete, index_cache, parent, is_scalar_noise, is_dde, isscheduled)
+            complete, index_cache, parent, is_scalar_noise, is_dde, isscheduled, tearing_state)
     end
 end
 

src/systems/systems.jl

@@ -155,7 +155,7 @@ function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = fal
             get_iv(ode_sys), unknowns(ode_sys), parameters(ode_sys);
             name = nameof(ode_sys), is_scalar_noise, observed = observed(ode_sys), defaults = defaults(sys),
             parameter_dependencies = parameter_dependencies(sys), assertions = assertions(sys),
-            guesses = guesses(sys), initialization_eqs = initialization_equations(sys))
+            guesses = guesses(sys), initialization_eqs = initialization_equations(sys), tearing_state = get_tearing_state(ode_sys))
     end
 end