Make sure structural_simplify is composable

Author: YingboMa

src/structural_transformation/StructuralTransformations.jl

@@ -16,7 +16,7 @@ using ModelingToolkit: ODESystem, var_from_nested_derivative, Differential,
                        states, equations, vars, Symbolic, diff2term, value,
                        operation, arguments, Sym, Term, simplify, solve_for,
                        isdiffeq, isdifferential,
-                       get_structure, default_u0, default_p
+                       get_structure, get_reduced_states, default_u0, default_p
 
 using ModelingToolkit.BipartiteGraphs
 using ModelingToolkit.SystemStructures

src/structural_transformation/tearing.jl

@@ -180,6 +180,7 @@ function tearing_reassemble(sys; simplify=false)
     @set! sys.structure.algeqs = newalgeqs
     @set! sys.eqs = neweqs
     @set! sys.states = newstates
+    @set! sys.reduced_states = [get_reduced_states(sys); solvars]
     @set! sys.observed = vcat(observed(sys), obseqs)
     return sys
 end

src/systems/abstractsystem.jl

@@ -154,6 +154,7 @@ for prop in [
              :inequality_constraints
              :controls
              :loss
+             :reduced_states
             ]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
@@ -471,9 +472,9 @@ Structurally simplify algebraic equations in a system and compute the
 topological sort of the observed equations.
 """
 function structural_simplify(sys::AbstractSystem)
-    ss = states(sys)
     sys = tearing(alias_elimination(sys))
     s = structure(sys)
-    @set! sys.observed = topsort_equations(observed(sys), ss)
+    fullstates = [get_reduced_states(sys); states(sys)]
+    @set! sys.observed = topsort_equations(observed(sys), fullstates)
     return sys
 end

src/systems/alias_elimination.jl

@@ -3,12 +3,13 @@ using SymbolicUtils: Rewriters
 const KEEP = typemin(Int)
 
 function alias_elimination(sys)
-    sys = flatten(sys)
-    s = get_structure(sys)
-    if !(s isa SystemStructure)
-        sys = initialize_system_structure(sys)
-        s = structure(sys)
-    end
+    # FIXME: update `structure` too
+    #sys = flatten(sys)
+    #s = get_structure(sys)
+    #if !(s isa SystemStructure)
+    sys = initialize_system_structure(sys)
+    s = structure(sys)
+    #end
     is_linear_equations, eadj, cadj = find_linear_equations(sys)
 
     v_eliminated, v_types, n_null_vars, degenerate_equations, linear_equations = alias_eliminate_graph(
@@ -18,9 +19,12 @@ function alias_elimination(sys)
     s = structure(sys)
     @unpack fullvars, graph = s
 
+    n_reduced_states = length(v_eliminated) - n_null_vars
+    reduced_states = similar(v_eliminated, Any, n_reduced_states)
     subs = OrderedDict()
-    if length(v_eliminated) - n_null_vars > 0
-        for v in v_eliminated[n_null_vars+1:end]
+    if n_reduced_states > 0
+        for (i, v) in enumerate(@view v_eliminated[n_null_vars+1:end])
+            reduced_states[i] = fullvars[v]
             subs[fullvars[v]] = iszeroterm(v_types, v) ? 0.0 :
                                 isalias(v_types, v) ? fullvars[alias(v_types, v)] :
                                 -fullvars[negalias(v_types, v)]
@@ -63,6 +67,7 @@ function alias_elimination(sys)
 
     @set! sys.eqs = eqs
     @set! sys.states = newstates
+    @set! sys.reduced_states = [get_reduced_states(sys); reduced_states]
     @set! sys.observed = [get_observed(sys); [lhs ~ rhs for (lhs, rhs) in pairs(subs)]]
     @set! sys.structure = nothing
     return sys

src/systems/diffeqs/odesystem.jl

@@ -74,6 +74,7 @@ struct ODESystem <: AbstractODESystem
     structure: structural information of the system
     """
     structure::Any
+    reduced_states::Vector
 end
 
 function ODESystem(
@@ -101,7 +102,7 @@ function ODESystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    ODESystem(deqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, default_u0, default_p, nothing)
+    ODESystem(deqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, default_u0, default_p, nothing, [])
 end
 
 var_from_nested_derivative(x, i=0) = (missing, missing)

src/systems/nonlinear/nonlinearsystem.jl

@@ -48,6 +48,7 @@ struct NonlinearSystem <: AbstractSystem
     structure: structural information of the system
     """
     structure::Any
+    reduced_states::Any
 end
 
 function NonlinearSystem(eqs, states, ps;
@@ -60,7 +61,7 @@ function NonlinearSystem(eqs, states, ps;
     default_p isa Dict || (default_p = Dict(default_p))
     default_u0 = Dict(value(k) => value(default_u0[k]) for k in keys(default_u0))
     default_p = Dict(value(k) => value(default_p[k]) for k in keys(default_p))
-    NonlinearSystem(eqs, value.(states), value.(ps), observed, name, systems, default_u0, default_p, nothing)
+    NonlinearSystem(eqs, value.(states), value.(ps), observed, name, systems, default_u0, default_p, nothing, [])
 end
 
 function calculate_jacobian(sys::NonlinearSystem;sparse=false,simplify=false)