fix: make the codegen work with shifted observables

Author: vyudu

src/structural_transformation/symbolics_tearing.jl

@@ -558,6 +558,10 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
             end
 
             total_sub[simplify_shifts(neweq.lhs)] = neweq.rhs
+            # Substitute unshifted variables x(k), y(k) on RHS of implicit equations
+            if is_only_discrete(structure)
+                var_to_diff[iv] === nothing && (total_sub[var] = neweq.rhs)
+            end
             push!(diff_eqs, neweq)
             push!(diffeq_idxs, ieq)
             push!(diff_vars, diff_to_var[iv])

src/systems/discrete_system/implicit_discrete_system.jl

@@ -264,13 +264,20 @@ end
 function generate_function(
         sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
     iv = get_iv(sys)
+    # Algebraic equations get shifted forward 1, to match with differential equations
     exprs = map(equations(sys)) do eq
-        _iszero(eq.lhs) ? eq.rhs : (distribute_shift(Shift(iv, -1)(eq.rhs)) - distribute_shift(Shift(iv, -1)(eq.lhs)))
+        _iszero(eq.lhs) ? distribute_shift(Shift(iv, 1)(eq.rhs)) : (eq.rhs - eq.lhs)
     end
 
-    u_next = dvs
-    u = map(Shift(iv, -1), u_next)
-    build_function_wrapper(sys, exprs, u_next, u, ps..., iv; p_start = 3, kwargs...)
+    # Handle observables in algebraic equations, since they are shifted
+    obs = observed(sys)
+    shifted_obs = [distribute_shift(Shift(iv, 1)(eq)) for eq in obs]
+    obsidxs = observed_equations_used_by(sys, exprs; obs = shifted_obs)
+    extra_assignments = [Assignment(shifted_obs[i].lhs, shifted_obs[i].rhs) for i in obsidxs]
+
+    u_next = map(Shift(iv, 1), dvs)
+    u = dvs
+    build_function_wrapper(sys, exprs, u_next, u, ps..., iv; p_start = 3, extra_assignments, kwargs...)
 end
 
 function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
@@ -279,15 +286,22 @@ function shift_u0map_forward(sys::ImplicitDiscreteSystem, 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[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)(k))).")
+        else
+            updated[k] = v
         end
-        updated[shift2term(k)] = v
     end
     for var in unknowns(sys)
         op = operation(var)
-        haskey(updated, var) && continue
         root = getunshifted(var)
+        shift = getshift(var)
         isnothing(root) && continue
+        (haskey(updated, Shift(iv, shift)(root)) || haskey(updated, var)) && continue
         haskey(defs, root) || error("Initial condition for $var not provided.")
         updated[var] = defs[root]
     end
@@ -317,7 +331,9 @@ function SciMLBase.ImplicitDiscreteProblem(
     u0map = to_varmap(u0map, dvs)
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
     f, u0, p = process_SciMLProblem(
-        ImplicitDiscreteFunction, sys, u0map, parammap; eval_expression, eval_module)
+        ImplicitDiscreteFunction, sys, u0map, parammap; eval_expression, eval_module, kwargs...)
+
+    kwargs = filter_kwargs(kwargs)
     ImplicitDiscreteProblem(f, u0, tspan, p; kwargs...)
 end
 

src/systems/systemstructure.jl

@@ -440,7 +440,7 @@ function TearingState(sys; quick_cancel = false, check = true)
         SystemStructure(complete(var_to_diff), complete(eq_to_diff),
             complete(graph), nothing, var_types, sys isa AbstractDiscreteSystem),
         Any[])
-    if sys isa AbstractDiscreteSystem
+    if sys isa DiscreteSystem
         ts = shift_discrete_system(ts)
     end
     return ts

test/implicit_discrete_system.jl

@@ -5,12 +5,13 @@ using StableRNGs
 k = ShiftIndex(t)
 rng = StableRNG(22525) 
 
-# Shift(t, -1)(x(t)) - x_{t-1}(t)
-# -3 - x(t) + x(t)*x_{t-1}
 @testset "Correct ImplicitDiscreteFunction" begin
     @variables x(t) = 1
     @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
     tspan = (0, 10)
+
+    # u[2] - u_next[1]
+    # -3 - u_next[2] + u_next[2]*u_next[1]
     f = ImplicitDiscreteFunction(sys)
     u_next = [3., 1.5]
     @test f(u_next, [2.,3.], [], t) ≈ [0., 0.]
@@ -30,7 +31,6 @@ rng = StableRNG(22525)
     @test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
 end
 
-# Test solvers
 @testset "System with algebraic equations" begin
     @variables x(t) y(t)
     eqs = [x(k) ~ x(k-1) + x(k-2), 
@@ -51,13 +51,24 @@ end
     end
 
     # Initialization is satisfied.
-    prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
+    prob = ImplicitDiscreteProblem(sys, [x(k-1) => .3, x(k-2) => .4], (0, 10), guesses = [y => 1])
     @test (prob.u0[1] + prob.u0[2])^2 + prob.u0[3]^2 ≈ 1
 end
 
-@testset "System with algebraic equations, implicit difference equations, explicit difference equations" begin
-    @variables x(t) y(t)
+@testset "Handle observables in function codegen" begin
+    # Observable appears in differential equation
+    @variables x(t) y(t) z(t)
     eqs = [x(k) ~ x(k-1) + x(k-2),
-           y(k) ~ x(k) + x(k-2)*y(k-1)]
+           y(k) ~ x(k) + x(k-2)*z(k-1), 
+           x + y + z ~ 2]
+    @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
+    @test length(unknowns(sys)) == length(equations(sys)) == 3
+    @test occursin("var\"y(t)\"", string(ImplicitDiscreteFunctionExpr(sys)))
+
+    # Shifted observable that appears in algebraic equation is properly handled.
+    eqs = [z(k) ~ x(k) + sin(x(k)), 
+           y(k) ~ x(k-1) + x(k-2),
+           z(k) * x(k) ~ 3]
     @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
+    @test occursin("var\"Shift(t, 1)(z(t))\"", string(ImplicitDiscreteFunctionExpr(sys)))
 end