From cb64a00cd58209f5c02cab7ae1750dc8e258342c Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 11 Feb 2025 22:43:18 -0500
Subject: [PATCH 01/16] refactor: sort systems tests into new folder

---
 test/{ => systems}/abstractsystem.jl           | 0
 test/{ => systems}/bigsystem.jl                | 0
 test/{ => systems}/discrete_system.jl          | 0
 test/{ => systems}/implicit_discrete_system.jl | 0
 test/{ => systems}/initializationsystem.jl     | 0
 test/{ => systems}/jumpsystem.jl               | 0
 test/{ => systems}/nonlinearsystem.jl          | 0
 test/{ => systems}/odesystem.jl                | 0
 test/{ => systems}/optimizationsystem.jl       | 0
 test/{pde.jl => systems/pdesystem.jl}          | 0
 test/{ => systems}/sdesystem.jl                | 0
 test/{ => systems}/steadystatesystems.jl       | 0
 12 files changed, 0 insertions(+), 0 deletions(-)
 rename test/{ => systems}/abstractsystem.jl (100%)
 rename test/{ => systems}/bigsystem.jl (100%)
 rename test/{ => systems}/discrete_system.jl (100%)
 rename test/{ => systems}/implicit_discrete_system.jl (100%)
 rename test/{ => systems}/initializationsystem.jl (100%)
 rename test/{ => systems}/jumpsystem.jl (100%)
 rename test/{ => systems}/nonlinearsystem.jl (100%)
 rename test/{ => systems}/odesystem.jl (100%)
 rename test/{ => systems}/optimizationsystem.jl (100%)
 rename test/{pde.jl => systems/pdesystem.jl} (100%)
 rename test/{ => systems}/sdesystem.jl (100%)
 rename test/{ => systems}/steadystatesystems.jl (100%)

diff --git a/test/abstractsystem.jl b/test/systems/abstractsystem.jl
similarity index 100%
rename from test/abstractsystem.jl
rename to test/systems/abstractsystem.jl
diff --git a/test/bigsystem.jl b/test/systems/bigsystem.jl
similarity index 100%
rename from test/bigsystem.jl
rename to test/systems/bigsystem.jl
diff --git a/test/discrete_system.jl b/test/systems/discrete_system.jl
similarity index 100%
rename from test/discrete_system.jl
rename to test/systems/discrete_system.jl
diff --git a/test/implicit_discrete_system.jl b/test/systems/implicit_discrete_system.jl
similarity index 100%
rename from test/implicit_discrete_system.jl
rename to test/systems/implicit_discrete_system.jl
diff --git a/test/initializationsystem.jl b/test/systems/initializationsystem.jl
similarity index 100%
rename from test/initializationsystem.jl
rename to test/systems/initializationsystem.jl
diff --git a/test/jumpsystem.jl b/test/systems/jumpsystem.jl
similarity index 100%
rename from test/jumpsystem.jl
rename to test/systems/jumpsystem.jl
diff --git a/test/nonlinearsystem.jl b/test/systems/nonlinearsystem.jl
similarity index 100%
rename from test/nonlinearsystem.jl
rename to test/systems/nonlinearsystem.jl
diff --git a/test/odesystem.jl b/test/systems/odesystem.jl
similarity index 100%
rename from test/odesystem.jl
rename to test/systems/odesystem.jl
diff --git a/test/optimizationsystem.jl b/test/systems/optimizationsystem.jl
similarity index 100%
rename from test/optimizationsystem.jl
rename to test/systems/optimizationsystem.jl
diff --git a/test/pde.jl b/test/systems/pdesystem.jl
similarity index 100%
rename from test/pde.jl
rename to test/systems/pdesystem.jl
diff --git a/test/sdesystem.jl b/test/systems/sdesystem.jl
similarity index 100%
rename from test/sdesystem.jl
rename to test/systems/sdesystem.jl
diff --git a/test/steadystatesystems.jl b/test/systems/steadystatesystems.jl
similarity index 100%
rename from test/steadystatesystems.jl
rename to test/systems/steadystatesystems.jl

From a901d097ac630ef6c553cbe947528af617a8ab64 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 11 Feb 2025 22:44:42 -0500
Subject: [PATCH 02/16] revert

---
 test/{systems => }/abstractsystem.jl           | 0
 test/{systems => }/bigsystem.jl                | 0
 test/{systems => }/discrete_system.jl          | 0
 test/{systems => }/implicit_discrete_system.jl | 0
 test/{systems => }/initializationsystem.jl     | 0
 test/{systems => }/jumpsystem.jl               | 0
 test/{systems => }/nonlinearsystem.jl          | 0
 test/{systems => }/odesystem.jl                | 0
 test/{systems => }/optimizationsystem.jl       | 0
 test/{systems => }/pdesystem.jl                | 0
 test/{systems => }/sdesystem.jl                | 0
 test/{systems => }/steadystatesystems.jl       | 0
 12 files changed, 0 insertions(+), 0 deletions(-)
 rename test/{systems => }/abstractsystem.jl (100%)
 rename test/{systems => }/bigsystem.jl (100%)
 rename test/{systems => }/discrete_system.jl (100%)
 rename test/{systems => }/implicit_discrete_system.jl (100%)
 rename test/{systems => }/initializationsystem.jl (100%)
 rename test/{systems => }/jumpsystem.jl (100%)
 rename test/{systems => }/nonlinearsystem.jl (100%)
 rename test/{systems => }/odesystem.jl (100%)
 rename test/{systems => }/optimizationsystem.jl (100%)
 rename test/{systems => }/pdesystem.jl (100%)
 rename test/{systems => }/sdesystem.jl (100%)
 rename test/{systems => }/steadystatesystems.jl (100%)

diff --git a/test/systems/abstractsystem.jl b/test/abstractsystem.jl
similarity index 100%
rename from test/systems/abstractsystem.jl
rename to test/abstractsystem.jl
diff --git a/test/systems/bigsystem.jl b/test/bigsystem.jl
similarity index 100%
rename from test/systems/bigsystem.jl
rename to test/bigsystem.jl
diff --git a/test/systems/discrete_system.jl b/test/discrete_system.jl
similarity index 100%
rename from test/systems/discrete_system.jl
rename to test/discrete_system.jl
diff --git a/test/systems/implicit_discrete_system.jl b/test/implicit_discrete_system.jl
similarity index 100%
rename from test/systems/implicit_discrete_system.jl
rename to test/implicit_discrete_system.jl
diff --git a/test/systems/initializationsystem.jl b/test/initializationsystem.jl
similarity index 100%
rename from test/systems/initializationsystem.jl
rename to test/initializationsystem.jl
diff --git a/test/systems/jumpsystem.jl b/test/jumpsystem.jl
similarity index 100%
rename from test/systems/jumpsystem.jl
rename to test/jumpsystem.jl
diff --git a/test/systems/nonlinearsystem.jl b/test/nonlinearsystem.jl
similarity index 100%
rename from test/systems/nonlinearsystem.jl
rename to test/nonlinearsystem.jl
diff --git a/test/systems/odesystem.jl b/test/odesystem.jl
similarity index 100%
rename from test/systems/odesystem.jl
rename to test/odesystem.jl
diff --git a/test/systems/optimizationsystem.jl b/test/optimizationsystem.jl
similarity index 100%
rename from test/systems/optimizationsystem.jl
rename to test/optimizationsystem.jl
diff --git a/test/systems/pdesystem.jl b/test/pdesystem.jl
similarity index 100%
rename from test/systems/pdesystem.jl
rename to test/pdesystem.jl
diff --git a/test/systems/sdesystem.jl b/test/sdesystem.jl
similarity index 100%
rename from test/systems/sdesystem.jl
rename to test/sdesystem.jl
diff --git a/test/systems/steadystatesystems.jl b/test/steadystatesystems.jl
similarity index 100%
rename from test/systems/steadystatesystems.jl
rename to test/steadystatesystems.jl

