sort states in initialize_system_structure

Author: YingboMa

src/structural_transformation/tearing.jl

@@ -119,7 +119,6 @@ function tearing_reassemble(sys; simplify=false)
         return nothing
     end
 
-
     ### update equations
     odestats = []
     for idx in eachindex(fullvars); isdervar(s, idx) && continue
@@ -191,7 +190,6 @@ function tearing_reassemble(sys; simplify=false)
     return sys
 end
 
-
 """
     algebraic_equations_scc(sys)
 

src/systems/diffeqs/odesystem.jl

@@ -281,36 +281,6 @@ function _eq_unordered(a, b)
     return true
 end
 
-"""
-$(SIGNATURES)
-
-Sort the states so that the mass matrix is as identity-like as possible.
-"""
-function sort_states(sys::ODESystem)
-    sys = flatten(sys)
-    eqs = equations(sys)
-    sts = states(sys)
-    sorted_states = OrderedSet()
-    var2idx = Dict(sts .=> eachindex(sts))
-    idx = 0
-    for eq in eqs
-        if isdiffeq(eq)
-            var = arguments(eq.lhs)[1]
-            if haskey(var2idx, var)
-                pop!(var2idx, var)
-                push!(sorted_states, var)
-            end
-        end
-    end
-
-    var2idx = collect(pairs(var2idx))
-    sorted_states = collect(sorted_states)
-    append!(sorted_states, map(x->x[1], sort(var2idx, by=x->x[2])))
-
-    @set! sys.states = sorted_states
-    @set! sys.structure = nothing
-end
-
 function collect_differential_variables(sys::ODESystem)
     eqs = equations(sys)
     vars = Set()

src/systems/systemstructure.jl

@@ -111,7 +111,7 @@ function initialize_system_structure(sys)
         else
             eq = 0 ~ eq′.rhs - eq′.lhs
         end
-        vars!(vars, eq)
+        vars!(vars, eq.rhs)
         isalgeq = true
         statevars = []
         for var in vars
@@ -140,6 +140,24 @@ function initialize_system_structure(sys)
         end
     end
 
+    # sort `fullvars` such that the mass matrix is as diagonal as possible.
+    sorted_fullvars = OrderedSet(fullvars[dervaridxs])
+    for dervaridx in dervaridxs
+        dervar = fullvars[dervaridx]
+        diffvar = arguments(dervar)[1]
+        if !(diffvar in sorted_fullvars)
+            push!(sorted_fullvars, diffvar)
+        end
+    end
+    for v in fullvars
+        if !(v in sorted_fullvars)
+            push!(sorted_fullvars, v)
+        end
+    end
+    fullvars = collect(sorted_fullvars)
+    var2idx = Dict(fullvars .=> eachindex(fullvars))
+    dervaridxs = 1:length(dervaridxs)
+
     nvars = length(fullvars)
     diffvars = []
     vartype = fill(DIFFERENTIAL_VARIABLE, nvars)