From d690f389b8c21177de70c5eaa28186e1d1ea033e Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 10:48:39 -0400
Subject: [PATCH 01/17] init: add jac_prototype

---
 src/systems/diffeqs/sdesystem.jl | 33 ++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 66044427eb..98c8aa6b25 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -593,6 +593,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         u0 = nothing;
         version = nothing, tgrad = false, sparse = false,
         jac = false, Wfact = false, eval_expression = false,
+        sparsity = false, analytic = nothing,
         eval_module = @__MODULE__,
         checkbounds = false, initialization_data = nothing,
         cse = true, kwargs...) where {iip, specialize}
@@ -641,6 +642,17 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         _Wfact, _Wfact_t = nothing, nothing
     end
 
+    jac_prototype = if sparse
+        uElType = u0 === nothing ? Float64 : eltype(u0)
+        if jac
+            similar(calculate_jacobian(sys, sparse = sparse), uElType)
+        else
+            similar(jacobian_sparsity(sys), uElType)
+        end
+    else
+        nothing
+    end
+
     M = calculate_massmatrix(sys)
     _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0', M)
 
@@ -651,10 +663,14 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         tgrad = _tgrad === nothing ? nothing : _tgrad,
+        mass_matrix = _M
+        jac_prototype = jac_prototype,
+        observed = observedfun,
+        sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
+        analytic = analytic,
         Wfact = _Wfact === nothing ? nothing : _Wfact,
         Wfact_t = _Wfact_t === nothing ? nothing : _Wfact_t,
