Iterate ssrm

Author: Keno

src/transform/codegen/dae_factory.jl

@@ -138,7 +138,31 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
     insert_node_here!(oc_compact,
         NewInstruction(Expr(:invoke, daef_ci, oc_sicm, (), in_u_mm, in_u_unassgn, in_du_unassgn, in_alg, out_du_mm, out_eq, Argument(6)), Nothing, line))
 
-    # Manually apply mass matrix
+    # TODO: We should not have to recompute this here
+    var_eq_matching = matching_for_key(result, key, state.structure)
+    (slot_assignments, var_assignment, eq_assignment) = assign_slots(state, key, var_eq_matching)
+
+    # Manually apply mass matrix and implicit equations between selected states
+    for v = 1:ndsts(state.structure.graph)
+        vdiff = state.structure.var_to_diff[v]
+        vdiff === nothing && continue
+
+        if var_eq_matching[v] !== SelectedState() || var_eq_matching[vdiff] !== SelectedState()
+            # Solved variables were already handled above
+            continue
+        end
+
+        (kind, slot) = var_assignment[v]
+        (dkind, dslot) = var_assignment[vdiff]
+        @assert kind == AssignedDiff
+        @assert dkind in (AssignedDiff, UnassignedDiff)
+
+        v_val = insert_node_here!(oc_compact,
+            NewInstruction(Expr(:call, Base.getindex, dkind == AssignedDiff ? in_u_mm : in_u_unassgn, dslot), Any, line))
+        insert_node_here!(oc_compact,
+            NewInstruction(Expr(:call, Base.setindex!, out_du_mm, v_val, slot), Any, line))
+    end
+
     bc = insert_node_here!(oc_compact,
         NewInstruction(Expr(:call, Base.Broadcast.broadcasted, -, out_du_mm, du_assgn), Any, line))
     insert_node_here!(oc_compact,

src/transform/codegen/rhs.jl

@@ -195,7 +195,7 @@ function rhs_finish!(
                 ir[SSAValue(i)][:inst] = Intrinsics.placeholder_equation
             elseif is_solved_variable(stmt)
                 var = stmt.args[end-1]
-                vint = invview(result.var_to_diff)[var]
+                vint = invview(structure.var_to_diff)[var]
                 if vint !== nothing && key.diff_states !== nothing && (vint in key.diff_states) && !(var in diff_states_in_callee)
                     handle_contribution!(ir, inst, StateDiff, var_assignment[vint][2], arg_range, stmt.args[end])
                 else

src/transform/state_selection.jl

@@ -12,7 +12,7 @@ function StateSelection.linear_subsys_adjmat!(state::TransformationState)
     eadj = Vector{Int}[]
     cadj = Vector{Int}[]
     linear_equations = Vector{Int}()
-    for (i, inc) in enumerate(state.result.total_incidence)
+    for (i, inc) in enumerate(state.total_incidence)
         isa(inc, Const) && continue
         any(x->x === nonlinear || !isinteger(x), nonzeros(inc.row)) && continue
         (isa(inc.typ, Const) && iszero(inc.typ.val)) || continue
@@ -56,13 +56,24 @@ function StateSelection.eq_derivative!(state::TransformationState, eq)
     push!(state.total_incidence, structural_inc_ddt(s.var_to_diff, nothing, nothing, state.total_incidence[eq]))
 end
 
+function StateSelection.var_derivative!(state::TransformationState, var)
+    return StateSelection.var_derivative_graph!(state.structure, var)
+end
+
 function baseeq(result, structure, eq)
     while eq > length(result.eqkinds)
         eq = invview(structure.eq_to_diff)[eq]
     end
     return eq
 end
 
+function basevar(result, structure, var)
+    while var > length(result.varkinds)
+        var = invview(structure.var_to_diff)[var]
+    end
+    return var
+end
+
 function eqkind(result, structure, eq)
     return result.eqkinds[baseeq(result, structure, eq)]
 end
@@ -72,7 +83,7 @@ function eqclassification(result, structure, eq)
     return result.eqclassification[baseeq(result, structure, eq)]
 end
 
-function structural_transformation!(state::TransformationState)
+function ssrm!(state::TransformationState)
     ils = StateSelection.structural_singularity_removal!(state)
 
     s = state.structure
@@ -82,11 +93,35 @@ function structural_transformation!(state::TransformationState)
         end
         state.total_incidence[e] = Incidence(Const(0.), IncidenceVector(MAX_EQS, map(x->x+1, ils.row_cols[ei]), Union{Float64, NonLinear}[Float64(x) for x in ils.row_vals[ei]]))
     end
+end
 
-    var_eq_matching = StateSelection.pantelides!(state;
-        varfilter = var->state.result.varkinds[var] == Intrinsics.Continuous && !(var <= state.result.nexternalvars),
-        eqfilter  = eq->eqkind(state, eq) == Intrinsics.Always)
-    return StateSelection.complete(var_eq_matching, nsrcs(state.structure.graph))
+function varkind(result, structure, var)
+    while var > length(result.varkinds)
+        var = invview(structure.var_to_diff)[var]
+    end
+    return result.varkinds[var]
+end
+varkind(state, var) = varkind(state.result, state.structure, var)
+
+function structural_transformation!(state::TransformationState)
+    first = true
+    # This loop is required to handle situations where additional structural signularities arise in the differentiated
+    # equations. We could probably do lot better here by not constantly recomputing the datastructures.
+    while true
+        neq_before = length(state.structure.eq_to_diff)
+        var_eq_matching = StateSelection.pantelides!(state;
+            varfilter = var->varkind(state, var) == Intrinsics.Continuous && !(var <= state.result.nexternalvars),
+            eqfilter  = eq->eqkind(state, eq) == Intrinsics.Always)
+
+        differentiated_any = neq_before != length(state.structure.eq_to_diff)
+        if differentiated_any || first
+            ssrm!(state)
+            first = false
+            continue
+        end
+        
+        return StateSelection.complete(var_eq_matching, nsrcs(state.structure.graph))
+    end
 end
 
 using StateSelection: Unassigned, SelectedState, unassigned
@@ -99,14 +134,6 @@ function top_level_state_selection!(tstate)
 
     StateSelection.complete!(structure)
 
-    diffvars = result.varkinds .== Intrinsics.Continuous
-    for param in param_vars
-        diffvars[param] = false
-    end
-
-    @assert length(diffvars) == ndsts(structure.graph) == length(structure.var_to_diff)
-    varwhitelist = StateSelection.computed_highest_diff_variables(structure, diffvars)
-
     ## Part 1: Perform the selection of differential states and subsequent tearing of the
     #          non-linear problem at every time step. 
 
@@ -134,16 +161,18 @@ function top_level_state_selection!(tstate)
     diff_key = TornCacheKey(diff_vars, alg_vars, param_vars, explicit_eqs, Vector{Pair{BitSet, BitSet}}())
     @assert matching_for_key(result, diff_key, structure) == var_eq_matching
 
+    varfilter(var) = varkind(result, structure, var) == Intrinsics.Continuous && !(var <= result.nexternalvars)
+
     ## Part 2: Perform the selection of differential states and subsequent tearing of the
     #          non-linear problem at every time step. 
     init_var_eq_matching = StateSelection.complete(StateSelection.maximal_matching(structure.graph;
-        dstfilter = var->diffvars[var], srcfilter = eq->eqkind(result, structure, eq) in (Intrinsics.Always, Intrinsics.Initial)), nsrcs(structure.graph))
+        dstfilter = varfilter, srcfilter = eq->eqkind(result, structure, eq) in (Intrinsics.Always, Intrinsics.Initial)), nsrcs(structure.graph))
     init_var_eq_matching = StateSelection.pss_graph_modia!(structure, init_var_eq_matching)
 
     init_state_vars = BitSet()
     init_explicit_eqs = BitSet()
     for (v, match) in enumerate(init_var_eq_matching)