From f4a9d1369555a35e3c59542ce259fbcf3e6ba5f2 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 13 Feb 2025 16:15:48 -0500
Subject: [PATCH 03/16] up

---
 .../implicit_discrete_system.jl               | 23 +++----------------
 1 file changed, 3 insertions(+), 20 deletions(-)

diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index 3a137eb8c1..82ef66d17e 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -238,8 +238,6 @@ function ImplicitDiscreteSystem(eqs, iv; kwargs...)
     return ImplicitDiscreteSystem(eqs, iv,
         collect(allunknowns), collect(new_ps); kwargs...)
 end
-# basically at every timestep it should build a nonlinear solve
-# Previous timesteps should be treated as parameters? is this right? 
 
 function flatten(sys::ImplicitDiscreteSystem, noeqs = false)
     systems = get_systems(sys)
@@ -263,25 +261,10 @@ end
 
 function generate_function(
         sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
-    if !iscomplete(sys)
-        error("A completed system is required. Call `complete` or `structural_simplify` on the system.")
-    end
-    p = (reorder_parameters(sys, unwrap.(ps))..., cachesyms...)
-    isscalar = !(exprs isa AbstractArray)
-    pre, sol_states = get_substitutions_and_solved_unknowns(sys, isscalar ? [exprs] : exprs)
-    if postprocess_fbody === nothing
-        postprocess_fbody = pre
-    end
-    if states === nothing
-        states = sol_states
-    end
     exprs = [eq.lhs - eq.rhs for eq in equations(sys)]
-    u = map(Shift(iv, -1), dvs)
-    u_next = dvs
-
-    wrap_code = wrap_code .∘ wrap_array_vars(sys, exprs) .∘ wrap_parameter_dependencies(sys, false)
-
-    build_function(exprs, u_next, u, p..., get_iv(sys))
+    u = dvs
+    u_next = map(Shift(iv, 1), u)
+    generate_custom_function(sys, exprs, u_next, u, ps..., get_iv(sys); kwargs...)
 end
 
 function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)

From 5f231c751a5def2e2a7d9911ae961793c8e2e0e2 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 14 Feb 2025 07:54:04 -0500
Subject: [PATCH 04/16] fix: fix initialization cases

---
 src/ModelingToolkit.jl                        |  3 ++
 .../StructuralTransformations.jl              |  2 +-
 src/structural_transformation/utils.jl        |  5 ++--
 .../discrete_system/discrete_system.jl        | 10 +++++--
 .../implicit_discrete_system.jl               | 25 +++++++++-------
 src/systems/systems.jl                        |  4 +--
 src/systems/systemstructure.jl                |  2 +-
 test/implicit_discrete_system.jl              | 30 +++++++++++++++++++
 test/runtests.jl                              |  1 +
 9 files changed, 62 insertions(+), 20 deletions(-)
 create mode 100644 test/implicit_discrete_system.jl

diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 09b59c4ed6..ce3db8bde7 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -123,6 +123,7 @@ abstract type AbstractTimeIndependentSystem <: AbstractSystem end
 abstract type AbstractODESystem <: AbstractTimeDependentSystem end
 abstract type AbstractMultivariateSystem <: AbstractSystem end
 abstract type AbstractOptimizationSystem <: AbstractTimeIndependentSystem end
+abstract type AbstractDiscreteSystem <: AbstractTimeDependentSystem end
 
 function independent_variable end
 
@@ -165,6 +166,7 @@ include("systems/diffeqs/modelingtoolkitize.jl")
 include("systems/diffeqs/basic_transformations.jl")
 
 include("systems/discrete_system/discrete_system.jl")
+include("systems/discrete_system/implicit_discrete_system.jl")
 
 include("systems/jumps/jumpsystem.jl")
 
@@ -230,6 +232,7 @@ export DAEFunctionExpr, DAEProblemExpr
 export SDESystem, SDEFunction, SDEFunctionExpr, SDEProblemExpr
 export SystemStructure
 export DiscreteSystem, DiscreteProblem, DiscreteFunction, DiscreteFunctionExpr
+export ImplicitDiscreteSystem, ImplicitDiscreteProblem, ImplicitDiscreteFunction, ImplicitDiscreteFunctionExpr
 export JumpSystem
 export ODEProblem, SDEProblem
 export NonlinearFunction, NonlinearFunctionExpr
diff --git a/src/structural_transformation/StructuralTransformations.jl b/src/structural_transformation/StructuralTransformations.jl
index 510352ad59..c7b3957c62 100644
--- a/src/structural_transformation/StructuralTransformations.jl
+++ b/src/structural_transformation/StructuralTransformations.jl
@@ -63,7 +63,7 @@ export torn_system_jacobian_sparsity
 export full_equations
 export but_ordered_incidence, lowest_order_variable_mask, highest_order_variable_mask
 export computed_highest_diff_variables
-export shift2term, lower_shift_varname
+export shift2term, lower_shift_varname, simplify_shifts
 
 include("utils.jl")
 include("pantelides.jl")
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index f64a8da132..db79c8a124 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -452,11 +452,10 @@ end
 # For discrete variables. Turn Shift(t, k)(x(t)) into xₜ₋ₖ(t)
 function lower_shift_varname(var, iv)
     op = operation(var)
-    op isa Shift || return Shift(iv, 0)(var, true) # hack to prevent simplification of x(t) - x(t)
-    if op.steps < 0
+    if op isa Shift && op.steps < 0
         return shift2term(var)
     else