-        mass_matrix = _M, initialization_data,
-        observed = observedfun)
+        initialization_data)
 end
 
 """
@@ -724,6 +740,17 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         _jac = :nothing
     end
 
+    jac_prototype = if sparse
+        uElType = u0 === nothing ? Float64 : eltype(u0)
+        if jac
+            similar(calculate_jacobian(sys, sparse = sparse), uElType)
+        else
+            similar(jacobian_sparsity(sys), uElType)
+        end
+    else
+        nothing
+    end
+
     if Wfact
         tmp_Wfact, tmp_Wfact_t = generate_factorized_W(
             sys, dvs, ps; expression = Val{true},
@@ -743,11 +770,13 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         g = $g
         tgrad = $_tgrad
         jac = $_jac
+        jac_prototype = $jac_prototype
         Wfact = $_Wfact
         Wfact_t = $_Wfact_t
         M = $_M
         SDEFunction{$iip}(f, g,
             jac = jac,
+            jac_prototype = jac_prototype,
             tgrad = tgrad,
             Wfact = Wfact,
             Wfact_t = Wfact_t,

From a71ba77435549da2ee8cf6dfd80c7879fae856de Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 10:54:59 -0400
Subject: [PATCH 02/17] typo: comma

---
 src/systems/diffeqs/sdesystem.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 98c8aa6b25..806dbdf38a 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -663,7 +663,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         tgrad = _tgrad === nothing ? nothing : _tgrad,
-        mass_matrix = _M
+        mass_matrix = _M,
         jac_prototype = jac_prototype,
         observed = observedfun,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,

From bd8913b23db9e9fe6320bb05b3708998dc2fbdfd Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 12:50:12 -0400
Subject: [PATCH 03/17] feat: add W sparsity/generation functions

---
 src/systems/diffeqs/abstractodesystem.jl | 46 +++++++++++++++++++++---
 src/systems/diffeqs/sdesystem.jl         | 13 ++++---
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 23f20b00ec..a231db040c 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -126,6 +126,33 @@ function generate_jacobian(sys::AbstractODESystem, dvs = unknowns(sys),
         dvs,
         p...,
         get_iv(sys);
+        wrap_code = sparse ? assert_jac_length_header(sys) : (identity, identity),
+        kwargs...)
+end
+
+function assert_jac_length_header(sys)
+    W = W_sparsity(sys)
+
+    identity, expr -> Func([expr.args...], [], LiteralExpr(quote
+        @assert nnz($(expr.args[1])) == nnz(W)
+        expr.body
+    end))
+end
+
+function generate_W(sys::AbstractODESystem, γ = 1., dvs = unknowns(sys),
+        ps = parameters(sys; initial_parameters = true); 
+        simplify = false, sparse = false, kwargs...)
+    @variables ˍ₋gamma
+    M = calculate_massmatrix(sys; simplify)
+    J = calculate_jacobian(sys; simplify, sparse, dvs)
+    W = ˍ₋gamma*M + J
+
+    p = reorder_parameters(sys, ps)
+    return build_function_wrapper(sys, W, 
+        dvs,
+        p...,
+        ˍ₋gamma,
+        get_iv(sys);
         kwargs...)
 end
 
@@ -264,6 +291,12 @@ function jacobian_dae_sparsity(sys::AbstractODESystem)
     J1 + J2
 end
 
+function W_sparsity(sys::AbstractODESystem) 
+    jac_sparsity = jacobian_sparsity(sys)
+    M_sparsity = sparse(iszero.(calculate_massmatrix(sys)))
+    jac_sparsity .|| M_sparsity
+end
+
 function isautonomous(sys::AbstractODESystem)
     tgrad = calculate_tgrad(sys; simplify = true)
     all(iszero, tgrad)
@@ -368,15 +401,17 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
     observedfun = ObservedFunctionCache(
         sys; steady_state, eval_expression, eval_module, checkbounds, cse)
 
-    jac_prototype = if sparse
+    if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
         if jac
-            similar(calculate_jacobian(sys, sparse = sparse), uElType)
+            jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
         else
-            similar(jacobian_sparsity(sys), uElType)
+            jac_prototype = similar(jacobian_sparsity(sys), uElType)
         end
+        W_prototype = similar(W_sparsity(sys), uElType)
     else
-        nothing
+        jac_prototype = nothing
+        W_prototype = nothing
     end
 
     @set! sys.split_idxs = split_idxs
@@ -386,7 +421,8 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         jac = _jac === nothing ? nothing : _jac,
         tgrad = _tgrad === nothing ? nothing : _tgrad,
         mass_matrix = _M,
-        jac_prototype = jac_prototype,
+        jac_prototype = W_prototype,
+        W_prototype = W_prototype,
         observed = observedfun,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
         analytic = analytic,
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 806dbdf38a..2f84b70dcd 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -642,15 +642,17 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         _Wfact, _Wfact_t = nothing, nothing
     end
 
-    jac_prototype = if sparse
+    if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
         if jac
-            similar(calculate_jacobian(sys, sparse = sparse), uElType)
+            jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
         else
-            similar(jacobian_sparsity(sys), uElType)
+            jac_prototype = similar(jacobian_sparsity(sys), uElType)
         end
+        W_prototype = similar(W_sparsity(sys), uElType)
     else
-        nothing
+        jac_prototype = nothing
+        W_prototype = nothing
     end
 
     M = calculate_massmatrix(sys)
@@ -664,7 +666,8 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         jac = _jac === nothing ? nothing : _jac,
         tgrad = _tgrad === nothing ? nothing : _tgrad,
         mass_matrix = _M,
-        jac_prototype = jac_prototype,
+        jac_prototype = W_prototype,
+        W_prototype = W_prototype,
         observed = observedfun,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
         analytic = analytic,

From 1e6b29606960b4a7cb8c11b3c33cdd63188e903e Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 13:08:24 -0400
Subject: [PATCH 04/17] fix: check actual sparsity pattern

---
 src/systems/diffeqs/abstractodesystem.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index a231db040c..63aa97ac20 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -132,9 +132,10 @@ end
 
 function assert_jac_length_header(sys)
     W = W_sparsity(sys)
-
+    (W_Is, W_Js, W_Vs) = findnz(W)
     identity, expr -> Func([expr.args...], [], LiteralExpr(quote
-        @assert nnz($(expr.args[1])) == nnz(W)
+        J_Is, J_Js, J_Vs = $(findnz)($(expr.args[1]))
+        @assert (J_Is, J_Js) == ($W_Is, $W_Js)
         expr.body
     end))
 end

From 9c05aa4999826fedd3d99e1798cab1be31a7f951 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 13:28:22 -0400
Subject: [PATCH 05/17] test: test W_sparsity

---
 src/systems/diffeqs/abstractodesystem.jl |  6 ++----
 test/jacobiansparsity.jl                 | 25 +++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 63aa97ac20..911762c765 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -132,10 +132,8 @@ end
 
 function assert_jac_length_header(sys)
     W = W_sparsity(sys)
-    (W_Is, W_Js, W_Vs) = findnz(W)
     identity, expr -> Func([expr.args...], [], LiteralExpr(quote
-        J_Is, J_Js, J_Vs = $(findnz)($(expr.args[1]))
-        @assert (J_Is, J_Js) == ($W_Is, $W_Js)
+        @assert $(findnz)($(expr.args[1]))[1:2] == $(findnz)($W)[1:2]
         expr.body
     end))
 end
@@ -294,7 +292,7 @@ end
 
 function W_sparsity(sys::AbstractODESystem) 
     jac_sparsity = jacobian_sparsity(sys)
-    M_sparsity = sparse(iszero.(calculate_massmatrix(sys)))
+    M_sparsity = sparse((!iszero).(calculate_massmatrix(sys)))
     jac_sparsity .|| M_sparsity
 end
 
diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl
index 24d47236b3..ef4a38c14f 100644
--- a/test/jacobiansparsity.jl
+++ b/test/jacobiansparsity.jl
@@ -1,4 +1,4 @@
-using OrdinaryDiffEq, ModelingToolkit, Test, SparseArrays
+using OrdinaryDiffEq, ModelingToolkit, SparseArrays
 
 N = 3
 xyd_brusselator = range(0, stop = 1, length = N)
@@ -82,3 +82,26 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = false)
 
 prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
 @test eltype(prob.f.jac_prototype) == Float32
+
+@testset "W matrix sparsity" begin
+    @parameters g
+    @variables x(t) y(t) [state_priority = 10] λ(t)
+    eqs = [D(D(x)) ~ λ * x
+           D(D(y)) ~ λ * y - g
+           x^2 + y^2 ~ 1]
+    @mtkbuild pend = ODESystem(eqs, t)
+
+    u0 = [x => 1, y => 0]
+    prob = ODEProblem(pend, u0, (0, 11.5), [g => 1], guesses = [λ => 1], sparse = true, jac = true)
+    jac, jac! = generate_jacobian(pend; expression = Val{false}, sparse = true)
+    jac_prototype = ModelingToolkit.jacobian_sparsity(pend)
+    W_prototype = ModelingToolkit.W_sparsity(pend)
+    @test nnz(W_prototype) == nnz(jac_prototype) + 2
+
+    @test findnz(prob.f.jac_prototype)[1:2] == findnz(W_prototype)[1:2]
+
+    u = zeros(5)
+    p = prob.p
+    t = 0.0
+    @test_throws AssertionError jac!(similar(jac_prototype, Float64), u, p, t)
+end

From 47c1418d0bcb7b95bd1540db77743f61aaa0ea1f Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 13:37:50 -0400
Subject: [PATCH 06/17] fix tests

---
 src/systems/diffeqs/abstractodesystem.jl | 7 +++++--
 test/jacobiansparsity.jl                 | 7 ++++---
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 911762c765..199c1de9bb 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -134,7 +134,7 @@ function assert_jac_length_header(sys)
     W = W_sparsity(sys)
     identity, expr -> Func([expr.args...], [], LiteralExpr(quote
         @assert $(findnz)($(expr.args[1]))[1:2] == $(findnz)($W)[1:2]
-        expr.body
+        $(expr.body)
     end))
 end
 
@@ -292,7 +292,10 @@ end
 
 function W_sparsity(sys::AbstractODESystem) 
     jac_sparsity = jacobian_sparsity(sys)
-    M_sparsity = sparse((!iszero).(calculate_massmatrix(sys)))
+    (n, n) = size(jac_sparsity)
+
+    M = calculate_massmatrix(sys)
+    M_sparsity = M isa UniformScaling ? sparse(I(n)) : sparse((!iszero).(M))
     jac_sparsity .|| M_sparsity
 end
 
diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl
index ef4a38c14f..28f8c66dc0 100644
--- a/test/jacobiansparsity.jl
+++ b/test/jacobiansparsity.jl
@@ -3,15 +3,15 @@ using OrdinaryDiffEq, ModelingToolkit, SparseArrays
 N = 3
 xyd_brusselator = range(0, stop = 1, length = N)
 brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0
-limit(a, N) = ModelingToolkit.ifelse(a == N + 1, 1, ModelingToolkit.ifelse(a == 0, N, a))
+lim(a, N) = ModelingToolkit.ifelse(a == N + 1, 1, ModelingToolkit.ifelse(a == 0, N, a))
 function brusselator_2d_loop(du, u, p, t)
     A, B, alpha, dx = p
     alpha = alpha / dx^2
     @inbounds for I in CartesianIndices((N, N))
         i, j = Tuple(I)
         x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
-        ip1, im1, jp1, jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N),
-        limit(j - 1, N)
+        ip1, im1, jp1, jm1 = lim(i + 1, N), lim(i - 1, N), lim(j + 1, N),
+        lim(j - 1, N)
         du[i, j, 1] = alpha * (u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] -
                        4u[i, j, 1]) +
                       B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] +
@@ -84,6 +84,7 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
 @test eltype(prob.f.jac_prototype) == Float32
 
 @testset "W matrix sparsity" begin
+    t = ModelingToolkit.t_nounits
     @parameters g
     @variables x(t) y(t) [state_priority = 10] λ(t)
     eqs = [D(D(x)) ~ λ * x

From 92c1cb90e84b7b07c4091dea4b4906b860e96d5d Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 16:24:35 -0400
Subject: [PATCH 07/17] fix: fix generate_W

---
 src/ModelingToolkit.jl                             | 2 +-
 src/structural_transformation/symbolics_tearing.jl | 5 +++++
 src/systems/diffeqs/abstractodesystem.jl           | 5 +++++
 test/jacobiansparsity.jl                           | 8 +++++++-
 4 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index b53d3ec098..bd77abfbb9 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -303,7 +303,7 @@ export structural_simplify, expand_connections, linearize, linearization_functio
        LinearizationProblem
 export solve
 
-export calculate_jacobian, generate_jacobian, generate_function, generate_custom_function
+export calculate_jacobian, generate_jacobian, generate_function, generate_custom_function, generate_W
 export calculate_control_jacobian, generate_control_jacobian
 export calculate_tgrad, generate_tgrad
 export calculate_gradient, generate_gradient
diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 6845351008..ea17b0d005 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -378,6 +378,7 @@ function generate_derivative_variables!(
     # For variable x, make dummy derivative x_t if the
     # derivative is in the system
     for v in 1:length(var_to_diff)
+        @show graph
         dv = var_to_diff[v]
         dv isa Int || continue
         solved = var_eq_matching[dv] isa Int
@@ -403,6 +404,10 @@ function generate_derivative_variables!(
         v_t = add_dd_variable!(structure, fullvars, x_t, dv)
         # Add `D(x) - x_t ~ 0` to the graph
         dummy_eq = add_dd_equation!(structure, neweqs, 0 ~ dx - x_t, dv, v_t)
+        # Update graph to say, all the equations featuring D(x) also feature x_t
+        for e in 𝑑neighbors(graph, dv)
+            add_edge!(graph, e, v_t)
+        end
 
         # Update matching
         push!(var_eq_matching, unassigned)
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 199c1de9bb..df3eb6cf92 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -143,15 +143,20 @@ function generate_W(sys::AbstractODESystem, γ = 1., dvs = unknowns(sys),
         simplify = false, sparse = false, kwargs...)
     @variables ˍ₋gamma
     M = calculate_massmatrix(sys; simplify)
+    sparse && (M = SparseArrays.sparse(M))
     J = calculate_jacobian(sys; simplify, sparse, dvs)
     W = ˍ₋gamma*M + J
+    @show W
 
     p = reorder_parameters(sys, ps)
+    @show length(p)
     return build_function_wrapper(sys, W, 
         dvs,
         p...,
         ˍ₋gamma,
         get_iv(sys);
+        wrap_code = sparse ? assert_jac_length_header(sys) : (identity, identity),
+        p_end = 1 + length(p),
         kwargs...)
 end
 
diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl
index 28f8c66dc0..ced723507b 100644
--- a/test/jacobiansparsity.jl
+++ b/test/jacobiansparsity.jl
@@ -86,7 +86,7 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
 @testset "W matrix sparsity" begin
     t = ModelingToolkit.t_nounits
     @parameters g
-    @variables x(t) y(t) [state_priority = 10] λ(t)
+    @variables x(t) y(t) λ(t)
     eqs = [D(D(x)) ~ λ * x
            D(D(y)) ~ λ * y - g
            x^2 + y^2 ~ 1]
@@ -105,4 +105,10 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
     p = prob.p
     t = 0.0
     @test_throws AssertionError jac!(similar(jac_prototype, Float64), u, p, t)
+
+    W, W! = generate_W(pend; expression = Val{false}, sparse = true)
+    γ = .1
+    M = sparse(calculate_massmatrix(pend))
+    @test_throws AssertionError W!(similar(jac_prototype, Float64), u, p, γ, t)
+    @test W!(similar(W_prototype, Float64), u, p, γ, t) == 0.1 * M + jac!(similar(W_prototype, Float64), u, p, t)
 end

From 1d3da35f166ef46dfdbe9c1c23e2efa688cdd612 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 17:01:10 -0400
Subject: [PATCH 08/17] fix

---
 src/linearization.jl                               | 1 -
 src/structural_transformation/symbolics_tearing.jl | 1 -
 src/systems/diffeqs/abstractodesystem.jl           | 2 --
 src/systems/diffeqs/odesystem.jl                   | 1 -
 test/jacobiansparsity.jl                           | 7 ++++---
 5 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/src/linearization.jl b/src/linearization.jl
index 57b171e874..77f4422b63 100644
--- a/src/linearization.jl
+++ b/src/linearization.jl
@@ -535,7 +535,6 @@ function linearize_symbolic(sys::AbstractSystem, inputs,
         if !iszero(Bs)
             if !allow_input_derivatives
                 der_inds = findall(vec(any(!iszero, Bs, dims = 1)))
-                @show typeof(der_inds)
                 error("Input derivatives appeared in expressions (-g_z\\g_u != 0), the following inputs appeared differentiated: $(ModelingToolkit.inputs(sys)[der_inds]). Call `linearize_symbolic` with keyword argument `allow_input_derivatives = true` to allow this and have the returned `B` matrix be of double width ($(2nu)), where the last $nu inputs are the derivatives of the first $nu inputs.")
             end
             B = [B [zeros(nx, nu); Bs]]
diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index ea17b0d005..552c6d13c3 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -378,7 +378,6 @@ function generate_derivative_variables!(
     # For variable x, make dummy derivative x_t if the
     # derivative is in the system
     for v in 1:length(var_to_diff)
-        @show graph
         dv = var_to_diff[v]
         dv isa Int || continue
         solved = var_eq_matching[dv] isa Int
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index df3eb6cf92..d006b8934b 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -146,10 +146,8 @@ function generate_W(sys::AbstractODESystem, γ = 1., dvs = unknowns(sys),
     sparse && (M = SparseArrays.sparse(M))
     J = calculate_jacobian(sys; simplify, sparse, dvs)
     W = ˍ₋gamma*M + J
-    @show W
 
     p = reorder_parameters(sys, ps)
-    @show length(p)
     return build_function_wrapper(sys, W, 
         dvs,
         p...,
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 57d71b18c8..599ea06b7b 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -323,7 +323,6 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
             cons = get_constraintsystem(sys)
             cons !== nothing && push!(conssystems, cons)
         end
-        @show conssystems
         @set! constraintsystem.systems = conssystems
     end
 
diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl
index ced723507b..1668788375 100644
--- a/test/jacobiansparsity.jl
+++ b/test/jacobiansparsity.jl
@@ -1,4 +1,4 @@
-using OrdinaryDiffEq, ModelingToolkit, SparseArrays
+using ModelingToolkit, SparseArrays#, OrdinaryDiffEq
 
 N = 3
 xyd_brusselator = range(0, stop = 1, length = N)
@@ -51,7 +51,7 @@ JP = prob.f.jac_prototype
 
 # test sparse jacobian
 prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
-@test_nowarn solve(prob, Rosenbrock23())
+#@test_nowarn solve(prob, Rosenbrock23())
 @test findnz(calculate_jacobian(sys, sparse = true))[1:2] ==
       findnz(prob.f.jac_prototype)[1:2]
 
@@ -74,7 +74,7 @@ f = DiffEqBase.ODEFunction(sys, u0 = nothing, sparse = true, jac = false)
 # test when u0 is not Float64
 u0 = similar(init_brusselator_2d(xyd_brusselator), Float32)
 prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
-    u0, (0.0, 11.5), p)
+                                     u0, (0.0, 11.5), p)
 sys = complete(modelingtoolkitize(prob_ode_brusselator_2d))
 
 prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = false)
@@ -85,6 +85,7 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
 
 @testset "W matrix sparsity" begin
     t = ModelingToolkit.t_nounits
+    D = ModelingToolkit.D_nounits
     @parameters g
     @variables x(t) y(t) λ(t)
     eqs = [D(D(x)) ~ λ * x

From d2e952f703a1ab75df2cee68ff9452864ee35ea7 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 18:26:06 -0400
Subject: [PATCH 09/17] fix: sdesystems don't have W_prototype

---
 src/systems/diffeqs/sdesystem.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 2f84b70dcd..78b7a02195 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -667,7 +667,6 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         tgrad = _tgrad === nothing ? nothing : _tgrad,
         mass_matrix = _M,
         jac_prototype = W_prototype,
-        W_prototype = W_prototype,
         observed = observedfun,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
         analytic = analytic,

From 7f37ea3179920bd2a8b0791ad5615ae2e6532447 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 4 Apr 2025 13:24:12 -0400
Subject: [PATCH 10/17] fix: SDE has no tearing state

---
 src/systems/diffeqs/sdesystem.jl | 25 +++++++++----------------
 test/jacobiansparsity.jl         |  3 ++-
 test/nonlinearsystem.jl          |  2 +-
 3 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 78b7a02195..605e340aa5 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -642,20 +642,16 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         _Wfact, _Wfact_t = nothing, nothing
     end
 
+    M = calculate_massmatrix(sys)
     if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
-        if jac
-            jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
-        else
-            jac_prototype = similar(jacobian_sparsity(sys), uElType)
-        end
-        W_prototype = similar(W_sparsity(sys), uElType)
+        jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
+        W_prototype = similar(jac_prototype .+ M, uElType)
     else
         jac_prototype = nothing
         W_prototype = nothing
     end
 
-    M = calculate_massmatrix(sys)
     _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0', M)
 
     observedfun = ObservedFunctionCache(
@@ -742,15 +738,14 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         _jac = :nothing
     end
 
-    jac_prototype = if sparse
+    M = calculate_massmatrix(sys)
+    if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
-        if jac
-            similar(calculate_jacobian(sys, sparse = sparse), uElType)
-        else
-            similar(jacobian_sparsity(sys), uElType)
-        end
+        jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
+        W_prototype = similar(jac_prototype .+ M, uElType)
     else
-        nothing
+        jac_prototype = nothing
+        W_prototype = nothing
     end
 
     if Wfact
@@ -763,8 +758,6 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         _Wfact, _Wfact_t = :nothing, :nothing
     end
 
-    M = calculate_massmatrix(sys)
-
     _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0', M)
 
     ex = quote
diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl
index 1668788375..f6f9853137 100644
--- a/test/jacobiansparsity.jl
+++ b/test/jacobiansparsity.jl
@@ -1,4 +1,4 @@
-using ModelingToolkit, SparseArrays#, OrdinaryDiffEq
+using ModelingToolkit, SparseArrays, OrdinaryDiffEq
 
 N = 3
 xyd_brusselator = range(0, stop = 1, length = N)
@@ -100,6 +100,7 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
     W_prototype = ModelingToolkit.W_sparsity(pend)
     @test nnz(W_prototype) == nnz(jac_prototype) + 2
 
+    # jac_prototype should be the same as W_prototype
     @test findnz(prob.f.jac_prototype)[1:2] == findnz(W_prototype)[1:2]
 
     u = zeros(5)
diff --git a/test/nonlinearsystem.jl b/test/nonlinearsystem.jl
index 9475f24006..a315371141 100644
--- a/test/nonlinearsystem.jl
+++ b/test/nonlinearsystem.jl
@@ -30,7 +30,7 @@ eqs = [0 ~ σ * (y - x) * h,
 @test eval(toexpr(ns)) == ns
 test_nlsys_inference("standard", ns, (x, y, z), (σ, ρ, β))
 @test begin
-    f = eval(generate_function(ns, [x, y, z], [σ, ρ, β])[2])
+    f = generate_function(ns, [x, y, z], [σ, ρ, β], expression = Val{false})[2]
     du = [0.0, 0.0, 0.0]
     f(du, [1, 2, 3], [1, 2, 3])
     du ≈ [1, -3, -7]

From 400dbb43d3bf79b71030a07741b73b40f526df6a Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 4 Apr 2025 17:04:32 -0400
Subject: [PATCH 11/17] fix: remove bad broadcast

---
 src/systems/diffeqs/sdesystem.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 605e340aa5..1d7dd321bc 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -739,10 +739,12 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
     end
 
     M = calculate_massmatrix(sys)
+    _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0', M)
+
     if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
         jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
-        W_prototype = similar(jac_prototype .+ M, uElType)
+        W_prototype = similar(jac_prototype + M, uElType)
     else
         jac_prototype = nothing
         W_prototype = nothing
@@ -758,7 +760,6 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         _Wfact, _Wfact_t = :nothing, :nothing
     end
 
-    _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0', M)
 
     ex = quote
         f = $f

From e21ac76764b7f6d3b662dce39f4d0b71c6f50d1c Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 4 Apr 2025 17:05:14 -0400
Subject: [PATCH 12/17] up

---
 src/systems/diffeqs/sdesystem.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 1d7dd321bc..af4273fc22 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -760,7 +760,6 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         _Wfact, _Wfact_t = :nothing, :nothing
     end
 
-
     ex = quote
         f = $f
         g = $g

From 8251a012f12966d9548f8684795fa990a3143020 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Sat, 5 Apr 2025 12:07:17 -0400
Subject: [PATCH 13/17] useW-sparsity

---
 src/systems/diffeqs/abstractodesystem.jl | 17 +++++++----------
 src/systems/diffeqs/sdesystem.jl         |  6 ++----
 2 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index d006b8934b..ca9681cd8b 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -294,9 +294,13 @@ function jacobian_dae_sparsity(sys::AbstractODESystem)
 end
 
 function W_sparsity(sys::AbstractODESystem) 
-    jac_sparsity = jacobian_sparsity(sys)
+    if has_tearing_state(sys)
+        jac_sparsity = jacobian_sparsity(sys)
+    else
+        jac = calculate_jacobian(sys; sparse = true)
+        jac_sparsity = sparse((!iszero).(jac))
+    end
     (n, n) = size(jac_sparsity)
-
     M = calculate_massmatrix(sys)
     M_sparsity = M isa UniformScaling ? sparse(I(n)) : sparse((!iszero).(M))
     jac_sparsity .|| M_sparsity
@@ -408,14 +412,8 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
 
     if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
-        if jac
-            jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
-        else
-            jac_prototype = similar(jacobian_sparsity(sys), uElType)
-        end
         W_prototype = similar(W_sparsity(sys), uElType)
     else
-        jac_prototype = nothing
         W_prototype = nothing
     end
 
@@ -427,9 +425,8 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         tgrad = _tgrad === nothing ? nothing : _tgrad,
         mass_matrix = _M,
         jac_prototype = W_prototype,
-        W_prototype = W_prototype,
         observed = observedfun,
-        sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
+        sparsity = sparsity ? W_sparsity(sys) : nothing,
         analytic = analytic,
         initialization_data)
 end
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index af4273fc22..e496b89cfa 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -645,10 +645,8 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
     M = calculate_massmatrix(sys)
     if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
-        jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
-        W_prototype = similar(jac_prototype .+ M, uElType)
+        W_prototype = similar(W_sparsity(sys), uElType)
     else
-        jac_prototype = nothing
         W_prototype = nothing
     end
 
@@ -664,7 +662,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         mass_matrix = _M,
         jac_prototype = W_prototype,
         observed = observedfun,
-        sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
+        sparsity = sparsity ? W_sparsity(sys) : nothing,
         analytic = analytic,
         Wfact = _Wfact === nothing ? nothing : _Wfact,
         Wfact_t = _Wfact_t === nothing ? nothing : _Wfact_t,

From ade11d39ffa3fefc4bdc2a2eeacd6dd6a9343b8d Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Sun, 6 Apr 2025 12:38:47 -0400
Subject: [PATCH 14/17] use W_sparsity in the SDEFunctionExpr

---
 src/systems/diffeqs/sdesystem.jl | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index e496b89cfa..2fc9749309 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -741,10 +741,8 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
 
     if sparse
         uElType = u0 === nothing ? Float64 : eltype(u0)
-        jac_prototype = similar(calculate_jacobian(sys; sparse), uElType)
-        W_prototype = similar(jac_prototype + M, uElType)
+        W_prototype = similar(W_sparsity(sys), uElType)
     else
-        jac_prototype = nothing
         W_prototype = nothing
     end
 
@@ -763,13 +761,13 @@ function SDEFunctionExpr{iip}(sys::SDESystem, dvs = unknowns(sys),
         g = $g
         tgrad = $_tgrad
         jac = $_jac
-        jac_prototype = $jac_prototype
+        W_prototype = $W_prototype
         Wfact = $_Wfact
         Wfact_t = $_Wfact_t
         M = $_M
         SDEFunction{$iip}(f, g,
             jac = jac,
-            jac_prototype = jac_prototype,
+            jac_prototype = W_prototype,
             tgrad = tgrad,
             Wfact = Wfact,
             Wfact_t = Wfact_t,

From e74e754d51eda3f6f4d5dc8926e549fd007078a8 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 7 Apr 2025 10:03:50 -0400
Subject: [PATCH 15/17] fix type instability

---
 src/systems/diffeqs/abstractodesystem.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index ca9681cd8b..986aa27508 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -298,12 +298,12 @@ function W_sparsity(sys::AbstractODESystem)
         jac_sparsity = jacobian_sparsity(sys)
     else
         jac = calculate_jacobian(sys; sparse = true)
-        jac_sparsity = sparse((!iszero).(jac))
+        jac_sparsity = SparseMatrixCSC{Bool, Int64}((!iszero).(jac))
     end
     (n, n) = size(jac_sparsity)
     M = calculate_massmatrix(sys)
-    M_sparsity = M isa UniformScaling ? sparse(I(n)) : sparse((!iszero).(M))
-    jac_sparsity .|| M_sparsity
+    M_sparsity = M isa UniformScaling ? sparse(I(n)) : SparseMatrixCSC{Bool, Int64}((!iszero).(M))
+    jac_sparsity .| M_sparsity
 end
 
 function isautonomous(sys::AbstractODESystem)

From b85ce1dafb1d32c69c467441892a8292c7039a58 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 7 Apr 2025 13:22:51 -0400
Subject: [PATCH 16/17] fix: calculate_jacobian(sparse) should use W-sparsity)

---
 src/systems/diffeqs/abstractodesystem.jl | 16 ++++++++++------
 src/systems/diffeqs/sdesystem.jl         |  6 ++++--
 src/systems/systems.jl                   |  2 +-
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 986aa27508..afe2c6d1b5 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/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))
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index 2fc9749309..3fa1302630 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/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
 
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 367e1cd4cd..2652572bfe 100644
--- a/src/systems/systems.jl
+++ b/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
 

From 66f5d594bb0d717133fae965cfeb63f7b80f55d5 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 7 Apr 2025 13:44:23 -0400
Subject: [PATCH 17/17] fix simplify

---
 src/systems/systems.jl | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 2652572bfe..0f8633f31f 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -151,11 +151,13 @@ function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = fal
         end
 
         noise_eqs = StructuralTransformations.tearing_substitute_expr(ode_sys, noise_eqs)
-        return SDESystem(Vector{Equation}(full_equations(ode_sys)), noise_eqs,
+        ssys = SDESystem(Vector{Equation}(full_equations(ode_sys)), noise_eqs,
             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), tearing_state = get_tearing_state(ode_sys))
+            guesses = guesses(sys), initialization_eqs = initialization_equations(sys))
+        @set! ssys.tearing_state = get_tearing_state(ode_sys)
+        return ssys
     end
 end