-        diffvars[v] || continue
+        varfilter(v) || continue
         if match === unassigned
             push!(init_state_vars, v)
         end

src/transform/tearing/schedule.jl

@@ -235,23 +235,27 @@ function compute_eq_schedule(key::TornCacheKey, result, mss::StateSelection.Matc
     frontier = BitSet()
 
     isempty(var_schedule) && (var_schedule = Pair{BitSet, BitSet}[available=>BitSet()])
+    vargraph = DiCMOBiGraph{true}(structure.graph, var_eq_matching)
+
+    function may_enqueue_frontier_var!(var)
+        isa(var_eq_matching[var], Int) || return
+        (var in available) && return # Already processed
+        (var in frontier) && return # Already queued
+        # Check if this neighbor is ready
+        if all(inn->(inn in available), inneighbors(vargraph, var))
+            push!(frontier, var)
+        end
+    end
 
     for sched in var_schedule
         eq_order = Union{Int, SSAValue}[]
         push!(eq_orders, eq_order)
-        vargraph = DiCMOBiGraph{true}(structure.graph, var_eq_matching)
 
         (in_vars, _) = sched
         union!(available, in_vars)
 
         for neighbor in 1:ndsts(structure.graph)
-            isa(var_eq_matching[neighbor], Int) || continue
-            (neighbor in available) && continue # Already processed
-            (neighbor in frontier) && continue # Already queued
-            # Check if this neighbor is ready
-            if all(inn->(inn in available), inneighbors(vargraph, neighbor))
-                push!(frontier, neighbor)
-            end
+            may_enqueue_frontier_var!(neighbor)
         end
 
         new_available = BitSet()
@@ -347,6 +351,11 @@ function compute_eq_schedule(key::TornCacheKey, result, mss::StateSelection.Matc
 
             if !isempty(new_available)
                 union!(available, new_available)
+                for var in new_available
+                    for neighbor in outneighbors(vargraph, var)
+                        may_enqueue_frontier_var!(neighbor)
+                    end
+                end
                 setdiff!(frontier, new_available)
                 empty!(new_available)
                 continue
@@ -443,7 +452,10 @@ function assign_slots(state::TransformationState, key::TornCacheKey, var_eq_matc
         if kind == AlgebraicDerivative
             var_assignment[i] = kind => slot
         elseif kind == AssignedDiff && var_eq_matching !== nothing
-            eq_assignment[var_eq_matching[state.structure.var_to_diff[i]]] = StateDiff => slot
+            eq = var_eq_matching[state.structure.var_to_diff[i]]
+            if isa(eq, Int)
+                eq_assignment[eq] = StateDiff => slot
+            end
         end
     end
 
@@ -475,7 +487,7 @@ function matching_for_key(result::DAEIPOResult, key::TornCacheKey, structure = m
 
     allow_init_eqs = key.diff_states === nothing
 
-    may_use_var(var) = var > result.nexternalvars && (diff_states === nothing || !(var in diff_states)) && !(var in alg_states) && result.varkinds[var] == Intrinsics.Continuous
+    may_use_var(var) = var > result.nexternalvars && (diff_states === nothing || !(var in diff_states)) && !(var in alg_states) && varkind(result, structure, var) == Intrinsics.Continuous
     may_use_eq(eq) = !(eq in explicit_eqs) && eqclassification(result, structure, eq) != External && eqkind(result, structure, eq) in (allow_init_eqs ? (Intrinsics.Initial, Intrinsics.Always) : (Intrinsics.Always,))
 
     # Max match is the (unique) tearing result given the choice of states
@@ -825,7 +837,7 @@ function tearing_schedule!(state::TransformationState, ci::CodeInstance, key::To
         ir_sicm = Compiler.finish(compact)
     end
 
-    var_sols = Vector{Any}(undef, length(result.var_to_diff))
+    var_sols = Vector{Any}(undef, length(structure.var_to_diff))
 
     for var in key.param_vars
         var_sols[var] = 0.0
@@ -847,7 +859,7 @@ function tearing_schedule!(state::TransformationState, ci::CodeInstance, key::To
     end
 
     function insert_solved_var_here!(compact1, var, curval, line)
-        if result.varclassification[var] != Owned
+        if result.varclassification[basevar(result, structure, var)] != Owned
             return
         end
         insert_node_here!(compact1, NewInstruction(Expr(:call, solved_variable, var, curval), Nothing, line))

test/basic.jl

@@ -47,16 +47,16 @@ sol = solve(DAECProblem(pantelides, (1,) .=> 0.), DFBDF(autodiff=false))
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], sol.t))
 
 #= Structural Singularity Removal =#
-function ssm()
+function ssrm()
     a = continuous()
     b = continuous()
     abdot = ddt(a +ᵢ b)
     always!(a -ᵢ abdot)
     always!(b -ᵢ abdot)
 end
 
-ssm()
-sol = solve(DAECProblem(ssm, (1,) .=> 1.), DFBDF(autodiff=false))
+ssrm()
+sol = solve(DAECProblem(ssrm, (1,) .=> 1.), DFBDF(autodiff=false))
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], exp.(0.5sol.t)))
 
 #= Pantelides from init =#