-        return var
+        return Shift(iv, 0)(var, true)
     end
 end
 
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index 9a9ac47853..703efcc462 100644
--- a/src/systems/discrete_system/discrete_system.jl
+++ b/src/systems/discrete_system/discrete_system.jl
@@ -17,7 +17,7 @@ eqs = [x(k+1) ~ σ*(y-x),
 @named de = DiscreteSystem(eqs)
 ```
 """
-struct DiscreteSystem <: AbstractTimeDependentSystem
+struct DiscreteSystem <: AbstractDiscreteSystem
     """
     A tag for the system. If two systems have the same tag, then they are
     structurally identical.
@@ -237,6 +237,8 @@ function DiscreteSystem(eqs, iv; kwargs...)
         collect(allunknowns), collect(new_ps); kwargs...)
 end
 
+DiscreteSystem(eq::Equation, args...; kwargs...) = DiscreteSystem([eq], args...; kwargs...)
+
 function flatten(sys::DiscreteSystem, noeqs = false)
     systems = get_systems(sys)
     if isempty(systems)
@@ -271,14 +273,16 @@ function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
         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)).")
         end
-        updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
+        k_next = Shift(iv, op.steps + 1)(arguments(k)[1])
+        operation(k_next) isa Shift ? updated[shift2term(k_next)] = v :
+                                      updated[k_next] = v
     end
     for var in unknowns(sys)
         op = operation(var)
         haskey(updated, var) && continue
         root = getunshifted(var)
         isnothing(root) && continue
-        haskey(defs, root) || error("Initial condition for $root not provided.")
+        haskey(defs, root) || error("Initial condition for $var not provided.")
         updated[var] = defs[root]
     end
     return updated
diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index 82ef66d17e..e0450a1ada 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -239,6 +239,8 @@ function ImplicitDiscreteSystem(eqs, iv; kwargs...)
         collect(allunknowns), collect(new_ps); kwargs...)
 end
 
+ImplicitDiscreteSystem(eq::Equation, args...; kwargs...) = ImplicitDiscreteSystem([eq], args...; kwargs...)
+
 function flatten(sys::ImplicitDiscreteSystem, noeqs = false)
     systems = get_systems(sys)
     if isempty(systems)
@@ -261,10 +263,14 @@ end
 
 function generate_function(
         sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
-    exprs = [eq.lhs - eq.rhs for eq in equations(sys)]
-    u = dvs
-    u_next = map(Shift(iv, 1), u)
-    generate_custom_function(sys, exprs, u_next, u, ps..., get_iv(sys); kwargs...)
+    iv = get_iv(sys)
+    exprs = map(equations(sys)) do eq
+        _iszero(eq.lhs) ? eq.rhs : (simplify_shifts(Shift(iv, -1)(eq.rhs)) - simplify_shifts(Shift(iv, -1)(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...)
 end
 
 function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
@@ -275,13 +281,13 @@ function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
         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)).")
         end
-        updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
+        updated[shift2term(k)] = v
     end
     for var in unknowns(sys)
         op = operation(var)
-        op isa Shift || continue
         haskey(updated, var) && continue
-        root = first(arguments(var))
+        root = getunshifted(var)
+        isnothing(root) && continue
         haskey(defs, root) || error("Initial condition for $var not provided.")
         updated[var] = defs[root]
     end
@@ -301,7 +307,7 @@ function SciMLBase.ImplicitDiscreteProblem(
         kwargs...
 )
     if !iscomplete(sys)
-        error("A completed `ImplicitDiscreteSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `ImplicitDiscreteProblem`")
+        error("A completed `ImplicitDiscreteSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `ImplicitDiscreteProblem`.")
     end
     dvs = unknowns(sys)
     ps = parameters(sys)
@@ -312,8 +318,7 @@ function SciMLBase.ImplicitDiscreteProblem(
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
     f, u0, p = process_SciMLProblem(
         ImplicitDiscreteFunction, sys, u0map, parammap; eval_expression, eval_module)
-    u0 = f(u0, p, tspan[1])
-    NonlinearProblem(f, u0, tspan, p; kwargs...)
+    ImplicitDiscreteProblem(f, u0, tspan, p; kwargs...)
 end
 
 function SciMLBase.ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...; kwargs...)
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index f8630f2d20..40f3353884 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -42,8 +42,8 @@ function structural_simplify(
     if newsys isa DiscreteSystem &&
        any(eq -> symbolic_type(eq.lhs) == NotSymbolic(), equations(newsys))
         error("""
-            Encountered algebraic equations when simplifying discrete system. This is \
-            not yet supported.
+            Encountered algebraic equations when simplifying discrete system. Please construct \
+            an ImplicitDiscreteSystem instead.
         """)
     end
     for pass in additional_passes
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 6fee78cfd6..a6e6b1d218 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -438,7 +438,7 @@ function TearingState(sys; quick_cancel = false, check = true)
 
     ts = TearingState(sys, fullvars,
         SystemStructure(complete(var_to_diff), complete(eq_to_diff),
-            complete(graph), nothing, var_types, sys isa DiscreteSystem),
+            complete(graph), nothing, var_types, sys isa AbstractDiscreteSystem),
         Any[])
     if sys isa DiscreteSystem
         ts = shift_discrete_system(ts)
diff --git a/test/implicit_discrete_system.jl b/test/implicit_discrete_system.jl
new file mode 100644
index 0000000000..88e79acc77
--- /dev/null
+++ b/test/implicit_discrete_system.jl
@@ -0,0 +1,30 @@
+using ModelingToolkit, Test
+using ModelingToolkit: t_nounits as t
+
+k = ShiftIndex(t)
+@variables x(t) = 1
+@mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+tspan = (0, 10)
+
+# Shift(t, -1)(x(t)) - x_{t-1}(t)
+# -3 - x(t) + x(t)*x_{t-1}
+f = ImplicitDiscreteFunction(sys)
+u_next = [3., 1.5]
+@test f(u_next, [2.,3.], [], t) ≈ [0., 0.]
+u_next = [0., 0.]
+@test f(u_next, [2.,3.], [], t) ≈ [3., -3.]
+
+resid = rand(2)
+f(resid, u_next, [2.,3.], [], t)
+@test resid ≈ [3., -3.]
+
+# Initialization cases.
+prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
+@test prob.u0 == [3., 1.]
+prob = ImplicitDiscreteProblem(sys, [], tspan)
+@test prob.u0 == [1., 1.]
+@variables x(t)
+@mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+@test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
+
+# Test solvers
diff --git a/test/runtests.jl b/test/runtests.jl
index 9537b1b44e..e774bbe4cc 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -80,6 +80,7 @@ end
             @safetestset "Variable Metadata Test" include("test_variable_metadata.jl")
             @safetestset "OptimizationSystem Test" include("optimizationsystem.jl")
             @safetestset "Discrete System" include("discrete_system.jl")
+            @safetestset "Implicit Discrete System" include("implicit_discrete_system.jl")
             @safetestset "SteadyStateSystem Test" include("steadystatesystems.jl")
             @safetestset "SDESystem Test" include("sdesystem.jl")
             @safetestset "DDESystem Test" include("dde.jl")

From a826a6241725db58ca87beb3f252f15d7c658273 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 14 Feb 2025 08:00:39 -0500
Subject: [PATCH 05/16] fix unit

---
 src/structural_transformation/symbolics_tearing.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index ffbd6e3452..a3a90abdfb 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -395,7 +395,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         dx = fullvars[dv]
         order, lv = var_order(dv, diff_to_var)
         x_t = is_discrete ? lower_shift_varname_with_unit(fullvars[dv], iv) : 
-                            Symbolics.diff2term(fullvars[dv])
+                            lower_varname_with_unit(fullvars[lv], iv, order)
 
         # Add `x_t` to the graph
         v_t = add_dd_variable!(structure, fullvars, x_t, dv)

From b58993313f0b1178b715ba1e9195a871cfba3a60 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 14 Feb 2025 08:02:01 -0500
Subject: [PATCH 06/16] add doc file

---
 docs/src/systems/ImplicitDiscreteSystem.md | 28 ++++++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 docs/src/systems/ImplicitDiscreteSystem.md

diff --git a/docs/src/systems/ImplicitDiscreteSystem.md b/docs/src/systems/ImplicitDiscreteSystem.md
new file mode 100644
index 0000000000..d02ccc5e42
--- /dev/null
+++ b/docs/src/systems/ImplicitDiscreteSystem.md
@@ -0,0 +1,28 @@
+# ImplicitDiscreteSystem
+
+## System Constructors
+
+```@docs
+ImplicitDiscreteSystem
+```
+
+## Composition and Accessor Functions
+
+  - `get_eqs(sys)` or `equations(sys)`: The equations that define the implicit discrete system.
+  - `get_unknowns(sys)` or `unknowns(sys)`: The set of unknowns in the implicit discrete system.
+  - `get_ps(sys)` or `parameters(sys)`: The parameters of the implicit discrete system.
+  - `get_iv(sys)`: The independent variable of the implicit discrete system
+  - `discrete_events(sys)`: The set of discrete events in the implicit discrete system.
+
+## Transformations
+
+```@docs; canonical=false
+structural_simplify
+```
+
+## Problem Constructors
+
+```@docs; canonical=false
+ImplicitDiscreteProblem(sys::ImplicitDiscreteSystem, u0map, tspan)
+ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...)
+```

From f5770844623e89d82d40e64b6489d92188d4b290 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 14 Feb 2025 08:11:58 -0500
Subject: [PATCH 07/16] add to docs/pages.jl

---
 docs/pages.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/pages.jl b/docs/pages.jl
index 6e18f403dd..0b2c867507 100644
--- a/docs/pages.jl
+++ b/docs/pages.jl
@@ -42,7 +42,8 @@ pages = [
         "systems/NonlinearSystem.md",
         "systems/OptimizationSystem.md",
         "systems/PDESystem.md",
-        "systems/DiscreteSystem.md"],
+        "systems/DiscreteSystem.md",
+        "systems/ImplicitDiscreteSystem.md"],
     "comparison.md",
     "internals.md"
 ]

