Merge pull request #391 from SciML/myb/ordering

Place algebraic equations after differential equations

Author: YingboMa

Project.toml

@@ -1,10 +1,11 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "3.6.1"
+version = "3.6.2"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqJump = "c894b116-72e5-5b58-be3c-e6d8d4ac2b12"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
@@ -28,6 +29,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
 ArrayInterface = "2.8"
+DataStructures = "0.17"
 DiffEqBase = "6.28"
 DiffEqJump = "6.7.5"
 DiffRules = "0.1, 1.0"

src/systems/diffeqs/first_order_transform.jl

@@ -1,3 +1,4 @@
+using DataStructures: OrderedDict
 function lower_varname(var::Variable, idv, order)
     order == 0 && return var
     name = Symbol(var.name, :ˍ, string(idv.name)^order)
@@ -24,43 +25,42 @@ function ode_order_lowering(sys::ODESystem)
 end
 
 function ode_order_lowering(eqs, iv, states)
-    var_order = Dict{Variable,Int}()
-    vars = Variable[]
-    new_eqs = Equation[]
-    new_vars = Variable[]
+    var_order = OrderedDict{Variable,Int}()
     D = Differential(iv())
+    diff_eqs = Equation[]
+    diff_vars = Variable[]
+    alge_eqs = Equation[]
+    alge_vars = Variable[]
 
     for (i, (eq, ss)) ∈ enumerate(zip(eqs, states))
         if isequal(eq.lhs, Constant(0))
-            push!(new_vars, ss)
-            push!(new_eqs, eq)
+            push!(alge_vars, ss)
+            push!(alge_eqs, eq)
         else
             var, maxorder = var_from_nested_derivative(eq.lhs)
-            if maxorder > get(var_order, var, 0)
-                var_order[var] = maxorder
-                any(isequal(var), vars) || push!(vars, var)
-            end
+            # only save to the dict when we need to lower the order to save memory
+            maxorder > get(var_order, var, 1) && (var_order[var] = maxorder)
             var′ = lower_varname(var, iv, maxorder - 1)
             rhs′ = rename_lower_order(eq.rhs)
-            push!(new_vars, var′)
-            push!(new_eqs, D(var′(iv())) ~ rhs′)
+            push!(diff_vars, var′)
+            push!(diff_eqs, D(var′(iv())) ~ rhs′)
         end
     end
 
-    for var ∈ vars
-        order = var_order[var]
+    for (var, order) ∈ var_order
         for o in (order-1):-1:1
             lvar = lower_varname(var, iv, o-1)
             rvar = lower_varname(var, iv, o)
-            push!(new_vars, lvar)
+            push!(diff_vars, lvar)
 
             rhs = rvar(iv())
             eq = Differential(iv())(lvar(iv())) ~ rhs
-            push!(new_eqs, eq)
+            push!(diff_eqs, eq)
         end
     end
 
-    return (new_eqs, new_vars)
+    # we want to order the equations and variables to be `(diff, alge)`
+    return (vcat(diff_eqs, alge_eqs), vcat(diff_vars, alge_vars))
 end
 
 function rename_lower_order(O::Expression)

test/lowering_solving.jl

@@ -1,15 +1,24 @@
-using ModelingToolkit, OrdinaryDiffEq, Test
+using ModelingToolkit, OrdinaryDiffEq, Test, LinearAlgebra
 
 @parameters t σ ρ β
-@variables x(t) y(t) z(t)
+@variables x(t) y(t) z(t) k(t)
 @derivatives D'~t
 
 eqs = [D(D(x)) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ x*y - β*z]
 
-sys = ODESystem(eqs)
-sys = ode_order_lowering(sys)
+sys′ = ODESystem(eqs)
+sys = ode_order_lowering(sys′)
+
+eqs2 = [0 ~ x*y - k,
+        D(D(x)) ~ σ*(y-x),
+        D(y) ~ x*(ρ-z)-y,
+        D(z) ~ x*y - β*z]
+sys2 = ODESystem(eqs2, t, [x, y, z, k], sys′.ps)
+sys2 = ode_order_lowering(sys2)
+# test equation/varible ordering
+ModelingToolkit.calculate_massmatrix(sys2) == Diagonal([1, 1, 1, 1, 0])
 
 u0 = [D(x) => 2.0,
 	  x => 1.0,