feat: allow shifting array and scalarized variables

ext/MTKInfiniteOptExt.jl

@@ -340,7 +340,7 @@ function constructDefault(T::Type = Float64)
     A = map(T, A)
     α = map(T, α)
     c = map(T, c)
-    
+
     DiffEqBase.ImplicitRKTableau(A, c, α, 5)
 end
 
@@ -422,7 +422,6 @@ function _solve(prob::AbstractDynamicOptProblem, jump_solver, solver)
     DynamicOptSolution(model, sol, input_sol)
 end
 
-
 import InfiniteOpt: JuMP, GeneralVariableRef
 
 for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt]

src/discretedomain.jl

@@ -39,9 +39,13 @@ SymbolicUtils.isbinop(::Shift) = false
 
 function (D::Shift)(x, allow_zero = false)
     !allow_zero && D.steps == 0 && return x
-    Term{symtype(x)}(D, Any[x])
+    if Symbolics.isarraysymbolic(x)
+        Symbolics.array_term(D, x)
+    else
+        term(D, x)
+    end
 end
-function (D::Shift)(x::Num, allow_zero = false)
+function (D::Shift)(x::Union{Num, Symbolics.Arr}, allow_zero = false)
     !allow_zero && D.steps == 0 && return x
     vt = value(x)
     if iscall(vt)
@@ -52,11 +56,11 @@ function (D::Shift)(x::Num, allow_zero = false)
             if D.t === nothing || isequal(D.t, op.t)
                 arg = arguments(vt)[1]
                 newsteps = D.steps + op.steps
-                return Num(newsteps == 0 ? arg : Shift(D.t, newsteps)(arg))
+                return wrap(newsteps == 0 ? arg : Shift(D.t, newsteps)(arg))
             end
         end
     end
-    Num(D(vt, allow_zero))
+    wrap(D(vt, allow_zero))
 end
 SymbolicUtils.promote_symtype(::Shift, t) = t
 
@@ -202,11 +206,19 @@ function (xn::Num)(k::ShiftIndex)
     x = value(xn)
     # Verify that the independent variables of k and x match and that the expression doesn't have multiple variables
     vars = Symbolics.get_variables(x)
-    length(vars) == 1 ||
+    if length(vars) != 1
         error("Cannot shift a multivariate expression $x. Either create a new unknown and shift this, or shift the individual variables in the expression.")
-    args = Symbolics.arguments(vars[]) # args should be one element vector with the t in x(t)
-    length(args) == 1 ||
+    end
+    var = only(vars)
+    if !iscall(var)
+        throw(ArgumentError("Cannot shift time-independent variable $var"))
+    end
+    if operation(var) == getindex
+        var = first(arguments(var))
+    end
+    if length(arguments(var)) != 1
         error("Cannot shift an expression with multiple independent variables $x.")
+    end
 
     # d, _ = propagate_time_domain(xn)
     # if d != clock # this is only required if the variable has another clock
@@ -220,6 +232,34 @@ function (xn::Num)(k::ShiftIndex)
     Shift(t, steps)(xn) # a shift of k steps
 end
 
+function (xn::Symbolics.Arr)(k::ShiftIndex)
+    @unpack clock, steps = k
+    x = value(xn)
+    # Verify that the independent variables of k and x match and that the expression doesn't have multiple variables
+    vars = ModelingToolkit.vars(x)
+    if length(vars) != 1
+        error("Cannot shift a multivariate expression $x. Either create a new unknown and shift this, or shift the individual variables in the expression.")
+    end
+    var = only(vars)
+    if !iscall(var)
+        throw(ArgumentError("Cannot shift time-independent variable $var"))
+    end
+    if length(arguments(var)) != 1
+        error("Cannot shift an expression with multiple independent variables $x.")
+    end
+
+    # d, _ = propagate_time_domain(xn)
+    # if d != clock # this is only required if the variable has another clock
+    #     xn = Sample(t, clock)(xn)
+    # end
+    # QUESTION: should we return a variable with time domain set to k.clock?
+    xn = wrap(setmetadata(unwrap(xn), VariableTimeDomain, k.clock))
+    if steps == 0
+        return xn # x(k) needs no shift operator if the step of k is 0
+    end
+    Shift(t, steps)(xn) # a shift of k steps
+end
+
 Base.:+(k::ShiftIndex, i::Int) = ShiftIndex(k.clock, k.steps + i)
 Base.:-(k::ShiftIndex, i::Int) = k + (-i)
 