From 6fb59de1cbec6ba8a2662202b3f5a8135efb3b9c Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 25 Feb 2025 19:25:08 -0500
Subject: [PATCH 08/16] feat: implement

---
 src/structural_transformation/utils.jl        | 35 +++++++++++++++++++
 .../implicit_discrete_system.jl               |  2 +-
 src/systems/systemstructure.jl                |  2 +-
 test/implicit_discrete_system.jl              | 15 +++++++-
 test/structural_transformation/utils.jl       | 20 +++++++++++
 5 files changed, 71 insertions(+), 3 deletions(-)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index db79c8a124..fb78d707d9 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -507,3 +507,38 @@ function simplify_shifts(var)
             unwrap(var).metadata)
     end
 end
+
+function distribute_shift(var) 
+    var = unwrap(var)
+    var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
+
+    ModelingToolkit.hasshift(var) || return var
+    shift = operation(var)
+    shift isa Shift || return var
+
+    shift = operation(var)
+    expr = only(arguments(var))
+    if expr isa Equation
+        return distribute_shift(shift(expr.lhs)) ~ distribute_shift(shift(expr.rhs))
+    end
+    shiftexpr = _distribute_shift(expr, shift)
+    return simplify_shifts(shiftexpr)
+end
+
+function _distribute_shift(expr, shift)
+    if iscall(expr)
+        op = operation(expr)
+        args = arguments(expr)
+
+        if ModelingToolkit.isvariable(expr)
+            (length(args) == 1 && isequal(shift.t, only(args))) ? (return shift(expr)) : (return expr)
+        elseif op isa Shift
+            return shift(expr)
+        else
+            return maketerm(typeof(expr), operation(expr), Base.Fix2(_distribute_shift, shift).(args),
+                unwrap(expr).metadata)
+        end
+    else
+        return expr
+    end
+end
diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index e0450a1ada..5ff9631b3e 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -265,7 +265,7 @@ function generate_function(
         sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
     iv = get_iv(sys)
     exprs = map(equations(sys)) do eq
-        _iszero(eq.lhs) ? eq.rhs : (simplify_shifts(Shift(iv, -1)(eq.rhs)) - simplify_shifts(Shift(iv, -1)(eq.lhs)))
+        _iszero(eq.lhs) ? eq.rhs : (distribute_shift(Shift(iv, -1)(eq.rhs)) - distribute_shift(Shift(iv, -1)(eq.lhs)))
     end
 
     u_next = dvs
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index a6e6b1d218..9bd3d6c8ab 100644
--- a/src/systems/systemstructure.jl
+++ b/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 DiscreteSystem
+    if sys isa AbstractDiscreteSystem
         ts = shift_discrete_system(ts)
     end
     return ts
diff --git a/test/implicit_discrete_system.jl b/test/implicit_discrete_system.jl
index 88e79acc77..ffe4ce500f 100644
--- a/test/implicit_discrete_system.jl
+++ b/test/implicit_discrete_system.jl
@@ -18,7 +18,6 @@ resid = rand(2)
 f(resid, u_next, [2.,3.], [], t)
 @test resid ≈ [3., -3.]
 
-# Initialization cases.
 prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
 @test prob.u0 == [3., 1.]
 prob = ImplicitDiscreteProblem(sys, [], tspan)
@@ -28,3 +27,17 @@ prob = ImplicitDiscreteProblem(sys, [], tspan)
 @test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
 
 # Test solvers
+@testset "System with algebraic equations" begin
+    @variables x(t) y(t)
+    eqs = [x(k) ~ x(k-1) + x(k-2), 
+           x^2 ~ 1 - y^2]
+    @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
+end
+
+@testset "System with algebraic equations, implicit difference equations, explicit difference equations" begin
+    @variables x(t) y(t)
+    eqs = [x(k) ~ x(k-1) + x(k-2),
+           y(k) ~ x(k) + x(k-2)*y(k-1)]
+    @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
+end
+
diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl
index 4da3d1e924..b7cc4b720b 100644
--- a/test/structural_transformation/utils.jl
+++ b/test/structural_transformation/utils.jl
@@ -162,3 +162,23 @@ end
     structural_simplify(sys; additional_passes = [pass])
     @test value[] == 1
 end
+
+@testset "Distribute shifts" begin
+    @variables x(t) y(t) z(t)
+    @parameters a b c
+
+    # Expand shifts
+    @test isequal(ST.distribute_shift(Shift(t, -1)(x + y)), Shift(t, -1)(x) + Shift(t, -1)(y))
+
+    expr = a * Shift(t, -2)(x) + Shift(t, 2)(y) + b
+    @test isequal(ST.simplify_shifts(ST.distribute_shift(Shift(t, 2)(expr))),
+                  a*x + Shift(t, 4)(y) + b)
+    @test isequal(ST.distribute_shift(Shift(t, 2)(exp(z))), exp(Shift(t, 2)(z)))
+    @test isequal(ST.distribute_shift(Shift(t, 2)(exp(a) + b)), exp(a) + b)
+
+    expr = a^x - log(b*y) + z*x
+    @test isequal(ST.distribute_shift(Shift(t, -3)(expr)), a^(Shift(t, -3)(x)) - log(b * Shift(t, -3)(y)) + Shift(t, -3)(z)*Shift(t, -3)(x))
+
+    expr = x(k+1) ~ x + x(k-1)
+    @test isequal(ST.distribute_shift(Shift(t, -1)(expr)), x ~ x(k-1) + x(k-2))
+end

From 00a8222960a6f3bcb147c0ae7e43e37c44bda3f7 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 25 Feb 2025 22:55:34 -0500
Subject: [PATCH 09/16] add shift2term for positive shifts

---
 src/discretedomain.jl                         | 13 ++++
 .../StructuralTransformations.jl              |  2 +-
 src/structural_transformation/utils.jl        | 20 ++++--
 test/implicit_discrete_system.jl              | 62 ++++++++++++-------
 4 files changed, 71 insertions(+), 26 deletions(-)

diff --git a/src/discretedomain.jl b/src/discretedomain.jl
index c7f90007c5..508cfcd216 100644
--- a/src/discretedomain.jl
+++ b/src/discretedomain.jl
@@ -70,6 +70,19 @@ Base.literal_pow(f::typeof(^), D::Shift, ::Val{n}) where {n} = Shift(D.t, D.step
 
 hasshift(eq::Equation) = hasshift(eq.lhs) || hasshift(eq.rhs)
 
+"""
+    Next(x)
+
+An alias for Shift(t, 1)(x).
+"""
+Next(x) = Shift(t, 1)(x)
+"""
+    Prev(x)
+
+An alias for Shift(t, -1)(x).
+"""
+Prev(x) = Shift(t, -1)(x)
+
 """
     hasshift(O)
 
diff --git a/src/structural_transformation/StructuralTransformations.jl b/src/structural_transformation/StructuralTransformations.jl
index e6be30891f..2e600e86e4 100644
--- a/src/structural_transformation/StructuralTransformations.jl
+++ b/src/structural_transformation/StructuralTransformations.jl
@@ -65,7 +65,7 @@ export torn_system_jacobian_sparsity
 export full_equations
 export but_ordered_incidence, lowest_order_variable_mask, highest_order_variable_mask
 export computed_highest_diff_variables
-export shift2term, lower_shift_varname, simplify_shifts
+export shift2term, lower_shift_varname, simplify_shifts, distribute_shift
 
 include("utils.jl")
 include("pantelides.jl")
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 31c58e797d..031ef6b2d8 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -457,7 +457,7 @@ Handle renaming variable names for discrete structural simplification. Three cas
 """
 function lower_shift_varname(var, iv)
     op = operation(var)
-    if op isa Shift && op.steps < 0
+    if op isa Shift
         return shift2term(var)
     else
         return Shift(iv, 0)(var, true)
@@ -475,10 +475,14 @@ function shift2term(var)
 
     backshift = is_lowered ? op.steps + ModelingToolkit.getshift(arg) : op.steps
 
-    num = join(Char(0x2080 + d) for d in reverse!(digits(-backshift))) # subscripted number, e.g. ₁
-    ds = join([Char(0x209c), Char(0x208b), num])
-    # Char(0x209c) = ₜ
     # Char(0x208b) = ₋ (subscripted minus)
+    # Char(0x208a) = ₊ (subscripted plus)
+    pm = backshift > 0 ? Char(0x208a) : Char(0x208b)
+    # subscripted number, e.g. ₁
+    num = join(Char(0x2080 + d) for d in reverse!(digits(abs(backshift))))
+    # Char(0x209c) = ₜ
+    # ds = ₜ₋₁
+    ds = join([Char(0x209c), pm, num])
 
     O = is_lowered ? ModelingToolkit.getunshifted(arg) : arg
     oldop = operation(O)
@@ -498,6 +502,9 @@ function isdoubleshift(var)
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
 end
 
+"""
+Simplify multiple shifts: Shift(t, k1)(Shift(t, k2)(x)) becomes Shift(t, k1+k2)(x).
+"""
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
     var isa Equation && return simplify_shifts(var.lhs) ~ simplify_shifts(var.rhs)
@@ -518,6 +525,11 @@ function simplify_shifts(var)
     end
 end
 
+"""
+Distribute a shift applied to a whole expression or equation. 
+Shift(t, 1)(x + y) will become Shift(t, 1)(x) + Shift(t, 1)(y).
+Only shifts variables whose independent variable is the same t that appears in the Shift (i.e. constants, time-independent parameters, etc. do not get shifted).
+"""
 function distribute_shift(var) 
     var = unwrap(var)
     var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
diff --git a/test/implicit_discrete_system.jl b/test/implicit_discrete_system.jl
index ffe4ce500f..f2e7deaa32 100644
--- a/test/implicit_discrete_system.jl
+++ b/test/implicit_discrete_system.jl
@@ -1,30 +1,34 @@
 using ModelingToolkit, Test
 using ModelingToolkit: t_nounits as t
+using StableRNGs
 
 k = ShiftIndex(t)
-@variables x(t) = 1
-@mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
-tspan = (0, 10)
+rng = StableRNG(22525) 
 
 # Shift(t, -1)(x(t)) - x_{t-1}(t)
 # -3 - x(t) + x(t)*x_{t-1}
-f = ImplicitDiscreteFunction(sys)
-u_next = [3., 1.5]
-@test f(u_next, [2.,3.], [], t) ≈ [0., 0.]
-u_next = [0., 0.]
-@test f(u_next, [2.,3.], [], t) ≈ [3., -3.]
-
-resid = rand(2)
-f(resid, u_next, [2.,3.], [], t)
-@test resid ≈ [3., -3.]
-
-prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
-@test prob.u0 == [3., 1.]
-prob = ImplicitDiscreteProblem(sys, [], tspan)
-@test prob.u0 == [1., 1.]
-@variables x(t)
-@mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
-@test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
+@testset "Correct ImplicitDiscreteFunction" begin
+    @variables x(t) = 1
+    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+    tspan = (0, 10)
+    f = ImplicitDiscreteFunction(sys)
+    u_next = [3., 1.5]
+    @test f(u_next, [2.,3.], [], t) ≈ [0., 0.]
+    u_next = [0., 0.]
+    @test f(u_next, [2.,3.], [], t) ≈ [3., -3.]
+    
+    resid = rand(2)
+    f(resid, u_next, [2.,3.], [], t)
+    @test resid ≈ [3., -3.]
+    
+    prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
+    @test prob.u0 == [3., 1.]
+    prob = ImplicitDiscreteProblem(sys, [], tspan)
+    @test prob.u0 == [1., 1.]
+    @variables x(t)
+    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+    @test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
+end
 
 # Test solvers
 @testset "System with algebraic equations" begin
@@ -32,6 +36,23 @@ prob = ImplicitDiscreteProblem(sys, [], tspan)
     eqs = [x(k) ~ x(k-1) + x(k-2), 
            x^2 ~ 1 - y^2]
     @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
+    f = ImplicitDiscreteFunction(sys)
+
+    function correct_f(u_next, u, p, t) 
+        [u[2] - u_next[1],
+         u[1] + u[2] - u_next[2],
+         1 - (u_next[1]+u_next[2])^2 - u_next[3]^2]
+    end
+
+    for _ in 1:10
+        u_next = rand(rng, 3)
+        u = rand(rng, 3)
+        @test correct_f(u_next, u, [], 0.) ≈ f(u_next, u, [], 0.)
+    end
+
+    # Initialization is satisfied.
+    prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
+    @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
@@ -40,4 +61,3 @@ end
            y(k) ~ x(k) + x(k-2)*y(k-1)]
     @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
 end
-

From c504aefa92bf974e24d3f1d4625a65ae25070157 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Wed, 26 Feb 2025 16:06:10 -0500
Subject: [PATCH 10/16] fix: make the codegen work with shifted observables

---
 .../symbolics_tearing.jl                      |  4 +++
 .../implicit_discrete_system.jl               | 32 ++++++++++++++-----
 src/systems/systemstructure.jl                |  2 +-
 test/implicit_discrete_system.jl              | 25 +++++++++++----
 4 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 9818bba361..f1cdd7ce9c 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/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])
diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index 5ff9631b3e..5f13e87be4 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/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
 
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index e7d5fa71f1..e5909ed892 100644
--- a/src/systems/systemstructure.jl
+++ b/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
diff --git a/test/implicit_discrete_system.jl b/test/implicit_discrete_system.jl
index f2e7deaa32..f298f43e7c 100644
--- a/test/implicit_discrete_system.jl
+++ b/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

From fc538244f1a9fff8afd3155646dbd33b4774902e Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 28 Feb 2025 10:17:48 -0500
Subject: [PATCH 11/16] add Prev and Next to docs

---
 docs/src/systems/DiscreteSystem.md         |  7 +++++++
 docs/src/systems/ImplicitDiscreteSystem.md |  7 +++++++
 docs/src/tutorials/SampledData.md          | 13 +++++++++++--
 src/ModelingToolkit.jl                     |  2 +-
 4 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/docs/src/systems/DiscreteSystem.md b/docs/src/systems/DiscreteSystem.md
index 5ede50c62a..1b2e6e81c5 100644
--- a/docs/src/systems/DiscreteSystem.md
+++ b/docs/src/systems/DiscreteSystem.md
@@ -26,3 +26,10 @@ structural_simplify
 DiscreteProblem(sys::DiscreteSystem, u0map, tspan)
 DiscreteFunction(sys::DiscreteSystem, args...)
 ```
+
+## Discrete Domain
+```@docs; canonical=false
+Shift
+Prev
+Next
+```
diff --git a/docs/src/systems/ImplicitDiscreteSystem.md b/docs/src/systems/ImplicitDiscreteSystem.md
index d02ccc5e42..189f4fe65c 100644
--- a/docs/src/systems/ImplicitDiscreteSystem.md
+++ b/docs/src/systems/ImplicitDiscreteSystem.md
@@ -26,3 +26,10 @@ structural_simplify
 ImplicitDiscreteProblem(sys::ImplicitDiscreteSystem, u0map, tspan)
 ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...)
 ```
