WIP: Partial State Selection

This implements the basic outline of partial state selection.
It is functional and resolves the example in #988, but we do not
currently implement any tearing priorities or handling of linear
subsystems, so it's pretty easy for state selection to end up
selecting states that are singular.

Co-authored-by: Yingbo Ma <[email protected]>

Author: Keno

src/structural_transformation/StructuralTransformations.jl

@@ -36,7 +36,7 @@ using SparseArrays
 
 using NonlinearSolve
 
-export tearing, dae_index_lowering, check_consistency
+export tearing, partial_state_selection, dae_index_lowering, check_consistency
 export build_torn_function, build_observed_function, ODAEProblem
 export sorted_incidence_matrix
 
@@ -45,6 +45,7 @@ include("pantelides.jl")
 include("bipartite_tearing/modia_tearing.jl")
 include("tearing.jl")
 include("symbolics_tearing.jl")
+include("partial_state_selection.jl")
 include("codegen.jl")
 
 end # module

src/structural_transformation/pantelides.jl

@@ -72,13 +72,14 @@ end
 Perform Pantelides algorithm.
 """
 function pantelides!(sys::ODESystem; maxiters = 8000)
+    find_solvables!(sys)
     s = structure(sys)
     # D(j) = assoc[j]
-    @unpack graph, var_to_diff = s
-    return (sys, pantelides!(graph, var_to_diff)...)
+    @unpack graph, var_to_diff, solvable_graph = s
+    return (sys, pantelides!(graph, solvable_graph, var_to_diff)...)
 end
 
-function pantelides!(graph, var_to_diff; maxiters = 8000)
+function pantelides!(graph, solvable_graph, var_to_diff; maxiters = 8000)
     neqs = nsrcs(graph)
     nvars = nv(var_to_diff)
     vcolor = falses(nvars)
@@ -106,7 +107,7 @@ function pantelides!(graph, var_to_diff; maxiters = 8000)
             for var in eachindex(vcolor); vcolor[var] || continue
                 # introduce a new variable
                 nvars += 1
-                add_vertex!(graph, DST)
+                add_vertex!(graph, DST); add_vertex!(solvable_graph, DST)
                 # the new variable is the derivative of `var`
 
                 add_edge!(var_to_diff, var, add_vertex!(var_to_diff))
@@ -116,13 +117,16 @@ function pantelides!(graph, var_to_diff; maxiters = 8000)
             for eq in eachindex(ecolor); ecolor[eq] || continue
                 # introduce a new equation
                 neqs += 1
-                add_vertex!(graph, SRC)
+                add_vertex!(graph, SRC); add_vertex!(solvable_graph, SRC)
                 # the new equation is created by differentiating `eq`
                 eq_diff = add_vertex!(eq_to_diff)
                 add_edge!(eq_to_diff, eq, eq_diff)
                 for var in 𝑠neighbors(graph, eq)
                     add_edge!(graph, eq_diff, var)
                     add_edge!(graph, eq_diff, var_to_diff[var])
+                    # If you have f(x) = 0, then the derivative is (∂f/∂x) ẋ = 0.
+                    # which is linear, thus solvable in ẋ.
+                    add_edge!(solvable_graph, eq_diff, var_to_diff[var])
                 end
             end
 

src/structural_transformation/partial_state_selection.jl

@@ -0,0 +1,75 @@
+function partial_state_selection_graph!(sys::ODESystem)
+    s = get_structure(sys)
+    (s isa SystemStructure) || (sys = initialize_system_structure(sys))
+    s = structure(sys)
+    find_solvables!(sys)
+    @set! s.graph = complete(s.graph)
+    @set! sys.structure = s
+    (sys, partial_state_selection_graph!(s.graph, s.solvable_graph, s.var_to_diff)...)
+end
+
+struct SelectedState; end
+function partial_state_selection_graph!(graph, solvable_graph, var_to_diff)
+    var_eq_matching, eq_to_diff = pantelides!(graph, solvable_graph, var_to_diff)
+    eq_to_diff = complete(eq_to_diff)
+
+    eqlevel = map(1:nsrcs(graph)) do eq
+        level = 0
+        while eq_to_diff[eq] !== nothing
+            eq = eq_to_diff[eq]
+            level += 1
+        end
+        level
+    end
+
+    varlevel = map(1:ndsts(graph)) do var
+        level = 0
+        while var_to_diff[var] !== nothing
+            var = var_to_diff[var]
+            level += 1
+        end
+        level
+    end
+
+    all_selected_states = Int[]
+
+    level = 0
+    level_vars = [var for var in 1:ndsts(graph) if varlevel[var] == 0 && invview(var_to_diff)[var] !== nothing]
+
+    # TODO: Is this actually useful or should we just compute another maximal matching?
+    for var in 1:ndsts(graph)
+        if !(var in level_vars)
+            var_eq_matching[var] = unassigned
+        end
+    end
+
+    while level < maximum(eqlevel)
+        var_eq_matching = tear_graph_modia(graph, solvable_graph;
+            eqfilter  = eq->eqlevel[eq] == level && invview(eq_to_diff)[eq] !== nothing,
+            varfilter = var->(var in level_vars && !(var in all_selected_states)))
+        for var in level_vars
+            if var_eq_matching[var] === unassigned
+                selected_state = invview(var_to_diff)[var]
+                push!(all_selected_states, selected_state)
+                #=
+                    # TODO: This is what the Matteson paper says, but it doesn't
+                    #       quite seem to work.
+                    while selected_state !== nothing
+                        push!(all_selected_states, selected_state)
+                        selected_state = invview(var_to_diff)[selected_state]
+                    end
+                =#
+            end
+        end
+        level += 1
+        level_vars = [var for var = 1:ndsts(graph) if varlevel[var] == level && invview(var_to_diff)[var] !== nothing]
+    end
+
+    var_eq_matching = tear_graph_modia(graph, solvable_graph;
+        varfilter = var->!(var in all_selected_states))
+    var_eq_matching = Matching{Union{Unassigned, SelectedState}}(var_eq_matching)
+    for var in all_selected_states
+        var_eq_matching[var] = SelectedState()
+    end
+    return var_eq_matching, eq_to_diff
+end

src/structural_transformation/symbolics_tearing.jl

@@ -3,16 +3,49 @@ function tearing_sub(expr, dict, s)
     s ? simplify(expr) : expr
 end
 
-function tearing_reassemble(sys, var_eq_matching; simplify=false)
+function tearing_reassemble(sys, var_eq_matching, eq_to_diff=nothing; simplify=false)
     s = structure(sys)
-    @unpack fullvars, solvable_graph, graph = s
+    @unpack fullvars, solvable_graph, var_to_diff, graph = s
 
     eqs = equations(sys)
 
+    ### Add the differentiated equations and variables
+    D = Differential(get_iv(sys))
+    if length(fullvars) != length(var_to_diff)
+        for i = (length(fullvars)+1):length(var_to_diff)
+            push!(fullvars, D(fullvars[invview(var_to_diff)[i]]))
+        end
+    end
+
+    ### Add the differentiated equations
+    neweqs = copy(eqs)
+    if eq_to_diff !== nothing
+        eq_to_diff = complete(eq_to_diff)
+        for i = (length(eqs)+1):length(eq_to_diff)
+            eq = neweqs[invview(eq_to_diff)[i]]
+            push!(neweqs, ModelingToolkit.expand_derivatives(0 ~ D(eq.rhs - eq.lhs)))
+        end
+
+        ### Replace derivatives of non-selected states by dumy derivatives
+        dummy_subs = Dict()
+        for var = 1:length(fullvars)
+            invview(var_to_diff)[var] === nothing && continue
+            if var_eq_matching[invview(var_to_diff)[var]] !== SelectedState()
+                fullvar = fullvars[var]
+                subst_fullvar = tearing_sub(fullvar, dummy_subs, simplify)
+                dummy_subs[fullvar] = fullvars[var] = diff2term(unwrap(subst_fullvar))
+                var_to_diff[invview(var_to_diff)[var]] = nothing
+            end
+        end
+        neweqs = map(neweqs) do eq
+            0 ~ tearing_sub(eq.rhs - eq.lhs, dummy_subs, simplify)
+        end
+    end
+
     ### extract partition information
     function solve_equation(ieq, iv)
         var = fullvars[iv]
-        eq = eqs[ieq]
+        eq = neweqs[ieq]
         rhs = value(solve_for(eq, var; simplify=simplify, check=false))
 
         if var in vars(rhs)
@@ -30,7 +63,7 @@ function tearing_reassemble(sys, var_eq_matching; simplify=false)
         end
         var => rhs
     end
-    is_solvable(eq, iv) = eq !== unassigned && BipartiteEdge(eq, iv) in solvable_graph
+    is_solvable(eq, iv) = isa(eq, Int) && BipartiteEdge(eq, iv) in solvable_graph
 
     solved_equations = Int[]
     solved_variables = Int[]
@@ -41,32 +74,31 @@ function tearing_reassemble(sys, var_eq_matching; simplify=false)
         push!(solved_equations, ieq); push!(solved_variables, iv)
     end
 
+    isdiffvar(var) = invview(var_to_diff)[var] !== nothing && var_eq_matching[invview(var_to_diff)[var]] === SelectedState()
     solved = Dict(solve_equation(ieq, iv) for (ieq, iv) in zip(solved_equations, solved_variables))
     obseqs = [var ~ rhs for (var, rhs) in solved]
 
     # Rewrite remaining equations in terms of solved variables
     function substitute_equation(ieq)
-        eq = eqs[ieq]
-        if isdiffeq(eq)
-            return eq.lhs ~ tearing_sub(eq.rhs, solved, simplify)
-        else
-            if !(eq.lhs isa Number && eq.lhs == 0)
-                eq = 0 ~ eq.rhs - eq.lhs
-            end
-            rhs = tearing_sub(eq.rhs, solved, simplify)
-            if rhs isa Symbolic
-                return 0 ~ rhs
-            else # a number
-                if abs(rhs) > 100eps(float(rhs))
-                    @warn "The equation $eq is not consistent. It simplifed to 0 == $rhs."
-                end
-                return nothing
+        eq = neweqs[ieq]
+        if !(eq.lhs isa Number && eq.lhs == 0)
+            eq = 0 ~ eq.rhs - eq.lhs
+        end
+        rhs = tearing_sub(eq.rhs, solved, simplify)
+        if rhs isa Symbolic
+            return 0 ~ rhs
+        else # a number
+            if abs(rhs) > 100eps(float(rhs))
+                @warn "The equation $eq is not consistent. It simplifed to 0 == $rhs."
             end
+            return nothing
         end
     end
 
-    neweqs = Any[substitute_equation(ieq) for ieq in 1:length(eqs) if !(ieq in solved_equations)]
+    diffeqs = [fullvars[iv] ~ tearing_sub(solved[fullvars[iv]], solved, simplify) for iv in solved_variables if isdiffvar(iv)]
+    neweqs = Any[substitute_equation(ieq) for ieq in 1:length(neweqs) if !(ieq in solved_equations)]
     filter!(!isnothing, neweqs)
+    prepend!(neweqs, diffeqs)
 
     # Contract the vertices in the structure graph to make the structure match
     # the new reality of the system we've just created.
@@ -78,30 +110,17 @@ function tearing_reassemble(sys, var_eq_matching; simplify=false)
 
     @set! s.graph = graph
     @set! s.fullvars = [v for (i, v) in enumerate(fullvars) if i in active_vars]
+    @set! s.var_to_diff = DiffGraph(Union{Int, Nothing}[v for (i, v) in enumerate(s.var_to_diff) if i in active_vars])
     @set! s.vartype = [v for (i, v) in enumerate(s.vartype) if i in active_vars]
     @set! s.algeqs = [e for (i, e) in enumerate(s.algeqs) if i in active_eqs]
 
     @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.states = [fullvars[i] for i in active_vars]
     @set! sys.observed = [observed(sys); obseqs]
     return sys
 end
 
-"""
-    tearing(sys; simplify=false)
-
-Tear the nonlinear equations in system. When `simplify=true`, we simplify the
-new residual residual equations after tearing. End users are encouraged to call [`structural_simplify`](@ref)
-instead, which calls this function internally.
-"""
-function tearing(sys; simplify=false)
-    sys = init_for_tearing(sys)
-    var_eq_matching = tear_graph(sys)
-
-    tearing_reassemble(sys, var_eq_matching; simplify=simplify)
-end
-
 function init_for_tearing(sys)
     s = get_structure(sys)
     if !(s isa SystemStructure)
@@ -119,6 +138,38 @@ end
 function tear_graph(sys)
     s = structure(sys)
     @unpack graph, solvable_graph = s
-    tear_graph_modia(graph, solvable_graph;
-        varfilter=var->isalgvar(s, var), eqfilter=eq->s.algeqs[eq])
+    var_eq_matching = Matching{Union{Unassigned, SelectedState}}(tear_graph_modia(graph, solvable_graph;
+        varfilter=var->isalgvar(s, var), eqfilter=eq->s.algeqs[eq]))
+    for var in 1:ndsts(graph)
+        if !isalgvar(s, var)
+            var_eq_matching[var] = SelectedState()
+        end
+    end
+    var_eq_matching
+end
+
+"""
+    tearing(sys; simplify=false)
+
+Tear the nonlinear equations in system. When `simplify=true`, we simplify the
+new residual residual equations after tearing. End users are encouraged to call [`structural_simplify`](@ref)
+instead, which calls this function internally.
+"""
+function tearing(sys; simplify=false)
+    sys = init_for_tearing(sys)
+    var_eq_matching = tear_graph(sys)
+
+    tearing_reassemble(sys, var_eq_matching; simplify=simplify)
+end
+
+"""
+    tearing(sys; simplify=false)
+
+Perform partial state selection and tearing.
+"""
+function partial_state_selection(sys; simplify=false)
+    sys = init_for_tearing(sys)
+    sys, var_eq_matching, eq_to_diff = partial_state_selection_graph!(sys)
+
+    tearing_reassemble(sys, var_eq_matching, eq_to_diff; simplify=simplify)
 end

src/structural_transformation/utils.jl

@@ -164,10 +164,8 @@ function find_solvables!(sys)
     eqs = equations(sys)
     empty!(solvable_graph)
     for (i, eq) in enumerate(eqs)
-        isdiffeq(eq) && continue
         term = value(eq.rhs - eq.lhs)
         for j in 𝑠neighbors(graph, i)
-            isalgvar(s, j) || continue
             var = fullvars[j]
             isinput(var) && continue
             a, b, islinear = linear_expansion(term, var)

src/systems/systemstructure.jl

@@ -62,6 +62,8 @@ struct DiffGraph <: Graphs.AbstractGraph{Int}
     primal_to_diff::Vector{Union{Int, Nothing}}
     diff_to_primal::Union{Nothing, Vector{Union{Int, Nothing}}}
 end
+DiffGraph(primal_to_diff::Vector{Union{Int, Nothing}}) =
+    DiffGraph(primal_to_diff, nothing)
 DiffGraph(n::Integer, with_badj::Bool=false) = DiffGraph(Union{Int, Nothing}[nothing for _=1:n],
     with_badj ? Union{Int, Nothing}[nothing for _=1:n] : nothing)
 

test/structural_transformation/index_reduction.jl

@@ -134,3 +134,20 @@ p = [
 prob_auto = ODEProblem(new_sys,u0,(0.0,10.0),p)
 sol = solve(prob_auto, Rodas5());
 #plot(sol, vars=(D(x), y))
+
+let pss_pendulum2 = partial_state_selection(pendulum2)
+    @test length(equations(pss_pendulum2)) == 3
+    @test length(equations(ModelingToolkit.ode_order_lowering(pss_pendulum2))) == 4
+end
+
+eqs = [D(x) ~ w,
+       D(y) ~ z,
+       D(w) ~ T*x,
+       D(z) ~ T*y - g,
+       0 ~ x^2 + y^2 - L^2]
+pendulum = ODESystem(eqs, t, [x, y, w, z, T], [L, g], name=:pendulum)
+
+let pss_pendulum = partial_state_selection(pendulum)
+    @test length(equations(pss_pendulum)) == 3
+    @test length(equations(ModelingToolkit.ode_order_lowering(pss_pendulum))) == 4
+end