From ab25ef57e3ac4f165801faf5d490ffc670ca588e Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sun, 1 Jun 2025 13:32:23 +0530
Subject: [PATCH 01/11] refactor: remove `partial_state_selection`

---
 .../StructuralTransformations.jl              |   2 +-
 .../partial_state_selection.jl                | 169 ------------------
 test/state_selection.jl                       |  21 ---
 .../index_reduction.jl                        |  11 --
 4 files changed, 1 insertion(+), 202 deletions(-)

diff --git a/src/structural_transformation/StructuralTransformations.jl b/src/structural_transformation/StructuralTransformations.jl
index 33269f88a7..df30f7d2d2 100644
--- a/src/structural_transformation/StructuralTransformations.jl
+++ b/src/structural_transformation/StructuralTransformations.jl
@@ -55,7 +55,7 @@ using SimpleNonlinearSolve
 
 using DocStringExtensions
 
-export tearing, partial_state_selection, dae_index_lowering, check_consistency
+export tearing, dae_index_lowering, check_consistency
 export dummy_derivative
 export sorted_incidence_matrix,
        pantelides!, pantelides_reassemble, tearing_reassemble, find_solvables!,
diff --git a/src/structural_transformation/partial_state_selection.jl b/src/structural_transformation/partial_state_selection.jl
index 8a0ae5276e..887571ceb1 100644
--- a/src/structural_transformation/partial_state_selection.jl
+++ b/src/structural_transformation/partial_state_selection.jl
@@ -1,173 +1,4 @@
-function partial_state_selection_graph!(state::TransformationState)
-    find_solvables!(state; allow_symbolic = true)
-    var_eq_matching = complete(pantelides!(state))
-    complete!(state.structure)
-    partial_state_selection_graph!(state.structure, var_eq_matching)
-end
-
-function ascend_dg(xs, dg, level)
-    while level > 0
-        xs = Int[dg[x] for x in xs]
-        level -= 1
-    end
-    return xs
-end
-
-function ascend_dg_all(xs, dg, level, maxlevel)
-    r = Int[]
-    while true
-        if level <= 0
-            append!(r, xs)
-        end
-        maxlevel <= 0 && break
-        xs = Int[dg[x] for x in xs if dg[x] !== nothing]
-        level -= 1
-        maxlevel -= 1
-    end
-    return r
-end
-
-function pss_graph_modia!(structure::SystemStructure, maximal_top_matching, varlevel,
-        inv_varlevel, inv_eqlevel)
-    @unpack eq_to_diff, var_to_diff, graph, solvable_graph = structure
-
-    # var_eq_matching is a maximal matching on the top-differentiated variables.
-    # Find Strongly connected components. Note that after pantelides, we expect
-    # a balanced system, so a maximal matching should be possible.
-    var_sccs::Vector{Union{Vector{Int}, Int}} = find_var_sccs(graph, maximal_top_matching)
-    var_eq_matching = Matching{Union{Unassigned, SelectedState}}(ndsts(graph))
-    for vars in var_sccs
-        # TODO: We should have a way to not have the scc code look at unassigned vars.
-        if length(vars) == 1 && maximal_top_matching[vars[1]] === unassigned
-            continue
-        end
-
-        # Now proceed level by level from lowest to highest and tear the graph.
-        eqs = [maximal_top_matching[var]
-               for var in vars if maximal_top_matching[var] !== unassigned]
-        isempty(eqs) && continue
-        maxeqlevel = maximum(map(x -> inv_eqlevel[x], eqs))
-        maxvarlevel = level = maximum(map(x -> inv_varlevel[x], vars))
-        old_level_vars = ()
-        ict = IncrementalCycleTracker(
-            DiCMOBiGraph{true}(graph,
-                complete(Matching(ndsts(graph)), nsrcs(graph))),
-            dir = :in)
-
-        while level >= 0
-            to_tear_eqs_toplevel = filter(eq -> inv_eqlevel[eq] >= level, eqs)
-            to_tear_eqs = ascend_dg(to_tear_eqs_toplevel, invview(eq_to_diff), level)
-
-            to_tear_vars_toplevel = filter(var -> inv_varlevel[var] >= level, vars)
-            to_tear_vars = ascend_dg(to_tear_vars_toplevel, invview(var_to_diff), level)
-
-            assigned_eqs = Int[]
-
-            if old_level_vars !== ()
-                # Inherit constraints from previous level.
-                # TODO: Is this actually a good idea or do we want full freedom
-                # to tear differently on each level? Does it make a difference
-                # whether we're using heuristic or optimal tearing?
-                removed_eqs = Int[]
-                removed_vars = Int[]
-                for var in old_level_vars
-                    old_assign = var_eq_matching[var]
-                    if isa(old_assign, SelectedState)
-                        push!(removed_vars, var)
-                        continue
-                    elseif !isa(old_assign, Int) ||
-                           ict.graph.matching[var_to_diff[var]] !== unassigned
-                        continue
-                    end
-                    # Make sure the ict knows about this edge, so it doesn't accidentally introduce
-                    # a cycle.
-                    assgned_eq = eq_to_diff[old_assign]
-                    ok = try_assign_eq!(ict, var_to_diff[var], assgned_eq)
-                    @assert ok
-                    var_eq_matching[var_to_diff[var]] = assgned_eq
-                    push!(removed_eqs, eq_to_diff[ict.graph.matching[var]])
-                    push!(removed_vars, var_to_diff[var])
-                    push!(removed_vars, var)
-                end
-                to_tear_eqs = setdiff(to_tear_eqs, removed_eqs)
-                to_tear_vars = setdiff(to_tear_vars, removed_vars)
-            end
-            tearEquations!(ict, solvable_graph.fadjlist, to_tear_eqs, BitSet(to_tear_vars),
-                nothing)
-
-            for var in to_tear_vars
-                @assert var_eq_matching[var] === unassigned
-                assgned_eq = ict.graph.matching[var]
-                var_eq_matching[var] = assgned_eq
-                isa(assgned_eq, Int) && push!(assigned_eqs, assgned_eq)
-            end
-
-            if level != 0
-                remaining_vars = collect(v
-                for v in to_tear_vars
-                if var_eq_matching[v] === unassigned)
-                if !isempty(remaining_vars)
-                    remaining_eqs = setdiff(to_tear_eqs, assigned_eqs)
-                    nlsolve_matching = maximal_matching(graph,
-                        Base.Fix2(in, remaining_eqs),
-                        Base.Fix2(in, remaining_vars))
-                    for var in remaining_vars
-                        if nlsolve_matching[var] === unassigned &&
-                           var_eq_matching[var] === unassigned
-                            var_eq_matching[var] = SelectedState()
-                        end
-                    end
-                end
-            end
-
-            old_level_vars = to_tear_vars
-            level -= 1
-        end
-    end
-    return complete(var_eq_matching, nsrcs(graph))
-end
-
 struct SelectedState end
-function partial_state_selection_graph!(structure::SystemStructure, var_eq_matching)
-    @unpack eq_to_diff, var_to_diff, graph, solvable_graph = structure
-    eq_to_diff = complete(eq_to_diff)
-
-    inv_eqlevel = map(1:nsrcs(graph)) do eq
-        level = 0
-        while invview(eq_to_diff)[eq] !== nothing
-            eq = invview(eq_to_diff)[eq]
-            level += 1
-        end
-        level
-    end
-
-    varlevel = map(1:ndsts(graph)) do var
-        graph_level = level = 0
-        while var_to_diff[var] !== nothing
-            var = var_to_diff[var]
-            level += 1
-            if !isempty(𝑑neighbors(graph, var))
-                graph_level = level
-            end
-        end
-        graph_level
-    end
-
-    inv_varlevel = map(1:ndsts(graph)) do var
-        level = 0
-        while invview(var_to_diff)[var] !== nothing
-            var = invview(var_to_diff)[var]
-            level += 1
-        end
-        level
-    end
-
-    var_eq_matching = pss_graph_modia!(structure,
-        complete(var_eq_matching), varlevel, inv_varlevel,
-        inv_eqlevel)
-
-    var_eq_matching
-end
 
 function dummy_derivative_graph!(state::TransformationState, jac = nothing;
         state_priority = nothing, log = Val(false), kwargs...)
diff --git a/test/state_selection.jl b/test/state_selection.jl
index 97741110e2..fd7a840798 100644
--- a/test/state_selection.jl
+++ b/test/state_selection.jl
@@ -20,27 +20,6 @@ let dd = dummy_derivative(sys)
     @test length(unknowns(dd)) == length(equations(dd)) < 9
 end
 