+
+## Discrete Domain
+```@docs; canonical=false
+Shift
+Prev
+Next
+```
diff --git a/docs/src/tutorials/SampledData.md b/docs/src/tutorials/SampledData.md
index a72fd1698b..c700bae5c2 100644
--- a/docs/src/tutorials/SampledData.md
+++ b/docs/src/tutorials/SampledData.md
@@ -25,8 +25,10 @@ The operators [`Sample`](@ref) and [`Hold`](@ref) are thus providing the interfa
 The [`ShiftIndex`](@ref) operator is used to refer to past and future values of discrete-time variables. The example below illustrates its use, implementing the discrete-time system
 
 ```math
-x(k+1) = 0.5x(k) + u(k)
-y(k) = x(k)
+\begin{align}
+    x(k+1) &= 0.5x(k) + u(k) \\ 
+    y(k) &= x(k)
+\end{align}
 ```
 
 ```@example clocks
@@ -187,3 +189,10 @@ connections = [r ~ sin(t)          # reference signal
 
 @named cl = ODESystem(connections, t, systems = [f, c, p])
 ```
+
+```@docs
+Sample
+Hold
+ShiftIndex
+Clock
+```
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 8adabb238f..11e0c0a390 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -301,7 +301,7 @@ export debug_system
 #export ContinuousClock, Discrete, sampletime, input_timedomain, output_timedomain
 #export has_discrete_domain, has_continuous_domain
 #export is_discrete_domain, is_continuous_domain, is_hybrid_domain
-export Sample, Hold, Shift, ShiftIndex, sampletime, SampleTime
+export Sample, Hold, Shift, ShiftIndex, sampletime, SampleTime, Next, Prev
 export Clock, SolverStepClock, TimeDomain
 
 export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunables

From 1e2e6672f8c555070202e7e81018e27ecca22460 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 28 Feb 2025 10:23:38 -0500
Subject: [PATCH 12/16] format

---
 docs/src/systems/DiscreteSystem.md            |  1 +
 docs/src/systems/ImplicitDiscreteSystem.md    |  1 +
 src/ModelingToolkit.jl                        |  3 +-
 src/structural_transformation/utils.jl        |  8 ++-
 .../implicit_discrete_system.jl               | 25 +++++---
 test/implicit_discrete_system.jl              | 57 ++++++++++---------
 test/structural_transformation/utils.jl       | 14 +++--
 7 files changed, 62 insertions(+), 47 deletions(-)

diff --git a/docs/src/systems/DiscreteSystem.md b/docs/src/systems/DiscreteSystem.md
index 1b2e6e81c5..c00cc7f3c9 100644
--- a/docs/src/systems/DiscreteSystem.md
+++ b/docs/src/systems/DiscreteSystem.md
@@ -28,6 +28,7 @@ DiscreteFunction(sys::DiscreteSystem, args...)
 ```
 
 ## Discrete Domain
+
 ```@docs; canonical=false
 Shift
 Prev
diff --git a/docs/src/systems/ImplicitDiscreteSystem.md b/docs/src/systems/ImplicitDiscreteSystem.md
index 189f4fe65c..268d8ec13c 100644
--- a/docs/src/systems/ImplicitDiscreteSystem.md
+++ b/docs/src/systems/ImplicitDiscreteSystem.md
@@ -28,6 +28,7 @@ ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...)
 ```
 
 ## Discrete Domain
+
 ```@docs; canonical=false
 Shift
 Prev
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 11e0c0a390..328e5b6df0 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -237,7 +237,8 @@ export DAEFunctionExpr, DAEProblemExpr
 export SDESystem, SDEFunction, SDEFunctionExpr, SDEProblemExpr
 export SystemStructure
 export DiscreteSystem, DiscreteProblem, DiscreteFunction, DiscreteFunctionExpr
