add shift2term for positive shifts

Author: vyudu

src/discretedomain.jl

@@ -70,6 +70,19 @@ Base.literal_pow(f::typeof(^), D::Shift, ::Val{n}) where {n} = Shift(D.t, D.step
 
 hasshift(eq::Equation) = hasshift(eq.lhs) || hasshift(eq.rhs)
 
+"""
+    Next(x)
+
+An alias for Shift(t, 1)(x).
+"""
+Next(x) = Shift(t, 1)(x)
+"""
+    Prev(x)
+
+An alias for Shift(t, -1)(x).
+"""
+Prev(x) = Shift(t, -1)(x)
+
 """
     hasshift(O)
 

src/structural_transformation/StructuralTransformations.jl

@@ -65,7 +65,7 @@ export torn_system_jacobian_sparsity
 export full_equations
 export but_ordered_incidence, lowest_order_variable_mask, highest_order_variable_mask
 export computed_highest_diff_variables
-export shift2term, lower_shift_varname, simplify_shifts
+export shift2term, lower_shift_varname, simplify_shifts, distribute_shift
 
 include("utils.jl")
 include("pantelides.jl")

src/structural_transformation/utils.jl

@@ -457,7 +457,7 @@ Handle renaming variable names for discrete structural simplification. Three cas
 """
 function lower_shift_varname(var, iv)
     op = operation(var)
-    if op isa Shift && op.steps < 0
+    if op isa Shift
         return shift2term(var)
     else
         return Shift(iv, 0)(var, true)
@@ -475,10 +475,14 @@ function shift2term(var)
 
     backshift = is_lowered ? op.steps + ModelingToolkit.getshift(arg) : op.steps
 
-    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)
+    # Char(0x208a) = ₊ (subscripted plus)
+    pm = backshift > 0 ? Char(0x208a) : Char(0x208b)
+    # subscripted number, e.g. ₁
+    num = join(Char(0x2080 + d) for d in reverse!(digits(abs(backshift))))
+    # Char(0x209c) = ₜ
+    # ds = ₜ₋₁
+    ds = join([Char(0x209c), pm, num])
 
     O = is_lowered ? ModelingToolkit.getunshifted(arg) : arg
     oldop = operation(O)
@@ -498,6 +502,9 @@ function isdoubleshift(var)
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
 end
 
+"""
+Simplify multiple shifts: Shift(t, k1)(Shift(t, k2)(x)) becomes Shift(t, k1+k2)(x).
+"""
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
     var isa Equation && return simplify_shifts(var.lhs) ~ simplify_shifts(var.rhs)
@@ -518,6 +525,11 @@ function simplify_shifts(var)
     end
 end
 
+"""
+Distribute a shift applied to a whole expression or equation. 
+Shift(t, 1)(x + y) will become Shift(t, 1)(x) + Shift(t, 1)(y).
+Only shifts variables whose independent variable is the same t that appears in the Shift (i.e. constants, time-independent parameters, etc. do not get shifted).
+"""
 function distribute_shift(var) 
     var = unwrap(var)
     var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)

test/implicit_discrete_system.jl

@@ -1,37 +1,58 @@
 using ModelingToolkit, Test
 using ModelingToolkit: t_nounits as t
+using StableRNGs
 
 k = ShiftIndex(t)
-@variables x(t) = 1
-@mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
-tspan = (0, 10)
+rng = StableRNG(22525) 
 
 # Shift(t, -1)(x(t)) - x_{t-1}(t)
 # -3 - x(t) + x(t)*x_{t-1}
-f = ImplicitDiscreteFunction(sys)
-u_next = [3., 1.5]
-@test f(u_next, [2.,3.], [], t) ≈ [0., 0.]
-u_next = [0., 0.]
-@test f(u_next, [2.,3.], [], t) ≈ [3., -3.]
-
-resid = rand(2)
-f(resid, u_next, [2.,3.], [], t)
-@test resid ≈ [3., -3.]
-
-prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
-@test prob.u0 == [3., 1.]
-prob = ImplicitDiscreteProblem(sys, [], tspan)
-@test prob.u0 == [1., 1.]
-@variables x(t)
-@mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
-@test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
+@testset "Correct ImplicitDiscreteFunction" begin
+    @variables x(t) = 1
+    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+    tspan = (0, 10)
+    f = ImplicitDiscreteFunction(sys)
+    u_next = [3., 1.5]
+    @test f(u_next, [2.,3.], [], t) ≈ [0., 0.]
+    u_next = [0., 0.]
+    @test f(u_next, [2.,3.], [], t) ≈ [3., -3.]
+    
+    resid = rand(2)
+    f(resid, u_next, [2.,3.], [], t)
+    @test resid ≈ [3., -3.]
+    
+    prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
+    @test prob.u0 == [3., 1.]
+    prob = ImplicitDiscreteProblem(sys, [], tspan)
+    @test prob.u0 == [1., 1.]
+    @variables x(t)
+    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+    @test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
+end
 
 # 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)
+    f = ImplicitDiscreteFunction(sys)
+
+    function correct_f(u_next, u, p, t) 
+        [u[2] - u_next[1],
+         u[1] + u[2] - u_next[2],
+         1 - (u_next[1]+u_next[2])^2 - u_next[3]^2]
+    end
+
+    for _ in 1:10
+        u_next = rand(rng, 3)
+        u = rand(rng, 3)
+        @test correct_f(u_next, u, [], 0.) ≈ f(u_next, u, [], 0.)
+    end
+
+    # Initialization is satisfied.
+    prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
+    @test (prob.u0[1] + prob.u0[2])^2 + prob.u0[3]^2 ≈ 1
 end
 
 @testset "System with algebraic equations, implicit difference equations, explicit difference equations" begin
@@ -40,4 +61,3 @@ end
            y(k) ~ x(k) + x(k-2)*y(k-1)]
     @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
 end
-