-@test_skip let pss = partial_state_selection(sys)
-    @test length(equations(pss)) == 1
-    @test length(unknowns(pss)) == 2
-end
-
-@parameters σ ρ β
-@variables x(t) y(t) z(t) a(t) u(t) F(t)
-
-eqs = [D(x) ~ σ * (y - x)
-       D(y) ~ x * (ρ - z) - y + β
-       0 ~ z - x + y
-       0 ~ a + z
-       u ~ z + a]
-
-lorenz1 = System(eqs, t, name = :lorenz1)
-let al1 = alias_elimination(lorenz1)
-    let lss = partial_state_selection(al1)
-        @test length(equations(lss)) == 2
-    end
-end
-
 # 1516
 let
     @connector function Fluid_port(; name, p = 101325.0, m = 0.0, T = 293.15)
diff --git a/test/structural_transformation/index_reduction.jl b/test/structural_transformation/index_reduction.jl
index 94dd1f884b..229b1f5736 100644
--- a/test/structural_transformation/index_reduction.jl
+++ b/test/structural_transformation/index_reduction.jl
@@ -31,12 +31,6 @@ eqs2 = [D(D(x)) ~ T * x,
     0 ~ x^2 + y^2 - L^2]
 pendulum2 = System(eqs2, t, [x, y, T], [L, g], name = :pendulum)
 
