Merge pull request #52 from JuliaComputing/as/backshift-additional-observed

fix: backshift `additional_observed` when generating equations

Author: AayushSabharwal

lib/ModelingToolkitTearing/src/reassemble.jl

@@ -939,9 +939,13 @@ function update_simplified_system!(
         MTKBase.isdiffeq(eq) || continue
         obs_sub[eq.lhs] = eq.rhs
     end
+    (; additional_observed) = state
+    if StateSelection.is_only_discrete(structure)
+        additional_observed = map(Base.Fix2(backshift_expr, iv), additional_observed)
+    end
     # TODO: compute the dependency correctly so that we don't have to do this
     obs = [substitute(observed(sys), obs_sub); solved_eqs;
-           substitute(state.additional_observed, obs_sub)]
+           substitute(additional_observed, obs_sub)]
 
     filterer = let diff_to_var = diff_to_var, ispresent = ispresent, fullvars = fullvars,
         solved_vars = solved_vars

lib/ModelingToolkitTearing/src/tearingstate.jl

@@ -591,7 +591,15 @@ function shift_discrete_system(ts::TearingState)
         eqs[i] = MTKBase.simplify_shifts(substitute(
             eqs[i], discmap; filterer = Symbolics.FPSubFilterer{Union{Sample, Hold, Pre}}()))
     end
+
+    original_eqs = copy(ts.original_eqs)
+    for i in eachindex(original_eqs)
+        original_eqs[i] = MTKBase.simplify_shifts(substitute(
+            original_eqs[i], discmap; filterer = Symbolics.FPSubFilterer{Union{Sample, Hold, Pre}}()))
+    end
+
     @set! ts.sys.eqs = eqs
+    @set! ts.original_eqs = original_eqs
     @set! ts.fullvars = fullvars
     return ts
 end

lib/ModelingToolkitTearing/test/runtests.jl

@@ -110,3 +110,22 @@ end
     @test isempty(unknowns(sys))
     @test length(observed(sys)) == 2
 end
+
+@testset "`additional_observed` works correctly for discrete systems" begin
+    @variables x(t) y(t)
+    k = ShiftIndex(t)
+    @named sys = System([x(k) ~ x(k-1) + 1, y(k) ~ x(k)], t)
+    ts = TearingState(sys)
+    # Original equations should mirror those of the system
+    @test issetequal(ts.original_eqs, equations(sys))
+    # `mark_discrete` should shift `original_eqs` forward too
+    ts = MTKTearing.mark_discrete(ts)
+    @test issetequal(ts.original_eqs, [x(k+1) ~ x(k) + 1, y(k+1) ~ x(k + 1)])
+    # `trivial_tearing!` should then move them to `additional_observed`
+    StateSelection.trivial_tearing!(ts)
+    @test issetequal(ts.additional_observed, [y(k+1) ~ x(k+1)])
+    # MTK calls `trivial_tearing!`, and the reassemble process should backshift
+    # `additional_observed`
+    ss = mtkcompile(sys)
+    @test any(isequal(y(k) ~ x(k)), observed(ss))
+end