Shift unscalarized array variables with their scalarized elements

lib/ModelingToolkitTearing/Project.toml

@@ -1,6 +1,6 @@
 name = "ModelingToolkitTearing"
 uuid = "6bb917b9-1269-42b9-9f7c-b0dca72083ab"
-version = "1.20.6"
+version = "1.20.7"
 authors = ["Aayush Sabharwal <aayush.sabharwal@juliahub.com>"]
 
 [deps]

lib/ModelingToolkitTearing/src/tearingstate.jl

@@ -1072,6 +1072,17 @@ function shift_discrete_system(ts::TearingState)
     # necessary for hybrid system handling.
     (; fullvars, sys) = ts
     fullvars_set = Set{SymbolicT}(fullvars)
+    # `fullvars` only contains scalarized array elements, so an array variable used
+    # without scalarization (as in `f(x(k - 1))` for an array `x`) is absent from
+    # `fullvars_set` and would be left unshifted while its elements are shifted, which
+    # puts the two forms one tick apart. Add the arrays the scalarized elements belong
+    # to, so that both forms are shifted alike; the `isoperator` check below discards
+    # the entries that are a `Sample`, `Hold` or `Pre` of an array.
+    for v in fullvars
+        arr, isidx = MTKBase.split_indexed_var(v)
+        isidx || continue
+        push!(fullvars_set, strip_shifts(arr))
+    end
     discvars = OrderedSet{SymbolicT}()
     eqs = equations(sys)
     for eq in eqs

lib/ModelingToolkitTearing/src/utils.jl

@@ -35,6 +35,18 @@ function descend_lower_shift_varname(var, iv)
     end
 end
 
+"""
+    $TYPEDSIGNATURES
+
+Remove all `Shift`s applied to `var` and return the variable they are applied to.
+"""
+function strip_shifts(var::SymbolicT)
+    @match var begin
+        BSImpl.Term(; f, args) && if f isa Shift end => strip_shifts(args[1])
+        _ => var
+    end
+end
+
 function is_time_dependent_parameter(p::SymbolicT, allps::Set{SymbolicT}, iv::SymbolicT)
     return p in allps && @match p begin
         BSImpl.Term(; f, args) => begin

lib/ModelingToolkitTearing/test/runtests.jl

@@ -735,3 +735,62 @@ end
     @test prob[y] == 3.0
     @test prob[z] == 3.0
 end
+
+# A function of a whole array which cannot be scalarized, so that the unscalarized form of
+# an array variable survives into the compiled system.
+combine_array(v) = sum(v)
+@register_symbolic combine_array(v::AbstractVector)
+
+@testset "Shifts applied to whole array variables" begin
+    # https://github.com/SciML/ModelingToolkit.jl/issues/5169
+    k = ShiftIndex(t)
+    @variables y(t) ud(t) (xd(t))[1:2]
+    # The elements of `xd` are aliased away, so `fullvars` only holds the scalarized
+    # elements while the equation for `ud` uses the whole array.
+    common_eqs = [y ~ y(k - 1) + 1, xd[1] ~ y, xd[2] ~ 2y]
+
+    @named sys = System([common_eqs; ud ~ combine_array(xd(k - 1))], t)
+    @named ref = System([common_eqs; ud ~ xd(k - 1)[1] + xd(k - 1)[2]], t)
+    ss = mtkcompile(sys)
+    ssref = mtkcompile(ref)
+
+    @test issetequal(unknowns(ss), unknowns(ssref))
+    udeq = only(filter(eq -> isequal(eq.lhs, unwrap(ud)), observed(ss)))
+    # `xd` is used one tick back, so the observed equation for `ud` must refer to the
+    # same shifted variable the unknowns are lowered to.
+    prev = only(arguments(udeq.rhs))
+    @test issubset(Set(collect(prev)), Set(unknowns(ss)))
+
+    u0 = [y(k - 1) => 0.0, xd(k - 1) => [3.0, 4.0]]
+    prob = DiscreteProblem(ss, u0, (0, 5))
+    probref = DiscreteProblem(ssref, u0, (0, 5))
+    @test prob[ud] == probref[ud]
+end
+
+@testset "Shifts applied to whole array variables in a clock partition" begin
+    # https://github.com/SciML/ModelingToolkit.jl/issues/5169
+    dt = 0.1
+    k = ShiftIndex(Clock(dt))
+    @variables x(t) y(t) u(t) ud(t) (xd(t))[1:2]
+    eqs = [
+        xd[1] ~ Sample(dt)(y)
+        xd[2] ~ 2 * Sample(dt)(y)
+        ud ~ combine_array(xd(k - 1))
+        u ~ Hold(ud)
+        D(x) ~ -x + u
+        y ~ x
+    ]
+    @named sys = System(eqs, t)
+    ci = MTKTearing.infer_clocks!(MTKTearing.ClockInference(TearingState(sys)))
+    tss, _, continuous_id, _ = MTKTearing.split_system(ci)
+    disc = tss[findfirst(!=(continuous_id), eachindex(tss))]
+
+    # `split_system` shifts the discrete partition forward by one tick, and the
+    # unscalarized `xd` must be shifted along with its elements.
+    vars = Set{Symbolics.SymbolicT}()
+    for eq in equations(disc)
+        SU.search_variables!(vars, eq; is_atomic = MTKBase.OperatorIsAtomic{SU.Operator}())
+    end
+    @test unwrap(xd) in vars
+    @test !(unwrap(xd(k - 1)) in vars)
+end