-@test_skip begin
-    let pss_pendulum2 = partial_state_selection(pendulum2)
-        length(equations(pss_pendulum2)) <= 6
-    end
-end
-
 eqs = [D(x) ~ w,
     D(y) ~ z,
     D(w) ~ T * x,
@@ -44,11 +38,6 @@ eqs = [D(x) ~ w,
     0 ~ x^2 + y^2 - L^2]
 pendulum = System(eqs, t, [x, y, w, z, T], [L, g], name = :pendulum)
 
-let pss_pendulum = partial_state_selection(pendulum)
-    # This currently selects `T` rather than `x` at top level. Needs tearing priorities to fix.
-    @test_broken length(equations(pss_pendulum)) == 3
-end
-
 let sys = mtkcompile(pendulum2)
     @test length(equations(sys)) == 5
     @test length(unknowns(sys)) == 5

From caf6f2ea3e99fd8740d6dd249306da72f1fa16e2 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sun, 1 Jun 2025 13:32:50 +0530
Subject: [PATCH 02/11] refactor: always return all information from
 `dummy_derivative_graph!`

---
 src/structural_transformation/partial_state_selection.jl | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/structural_transformation/partial_state_selection.jl b/src/structural_transformation/partial_state_selection.jl
index 887571ceb1..89447a4e16 100644
--- a/src/structural_transformation/partial_state_selection.jl
+++ b/src/structural_transformation/partial_state_selection.jl
@@ -174,11 +174,7 @@ function dummy_derivative_graph!(
     end
 
     ret = tearing_with_dummy_derivatives(structure, BitSet(dummy_derivatives))
-    if log
-        (ret..., DummyDerivativeSummary(var_dummy_scc, var_state_priority))
-    else
-        ret[1]
-    end
+    (ret..., DummyDerivativeSummary(var_dummy_scc, var_state_priority))
 end
 
 function is_present(structure, v)::Bool

From 9db7e6aa0d898f59ee3554d24388e266f25b40fa Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sun, 1 Jun 2025 13:33:25 +0530
Subject: [PATCH 03/11] refactor: sort SCCs in `find_var_sccs`

---
 src/structural_transformation/utils.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 131010550e..f36584744a 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -153,6 +153,9 @@ function find_var_sccs(g::BipartiteGraph, assign = nothing)
     cmog = DiCMOBiGraph{true}(g,
         Matching(assign === nothing ? Base.OneTo(nsrcs(g)) : assign))
     sccs = Graphs.strongly_connected_components(cmog)
+    cgraph = MatchedCondensationGraph(cmog, sccs)
+    toporder = topological_sort(cgraph)
+    permute!(sccs, toporder)
     foreach(sort!, sccs)
     return sccs
 end

From eac59a14d384f19c95fda130ccfd244cd5c339b5 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sun, 1 Jun 2025 13:33:42 +0530
Subject: [PATCH 04/11] fix: fix `but_sorted_incidence`

---
 src/structural_transformation/utils.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index f36584744a..e900764698 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -352,6 +352,7 @@ function but_ordered_incidence(ts::TearingState, varmask = highest_order_variabl
         isemptyc || push!(bb, l)
     end
     mm = incidence_matrix(graph)
+    reverse!(vordering)
     mm[[var_eq_matching[v] for v in vordering if var_eq_matching[v] isa Int], vordering], bb
 end
 

From 4ecbc2c8a3bdddc40ba5ebd5f335a85992c4a9ac Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sun, 1 Jun 2025 13:37:03 +0530
Subject: [PATCH 05/11] refactor: BLT sort equations and variables in tearing

---
 .../symbolics_tearing.jl                      | 537 +++++++++++++-----
 src/systems/system.jl                         |   7 +-
 src/systems/systemstructure.jl                |   4 +-
 3 files changed, 404 insertions(+), 144 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index f6c394ee1c..0014a947c8 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -317,11 +317,16 @@ Effects on the system structure:
 - fullvars: add the new derivative variables x_t
 - neweqs: add the identity equations for the new variables, D(x) ~ x_t
 - graph: update graph with the new equations and variables, and their connections
-- solvable_graph:
-- var_eq_matching: match D(x) to the added identity equation D(x) ~ x_t
+- solvable_graph: mark the new equation as solvable for `D(x)`
+- var_eq_matching: match D(x) to the added identity equation `D(x) ~ x_t`
+- full_var_eq_matching: match `x_t` to the equation that `D(x)` used to match to, and
+  match `D(x)` to `D(x) ~ x_t`
+- var_sccs: Replace `D(x)` in its SCC by `x_t`, and add `D(x)` in its own SCC. Return
+  the new list of SCCs.
 """
 function generate_derivative_variables!(
-        ts::TearingState, neweqs, var_eq_matching; mm = nothing, iv = nothing, D = nothing)
+        ts::TearingState, neweqs, var_eq_matching, full_var_eq_matching,
+        var_sccs; mm, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
@@ -330,44 +335,117 @@ function generate_derivative_variables!(
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
+    # We need the inverse mapping of `var_sccs` to update it efficiently later.
+    v_to_scc = Vector{NTuple{2, Int}}(undef, ndsts(graph))
+    for (i, scc) in enumerate(var_sccs), (j, v) in enumerate(scc)
+        v_to_scc[v] = (i, j)
+    end
+    # Pairs of `(x_t, dx)` added below
+    v_t_dvs = NTuple{2, Int}[]
+
     # For variable x, make dummy derivative x_t if the
     # derivative is in the system
     for v in 1:length(var_to_diff)
         dv = var_to_diff[v]
+        # if the variable is not differentiated, there is nothing to do
         dv isa Int || continue
+        # if we will solve for the differentiated variable, there is nothing to do
         solved = var_eq_matching[dv] isa Int
         solved && continue
 
         # If there's `D(x) = x_t` already, update mappings and continue without
         # adding new equations/variables
         dd = find_duplicate_dd(dv, solvable_graph, diff_to_var, linear_eqs, mm)
-        if !isnothing(dd)
-            dummy_eq, v_t = dd
-            var_to_diff[v_t] = var_to_diff[dv]
-            var_eq_matching[dv] = unassigned
-            eq_var_matching[dummy_eq] = dv
-            continue
+        if dd === nothing
+            # there is no such pre-existing equation
+            # generate the dummy derivative variable
+            dx = fullvars[dv]
+            order, lv = var_order(dv, diff_to_var)
+            x_t = is_discrete ? lower_shift_varname_with_unit(fullvars[dv], iv) :
+                  lower_varname_with_unit(fullvars[lv], iv, order)
+
+            # Add `x_t` to the graph
+            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)
+            push!(full_var_eq_matching, unassigned)
+
+            # We also need to substitute all occurrences of `D(x)` with `x_t` in all equations
+            # except `dummy_eq`, but that is handled in `generate_system_equations!` since
+            # we will solve for `D(x) ~ x_t` and add it to the substitution map.
+            dd = dummy_eq, v_t
         end
+        # there is a duplicate `D(x) ~ x_t` equation
+        # `dummy_eq` is the index of the equation
+        # `v_t` is the dummy derivative variable
+        dummy_eq, v_t = dd
+        var_to_diff[v_t] = var_to_diff[dv]
+        old_matched_eq = full_var_eq_matching[dv]
+        full_var_eq_matching[dv] = var_eq_matching[dv] = dummy_eq
+        full_var_eq_matching[v_t] = old_matched_eq
+        eq_var_matching[dummy_eq] = dv
+        push!(v_t_dvs, (v_t, dv))
+    end
 
-        dx = fullvars[dv]
-        order, lv = var_order(dv, diff_to_var)
-        x_t = is_discrete ? lower_shift_varname_with_unit(fullvars[dv], iv) :
-              lower_varname_with_unit(fullvars[lv], iv, order)
-
-        # Add `x_t` to the graph
-        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)
+    # tuples of (index, scc) indicating that `scc` has to be inserted at
+    # index `index` in `var_sccs`. Same length as `v_t_dvs` because we will
+    # have one new SCC per new variable.
+    sccs_to_insert = similar(v_t_dvs, Tuple{Int, Vector{Int}})
+    # mapping of SCC index to indexes in the SCC to delete
+    idxs_to_remove = Dict{Int, Vector{Int}}()
+    for (k, (v_t, dv)) in enumerate(v_t_dvs)
+        # replace `dv` with `v_t`
+        i, j = v_to_scc[dv]
+        var_sccs[i][j] = v_t
+        if v_t <= length(v_to_scc)
+            # v_t wasn't added by this process, it was already present. Which
+            # means we need to remove it from whatever SCC it is in, since it is
+            # now in this one
+            i_, j_ = v_to_scc[v_t]
+            scc_del_idxs = get!(() -> Int[], idxs_to_remove, i_)
+            push!(scc_del_idxs, j_)
         end
+        # `dv` still needs to be present in some SCC. Since we solve for `dv` from
+        # `0 ~ D(x) - x_t`, it is in its own SCC. This new singleton SCC is solved
+        # immediately before the one that `dv` used to be in (`i`)
+        sccs_to_insert[k] = (i, [dv])
+    end
+    sort!(sccs_to_insert, by = first)
+    # remove the idxs we need to remove
+    for (i, idxs) in idxs_to_remove
+        deleteat!(var_sccs[i], idxs)
+    end
+    # insert the new SCCs, accounting for the fact that we might have multiple entries
+    # in `sccs_to_insert` to be inserted at the same index.
+    old_idx = 1
+    insert_idx = 1
+    new_sccs = similar(var_sccs, length(var_sccs) + length(sccs_to_insert))
+    for i in eachindex(new_sccs)
+        # if we have SCCs to insert, and the index we have to insert them at is the current
+        # one in the old list of SCCs
+        if insert_idx <= length(sccs_to_insert) && sccs_to_insert[insert_idx][1] == old_idx
+            # insert it
+            new_sccs[i] = sccs_to_insert[insert_idx][2]
+            insert_idx += 1
+        else
+            # otherwise, insert the old SCC
+            new_sccs[i] = copy(var_sccs[old_idx])
+            old_idx += 1
+        end
+    end
 
-        # Update matching
-        push!(var_eq_matching, unassigned)
-        var_eq_matching[dv] = unassigned
-        eq_var_matching[dummy_eq] = dv
+    filter!(!isempty, new_sccs)
+    if mm !== nothing
+        @set! mm.ncols = ndsts(graph)
     end
+
+    return new_sccs
 end
 
 """
@@ -442,23 +520,20 @@ or solved variables on the RHS of the final equations will get substituted.
 The topological sort of the equations ensures that variables are solved for
 before they appear in equations. 
 
-Reorder the equations and unknowns to be:
-   [diffeqs; ...]
-   [diffvars; ...]
-such that the mass matrix is:
-   [I  0
-    0  0].
+Reorder the equations and unknowns to be in the BLT sorted form.
 
-Order the new equations and variables such that the differential equations
-and variables come first. Return the new equations, the solved equations,
+Return the new equations, the solved equations,
 the new orderings, and the number of solved variables and equations.
 """
-function generate_system_equations!(state::TearingState, neweqs, var_eq_matching;
+function generate_system_equations!(state::TearingState, neweqs, var_eq_matching,
+        full_var_eq_matching, var_sccs, extra_eqs_vars;
         simplify = false, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = state
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
+    full_eq_var_matching = invview(full_var_eq_matching)
     diff_to_var = invview(var_to_diff)
+    extra_eqs, extra_vars = extra_eqs_vars
 
     total_sub = Dict()
     if is_only_discrete(structure)
@@ -473,88 +548,236 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
         end
     end
 
-    # if var is like D(x) or Shift(t, 1)(x)
-    isdervar = let diff_to_var = diff_to_var
-        var -> diff_to_var[var] !== nothing
-    end
+    eq_generator = EquationGenerator(state, total_sub, D, iv)
 
-    # Extract partition information
-    is_solvable = let solvable_graph = solvable_graph
-        (eq, iv) -> eq isa Int && iv isa Int && BipartiteEdge(eq, iv) in solvable_graph
+    # We need to solve extra equations before everything to repsect
+    # topological order.
+    for eq in extra_eqs
+        var = eq_var_matching[eq]
+        var isa Int || continue
+        codegen_equation!(eq_generator, neweqs[eq], eq, var; simplify)
     end
 
-    diff_eqs = Equation[]
-    diffeq_idxs = Int[]
-    diff_vars = Int[]
-    alge_eqs = Equation[]
-    algeeq_idxs = Int[]
-    solved_eqs = Equation[]
-    solvedeq_idxs = Int[]
-    solved_vars = Int[]
+    # if the variable is present in the equations either as-is or differentiated
+    ispresent = let var_to_diff = var_to_diff, graph = graph
+        i -> (!isempty(𝑑neighbors(graph, i)) ||
+              (var_to_diff[i] !== nothing && !isempty(𝑑neighbors(graph, var_to_diff[i]))))
+    end
 
-    toporder = topological_sort(DiCMOBiGraph{false}(graph, var_eq_matching))
-    eqs = Iterators.reverse(toporder)
+    digraph = DiCMOBiGraph{false}(graph, var_eq_matching)
     idep = iv
+    for (i, scc) in enumerate(var_sccs)
+        # note that the `vscc <-> escc` relation is a set-to-set mapping, and not
+        # point-to-point.
+        vscc, escc = get_sorted_scc(digraph, full_var_eq_matching, var_eq_matching, scc)
+        var_sccs[i] = vscc
+
+        if length(escc) != length(vscc)
+            isempty(escc) && continue
+            escc = setdiff(escc, extra_eqs)
+            isempty(escc) && continue
+            vscc = setdiff(vscc, extra_vars)
+            isempty(vscc) && continue
+        end
 
-    # Generate equations.
-    #   Solvable equations of differential variables D(x) become differential equations
-    #   Solvable equations of non-differential variables become observable equations
-    #   Non-solvable equations become algebraic equations.
-    for ieq in eqs
-        iv = eq_var_matching[ieq]
-        eq = neweqs[ieq]
-
-        if is_solvable(ieq, iv) && isdervar(iv)
-            var = fullvars[iv]
-            isnothing(D) && throw(UnexpectedDifferentialError(equations(sys)[ieq]))
-            order, lv = var_order(iv, diff_to_var)
-            dx = D(simplify_shifts(fullvars[lv]))
-
-            neweq = make_differential_equation(var, dx, eq, total_sub)
-            for e in 𝑑neighbors(graph, iv)
-                e == ieq && continue
-                rem_edge!(graph, e, iv)
-            end
-
-            total_sub[simplify_shifts(neweq.lhs)] = neweq.rhs
-            # Substitute unshifted variables x(k), y(k) on RHS of implicit equations
-            if is_only_discrete(structure)
-                var_to_diff[iv] === nothing && (total_sub[var] = neweq.rhs)
-            end
-            push!(diff_eqs, neweq)
-            push!(diffeq_idxs, ieq)
-            push!(diff_vars, diff_to_var[iv])
-        elseif is_solvable(ieq, iv)
-            var = fullvars[iv]
-            neweq = make_solved_equation(var, eq, total_sub; simplify)
-            !isnothing(neweq) && begin
-                push!(solved_eqs, neweq)
-                push!(solvedeq_idxs, ieq)
-                push!(solved_vars, iv)
-            end
-        else
-            neweq = make_algebraic_equation(eq, total_sub)
-            push!(alge_eqs, neweq)
-            push!(algeeq_idxs, ieq)
+        offset = 1
+        for ieq in escc
+            iv = eq_var_matching[ieq]
+            eq = neweqs[ieq]
+            codegen_equation!(eq_generator, neweqs[ieq], ieq, iv; simplify)
         end
     end
 
+    for eq in extra_eqs
+        var = eq_var_matching[eq]
+        var isa Int && continue
+        codegen_equation!(eq_generator, neweqs[eq], eq, var; simplify)
+    end
+
+    @unpack neweqs′, eq_ordering, var_ordering, solved_eqs, solved_vars = eq_generator
+
+    is_diff_eq = .!iszero.(var_ordering)
     # Generate new equations and orderings 
-    neweqs = [diff_eqs; alge_eqs]
-    eq_ordering = [diffeq_idxs; algeeq_idxs]
+    diff_vars = var_ordering[is_diff_eq]
     diff_vars_set = BitSet(diff_vars)
     if length(diff_vars_set) != length(diff_vars)
         error("Tearing internal error: lowering DAE into semi-implicit ODE failed!")
     end
     solved_vars_set = BitSet(solved_vars)
-    var_ordering = [diff_vars;
-                    setdiff!(setdiff(1:ndsts(graph), diff_vars_set),
-                        solved_vars_set)]
-
+    # We filled zeros for algebraic variables, so fill them properly here
+    offset = 1
+    for (i, v) in enumerate(var_ordering)
+        v == 0 || continue
+        # find the next variable which is not differential or solved, is not the
+        # derivative of another variable and is present in the equations
+        index = findnext(1:ndsts(graph), offset) do j
+            !(j in diff_vars_set || j in solved_vars_set) && diff_to_var[j] === nothing &&
+                ispresent(j)
+        end
+        # in case of overdetermined systems, this may not be present
+        index === nothing && break
+        var_ordering[i] = index
+        offset = index + 1
+    end
+    filter!(!iszero, var_ordering)
+    var_ordering = [var_ordering; setdiff(1:ndsts(graph), var_ordering, solved_vars_set)]
+    neweqs = neweqs′
     return neweqs, solved_eqs, eq_ordering, var_ordering, length(solved_vars),
     length(solved_vars_set)
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Sort the provided SCC `scc`, given the `digraph` of the system constructed using
+`var_eq_matching` along with both the matchings of the system.
+"""
+function get_sorted_scc(
+        digraph::DiCMOBiGraph, full_var_eq_matching::Matching, var_eq_matching::Matching, scc::Vector{Int})
+    eq_var_matching = invview(var_eq_matching)
+    full_eq_var_matching = invview(full_var_eq_matching)
+    # obtain the matched equations in the SCC
+    scc_eqs = Int[full_var_eq_matching[v] for v in scc if full_var_eq_matching[v] isa Int]
+    # obtain the equations in the SCC that are linearly solvable
+    scc_solved_eqs = Int[var_eq_matching[v] for v in scc if var_eq_matching[v] isa Int]
+    # obtain the subgraph of the contracted graph involving the solved equations
+    subgraph, varmap = Graphs.induced_subgraph(digraph, scc_solved_eqs)
+    # topologically sort the solved equations and append the remainder
+    scc_eqs = [varmap[reverse(topological_sort(subgraph))];
+               setdiff(scc_eqs, scc_solved_eqs)]
+    # the variables of the SCC are obtained by inverse mapping the sorted equations
+    # and appending the rest
+    scc_vars = [eq_var_matching[e] for e in scc_eqs if eq_var_matching[e] isa Int]
+    append!(scc_vars, setdiff(scc, scc_vars))
+    return scc_vars, scc_eqs
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Struct containing the information required to generate equations of a system, as well as
+the generated equations and associated metadata.
+"""
+struct EquationGenerator{S, D, I}
+    """
+    `TearingState` of the system.
+    """
+    state::S
+    """
+    Substitutions to perform in all subsequent equations. For each differential equation
+    `D(x) ~ f(..)`, the substitution `D(x) => f(..)` is added to the rules.
+    """
+    total_sub::Dict{Any, Any}
+    """
+    The differential operator, or `nothing` if not applicable.
+    """
+    D::D
+    """
+    The independent variable, or `nothing` if not applicable.
+    """
+    idep::I
+    """
+    The new generated equations of the system.
+    """
+    neweqs′::Vector{Equation}
+    """
+    `eq_ordering[i]` is the index `neweqs′[i]` was originally at in the untorn equations of
+    the system. This is used to permute the state of the system into BLT sorted form. 
+    """
+    eq_ordering::Vector{Int}
+    """
+    `var_ordering[i]` is the index in `state.fullvars` of the variable at the `i`th index in
+    the BLT sorted form.
+    """
+    var_ordering::Vector{Int}
+    """
+    List of linearly solved (observed) equations.
+    """
+    solved_eqs::Vector{Equation}
+    """
+    `eq_ordering` for `solved_eqs`.
+    """
+    solved_vars::Vector{Int}
+end
+
+function EquationGenerator(state, total_sub, D, idep)
+    EquationGenerator(
+        state, total_sub, D, idep, Equation[], Int[], Int[], Equation[], Int[])
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Check if equation at index `ieq` is linearly solvable for variable at index `iv`.
+"""
+function is_solvable(eg::EquationGenerator, ieq, iv)
+    solvable_graph = eg.state.structure.solvable_graph
+    return ieq isa Int && iv isa Int && BipartiteEdge(ieq, iv) in solvable_graph
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+    If `iv` is like D(x) or Shift(t, 1)(x)
+"""
+function is_dervar(eg::EquationGenerator, iv::Int)
+    diff_to_var = invview(eg.state.structure.var_to_diff)
+    diff_to_var[iv] !== nothing
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Appropriately codegen the given equation `eq`, which occurs at index `ieq` in the untorn
+list of equations and is matched to variable at index `iv`.
+"""
+function codegen_equation!(eg::EquationGenerator,
+        eq::Equation, ieq::Int, iv::Union{Int, Unassigned}; simplify = false)
+    # We generate equations ordered by the matched variables
+    #   Solvable equations of differential variables D(x) become differential equations
+    #   Solvable equations of non-differential variables become observable equations
+    #   Non-solvable equations become algebraic equations.
+    @unpack state, total_sub, neweqs′, eq_ordering, var_ordering = eg
+    @unpack solved_eqs, solved_vars, D, idep = eg
+    @unpack fullvars, sys, structure = state
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
+    diff_to_var = invview(var_to_diff)
+    if is_solvable(eg, ieq, iv) && is_dervar(eg, iv)
+        var = fullvars[iv]
+        isnothing(D) && throw(UnexpectedDifferentialError(equations(sys)[ieq]))
+        order, lv = var_order(iv, diff_to_var)
+        dx = D(simplify_shifts(fullvars[lv]))
+
+        neweq = make_differential_equation(var, dx, eq, total_sub)
+        for e in 𝑑neighbors(graph, iv)
+            e == ieq && continue
+            rem_edge!(graph, e, iv)
+        end
+
+        total_sub[simplify_shifts(neweq.lhs)] = neweq.rhs
+        # Substitute unshifted variables x(k), y(k) on RHS of implicit equations
+        if is_only_discrete(structure)
+            var_to_diff[iv] === nothing && (total_sub[var] = neweq.rhs)
+        end
+        push!(neweqs′, neweq)
+        push!(eq_ordering, ieq)
+        push!(var_ordering, diff_to_var[iv])
+    elseif is_solvable(eg, ieq, iv)
+        var = fullvars[iv]
+        neweq = make_solved_equation(var, eq, total_sub; simplify)
+        if neweq !== nothing
+            push!(solved_eqs, neweq)
+            push!(solved_vars, iv)
+        end
+    else
+        neweq = make_algebraic_equation(eq, total_sub)
+        push!(neweqs′, neweq)
+        push!(eq_ordering, ieq)
+        # we push a dummy to `var_ordering` here because `iv` is `unassigned`
+        push!(var_ordering, 0)
+    end
+end
+
 """
 Occurs when a variable D(x) occurs in a non-differential system.
 """
@@ -615,8 +838,7 @@ tearing state to account for the new order. Permute the variables and equations.
 Eliminate the solved variables and equations from the graph and permute the
 graph's vertices to account for the new variable/equation ordering.
 """
-# TODO: BLT sorting
-function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering,
+function reorder_vars!(state::TearingState, var_eq_matching, var_sccs, eq_ordering,
         var_ordering, nsolved_eq, nsolved_var)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
 
@@ -650,6 +872,17 @@ function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering,
     end
     new_fullvars = state.fullvars[var_ordering]
 
+    # Update the SCCs
+    var_ordering_set = BitSet(var_ordering)
+    for scc in var_sccs
+        # Map variables to their new indices
+        map!(v -> varsperm[v], scc, scc)
+        # Remove variables not in the reduced set
+        filter!(!iszero, scc)
+    end
+    # Remove empty SCCs
+    filter!(!isempty, var_sccs)
+
     # Update system structure 
     @set! state.structure.graph = complete(new_graph)
     @set! state.structure.var_to_diff = new_var_to_diff
@@ -662,7 +895,7 @@ end
 Update the system equations, unknowns, and observables after simplification.
 """
 function update_simplified_system!(
-        state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns;
+        state::TearingState, neweqs, solved_eqs, dummy_sub, var_sccs, extra_unknowns;
         cse_hack = true, array_hack = true)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
     diff_to_var = invview(var_to_diff)
@@ -681,9 +914,9 @@ function update_simplified_system!(
     # TODO: compute the dependency correctly so that we don't have to do this
     obs = [fast_substitute(observed(sys), obs_sub); solved_eqs]
 
-    unknowns = Any[v
-                   for (i, v) in enumerate(state.fullvars)
-                   if diff_to_var[i] === nothing && ispresent(i)]
+    unknown_idxs = filter(
+        i -> diff_to_var[i] === nothing && ispresent(i), eachindex(state.fullvars))
+    unknowns = state.fullvars[unknown_idxs]
     unknowns = [unknowns; extra_unknowns]
     @set! sys.unknowns = unknowns
 
@@ -695,7 +928,12 @@ function update_simplified_system!(
 
     # Only makes sense for time-dependent
     if ModelingToolkit.has_schedule(sys)
-        @set! sys.schedule = Schedule(var_eq_matching, dummy_sub)
+        unknowns_set = BitSet(unknown_idxs)
+        for scc in var_sccs
+            intersect!(scc, unknowns_set)
+        end
+        filter!(!isempty, var_sccs)
+        @set! sys.schedule = Schedule(var_sccs, dummy_sub)
     end
     if ModelingToolkit.has_isscheduled(sys)
         @set! sys.isscheduled = true
@@ -729,18 +967,19 @@ variables are marked as `SelectedState` and they are differentiated in the
 DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
 appear in the system. Algebraic variables are variables that are not
 differential variables.
+
+# Arguments
+
+- `state`: The `TearingState` of the system.
+- `var_eq_matching`: The maximal matching after state selection.
+- `full_var_eq_matching`: The maximal matching prior to state selection.
+- `var_sccs`: The topologically sorted strongly connected components of the system
+  according to `full_var_eq_matching`.
 """
-function tearing_reassemble(state::TearingState, var_eq_matching,
-        full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
-    extra_vars = Int[]
-    if full_var_eq_matching !== nothing
-        for v in 𝑑vertices(state.structure.graph)
-            eq = full_var_eq_matching[v]
-            eq isa Int && continue
-            push!(extra_vars, v)
-        end
-    end
-    extra_unknowns = state.fullvars[extra_vars]
+function tearing_reassemble(state::TearingState, var_eq_matching::Matching,
+        full_var_eq_matching::Matching, var_sccs::Vector{Vector{Int}}; simplify = false, mm, cse_hack = true,
+        array_hack = true, fully_determined = true)
+    extra_eqs_vars = get_extra_eqs_vars(state, full_var_eq_matching, fully_determined)
     neweqs = collect(equations(state))
     dummy_sub = Dict()
 
@@ -755,18 +994,22 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         iv = D = nothing
     end
 
+    extra_unknowns = state.fullvars[extra_eqs_vars[2]]
     # Structural simplification 
     substitute_derivatives_algevars!(state, neweqs, var_eq_matching, dummy_sub; iv, D)
 
-    generate_derivative_variables!(state, neweqs, var_eq_matching; mm, iv, D)
+    var_sccs = generate_derivative_variables!(
+        state, neweqs, var_eq_matching, full_var_eq_matching, var_sccs; mm, iv, D)
 
     neweqs, solved_eqs, eq_ordering, var_ordering, nelim_eq, nelim_var = generate_system_equations!(
-        state, neweqs, var_eq_matching; simplify, iv, D)
+        state, neweqs, var_eq_matching, full_var_eq_matching,
+        var_sccs, extra_eqs_vars; simplify, iv, D)
 
     state = reorder_vars!(
-        state, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
+        state, var_eq_matching, var_sccs, eq_ordering, var_ordering, nelim_eq, nelim_var)
+    # var_eq_matching and full_var_eq_matching are now invalidated
 
-    sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_eq_matching,
+    sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_sccs,
         extra_unknowns; cse_hack, array_hack)
 
     @set! state.sys = sys
@@ -774,6 +1017,35 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     return invalidate_cache!(sys)
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Return a 2-tuple of integer vectors containing indices of extra equations and variables
+respectively. For fully-determined systems, both of these are empty. Overdetermined systems
+have extra equations, and underdetermined systems have extra variables.
+"""
+function get_extra_eqs_vars(
+        state::TearingState, full_var_eq_matching::Matching, fully_determined::Bool)
+    fully_determined && return Int[], Int[]
+
+    extra_eqs = Int[]
+    extra_vars = Int[]
+    full_eq_var_matching = invview(full_var_eq_matching)
+
+    for v in 𝑑vertices(state.structure.graph)
+        eq = full_var_eq_matching[v]
+        eq isa Int && continue
+        push!(extra_vars, v)
+    end
+    for eq in 𝑠vertices(state.structure.graph)
+        v = full_eq_var_matching[eq]
+        v isa Int && continue
+        push!(extra_eqs, eq)
+    end
+
+    return extra_eqs, extra_vars
+end
+
 """
 # HACK 1
 
@@ -937,22 +1209,11 @@ new residual equations after tearing. End users are encouraged to call [`mtkcomp
 instead, which calls this function internally.
 """
 function tearing(sys::AbstractSystem, state = TearingState(sys); mm = nothing,
-        simplify = false, cse_hack = true, array_hack = true, kwargs...)
-    var_eq_matching, full_var_eq_matching = tearing(state)
+        simplify = false, cse_hack = true, array_hack = true, fully_determined = true, kwargs...)
+    var_eq_matching, full_var_eq_matching, var_sccs, can_eliminate = tearing(state)
     invalidate_cache!(tearing_reassemble(
-        state, var_eq_matching, full_var_eq_matching; mm, simplify, cse_hack, array_hack))
-end
-
-"""
-    partial_state_selection(sys; simplify=false)
-
-Perform partial state selection and tearing.
-"""
-function partial_state_selection(sys; simplify = false)
-    state = TearingState(sys)
-    var_eq_matching = partial_state_selection_graph!(state)
-
-    tearing_reassemble(state, var_eq_matching; simplify = simplify)
+        state, var_eq_matching, full_var_eq_matching, var_sccs; mm,
+        simplify, cse_hack, array_hack, fully_determined))
 end
 
 """
@@ -962,7 +1223,7 @@ Perform index reduction and use the dummy derivative technique to ensure that
 the system is balanced.
 """
 function dummy_derivative(sys, state = TearingState(sys); simplify = false,
-        mm = nothing, cse_hack = true, array_hack = true, kwargs...)
+        mm = nothing, cse_hack = true, array_hack = true, fully_determined = true, kwargs...)
     jac = let state = state
         (eqs, vars) -> begin
             symeqs = EquationsView(state)[eqs]
@@ -984,7 +1245,9 @@ function dummy_derivative(sys, state = TearingState(sys); simplify = false,
             p
         end
     end
-    var_eq_matching = dummy_derivative_graph!(state, jac; state_priority,
+    var_eq_matching, full_var_eq_matching, var_sccs, can_eliminate, summary = dummy_derivative_graph!(
+        state, jac; state_priority,
         kwargs...)
-    tearing_reassemble(state, var_eq_matching; simplify, mm, cse_hack, array_hack)
+    tearing_reassemble(state, var_eq_matching, full_var_eq_matching, var_sccs;
+        simplify, mm, cse_hack, array_hack, fully_determined)
 end
diff --git a/src/systems/system.jl b/src/systems/system.jl
index 5c01557cb5..7cab43156b 100644
--- a/src/systems/system.jl
+++ b/src/systems/system.jl
@@ -1,8 +1,5 @@
-struct Schedule{V <: BipartiteGraphs.Matching}
-    """
-    Maximal matching of variables to equations calculated during structural simplification.
-    """
-    var_eq_matching::V
+struct Schedule
+    var_sccs::Vector{Vector{Int}}
     """
     Mapping of `Differential`s of variables to corresponding derivative expressions.
     """
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 32da7a81bf..8b0e020cc4 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -762,10 +762,10 @@ function _mtkcompile!(state::TearingState; simplify = false,
         state = TearingState(sys)
         sys, mm = ModelingToolkit.alias_elimination!(state; kwargs...)
         sys = ModelingToolkit.dummy_derivative(
-            sys, state; simplify, mm, check_consistency, kwargs...)
+            sys, state; simplify, mm, check_consistency, fully_determined, kwargs...)
     else
         sys = ModelingToolkit.tearing(
-            sys, state; simplify, mm, check_consistency, kwargs...)
+            sys, state; simplify, mm, check_consistency, fully_determined, kwargs...)
     end
     fullunknowns = [observables(sys); unknowns(sys)]
     @set! sys.observed = ModelingToolkit.topsort_equations(observed(sys), fullunknowns)

From 6d4edba320f9bde164643b5efc53f68e8cf3eb27 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Sun, 1 Jun 2025 13:37:53 +0530
Subject: [PATCH 06/11] test: update tests to account for new tearing

---
 test/bvproblem.jl                         | 13 +++++++-----
 test/dde.jl                               |  4 ++--
 test/extensions/dynamic_optimization.jl   | 24 +++++++++++------------
 test/fmi/fmi.jl                           |  3 ++-
 test/initial_values.jl                    |  9 +++++----
 test/initializationsystem.jl              |  2 +-
 test/input_output_handling.jl             |  4 ++--
 test/odesystem.jl                         | 13 ++++++------
 test/reduction.jl                         |  2 +-
 test/structural_transformation/tearing.jl | 13 ++++++++----
 10 files changed, 49 insertions(+), 38 deletions(-)

diff --git a/test/bvproblem.jl b/test/bvproblem.jl
index abc925ccfa..d86ba251c7 100644
--- a/test/bvproblem.jl
+++ b/test/bvproblem.jl
@@ -30,7 +30,7 @@ daesolvers = [Ascher2, Ascher4, Ascher6]
     for solver in solvers
         sol = solve(bvp, solver(), dt = 0.01)
         @test isapprox(sol.u[end], osol.u[end]; atol = 0.01)
-        @test sol.u[1] == [1.0, 2.0]
+        @test sol[[x, y], 1] == [1.0, 2.0]
     end
 
     # Test out of place
@@ -40,7 +40,7 @@ daesolvers = [Ascher2, Ascher4, Ascher6]
     for solver in solvers
         sol = solve(bvp2, solver(), dt = 0.01)
         @test isapprox(sol.u[end], osol.u[end]; atol = 0.01)
-        @test sol.u[1] == [1.0, 2.0]
+        @test sol[[x, y], 1] == [1.0, 2.0]
     end
 end
 
@@ -129,7 +129,9 @@ end
     sol2 = solve(bvpi2, MIRK4(), dt = 0.01)
     sol3 = solve(bvpi3, MIRK4(), dt = 0.01)
     sol4 = solve(bvpi4, MIRK4(), dt = 0.01)
-    @test sol1 ≈ sol2 ≈ sol3 ≈ sol4 # don't get true equality here, not sure why
+    @test sol1.t ≈ sol2.t ≈ sol3.t ≈ sol4.t
+    @test sol1.u ≈ sol2.u ≈ sol3[[x(t), y(t)]] ≈ sol4[[x(t), y(t)]]
+    # @test sol1 ≈ sol2 ≈ sol3 ≈ sol4 # don't get true equality here, not sure why
 end
 
 function test_solvers(
@@ -296,7 +298,7 @@ end
     costfn = ModelingToolkit.generate_cost(
         lksys; expression = Val{false}, wrap_gfw = Val{true})
     _t = tspan[2]
-    @test costfn(sol, prob.p, _t) ≈ (sol(0.6)[1] + 3)^2 + sol(0.3)[1]^2
+    @test costfn(sol, prob.p, _t) ≈ (sol(0.6; idxs = x(t)) + 3)^2 + sol(0.3; idxs = x(t))^2
 
     ### With a parameter
     @parameters t_c
@@ -309,5 +311,6 @@ end
     sol = solve(prob, Tsit5())
     costfn = ModelingToolkit.generate_cost(
         lksys; expression = Val{false}, wrap_gfw = Val{true})
-    @test costfn(sol, prob.p, _t) ≈ log(sol(0.56)[2] + sol(0.0)[1]) - sol(0.4)[1]^2
+    @test costfn(sol, prob.p, _t) ≈
+          log(sol(0.56; idxs = y(t)) + sol(0.0; idxs = x(t))) - sol(0.4; idxs = x(t))^2
 end
diff --git a/test/dde.jl b/test/dde.jl
index 211b59075c..fb92bf61eb 100644
--- a/test/dde.jl
+++ b/test/dde.jl
@@ -46,13 +46,13 @@ prob = DDEProblem(sys,
     tspan,
     constant_lags = [tau])
 sol_mtk = solve(prob, alg, reltol = 1e-7, abstol = 1e-10)
-@test sol_mtk.u[end] ≈ sol.u[end]
+@test sol_mtk[[x₀, x₁, x₂(t)]][end] ≈ sol.u[end]
 prob2 = DDEProblem(sys,
     [x₀ => 1.0 - t * q1 * 10, x₁ => 1.0 - t * q1 * 10, x₂(t) => 1.0 - t * q1 * 10],
     tspan,
     constant_lags = [tau])
 sol2_mtk = solve(prob2, alg, reltol = 1e-7, abstol = 1e-10)
-@test sol2_mtk.u[end] ≈ sol2.u[end]
+@test sol2_mtk[[x₀, x₁, x₂(t)]][end] ≈ sol2.u[end]
 @test_nowarn sol2_mtk[[x₀, x₁, x₂(t)]]
 @test_nowarn sol2_mtk[[x₀, x₁, x₂(t - 0.1)]]
 
diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl
index 62553121fe..dd2368b558 100644
--- a/test/extensions/dynamic_optimization.jl
+++ b/test/extensions/dynamic_optimization.jl
@@ -58,22 +58,22 @@ const M = ModelingToolkit
     jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     @test JuMP.num_constraints(jprob.model) == 2
     jsol = solve(jprob, Ipopt.Optimizer, constructTsitouras5, silent = true) # 12.190 s, 9.68 GiB
-    @test jsol.sol(0.6)[1] ≈ 3.5
-    @test jsol.sol(0.3)[1] ≈ 7.0
+    @test jsol.sol(0.6; idxs = x(t)) ≈ 3.5
+    @test jsol.sol(0.3; idxs = x(t)) ≈ 7.0
 
     cprob = CasADiDynamicOptProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     csol = solve(cprob, "ipopt", constructTsitouras5, silent = true)
-    @test csol.sol(0.6)[1] ≈ 3.5
-    @test csol.sol(0.3)[1] ≈ 7.0
+    @test csol.sol(0.6; idxs = x(t)) ≈ 3.5
+    @test csol.sol(0.3; idxs = x(t)) ≈ 7.0
 
     iprob = InfiniteOptDynamicOptProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer,
         derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
     sol = isol.sol
-    @test sol(0.6)[1] ≈ 3.5
-    @test sol(0.3)[1] ≈ 7.0
+    @test sol(0.6; idxs = x(t)) ≈ 3.5
+    @test sol(0.3; idxs = x(t)) ≈ 7.0
 
     # Test whole-interval constraints
     constr = [x(t) ≳ 1, y(t) ≳ 1]
@@ -128,14 +128,14 @@ end
     # Linear systems have bang-bang controls
     @test is_bangbang(jsol.input_sol, [-1.0], [1.0])
     # Test reached final position.
-    @test ≈(jsol.sol.u[end][2], 0.25, rtol = 1e-5)
+    @test ≈(jsol.sol[x(t)][end], 0.25, rtol = 1e-5)
 
     cprob = CasADiDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01)
     csol = solve(cprob, "ipopt", constructVerner8, silent = true)
     # Linear systems have bang-bang controls
     @test is_bangbang(csol.input_sol, [-1.0], [1.0])
     # Test reached final position.
-    @test ≈(csol.sol.u[end][2], 0.25, rtol = 1e-5)
+    @test ≈(csol.sol[x(t)][end], 0.25, rtol = 1e-5)
 
     # Test dynamics
     @parameters (u_interp::ConstantInterpolation)(..)
@@ -149,7 +149,7 @@ end
     iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [-1.0], [1.0])
-    @test ≈(isol.sol.u[end][2], 0.25, rtol = 1e-5)
+    @test ≈(isol.sol[x(t)][end], 0.25, rtol = 1e-5)
     osol = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false)
     @test ≈(isol.sol.u, osol.u, rtol = 0.05)
 
