Cache subed equations and only use full_equations in calculate_* functions

Author: YingboMa

src/ModelingToolkit.jl

@@ -181,7 +181,7 @@ export Differential, expand_derivatives, @derivatives
 export Equation, ConstrainedEquation
 export Term, Sym
 export SymScope, LocalScope, ParentScope, GlobalScope
-export independent_variables, independent_variable, states, parameters, equations, controls, observed, structure
+export independent_variables, independent_variable, states, parameters, equations, controls, observed, structure, full_equations
 export structural_simplify, expand_connections
 export DiscreteSystem, DiscreteProblem
 

src/structural_transformation/StructuralTransformations.jl

@@ -21,7 +21,7 @@ using ModelingToolkit: ODESystem, AbstractSystem,var_from_nested_derivative, Dif
                        ExtraVariablesSystemException,
                        get_postprocess_fbody, vars!,
                        IncrementalCycleTracker, add_edge_checked!, topological_sort,
-                       invalidate_cache!
+                       invalidate_cache!, Substitutions
 
 using ModelingToolkit.BipartiteGraphs
 import .BipartiteGraphs: invview
@@ -43,6 +43,7 @@ export tearing_assignments, tearing_substitution
 export build_torn_function, build_observed_function, ODAEProblem
 export sorted_incidence_matrix
 export torn_system_jacobian_sparsity
+export full_equations
 
 include("utils.jl")
 include("pantelides.jl")

src/structural_transformation/symbolics_tearing.jl

@@ -34,9 +34,11 @@ function tearing_sub(expr, dict, s)
     s ? simplify(expr) : expr
 end
 
-function tearing_substitution(sys::AbstractSystem; simplify=false)
-    empty_substitutions(sys) && return sys
-    subs, = get_substitutions(sys)
+function full_equations(sys::AbstractSystem; simplify=false)
+    empty_substitutions(sys) && return equations(sys)
+    substitutions = get_substitutions(sys)
+    substitutions.subed_eqs === nothing && return substitutions.subed_eqs
+    @unpack subs = substitutions
     solved = Dict(eq.lhs => eq.rhs for eq in subs)
     neweqs = map(equations(sys)) do eq
         if isdiffeq(eq)
@@ -54,6 +56,12 @@ function tearing_substitution(sys::AbstractSystem; simplify=false)
         end
         eq
     end
+    substitutions.subed_eqs = neweqs
+    return neweqs
+end
+
+function tearing_substitution(sys::AbstractSystem; kwargs...)
+    neweqs = full_equations(sys::AbstractSystem; kwargs...)
     @set! sys.eqs = neweqs
     @set! sys.substitutions = nothing
 end
@@ -64,7 +72,7 @@ function tearing_assignments(sys::AbstractSystem)
         deps = Int[]
         sol_states = Code.LazyState()
     else
-        subs, deps = get_substitutions(sys)
+        @unpack subs, deps = get_substitutions(sys)
         assignments = [Assignment(eq.lhs, eq.rhs) for eq in subs]
         sol_states = Code.NameState(Dict(eq.lhs => Symbol(eq.lhs) for eq in subs))
     end
@@ -105,13 +113,13 @@ function tearing_reassemble(sys, var_eq_matching; simplify=false)
     end
     subgraph = substitution_graph(graph, solved_equations, solved_variables, var_eq_matching)
     toporder = topological_sort_by_dfs(subgraph)
-    substitutions = [solve_equation(
-                                    eqs[solved_equations[i]],
-                                    fullvars[solved_variables[i]],
-                                    simplify
-                                   ) for i in toporder]
+    subeqs = [solve_equation(
+                             eqs[solved_equations[i]],
+                             fullvars[solved_variables[i]],
+                             simplify
+                            ) for i in toporder]
     invtoporder = invperm(toporder)
-    deps = [[invtoporder[n] for n in neighborhood(subgraph, j, Inf, dir=:in) if n!=j] for (i, j) in enumerate(toporder)]
+    deps = [Int[invtoporder[n] for n in neighborhood(subgraph, j, Inf, dir=:in) if n!=j] for (i, j) in enumerate(toporder)]
 
     # Rewrite remaining equations in terms of solved variables
 
