Add state_priority to prefer some states in dummy derivative

Author: YingboMa

src/structural_transformation/partial_state_selection.jl

@@ -147,11 +147,11 @@ function partial_state_selection_graph!(structure::SystemStructure, var_eq_match
     var_eq_matching
 end
 
-function dummy_derivative_graph!(state::TransformationState, jac = nothing; kwargs...)
+function dummy_derivative_graph!(state::TransformationState, jac = nothing, state_priority = nothing; kwargs...)
     state.structure.solvable_graph === nothing && find_solvables!(state; kwargs...)
     var_eq_matching = complete(pantelides!(state))
     complete!(state.structure)
-    dummy_derivative_graph!(state.structure, var_eq_matching, jac)
+    dummy_derivative_graph!(state.structure, var_eq_matching, jac, state_priority)
 end
 
 function compute_diff_level(diff_to_x)
@@ -171,7 +171,7 @@ function compute_diff_level(diff_to_x)
     return xlevel, maxlevel
 end
 
-function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, jac)
+function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, jac, state_priority)
     @unpack eq_to_diff, var_to_diff, graph = structure
     diff_to_eq = invview(eq_to_diff)
     diff_to_var = invview(var_to_diff)
@@ -191,12 +191,17 @@ function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, ja
         iszero(maxlevel) && continue
 
         rank_matching = Matching(nvars)
+        isfirst = true
         for _ in maxlevel:-1:1
             eqs = filter(eq -> diff_to_eq[eq] !== nothing, eqs)
             nrows = length(eqs)
             iszero(nrows) && break
             eqs_set = BitSet(eqs)
 
+            if state_priority !== nothing && isfirst
+                sort!(vars, by = state_priority)
+            end
+            isfirst = false
             # TODO: making the algorithm more robust
             # 1. If the Jacobian is a integer matrix, use Bareiss to check
             # linear independence. (done)

src/structural_transformation/symbolics_tearing.jl

@@ -648,10 +648,27 @@ 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, kwargs...)
-    function jac(eqs, vars)
-        symeqs = EquationsView(state)[eqs]
-        Symbolics.jacobian((x -> x.rhs).(symeqs), state.fullvars[vars])
+    jac = let state = state
+        (eqs, vars) -> begin
+            symeqs = EquationsView(state)[eqs]
+            Symbolics.jacobian((x -> x.rhs).(symeqs), state.fullvars[vars])
+        end
+    end
+    state_priority = let state = state
+        var -> begin
+            p = 0.0
+            var_to_diff = state.structure.var_to_diff
+            diff_to_var = invview(var_to_diff)
+            while var_to_diff[var] !== nothing
+                var = var_to_diff[var]
+            end
+            while true
+                p = max(p, ModelingToolkit.state_priority(state.fullvars[var]))
+                (var = diff_to_var[var]) === nothing && break
+            end
+            p
+        end
     end
-    var_eq_matching = dummy_derivative_graph!(state, jac; kwargs...)
+    var_eq_matching = dummy_derivative_graph!(state, jac, state_priority; kwargs...)
     tearing_reassemble(state, var_eq_matching; simplify = simplify)
 end

src/variables.jl

@@ -4,12 +4,14 @@ struct VariableNoiseType end
 struct VariableInput end
 struct VariableOutput end
 struct VariableIrreducible end
+struct VariableStatePriority end
 Symbolics.option_to_metadata_type(::Val{:unit}) = VariableUnit
 Symbolics.option_to_metadata_type(::Val{:connect}) = VariableConnectType
 Symbolics.option_to_metadata_type(::Val{:noise}) = VariableNoiseType
 Symbolics.option_to_metadata_type(::Val{:input}) = VariableInput
 Symbolics.option_to_metadata_type(::Val{:output}) = VariableOutput
 Symbolics.option_to_metadata_type(::Val{:irreducible}) = VariableIrreducible
+Symbolics.option_to_metadata_type(::Val{:state_priority}) = VariableStatePriority
 
 abstract type AbstractConnectType end
 struct Equality <: AbstractConnectType end # Equality connection
@@ -26,6 +28,7 @@ end
 isinput(x) = isvarkind(VariableInput, x)
 isoutput(x) = isvarkind(VariableOutput, x)
 isirreducible(x) = isvarkind(VariableIrreducible, x) || isinput(x)
+state_priority(x) = convert(Float64, getmetadata(x, VariableStatePriority, 0.0))::Float64
 
 """
 $(SIGNATURES)