Merge pull request #1872 from SciML/myb/one_tearing_state

Update structural info after alias elimination

Author: YingboMa

src/structural_transformation/partial_state_selection.jl

@@ -248,6 +248,7 @@ function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, ja
         end
     end
     if diff_va !== nothing
+        # differentiated alias
         n_dummys = length(dummy_derivatives)
         needed = count(x -> x isa Int, diff_to_eq) - n_dummys
         n = 0
@@ -265,36 +266,68 @@ function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, ja
         @warn "The number of dummy derivatives ($n_dummys) does not match the number of differentiated equations ($n_diff_eqs)."
     end
     dummy_derivatives_set = BitSet(dummy_derivatives)
+
+    if ag !== nothing
+        function isreducible(x)
+            # `k` is reducible if all lower differentiated variables are.
+            isred = true
+            while isred
+                if x in dummy_derivatives_set
+                    break
+                end
+                x = diff_to_var[x]
+                x === nothing && break
+                if !haskey(ag, x)
+                    isred = false
+                end
+            end
+            isred
+        end
+        irreducible_set = BitSet()
+        for (k, (_, v)) in ag
+            isreducible(k) || push!(irreducible_set, k)
+            isreducible(k) || push!(irreducible_set, k)
+            push!(irreducible_set, v)
+        end
+    end
+
+    is_not_present = v -> isempty(𝑑neighbors(graph, v)) &&
+        (ag === nothing || (haskey(ag, v) && !(v in irreducible_set)))
+    # Derivatives that are either in the dummy derivatives set or ended up not
+    # participating in the system at all are not considered differential
+    is_some_diff = let dummy_derivatives_set = dummy_derivatives_set
+        v -> !(v in dummy_derivatives_set) &&
+            !(var_to_diff[v] === nothing && is_not_present(v))
+    end
+
+    # We don't want tearing to give us `y_t ~ D(y)`, so we skip equations with
+    # actually differentiated variables.
+    isdiffed = let diff_to_var = diff_to_var
+        v -> diff_to_var[v] !== nothing && is_some_diff(v)
+    end
+
     # We can eliminate variables that are not a selected state (differential
     # variables). Selected states are differentiated variables that are not
     # dummy derivatives.
-    can_eliminate = let var_to_diff = var_to_diff,
-        dummy_derivatives_set = dummy_derivatives_set
-
+    can_eliminate = let var_to_diff = var_to_diff
         v -> begin
             if ag !== nothing
                 haskey(ag, v) && return false
             end
             dv = var_to_diff[v]
-            dv === nothing || dv in dummy_derivatives_set
+            dv === nothing && return true
+            is_some_diff(dv) || return true
+            return false
         end
     end
 
-    # We don't want tearing to give us `y_t ~ D(y)`, so we skip equations with
-    # actually differentiated variables.
-    isdiffed = let diff_to_var = diff_to_var, dummy_derivatives_set = dummy_derivatives_set
-        v -> diff_to_var[v] !== nothing && !(v in dummy_derivatives_set)
-    end
-
     var_eq_matching = tear_graph_modia(structure, isdiffed,
                                        Union{Unassigned, SelectedState};
                                        varfilter = can_eliminate)
     for v in eachindex(var_eq_matching)
-        if ag !== nothing && haskey(ag, v) && iszero(ag[v][1])
-            continue
-        end
+        is_not_present(v) && continue
         dv = var_to_diff[v]
-        (dv === nothing || dv in dummy_derivatives_set) && continue
+        (dv === nothing || !is_some_diff(dv)) && continue
         var_eq_matching[v] = SelectedState()
     end
 

src/structural_transformation/symbolics_tearing.jl

@@ -217,7 +217,8 @@ function check_diff_graph(var_to_diff, fullvars)
 end
 =#
 
-function tearing_reassemble(state::TearingState, var_eq_matching; simplify = false)
+function tearing_reassemble(state::TearingState, var_eq_matching, ag = nothing;
+                            simplify = false)
     @unpack fullvars, sys, structure = state
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
 
@@ -318,6 +319,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
         var -> diff_to_var[var] !== nothing
     end
 
+    #retear = BitSet()
     # There are three cases where we want to generate new variables to convert
     # the system into first order (semi-implicit) ODEs.
     #