@@ -220,16 +220,16 @@ end
         Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
     jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     jsol = solve(jprob, Ipopt.Optimizer, constructRadauIIA5, silent = true)
-    @test jsol.sol.u[end][1] > 1.012
+    @test jsol.sol[h(t)][end] > 1.012
 
     cprob = CasADiDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     csol = solve(cprob, "ipopt"; silent = true)
-    @test csol.sol.u[end][1] > 1.012
+    @test csol.sol[h(t)][end] > 1.012
 
     iprob = InfiniteOptDynamicOptProblem(
         rocket, u0map, (ts, te), pmap; dt = 0.001)
     isol = solve(iprob, Ipopt.Optimizer, silent = true)
-    @test isol.sol.u[end][1] > 1.012
+    @test isol.sol[h(t)][end] > 1.012
 
     # Test solution
     @parameters (T_interp::CubicSpline)(..)
diff --git a/test/fmi/fmi.jl b/test/fmi/fmi.jl
index e0eaf373d0..85e1830650 100644
--- a/test/fmi/fmi.jl
+++ b/test/fmi/fmi.jl
@@ -170,7 +170,8 @@ end
         sol = solve(prob, Rodas5P(autodiff = false), abstol = 1e-8, reltol = 1e-8)
         @test SciMLBase.successful_retcode(sol)
 
