From bb2eb3f25008e3ff1a840051acf159c7c3086335 Mon Sep 17 00:00:00 2001
From: Yingbo Ma <mayingbo5@gmail.com>
Date: Thu, 8 Aug 2024 20:19:56 -0400
Subject: [PATCH] Experimental variable filtering in dummy derivative

---
 docs/src/basics/InputOutput.md                    |  5 -----
 .../partial_state_selection.jl                    | 15 +++++++++++----
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/docs/src/basics/InputOutput.md b/docs/src/basics/InputOutput.md
index 09bf47331a..84d8939ac7 100644
--- a/docs/src/basics/InputOutput.md
+++ b/docs/src/basics/InputOutput.md
@@ -34,10 +34,8 @@ This function takes a vector of variables that are to be considered inputs, i.e.
 
 ### Example:
 
-
 The following example implements a simple first-order system with an input `u` and state `x`. The function `f` is generated using `generate_control_function`, and the function `f` is then tested with random input and state values.
 
-
 ```@example inputoutput
 import ModelingToolkit: t_nounits as t, D_nounits as D
 @variables x(t)=0 u(t)=0 y(t)
@@ -50,10 +48,8 @@ f, x_sym, ps = ModelingToolkit.generate_control_function(sys, [u], simplify = tr
 nothing # hide
 ```
 
-
 We can inspect the state realization chosen by MTK
 
-
 ```@example inputoutput
 x_sym
 ```
@@ -70,7 +66,6 @@ nothing  # hide
 
 Now we can test the generated function `f` with random input and state values
 
-
 ```@example inputoutput
 p = [1]
 x = [rand()]
diff --git a/src/structural_transformation/partial_state_selection.jl b/src/structural_transformation/partial_state_selection.jl
index 9fcfdd48bc..405165beac 100644
--- a/src/structural_transformation/partial_state_selection.jl
+++ b/src/structural_transformation/partial_state_selection.jl
@@ -208,18 +208,23 @@ function dummy_derivative_graph!(
     col_order = Int[]
     nvars = ndsts(graph)
     eqs = Int[]
+    vars = Int[]
     next_eq_idxs = Int[]
     next_var_idxs = Int[]
     new_eqs = Int[]
     new_vars = Int[]
     eqs_set = BitSet()
-    for vars in var_sccs
+    for vars′ in var_sccs
         empty!(eqs)
-        for var in vars
+        empty!(vars)
+        for var in vars′
             eq = var_eq_matching[var]
             eq isa Int || continue
-            diff_to_eq[eq] === nothing && continue
-            push!(eqs, eq)
+            diff_to_eq[eq] === nothing || push!(eqs, eq)
+            if var_to_diff[var] !== nothing
+                error("Invalid SCC")
+            end
+            (diff_to_var[var] !== nothing && is_present(structure, var)) && push!(vars, var)
         end
         isempty(eqs) && continue
 
@@ -304,6 +309,8 @@ function dummy_derivative_graph!(
             for (i, var) in enumerate(vars)
                 ∫var = diff_to_var[var]
                 ∫var === nothing && continue
+                ∫∫var = diff_to_var[∫var]
+                ∫∫var === nothing && continue
                 if J !== nothing
                     push!(next_var_idxs, i)
                 end