test/runtests.jl

@@ -1,3 +1,3 @@
 include("basic.jl")
 include("ipo.jl")
-include("ssm.jl")
+include("ssrm.jl")

test/ssm.jl renamed to test/ssrm.jl

@@ -10,22 +10,22 @@ using OrdinaryDiffEq
 const +ᵢ = Core.Intrinsics.add_float
 const -ᵢ = Core.Intrinsics.sub_float
 
-function ssm2()
+function ssrm2()
     a = continuous()
     b = continuous()
     abdot = ddt(a +ᵢ b)
     always!(a -ᵢ abdot)
     always!(b +ᵢ abdot)
 end
 
-ssm2()
-
-# TODO: Currently broken
-# solve(DAECProblem(ssm2, (1,) .=> 1.), DFBDF(autodiff=false))
+ssrm2()
+# Doesn't DFBDF doens't like the degenerate case
+@test isempty(DAECompiler.factory(ssrm2)[2])
+@test_broken solve(DAECProblem(ssrm2, ()), DFBDF())
 
 # This is example (7.30) from Taihei Oki "Computing Valuations of Determinants via Combinatorial Optimization: Applications to Differential Equations".
 # The system is index 4 and requires iterating pantelides/ssm.
-function ssm4()
+function ssrm4()
     x₁ = continuous()
     x₂ = continuous()
     x₃ = continuous()
