feat: initialization of DiscreteSystem

Author: vyudu

src/structural_transformation/symbolics_tearing.jl

@@ -341,14 +341,18 @@ in order to properly generate the difference equations.
 
 In the system x(k) ~ x(k-1) + x(k-2), becomes Shift(t, 1)(x(t)) ~ x(t) + Shift(t, -1)(x(t))
 
-The lowest-order term is Shift(t, k)(x(t)), instead of x(t).
-As such we actually want dummy variables for the k-1 lowest order terms instead of the k-1 highest order terms.
+The lowest-order term is Shift(t, k)(x(t)), instead of x(t). As such we actually want 
+dummy variables for the k-1 lowest order terms instead of the k-1 highest order terms.
 
 Shift(t, -1)(x(t)) -> x\_{t-1}(t)
 
-Since Shift(t, -1)(x) is not a derivative, it is directly substituted in `fullvars`. No equation or variable is added for it. 
+Since Shift(t, -1)(x) is not a derivative, it is directly substituted in `fullvars`. 
+No equation or variable is added for it. 
 
-For ODESystems D(D(D(x))) in equations is recursively substituted as D(x) ~ x_t, D(x_t) ~ x_tt, etc. The analogue for discrete systems, Shift(t, 1)(Shift(t,1)(Shift(t,1)(Shift(t, -3)(x(t))))) does not actually appear. So `total_sub` in generate_system_equations` is directly initialized with all of the lowered variables `Shift(t, -3)(x) -> x_t-3(t)`, etc. 
+For ODESystems D(D(D(x))) in equations is recursively substituted as D(x) ~ x_t, D(x_t) ~ x_tt, etc. 
+The analogue for discrete systems, Shift(t, 1)(Shift(t,1)(Shift(t,1)(Shift(t, -3)(x(t))))) 
+does not actually appear. So `total_sub` in generate_system_equations` is directly 
+initialized with all of the lowered variables `Shift(t, -3)(x) -> x_t-3(t)`, etc. 
 =#
 """
 Generate new derivative variables for the system.
@@ -358,7 +362,7 @@ Effects on the system structure:
 - neweqs: add the identity equations for the new variables, D(x) ~ x_t
 - graph: update graph with the new equations and variables, and their connections
 - solvable_graph:
-- var_eq_matching: match D(x) to the added identity equation
+- var_eq_matching: match D(x) to the added identity equation D(x) ~ x_t
 """
 function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching; mm = nothing, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = ts
@@ -406,7 +410,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
 end
 
 """
-Check if there's `D(x) = x_t` already.
+Check if there's `D(x) ~ x_t` already.
 """
 function find_duplicate_dd(dv, solvable_graph, diff_to_var, linear_eqs, mm)
     for eq in 𝑑neighbors(solvable_graph, dv)
@@ -427,6 +431,10 @@ function find_duplicate_dd(dv, solvable_graph, diff_to_var, linear_eqs, mm)
     return nothing 
 end
 
+"""
+Add a dummy derivative variable x_t corresponding to symbolic variable D(x) 
+which has index dv in `fullvars`. Return the new index of x_t.
+"""
 function add_dd_variable!(s::SystemStructure, fullvars, x_t, dv)
     push!(fullvars, simplify_shifts(x_t))
     v_t = length(fullvars)
@@ -439,7 +447,11 @@ function add_dd_variable!(s::SystemStructure, fullvars, x_t, dv)
     v_t
 end
 
-# dv = index of D(x), v_t = index of x_t
+"""
+Add the equation D(x) - x_t ~ 0 to `neweqs`. `dv` and `v_t` are the indices
+of the higher-order derivative variable and the newly-introduced dummy
+derivative variable. Return the index of the new equation in `neweqs`.
+"""
 function add_dd_equation!(s::SystemStructure, neweqs, eq, dv, v_t)
     push!(neweqs, eq)
     add_vertex!(s.graph, SRC)