-        @test truesol(sol.t; idxs = [truesys.a, truesys.b, truesys.c]).u≈sol.u rtol=4e-2
+        @test truesol(sol.t;
+            idxs = [truesys.a, truesys.b, truesys.c]).u≈sol[[sys.a, sys.b, sys.c]] rtol=4e-2
         # sys.adder.c is a discrete variable
         @test truesol(sol.t; idxs = truesys.adder.c).u≈sol(sol.t; idxs = sys.adder.c).u rtol=4e-2
     end
diff --git a/test/initial_values.jl b/test/initial_values.jl
index d15f30c280..1b32dc7b51 100644
--- a/test/initial_values.jl
+++ b/test/initial_values.jl
@@ -8,9 +8,10 @@ using SymbolicIndexingInterface
 @variables x(t)[1:3]=[1.0, 2.0, 3.0] y(t) z(t)[1:2]
 
 @mtkcompile sys=System([D(x) ~ t * x], t) simplify=false
-@test get_u0(sys, []) == [1.0, 2.0, 3.0]
-@test get_u0(sys, [x => [2.0, 3.0, 4.0]]) == [2.0, 3.0, 4.0]
-@test get_u0(sys, [x[1] => 2.0, x[2] => 3.0, x[3] => 4.0]) == [2.0, 3.0, 4.0]
+reorderer = getsym(sys, x)
+@test reorderer(get_u0(sys, [])) == [1.0, 2.0, 3.0]
+@test reorderer(get_u0(sys, [x => [2.0, 3.0, 4.0]])) == [2.0, 3.0, 4.0]
+@test reorderer(get_u0(sys, [x[1] => 2.0, x[2] => 3.0, x[3] => 4.0])) == [2.0, 3.0, 4.0]
 @test get_u0(sys, [2.0, 3.0, 4.0]) == [2.0, 3.0, 4.0]
 
 @mtkcompile sys=System([
@@ -182,7 +183,7 @@ end
     prob = ODEProblem(sys, [], (0.0, 1.0))
     sol = solve(prob)
     @test SciMLBase.successful_retcode(sol)
-    @test sol.u[1] ≈ [1.0, 1.0, 0.5, 0.5]
+    @test sol[x, 1] ≈ [1.0, 1.0, 0.5, 0.5]
 end
 
 @testset "Missing/cyclic guesses throws error" begin
diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
index 1e83ec579f..fc01e9e2e6 100644
--- a/test/initializationsystem.jl
+++ b/test/initializationsystem.jl
@@ -459,7 +459,7 @@ sol = solve(prob, Tsit5())
 
 prob = ODEProblem(simpsys, [z => 1.0, y => 1.0], tspan, guesses = [x => 2.0])
 sol = solve(prob, Tsit5())
-@test sol.u[1] == [0.0, 1.0]
+@test sol[[x, y], 1] == [0.0, 1.0]
 
 # This should warn, but logging tests can't be marked as broken
 @test_logs prob = ODEProblem(simpsys, [], tspan, guesses = [x => 2.0])
diff --git a/test/input_output_handling.jl b/test/input_output_handling.jl
index c295546248..f30949ed1b 100644
--- a/test/input_output_handling.jl
+++ b/test/input_output_handling.jl
@@ -369,7 +369,7 @@ eqs = [D(y₁) ~ -k₁ * y₁ + k₃ * y₂ * y₃ + u1
 m_inputs = [u[1], u[2]]
 m_outputs = [y₂]
 sys_simp = mtkcompile(sys, inputs = m_inputs, outputs = m_outputs)
-@test isequal(unknowns(sys_simp), collect(x[1:2]))
+@test issetequal(unknowns(sys_simp), collect(x[1:2]))
 @test length(inputs(sys_simp)) == 2
 
 # https://github.com/SciML/ModelingToolkit.jl/issues/1577
@@ -417,7 +417,7 @@ matrices, ssys = linearize(augmented_sys,
     ], outs;
     op = [augmented_sys.u => 0.0, augmented_sys.input.u[2] => 0.0, augmented_sys.d => 0.0])
 matrices = ModelingToolkit.reorder_unknowns(
-    matrices, unknowns(ssys), [ssys.x[2], ssys.integrator.x[1], ssys.x[1]])
+    matrices, unknowns(ssys), [ssys.x[1], ssys.x[2], ssys.integrator.x[1]])
 @test matrices.A ≈ [A [1; 0]; zeros(1, 2) -0.001]
 @test matrices.B == I
 @test matrices.C == [C zeros(2)]
diff --git a/test/odesystem.jl b/test/odesystem.jl
index 58315f0cff..fed0ac28fe 100644
--- a/test/odesystem.jl
+++ b/test/odesystem.jl
@@ -568,10 +568,10 @@ let
     @named sys = System(eqs, t, vcat(x, [y]), [k], defaults = Dict(x .=> 0))
     sys = mtkcompile(sys)
 
-    u0 = unknowns(sys) .=> [0.5, 0]
-    du0 = D.(unknowns(sys)) .=> 0.0
+    u0 = x .=> [0.5, 0]
+    du0 = D.(x) .=> 0.0
     prob = DAEProblem(sys, [du0; u0], (0, 50))
-    @test prob.u0 ≈ [0.5, 0.0]
+    @test prob[x] ≈ [0.5, 0.0]
     @test prob.du0 ≈ [0.0, 0.0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 1
@@ -581,7 +581,8 @@ let
 
     prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0, x[1] => 0.5],
         (0, 50))