-export ImplicitDiscreteSystem, ImplicitDiscreteProblem, ImplicitDiscreteFunction, ImplicitDiscreteFunctionExpr
+export ImplicitDiscreteSystem, ImplicitDiscreteProblem, ImplicitDiscreteFunction,
+       ImplicitDiscreteFunctionExpr
 export JumpSystem
 export ODEProblem, SDEProblem
 export NonlinearFunction, NonlinearFunctionExpr
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 031ef6b2d8..14628f2958 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -530,7 +530,7 @@ Distribute a shift applied to a whole expression or equation.
 Shift(t, 1)(x + y) will become Shift(t, 1)(x) + Shift(t, 1)(y).
 Only shifts variables whose independent variable is the same t that appears in the Shift (i.e. constants, time-independent parameters, etc. do not get shifted).
 """
-function distribute_shift(var) 
+function distribute_shift(var)
     var = unwrap(var)
     var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
 
@@ -553,11 +553,13 @@ function _distribute_shift(expr, shift)
         args = arguments(expr)
 
         if ModelingToolkit.isvariable(expr)
-            (length(args) == 1 && isequal(shift.t, only(args))) ? (return shift(expr)) : (return expr)
+            (length(args) == 1 && isequal(shift.t, only(args))) ? (return shift(expr)) :
+            (return expr)
         elseif op isa Shift
             return shift(expr)
         else
-            return maketerm(typeof(expr), operation(expr), Base.Fix2(_distribute_shift, shift).(args),
+            return maketerm(
+                typeof(expr), operation(expr), Base.Fix2(_distribute_shift, shift).(args),
                 unwrap(expr).metadata)
         end
     else
diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index bf4e6037d1..cbc2555c36 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -239,7 +239,9 @@ function ImplicitDiscreteSystem(eqs, iv; kwargs...)
         collect(allunknowns), collect(new_ps); kwargs...)
 end
 
-ImplicitDiscreteSystem(eq::Equation, args...; kwargs...) = ImplicitDiscreteSystem([eq], args...; kwargs...)
+function ImplicitDiscreteSystem(eq::Equation, args...; kwargs...)
+    ImplicitDiscreteSystem([eq], args...; kwargs...)
+end
 
 function flatten(sys::ImplicitDiscreteSystem, noeqs = false)
     systems = get_systems(sys)
@@ -273,11 +275,13 @@ function generate_function(
     obs = observed(sys)
     shifted_obs = Symbolics.Equation[distribute_shift(Next(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]
+    extra_assignments = [Assignment(shifted_obs[i].lhs, shifted_obs[i].rhs)
+                         for i in obsidxs]
 
     u_next = map(Next, dvs)
     u = dvs
-    build_function_wrapper(sys, exprs, u_next, u, ps..., iv; p_start = 3, extra_assignments, kwargs...)
+    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)
@@ -341,11 +345,13 @@ function SciMLBase.ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...
     ImplicitDiscreteFunction{true}(sys, args...; kwargs...)
 end
 
-function SciMLBase.ImplicitDiscreteFunction{true}(sys::ImplicitDiscreteSystem, args...; kwargs...)
+function SciMLBase.ImplicitDiscreteFunction{true}(
+        sys::ImplicitDiscreteSystem, args...; kwargs...)
     ImplicitDiscreteFunction{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
 end
 
-function SciMLBase.ImplicitDiscreteFunction{false}(sys::ImplicitDiscreteSystem, args...; kwargs...)
+function SciMLBase.ImplicitDiscreteFunction{false}(
+        sys::ImplicitDiscreteSystem, args...; kwargs...)
     ImplicitDiscreteFunction{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
 end
 function SciMLBase.ImplicitDiscreteFunction{iip, specialize}(
@@ -360,7 +366,6 @@ function SciMLBase.ImplicitDiscreteFunction{iip, specialize}(
         eval_module = @__MODULE__,
         analytic = nothing,
         kwargs...) where {iip, specialize}
-
     if !iscomplete(sys)
         error("A completed `ImplicitDiscreteSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `ImplicitDiscreteProblem`")
     end
@@ -404,9 +409,12 @@ struct ImplicitDiscreteFunctionClosure{O, I} <: Function
     f_iip::I
 end
 (f::ImplicitDiscreteFunctionClosure)(u_next, u, p, t) = f.f_oop(u_next, u, p, t)
-(f::ImplicitDiscreteFunctionClosure)(resid, u_next, u, p, t) = f.f_iip(resid, u_next, u, p, t)
+function (f::ImplicitDiscreteFunctionClosure)(resid, u_next, u, p, t)
+    f.f_iip(resid, u_next, u, p, t)
+end
 
-function ImplicitDiscreteFunctionExpr{iip}(sys::ImplicitDiscreteSystem, dvs = unknowns(sys),
+function ImplicitDiscreteFunctionExpr{iip}(
+        sys::ImplicitDiscreteSystem, dvs = unknowns(sys),
         ps = parameters(sys), u0 = nothing;
         version = nothing, p = nothing,
         linenumbers = false,
@@ -427,4 +435,3 @@ end
 function ImplicitDiscreteFunctionExpr(sys::ImplicitDiscreteSystem, args...; kwargs...)
     ImplicitDiscreteFunctionExpr{true}(sys, args...; kwargs...)
 end
-
diff --git a/test/implicit_discrete_system.jl b/test/implicit_discrete_system.jl
index f298f43e7c..70cc0b792e 100644
--- a/test/implicit_discrete_system.jl
+++ b/test/implicit_discrete_system.jl
@@ -3,72 +3,73 @@ using ModelingToolkit: t_nounits as t
 using StableRNGs
 
 k = ShiftIndex(t)
-rng = StableRNG(22525) 
+rng = StableRNG(22525)
 
 @testset "Correct ImplicitDiscreteFunction" begin
     @variables x(t) = 1
-    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+    @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.]
-    u_next = [0., 0.]
-    @test f(u_next, [2.,3.], [], t) ≈ [3., -3.]
-    
+    u_next = [3.0, 1.5]
+    @test f(u_next, [2.0, 3.0], [], t) ≈ [0.0, 0.0]
+    u_next = [0.0, 0.0]
+    @test f(u_next, [2.0, 3.0], [], t) ≈ [3.0, -3.0]
+
     resid = rand(2)
-    f(resid, u_next, [2.,3.], [], t)
-    @test resid ≈ [3., -3.]
-    
-    prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
-    @test prob.u0 == [3., 1.]
+    f(resid, u_next, [2.0, 3.0], [], t)
+    @test resid ≈ [3.0, -3.0]
+
+    prob = ImplicitDiscreteProblem(sys, [x(k - 1) => 3.0], tspan)
+    @test prob.u0 == [3.0, 1.0]
     prob = ImplicitDiscreteProblem(sys, [], tspan)
-    @test prob.u0 == [1., 1.]
+    @test prob.u0 == [1.0, 1.0]
     @variables x(t)
-    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
-    @test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
+    @mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k) * x(k - 1) - 3], t)
+    @test_throws ErrorException prob=ImplicitDiscreteProblem(sys, [], tspan)
 end
 
 @testset "System with algebraic equations" begin
     @variables x(t) y(t)
-    eqs = [x(k) ~ x(k-1) + x(k-2), 
-           x^2 ~ 1 - y^2]
+    eqs = [x(k) ~ x(k - 1) + x(k - 2),
+        x^2 ~ 1 - y^2]
     @mtkbuild sys = ImplicitDiscreteSystem(eqs, t)
     f = ImplicitDiscreteFunction(sys)
 
-    function correct_f(u_next, u, p, t) 
+    function correct_f(u_next, u, p, t)
         [u[2] - u_next[1],
-         u[1] + u[2] - u_next[2],
-         1 - (u_next[1]+u_next[2])^2 - u_next[3]^2]
+            u[1] + u[2] - u_next[2],
+            1 - (u_next[1] + u_next[2])^2 - u_next[3]^2]
     end
 
     for _ in 1:10
         u_next = rand(rng, 3)
         u = rand(rng, 3)
-        @test correct_f(u_next, u, [], 0.) ≈ f(u_next, u, [], 0.)
+        @test correct_f(u_next, u, [], 0.0) ≈ f(u_next, u, [], 0.0)
     end
 
     # Initialization is satisfied.
-    prob = ImplicitDiscreteProblem(sys, [x(k-1) => .3, x(k-2) => .4], (0, 10), guesses = [y => 1])
+    prob = ImplicitDiscreteProblem(
+        sys, [x(k - 1) => 0.3, x(k - 2) => 0.4], (0, 10), guesses = [y => 1])
     @test (prob.u0[1] + prob.u0[2])^2 + prob.u0[3]^2 ≈ 1
 end
 
 @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)*z(k-1), 
-           x + y + z ~ 2]
+    eqs = [x(k) ~ x(k - 1) + x(k - 2),
+        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]
+    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
diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl
index 0013fb4bbe..646b9ce74e 100644
--- a/test/structural_transformation/utils.jl
+++ b/test/structural_transformation/utils.jl
@@ -169,19 +169,21 @@ end
     k = ShiftIndex(t)
 
     # Expand shifts
-    @test isequal(ST.distribute_shift(Shift(t, -1)(x + y)), Shift(t, -1)(x) + Shift(t, -1)(y))
+    @test isequal(
+        ST.distribute_shift(Shift(t, -1)(x + y)), Shift(t, -1)(x) + Shift(t, -1)(y))
 
     expr = a * Shift(t, -2)(x) + Shift(t, 2)(y) + b
     @test isequal(ST.simplify_shifts(ST.distribute_shift(Shift(t, 2)(expr))),
-                  a*x + Shift(t, 4)(y) + b)
+        a * x + Shift(t, 4)(y) + b)
     @test isequal(ST.distribute_shift(Shift(t, 2)(exp(z))), exp(Shift(t, 2)(z)))
     @test isequal(ST.distribute_shift(Shift(t, 2)(exp(a) + b)), exp(a) + b)
 
-    expr = a^x - log(b*y) + z*x
-    @test isequal(ST.distribute_shift(Shift(t, -3)(expr)), a^(Shift(t, -3)(x)) - log(b * Shift(t, -3)(y)) + Shift(t, -3)(z)*Shift(t, -3)(x))
+    expr = a^x - log(b * y) + z * x
+    @test isequal(ST.distribute_shift(Shift(t, -3)(expr)),
+        a^(Shift(t, -3)(x)) - log(b * Shift(t, -3)(y)) + Shift(t, -3)(z) * Shift(t, -3)(x))
 
-    expr = x(k+1) ~ x + x(k-1)
-    @test isequal(ST.distribute_shift(Shift(t, -1)(expr)), x ~ x(k-1) + x(k-2))
+    expr = x(k + 1) ~ x + x(k - 1)
+    @test isequal(ST.distribute_shift(Shift(t, -1)(expr)), x ~ x(k - 1) + x(k - 2))
 end
 
 @testset "`map_variables_to_equations`" begin

From 506522c7f3b8ca2fafe4d205b70542301f26b26d Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 28 Feb 2025 11:01:27 -0500
Subject: [PATCH 13/16] fix missing var

---
 test/structural_transformation/utils.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl
index 646b9ce74e..9bf239a574 100644
--- a/test/structural_transformation/utils.jl
+++ b/test/structural_transformation/utils.jl
@@ -5,6 +5,7 @@ using SparseArrays
 using UnPack
 using ModelingToolkit: t_nounits as t, D_nounits as D, default_toterm
 using Symbolics: unwrap
+const ST = StructuralTransformations
 
 # Define some variables
 @parameters L g

From af3995c842cca7d342a8a642797795ea83f859d4 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 28 Feb 2025 11:45:02 -0500
Subject: [PATCH 14/16] make initialization_data for IMplicitDIscreteFunction,
 fix bad initialization test

---
 src/systems/discrete_system/implicit_discrete_system.jl | 5 ++++-
 test/implicit_discrete_system.jl                        | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index cbc2555c36..ae5524768b 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -334,8 +334,10 @@ function SciMLBase.ImplicitDiscreteProblem(
 
     u0map = to_varmap(u0map, dvs)
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
+    @show u0map
     f, u0, p = process_SciMLProblem(
         ImplicitDiscreteFunction, sys, u0map, parammap; eval_expression, eval_module, kwargs...)
+    @show u0
 
     kwargs = filter_kwargs(kwargs)
     ImplicitDiscreteProblem(f, u0, tspan, p; kwargs...)
@@ -388,7 +390,8 @@ function SciMLBase.ImplicitDiscreteFunction{iip, specialize}(
     ImplicitDiscreteFunction{iip, specialize}(f;
         sys = sys,
         observed = observedfun,
-        analytic = analytic)
+        analytic = analytic,
+        kwargs...)
 end
 
 """