@@ -452,8 +464,33 @@ function add_dd_equation!(s::SystemStructure, neweqs, eq, dv, v_t)
 end
 
 """
-Solve the solvable equations of the system and generate differential (or discrete)
-equations in terms of the selected states.
+Solve the equations in `neweqs` to obtain the final equations of the 
+system.
+
+For each equation of `neweqs`, do one of the following: 
+   1. If the equation is solvable for a differentiated variable D(x),
+      then solve for D(x), and add D(x) ~ sol as a differential equation
+      of the system.
+   2. If the equation is solvable for an un-differentiated variable x, 
+      solve for x and then add x ~ sol as a solved equation. These will
+      become observables.
+   3. If the equation is not solvable, add it as an algebraic equation.
+
+Solved equations are added to `total_sub`. Occurrences of differential
+or solved variables on the RHS of the final equations will get substituted.
+The topological sort of the equations ensures that variables are solved for
+before they appear in equations. 
+
+Reorder the equations and unknowns to be:
+   [diffeqs; ...]
+   [diffvars; ...]
+such that the mass matrix is:
+   [I  0
+    0  0].
+
+Order the new equations and variables such that the differential equations
+and variables come first. Return the new equations, the solved equations,
+the new orderings, and the number of solved variables and equations.
 """
 function generate_system_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = state 
@@ -550,6 +587,9 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
     return neweqs, solved_eqs, eq_ordering, var_ordering, length(solved_vars), length(solved_vars_set)
 end
 
+"""
+Occurs when a variable D(x) occurs in a non-differential system.
+"""
 struct UnexpectedDifferentialError
     eq::Equation
 end
@@ -558,12 +598,20 @@ function Base.showerror(io::IO, err::UnexpectedDifferentialError)
     error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(err.eq)")
 end
 
+"""
+Generate a first-order differential equation whose LHS is `dx`.
+
+`var` and `dx` represent the same variable, but `var` may be a higher-order differential and `dx` is always first-order. For example, if `var` is D(D(x)), then `dx` would be `D(x_t)`. Solve `eq` for `var`, substitute previously solved variables, and return the differential equation.
+"""
 function make_differential_equation(var, dx, eq, total_sub)
     dx ~ simplify_shifts(Symbolics.fixpoint_sub(
         Symbolics.symbolic_linear_solve(eq, var),
         total_sub; operator = ModelingToolkit.Shift))
 end
 
+"""
+Generate an algebraic equation. Substitute solved variables into `eq` and return the equation.
+"""
 function make_algebraic_equation(eq, total_sub)
     rhs = eq.rhs
     if !(eq.lhs isa Number && eq.lhs == 0)
@@ -572,6 +620,9 @@ function make_algebraic_equation(eq, total_sub)
     0 ~ simplify_shifts(Symbolics.fixpoint_sub(rhs, total_sub))
 end
 
+"""
+Solve equation `eq` for `var`, substitute previously solved variables, and return the solved equation.
+"""
 function make_solved_equation(var, eq, total_sub; simplify = false)
     residual = eq.lhs - eq.rhs
     a, b, islinear = linear_expansion(residual, var)
@@ -591,17 +642,13 @@ function make_solved_equation(var, eq, total_sub; simplify = false)
 end
 
 """
-Reorder the equations and unknowns to be:
-   [diffeqs; ...]
-   [diffvars; ...]
-such that the mass matrix is:
-   [I  0
-    0  0].
-
-Update the state to account for the new ordering and equations.
+Given the ordering returned by `generate_system_equations!`, update the 
+tearing state to account for the new order. Permute the variables and equations.
+Eliminate the solved variables and equations from the graph and permute the
+graph's vertices to account for the new variable/equation ordering.
 """
 # TODO: BLT sorting
-function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
+function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_ordering, nsolved_eq, nsolved_var)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
 
     eqsperm = zeros(Int, nsrcs(graph))
@@ -616,7 +663,7 @@ function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_or
     # Contract the vertices in the structure graph to make the structure match
     # the new reality of the system we've just created.
     new_graph = contract_variables(graph, var_eq_matching, varsperm, eqsperm,
-        nelim_eq, nelim_var)
+        nsolved_eq, nsolved_var)
 
     new_var_to_diff = complete(DiffGraph(length(var_ordering)))
     for (v, d) in enumerate(var_to_diff)
@@ -643,7 +690,7 @@ function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_or
 end
 
 """
-Set the system equations, unknowns, observables post-tearing.
+Update the system equations, unknowns, and observables after simplification.
 """
 function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns; 
         cse_hack = true, array_hack = true)
