feat: implement

Author: vyudu

src/structural_transformation/utils.jl

@@ -507,3 +507,38 @@ function simplify_shifts(var)
             unwrap(var).metadata)
     end
 end
+
+function distribute_shift(var) 
+    var = unwrap(var)
+    var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
+
+    ModelingToolkit.hasshift(var) || return var
+    shift = operation(var)
+    shift isa Shift || return var
+
+    shift = operation(var)
+    expr = only(arguments(var))
+    if expr isa Equation
+        return distribute_shift(shift(expr.lhs)) ~ distribute_shift(shift(expr.rhs))
+    end
+    shiftexpr = _distribute_shift(expr, shift)
+    return simplify_shifts(shiftexpr)
+end
+
+function _distribute_shift(expr, shift)
+    if iscall(expr)
+        op = operation(expr)
+        args = arguments(expr)
+
+        if ModelingToolkit.isvariable(expr)
+            (length(args) == 1 && isequal(shift.t, only(args))) ? (return shift(expr)) : (return expr)
+        elseif op isa Shift
+            return shift(expr)
+        else
+            return maketerm(typeof(expr), operation(expr), Base.Fix2(_distribute_shift, shift).(args),
+                unwrap(expr).metadata)
+        end
+    else
+        return expr
+    end
+end

src/systems/discrete_system/implicit_discrete_system.jl

@@ -265,7 +265,7 @@ function generate_function(
         sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
     iv = get_iv(sys)
     exprs = map(equations(sys)) do eq
-        _iszero(eq.lhs) ? eq.rhs : (simplify_shifts(Shift(iv, -1)(eq.rhs)) - simplify_shifts(Shift(iv, -1)(eq.lhs)))
+        _iszero(eq.lhs) ? eq.rhs : (distribute_shift(Shift(iv, -1)(eq.rhs)) - distribute_shift(Shift(iv, -1)(eq.lhs)))
     end
 
     u_next = dvs

src/systems/systemstructure.jl

@@ -440,7 +440,7 @@ function TearingState(sys; quick_cancel = false, check = true)
         SystemStructure(complete(var_to_diff), complete(eq_to_diff),
             complete(graph), nothing, var_types, sys isa AbstractDiscreteSystem),
         Any[])
-    if sys isa DiscreteSystem
+    if sys isa AbstractDiscreteSystem
         ts = shift_discrete_system(ts)
     end
     return ts

test/implicit_discrete_system.jl

@@ -18,7 +18,6 @@ resid = rand(2)
 f(resid, u_next, [2.,3.], [], t)
 @test resid ≈ [3., -3.]
 
-# Initialization cases.
 prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
 @test prob.u0 == [3., 1.]
 prob = ImplicitDiscreteProblem(sys, [], tspan)
@@ -28,3 +27,17 @@ prob = ImplicitDiscreteProblem(sys, [], tspan)
 @test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
 
 # Test solvers
+@testset "System with algebraic equations" begin
+    @variables x(t) y(t)
+    eqs = [x(k) ~ x(k-1) + x(k-2), 
+           x^2 ~ 1 - y^2]
+    @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
+end
+
+@testset "System with algebraic equations, implicit difference equations, explicit difference equations" begin
+    @variables x(t) y(t)
+    eqs = [x(k) ~ x(k-1) + x(k-2),
+           y(k) ~ x(k) + x(k-2)*y(k-1)]
+    @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
+end
+

test/structural_transformation/utils.jl

@@ -162,3 +162,23 @@ end
     structural_simplify(sys; additional_passes = [pass])
     @test value[] == 1
 end
+
+@testset "Distribute shifts" begin
+    @variables x(t) y(t) z(t)
+    @parameters a b c
+
+    # Expand shifts
+    @test isequal(ST.distribute_shift(Shift(t, -1)(x + y)), Shift(t, -1)(x) + Shift(t, -1)(y))
+
+    expr = a * Shift(t, -2)(x) + Shift(t, 2)(y) + b
+    @test isequal(ST.simplify_shifts(ST.distribute_shift(Shift(t, 2)(expr))),
+                  a*x + Shift(t, 4)(y) + b)
+    @test isequal(ST.distribute_shift(Shift(t, 2)(exp(z))), exp(Shift(t, 2)(z)))
+    @test isequal(ST.distribute_shift(Shift(t, 2)(exp(a) + b)), exp(a) + b)
+
+    expr = a^x - log(b*y) + z*x
+    @test isequal(ST.distribute_shift(Shift(t, -3)(expr)), a^(Shift(t, -3)(x)) - log(b * Shift(t, -3)(y)) + Shift(t, -3)(z)*Shift(t, -3)(x))
+
+    expr = x(k+1) ~ x + x(k-1)
+    @test isequal(ST.distribute_shift(Shift(t, -1)(expr)), x ~ x(k-1) + x(k-2))
+end