src/structural_transformation/utils.jl

@@ -471,6 +471,13 @@ function shift2term(var)
     op = operation(var)
     iv = op.t
     arg = only(arguments(var))
+    if operation(arg) === getindex
+        idxs = arguments(arg)[2:end]
+        newvar = shift2term(op(first(arguments(arg))))[idxs...]
+        unshifted = ModelingToolkit.getunshifted(newvar)[idxs...]
+        newvar = setmetadata(newvar, ModelingToolkit.VariableUnshifted, unshifted)
+        return newvar
+    end
     is_lowered = !isnothing(ModelingToolkit.getunshifted(arg))
 
     backshift = is_lowered ? op.steps + ModelingToolkit.getshift(arg) : op.steps

src/systems/discrete_system/discrete_system.jl

@@ -319,6 +319,7 @@ function SciMLBase.DiscreteProblem(
     iv = get_iv(sys)
 
     u0map = to_varmap(u0map, dvs)
+    scalarize_varmap!(u0map)
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
     f, u0, p = process_SciMLProblem(
         DiscreteFunction, sys, u0map, parammap; eval_expression, eval_module, build_initializeprob = false)

src/utils.jl

@@ -130,6 +130,9 @@ function check_parameters(ps, iv)
 end
 
 function is_delay_var(iv, var)
+    if Symbolics.isarraysymbolic(var)
+        return is_delay_var(iv, first(collect(var)))
+    end
     args = nothing
     try
         args = arguments(var)

test/discrete_system.jl

@@ -325,3 +325,21 @@ end
     @test isequal(shift2term(Shift(t, -1)(vars[4])), vars[5])
     @test isequal(shift2term(Shift(t, -2)(vars[1])), vars[3])
 end
+
+@testset "Shifted array variables" begin
+    @variables x(t)[1:2] y(t)[1:2]
+    k = ShiftIndex(t)
+    eqs = [
+        x(k) ~ x(k - 1) + x(k - 2),
+        y[1](k) ~ y[1](k - 1) + y[1](k - 2),
+        y[2](k) ~ y[2](k - 1) + y[2](k - 2)
+    ]
+    @mtkbuild sys = DiscreteSystem(eqs, t)
+    prob = DiscreteProblem(sys,
+        [x(k - 1) => ones(2), x(k - 2) => zeros(2), y[1](k - 1) => 1.0,
+            y[1](k - 2) => 0.0, y[2](k - 1) => 1.0, y[2](k - 2) => 0.0],
+        (0, 4))
+    @test all(isone, prob.u0)
+    sol = solve(prob, FunctionMap())
+    @test sol[[x..., y...], end] == 8ones(4)
+end

test/initializationsystem.jl

@@ -1194,13 +1194,13 @@ end
     @test integ[x] ≈ 1 / cbrt(3)
     @test integ[y] ≈ 2 / cbrt(3)
     @test integ.ps[p] == 1.0
-    @test integ.ps[q] ≈ 3 / cbrt(3) atol=1e-5
+    @test integ.ps[q]≈3 / cbrt(3) atol=1e-5
     prob2 = remake(prob; u0 = [y => 3x], p = [q => 2x])
     integ2 = init(prob2)
-    @test integ2[x] ≈ cbrt(3 / 28) atol=1e-5
-    @test integ2[y] ≈ 3cbrt(3 / 28) atol=1e-5
+    @test integ2[x]≈cbrt(3 / 28) atol=1e-5
+    @test integ2[y]≈3cbrt(3 / 28) atol=1e-5
     @test integ2.ps[p] == 1.0
-    @test integ2.ps[q] ≈ 2cbrt(3 / 28) atol=1e-5
+    @test integ2.ps[q]≈2cbrt(3 / 28) atol=1e-5
 end
 
 function test_dummy_initialization_equation(prob, var)