fix tests and shift2term

Author: vyudu

src/structural_transformation/utils.jl

@@ -469,16 +469,21 @@ end
 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
+    op = operation(var)
+    iv = op.t
+    arg = only(arguments(var))
+    is_lowered = !isnothing(ModelingToolkit.getunshifted(arg))
+
+    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)
 
-    O = only(arguments(var))
+    O = is_lowered ? ModelingToolkit.getunshifted(arg) : arg
     oldop = operation(O)
-    newname = Symbol(string(nameof(oldop)), ds)
+    newname = backshift != 0 ? Symbol(string(nameof(oldop)), ds) : Symbol(string(nameof(oldop)))
 
     newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
     newvar = setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
@@ -487,24 +492,6 @@ function shift2term(var)
     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

@@ -269,12 +269,15 @@ function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
     for k in collect(keys(u0map))
         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)).")
+            isnothing(getunshifted(k)) && 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)).")
+            
+            updated[Shift(iv, 1)(k)] = v
         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)))).")
+        else
+            updated[Shift(iv, op.steps + 1)(only(arguments(k)))] = v
         end
 
-        updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
     end
     for var in unknowns(sys)
         op = operation(var)

test/discrete_system.jl

@@ -282,10 +282,10 @@ end
     prob = DiscreteProblem(de, [x(k - 1) => 2.0], (0, 10))
     @test prob[x] == 3.0
     @test prob[x(k - 1)] == 2.0
+    @variables xₜ₋₁(t)
     @test prob[xₜ₋₁] == 2.0
 
     # Test initial assignment with lowered variable
-    @variables xₜ₋₁(t)
     prob = DiscreteProblem(de, [xₜ₋₁(k-1) => 4.0], (0, 10))
     @test prob[x(k-1)] == prob[xₜ₋₁] == 1.0
     @test prob[x] == 5.
@@ -298,16 +298,17 @@ end
 
     # Test non-assigned initials are given default value
     @variables x(t) = 2.
+    @mtkbuild de = DiscreteSystem([x ~ x(k-1) + x(k-2)*x(k-3)], t)
     prob = DiscreteProblem(de, [x(k-3) => 12.], (0, 10))
     @test prob[x] == 26.0
     @test prob[x(k-1)] == 2.0
     @test prob[x(k-2)] == 2.0
 
     # Elaborate test
+    @variables xₜ₋₂(t) zₜ₋₁(t) z(t)
     eqs = [x ~ x(k-1) + z(k-2), 
            z ~ x(k-2) * x(k-3) - z(k-1)^2]
     @mtkbuild de = DiscreteSystem(eqs, t)
-    @variables xₜ₋₂(t) zₜ₋₁(t)
     u0 = [x(k-1) => 3, 
           xₜ₋₂(k-1) => 4, 
           x(k-2) => 1, 
@@ -316,4 +317,12 @@ end
     prob = DiscreteProblem(de, u0, (0, 10))
     @test prob[x] == 15
     @test prob[z] == -21
+
+    import ModelingToolkit: shift2term
+    # unknowns(de) = xₜ₋₁, x, zₜ₋₁, xₜ₋₂, z
+    vars = ModelingToolkit.value.(unknowns(de))
+    @test isequal(shift2term(Shift(t, 1)(vars[1])), vars[2])
+    @test isequal(shift2term(Shift(t, 1)(vars[4])), vars[1])
+    @test isequal(shift2term(Shift(t, -1)(vars[5])), vars[3])
+    @test isequal(shift2term(Shift(t, -2)(vars[2])), vars[4])
 end