Optimize and simplify ode_order_lowering implementation

Author: YingboMa

Project.toml

@@ -5,6 +5,7 @@ version = "3.6.1"
 
 [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,8 +25,7 @@ function ode_order_lowering(sys::ODESystem)
 end
 
 function ode_order_lowering(eqs, iv, states)
-    var_order = Dict{Variable,Int}()
-    vars = Variable[]
+    var_order = OrderedDict{Variable,Int}()
     D = Differential(iv())
     diff_eqs = Equation[]
     diff_vars = Variable[]
@@ -38,19 +38,16 @@ function ode_order_lowering(eqs, iv, states)
             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!(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)