rename functions

Author: vyudu

src/structural_transformation/symbolics_tearing.jl

@@ -254,7 +254,9 @@ function substitute_dummy_derivatives!(ts::TearingState, neweqs, dummy_sub, var_
     diff_to_var = invview(var_to_diff)
     iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
 
+    @show neweqs
     for var in 1:length(fullvars)
+        #@show neweqs
         dv = var_to_diff[var]
         dv === nothing && continue
         if var_eq_matching[var] !== SelectedState()
@@ -286,9 +288,7 @@ function substitute_dummy_derivatives!(ts::TearingState, neweqs, dummy_sub, var_
     end
 end
 
-"""
-Generate new derivative variables for the system.
-
+#= 
 There are three cases where we want to generate new variables to convert
 the system into first order (semi-implicit) ODEs.
     
@@ -361,6 +361,16 @@ This is different from the continuous case where D(D(x)) can be substituted for
 by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
 shift_sub dict is updated at the time that the renamed variables are written,
 inside the loop where new variables are generated.
+=#
+"""
+Generate new derivative variables for the system.
+
+Effects on the state: 
+- fullvars: add the new derivative variables x_t
+- 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: solvable equations
 """
 function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching;
         is_discrete = false, mm = nothing, shift_sub = nothing)
@@ -406,7 +416,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
             dx = fullvars[dv]
             order, lv = var_order(diff_to_var, dv)
             @show fullvars[uv]
-            x_t = lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv])
+            x_t = lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv], allow_zero = true)
             shift_sub[dx] = x_t
             (var_eq_matching[dv] isa Int) ? continue : @goto DISCRETE_VARIABLE
         end
@@ -439,7 +449,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         x_t = lower_name(fullvars[lv], iv, order)
 
         @label DISCRETE_VARIABLE
-        push!(fullvars, x_t)
+        push!(fullvars, simplify_shifts(x_t))
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
         add_vertex!(graph, DST)
@@ -448,7 +458,6 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         add_vertex!(solvable_graph, DST)
         push!(var_eq_matching, unassigned)
 
-        # Add discrete substitutions to total_sub directly.  
         is_discrete && begin
             idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
             for e in 𝑑neighbors(graph, dv)
@@ -463,7 +472,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         end
 
         # add `D(x) - x_t ~ 0`
-        push!(neweqs, 0 ~ x_t - dx)
+        push!(neweqs, 0 ~ dx - x_t)
         add_vertex!(graph, SRC)
         dummy_eq = length(neweqs)
         add_edge!(graph, dummy_eq, dv)
@@ -478,12 +487,11 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
     end
 end
 
-function add_solvable_variable!() 
+function add_solvable_variable!(state::TearingState)
 
 end
 
-function add_solvable_equation!() 
-    
+function add_solvable_equation!(s::SystemStructure, neweqs, eq)
 end
 
 """
@@ -500,7 +508,8 @@ such that the mass matrix is:
 Update the state to account for the new ordering and equations.
 
 ####### DISCRETE CASE
-- only substitute Shift(t, -2) 
+- Differential equations: substitute variables with everything shifted forward one timestep.
+- Algebraic and observable equations: substitute variables with everything shifted back one timestep.
 """
 function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false, shift_sub = Dict())
     @unpack fullvars, sys, structure = state 
@@ -562,8 +571,6 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
                 isnothing(D) &&
                     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: $(equations(sys)[ieq])")
                 order, lv = var_order(diff_to_var, iv)
-                @show fullvars[lv]
-                @show simplify_shifts(fullvars[lv])
                 dx = D(simplify_shifts(fullvars[lv]))
                 eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
                     Symbolics.symbolic_linear_solve(neweqs[ieq],
@@ -576,8 +583,6 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
                 push!(diff_eqs, eq)
                 total_sub[simplify_shifts(eq.lhs)] = eq.rhs
                 dx_sub[simplify_shifts(eq.lhs)] = eq.rhs
-                @show total_sub
-                @show eq
                 push!(diffeq_idxs, ieq)
                 push!(diff_vars, diff_to_var[iv])
                 continue
@@ -611,6 +616,8 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
             push!(algeeq_idxs, ieq)
         end
     end
+    @show neweqs
+    @show subeqs
 
     # TODO: BLT sorting
     neweqs = [diff_eqs; alge_eqs]

src/structural_transformation/utils.jl

@@ -271,13 +271,16 @@ end
 
 function find_solvables!(state::TearingState; kwargs...)
     @assert state.structure.solvable_graph === nothing
+    println("in find_solvables")
+    @show eqs
     eqs = equations(state)
     graph = state.structure.graph
     state.structure.solvable_graph = BipartiteGraph(nsrcs(graph), ndsts(graph))
     to_rm = Int[]
     for ieq in 1:length(eqs)
         find_eq_solvables!(state, ieq, to_rm; kwargs...)
     end
+    @show eqs
     return nothing
 end
 
@@ -477,7 +480,7 @@ end
 ### Rules
 # 1. x(t) -> x(t)
 # 2. Shift(t, 0)(x(t)) -> x(t)
-# 3. Shift(t, 1)(x + z) -> Shift(t, 1)(x) + Shift(t, 1)(z)
+# 3. Shift(t, 3)(Shift(t, 2)(x(t)) -> Shift(t, 5)(x(t))
 
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
@@ -498,3 +501,58 @@ function simplify_shifts(var)
             unwrap(var).metadata)
     end
 end
+
+"""
+Power expand the shifts. Used for substitution.
+
+Shift(t, -3)(x(t)) -> Shift(t, -1)(Shift(t, -1)(Shift(t, -1)(x)))
+"""
+function expand_shifts(var)
+    ModelingToolkit.hasshift(var) || return var
+    var = ModelingToolkit.value(var)
+
+    var isa Equation && return expand_shifts(var.lhs) ~ expand_shifts(var.rhs)
+    op = operation(var)
+    s = sign(op.steps)
+    arg = only(arguments(var))
+
+    if ModelingToolkit.isvariable(arg) && (ModelingToolkit.getvariabletype(arg) === VARIABLE) && isequal(op.t, only(arguments(arg)))
+        out = arg
+        for i in 1:op.steps
+            out = Shift(op.t, s)(out)
+        end
+        return out
+    elseif iscall(arg)
+        return maketerm(typeof(var), operation(var), expand_shifts.(arguments(var)),
+            unwrap(var).metadata)
+    else
+        return arg
+    end
+end
+
+"""
+Shift(t, 1)(x + z) -> Shift(t, 1)(x) + Shift(t, 1)(z)
+"""
+function distribute_shift(var) 
+    ModelingToolkit.hasshift(var) || return var
+    var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
+    shift = operation(var)
+    expr = only(arguments(var))
+    _distribute_shift(expr, shift)
+end
+
+function _distribute_shift(expr, shift)
+    op = operation(expr)
+    args = arguments(expr)
+
+    if length(args) == 1 
+        if ModelingToolkit.isvariable(only(args)) && isequal(op.t, only(args))
+            return shift(only(args))
+        else
+            return only(args)
+        end
+    else iscall(op)
+        return maketerm(typeof(expr), operation(expr), _distribute_shift.(args, shift),
+            unwrap(var).metadata)
+    end
+end

test/structural_transformation/utils.jl

@@ -4,6 +4,7 @@ using Graphs
 using SparseArrays
 using UnPack
 using ModelingToolkit: t_nounits as t, D_nounits as D
+const ST = StructuralTransformations
 
 # Define some variables
 @parameters L g
@@ -161,3 +162,19 @@ end
     structural_simplify(sys; additional_passes = [pass])
     @test value[] == 1
 end
+
+@testset "Shift simplification" begin
+    @variables x(t) y(t) z(t)
+    @parameters a b c
+    
+    # Expand shifts
+    @test isequal(ST.expand_shifts(Shift(t, -3)(x)), Shift(t, -1)(Shift(t, -1)(Shift(t, -1)(x))))
+    expr = a * Shift(t, -2)(x) + Shift(t, 2)(y) + b
+    @test isequal(ST.expand_shifts(expr), 
+                    a * Shift(t, -1)(Shift(t, -1)(x)) + Shift(t, 1)(Shift(t, 1)(y)) + b)
+    @test isequal(ST.expand_shifts(Shift(t, 2)(Shift(t, 1)(a))), a)
+
+
+    # Distribute shifts
+
+end