@@ -685,16 +732,10 @@ function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dumm
     sys = schedule(sys)
 end
 
-# Terminology and Definition:
-# A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
-# characterize variables in `u(t)` into two classes: differential variables
-# (denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
-# variables are marked as `SelectedState` and they are differentiated in the
-# DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
-# appear in the system. Algebraic variables are variables that are not
-# differential variables.
 
-# Give the order of the variable indexed by dv
+"""
+Give the order of the variable indexed by dv.
+"""
 function var_order(dv, diff_to_var)
     order = 0
     while (dv′ = diff_to_var[dv]) !== nothing
@@ -704,6 +745,21 @@ function var_order(dv, diff_to_var)
     order, dv
 end
 
+"""
+Main internal function for structural simplification for DAE systems and discrete systems.
+Generate dummy derivative variables, new equations in terms of variables, return updated
+system and tearing state.
+
+Terminology and Definition:
+
+A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
+characterize variables in `u(t)` into two classes: differential variables
+(denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
+variables are marked as `SelectedState` and they are differentiated in the
+DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
+appear in the system. Algebraic variables are variables that are not
+differential variables.
+"""
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
     extra_vars = Int[]

src/structural_transformation/utils.jl

@@ -449,7 +449,12 @@ end
 ### Misc
 ###
 
-# For discrete variables. Turn Shift(t, k)(x(t)) into xₜ₋ₖ(t)
+"""
+Handle renaming variable names for discrete structural simplification. Three cases: 
+- positive shift: do nothing
+- zero shift: x(t) => Shift(t, 0)(x(t))
+- negative shift: rename the variable
+"""
 function lower_shift_varname(var, iv)
     op = operation(var)
     op isa Shift || return Shift(iv, 0)(var, true) # hack to prevent simplification of x(t) - x(t)
@@ -460,30 +465,46 @@ function lower_shift_varname(var, iv)
     end
 end
 
-function shift2term(var) 
+"""
+Rename a Shift variable with negative shift, Shift(t, k)(x(t)) to xₜ₋ₖ(t).
+"""
+function shift2term(var)
     backshift = operation(var).steps
     iv = operation(var).t
-    num = join(Char(0x2080 + d) for d in reverse!(digits(-backshift)))
-    ds = join([Char(0x209c), Char(0x208b), num])
-    #ds = "$iv-$(-backshift)"
-    #d_separator = 'ˍ'
-
-    if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
-        O = only(arguments(var))
-        oldop = operation(O)
-        newname = Symbol(string(nameof(oldop)), ds)
-    else
-        O = var
-        oldop = operation(var) 
-        varname = split(string(nameof(oldop)), d_separator)[1]
-        newname = Symbol(varname, d_separator, ds)
-    end
+    num = join(Char(0x2080 + d) for d in reverse!(digits(-backshift))) # subscripted number, e.g. ₁
+    ds = join([Char(0x209c), Char(0x208b), num]) 
+    # Char(0x209c) = ₜ
+    # Char(0x208b) = ₋ (subscripted minus)
+
+    O = only(arguments(var))
+    oldop = operation(O)
+    newname = Symbol(string(nameof(oldop)), ds)
+
     newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
     newvar = setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
     newvar = setmetadata(newvar, ModelingToolkit.VariableUnshifted, O)
+    newvar = setmetadata(newvar, ModelingToolkit.VariableShift, backshift)
     return newvar
 end
 
+function term2shift(var)
+    var = Symbolics.unwrap(var)
+    name = Symbolics.getname(var)
+    O = only(arguments(var))
+    oldop = operation(O)
+    iv = only(arguments(x))
+    # Split on ₋
+    if occursin(Char(0x208b), name)
+        substrings = split(name, Char(0x208b))
+        shift = last(split(name, Char(0x208b)))
+        newname = join(substrings[1:end-1])[1:end-1]
+        newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
+        return Shift(iv, -shift)(newvar)
+    else
+        return var
+    end
+end
+
 function isdoubleshift(var)
     return ModelingToolkit.isoperator(var, ModelingToolkit.Shift) &&
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)

src/systems/discrete_system/discrete_system.jl

@@ -270,15 +270,19 @@ function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
         v = u0map[k]
         if !((op = operation(k)) isa Shift)
             error("Initial conditions must be for the past state of the unknowns. Instead of providing the condition for $k, provide the condition for $(Shift(iv, -1)(k)).")