-    @test prob.u0 ≈ [0.5, 0]
+
+    @test prob[x] ≈ [0.5, 0]
     @test prob.du0 ≈ [0, 0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 1
@@ -590,7 +591,7 @@ let
 
     prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0, x[1] => 0.5, k => 2],
         (0, 50))
-    @test prob.u0 ≈ [0.5, 0]
+    @test prob[x] ≈ [0.5, 0]
     @test prob.du0 ≈ [0, 0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 2
@@ -769,7 +770,7 @@ let
     @named sys4 = System([D(x) ~ -y; D(y) ~ 1 + pp * y + x], t)
     sys4s = mtkcompile(sys4)
     prob = ODEProblem(sys4s, [x => 1.0, D(x) => 1.0], (0, 1.0))
-    @test string.(unknowns(prob.f.sys)) == ["x(t)", "y(t)"]
+    @test issetequal(string.(unknowns(prob.f.sys)), ["x(t)", "y(t)"])
     @test string.(parameters(prob.f.sys)) == ["pp"]
     @test string.(independent_variables(prob.f.sys)) == ["t"]
 end
diff --git a/test/reduction.jl b/test/reduction.jl
index 43fef0fb63..97d360fd67 100644
--- a/test/reduction.jl
+++ b/test/reduction.jl
@@ -245,7 +245,7 @@ eqs = [D(x) ~ x
 sys = mtkcompile(model)
 Js = ModelingToolkit.jacobian_sparsity(sys)
 @test size(Js) == (3, 3)
-@test Js == Diagonal([1, 1, 0])
+@test Js == Diagonal([0, 1, 1])
 
 # MWE for #1722
 vars = @variables a(t) w(t) phi(t)
diff --git a/test/structural_transformation/tearing.jl b/test/structural_transformation/tearing.jl
index 7f18cbb703..4025f7a298 100644
--- a/test/structural_transformation/tearing.jl
+++ b/test/structural_transformation/tearing.jl
@@ -5,6 +5,7 @@ using ModelingToolkit.StructuralTransformations: SystemStructure, find_solvables
 using NonlinearSolve
 using LinearAlgebra
 using UnPack
+using SymbolicIndexingInterface
 using ModelingToolkit: t_nounits as t, D_nounits as D
 ###
 ### Nonlinear system
@@ -147,9 +148,10 @@ eqs = [D(x) ~ z * h
        0 ~ sin(z) + y - p * t]
 @named daesys = System(eqs, t)
 newdaesys = mtkcompile(daesys)
-@test equations(newdaesys) == [D(x) ~ h * z; 0 ~ y + sin(z) - p * t]
-@test equations(tearing_substitution(newdaesys)) == [D(x) ~ h * z; 0 ~ x + sin(z) - p * t]
-@test isequal(unknowns(newdaesys), [x, z])
+@test issetequal(equations(newdaesys), [D(x) ~ h * z; 0 ~ y + sin(z) - p * t])
+@test issetequal(
+    equations(tearing_substitution(newdaesys)), [D(x) ~ h * z; 0 ~ x + sin(z) - p * t])
+@test issetequal(unknowns(newdaesys), [x, z])
 prob = ODEProblem(newdaesys, [x => 1.0, z => -0.5π, p => 0.2], (0, 1.0))
 du = [0.0, 0.0];
 u = [1.0, -0.5π];
@@ -157,7 +159,10 @@ pr = prob.p;
 tt = 0.1;
 @test (@ballocated $(prob.f)($du, $u, $pr, $tt)) == 0
 prob.f(du, u, pr, tt)
-@test du≈[u[2], u[1] + sin(u[2]) - prob.ps[p] * tt] atol=1e-5
+xgetter = getsym(prob, x)
+zgetter = getsym(prob, z)
+@test xgetter(du)≈zgetter(u) atol=1e-5
+@test zgetter(du)≈xgetter(u) + sin(zgetter(u)) - prob.ps[p] * tt atol=1e-5
 
 # test the initial guess is respected
 @named sys = System(eqs, t, defaults = Dict(z => NaN))

From bea3914be5f8afb6f5bb3d8d1ce4b88c72ecc66c Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Jun 2025 11:46:59 +0530
Subject: [PATCH 07/11] test: optimization test is no longer broken

---
 test/optimizationsystem.jl | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/test/optimizationsystem.jl b/test/optimizationsystem.jl
index 0ff85a518d..ad77dbc416 100644
--- a/test/optimizationsystem.jl
+++ b/test/optimizationsystem.jl
@@ -109,11 +109,9 @@ end
     x0 = zeros(2)
     p = [1.0, 100.0]
     f = OptimizationFunction(rosenbrock, Optimization.AutoSymbolics())
-    @test_broken begin
-        prob = OptimizationProblem(f, x0, p)
-        sol = solve(prob, Newton())
-        @test sol.u ≈ [1.0, 1.0]
-    end
+    prob = OptimizationProblem(f, x0, p)
+    sol = solve(prob, Newton())
+    @test sol.u ≈ [1.0, 1.0]
 end
 
 # issue #819

From a6ced9a1ee2b2d9c4f479069b78cfa3f639d05f6 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Jun 2025 13:16:48 +0530
Subject: [PATCH 08/11] refactor: use SCC information calculated during
 `mtkcompile` in `SCCNonlinearProblem`

---
 src/problems/sccnonlinearproblem.jl | 29 ++++++++++++++++++++---------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/problems/sccnonlinearproblem.jl b/src/problems/sccnonlinearproblem.jl
index 3e00170ad9..d2afb7b315 100644
--- a/src/problems/sccnonlinearproblem.jl
+++ b/src/problems/sccnonlinearproblem.jl
@@ -80,21 +80,32 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::System, op; eval_expression = f
     end
 
     ts = get_tearing_state(sys)
-    var_eq_matching, var_sccs = StructuralTransformations.algebraic_variables_scc(ts)
+    sched = get_schedule(sys)
+    if sched === nothing
+        @warn "System is simplified but does not have a schedule. This should not happen."
+        var_eq_matching, var_sccs = StructuralTransformations.algebraic_variables_scc(ts)
+        condensed_graph = MatchedCondensationGraph(
+            DiCMOBiGraph{true}(complete(ts.structure.graph),
+                complete(var_eq_matching)),
+            var_sccs)
+        toporder = topological_sort_by_dfs(condensed_graph)
+        var_sccs = var_sccs[toporder]
+        eq_sccs = map(Base.Fix1(getindex, var_eq_matching), var_sccs)
+    else
+        var_sccs = sched.var_sccs
+        # Equations are already in the order of SCCs
+        eq_sccs = length.(var_sccs)
+        cumsum!(eq_sccs, eq_sccs)
+        eq_sccs = map(enumerate(eq_sccs)) do (i, lasti)
+            i == 1 ? (1:lasti) : ((eq_sccs[i - 1] + 1):lasti)
+        end
+    end
 
     if length(var_sccs) == 1
         return NonlinearProblem{iip}(
             sys, op; eval_expression, eval_module, kwargs...)
     end
 
-    condensed_graph = MatchedCondensationGraph(
-        DiCMOBiGraph{true}(complete(ts.structure.graph),
-            complete(var_eq_matching)),
-        var_sccs)
-    toporder = topological_sort_by_dfs(condensed_graph)
-    var_sccs = var_sccs[toporder]
-    eq_sccs = map(Base.Fix1(getindex, var_eq_matching), var_sccs)
-
     dvs = unknowns(sys)
     ps = parameters(sys)
     eqs = equations(sys)

From be49462e0547a7fdf38e90f7fa79b9286817ebdf Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Jun 2025 13:17:01 +0530
Subject: [PATCH 09/11] test: test SCC sorted order

---
 test/scc_nonlinear_problem.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/test/scc_nonlinear_problem.jl b/test/scc_nonlinear_problem.jl
index ded9852670..fad30dbcea 100644
--- a/test/scc_nonlinear_problem.jl
+++ b/test/scc_nonlinear_problem.jl
@@ -39,6 +39,9 @@ using ModelingToolkit: t_nounits as t, D_nounits as D
         @test prob.u0 isa SVector
         @test !SciMLBase.isinplace(prob)
     end
+
+    # Test BLT sorted
+    @test istril(StructuralTransformations.sorted_incidence_matrix(model), 2)
 end
 
 @testset "With parameters" begin
@@ -90,6 +93,9 @@ end
     sccsol = solve(sccprob, SimpleNewtonRaphson(); abstol = 1e-9)
     @test SciMLBase.successful_retcode(sccsol)
     @test norm(sccsol.resid) < norm(sol.resid)
+
+    # Test BLT sorted
+    @test istril(StructuralTransformations.sorted_incidence_matrix(sys), 1)
 end
 
 @testset "Transistor amplifier" begin
@@ -149,6 +155,9 @@ end
     sccsol = solve(sccprob, NewtonRaphson(); abstol = 1e-12)
 
     @test sol.u≈sccsol.u atol=1e-10
+
+    # Test BLT sorted
+    @test istril(StructuralTransformations.sorted_incidence_matrix(sys), 1)
 end
 
 @testset "Expression caching" begin

From 1e564207241814265fc9870a100625f382f84a60 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Jun 2025 14:50:09 +0530
Subject: [PATCH 10/11] feat: implement `sorted_incidence_matrix(sys)`

---
 .../StructuralTransformations.jl              |  2 +-
 src/structural_transformation/utils.jl        | 21 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/structural_transformation/StructuralTransformations.jl b/src/structural_transformation/StructuralTransformations.jl
index df30f7d2d2..eeb480e0f7 100644
--- a/src/structural_transformation/StructuralTransformations.jl
+++ b/src/structural_transformation/StructuralTransformations.jl
@@ -26,7 +26,7 @@ using ModelingToolkit: System, AbstractSystem, var_from_nested_derivative, Diffe
                        filter_kwargs, lower_varname_with_unit,
                        lower_shift_varname_with_unit, setio, SparseMatrixCLIL,
                        get_fullvars, has_equations, observed,
-                       Schedule, schedule
+                       Schedule, schedule, iscomplete, get_schedule
 
 using ModelingToolkit.BipartiteGraphs
 import .BipartiteGraphs: invview, complete
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index e900764698..50f9aaf887 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -205,6 +205,27 @@ function sorted_incidence_matrix(ts::TransformationState, val = true; only_algeq
     sparse(I, J, val, nsrcs(graph), ndsts(graph))
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Obtain the incidence matrix of the system sorted by the SCCs. Requires that the system is
+simplified and has a `schedule`.
+"""
+function sorted_incidence_matrix(sys::AbstractSystem)
+    if !iscomplete(sys) || get_tearing_state(sys) === nothing ||
+       get_schedule(sys) === nothing
+        error("A simplified `System` is required. Call `mtkcompile` on the system before creating an `SCCNonlinearProblem`.")
+    end
+    sched = get_schedule(sys)
+    var_sccs = sched.var_sccs
+
+    ts = get_tearing_state(sys)
+    imat = Graphs.incidence_matrix(ts.structure.graph)
+    buffer = similar(imat)
+    permute!(buffer, imat, 1:size(imat, 2), reduce(vcat, var_sccs))
+    buffer
+end
+
 ###
 ### Structural and symbolic utilities
 ###

From 26437e211b80588e5a7b2b6bcb9ed6584e0853da Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 2 Jun 2025 17:12:18 +0530
Subject: [PATCH 11/11] test: fix undefined import in tests

---
 test/symbolic_events.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/symbolic_events.jl b/test/symbolic_events.jl
index 2c6253e70d..5968976653 100644
--- a/test/symbolic_events.jl
+++ b/test/symbolic_events.jl
@@ -2,7 +2,6 @@ using ModelingToolkit, OrdinaryDiffEq, StochasticDiffEq, JumpProcesses, Test
 using SciMLStructures: canonicalize, Discrete
 using ModelingToolkit: SymbolicContinuousCallback,
                        SymbolicDiscreteCallback,
-                       get_callback,
                        t_nounits as t,
                        D_nounits as D,
                        affects, affect_negs, system, observed, AffectSystem