Add some convenient interfaces

Author: YingboMa

src/systems/systemstructure.jl

@@ -1,3 +1,11 @@
+module SystemStructures
+
+using ..ModelingToolkit
+using ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!
+using SymbolicUtils: arguments
+using ..BipartiteGraphs
+using UnPack
+using Setfield
 using SparseArrays
 
 #=
@@ -27,10 +35,18 @@ for v in 𝑣vertices(graph); active_𝑣vertices[v] || continue
 end
 =#
 
+export SystemStructure, initialize_system_structure
+export diffvars_range, dervars_range, algvars_range
+export isdiffvars, isdervars, isalgvars
+export DIFFERENTIAL_VARIABLE, ALGEBRAIC_VARIABLE, DERIVATIVE_VARIABLE
+export DIFFERENTIAL_EQUATION, ALGEBRAIC_EQUATION
+export vartype, eqtype
+
 struct SystemStructure
     dxvar_offset::Int
     fullvars::Vector # [xvar; dxvars; algvars]
     varassoc::Vector{Int}
+    algeqs::BitVector
     graph::BipartiteGraph{Int}
     solvable_graph::BipartiteGraph{Int}
     assign::Vector{Int}
@@ -39,12 +55,33 @@ struct SystemStructure
     partitions::Vector{NTuple{4, Vector{Int}}}
 end
 
+diffvars_range(s::SystemStructure) = 1:s.dxvar_offset
+dervars_range(s::SystemStructure) = s.dxvar_offset+1:2s.dxvar_offset
+algvars_range(s::SystemStructure) = 2s.dxvar_offset+1:length(s.fullvars)
+
+isdiffvar(s::SystemStructure, var::Integer) = var in diffvars_range(s)
+isdervar(s::SystemStructure, var::Integer) = var in dervars_range(s)
+isalgvar(s::SystemStructure, var::Integer) = var in algvars_range(s)
+
+@enum VariableType DIFFERENTIAL_VARIABLE ALGEBRAIC_VARIABLE DERIVATIVE_VARIABLE
+
+function vartype(s::SystemStructure, var::Integer)::VariableType
+    isdiffvar(s, var) ? DIFFERENTIAL_VARIABLE :
+    isdervar(s, var)  ? DERIVATIVE_VARIABLE :
+    isalgvar(s, var)  ? ALGEBRAIC_VARIABLE : error("Variable $var out of bounds")
+end
+
+@enum EquationType DIFFERENTIAL_EQUATION ALGEBRAIC_EQUATION
+
+eqtype(s::SystemStructure, eq::Integer)::EquationType = s.algeqs[eq] ? ALGEBRAIC_EQUATION : DIFFERENTIAL_EQUATION
+
 function initialize_system_structure(sys)
-    sys, dxvar_offset, fullvars, varassoc, graph, solvable_graph = init_graph(flatten(sys))
+    sys, dxvar_offset, fullvars, varassoc, algeqs, graph, solvable_graph = init_graph(sys)
     @set sys.structure = SystemStructure(
                                          dxvar_offset,
                                          fullvars,
                                          varassoc,
+                                         algeqs,
                                          graph,
                                          solvable_graph,
                                          Int[],
@@ -74,8 +111,10 @@ end
 function collect_variables(sys)
     dxvars = []
     eqs = equations(sys)
+    algeqs = falses(length(eqs))
     for (i, eq) in enumerate(eqs)
         if isdiffeq(eq)
+            algeqs[i] = true
             lhs = eq.lhs
             # Make sure that the LHS is a first order derivative of a var.
             @assert !(arguments(lhs)[1] isa Differential) "The equation $eq is not first order"
@@ -86,11 +125,11 @@ function collect_variables(sys)
 
     xvars = (first ∘ var_from_nested_derivative).(dxvars)
     algvars  = setdiff(states(sys), xvars)
-    return xvars, dxvars, algvars
+    return xvars, dxvars, algvars, algeqs
 end
 
 function init_graph(sys)
-    xvars, dxvars, algvars = collect_variables(sys)
+    xvars, dxvars, algvars, algeqs = collect_variables(sys)
     dxvar_offset = length(xvars)
     algvar_offset = 2dxvar_offset
 
@@ -119,5 +158,7 @@ function init_graph(sys)
     end
 
     varassoc = Int[(1:dxvar_offset) .+ dxvar_offset; zeros(Int, length(fullvars) - dxvar_offset)] # variable association list
-    sys, dxvar_offset, fullvars, varassoc, graph, solvable_graph
+    sys, dxvar_offset, fullvars, varassoc, algeqs, graph, solvable_graph
 end
+
+end # module