@@ -468,20 +470,38 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
         for (ogidx, dx_idx, x_t_idx) in Iterators.reverse(subinfo)
             # We need a loop here because both `D(D(x))` and `D(x_t)` need to be
             # substituted to `x_tt`.
-            for idx in (ogidx, dx_idx)
+            for idx in (ogidx == dx_idx ? ogidx : (ogidx, dx_idx))
                 subidx = ((idx => x_t_idx),)
                 # This handles case 2.2
-                if var_eq_matching[idx] isa Int
-                    var_eq_matching[x_t_idx] = var_eq_matching[idx]
-                end
                 substitute_vars!(structure, subidx, idx_buffer, sub_callback!;
                                  exclude = order_lowering_eqs)
+                if var_eq_matching[idx] isa Int
+                    original_assigned_eq = var_eq_matching[idx]
+                    # This removes the assignment of the variable `idx`, so we
+                    # should consider assign them again later.
+                    var_eq_matching[x_t_idx] = original_assigned_eq
+                    #if !isempty(𝑑neighbors(graph, idx))
+                    #    push!(retear, idx)
+                    #end
+                end
             end
         end
         empty!(subinfo)
         empty!(subs)
     end
 
+    #ict = IncrementalCycleTracker(DiCMOBiGraph{true}(graph, var_eq_matching); dir = :in)
+    #for idx in retear
+    #    for alternative_eq in 𝑑neighbors(solvable_graph, idx)
+    #        # skip actually differentiated variables
+    #        any(𝑠neighbors(graph, alternative_eq)) do alternative_v
+    #            ((vv = diff_to_var[alternative_v]) !== nothing &&
+    #             var_eq_matching[vv] === SelectedState())
+    #        end && continue
+    #        try_assign_eq!(ict, idx, alternative_eq) && break
+    #    end
+    #end
+
     # Will reorder equations and states to be:
     # [diffeqs; ...]
     # [diffvars; ...]
@@ -555,9 +575,11 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
     if length(diff_vars_set) != length(diff_vars)
         error("Tearing internal error: lowering DAE into semi-implicit ODE failed!")
     end
+    solved_variables_set = BitSet(solved_variables)
+    ag === nothing || union!(solved_variables_set, keys(ag))
     invvarsperm = [diff_vars;
                    setdiff!(setdiff(1:ndsts(graph), diff_vars_set),
-                            BitSet(solved_variables))]
+                            solved_variables_set)]
     varsperm = zeros(Int, ndsts(graph))
     for (i, v) in enumerate(invvarsperm)
         varsperm[v] = i
@@ -580,7 +602,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
     # 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, varsperm, eqsperm,
-                               length(solved_variables))
+                               length(solved_variables), length(solved_variables_set))
 
     # Update system
     new_var_to_diff = complete(DiffGraph(length(invvarsperm)))
@@ -632,6 +654,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
         isdiffeq(eq) || continue
         obs_sub[eq.lhs] = eq.rhs
     end
+    # TODO: compute the dependency correctly so that we don't have to do this
     obs = substitute.([oldobs; subeqs], (obs_sub,))
     @set! sys.observed = obs
     @set! state.sys = sys