@@ -37,14 +37,24 @@ function ssm4()
     ẍ₁ = ddt(ẋ₁)
     ẍ₂ = ddt(ẋ₂)
     ẍ₃ = ddt(ẋ₃)
-    always!(ẍ₁+ẍ₂-(ẋ₁+ẋ₂)+x₄)
-    always!(ẍ₁+ẍ₂+x₃)
+    always!((ẍ₁+ẍ₂)-(ẋ₁+ẋ₂)+x₄)
+    always!((ẍ₁+ẍ₂)+x₃)
     always!(x₂+ẍ₃+ẋ₄)
     always!(x₃+ẋ₄)
 end
 
-ssm4()
+ssrm4()
+# We expect state selection here to pick (x₁, x₄, ẋ₁)
+# The system simplifies to ẍ₁ = ẋ₄ = ẋ₁ - x₄
+init = (1.,0.,1.)
+function analytic(init, t)
+    c = init[3] - init[2]
+    ẋ₁ = init[3] + c*t
+    x₄ = init[2] + c*t
+    x₁ = init[1] + init[3]*t + 1/2*c*t^2
+    return (x₁, x₄, ẋ₁)
+end
+sol = solve(DAECProblem(ssrm4, (1,2,3) .=> init), DFBDF(autodiff=false))
+@test isapprox(sol[:, :]', mapreduce(t->[analytic(init, t)...]', vcat, sol.t), atol=1e-4)
 
-# TODO: Currently broken
-# solve(DAECProblem(ssm4, (1,) .=> 1.), DFBDF(autodiff=false))
 end