feat: implement new trivial_tearing! interface for MTK

Author: AayushSabharwal

lib/ModelingToolkitTearing/src/stateselection_interface.jl

@@ -326,3 +326,38 @@ function manual_dispatch_is_small_int(@nospecialize(x::Number))::Int
         end
     end
 end
+
+function StateSelection.rm_eqs_vars!(
+        structure::SystemStructure, eqs_to_rm::Vector{Int}, vars_to_rm::Vector{Int};
+        eqs_sorted_and_uniqued::Bool = false, vars_sorted_and_uniqued::Bool = false
+    )
+    old_to_new_eq, old_to_new_var = StateSelection.default_rm_eqs_vars!(
+        structure, eqs_to_rm, vars_to_rm; eqs_sorted_and_uniqued, vars_sorted_and_uniqued
+    )
+    deleteat!(structure.state_priorities, vars_to_rm)
+    deleteat!(structure.var_types, vars_to_rm)
+    return old_to_new_eq, old_to_new_var
+end
+
+function StateSelection.rm_eqs_vars!(
+        state::TearingState, eqs_to_rm::Vector{Int}, vars_to_rm::Vector{Int};
+        eqs_sorted_and_uniqued::Bool = false, vars_sorted_and_uniqued::Bool = false
+    )
+    (; structure, sys) = state
+    old_to_new_eq, old_to_new_var = StateSelection.rm_eqs_vars!(
+        structure, eqs_to_rm, vars_to_rm; eqs_sorted_and_uniqued, vars_sorted_and_uniqued
+    )
+    deleteat!(state.fullvars, vars_to_rm)
+    eqs = copy(MTKBase.get_eqs(state.sys))
+    deleteat!(eqs, eqs_to_rm)
+    deleteat!(state.original_eqs, eqs_to_rm)
+    if !isempty(state.eqs_source)
+        deleteat!(state.eqs_source, eqs_to_rm)
+    end
+
+    @set! sys.eqs = eqs
+    state.sys = sys
+
+    return old_to_new_eq, old_to_new_var
+end
+

lib/ModelingToolkitTearing/src/trivial_tearing_interface.jl

@@ -23,21 +23,8 @@ function StateSelection.possibly_explicit_equations(state::TearingState)
     return Iterators.map(mapfn, Iterators.filter(filterfn, eachindex(state.original_eqs)))
 end
 
-function StateSelection.trivial_tearing_postprocess!(ts::TearingState, torn_eqs::OrderedSet{Int}, torn_vars::OrderedSet{Int})
-    append!(ts.additional_observed, @view ts.original_eqs[collect(torn_eqs)])
-    sort!(torn_vars)
-    sort!(torn_eqs)
-    if ts.structure.var_types !== nothing
-        deleteat!(ts.structure.var_types, torn_vars)
-    end
-    deleteat!(ts.fullvars, torn_vars)
-    deleteat!(ts.structure.state_priorities, torn_vars)
-    deleteat!(ts.original_eqs, torn_eqs)
-    sys = ts.sys
-    eqs = copy(MTKBase.get_eqs(sys))
-    deleteat!(eqs, torn_eqs)
-    @set! sys.eqs = eqs
-    ts.sys = sys
+function StateSelection.trivial_tearing_postprocess!(ts::TearingState, torn_eqs::Vector{Int}, torn_vars::Vector{Int})
+    append!(ts.additional_observed, @view ts.original_eqs[torn_eqs])
     return ts
 end