feat: support negative shifts in structural_simplify

Author: AayushSabharwal

src/structural_transformation/symbolics_tearing.jl

@@ -382,8 +382,8 @@ function tearing_reassemble(state::TearingState, var_eq_matching;
         dx = fullvars[dv]
         # add `x_t`
         order, lv = var_order(dv)
-        x_t = lower_varname(fullvars[lv], iv, order)
-        push!(fullvars, x_t)
+        x_t = lower_varname_withshift(fullvars[lv], iv, order)
+        push!(fullvars, simplify_shifts(x_t))
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
         add_vertex!(graph, DST)
@@ -437,11 +437,12 @@ function tearing_reassemble(state::TearingState, var_eq_matching;
             # We cannot solve the differential variable like D(x)
             if isdervar(iv)
                 order, lv = var_order(iv)
-                dx = D(lower_varname(fullvars[lv], idep, order - 1))
-                eq = dx ~ ModelingToolkit.fixpoint_sub(
+                dx = D(simplify_shifts(lower_varname_withshift(
+                    fullvars[lv], idep, order - 1)))
+                eq = dx ~ simplify_shifts(ModelingToolkit.fixpoint_sub(
                     Symbolics.solve_for(neweqs[ieq],
                         fullvars[iv]),
-                    total_sub)
+                    total_sub; operator = ModelingToolkit.Shift))
                 for e in 𝑑neighbors(graph, iv)
                     e == ieq && continue
                     for v in 𝑠neighbors(graph, e)
@@ -450,7 +451,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching;
                     rem_edge!(graph, e, iv)
                 end
                 push!(diff_eqs, eq)
-                total_sub[eq.lhs] = eq.rhs
+                total_sub[simplify_shifts(eq.lhs)] = eq.rhs
                 push!(diffeq_idxs, ieq)
                 push!(diff_vars, diff_to_var[iv])
                 continue
@@ -469,7 +470,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching;
                 neweq = var ~ ModelingToolkit.fixpoint_sub(
                     simplify ?
                     Symbolics.simplify(rhs) : rhs,
-                    total_sub)
+                    total_sub; operator = ModelingToolkit.Shift)
                 push!(subeqs, neweq)
                 push!(solved_equations, ieq)
                 push!(solved_variables, iv)

src/structural_transformation/utils.jl

@@ -412,3 +412,58 @@ function numerical_nlsolve(f, u0, p)
     # TODO: robust initial guess, better debugging info, and residual check
     sol.u
 end
+
+###
+### Misc
+###
+
+function lower_varname_withshift(var, iv, order)
+    order == 0 && return var
+    if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
+        op = operation(var)
+        return Shift(op.t, order)(var)
+    end
+    return lower_varname(var, iv, order)
+end
+
+function isdoubleshift(var)
+    return ModelingToolkit.isoperator(var, ModelingToolkit.Shift) &&
+           ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
+end
+
+function simplify_shifts(var)
+    ModelingToolkit.hasshift(var) || return var
+    r = @rule ~x::isdoubleshift => begin
+        op1 = operation(~x)
+        vv1 = arguments(~x)[1]
+        op2 = operation(vv1)
+        vv2 = arguments(vv1)[1]
+        s1 = op1.steps
+        s2 = op2.steps
+        t1 = op1.t
+        t2 = op2.t
+        if t1 === nothing
+            ModelingToolkit.Shift(t2, s1 + s2)(vv2)
+        else
+            ModelingToolkit.Shift(t1, s1 + s2)(vv2)
+        end
+    end
+    return Postwalk(PassThrough(r))(var)
+    while ModelingToolkit.isoperator(var, ModelingToolkit.Shift) &&
+        ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
+        op1 = operation(var)
+        vv1 = arguments(var)[1]
+        op2 = operation(vv1)
+        vv2 = arguments(vv1)[1]
+        s1 = op1.steps
+        s2 = op2.steps
+        t1 = op1.t
+        t2 = op2.t
+        if t1 === nothing
+            var = Shift(t2, s1 + s2)(vv2)
+        else
+            var = Shift(t1, s1 + s2)(vv2)
+        end
+    end
+    return var
+end

src/systems/alias_elimination.jl

@@ -463,11 +463,11 @@ function observed2graph(eqs, unknowns)
     return graph, assigns
 end
 
-function fixpoint_sub(x, dict)
-    y = fast_substitute(x, dict)
+function fixpoint_sub(x, dict; operator = Nothing)
+    y = fast_substitute(x, dict; operator)
     while !isequal(x, y)
         y = x
-        x = fast_substitute(y, dict)
+        x = fast_substitute(y, dict; operator)
     end
 
     return x

src/systems/discrete_system/discrete_system.jl

@@ -179,8 +179,6 @@ function DiscreteSystem(eqs, iv; kwargs...)
             eq.lhs in diffvars &&
                 throw(ArgumentError("The shift variable $(eq.lhs) is not unique in the system of equations."))
             push!(diffvars, eq.lhs)
-        else
-            throw(ArgumentError("All equations in a `DiscreteSystem` must be difference equations with positive shifts"))
         end
     end
     new_ps = OrderedSet()

src/systems/systemstructure.jl

@@ -327,24 +327,6 @@ function TearingState(sys; quick_cancel = false, check = true)
 
             dvar = var
             idx = varidx
-            if ModelingToolkit.isoperator(dvar, ModelingToolkit.Shift)
-                if !(idx in dervaridxs)
-                    push!(dervaridxs, idx)
-                end
-                op = operation(dvar)
-                tt = op.t
-                steps = op.steps
-                v = arguments(dvar)[1]
-                for s in (steps - 1):-1:1
-                    sf = Shift(tt, s)
-                    dvar = sf(v)
-                    idx = addvar!(dvar)
-                    if !(idx in dervaridxs)
-                        push!(dervaridxs, idx)
-                    end
-                end
-                idx = addvar!(v)
-            end
 
             if istree(var) && operation(var) isa Symbolics.Operator &&
                !isdifferential(var) && (it = input_timedomain(var)) !== nothing
@@ -364,6 +346,47 @@ function TearingState(sys; quick_cancel = false, check = true)
             eqs[i] = eqs[i].lhs ~ rhs
         end
     end
+    lowest_shift = Dict()
+    for var in fullvars
+        if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
+            steps = operation(var).steps
+            v = arguments(var)[1]
+            lowest_shift[v] = min(get(lowest_shift, v, 0), steps)
+        end
+    end
+    for var in fullvars
+        if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
+            op = operation(var)
+            steps = op.steps
+            v = arguments(var)[1]
+            lshift = lowest_shift[v]
+            tt = op.t
+        elseif haskey(lowest_shift, var)
+            lshift = lowest_shift[var]
+            steps = 0
+            tt = iv
+            v = var
+            if lshift < 0
+                defs = ModelingToolkit.get_defaults(sys)
+                if (_val = get(defs, var, nothing)) !== nothing
+                    defs[Shift(tt, lshift)(v)] = _val
+                end
+            end
+        else
+            continue
+        end
+        if lshift < steps
+            push!(dervaridxs, var2idx[var])
+        end
+        for s in (steps - 1):-1:(lshift + 1)
+            sf = Shift(tt, s)
+            dvar = sf(v)
+            idx = addvar!(dvar)
+            if !(idx in dervaridxs)
+                push!(dervaridxs, idx)
+            end
+        end
+    end
 
     # sort `fullvars` such that the mass matrix is as diagonal as possible.
     dervaridxs = collect(dervaridxs)
@@ -418,15 +441,18 @@ end
 function lower_order_var(dervar)
     if isdifferential(dervar)
         diffvar = arguments(dervar)[1]
-    else # shift
+    elseif ModelingToolkit.isoperator(dervar, ModelingToolkit.Shift)
         s = operation(dervar)
         step = s.steps - 1
         vv = arguments(dervar)[1]
-        if step >= 1
+        if step != 0
             diffvar = Shift(s.t, step)(vv)
         else
             diffvar = vv
         end
+    else
+        iv = only(arguments(dervar))
+        return Shift(iv, -1)(dervar)
     end
     diffvar
 end

src/utils.jl

@@ -799,35 +799,51 @@ end
 # Symbolics needs to call unwrap on the substitution rules, but most of the time
 # we don't want to do that in MTK.
 const Eq = Union{Equation, Inequality}
-function fast_substitute(eq::Eq, subs)
+function fast_substitute(eq::Eq, subs; operator = Nothing)
     if eq isa Inequality
-        Inequality(fast_substitute(eq.lhs, subs), fast_substitute(eq.rhs, subs),
+        Inequality(fast_substitute(eq.lhs, subs; operator),
+            fast_substitute(eq.rhs, subs; operator),
             eq.relational_op)
     else
-        Equation(fast_substitute(eq.lhs, subs), fast_substitute(eq.rhs, subs))
+        Equation(fast_substitute(eq.lhs, subs; operator),
+            fast_substitute(eq.rhs, subs; operator))
     end
 end
-function fast_substitute(eq::T, subs::Pair) where {T <: Eq}
-    T(fast_substitute(eq.lhs, subs), fast_substitute(eq.rhs, subs))
+function fast_substitute(eq::T, subs::Pair; operator = Nothing) where {T <: Eq}
+    T(fast_substitute(eq.lhs, subs; operator), fast_substitute(eq.rhs, subs; operator))
 end
-fast_substitute(eqs::AbstractArray, subs) = fast_substitute.(eqs, (subs,))
-fast_substitute(a, b) = substitute(a, b)
-function fast_substitute(expr, pair::Pair)
+function fast_substitute(eqs::AbstractArray, subs; operator = Nothing)
+    fast_substitute.(eqs, (subs,); operator)
+end
+function fast_substitute(a, b; operator = Nothing)
+    b = Dict(value(k) => value(v) for (k, v) in b)
+    a = value(a)
+    haskey(b, a) && return b[a]
+    for _b in b
+        a = fast_substitute(a, _b; operator)
+    end
+    a
+end
+function fast_substitute(expr, pair::Pair; operator = Nothing)
     a, b = pair
+    a = value(a)
+    b = value(b)
     isequal(expr, a) && return b
 
     istree(expr) || return expr
-    op = fast_substitute(operation(expr), pair)
-    canfold = Ref(!(op isa Symbolic))
-    args = let canfold = canfold
-        map(SymbolicUtils.unsorted_arguments(expr)) do x
-            x′ = fast_substitute(x, pair)
-            canfold[] = canfold[] && !(x′ isa Symbolic)
-            x′
+    op = fast_substitute(operation(expr), pair; operator)
+    args = SymbolicUtils.unsorted_arguments(expr)
+    if !(op isa operator)
+        canfold = Ref(!(op isa Symbolic))
+        args = let canfold = canfold
+            map(args) do x
+                x′ = fast_substitute(x, pair; operator)
+                canfold[] = canfold[] && !(x′ isa Symbolic)
+                x′
+            end
         end
+        canfold[] && return op(args...)
     end
-    canfold[] && return op(args...)
-
     similarterm(expr,
         op,
         args,

test/clock.jl

@@ -423,7 +423,7 @@ ci, varmap = infer_clocks(expand_connections(_model))
 @test varmap[_model.feedback.output.u] == d
 @test varmap[_model.feedback.input2.u] == d
 
-@test_skip ssys = structural_simplify(model)
+ssys = structural_simplify(model)
 
 Tf = 0.2
 timevec = 0:(d.dt):Tf
@@ -445,20 +445,20 @@ y = res.y[:]
 # ref = Constant(k = 0.5)
 
 # ; model.controller.x(k-1) => 0.0
+prob = ODEProblem(ssys,
+    [model.plant.x => 0.0; model.controller.kp => 2.0; model.controller.ki => 2.0],
+    (0.0, Tf))
+
+@test prob.ps[Hold(ss.holder.input.u)] == 1 # constant output * kp issue https://github.com/SciML/ModelingToolkit.jl/issues/2356
+@test prob.ps[ss.controller.x(k - 1)] == 0 # c2d
+@test prob.ps[Sample(d)(ss.sampler.input.u)] == 0 # disc state
+sol = solve(prob,
+    Tsit5(),
+    kwargshandle = KeywordArgSilent,
+    abstol = 1e-8,
+    reltol = 1e-8)
 @test_skip begin
-    prob = ODEProblem(ssys,
-        [model.plant.x => 0.0; model.controller.kp => 2.0; model.controller.ki => 2.0],
-        (0.0, Tf))
-
-    @test prob.p[9] == 1 # constant output * kp issue https://github.com/SciML/ModelingToolkit.jl/issues/2356
-    @test prob.p[10] == 0 # c2d
-    @test prob.p[11] == 0 # disc state
-    sol = solve(prob,
-        Tsit5(),
-        kwargshandle = KeywordArgSilent,
-        abstol = 1e-8,
-        reltol = 1e-8)
-    plot([y sol(timevec, idxs = model.plant.output.u)], m = :o, lab = ["CS" "MTK"])
+    # plot([y sol(timevec, idxs = model.plant.output.u)], m = :o, lab = ["CS" "MTK"])
 
     ##