feat: allow simplifying DAEs to index zero, better DAE initialization

src/structural_transformation/pantelides.jl

@@ -71,7 +71,7 @@ function pantelides_reassemble(state::TearingState, var_eq_matching)
 end
 
 """
-    computed_highest_diff_variables(structure)
+    computed_highest_diff_variables(structure; whitelisted_vars = ())
 
 Computes which variables are the "highest-differentiated" for purposes of
 pantelides. Ordinarily this is relatively straightforward. However, in our
@@ -83,12 +83,18 @@ case, there is one complicating condition:
 
 This function takes care of these complications are returns a boolean array
 for every variable, indicating whether it is considered "highest-differentiated".
+
+For each index `i` in `whitelisted_vars`, the `i`th variable is included if it
+is the highest differentiated variable even if it doesn't appear in the system.
 """
-function computed_highest_diff_variables(structure)
+function computed_highest_diff_variables(structure; whitelisted_vars = ())
     @unpack graph, var_to_diff = structure
 
     nvars = length(var_to_diff)
     varwhitelist = falses(nvars)
+    for i in whitelisted_vars
+        varwhitelist[i] = true
+    end
     for var in 1:nvars
         if var_to_diff[var] === nothing && !varwhitelist[var]
             # This variable is structurally highest-differentiated, but may not actually appear in the
@@ -125,7 +131,7 @@ end
 Perform Pantelides algorithm.
 """
 function pantelides!(
-        state::TransformationState; finalize = true, maxiters = 8000, kwargs...)
+        state::TransformationState; finalize = true, maxiters = 8000, whitelisted_vars = (), kwargs...)
     @unpack graph, solvable_graph, var_to_diff, eq_to_diff = state.structure
     neqs = nsrcs(graph)
     nvars = nv(var_to_diff)
@@ -137,8 +143,7 @@ function pantelides!(
         eq -> !isempty(𝑠neighbors(graph, eq)) && eq_to_diff[eq] === nothing,
         1:neqs′)
 
-    varwhitelist = computed_highest_diff_variables(state.structure)
-
+    varwhitelist = computed_highest_diff_variables(state.structure; whitelisted_vars)
     if nnonemptyeqs > count(varwhitelist)
         throw(InvalidSystemException("System is structurally singular"))
     end
@@ -206,14 +211,19 @@ function pantelides!(
 end
 
 """
-    dae_index_lowering(sys::ODESystem; kwargs...) -> ODESystem
+    dae_index_lowering(sys::ODESystem; to_index_zero = false, kwargs...) -> ODESystem
 
 Perform the Pantelides algorithm to transform a higher index DAE to an index 1
 DAE. `kwargs` are forwarded to [`pantelides!`](@ref). End users are encouraged to call [`structural_simplify`](@ref)
