diff --git a/lib/ModelingToolkitTearing/Project.toml b/lib/ModelingToolkitTearing/Project.toml index cb49a90..4a6442b 100644 --- a/lib/ModelingToolkitTearing/Project.toml +++ b/lib/ModelingToolkitTearing/Project.toml @@ -25,7 +25,7 @@ BipartiteGraphs = "0.1.3" CommonSolve = "0.2" DocStringExtensions = "0.7, 0.8, 0.9" Graphs = "1" -ModelingToolkitBase = "1.0.0" +ModelingToolkitBase = "1.2.0" Moshi = "0.3" OffsetArrays = "1" OrderedCollections = "1.8.1" diff --git a/lib/ModelingToolkitTearing/src/clock_inference/interface.jl b/lib/ModelingToolkitTearing/src/clock_inference/interface.jl index 2024ff1..e8ca3c5 100644 --- a/lib/ModelingToolkitTearing/src/clock_inference/interface.jl +++ b/lib/ModelingToolkitTearing/src/clock_inference/interface.jl @@ -73,6 +73,9 @@ function system_subset(ts::TearingState, ieqs::Vector{Int}, iieqs::Vector{Int}, @set! ts.sys.initialization_eqs = initeqs[iieqs] @set! ts.original_eqs = ts.original_eqs[ieqs] @set! ts.structure = system_subset(ts.structure, ieqs, ivars) + if !isempty(ts.eqs_source) + @set! ts.eqs_source = ts.eqs_source[ieqs] + end if all(eq -> eq.rhs isa StateMachineOperator, MTKBase.get_eqs(ts.sys)) names = Symbol[] for eq in MTKBase.get_eqs(ts.sys) diff --git a/lib/ModelingToolkitTearing/src/tearingstate.jl b/lib/ModelingToolkitTearing/src/tearingstate.jl index 066c622..b9066f1 100644 --- a/lib/ModelingToolkitTearing/src/tearingstate.jl +++ b/lib/ModelingToolkitTearing/src/tearingstate.jl @@ -58,6 +58,13 @@ mutable struct TearingState <: StateSelection.TransformationState{System} """ additional_observed::Vector{Equation} statemachines::Vector{System} + """ + Source information for each equation in the `TearingState`. `Vector{Symbol}` for each + equation representing the path of the subsystem to which it belongs. Empty entries + indicate unknown source. If this field is empty, either the system has no equations + or source information is unknown. + """ + eqs_source::Vector{Vector{Symbol}} end function Base.show(io::IO, state::TearingState) @@ -83,15 +90,37 @@ function Base.push!(ev::EquationsView, eq) push!(ev.ts.extra_eqs, eq) end -function TearingState(sys::System; check::Bool = true, sort_eqs::Bool = true) +function TearingState(sys::System, source_info::Union{Nothing, MTKBase.EquationSourceInformation} = nothing; check::Bool = true, sort_eqs::Bool = true) # flatten system - sys = MTKBase.flatten(sys) + if source_info === nothing + sys = MTKBase.flatten(sys) + else + @assert isempty(MTKBase.get_systems(sys)) """ + If `source_info` is provided to `TearingState`, the system must be flattened. + """ + end sys = MTKBase.discrete_unknowns_to_parameters(sys) sys = MTKBase.discover_globalscoped(sys) MTKBase.check_no_parameter_equations(sys) iv = MTKBase.get_iv(sys) + sources = Vector{Symbol}[] # flatten array equations - eqs = MTKBase.flatten_equations(equations(sys)) + if source_info === nothing + eqs = MTKBase.flatten_equations(equations(sys)) + else + eqs = Equation[] + @assert length(equations(sys)) == length(source_info.eqs_source) """ + Mismatch between source information provided to `TearingState` and the structure \ + of the system. + """ + for (eq, src) in zip(equations(sys), source_info.eqs_source) + scal_eq = MTKBase.flatten_equation(eq) + append!(eqs, scal_eq) + for _ in scal_eq + push!(sources, src) + end + end + end original_eqs = copy(eqs) neqs = length(eqs) param_derivative_map = Dict{SymbolicT, SymbolicT}() @@ -234,6 +263,9 @@ function TearingState(sys::System; check::Bool = true, sort_eqs::Bool = true) eqs = eqs[sortidxs] original_eqs = original_eqs[sortidxs] symbolic_incidence = symbolic_incidence[sortidxs] + if !isempty(sources) + sources = sources[sortidxs] + end end dervaridxs = OrderedSet{Int}() @@ -262,7 +294,7 @@ function TearingState(sys::System; check::Bool = true, sort_eqs::Bool = true) structure = SystemStructure(complete(var_to_diff), complete(eq_to_diff), complete(graph), nothing, var_types, false) return TearingState(sys, fullvars, structure, Equation[], param_derivative_map, - no_deriv_params, original_eqs, Equation[], typeof(sys)[]) + no_deriv_params, original_eqs, Equation[], typeof(sys)[], sources) end function sort_fullvars(fullvars::Vector{SymbolicT}, dervaridxs::Vector{Int}, var_types::Vector{VariableType}, @nospecialize(iv::Union{SymbolicT, Nothing}))