Merge pull request #1844 from SciML/myb/state_priority

Add `state_priority` to prefer some states in dummy derivative

Author: YingboMa

src/structural_transformation/partial_state_selection.jl

@@ -159,11 +159,12 @@ end
 
 function dummy_derivative_graph!(state::TransformationState, jac = nothing,
                                  (ag, diff_va) = (nothing, nothing);
-                                 kwargs...)
+                                 state_priority = nothing, kwargs...)
     state.structure.solvable_graph === nothing && find_solvables!(state; kwargs...)
     var_eq_matching = complete(pantelides!(state, ag))
     complete!(state.structure)
-    dummy_derivative_graph!(state.structure, var_eq_matching, jac, (ag, diff_va))
+    dummy_derivative_graph!(state.structure, var_eq_matching, jac, (ag, diff_va),
+                            state_priority)
 end
 
 function compute_diff_level(diff_to_x)
@@ -184,7 +185,7 @@ function compute_diff_level(diff_to_x)
 end
 
 function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, jac,
-                                 (ag, diff_va))
+                                 (ag, diff_va), 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)
@@ -204,12 +205,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)