-instead, which calls this function internally.
+instead, which calls this function internally. If `to_index_zero` is true, the DAE will be reduced to an index 1 DAE.
 """
-function dae_index_lowering(sys::ODESystem; kwargs...)
+function dae_index_lowering(sys::ODESystem; to_index_zero = false, kwargs...)
     state = TearingState(sys)
-    var_eq_matching = pantelides!(state; finalize = false, kwargs...)
+    if to_index_zero
+        newvars = ModelingToolkit.add_missing_differentials!(state)
+    else
+        newvars = ()
+    end
+    var_eq_matching = pantelides!(state; finalize = false, whitelisted_vars = newvars, kwargs...)
     return invalidate_cache!(pantelides_reassemble(state, var_eq_matching))
 end

src/structural_transformation/symbolics_tearing.jl

@@ -812,7 +812,13 @@ 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,
-        mm = nothing, cse_hack = true, array_hack = true, kwargs...)
+        mm = nothing, cse_hack = true, array_hack = true, to_index_zero = false, kwargs...)
+    if to_index_zero
+        newvars = ModelingToolkit.add_missing_differentials!(state)
+    else
+        newvars = ()
+    end
+
     jac = let state = state
         (eqs, vars) -> begin
             symeqs = EquationsView(state)[eqs]
@@ -834,7 +840,7 @@ function dummy_derivative(sys, state = TearingState(sys); simplify = false,
             p
         end
     end
-    var_eq_matching = dummy_derivative_graph!(state, jac; state_priority,
+    var_eq_matching = dummy_derivative_graph!(state, jac; state_priority, whitelisted_vars = newvars,
         kwargs...)
     tearing_reassemble(state, var_eq_matching; simplify, mm, cse_hack, array_hack)
 end

src/systems/nonlinear/initializesystem.jl

@@ -11,9 +11,35 @@ function generate_initializesystem(sys::ODESystem;
         default_dd_guess = 0.0,
         algebraic_only = false,
         check_units = true, check_defguess = false,
+        implicit_dae = false,
         name = nameof(sys), kwargs...)
     trueobs, eqs = unhack_observed(observed(sys), equations(sys))
     vars = unique([unknowns(sys); getfield.(trueobs, :lhs)])
+
+    if implicit_dae
+        pre_simplification_sys = sys
+        while get_parent(pre_simplification_sys) !== nothing
+            pre_simplification_sys = get_parent(pre_simplification_sys)
+        end
+        schedule = get_schedule(sys)
+        if schedule === nothing
+            throw(ArgumentError("The system must be structurally simplified to create an initialization system for an implicit DAE."))
+        end
+        old_eqs = equations(pre_simplification_sys)
+        inv_dummy_sub = Dict()
+        for (k, v) in schedule.dummy_sub
+            if isequal(default_toterm(k), v)
+                inv_dummy_sub[v] = k
+            end
+        end
+        new_eqs = Symbolics.fast_substitute.([trueobs; eqs], (inv_dummy_sub,))
+        filter!(eq -> !isequal(eq.lhs, eq.rhs), new_eqs)
+        new_sys = ODESystem(new_eqs, get_iv(sys); name = nameof(sys))
+        new_sys = dummy_derivative(new_sys; to_index_zero = true, array_hack = false, cse_hack = false)
+        trueobs = observed(new_sys)
+        eqs = equations(new_sys)
+        vars = unique([unknowns(new_sys); getfield.(trueobs, :lhs)])
+    end
     vars_set = Set(vars) # for efficient in-lookup
 
     idxs_diff = isdiffeq.(eqs)

src/systems/systemstructure.jl

@@ -480,6 +480,37 @@ function shift_discrete_system(ts::TearingState)
     return ts
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+For each variable in `ts.fullvars` which does not have a derivative in `ts.fullvars`
+and is not the derivative of a variable in `ts.fullvars`, add its derivative to `ts`.
+Returns the indexes of added differential variables.
+"""
+function add_missing_differentials!(ts::TearingState)
+    sys = ts.sys
+    D = Differential(get_iv(sys))
+    newvars = Int[]
+    for (i, v) in enumerate(ts.fullvars)
+        # ignore variables that have a derivative...
+        ts.structure.var_to_diff[i] === nothing || continue
+        # or are the derivative
+        invview(ts.structure.var_to_diff)[i] === nothing || continue
+        # add to fullvars
+        push!(ts.fullvars, D(v))
+        push!(newvars, length(ts.fullvars))
+        # update diffgraph
+        add_vertex!(ts.structure.var_to_diff)
+        add_edge!(ts.structure.var_to_diff, i, length(ts.fullvars))
+        # update bipartite graphs
+        add_vertex!(ts.structure.graph, DST)
+        if ts.structure.solvable_graph !== nothing
+            add_vertex!(ts.structure.solvable_graph, DST)
+        end
+    end
+    return newvars
+end
+
 using .BipartiteGraphs: Label, BipartiteAdjacencyList
 struct SystemStructurePrintMatrix <:
        AbstractMatrix{Union{Label, BipartiteAdjacencyList}}
@@ -676,6 +707,7 @@ end
 function _structural_simplify!(state::TearingState, io; simplify = false,
         check_consistency = true, fully_determined = true, warn_initialize_determined = false,
         dummy_derivative = true,
+        to_index_zero = false,
         kwargs...)
     if fully_determined isa Bool
         check_consistency &= fully_determined
@@ -699,9 +731,14 @@ function _structural_simplify!(state::TearingState, io; simplify = false,
     end
     if fully_determined && dummy_derivative
         sys = ModelingToolkit.dummy_derivative(
-            sys, state; simplify, mm, check_consistency, kwargs...)
+            sys, state; simplify, mm, check_consistency, to_index_zero, kwargs...)
     elseif fully_determined
-        var_eq_matching = pantelides!(state; finalize = false, kwargs...)
+        if to_index_zero
+            newvars = add_missing_differentials!(state)
+        else
+            newvars = ()
+        end
+        var_eq_matching = pantelides!(state; finalize = false, whitelisted_vars = newvars, kwargs...)
         sys = pantelides_reassemble(state, var_eq_matching)
         state = TearingState(sys)
         sys, mm = ModelingToolkit.alias_elimination!(state; kwargs...)