Add sort_states to make mass matrix identity-like

Author: YingboMa

src/systems/abstractsystem.jl

@@ -527,7 +527,7 @@ function structural_simplify(sys::AbstractSystem)
     if sys isa ODESystem
         sys = dae_index_lowering(sys)
     end
-    sys = tearing(sys)
+    sys = sort_states(tearing(sys))
     fullstates = [map(eq->eq.lhs, observed(sys)); states(sys)]
     @set! sys.observed = topsort_equations(observed(sys), fullstates)
     return sys

src/systems/diffeqs/odesystem.jl

@@ -279,3 +279,33 @@ function _eq_unordered(a, b)
     end
     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