+        elseif op.steps > 0
+            error("Initial conditions must be for the past state of the unknowns. Instead of providing the condition for $k, provide the condition for $(Shift(iv, -1)(only(arguments(k)))).")
         end
+
         updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
     end
     for var in unknowns(sys)
         op = operation(var)
-        haskey(updated, var) && continue
         root = getunshifted(var)
+        shift = getshift(var)
         isnothing(root) && continue
-        haskey(defs, root) || error("Initial condition for $root not provided.")
+        (haskey(updated, Shift(iv, shift)(root)) || haskey(updated, var)) && continue
+        haskey(defs, root) || error("Initial condition for $var not provided.")
         updated[var] = defs[root]
     end
     return updated

src/variables.jl

@@ -6,7 +6,9 @@ struct VariableOutput end
 struct VariableIrreducible end
 struct VariableStatePriority end
 struct VariableMisc end
+# Metadata for renamed shift variables xₜ₋₁
 struct VariableUnshifted end
+struct VariableShift end
 Symbolics.option_to_metadata_type(::Val{:unit}) = VariableUnit
 Symbolics.option_to_metadata_type(::Val{:connect}) = VariableConnectType
 Symbolics.option_to_metadata_type(::Val{:input}) = VariableInput
@@ -15,6 +17,7 @@ Symbolics.option_to_metadata_type(::Val{:irreducible}) = VariableIrreducible
 Symbolics.option_to_metadata_type(::Val{:state_priority}) = VariableStatePriority
 Symbolics.option_to_metadata_type(::Val{:misc}) = VariableMisc
 Symbolics.option_to_metadata_type(::Val{:unshifted}) = VariableUnshifted
+Symbolics.option_to_metadata_type(::Val{:shift}) = VariableShift
 
 """
     dump_variable_metadata(var)
@@ -97,7 +100,7 @@ struct Stream <: AbstractConnectType end   # special stream connector
 
 Get the connect type of x. See also [`hasconnect`](@ref).
 """
-getconnect(x) = getconnect(unwrap(x))
+getconnect(x::Num) = getconnect(unwrap(x))
 getconnect(x::Symbolic) = Symbolics.getmetadata(x, VariableConnectType, nothing)
 """
     hasconnect(x)
@@ -264,7 +267,7 @@ end
 end
 
 struct IsHistory end
-ishistory(x) = ishistory(unwrap(x))
+ishistory(x::Num) = ishistory(unwrap(x))
 ishistory(x::Symbolic) = getmetadata(x, IsHistory, false)
 hist(x, t) = wrap(hist(unwrap(x), t))
 function hist(x::Symbolic, t)
@@ -575,7 +578,7 @@ end
 Fetch any miscellaneous data associated with symbolic variable `x`.
 See also [`hasmisc(x)`](@ref).
 """
-getmisc(x) = getmisc(unwrap(x))
+getmisc(x::Num) = getmisc(unwrap(x))
 getmisc(x::Symbolic) = Symbolics.getmetadata(x, VariableMisc, nothing)
 """
     hasmisc(x)
@@ -594,7 +597,7 @@ setmisc(x, miscdata) = setmetadata(x, VariableMisc, miscdata)
 
 Fetch the unit associated with variable `x`. This function is a metadata getter for an individual variable, while `get_unit` is used for unit inference on more complicated sdymbolic expressions.
 """
-getunit(x) = getunit(unwrap(x))
+getunit(x::Num) = getunit(unwrap(x))
 getunit(x::Symbolic) = Symbolics.getmetadata(x, VariableUnit, nothing)
 """
     hasunit(x)
@@ -603,5 +606,8 @@ Check if the variable `x` has a unit.
 """
 hasunit(x) = getunit(x) !== nothing
 
-getunshifted(x) = getunshifted(unwrap(x))
+getunshifted(x::Num) = getunshifted(unwrap(x))
 getunshifted(x::Symbolic) = Symbolics.getmetadata(x, VariableUnshifted, nothing)
+
+getshift(x::Num) = getshift(unwrap(x))
+getshift(x::Symbolic) = Symbolics.getmetadata(x, VariableShift, 0)