Automatic state detection and some performance optimization

Author: YingboMa

src/ModelingToolkit.jl

@@ -6,6 +6,7 @@ using Latexify, Unitful, ArrayInterface
 using MacroTools
 using UnPack: @unpack
 using DiffEqJump
+using DataStructures: OrderedDict, OrderedSet
 
 using Base.Threads
 import MacroTools: splitdef, combinedef, postwalk, striplines

src/systems/diffeqs/first_order_transform.jl

@@ -1,4 +1,3 @@
-using DataStructures: OrderedDict
 function lower_varname(var::Variable, idv, order)
     order == 0 && return var
     name = Symbol(var.name, :ˍ, string(idv.name)^order)

src/systems/diffeqs/odesystem.jl

@@ -77,19 +77,47 @@ end
 var_from_nested_derivative(x) = var_from_nested_derivative(x,0)
 var_from_nested_derivative(x::Constant) = (missing, missing)
 var_from_nested_derivative(x,i) = x.op isa Differential ? var_from_nested_derivative(x.args[1],i+1) : (x.op,i)
+
+function extract_eqs_states_ps(eqs::AbstractArray{<:Equation}, iv)
+    # NOTE: this assumes that the order of algebric equations doesn't matter
+    diffvars = OrderedSet{Variable}()
+    allstates = OrderedSet{Variable}()
+    ps = OrderedSet{Variable}()
+    # reorder equations such that it is in the form of `diffeq, algeeq`
+    diffeq = Equation[]
+    algeeq = Equation[]
+    for eq in eqs
+        for var in vars(eq.rhs for eq ∈ eqs)
+            var isa Variable || continue
+            if isparameter(var)
+                isequal(var, iv) || push!(ps, var)
+            else
+                push!(allstates, var)
+            end
+        end
+        if eq.lhs isa Constant
+            push!(algeeq, eq)
+        else
+            diffvar = first(var_from_nested_derivative(eq.lhs))
+            diffvar in diffvars && throw(ArgumentError("The differential variable $diffvar is not unique in the system of equations."))
+            push!(diffvars, diffvar)
+            push!(diffeq, eq)
+        end
+    end
+    algevars = setdiff(allstates, diffvars)
+    # the orders here are very important!
+    return append!(diffeq, algeeq), vcat(collect(diffvars), collect(algevars)), ps
+end
+
 iv_from_nested_derivative(x) = x.op isa Differential ? iv_from_nested_derivative(x.args[1]) : x.args[1].op
 iv_from_nested_derivative(x::Constant) = missing
 
 function ODESystem(eqs; kwargs...)
     ivs = unique(skipmissing(iv_from_nested_derivative(eq.lhs) for eq ∈ eqs))
-    length(ivs) == 1 || throw(ArgumentError("one independent variable currently supported"))
+    length(ivs) == 1 || throw(ArgumentError("An ODESystem can only have one independent variable."))
     iv = first(ivs)
-
-    dvs = unique(skipmissing(var_from_nested_derivative(eq.lhs)[1] for eq ∈ eqs))
-    ps = filter(vars(eq.rhs for eq ∈ eqs)) do x
-        isparameter(x) & !isequal(x, iv)
-    end |> collect
-    ODESystem(eqs, iv, dvs, ps; kwargs...)
+    eqs, dvs, ps = extract_eqs_states_ps(eqs, iv)
+    return ODESystem(eqs, iv, dvs, ps; kwargs...)
 end
 
 Base.:(==)(sys1::ODESystem, sys2::ODESystem) =