@@ -134,8 +142,8 @@ function tearing_reassemble(sys, var_eq_matching; simplify=false)
     @set! sys.structure = s
     @set! sys.eqs = neweqs
     @set! sys.states = [s.fullvars[idx] for idx in 1:length(s.fullvars) if !isdervar(s, idx)]
-    @set! sys.observed = [observed(sys); substitutions]
-    @set! sys.substitutions = substitutions, deps
+    @set! sys.observed = [observed(sys); subeqs]
+    @set! sys.substitutions = Substitutions(subeqs, deps)
     return sys
 end
 

src/systems/abstractsystem.jl

@@ -129,6 +129,14 @@ Generate a function to evaluate the system's equations.
 """
 function generate_function end
 
+
+mutable struct Substitutions
+    subs::Vector{Equation}
+    deps::Vector{Vector{Int}}
+    subed_eqs::Union{Nothing,Vector{Equation}}
+end
+Substitutions(subs, deps) = Substitutions(subs, deps, nothing)
+
 Base.nameof(sys::AbstractSystem) = getfield(sys, :name)
 
 #Deprecated

src/systems/diffeqs/abstractodesystem.jl

@@ -5,7 +5,7 @@ function calculate_tgrad(sys::AbstractODESystem;
   # We need to remove explicit time dependence on the state because when we
   # have `u(t) * t` we want to have the tgrad to be `u(t)` instead of `u'(t) *
   # t + u(t)`.
-  rhs = [detime_dvs(eq.rhs) for eq ∈ equations(sys)]
+  rhs = [detime_dvs(eq.rhs) for eq ∈ full_equations(sys)]
   iv = get_iv(sys)
   xs = states(sys)
   rule = Dict(map((x, xt) -> xt=>x, detime_dvs.(xs), xs))
@@ -23,7 +23,7 @@ function calculate_jacobian(sys::AbstractODESystem;
     if cache isa Tuple && cache[2] == (sparse, simplify)
         return cache[1]
     end
-    rhs = [eq.rhs for eq ∈ equations(sys)]
+    rhs = [eq.rhs for eq ∈ full_equations(sys)]
 
     iv = get_iv(sys)
     dvs = states(sys)
@@ -45,7 +45,7 @@ function calculate_control_jacobian(sys::AbstractODESystem;
         return cache[1]
     end
 
-    rhs = [eq.rhs for eq ∈ equations(sys)]
+    rhs = [eq.rhs for eq ∈ full_equations(sys)]
 
     iv = get_iv(sys)
     ctrls = controls(sys)
@@ -263,7 +263,7 @@ function time_varying_as_func(x, sys::AbstractTimeDependentSystem)
 end
 
 function calculate_massmatrix(sys::AbstractODESystem; simplify=false)
-    eqs = [eq for eq in equations(sys) if !isdifferenceeq(eq)]
+    eqs = [eq for eq in full_equations(sys) if !isdifferenceeq(eq)]
     dvs = states(sys)
     M = zeros(length(eqs),length(eqs))
     state2idx = Dict(s => i for (i, s) in enumerate(dvs))
@@ -285,7 +285,7 @@ function jacobian_sparsity(sys::AbstractODESystem)
     sparsity = torn_system_jacobian_sparsity(sys)
     sparsity === nothing || return sparsity
 
-    jacobian_sparsity([eq.rhs for eq ∈ equations(sys)],
+    jacobian_sparsity([eq.rhs for eq ∈ full_equations(sys)],
                       [dv for dv in states(sys)])
 end
 

src/utils.jl

@@ -430,15 +430,15 @@ isarray(x) = x isa AbstractArray || x isa Symbolics.Arr
 function empty_substitutions(sys)
     has_substitutions(sys) || return true
     subs = get_substitutions(sys)
-    isnothing(subs) || isempty(last(subs))
+    isnothing(subs) || isempty(subs.deps)
 end
 
 function get_substitutions_and_solved_states(sys; no_postprocess=false)
     if empty_substitutions(sys)
         sol_states = Code.LazyState()
         pre = no_postprocess ? (ex -> ex) : get_postprocess_fbody(sys)
     else
-        subs, = get_substitutions(sys)
+        @unpack subs = get_substitutions(sys)
         sol_states = Code.NameState(Dict(eq.lhs => Symbol(eq.lhs) for eq in subs))
         if no_postprocess
             pre = ex -> Let(Assignment[Assignment(eq.lhs, eq.rhs) for eq in subs], ex)