@@ -688,7 +711,8 @@ end
 Perform index reduction and use the dummy derivative technique to ensure that
 the system is balanced.
 """
-function dummy_derivative(sys, state = TearingState(sys); simplify = false, kwargs...)
+function dummy_derivative(sys, state = TearingState(sys), ag = nothing; simplify = false,
+                          kwargs...)
     jac = let state = state
         (eqs, vars) -> begin
             symeqs = EquationsView(state)[eqs]
@@ -710,6 +734,7 @@ function dummy_derivative(sys, state = TearingState(sys); simplify = false, kwar
             p
         end
     end
-    var_eq_matching = dummy_derivative_graph!(state, jac; state_priority, kwargs...)
-    tearing_reassemble(state, var_eq_matching; simplify = simplify)
+    var_eq_matching = dummy_derivative_graph!(state, jac, (ag, nothing); state_priority,
+                                              kwargs...)
+    tearing_reassemble(state, var_eq_matching, ag; simplify = simplify)
 end

src/structural_transformation/tearing.jl

@@ -22,7 +22,7 @@ function masked_cumsum!(A::Vector)
 end
 
 function contract_variables(graph::BipartiteGraph, var_eq_matching::Matching,
-                            var_rename, eq_rename, nelim)
+                            var_rename, eq_rename, nelim_eq, nelim_var)
     dig = DiCMOBiGraph{true}(graph, var_eq_matching)
 
     # Update bipartite graph
@@ -31,7 +31,7 @@ function contract_variables(graph::BipartiteGraph, var_eq_matching::Matching,
          for v′ in neighborhood(dig, v, Inf; dir = :in) if var_rename[v′] != 0]
     end
 
-    newgraph = BipartiteGraph(nsrcs(graph) - nelim, ndsts(graph) - nelim)
+    newgraph = BipartiteGraph(nsrcs(graph) - nelim_eq, ndsts(graph) - nelim_var)
     for e in 𝑠vertices(graph)
         ne = eq_rename[e]
         ne == 0 && continue

src/structural_transformation/utils.jl

@@ -43,10 +43,12 @@ end
 ###
 ### Structural check
 ###
-function check_consistency(state::TearingState)
+function check_consistency(state::TearingState, ag = nothing)
     fullvars = state.fullvars
     @unpack graph, var_to_diff = state.structure
-    n_highest_vars = count(v -> length(outneighbors(var_to_diff, v)) == 0,
+    n_highest_vars = count(v -> var_to_diff[v] === nothing &&
+                                    !isempty(𝑑neighbors(graph, v)) &&
+                                    (ag === nothing || !haskey(ag, v) || ag[v] != v),
                            vertices(var_to_diff))
     neqs = nsrcs(graph)
     is_balanced = n_highest_vars == neqs
@@ -69,11 +71,12 @@ function check_consistency(state::TearingState)
     # details, check the equation (15) of the original paper.
     extended_graph = (@set graph.fadjlist = Vector{Int}[graph.fadjlist;
                                                         map(collect, edges(var_to_diff))])
-    extended_var_eq_matching = maximal_matching(extended_graph)
+    extended_var_eq_matching = maximal_matching(extended_graph, eq -> true,
+                                                v -> ag === nothing || !haskey(ag, v))
 
     unassigned_var = []
     for (vj, eq) in enumerate(extended_var_eq_matching)
-        if eq === unassigned
+        if eq === unassigned && (ag === nothing || !haskey(ag, vj))
             push!(unassigned_var, fullvars[vj])
         end
     end
@@ -228,7 +231,7 @@ function find_solvables!(state::TearingState; kwargs...)
     return nothing
 end
 
-function linear_subsys_adjmat!(state::TransformationState)
+function linear_subsys_adjmat!(state::TransformationState; kwargs...)
     graph = state.structure.graph
     if state.structure.solvable_graph === nothing
         state.structure.solvable_graph = BipartiteGraph(nsrcs(graph), ndsts(graph))
@@ -240,7 +243,7 @@ function linear_subsys_adjmat!(state::TransformationState)
     coeffs = Int[]
     to_rm = Int[]
     for i in eachindex(eqs)
-        all_int_vars, rhs = find_eq_solvables!(state, i, to_rm, coeffs)
+        all_int_vars, rhs = find_eq_solvables!(state, i, to_rm, coeffs; kwargs...)
 
         # Check if all states in the equation is both linear and homogeneous,
         # i.e. it is in the form of

src/systems/abstractsystem.jl

@@ -1035,13 +1035,14 @@ function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false
     has_io = io !== nothing
     has_io && markio!(state, io...)
     state, input_idxs = inputs_to_parameters!(state, io)
-    sys = alias_elimination!(state)
+    sys, ag = alias_elimination!(state; kwargs...)
+    #ag = AliasGraph(length(ag))
     # TODO: avoid construct `TearingState` again.
-    state = TearingState(sys)
-    has_io && markio!(state, io..., check = false)
-    check_consistency(state)
-    find_solvables!(state; kwargs...)
-    sys = dummy_derivative(sys, state; simplify)
+    #state = TearingState(sys)
+    #has_io && markio!(state, io..., check = false)
+    check_consistency(state, ag)
+    #find_solvables!(state; kwargs...)
+    sys = dummy_derivative(sys, state, ag; simplify)
     fullstates = [map(eq -> eq.lhs, observed(sys)); states(sys)]
     @set! sys.observed = topsort_equations(observed(sys), fullstates)
     invalidate_cache!(sys)