From 70685b78a9332ecf30efafe1ef16460237eff8f7 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Mon, 21 Sep 2026 10:38:33 +0000
Subject: [PATCH] Build shifted variables in `var_derivative!` for discrete
 systems

`StateSelection.var_derivative!(::TearingState, ::Int)` built the new variable as
`Differential(iv)(v)` for every system. The variables of a discrete system are related by
shifts, and the differential form is not recognized by the discrete reassembly: the alias
elimination of `[u ~ p, y ~ u(k - 1), z ~ y]` requests the shifted form of the alias target
`y(k + 1)` and left `Differential(t)(y)` in the system, so that `mtkcompile` returned two
unknowns and one equation. The method now builds `simplify_shifts(Shift(iv, 1)(v))` when the
structure is discrete.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QBgBrXDidWG1G2KF64gAte
---
 .../src/stateselection_interface.jl           | 10 +++++++--
 lib/ModelingToolkitTearing/test/runtests.jl   | 21 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/lib/ModelingToolkitTearing/src/stateselection_interface.jl b/lib/ModelingToolkitTearing/src/stateselection_interface.jl
index 4b1e8ef..e86f8ca 100644
--- a/lib/ModelingToolkitTearing/src/stateselection_interface.jl
+++ b/lib/ModelingToolkitTearing/src/stateselection_interface.jl
@@ -2,8 +2,14 @@ function StateSelection.var_derivative!(ts::TearingState, v::Int)
     s = ts.structure
     var_diff = StateSelection.var_derivative_graph!(s, v)
     sys = ts.sys
-    D = Differential(MTKBase.get_iv(sys))
-    push!(ts.fullvars, D(ts.fullvars[v]))
+    iv = MTKBase.get_iv(sys)
+    # The variables of a discrete system are related by shifts, those of a continuous
+    # system by differentiation.
+    if StateSelection.is_only_discrete(s)
+        push!(ts.fullvars, MTKBase.simplify_shifts(Shift(iv, 1)(ts.fullvars[v])))
+    else
+        push!(ts.fullvars, Differential(iv)(ts.fullvars[v]))
+    end
     push!(ts.structure.state_priorities, ts.structure.state_priorities[v])
     push!(ts.structure.canonical_ranks, ts.structure.canonical_ranks[v] + 1)
     push!(ts.structure.var_types, ts.structure.var_types[v])
diff --git a/lib/ModelingToolkitTearing/test/runtests.jl b/lib/ModelingToolkitTearing/test/runtests.jl
index 3de3e54..1c5a19c 100644
--- a/lib/ModelingToolkitTearing/test/runtests.jl
+++ b/lib/ModelingToolkitTearing/test/runtests.jl
@@ -714,3 +714,24 @@ end
     MTKTearing.scalarize_tearing_state_eqs!(tss[cid])
     @test !iszero(Graphs.ne(tss[cid].structure.graph))
 end
+
+@testset "Alias of a variable shifted in a discrete system" begin
+    # The alias `z ~ y` with `y ~ u(k - 1)` eliminates `u` in favour of `Shift(t, 1)(y)`,
+    # and the shifted form of the alias target must be built as a shift, not as a
+    # differential.
+    k = ShiftIndex(t)
+    @parameters p = 3.0
+    @variables u(t) y(t) z(t)
+    @named sys = System([u ~ p, y ~ u(k - 1), z ~ y], t)
+    ss = mtkcompile(sys)
+    @test length(unknowns(ss)) == 1
+    for eq in [equations(ss); observed(ss)]
+        @test !MTKBase.isdifferential(eq.lhs)
+        @test !MTKBase.isdifferential(eq.rhs)
+    end
+    # `DiscreteProblem` advances the state one tick during construction, so `y` takes
+    # the value of `u` at the initial tick.
+    prob = DiscreteProblem(ss, [y => 0.0], (0, 5))
+    @test prob[y] == 3.0
+    @test prob[z] == 3.0
+end