diff --git a/test/implicit_discrete_system.jl b/test/implicit_discrete_system.jl
index 70cc0b792e..932b6c6981 100644
--- a/test/implicit_discrete_system.jl
+++ b/test/implicit_discrete_system.jl
@@ -53,7 +53,7 @@ end
     # Initialization is satisfied.
     prob = ImplicitDiscreteProblem(
         sys, [x(k - 1) => 0.3, x(k - 2) => 0.4], (0, 10), guesses = [y => 1])
-    @test (prob.u0[1] + prob.u0[2])^2 + prob.u0[3]^2 ≈ 1
+    @test length(equations(prob.f.initialization_data.initializeprob.f.sys)) == 1
 end
 
 @testset "Handle observables in function codegen" begin

From a98397ca205309603e87f2c12433ffd49cfd0296 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 10 Mar 2025 15:15:38 -0400
Subject: [PATCH 15/16] remove Next and Prev

---
 src/discretedomain.jl                               | 13 -------------
 .../discrete_system/implicit_discrete_system.jl     |  6 +++---
 2 files changed, 3 insertions(+), 16 deletions(-)

diff --git a/src/discretedomain.jl b/src/discretedomain.jl
index 0d8d38e15d..7260237053 100644
--- a/src/discretedomain.jl
+++ b/src/discretedomain.jl
@@ -70,19 +70,6 @@ Base.literal_pow(f::typeof(^), D::Shift, ::Val{n}) where {n} = Shift(D.t, D.step
 
 hasshift(eq::Equation) = hasshift(eq.lhs) || hasshift(eq.rhs)
 
-"""
-    Next(x)
-
-An alias for Shift(t, 1)(x).
-"""
-Next(x) = Shift(t, 1)(x)
-"""
-    Prev(x)
-
-An alias for Shift(t, -1)(x).
-"""
-Prev(x) = Shift(t, -1)(x)
-
 """
     hasshift(O)
 
diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index ae5524768b..327c82c47f 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -268,17 +268,17 @@ function generate_function(
     iv = get_iv(sys)
     # Algebraic equations get shifted forward 1, to match with differential equations
     exprs = map(equations(sys)) do eq
-        _iszero(eq.lhs) ? distribute_shift(Next(eq.rhs)) : (eq.rhs - eq.lhs)
+        _iszero(eq.lhs) ? distribute_shift(Shift(iv, 1)(eq.rhs)) : (eq.rhs - eq.lhs)
     end
 
     # Handle observables in algebraic equations, since they are shifted
     obs = observed(sys)
-    shifted_obs = Symbolics.Equation[distribute_shift(Next(eq)) for eq in obs]
+    shifted_obs = Symbolics.Equation[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(Next, dvs)
+    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...)

From d40fd6e69f21919d0984f61bf990100653121b99 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 10 Mar 2025 15:16:42 -0400
Subject: [PATCH 16/16] remove Next and Prev docs

---
 docs/src/systems/DiscreteSystem.md         | 2 --
 docs/src/systems/ImplicitDiscreteSystem.md | 2 --
 src/doc                                    | 1 +
 3 files changed, 1 insertion(+), 4 deletions(-)
 create mode 100644 src/doc

diff --git a/docs/src/systems/DiscreteSystem.md b/docs/src/systems/DiscreteSystem.md
index c00cc7f3c9..f8a71043ab 100644
--- a/docs/src/systems/DiscreteSystem.md
+++ b/docs/src/systems/DiscreteSystem.md
@@ -31,6 +31,4 @@ DiscreteFunction(sys::DiscreteSystem, args...)
 
 ```@docs; canonical=false
 Shift
-Prev
-Next
 ```
diff --git a/docs/src/systems/ImplicitDiscreteSystem.md b/docs/src/systems/ImplicitDiscreteSystem.md
index 268d8ec13c..d69f88f106 100644
--- a/docs/src/systems/ImplicitDiscreteSystem.md
+++ b/docs/src/systems/ImplicitDiscreteSystem.md
@@ -31,6 +31,4 @@ ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...)
 
 ```@docs; canonical=false
 Shift
-Prev
-Next
 ```
diff --git a/src/doc b/src/doc
new file mode 100644
index 0000000000..8b13789179
--- /dev/null
+++ b/src/doc
@@ -0,0 +1 @@
+