Cache tearing state

Author: YingboMa

src/structural_transformation/StructuralTransformations.jl

@@ -15,13 +15,13 @@ using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Di
                        operation, arguments, Sym, Term, simplify, solve_for,
                        isdiffeq, isdifferential, isinput,
                        empty_substitutions, get_substitutions,
-                       get_structure, get_iv, independent_variables,
-                       has_structure, defaults, InvalidSystemException,
+                       get_tearing_state, get_iv, independent_variables,
+                       has_tearing_state, defaults, InvalidSystemException,
                        ExtraEquationsSystemException,
                        ExtraVariablesSystemException,
                        get_postprocess_fbody, vars!,
                        IncrementalCycleTracker, add_edge_checked!, topological_sort,
-                       invalidate_cache!, Substitutions
+                       invalidate_cache!, Substitutions, get_or_construct_tearing_state
 
 using ModelingToolkit.BipartiteGraphs
 import .BipartiteGraphs: invview

src/structural_transformation/codegen.jl

@@ -245,7 +245,7 @@ function build_torn_function(
         push!(rhss, eq.rhs)
     end
 
-    state = TearingState(sys)
+    state = get_or_construct_tearing_state(sys)
     fullvars = state.fullvars
     var_eq_matching, var_sccs = algebraic_variables_scc(state)
     condensed_graph = MatchedCondensationGraph(

src/structural_transformation/symbolics_tearing.jl

@@ -220,19 +220,24 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify=false
 
     # Contract the vertices in the structure graph to make the structure match
     # the new reality of the system we've just created.
-    #graph = contract_variables(graph, var_eq_matching, solved_variables)
+    graph = contract_variables(graph, var_eq_matching, solved_variables)
 
     # Update system
     active_vars = setdiff(BitSet(1:length(fullvars)), solved_variables)
 
+    @set! state.structure.graph = graph
+    @set! state.fullvars = [v for (i, v) in enumerate(fullvars) if i in active_vars]
+
     sys = state.sys
     @set! sys.eqs = neweqs
     isstatediff(i) = var_eq_matching[i] !== SelectedState() && invview(var_to_diff)[i] !== nothing && var_eq_matching[invview(var_to_diff)[i]] === SelectedState()
     @set! sys.states = [fullvars[i] for i in active_vars if !isstatediff(i)]
     @set! sys.observed = [observed(sys); subeqs]
     @set! sys.substitutions = Substitutions(subeqs, deps)
+    @set! state.sys = sys
+    @set! sys.tearing_state = state
 
-    return sys
+    return invalidate_cache!(sys)
 end
 
 function tearing(state::TearingState)

src/structural_transformation/utils.jl

@@ -258,10 +258,11 @@ function uneven_invmap(n::Int, list)
 end
 
 function torn_system_jacobian_sparsity(sys)
-    has_structure(sys) || return nothing
-    get_structure(sys) isa SystemStructure || return nothing
+    state = get_tearing_state(sys)
+    state isa TearingState || return nothing
     s = structure(sys)
-    @unpack fullvars, graph = s
+    graph = state.structure.graph
+    fullvars = state.fullvars
 
     states_idxs = findall(!isdifferential, fullvars)
     var2idx = Dict{Int,Int}(v => i for (i, v) in enumerate(states_idxs))

src/systems/abstractsystem.jl

@@ -225,6 +225,7 @@ for prop in [
              :connections
              :preface
              :torn_matching
+             :tearing_state
              :substitutions
             ]
     fname1 = Symbol(:get_, prop)
@@ -683,6 +684,18 @@ end
 
 Base.write(io::IO, sys::AbstractSystem) = write(io, readable_code(toexpr(sys)))
 
+function get_or_construct_tearing_state(sys)
+    if has_tearing_state(sys)
+        state = get_tearing_state(sys)
+        if state === nothing
+            state = TearingState(sys)
+        end
+    else
+        state = nothing
+    end
+    state
+end
+
 function Base.show(io::IO, ::MIME"text/plain", sys::AbstractSystem)
     eqs = equations(sys)
     if eqs isa AbstractArray
@@ -741,7 +754,7 @@ function Base.show(io::IO, ::MIME"text/plain", sys::AbstractSystem)
     if has_torn_matching(sys)
         # If the system can take a torn matching, then we can initialize a tearing
         # state on it. Do so and get show the structure.
-        state = TearingState(sys; check=false)
+        state = get_or_construct_tearing_state(sys)
         if state !== nothing
             Base.printstyled(io, "\nIncidence matrix:"; color=:magenta)
             show(io, incidence_matrix(state.structure.graph, Num(Sym{Real}(:×))))

src/systems/diffeqs/odesystem.jl

@@ -102,14 +102,18 @@ struct ODESystem <: AbstractODESystem
     """
     continuous_events::Vector{SymbolicContinuousCallback}
     """
+    tearing_state: cache for intermediate tearing state
+    """
+    tearing_state::Any
+    """
     substitutions: substitutions generated by tearing.
     """
     substitutions::Any
 
     function ODESystem(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad,
             jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults,
             torn_matching, connector_type, connections, preface, events,
-            substitutions=nothing; checks::Bool=true)
+            tearing_state=nothing, substitutions=nothing; checks::Bool=true)
         if checks
             check_variables(dvs,iv)
             check_parameters(ps,iv)
@@ -119,7 +123,7 @@ struct ODESystem <: AbstractODESystem
         end
         new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, torn_matching,
-            connector_type, connections, preface, events, substitutions)
+            connector_type, connections, preface, events, tearing_state, substitutions)
     end
 end
 

src/systems/discrete_system/discrete_system.jl

@@ -52,10 +52,6 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     """
     defaults::Dict
     """
-    structure: structural information of the system
-    """
-    structure::Any
-    """
     preface: inject assignment statements before the evaluation of the RHS function.
     """
     preface::Any
@@ -64,17 +60,21 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     """
     connector_type::Any
     """
+    tearing_state: cache for intermediate tearing state
+    """
+    tearing_state::Any
+    """
     substitutions: substitutions generated by tearing.
     """
     substitutions::Any
 
-    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, preface, connector_type, substitutions=nothing; checks::Bool = true)
+    function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, preface, connector_type, tearing_state=nothing, substitutions=nothing; checks::Bool = true)
         if checks
             check_variables(dvs, iv)
             check_parameters(ps, iv)
             all_dimensionless([dvs;ps;iv;ctrls]) || check_units(discreteEqs)
         end
-        new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, structure, preface, connector_type, substitutions)
+        new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults, preface, connector_type, tearing_state, substitutions)
     end
 end
 
@@ -118,7 +118,7 @@ function DiscreteSystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    DiscreteSystem(eqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, name, systems, defaults, nothing, preface, connector_type, kwargs...)
+    DiscreteSystem(eqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, name, systems, defaults, preface, connector_type, kwargs...)
 end
 
 

src/systems/nonlinear/nonlinearsystem.jl

@@ -51,15 +51,19 @@ struct NonlinearSystem <: AbstractTimeIndependentSystem
     """
     connector_type::Any
     """
+    tearing_state: cache for intermediate tearing state
+    """
+    tearing_state::Any
+    """
     substitutions: substitutions generated by tearing.
     """
     substitutions::Any
 
-    function NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, connector_type, substitutions=nothing; checks::Bool = true)
+    function NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, connector_type, tearing_state=nothing, substitutions=nothing; checks::Bool = true)
         if checks
             all_dimensionless([states;ps]) ||check_units(eqs)
         end
-        new(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, connector_type, substitutions)
+        new(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, connector_type, tearing_state, substitutions)
     end
 end