From 76e515c4c55af8b9ca3bf8a3318d1dd3339702a3 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 23 Jan 2025 12:33:15 -0500
Subject: [PATCH 01/41] Init

---
 .../discrete_system/implicitdiscretesystem.jl | 406 ++++++++++++++++++
 1 file changed, 406 insertions(+)
 create mode 100644 src/systems/discrete_system/implicitdiscretesystem.jl

diff --git a/src/systems/discrete_system/implicitdiscretesystem.jl b/src/systems/discrete_system/implicitdiscretesystem.jl
new file mode 100644
index 0000000000..404fa7ff49
--- /dev/null
+++ b/src/systems/discrete_system/implicitdiscretesystem.jl
@@ -0,0 +1,406 @@
+"""
+$(TYPEDEF)
+An implicit system of difference equations.
+# Fields
+$(FIELDS)
+# Example
+```
+using ModelingToolkit
+using ModelingToolkit: t_nounits as t
+@parameters σ=28.0 ρ=10.0 β=8/3 δt=0.1
+@variables x(t)=1.0 y(t)=0.0 z(t)=0.0
+k = ShiftIndex(t)
+eqs = [x(k+1) ~ σ*(y-x),
+       y(k+1) ~ x*(ρ-z)-y,
+       z(k+1) ~ x*y - β*z]
+@named de = ImplicitDiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β]; tspan = (0, 1000.0)) # or
+@named de = ImplicitDiscreteSystem(eqs)
+```
+"""
+struct ImplicitDiscreteSystem <: AbstractTimeDependentSystem
+    """
+    A tag for the system. If two systems have the same tag, then they are
+    structurally identical.
+    """
+    tag::UInt
+    """The differential equations defining the discrete system."""
+    eqs::Vector{Equation}
+    """Independent variable."""
+    iv::BasicSymbolic{Real}
+    """Dependent (state) variables. Must not contain the independent variable."""
+    unknowns::Vector
+    """Parameter variables. Must not contain the independent variable."""
+    ps::Vector
+    """Time span."""
+    tspan::Union{NTuple{2, Any}, Nothing}
+    """Array variables."""
+    var_to_name::Any
+    """Observed states."""
+    observed::Vector{Equation}
+    """
+    The name of the system
+    """
+    name::Symbol
+    """
+    A description of the system.
+    """
+    description::String
+    """
+    The internal systems. These are required to have unique names.
+    """
+    systems::Vector{DiscreteSystem}
+    """
+    The default values to use when initial conditions and/or
+    parameters are not supplied in `DiscreteProblem`.
+    """
+    defaults::Dict
+    """
+    The guesses to use as the initial conditions for the
+    initialization system.
+    """
+    guesses::Dict
+    """
+    The system for performing the initialization.
+    """
+    initializesystem::Union{Nothing, NonlinearSystem}
+    """
+    Extra equations to be enforced during the initialization sequence.
+    """
+    initialization_eqs::Vector{Equation}
+    """
+    Inject assignment statements before the evaluation of the RHS function.
+    """
+    preface::Any
+    """
+    Type of the system.
+    """
+    connector_type::Any
+    """
+    Topologically sorted parameter dependency equations, where all symbols are parameters and
+    the LHS is a single parameter.
+    """
+    parameter_dependencies::Vector{Equation}
+    """
+    Metadata for the system, to be used by downstream packages.
+    """
+    metadata::Any
+    """
+    Metadata for MTK GUI.
+    """
+    gui_metadata::Union{Nothing, GUIMetadata}
+    """
+    Cache for intermediate tearing state.
+    """
+    tearing_state::Any
+    """
+    Substitutions generated by tearing.
+    """
+    substitutions::Any
+    """
+    If a model `sys` is complete, then `sys.x` no longer performs namespacing.
+    """
+    complete::Bool
+    """
+    Cached data for fast symbolic indexing.
+    """
+    index_cache::Union{Nothing, IndexCache}
+    """
+    The hierarchical parent system before simplification.
+    """
+    parent::Any
+    isscheduled::Bool
+
+    function ImplicitDiscreteSystem(tag, discreteEqs, iv, dvs, ps, tspan, var_to_name,
+            observed, name, description, systems, defaults, guesses, initializesystem,
+            initialization_eqs, preface, connector_type, parameter_dependencies = Equation[],
+            metadata = nothing, gui_metadata = nothing,
+            tearing_state = nothing, substitutions = nothing,
+            complete = false, index_cache = nothing, parent = nothing,
+            isscheduled = false;
+            checks::Union{Bool, Int} = true)
+        if checks == true || (checks & CheckComponents) > 0
+            check_independent_variables([iv])
+            check_variables(dvs, iv)
+            check_parameters(ps, iv)
+        end
+        if checks == true || (checks & CheckUnits) > 0
+            u = __get_unit_type(dvs, ps, iv)
+            check_units(u, discreteEqs)
+        end
+        new(tag, discreteEqs, iv, dvs, ps, tspan, var_to_name, observed, name, description,
+            systems, defaults, guesses, initializesystem, initialization_eqs,
+            preface, connector_type, parameter_dependencies, metadata, gui_metadata,
+            tearing_state, substitutions, complete, index_cache, parent, isscheduled)
+    end
+end
+
+"""
+    $(TYPEDSIGNATURES)
+Constructs a DiscreteSystem.
+"""
+function ImplicitDiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
+        observed = Num[],
+        systems = DiscreteSystem[],
+        tspan = nothing,
+        name = nothing,
+        description = "",
+        default_u0 = Dict(),
+        default_p = Dict(),
+        guesses = Dict(),
+        initializesystem = nothing,
+        initialization_eqs = Equation[],
+        defaults = _merge(Dict(default_u0), Dict(default_p)),
+        preface = nothing,
+        connector_type = nothing,
+        parameter_dependencies = Equation[],
+        metadata = nothing,
+        gui_metadata = nothing,
+        kwargs...)
+    name === nothing &&
+        throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
+    iv′ = value(iv)
+    dvs′ = value.(dvs)
+    ps′ = value.(ps)
+    if any(hasderiv, eqs) || any(hashold, eqs) || any(hassample, eqs) || any(hasdiff, eqs)
+        error("Equations in a `DiscreteSystem` can only have `Shift` operators.")
+    end
+    if !(isempty(default_u0) && isempty(default_p))
+        Base.depwarn(
+            "`default_u0` and `default_p` are deprecated. Use `defaults` instead.",
+            :DiscreteSystem, force = true)
+    end
+
+    defaults = Dict{Any, Any}(todict(defaults))
+    guesses = Dict{Any, Any}(todict(guesses))
+    var_to_name = Dict()
+    process_variables!(var_to_name, defaults, guesses, dvs′)
+    process_variables!(var_to_name, defaults, guesses, ps′)
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.lhs for eq in parameter_dependencies])
+    process_variables!(
+        var_to_name, defaults, guesses, [eq.rhs for eq in parameter_dependencies])
+    defaults = Dict{Any, Any}(value(k) => value(v)
+    for (k, v) in pairs(defaults) if v !== nothing)
+    guesses = Dict{Any, Any}(value(k) => value(v)
+    for (k, v) in pairs(guesses) if v !== nothing)
+
+    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
+
+    sysnames = nameof.(systems)
+    if length(unique(sysnames)) != length(sysnames)
+        throw(ArgumentError("System names must be unique."))
+    end
+    DiscreteSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
+        eqs, iv′, dvs′, ps′, tspan, var_to_name, observed, name, description, systems,
+        defaults, guesses, initializesystem, initialization_eqs, preface, connector_type,
+        parameter_dependencies, metadata, gui_metadata, kwargs...)
+end
+
+function ImplicitDiscreteSystem(eqs, iv; kwargs...)
+    eqs = collect(eqs)
+    diffvars = OrderedSet()
+    allunknowns = OrderedSet()
+    ps = OrderedSet()
+    iv = value(iv)
+    for eq in eqs
+        collect_vars!(allunknowns, ps, eq, iv; op = Shift)
+        if iscall(eq.lhs) && operation(eq.lhs) isa Shift
+            isequal(iv, operation(eq.lhs).t) ||
+                throw(ArgumentError("A DiscreteSystem can only have one independent variable."))
+            eq.lhs in diffvars &&
+                throw(ArgumentError("The shift variable $(eq.lhs) is not unique in the system of equations."))
+            push!(diffvars, eq.lhs)
+        end
+    end
+    for eq in get(kwargs, :parameter_dependencies, Equation[])
+        if eq isa Pair
+            collect_vars!(allunknowns, ps, eq, iv)
+        else
+            collect_vars!(allunknowns, ps, eq, iv)
+        end
+    end
+    new_ps = OrderedSet()
+    for p in ps
+        if iscall(p) && operation(p) === getindex
+            par = arguments(p)[begin]
+            if Symbolics.shape(Symbolics.unwrap(par)) !== Symbolics.Unknown() &&
+               all(par[i] in ps for i in eachindex(par))
+                push!(new_ps, par)
+            else
+                push!(new_ps, p)
+            end
+        else
+            push!(new_ps, p)
+        end
+    end
+    return DiscreteSystem(eqs, iv,
+        collect(allunknowns), collect(new_ps); kwargs...)
+end
+
+function flatten(sys::DiscreteSystem, noeqs = false)
+    systems = get_systems(sys)
+    if isempty(systems)
+        return sys
+    else
+        return DiscreteSystem(noeqs ? Equation[] : equations(sys),
+            get_iv(sys),
+            unknowns(sys),
+            parameters(sys),
+            observed = observed(sys),
+            defaults = defaults(sys),
+            guesses = guesses(sys),
+            initialization_eqs = initialization_equations(sys),
+            name = nameof(sys),
+            description = description(sys),
+            metadata = get_metadata(sys),
+            checks = false)
+    end
+end
+
+function generate_function(
+        sys::DiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
+    exprs = [eq.rhs for eq in equations(sys)]
+    wrap_code = wrap_code .∘ wrap_array_vars(sys, exprs) .∘
+                wrap_parameter_dependencies(sys, false)
+    generate_custom_function(sys, exprs, dvs, ps; wrap_code, kwargs...)
+end
+
+function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
+    iv = get_iv(sys)
+    updated = AnyDict()
+    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)).")
+        end
+        updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
+    end
+    for var in unknowns(sys)
+        op = operation(var)
+        op isa Shift || continue
+        haskey(updated, var) && continue
+        root = first(arguments(var))
+        haskey(defs, root) || error("Initial condition for $var not provided.")
+        updated[var] = defs[root]
+    end
+    return updated
+end
+
+"""
+    $(TYPEDSIGNATURES)
+Generates an ImplicitDiscreteProblem from an ImplicitDiscreteSystem.
+"""
+function SciMLBase.ImplicitDiscreteProblem(
+        sys::DiscreteSystem, u0map = [], tspan = get_tspan(sys),
+        parammap = SciMLBase.NullParameters();
+        eval_module = @__MODULE__,
+        eval_expression = false,
+        use_union = false,
+        kwargs...
+)
+    if !iscomplete(sys)
+        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)
+    eqs = equations(sys)
+    iv = get_iv(sys)
+
+    u0map = to_varmap(u0map, dvs)
+    u0map = shift_u0map_forward(sys, u0map, defaults(sys))
+    f, u0, p = process_SciMLProblem(
+        DiscreteFunction, sys, u0map, parammap; eval_expression, eval_module)
+    u0 = f(u0, p, tspan[1])
+    ImplicitDiscreteProblem(f, u0, tspan, p; kwargs...)
+end
+
+function SciMLBase.ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...; kwargs...)
+    ImplicitDiscreteFunction{true}(sys, args...; kwargs...)
+end
+
+function SciMLBase.ImplicitDiscreteFunction{true}(sys::ImplicitDiscreteSystem, args...; kwargs...)
+    ImplicitDiscreteFunction{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
+end
+
+function SciMLBase.ImplicitDiscreteFunction{false}(sys::ImplicitDiscreteSystem, args...; kwargs...)
+    ImplicitDiscreteFunction{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
+end
+function SciMLBase.ImplicitDiscreteFunction{iip, specialize}(
+        sys::ImplicitDiscreteSystem,
+        dvs = unknowns(sys),
+        ps = parameters(sys),
+        u0 = nothing;
+        version = nothing,
+        p = nothing,
+        t = nothing,
+        eval_expression = false,
+        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
+    f_gen = generate_function(sys, dvs, ps; expression = Val{true},
+        expression_module = eval_module, kwargs...)
+    f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
+    f(u, p, t) = f_oop(u, p, t)
+    f(du, u, p, t) = f_iip(du, u, p, t)
+
+    if specialize === SciMLBase.FunctionWrapperSpecialize && iip
+        if u0 === nothing || p === nothing || t === nothing
+            error("u0, p, and t must be specified for FunctionWrapperSpecialize on DiscreteFunction.")
+        end
+        f = SciMLBase.wrapfun_iip(f, (u0, u0, p, t))
+    end
+
+    observedfun = ObservedFunctionCache(
+        sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false))
+
+    ImplicitDiscreteFunction{iip, specialize}(f;
+        sys = sys,
+        observed = observedfun,
+        analytic = analytic)
+end
+
+"""
+```julia
+ImplicitDiscreteFunctionExpr{iip}(sys::ImplicitDiscreteSystem, dvs = states(sys),
+                                  ps = parameters(sys);
+                                  version = nothing,
+                                  kwargs...) where {iip}
+```
+
+Create a Julia expression for an `ImplicitDiscreteFunction` from the [`ImplicitDiscreteSystem`](@ref).
+The arguments `dvs` and `ps` are used to set the order of the dependent
+variable and parameter vectors, respectively.
+"""
+struct ImplicitDiscreteFunctionExpr{iip} end
+struct ImplicitDiscreteFunctionClosure{O, I} <: Function
+    f_oop::O
+    f_iip::I
+end
+(f::ImplicitDiscreteFunctionClosure)(u, p, t) = f.f_oop(u, p, t)
+(f::ImplicitDiscreteFunctionClosure)(du, u, p, t) = f.f_iip(du, u, p, t)
+
+function ImplicitDiscreteFunctionExpr{iip}(sys::ImplicitDiscreteSystem, dvs = unknowns(sys),
+        ps = parameters(sys), u0 = nothing;
+        version = nothing, p = nothing,
+        linenumbers = false,
+        simplify = false,
+        kwargs...) where {iip}
+    f_oop, f_iip = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)
+
+    fsym = gensym(:f)
+    _f = :($fsym = $ImplicitDiscreteFunctionClosure($f_oop, $f_iip))
+
+    ex = quote
+        $_f
+        DiscreteFunction{$iip}($fsym)
+    end
+    !linenumbers ? Base.remove_linenums!(ex) : ex
+end
+
+function ImplicitDiscreteFunctionExpr(sys::ImplicitDiscreteSystem, args...; kwargs...)
+    ImplicitDiscreteFunctionExpr{true}(sys, args...; kwargs...)
+end
+

From f1d2a0754ae33524a063d028b6d2515a9cabc06c Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 24 Jan 2025 11:20:11 -0500
Subject: [PATCH 02/41] rename

---
 .../{implicitdiscretesystem.jl => implicit_discrete_system.jl}    | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename src/systems/discrete_system/{implicitdiscretesystem.jl => implicit_discrete_system.jl} (100%)

diff --git a/src/systems/discrete_system/implicitdiscretesystem.jl b/src/systems/discrete_system/implicit_discrete_system.jl
similarity index 100%
rename from src/systems/discrete_system/implicitdiscretesystem.jl
rename to src/systems/discrete_system/implicit_discrete_system.jl

From 2bd65e2702544b1ed561449665115247b65de121 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 24 Jan 2025 16:06:16 -0500
Subject: [PATCH 03/41] up

---
 src/systems/systems.jl | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 04c50bc766..ffbee75b33 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -39,13 +39,13 @@ function structural_simplify(
     else
         newsys = newsys′
     end
-    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.
-        """)
-    end
+    # 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.
+    #     """)
+    # end
     for pass in additional_passes
         newsys = pass(newsys)
     end

From 4142923e7e7e38e90c65952265c02e93fc081472 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Wed, 29 Jan 2025 15:04:39 -0500
Subject: [PATCH 04/41] up

---
 .../implicit_discrete_system.jl               | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index 404fa7ff49..56e7545c6a 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -23,7 +23,7 @@ struct ImplicitDiscreteSystem <: AbstractTimeDependentSystem
     structurally identical.
     """
     tag::UInt
-    """The differential equations defining the discrete system."""
+    """The difference equations defining the discrete system."""
     eqs::Vector{Equation}
     """Independent variable."""
     iv::BasicSymbolic{Real}
@@ -48,10 +48,10 @@ struct ImplicitDiscreteSystem <: AbstractTimeDependentSystem
     """
     The internal systems. These are required to have unique names.
     """
-    systems::Vector{DiscreteSystem}
+    systems::Vector{ImplicitDiscreteSystem}
     """
     The default values to use when initial conditions and/or
-    parameters are not supplied in `DiscreteProblem`.
+    parameters are not supplied in `ImplicitDiscreteProblem`.
     """
     defaults::Dict
     """
@@ -136,11 +136,11 @@ end
 
 """
     $(TYPEDSIGNATURES)
-Constructs a DiscreteSystem.
+Constructs a ImplicitDiscreteSystem.
 """
 function ImplicitDiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
         observed = Num[],
-        systems = DiscreteSystem[],
+        systems = ImplicitDiscreteSystem[],
         tspan = nothing,
         name = nothing,
         description = "",
@@ -162,12 +162,12 @@ function ImplicitDiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
     dvs′ = value.(dvs)
     ps′ = value.(ps)
     if any(hasderiv, eqs) || any(hashold, eqs) || any(hassample, eqs) || any(hasdiff, eqs)
-        error("Equations in a `DiscreteSystem` can only have `Shift` operators.")
+        error("Equations in a `ImplicitDiscreteSystem` can only have `Shift` operators.")
     end
     if !(isempty(default_u0) && isempty(default_p))
         Base.depwarn(
             "`default_u0` and `default_p` are deprecated. Use `defaults` instead.",
-            :DiscreteSystem, force = true)
+            :ImplicitDiscreteSystem, force = true)
     end
 
     defaults = Dict{Any, Any}(todict(defaults))
@@ -190,7 +190,7 @@ function ImplicitDiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    DiscreteSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
+    ImplicitDiscreteSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         eqs, iv′, dvs′, ps′, tspan, var_to_name, observed, name, description, systems,
         defaults, guesses, initializesystem, initialization_eqs, preface, connector_type,
         parameter_dependencies, metadata, gui_metadata, kwargs...)
@@ -206,7 +206,7 @@ function ImplicitDiscreteSystem(eqs, iv; kwargs...)
         collect_vars!(allunknowns, ps, eq, iv; op = Shift)
         if iscall(eq.lhs) && operation(eq.lhs) isa Shift
             isequal(iv, operation(eq.lhs).t) ||
-                throw(ArgumentError("A DiscreteSystem can only have one independent variable."))
+                throw(ArgumentError("An ImplicitDiscreteSystem can only have one independent variable."))
             eq.lhs in diffvars &&
                 throw(ArgumentError("The shift variable $(eq.lhs) is not unique in the system of equations."))
             push!(diffvars, eq.lhs)
@@ -233,16 +233,16 @@ function ImplicitDiscreteSystem(eqs, iv; kwargs...)
             push!(new_ps, p)
         end
     end
-    return DiscreteSystem(eqs, iv,
+    return ImplicitDiscreteSystem(eqs, iv,
         collect(allunknowns), collect(new_ps); kwargs...)
 end
 
-function flatten(sys::DiscreteSystem, noeqs = false)
+function flatten(sys::ImplicitDiscreteSystem, noeqs = false)
     systems = get_systems(sys)
     if isempty(systems)
         return sys
     else
-        return DiscreteSystem(noeqs ? Equation[] : equations(sys),
+        return ImplicitDiscreteSystem(noeqs ? Equation[] : equations(sys),
             get_iv(sys),
             unknowns(sys),
             parameters(sys),
@@ -258,14 +258,14 @@ function flatten(sys::DiscreteSystem, noeqs = false)
 end
 
 function generate_function(
-        sys::DiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
+        sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
     exprs = [eq.rhs for eq in equations(sys)]
     wrap_code = wrap_code .∘ wrap_array_vars(sys, exprs) .∘
                 wrap_parameter_dependencies(sys, false)
     generate_custom_function(sys, exprs, dvs, ps; wrap_code, kwargs...)
 end
 
-function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
+function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
     iv = get_iv(sys)
     updated = AnyDict()
     for k in collect(keys(u0map))
@@ -291,7 +291,7 @@ end
 Generates an ImplicitDiscreteProblem from an ImplicitDiscreteSystem.
 """
 function SciMLBase.ImplicitDiscreteProblem(
-        sys::DiscreteSystem, u0map = [], tspan = get_tspan(sys),
+        sys::ImplicitDiscreteSystem, u0map = [], tspan = get_tspan(sys),
         parammap = SciMLBase.NullParameters();
         eval_module = @__MODULE__,
         eval_expression = false,
@@ -309,7 +309,7 @@ function SciMLBase.ImplicitDiscreteProblem(
     u0map = to_varmap(u0map, dvs)
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
     f, u0, p = process_SciMLProblem(
-        DiscreteFunction, sys, u0map, parammap; eval_expression, eval_module)
+        ImplicitDiscreteFunction, sys, u0map, parammap; eval_expression, eval_module)
     u0 = f(u0, p, tspan[1])
     ImplicitDiscreteProblem(f, u0, tspan, p; kwargs...)
 end
@@ -348,7 +348,7 @@ function SciMLBase.ImplicitDiscreteFunction{iip, specialize}(
 
     if specialize === SciMLBase.FunctionWrapperSpecialize && iip
         if u0 === nothing || p === nothing || t === nothing
-            error("u0, p, and t must be specified for FunctionWrapperSpecialize on DiscreteFunction.")
+            error("u0, p, and t must be specified for FunctionWrapperSpecialize on ImplicitDiscreteFunction.")
         end
         f = SciMLBase.wrapfun_iip(f, (u0, u0, p, t))
     end
@@ -395,7 +395,7 @@ function ImplicitDiscreteFunctionExpr{iip}(sys::ImplicitDiscreteSystem, dvs = un
 
     ex = quote
         $_f
-        DiscreteFunction{$iip}($fsym)
+        ImplicitDiscreteFunction{$iip}($fsym)
     end
     !linenumbers ? Base.remove_linenums!(ex) : ex
 end

From afc468996d65bff484d66327149f2f8ccabfa83e Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 30 Jan 2025 11:03:18 -0500
Subject: [PATCH 05/41] add test file

---
 src/ModelingToolkit.jl                        |  2 +
 .../discrete_system/discrete_system.jl        |  1 +
 .../implicit_discrete_system.jl               | 48 +++++++++++++------
 src/systems/systemstructure.jl                |  4 +-
 test/implicit_discrete_system.jl              |  2 +
 test/runtests.jl                              |  3 +-
 6 files changed, 44 insertions(+), 16 deletions(-)
 create mode 100644 test/implicit_discrete_system.jl

diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 2710d7d1e4..4083f2609d 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -164,6 +164,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")
 
@@ -229,6 +230,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/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index bd5c72eec7..0ad3317ea1 100644
--- a/src/systems/discrete_system/discrete_system.jl
+++ b/src/systems/discrete_system/discrete_system.jl
@@ -233,6 +233,7 @@ function DiscreteSystem(eqs, iv; kwargs...)
             push!(new_ps, p)
         end
     end
+    @show allunknowns
     return DiscreteSystem(eqs, iv,
         collect(allunknowns), collect(new_ps); kwargs...)
 end
diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index 56e7545c6a..1e1b5d5f6f 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -10,11 +10,10 @@ using ModelingToolkit: t_nounits as t
 @parameters σ=28.0 ρ=10.0 β=8/3 δt=0.1
 @variables x(t)=1.0 y(t)=0.0 z(t)=0.0
 k = ShiftIndex(t)
-eqs = [x(k+1) ~ σ*(y-x),
-       y(k+1) ~ x*(ρ-z)-y,
-       z(k+1) ~ x*y - β*z]
-@named de = ImplicitDiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β]; tspan = (0, 1000.0)) # or
-@named de = ImplicitDiscreteSystem(eqs)
+eqs = [x ~ σ*(y(k-1)-x),
+       y ~ x*(ρ-z(k-1))-y,
+       z ~ x(k-1)*y - β*z]
+@named ide = ImplicitDiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β]; tspan = (0, 1000.0))
 ```
 """
 struct ImplicitDiscreteSystem <: AbstractTimeDependentSystem
@@ -136,6 +135,7 @@ end
 
 """
     $(TYPEDSIGNATURES)
+
 Constructs a ImplicitDiscreteSystem.
 """
 function ImplicitDiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
@@ -170,6 +170,8 @@ function ImplicitDiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
             :ImplicitDiscreteSystem, force = true)
     end
 
+    # Copy equations to canonical form, but do not touch array expressions
+    eqs = [wrap(eq.lhs) isa Symbolics.Arr ? eq : 0 ~ eq.rhs - eq.lhs for eq in eqs]
     defaults = Dict{Any, Any}(todict(defaults))
     guesses = Dict{Any, Any}(todict(guesses))
     var_to_name = Dict()
@@ -236,6 +238,8 @@ 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)
@@ -259,10 +263,25 @@ end
 
 function generate_function(
         sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
-    exprs = [eq.rhs for eq in equations(sys)]
-    wrap_code = wrap_code .∘ wrap_array_vars(sys, exprs) .∘
-                wrap_parameter_dependencies(sys, false)
-    generate_custom_function(sys, exprs, dvs, ps; wrap_code, 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))
 end
 
 function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
@@ -311,7 +330,7 @@ function SciMLBase.ImplicitDiscreteProblem(
     f, u0, p = process_SciMLProblem(
         ImplicitDiscreteFunction, sys, u0map, parammap; eval_expression, eval_module)
     u0 = f(u0, p, tspan[1])
-    ImplicitDiscreteProblem(f, u0, tspan, p; kwargs...)
+    NonlinearProblem(f, u0, tspan, p; kwargs...)
 end
 
 function SciMLBase.ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...; kwargs...)
@@ -337,14 +356,15 @@ 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
     f_gen = generate_function(sys, dvs, ps; expression = Val{true},
         expression_module = eval_module, kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-    f(u, p, t) = f_oop(u, p, t)
-    f(du, u, p, t) = f_iip(du, u, p, t)
+    f(u_next, u, p, t) = f_oop(u_next, u, p, t)
+    f(resid, u_next, u, p, t) = f_iip(resid, u_next, u, p, t)
 
     if specialize === SciMLBase.FunctionWrapperSpecialize && iip
         if u0 === nothing || p === nothing || t === nothing
@@ -379,8 +399,8 @@ struct ImplicitDiscreteFunctionClosure{O, I} <: Function
     f_oop::O
     f_iip::I
 end
-(f::ImplicitDiscreteFunctionClosure)(u, p, t) = f.f_oop(u, p, t)
-(f::ImplicitDiscreteFunctionClosure)(du, u, p, t) = f.f_iip(du, u, p, t)
+(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 ImplicitDiscreteFunctionExpr{iip}(sys::ImplicitDiscreteSystem, dvs = unknowns(sys),
         ps = parameters(sys), u0 = nothing;
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 1bdc11f06a..215d843c2e 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -432,7 +432,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 DiscreteSystem),
         Any[])
-    if sys isa DiscreteSystem
+    if sys isa DiscreteSystem || sys isa ImplicitDiscreteSystem
         ts = shift_discrete_system(ts)
     end
     return ts
@@ -456,6 +456,8 @@ function lower_order_var(dervar, t)
     diffvar
 end
 
+"""
+"""
 function shift_discrete_system(ts::TearingState)
     @unpack fullvars, sys = ts
     discvars = OrderedSet()
diff --git a/test/implicit_discrete_system.jl b/test/implicit_discrete_system.jl
new file mode 100644
index 0000000000..adfb9d2fc1
--- /dev/null
+++ b/test/implicit_discrete_system.jl
@@ -0,0 +1,2 @@
+
+#init
diff --git a/test/runtests.jl b/test/runtests.jl
index 52875fdae5..bd5eecacd0 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -79,7 +79,8 @@ end
             @safetestset "Variable Utils Test" include("variable_utils.jl")
             @safetestset "Variable Metadata Test" include("test_variable_metadata.jl")
             @safetestset "OptimizationSystem Test" include("optimizationsystem.jl")
-            @safetestset "Discrete System" include("discrete_system.jl")
+            @safetestset "DiscreteSystem Test" include("discrete_system.jl")
+            @safetestset "ImplicitDiscreteSystem Test" include("implicit_discrete_system.jl")
             @safetestset "SteadyStateSystem Test" include("steadystatesystems.jl")
             @safetestset "SDESystem Test" include("sdesystem.jl")
             @safetestset "DDESystem Test" include("dde.jl")

From 397298fea4c78b4ccfbfc6f8d1c61fde5aabe16d Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 30 Jan 2025 11:34:02 -0500
Subject: [PATCH 06/41] working on structural simplification

---
 src/ModelingToolkit.jl                                  | 1 +
 src/systems/discrete_system/discrete_system.jl          | 2 +-
 src/systems/discrete_system/implicit_discrete_system.jl | 8 ++++----
 src/systems/systemstructure.jl                          | 4 ++--
 4 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 4083f2609d..2dbef8f272 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
 
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index 0ad3317ea1..f2d271dfe4 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.
diff --git a/src/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
index 1e1b5d5f6f..3a137eb8c1 100644
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ b/src/systems/discrete_system/implicit_discrete_system.jl
@@ -10,13 +10,13 @@ using ModelingToolkit: t_nounits as t
 @parameters σ=28.0 ρ=10.0 β=8/3 δt=0.1
 @variables x(t)=1.0 y(t)=0.0 z(t)=0.0
 k = ShiftIndex(t)
-eqs = [x ~ σ*(y(k-1)-x),
-       y ~ x*(ρ-z(k-1))-y,
-       z ~ x(k-1)*y - β*z]
+eqs = [x ~ σ*(y-x),
+       y ~ x*(ρ-z)-y,
+       z ~ x*y - β*z]
 @named ide = ImplicitDiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β]; tspan = (0, 1000.0))
 ```
 """
-struct ImplicitDiscreteSystem <: AbstractTimeDependentSystem
+struct ImplicitDiscreteSystem <: AbstractDiscreteSystem
     """
     A tag for the system. If two systems have the same tag, then they are
     structurally identical.
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 215d843c2e..86d1702041 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -430,9 +430,9 @@ 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 || sys isa ImplicitDiscreteSystem
+    if sys isa AbstractDiscreteSystem 
         ts = shift_discrete_system(ts)
     end
     return ts

From be1b270298f11b23306cf4b5a9ea677a0ba26894 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 30 Jan 2025 14:26:04 -0500
Subject: [PATCH 07/41] change variable renaming

---
 src/structural_transformation/utils.jl | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index bd24a1d017..29c0f66756 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -452,8 +452,15 @@ end
 function lower_varname_withshift(var, iv, order)
     order == 0 && return var
     if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
-        op = operation(var)
-        return Shift(op.t, order)(var)
+        O = only(arguments(var))
+        oldop = operation(O)
+        ds = "$iv-$order"
+        d_separator = 'ˍ'
+        newname = Symbol(string(nameof(oldop)), d_separator, ds)
+
+        newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
+        setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
+        return ModelingToolkit._with_unit(identity, newvar, iv)
     end
     return lower_varname_with_unit(var, iv, order)
 end

From cc23c888984a3e811d8be32f7f229b214dd7c7b0 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 30 Jan 2025 14:29:14 -0500
Subject: [PATCH 08/41] up

---
 src/systems/discrete_system/discrete_system.jl |  3 +--
 src/systems/systems.jl                         | 14 +++++++-------
 src/systems/systemstructure.jl                 |  6 ++----
 3 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index f2d271dfe4..bd5c72eec7 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 <: AbstractDiscreteSystem
+struct DiscreteSystem <: AbstractTimeDependentSystem
     """
     A tag for the system. If two systems have the same tag, then they are
     structurally identical.
@@ -233,7 +233,6 @@ function DiscreteSystem(eqs, iv; kwargs...)
             push!(new_ps, p)
         end
     end
-    @show allunknowns
     return DiscreteSystem(eqs, iv,
         collect(allunknowns), collect(new_ps); kwargs...)
 end
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 5ff28eba26..9c8c272c5c 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -39,13 +39,13 @@ function structural_simplify(
     else
         newsys = newsys′
     end
-    # 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.
-    #     """)
-    # end
+    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.
+        """)
+    end
     for pass in additional_passes
         newsys = pass(newsys)
     end
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 86d1702041..1bdc11f06a 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -430,9 +430,9 @@ 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 AbstractDiscreteSystem),
+            complete(graph), nothing, var_types, sys isa DiscreteSystem),
         Any[])
-    if sys isa AbstractDiscreteSystem 
+    if sys isa DiscreteSystem
         ts = shift_discrete_system(ts)
     end
     return ts
@@ -456,8 +456,6 @@ function lower_order_var(dervar, t)
     diffvar
 end
 
-"""
-"""
 function shift_discrete_system(ts::TearingState)
     @unpack fullvars, sys = ts
     discvars = OrderedSet()

From f5e5aff51476b66e8fc152ed9b8c8fb4307e8f4c Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 30 Jan 2025 14:30:08 -0500
Subject: [PATCH 09/41] revert all

---
 src/ModelingToolkit.jl                        |   3 -
 .../implicit_discrete_system.jl               | 426 ------------------
 test/implicit_discrete_system.jl              |   2 -
 test/runtests.jl                              |   3 +-
 4 files changed, 1 insertion(+), 433 deletions(-)
 delete mode 100644 src/systems/discrete_system/implicit_discrete_system.jl
 delete mode 100644 test/implicit_discrete_system.jl

diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 2dbef8f272..2710d7d1e4 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -123,7 +123,6 @@ 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,7 +164,6 @@ 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")
 
@@ -231,7 +229,6 @@ 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/systems/discrete_system/implicit_discrete_system.jl b/src/systems/discrete_system/implicit_discrete_system.jl
deleted file mode 100644
index 3a137eb8c1..0000000000
--- a/src/systems/discrete_system/implicit_discrete_system.jl
+++ /dev/null
@@ -1,426 +0,0 @@
-"""
-$(TYPEDEF)
-An implicit system of difference equations.
-# Fields
-$(FIELDS)
-# Example
-```
-using ModelingToolkit
-using ModelingToolkit: t_nounits as t
-@parameters σ=28.0 ρ=10.0 β=8/3 δt=0.1
-@variables x(t)=1.0 y(t)=0.0 z(t)=0.0
-k = ShiftIndex(t)
-eqs = [x ~ σ*(y-x),
-       y ~ x*(ρ-z)-y,
-       z ~ x*y - β*z]
-@named ide = ImplicitDiscreteSystem(eqs,t,[x,y,z],[σ,ρ,β]; tspan = (0, 1000.0))
-```
-"""
-struct ImplicitDiscreteSystem <: AbstractDiscreteSystem
-    """
-    A tag for the system. If two systems have the same tag, then they are
-    structurally identical.
-    """
-    tag::UInt
-    """The difference equations defining the discrete system."""
-    eqs::Vector{Equation}
-    """Independent variable."""
-    iv::BasicSymbolic{Real}
-    """Dependent (state) variables. Must not contain the independent variable."""
-    unknowns::Vector
-    """Parameter variables. Must not contain the independent variable."""
-    ps::Vector
-    """Time span."""
-    tspan::Union{NTuple{2, Any}, Nothing}
-    """Array variables."""
-    var_to_name::Any
-    """Observed states."""
-    observed::Vector{Equation}
-    """
-    The name of the system
-    """
-    name::Symbol
-    """
-    A description of the system.
-    """
-    description::String
-    """
-    The internal systems. These are required to have unique names.
-    """
-    systems::Vector{ImplicitDiscreteSystem}
-    """
-    The default values to use when initial conditions and/or
-    parameters are not supplied in `ImplicitDiscreteProblem`.
-    """
-    defaults::Dict
-    """
-    The guesses to use as the initial conditions for the
-    initialization system.
-    """
-    guesses::Dict
-    """
-    The system for performing the initialization.
-    """
-    initializesystem::Union{Nothing, NonlinearSystem}
-    """
-    Extra equations to be enforced during the initialization sequence.
-    """
-    initialization_eqs::Vector{Equation}
-    """
-    Inject assignment statements before the evaluation of the RHS function.
-    """
-    preface::Any
-    """
-    Type of the system.
-    """
-    connector_type::Any
-    """
-    Topologically sorted parameter dependency equations, where all symbols are parameters and
-    the LHS is a single parameter.
-    """
-    parameter_dependencies::Vector{Equation}
-    """
-    Metadata for the system, to be used by downstream packages.
-    """
-    metadata::Any
-    """
-    Metadata for MTK GUI.
-    """
-    gui_metadata::Union{Nothing, GUIMetadata}
-    """
-    Cache for intermediate tearing state.
-    """
-    tearing_state::Any
-    """
-    Substitutions generated by tearing.
-    """
-    substitutions::Any
-    """
-    If a model `sys` is complete, then `sys.x` no longer performs namespacing.
-    """
-    complete::Bool
-    """
-    Cached data for fast symbolic indexing.
-    """
-    index_cache::Union{Nothing, IndexCache}
-    """
-    The hierarchical parent system before simplification.
-    """
-    parent::Any
-    isscheduled::Bool
-
-    function ImplicitDiscreteSystem(tag, discreteEqs, iv, dvs, ps, tspan, var_to_name,
-            observed, name, description, systems, defaults, guesses, initializesystem,
-            initialization_eqs, preface, connector_type, parameter_dependencies = Equation[],
-            metadata = nothing, gui_metadata = nothing,
-            tearing_state = nothing, substitutions = nothing,
-            complete = false, index_cache = nothing, parent = nothing,
-            isscheduled = false;
-            checks::Union{Bool, Int} = true)
-        if checks == true || (checks & CheckComponents) > 0
-            check_independent_variables([iv])
-            check_variables(dvs, iv)
-            check_parameters(ps, iv)
-        end
-        if checks == true || (checks & CheckUnits) > 0
-            u = __get_unit_type(dvs, ps, iv)
-            check_units(u, discreteEqs)
-        end
-        new(tag, discreteEqs, iv, dvs, ps, tspan, var_to_name, observed, name, description,
-            systems, defaults, guesses, initializesystem, initialization_eqs,
-            preface, connector_type, parameter_dependencies, metadata, gui_metadata,
-            tearing_state, substitutions, complete, index_cache, parent, isscheduled)
-    end
-end
-
-"""
-    $(TYPEDSIGNATURES)
-
-Constructs a ImplicitDiscreteSystem.
-"""
-function ImplicitDiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
-        observed = Num[],
-        systems = ImplicitDiscreteSystem[],
-        tspan = nothing,
-        name = nothing,
-        description = "",
-        default_u0 = Dict(),
-        default_p = Dict(),
-        guesses = Dict(),
-        initializesystem = nothing,
-        initialization_eqs = Equation[],
-        defaults = _merge(Dict(default_u0), Dict(default_p)),
-        preface = nothing,
-        connector_type = nothing,
-        parameter_dependencies = Equation[],
-        metadata = nothing,
-        gui_metadata = nothing,
-        kwargs...)
-    name === nothing &&
-        throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
-    iv′ = value(iv)
-    dvs′ = value.(dvs)
-    ps′ = value.(ps)
-    if any(hasderiv, eqs) || any(hashold, eqs) || any(hassample, eqs) || any(hasdiff, eqs)
-        error("Equations in a `ImplicitDiscreteSystem` can only have `Shift` operators.")
-    end
-    if !(isempty(default_u0) && isempty(default_p))
-        Base.depwarn(
-            "`default_u0` and `default_p` are deprecated. Use `defaults` instead.",
-            :ImplicitDiscreteSystem, force = true)
-    end
-
-    # Copy equations to canonical form, but do not touch array expressions
-    eqs = [wrap(eq.lhs) isa Symbolics.Arr ? eq : 0 ~ eq.rhs - eq.lhs for eq in eqs]
-    defaults = Dict{Any, Any}(todict(defaults))
-    guesses = Dict{Any, Any}(todict(guesses))
-    var_to_name = Dict()
-    process_variables!(var_to_name, defaults, guesses, dvs′)
-    process_variables!(var_to_name, defaults, guesses, ps′)
-    process_variables!(
-        var_to_name, defaults, guesses, [eq.lhs for eq in parameter_dependencies])
-    process_variables!(
-        var_to_name, defaults, guesses, [eq.rhs for eq in parameter_dependencies])
-    defaults = Dict{Any, Any}(value(k) => value(v)
-    for (k, v) in pairs(defaults) if v !== nothing)
-    guesses = Dict{Any, Any}(value(k) => value(v)
-    for (k, v) in pairs(guesses) if v !== nothing)
-
-    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
-
-    sysnames = nameof.(systems)
-    if length(unique(sysnames)) != length(sysnames)
-        throw(ArgumentError("System names must be unique."))
-    end
-    ImplicitDiscreteSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
-        eqs, iv′, dvs′, ps′, tspan, var_to_name, observed, name, description, systems,
-        defaults, guesses, initializesystem, initialization_eqs, preface, connector_type,
-        parameter_dependencies, metadata, gui_metadata, kwargs...)
-end
-
-function ImplicitDiscreteSystem(eqs, iv; kwargs...)
-    eqs = collect(eqs)
-    diffvars = OrderedSet()
-    allunknowns = OrderedSet()
-    ps = OrderedSet()
-    iv = value(iv)
-    for eq in eqs
-        collect_vars!(allunknowns, ps, eq, iv; op = Shift)
-        if iscall(eq.lhs) && operation(eq.lhs) isa Shift
-            isequal(iv, operation(eq.lhs).t) ||
-                throw(ArgumentError("An ImplicitDiscreteSystem can only have one independent variable."))
-            eq.lhs in diffvars &&
-                throw(ArgumentError("The shift variable $(eq.lhs) is not unique in the system of equations."))
-            push!(diffvars, eq.lhs)
-        end
-    end
-    for eq in get(kwargs, :parameter_dependencies, Equation[])
-        if eq isa Pair
-            collect_vars!(allunknowns, ps, eq, iv)
-        else
-            collect_vars!(allunknowns, ps, eq, iv)
-        end
-    end
-    new_ps = OrderedSet()
-    for p in ps
-        if iscall(p) && operation(p) === getindex
-            par = arguments(p)[begin]
-            if Symbolics.shape(Symbolics.unwrap(par)) !== Symbolics.Unknown() &&
-               all(par[i] in ps for i in eachindex(par))
-                push!(new_ps, par)
-            else
-                push!(new_ps, p)
-            end
-        else
-            push!(new_ps, p)
-        end
-    end
-    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)
-    if isempty(systems)
-        return sys
-    else
-        return ImplicitDiscreteSystem(noeqs ? Equation[] : equations(sys),
-            get_iv(sys),
-            unknowns(sys),
-            parameters(sys),
-            observed = observed(sys),
-            defaults = defaults(sys),
-            guesses = guesses(sys),
-            initialization_eqs = initialization_equations(sys),
-            name = nameof(sys),
-            description = description(sys),
-            metadata = get_metadata(sys),
-            checks = false)
-    end
-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))
-end
-
-function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
-    iv = get_iv(sys)
-    updated = AnyDict()
-    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)).")
-        end
-        updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
-    end
-    for var in unknowns(sys)
-        op = operation(var)
-        op isa Shift || continue
-        haskey(updated, var) && continue
-        root = first(arguments(var))
-        haskey(defs, root) || error("Initial condition for $var not provided.")
-        updated[var] = defs[root]
-    end
-    return updated
-end
-
-"""
-    $(TYPEDSIGNATURES)
-Generates an ImplicitDiscreteProblem from an ImplicitDiscreteSystem.
-"""
-function SciMLBase.ImplicitDiscreteProblem(
-        sys::ImplicitDiscreteSystem, u0map = [], tspan = get_tspan(sys),
-        parammap = SciMLBase.NullParameters();
-        eval_module = @__MODULE__,
-        eval_expression = false,
-        use_union = false,
-        kwargs...
-)
-    if !iscomplete(sys)
-        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)
-    eqs = equations(sys)
-    iv = get_iv(sys)
-
-    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)
-    u0 = f(u0, p, tspan[1])
-    NonlinearProblem(f, u0, tspan, p; kwargs...)
-end
-
-function SciMLBase.ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...; kwargs...)
-    ImplicitDiscreteFunction{true}(sys, args...; kwargs...)
-end
-
-function SciMLBase.ImplicitDiscreteFunction{true}(sys::ImplicitDiscreteSystem, args...; kwargs...)
-    ImplicitDiscreteFunction{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
-end
-
-function SciMLBase.ImplicitDiscreteFunction{false}(sys::ImplicitDiscreteSystem, args...; kwargs...)
-    ImplicitDiscreteFunction{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
-end
-function SciMLBase.ImplicitDiscreteFunction{iip, specialize}(
-        sys::ImplicitDiscreteSystem,
-        dvs = unknowns(sys),
-        ps = parameters(sys),
-        u0 = nothing;
-        version = nothing,
-        p = nothing,
-        t = nothing,
-        eval_expression = false,
-        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
-    f_gen = generate_function(sys, dvs, ps; expression = Val{true},
-        expression_module = eval_module, kwargs...)
-    f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-    f(u_next, u, p, t) = f_oop(u_next, u, p, t)
-    f(resid, u_next, u, p, t) = f_iip(resid, u_next, u, p, t)
-
-    if specialize === SciMLBase.FunctionWrapperSpecialize && iip
-        if u0 === nothing || p === nothing || t === nothing
-            error("u0, p, and t must be specified for FunctionWrapperSpecialize on ImplicitDiscreteFunction.")
-        end
-        f = SciMLBase.wrapfun_iip(f, (u0, u0, p, t))
-    end
-
-    observedfun = ObservedFunctionCache(
-        sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false))
-
-    ImplicitDiscreteFunction{iip, specialize}(f;
-        sys = sys,
-        observed = observedfun,
-        analytic = analytic)
-end
-
-"""
-```julia
-ImplicitDiscreteFunctionExpr{iip}(sys::ImplicitDiscreteSystem, dvs = states(sys),
-                                  ps = parameters(sys);
-                                  version = nothing,
-                                  kwargs...) where {iip}
-```
-
-Create a Julia expression for an `ImplicitDiscreteFunction` from the [`ImplicitDiscreteSystem`](@ref).
-The arguments `dvs` and `ps` are used to set the order of the dependent
-variable and parameter vectors, respectively.
-"""
-struct ImplicitDiscreteFunctionExpr{iip} end
-struct ImplicitDiscreteFunctionClosure{O, I} <: Function
-    f_oop::O
-    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 ImplicitDiscreteFunctionExpr{iip}(sys::ImplicitDiscreteSystem, dvs = unknowns(sys),
-        ps = parameters(sys), u0 = nothing;
-        version = nothing, p = nothing,
-        linenumbers = false,
-        simplify = false,
-        kwargs...) where {iip}
-    f_oop, f_iip = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)
-
-    fsym = gensym(:f)
-    _f = :($fsym = $ImplicitDiscreteFunctionClosure($f_oop, $f_iip))
-
-    ex = quote
-        $_f
-        ImplicitDiscreteFunction{$iip}($fsym)
-    end
-    !linenumbers ? Base.remove_linenums!(ex) : ex
-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
deleted file mode 100644
index adfb9d2fc1..0000000000
--- a/test/implicit_discrete_system.jl
+++ /dev/null
@@ -1,2 +0,0 @@
-
-#init
diff --git a/test/runtests.jl b/test/runtests.jl
index bd5eecacd0..52875fdae5 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -79,8 +79,7 @@ end
             @safetestset "Variable Utils Test" include("variable_utils.jl")
             @safetestset "Variable Metadata Test" include("test_variable_metadata.jl")
             @safetestset "OptimizationSystem Test" include("optimizationsystem.jl")
-            @safetestset "DiscreteSystem Test" include("discrete_system.jl")
-            @safetestset "ImplicitDiscreteSystem Test" include("implicit_discrete_system.jl")
+            @safetestset "Discrete System" include("discrete_system.jl")
             @safetestset "SteadyStateSystem Test" include("steadystatesystems.jl")
             @safetestset "SDESystem Test" include("sdesystem.jl")
             @safetestset "DDESystem Test" include("dde.jl")

From b5305950778d256c70462e7858f4cb3b79ff835e Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 3 Feb 2025 07:58:14 -0500
Subject: [PATCH 10/41] working simplification

---
 .../symbolics_tearing.jl                      | 58 +++++++++++++++++--
 src/structural_transformation/utils.jl        | 18 +++---
 src/systems/systemstructure.jl                | 30 +++++++---
 3 files changed, 88 insertions(+), 18 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index d95951c41e..8c802d9765 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -240,7 +240,7 @@ end
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
     @unpack fullvars, sys, structure = state
-    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
     extra_vars = Int[]
     if full_var_eq_matching !== nothing
         for v in 𝑑vertices(state.structure.graph)
@@ -279,6 +279,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         iv = D = nothing
     end
     diff_to_var = invview(var_to_diff)
+
     dummy_sub = Dict()
     for var in 1:length(fullvars)
         dv = var_to_diff[var]
@@ -310,7 +311,10 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
             diff_to_var[dv] = nothing
         end
     end
+    @show neweqs
 
+    println("Post state selection.")
+    
     # `SelectedState` information is no longer needed past here. State selection
     # is done. All non-differentiated variables are algebraic variables, and all
     # variables that appear differentiated are differential variables.
@@ -331,10 +335,28 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                 order += 1
                 dv = dv′
             end
+            println("Order")
+            @show fullvars[dv]
+            is_only_discrete(state.structure) && begin
+                var = fullvars[dv]
+                key = operation(var) isa Shift ? only(arguments(var)) : var
+                order = -get(lowest_shift, key, 0) - order
+            end
             order, dv
         end
     end
 
+    lower_name = is_only_discrete(state.structure) ? lower_varname_withshift : lower_varname_with_unit
+    # is_only_discrete(state.structure) && for v in 1:length(fullvars)
+    #     var = fullvars[v]
+    #     op = operation(var)
+    #     if op isa Shift
+    #         x = only(arguments(var))
+    #         lowest_shift_idxs[v]
+    #         op.steps == lowest_shift[x] && (fullvars[v] = lower_varname_withshift(var, iv, -op.steps))
+    #     end
+    # end
+
     #retear = BitSet()
     # There are three cases where we want to generate new variables to convert
     # the system into first order (semi-implicit) ODEs.
@@ -384,9 +406,28 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     eq_var_matching = invview(var_eq_matching)
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
+
     for v in 1:length(var_to_diff)
-        dv = var_to_diff[v]
+        println()
+        @show fullvars
+        @show diff_to_var
+        is_highest_discrete = begin
+            var = fullvars[v]
+            op = operation(var)
+            if (!is_only_discrete(state.structure) || op isa Shift) 
+                false
+            elseif !haskey(lowest_shift, var)
+                false
+            else
+                low = lowest_shift[var]
+                idx = findfirst(x -> isequal(x, Shift(iv, low)(var)), fullvars)
+                true
+            end
+        end
+        dv = is_highest_discrete ? idx : var_to_diff[v]
+        @show (v, fullvars[v], dv)
         dv isa Int || continue
+
         solved = var_eq_matching[dv] isa Int
         solved && continue
         # check if there's `D(x) = x_t` already
@@ -404,17 +445,19 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                diff_to_var[v_t] === nothing)
                 @assert dv in rvs
                 dummy_eq = eq
+                @show "FOUND DUMMY EQ"
                 @goto FOUND_DUMMY_EQ
             end
         end
         dx = fullvars[dv]
         # add `x_t`
-        order, lv = var_order(dv)
-        x_t = lower_varname_withshift(fullvars[lv], iv, order)
+        @show order, lv = var_order(dv)
+        x_t = lower_name(fullvars[lv], iv, order)
         push!(fullvars, simplify_shifts(x_t))
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
         add_vertex!(graph, DST)
+        @show x_t, dx
         # TODO: do we care about solvable_graph? We don't use them after
         # `dummy_derivative_graph`.
         add_vertex!(solvable_graph, DST)
@@ -433,10 +476,16 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         add_edge!(solvable_graph, dummy_eq, dv)
         @assert nsrcs(graph) == nsrcs(solvable_graph) == dummy_eq
         @label FOUND_DUMMY_EQ
+        @show is_highest_discrete
+        @show diff_to_var
+        @show v_t, dv
+        # If var = x with no shift, then 
+        is_highest_discrete && (lowest_shift[x_t] = lowest_shift[fullvars[v]])
         var_to_diff[v_t] = var_to_diff[dv]
         var_eq_matching[dv] = unassigned
         eq_var_matching[dummy_eq] = dv
     end
+    @show neweqs
 
     # Will reorder equations and unknowns to be:
     # [diffeqs; ...]
@@ -537,6 +586,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
 
     deps = Vector{Int}[i == 1 ? Int[] : collect(1:(i - 1))
                        for i in 1:length(solved_equations)]
+
     # Contract the vertices in the structure graph to make the structure match
     # the new reality of the system we've just created.
     graph = contract_variables(graph, var_eq_matching, varsperm, eqsperm,
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 29c0f66756..dd063c0b07 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -451,18 +451,22 @@ end
 
 function lower_varname_withshift(var, iv, order)
     order == 0 && return var
+    ds = "$iv-$order"
+    d_separator = 'ˍ'
+
     if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
         O = only(arguments(var))
         oldop = operation(O)
-        ds = "$iv-$order"
-        d_separator = 'ˍ'
         newname = Symbol(string(nameof(oldop)), d_separator, ds)
-
-        newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
-        setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
-        return ModelingToolkit._with_unit(identity, newvar, iv)
+    else
+        O = var
+        oldop = operation(var) 
+        varname = split(string(nameof(oldop)), d_separator)[1]
+        newname = Symbol(varname, d_separator, ds)
     end
-    return lower_varname_with_unit(var, iv, order)
+    newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
+    setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
+    return ModelingToolkit._with_unit(identity, newvar, iv)
 end
 
 function isdoubleshift(var)
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 1bdc11f06a..44558ade6c 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -140,17 +140,21 @@ get_fullvars(ts::TransformationState) = ts.fullvars
 has_equations(::TransformationState) = true
 
 Base.@kwdef mutable struct SystemStructure
-    # Maps the (index of) a variable to the (index of) the variable describing
-    # its derivative.
+    """Maps the (index of) a variable to the (index of) the variable describing its derivative."""
     var_to_diff::DiffGraph
+    """Maps the (index of) a """
     eq_to_diff::DiffGraph
     # Can be access as
     # `graph` to automatically look at the bipartite graph
     # or as `torn` to assert that tearing has run.
+    """Incidence graph of the system of equations. An edge from equation x to variable y exists if variable y appears in equation x."""
     graph::BipartiteGraph{Int, Nothing}
+    """."""
     solvable_graph::Union{BipartiteGraph{Int, Nothing}, Nothing}
     var_types::Union{Vector{VariableType}, Nothing}
+    """Whether the system is discrete."""
     only_discrete::Bool
+    lowest_shift::Union{Dict, Nothing}
 end
 
 function Base.copy(structure::SystemStructure)
@@ -346,6 +350,8 @@ function TearingState(sys; quick_cancel = false, check = true)
             eqs[i] = eqs[i].lhs ~ rhs
         end
     end
+
+    ### Handle discrete variables
     lowest_shift = Dict()
     for var in fullvars
         if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
@@ -430,10 +436,10 @@ 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 DiscreteSystem, lowest_shift),
         Any[])
     if sys isa DiscreteSystem
-        ts = shift_discrete_system(ts)
+        ts = shift_discrete_system(ts, lowest_shift)
     end
     return ts
 end
@@ -456,17 +462,27 @@ function lower_order_var(dervar, t)
     diffvar
 end
 
-function shift_discrete_system(ts::TearingState)
+"""
+    Shift variable x by the largest shift s such that x(k-s) appears in the system of equations.
+    The lowest-shift term will have.
+"""
+function shift_discrete_system(ts::TearingState, lowest_shift)
     @unpack fullvars, sys = ts
+    return ts
     discvars = OrderedSet()
     eqs = equations(sys)
+
     for eq in eqs
         vars!(discvars, eq; op = Union{Sample, Hold})
     end
     iv = get_iv(sys)
-    discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, 1)(k))
+    discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, -get(lowest_shift, k, 0))(k))
     for k in discvars
-    if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold}))
+    if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold})) 
+
+    discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, 1)(k))    for k in discvars 
+    if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold})) 
+
     for i in eachindex(fullvars)
         fullvars[i] = StructuralTransformations.simplify_shifts(fast_substitute(
             fullvars[i], discmap; operator = Union{Sample, Hold}))

From fe372865a7dd79d1e576f19d50677b1977e70fd1 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 3 Feb 2025 14:21:03 -0500
Subject: [PATCH 11/41] solving equations

---
 .../symbolics_tearing.jl                      | 64 ++++++++++---------
 1 file changed, 34 insertions(+), 30 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 8c802d9765..599a75f4f7 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -237,6 +237,18 @@ function check_diff_graph(var_to_diff, fullvars)
 end
 =#
 
+function state_selection() 
+    
+end
+
+function create_new_deriv_variables() 
+    
+end
+
+function solve_solvable_equations() 
+    
+end
+
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
     @unpack fullvars, sys, structure = state
@@ -323,6 +335,11 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     is_solvable = let solvable_graph = solvable_graph
         (eq, iv) -> eq isa Int && iv isa Int && BipartiteEdge(eq, iv) in solvable_graph
     end
+    idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
+    for (i,var) in enumerate(fullvars)
+        key = operation(var) isa Shift ? only(arguments(var)) : var
+        idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
+    end
 
     # if var is like D(x)
     isdervar = let diff_to_var = diff_to_var
@@ -335,27 +352,11 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                 order += 1
                 dv = dv′
             end
-            println("Order")
-            @show fullvars[dv]
-            is_only_discrete(state.structure) && begin
-                var = fullvars[dv]
-                key = operation(var) isa Shift ? only(arguments(var)) : var
-                order = -get(lowest_shift, key, 0) - order
-            end
+            is_only_discrete(state.structure) && (order = -idx_to_lowest_shift[dv] - order - 1)
             order, dv
         end
     end
-
     lower_name = is_only_discrete(state.structure) ? lower_varname_withshift : lower_varname_with_unit
-    # is_only_discrete(state.structure) && for v in 1:length(fullvars)
-    #     var = fullvars[v]
-    #     op = operation(var)
-    #     if op isa Shift
-    #         x = only(arguments(var))
-    #         lowest_shift_idxs[v]
-    #         op.steps == lowest_shift[x] && (fullvars[v] = lower_varname_withshift(var, iv, -op.steps))
-    #     end
-    # end
 
     #retear = BitSet()
     # There are three cases where we want to generate new variables to convert
@@ -408,9 +409,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
     for v in 1:length(var_to_diff)
-        println()
-        @show fullvars
-        @show diff_to_var
         is_highest_discrete = begin
             var = fullvars[v]
             op = operation(var)
@@ -425,7 +423,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
             end
         end
         dv = is_highest_discrete ? idx : var_to_diff[v]
-        @show (v, fullvars[v], dv)
         dv isa Int || continue
 
         solved = var_eq_matching[dv] isa Int
@@ -445,19 +442,17 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                diff_to_var[v_t] === nothing)
                 @assert dv in rvs
                 dummy_eq = eq
-                @show "FOUND DUMMY EQ"
                 @goto FOUND_DUMMY_EQ
             end
         end
         dx = fullvars[dv]
         # add `x_t`
-        @show order, lv = var_order(dv)
+        order, lv = var_order(dv)
         x_t = lower_name(fullvars[lv], iv, order)
         push!(fullvars, simplify_shifts(x_t))
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
         add_vertex!(graph, DST)
-        @show x_t, dx
         # TODO: do we care about solvable_graph? We don't use them after
         # `dummy_derivative_graph`.
         add_vertex!(solvable_graph, DST)
@@ -476,16 +471,12 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         add_edge!(solvable_graph, dummy_eq, dv)
         @assert nsrcs(graph) == nsrcs(solvable_graph) == dummy_eq
         @label FOUND_DUMMY_EQ
-        @show is_highest_discrete
-        @show diff_to_var
-        @show v_t, dv
         # If var = x with no shift, then 
-        is_highest_discrete && (lowest_shift[x_t] = lowest_shift[fullvars[v]])
+        is_only_discrete(state.structure) && (idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv])
         var_to_diff[v_t] = var_to_diff[dv]
         var_eq_matching[dv] = unassigned
         eq_var_matching[dummy_eq] = dv
     end
-    @show neweqs
 
     # Will reorder equations and unknowns to be:
     # [diffeqs; ...]
@@ -501,9 +492,16 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     subeqs = Equation[]
     solved_equations = Int[]
     solved_variables = Int[]
+
     # Solve solvable equations
+    println()
+    println("SOLVING SOLVABLE EQUATIONS.")
+    @show eq_var_matching
     toporder = topological_sort(DiCMOBiGraph{false}(graph, var_eq_matching))
     eqs = Iterators.reverse(toporder)
+    @show eqs
+    @show neweqs
+    @show fullvars
     total_sub = Dict()
     idep = iv
     for ieq in eqs
@@ -516,12 +514,18 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                 isnothing(D) &&
                     error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(equations(sys)[ieq])")
                 order, lv = var_order(iv)
-                dx = D(simplify_shifts(lower_varname_withshift(
+                @show fullvars[iv]
+                @show (order, lv)
+                dx = D(simplify_shifts(lower_name(
                     fullvars[lv], idep, order - 1)))
+                @show dx
+                @show neweqs[ieq]
                 eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
                     Symbolics.symbolic_linear_solve(neweqs[ieq],
                         fullvars[iv]),
                     total_sub; operator = ModelingToolkit.Shift))
+                @show total_sub
+                @show eq
                 for e in 𝑑neighbors(graph, iv)
                     e == ieq && continue
                     for v in 𝑠neighbors(graph, e)

From dd71d2349f0d8eca81b04fc212209045bfb18018 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 3 Feb 2025 22:10:22 -0500
Subject: [PATCH 12/41] add docs

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

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

From 4475a5a6ca7b61a897f3c90041895ed23b3df5d4 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 3 Feb 2025 22:37:56 -0500
Subject: [PATCH 13/41] up

---
 .../symbolics_tearing.jl                      | 19 ++++---------------
 1 file changed, 4 insertions(+), 15 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 599a75f4f7..059c7ca9e4 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -237,18 +237,6 @@ function check_diff_graph(var_to_diff, fullvars)
 end
 =#
 
-function state_selection() 
-    
-end
-
-function create_new_deriv_variables() 
-    
-end
-
-function solve_solvable_equations() 
-    
-end
-
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
     @unpack fullvars, sys, structure = state
@@ -323,7 +311,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
             diff_to_var[dv] = nothing
         end
     end
-    @show neweqs
+    @show var_eq_matching
 
     println("Post state selection.")
     
@@ -337,7 +325,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     end
     idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
     for (i,var) in enumerate(fullvars)
-        key = operation(var) isa Shift ? only(arguments(var)) : var
+        key = (operation(var) isa Shift) ? only(arguments(var)) : var
         idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
     end
 
@@ -345,6 +333,8 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     isdervar = let diff_to_var = diff_to_var
         var -> diff_to_var[var] !== nothing
     end
+    # For discrete variables, we want the substitution to turn 
+    # Shift(t, k)(x(t)) => x_t-k(t)
     var_order = let diff_to_var = diff_to_var
         dv -> begin
             order = 0
@@ -358,7 +348,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     end
     lower_name = is_only_discrete(state.structure) ? lower_varname_withshift : lower_varname_with_unit
 
-    #retear = BitSet()
     # There are three cases where we want to generate new variables to convert
     # the system into first order (semi-implicit) ODEs.
     #

From a964dd7e1660f599afdbf0e9b2c9ddbfa717ea83 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 3 Feb 2025 23:01:48 -0500
Subject: [PATCH 14/41] up

---
 src/structural_transformation/symbolics_tearing.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 059c7ca9e4..a80bebd436 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -237,6 +237,10 @@ function check_diff_graph(var_to_diff, fullvars)
 end
 =#
 
+function substitute_lower_order!(state::TearingState) 
+    
+end
+
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
     @unpack fullvars, sys, structure = state

From 26a545dbaaf4ced4b59ae45cc02377523dec5f10 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 3 Feb 2025 23:07:07 -0500
Subject: [PATCH 15/41] up

---
 src/structural_transformation/symbolics_tearing.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index a80bebd436..e612369498 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -241,6 +241,7 @@ function substitute_lower_order!(state::TearingState)
     
 end
 
+import ModelingToolkit: Shift
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
     @unpack fullvars, sys, structure = state
@@ -255,6 +256,8 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     end
 
     neweqs = collect(equations(state))
+    lower_name = is_only_discrete(state.structure) ? lower_varname_withshift : lower_varname_with_unit
+
     # Terminology and Definition:
     #
     # A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
@@ -350,7 +353,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
             order, dv
         end
     end
-    lower_name = is_only_discrete(state.structure) ? lower_varname_withshift : lower_varname_with_unit
 
     # There are three cases where we want to generate new variables to convert
     # the system into first order (semi-implicit) ODEs.
@@ -464,7 +466,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         add_edge!(solvable_graph, dummy_eq, dv)
         @assert nsrcs(graph) == nsrcs(solvable_graph) == dummy_eq
         @label FOUND_DUMMY_EQ
-        # If var = x with no shift, then 
         is_only_discrete(state.structure) && (idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv])
         var_to_diff[v_t] = var_to_diff[dv]
         var_eq_matching[dv] = unassigned

From 1656277f45e7a12229d8ed0545a979cb8bddb273 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 4 Feb 2025 02:49:11 -0500
Subject: [PATCH 16/41] maadd comments

---
 .../symbolics_tearing.jl                      | 68 ++++++++++++++-----
 src/structural_transformation/utils.jl        |  3 +-
 src/systems/systemstructure.jl                |  3 -
 3 files changed, 52 insertions(+), 22 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index e612369498..cf6736b419 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -241,6 +241,32 @@ function substitute_lower_order!(state::TearingState)
     
 end
 
+# Documenting the differences to structural simplification for discrete systems:
+# In discrete systems the lowest-order term is x_k-i, instead of x(t). In order
+# to substitute dummy variables for x_k-1, x_k-2, ... instead you need to reverse 
+# the order. So for discrete systems `var_order` is defined a little differently.
+#
+# The orders will also be off by one. The reason this is is that the dynamics of
+# the system should be given in terms of Shift(t, 1)(x(t), x(t-1), ...). But 
+# having the observables be indexed by the next time step is not so nice. So we 
+# handle the shifts in the renaming, rather than explicitly.
+#
+# The substitution should look like the following: 
+#   x(t) -> Shift(t, 1)(x(t))
+#   x(k-1) -> x(t)
+#   x(k-2) -> x_{t-1}(t)
+#   x(k-3) -> x_{t-2}(t)
+#   and so on...
+#
+# In the implicit discrete case this shouldn't happen. The simplification should 
+# look like a NonlinearSystem.
+#
+# For discrete systems Shift(t, 2)(x(t)) is not equivalent to Shift(t, 1)(Shift(t,1)(x(t))
+# This is different from the continuous case where D(D(x)) can be substituted for 
+# by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
+# total_sub dict is updated at the time that the renamed variables are written,
+# inside the loop where new variables are generated.
+
 import ModelingToolkit: Shift
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
@@ -318,9 +344,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
             diff_to_var[dv] = nothing
         end
     end
-    @show var_eq_matching
-
-    println("Post state selection.")
     
     # `SelectedState` information is no longer needed past here. State selection
     # is done. All non-differentiated variables are algebraic variables, and all
@@ -330,18 +353,19 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     is_solvable = let solvable_graph = solvable_graph
         (eq, iv) -> eq isa Int && iv isa Int && BipartiteEdge(eq, iv) in solvable_graph
     end
-    idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
-    for (i,var) in enumerate(fullvars)
-        key = (operation(var) isa Shift) ? only(arguments(var)) : var
-        idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
+
+    if is_only_discrete(state.structure)
+         idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
+         for (i,var) in enumerate(fullvars)
+             key = (operation(var) isa Shift) ? only(arguments(var)) : var
+             idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
+         end
     end
 
     # if var is like D(x)
     isdervar = let diff_to_var = diff_to_var
         var -> diff_to_var[var] !== nothing
     end
-    # For discrete variables, we want the substitution to turn 
-    # Shift(t, k)(x(t)) => x_t-k(t)
     var_order = let diff_to_var = diff_to_var
         dv -> begin
             order = 0
@@ -349,7 +373,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                 order += 1
                 dv = dv′
             end
-            is_only_discrete(state.structure) && (order = -idx_to_lowest_shift[dv] - order - 1)
+            is_only_discrete(state.structure) && (order = -idx_to_lowest_shift[dv] - order)
             order, dv
         end
     end
@@ -403,6 +427,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
+    total_sub = Dict()
     for v in 1:length(var_to_diff)
         is_highest_discrete = begin
             var = fullvars[v]
@@ -442,8 +467,13 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         end
         dx = fullvars[dv]
         # add `x_t`
-        order, lv = var_order(dv)
+        println()
+        @show order, lv = var_order(dv)
         x_t = lower_name(fullvars[lv], iv, order)
+        @show fullvars[v]
+        @show fullvars[dv]
+        @show fullvars[lv]
+        @show dx, x_t
         push!(fullvars, simplify_shifts(x_t))
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
@@ -466,11 +496,16 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         add_edge!(solvable_graph, dummy_eq, dv)
         @assert nsrcs(graph) == nsrcs(solvable_graph) == dummy_eq
         @label FOUND_DUMMY_EQ
-        is_only_discrete(state.structure) && (idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv])
+        is_only_discrete(state.structure) && begin
+            idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
+            operation(dx) isa Shift && (total_sub[dx] = x_t)
+            order == 1 && (total_sub[x_t] = fullvars[var_to_diff[dv]])
+        end
         var_to_diff[v_t] = var_to_diff[dv]
         var_eq_matching[dv] = unassigned
         eq_var_matching[dummy_eq] = dv
     end
+    @show total_sub
 
     # Will reorder equations and unknowns to be:
     # [diffeqs; ...]
@@ -490,15 +525,13 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     # Solve solvable equations
     println()
     println("SOLVING SOLVABLE EQUATIONS.")
-    @show eq_var_matching
     toporder = topological_sort(DiCMOBiGraph{false}(graph, var_eq_matching))
     eqs = Iterators.reverse(toporder)
-    @show eqs
-    @show neweqs
-    @show fullvars
-    total_sub = Dict()
     idep = iv
+    @show eq_var_matching
+
     for ieq in eqs
+        println()
         iv = eq_var_matching[ieq]
         if is_solvable(ieq, iv)
             # We don't solve differential equations, but we will need to try to
@@ -513,7 +546,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                 dx = D(simplify_shifts(lower_name(
                     fullvars[lv], idep, order - 1)))
                 @show dx
-                @show neweqs[ieq]
                 eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
                     Symbolics.symbolic_linear_solve(neweqs[ieq],
                         fullvars[iv]),
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index dd063c0b07..707a57097b 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -451,7 +451,8 @@ end
 
 function lower_varname_withshift(var, iv, order)
     order == 0 && return var
-    ds = "$iv-$order"
+    #order == -1 && return Shift(iv, 1)(var)
+    ds = "$iv-$(order-1)"
     d_separator = 'ˍ'
 
     if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 44558ade6c..89fcebf549 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -476,9 +476,6 @@ function shift_discrete_system(ts::TearingState, lowest_shift)
         vars!(discvars, eq; op = Union{Sample, Hold})
     end
     iv = get_iv(sys)
-    discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, -get(lowest_shift, k, 0))(k))
-    for k in discvars
-    if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold})) 
 
     discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, 1)(k))    for k in discvars 
     if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold})) 

From 16ea18bbde45a0c50db93d0922375a5f54312fe9 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 4 Feb 2025 03:27:27 -0500
Subject: [PATCH 17/41] begin refactor

---
 .../symbolics_tearing.jl                      | 129 +++++++++---------
 1 file changed, 66 insertions(+), 63 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index cf6736b419..6340087003 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -237,8 +237,49 @@ function check_diff_graph(var_to_diff, fullvars)
 end
 =#
 
-function substitute_lower_order!(state::TearingState) 
-    
+
+function substitute_dummy_derivatives!(fullvars, var_to_diff, diff_to_var, var_eq_matching, neweqs)
+    # Dummy derivatives may determine that some differential variables are
+    # algebraic variables in disguise. The derivative of such variables are
+    # called dummy derivatives.
+
+    # Step 1:
+    # Replace derivatives of non-selected unknown variables by dummy derivatives
+
+    dummy_sub = Dict()
+    for var in 1:length(fullvars)
+        dv = var_to_diff[var]
+        dv === nothing && continue
+        if var_eq_matching[var] !== SelectedState()
+            dd = fullvars[dv]
+            v_t = setio(diff2term_with_unit(unwrap(dd), unwrap(iv)), false, false)
+            for eq in 𝑑neighbors(graph, dv)
+                dummy_sub[dd] = v_t
+                neweqs[eq] = fast_substitute(neweqs[eq], dd => v_t)
+            end
+            fullvars[dv] = v_t
+            # If we have:
+            # x -> D(x) -> D(D(x))
+            # We need to to transform it to:
+            # x   x_t -> D(x_t)
+            # update the structural information
+            dx = dv
+            x_t = v_t
+            while (ddx = var_to_diff[dx]) !== nothing
+                dx_t = D(x_t)
+                for eq in 𝑑neighbors(graph, ddx)
+                    neweqs[eq] = fast_substitute(neweqs[eq], fullvars[ddx] => dx_t)
+                end
+                fullvars[ddx] = dx_t
+                dx = ddx
+                x_t = dx_t
+            end
+            diff_to_var[dv] = nothing
+        end
+    end
+end
+
+function generate_derivative_variables!()
 end
 
 # Documenting the differences to structural simplification for discrete systems:
@@ -281,88 +322,50 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         end
     end
 
-    neweqs = collect(equations(state))
-    lower_name = is_only_discrete(state.structure) ? lower_varname_withshift : lower_varname_with_unit
-
-    # Terminology and Definition:
-    #
-    # A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
-    # characterize variables in `u(t)` into two classes: differential variables
-    # (denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
-    # variables are marked as `SelectedState` and they are differentiated in the
-    # DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
-    # appear in the system. Algebraic variables are variables that are not
-    # differential variables.
-    #
-    # Dummy derivatives may determine that some differential variables are
-    # algebraic variables in disguise. The derivative of such variables are
-    # called dummy derivatives.
-
-    # Step 1:
-    # Replace derivatives of non-selected unknown variables by dummy derivatives
-
     if ModelingToolkit.has_iv(state.sys)
         iv = get_iv(state.sys)
         if is_only_discrete(state.structure)
             D = Shift(iv, 1)
+            lower_name = lower_varname_withshift
+            idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
+            for (i,var) in enumerate(fullvars)
+                key = (operation(var) isa Shift) ? only(arguments(var)) : var
+                idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
+            end
         else
             D = Differential(iv)
+            lower_name = lower_varname_with_unit
         end
     else
         iv = D = nothing
+        lower_name = lower_varname_with_unit
     end
+
     diff_to_var = invview(var_to_diff)
+    neweqs = collect(equations(state))
 
-    dummy_sub = Dict()
-    for var in 1:length(fullvars)
-        dv = var_to_diff[var]
-        dv === nothing && continue
-        if var_eq_matching[var] !== SelectedState()
-            dd = fullvars[dv]
-            v_t = setio(diff2term_with_unit(unwrap(dd), unwrap(iv)), false, false)
-            for eq in 𝑑neighbors(graph, dv)
-                dummy_sub[dd] = v_t
-                neweqs[eq] = fast_substitute(neweqs[eq], dd => v_t)
-            end
-            fullvars[dv] = v_t
-            # If we have:
-            # x -> D(x) -> D(D(x))
-            # We need to to transform it to:
-            # x   x_t -> D(x_t)
-            # update the structural information
-            dx = dv
-            x_t = v_t
-            while (ddx = var_to_diff[dx]) !== nothing
-                dx_t = D(x_t)
-                for eq in 𝑑neighbors(graph, ddx)
-                    neweqs[eq] = fast_substitute(neweqs[eq], fullvars[ddx] => dx_t)
-                end
-                fullvars[ddx] = dx_t
-                dx = ddx
-                x_t = dx_t
-            end
-            diff_to_var[dv] = nothing
-        end
-    end
+    # Terminology and Definition:
+    #
+    # A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
+    # characterize variables in `u(t)` into two classes: differential variables
+    # (denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
+    # variables are marked as `SelectedState` and they are differentiated in the
+    # DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
+    # appear in the system. Algebraic variables are variables that are not
+    # differential variables.
     
+    substitute_dummy_derivatives!(fullvars, var_to_diff, diff_to_var, var_eq_matching, neweqs)
+
     # `SelectedState` information is no longer needed past here. State selection
     # is done. All non-differentiated variables are algebraic variables, and all
     # variables that appear differentiated are differential variables.
 
-    ### extract partition information
+    ### Extract partition information
     is_solvable = let solvable_graph = solvable_graph
         (eq, iv) -> eq isa Int && iv isa Int && BipartiteEdge(eq, iv) in solvable_graph
     end
 
-    if is_only_discrete(state.structure)
-         idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
-         for (i,var) in enumerate(fullvars)
-             key = (operation(var) isa Shift) ? only(arguments(var)) : var
-             idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
-         end
-    end
-
-    # if var is like D(x)
+    # if var is like D(x) or Shift(t, 1)(x)
     isdervar = let diff_to_var = diff_to_var
         var -> diff_to_var[var] !== nothing
     end

From bf3cd3394212acd34678e1927f79b958b5320bcd Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 4 Feb 2025 14:52:43 -0500
Subject: [PATCH 18/41] reorganization into functions

---
 .../symbolics_tearing.jl                      | 376 ++++++++++--------
 1 file changed, 208 insertions(+), 168 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 6340087003..cf70914a51 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -237,16 +237,22 @@ function check_diff_graph(var_to_diff, fullvars)
 end
 =#
 
+"""
+Replace derivatives of non-selected unknown variables by dummy derivatives. 
 
-function substitute_dummy_derivatives!(fullvars, var_to_diff, diff_to_var, var_eq_matching, neweqs)
-    # Dummy derivatives may determine that some differential variables are
-    # algebraic variables in disguise. The derivative of such variables are
-    # called dummy derivatives.
+State selection may determine that some differential variables are
+algebraic variables in disguise. The derivative of such variables are
+called dummy derivatives.
 
-    # Step 1:
-    # Replace derivatives of non-selected unknown variables by dummy derivatives
+`SelectedState` information is no longer needed past here. State selection
+is done. All non-differentiated variables are algebraic variables, and all
+variables that appear differentiated are differential variables.
+"""
+function substitute_dummy_derivatives!(ts::TearingState, neweqs, dummy_sub, var_eq_matching)
+    @unpack fullvars, sys, structure = ts
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
+    diff_to_var = invview(var_to_diff)
 
-    dummy_sub = Dict()
     for var in 1:length(fullvars)
         dv = var_to_diff[var]
         dv === nothing && continue
@@ -279,163 +285,73 @@ function substitute_dummy_derivatives!(fullvars, var_to_diff, diff_to_var, var_e
     end
 end
 
-function generate_derivative_variables!()
-end
+"""
+Generate new derivative variables for the system.
 
-# Documenting the differences to structural simplification for discrete systems:
-# In discrete systems the lowest-order term is x_k-i, instead of x(t). In order
-# to substitute dummy variables for x_k-1, x_k-2, ... instead you need to reverse 
-# the order. So for discrete systems `var_order` is defined a little differently.
-#
-# The orders will also be off by one. The reason this is is that the dynamics of
-# the system should be given in terms of Shift(t, 1)(x(t), x(t-1), ...). But 
-# having the observables be indexed by the next time step is not so nice. So we 
-# handle the shifts in the renaming, rather than explicitly.
-#
-# The substitution should look like the following: 
-#   x(t) -> Shift(t, 1)(x(t))
-#   x(k-1) -> x(t)
-#   x(k-2) -> x_{t-1}(t)
-#   x(k-3) -> x_{t-2}(t)
-#   and so on...
-#
-# In the implicit discrete case this shouldn't happen. The simplification should 
-# look like a NonlinearSystem.
-#
-# For discrete systems Shift(t, 2)(x(t)) is not equivalent to Shift(t, 1)(Shift(t,1)(x(t))
-# This is different from the continuous case where D(D(x)) can be substituted for 
-# by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
-# total_sub dict is updated at the time that the renamed variables are written,
-# inside the loop where new variables are generated.
+There are three cases where we want to generate new variables to convert
+the system into first order (semi-implicit) ODEs.
+    
+1. To first order:
+Whenever higher order differentiated variable like `D(D(D(x)))` appears,
+we introduce new variables `x_t`, `x_tt`, and `x_ttt` and new equations
+```
+D(x_tt) = x_ttt
+D(x_t) = x_tt
+D(x) = x_t
+```
+and replace `D(x)` to `x_t`, `D(D(x))` to `x_tt`, and `D(D(D(x)))` to
+`x_ttt`.
+
+2. To implicit to semi-implicit ODEs:
+2.1: Unsolvable derivative:
+If one derivative variable `D(x)` is unsolvable in all the equations it
+appears in, then we introduce a new variable `x_t`, a new equation
+```
+D(x) ~ x_t
+```
+and replace all other `D(x)` to `x_t`.
+
+2.2: Solvable derivative:
+If one derivative variable `D(x)` is solvable in at least one of the
+equations it appears in, then we introduce a new variable `x_t`. One of
+the solvable equations must be in the form of `0 ~ L(D(x), u...)` and
+there exists a function `l` such that `D(x) ~ l(u...)`. We should replace
+it to
+```
+0 ~ x_t - l(u...)
+D(x) ~ x_t
+```
+and replace all other `D(x)` to `x_t`.
+
+Observe that we don't need to actually introduce a new variable `x_t`, as
+the above equations can be lowered to
+```
+x_t := l(u...)
+D(x) ~ x_t
+```
+where `:=` denotes assignment.
+
+As a final note, in all the above cases where we need to introduce new
+variables and equations, don't add them when they already exist.
+"""
+function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var_eq_matching, var_order;
+        is_discrete = false, mm = nothing)
 
-import ModelingToolkit: Shift
-function tearing_reassemble(state::TearingState, var_eq_matching,
-        full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
-    @unpack fullvars, sys, structure = state
+    @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
-    extra_vars = Int[]
-    if full_var_eq_matching !== nothing
-        for v in 𝑑vertices(state.structure.graph)
-            eq = full_var_eq_matching[v]
-            eq isa Int && continue
-            push!(extra_vars, v)
-        end
-    end
-
-    if ModelingToolkit.has_iv(state.sys)
-        iv = get_iv(state.sys)
-        if is_only_discrete(state.structure)
-            D = Shift(iv, 1)
-            lower_name = lower_varname_withshift
-            idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
-            for (i,var) in enumerate(fullvars)
-                key = (operation(var) isa Shift) ? only(arguments(var)) : var
-                idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
-            end
-        else
-            D = Differential(iv)
-            lower_name = lower_varname_with_unit
-        end
-    else
-        iv = D = nothing
-        lower_name = lower_varname_with_unit
-    end
-
+    eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
-    neweqs = collect(equations(state))
-
-    # Terminology and Definition:
-    #
-    # A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
-    # characterize variables in `u(t)` into two classes: differential variables
-    # (denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
-    # variables are marked as `SelectedState` and they are differentiated in the
-    # DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
-    # appear in the system. Algebraic variables are variables that are not
-    # differential variables.
-    
-    substitute_dummy_derivatives!(fullvars, var_to_diff, diff_to_var, var_eq_matching, neweqs)
-
-    # `SelectedState` information is no longer needed past here. State selection
-    # is done. All non-differentiated variables are algebraic variables, and all
-    # variables that appear differentiated are differential variables.
-
-    ### Extract partition information
-    is_solvable = let solvable_graph = solvable_graph
-        (eq, iv) -> eq isa Int && iv isa Int && BipartiteEdge(eq, iv) in solvable_graph
-    end
-
-    # if var is like D(x) or Shift(t, 1)(x)
-    isdervar = let diff_to_var = diff_to_var
-        var -> diff_to_var[var] !== nothing
-    end
-    var_order = let diff_to_var = diff_to_var
-        dv -> begin
-            order = 0
-            while (dv′ = diff_to_var[dv]) !== nothing
-                order += 1
-                dv = dv′
-            end
-            is_only_discrete(state.structure) && (order = -idx_to_lowest_shift[dv] - order)
-            order, dv
-        end
-    end
+    iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
+    lower_name = is_discrete ? lower_name_withshift : lower_name_with_unit
 
-    # There are three cases where we want to generate new variables to convert
-    # the system into first order (semi-implicit) ODEs.
-    #
-    # 1. To first order:
-    # Whenever higher order differentiated variable like `D(D(D(x)))` appears,
-    # we introduce new variables `x_t`, `x_tt`, and `x_ttt` and new equations
-    # ```
-    # D(x_tt) = x_ttt
-    # D(x_t) = x_tt
-    # D(x) = x_t
-    # ```
-    # and replace `D(x)` to `x_t`, `D(D(x))` to `x_tt`, and `D(D(D(x)))` to
-    # `x_ttt`.
-    #
-    # 2. To implicit to semi-implicit ODEs:
-    # 2.1: Unsolvable derivative:
-    # If one derivative variable `D(x)` is unsolvable in all the equations it
-    # appears in, then we introduce a new variable `x_t`, a new equation
-    # ```
-    # D(x) ~ x_t
-    # ```
-    # and replace all other `D(x)` to `x_t`.
-    #
-    # 2.2: Solvable derivative:
-    # If one derivative variable `D(x)` is solvable in at least one of the
-    # equations it appears in, then we introduce a new variable `x_t`. One of
-    # the solvable equations must be in the form of `0 ~ L(D(x), u...)` and
-    # there exists a function `l` such that `D(x) ~ l(u...)`. We should replace
-    # it to
-    # ```
-    # 0 ~ x_t - l(u...)
-    # D(x) ~ x_t
-    # ```
-    # and replace all other `D(x)` to `x_t`.
-    #
-    # Observe that we don't need to actually introduce a new variable `x_t`, as
-    # the above equations can be lowered to
-    # ```
-    # x_t := l(u...)
-    # D(x) ~ x_t
-    # ```
-    # where `:=` denotes assignment.
-    #
-    # As a final note, in all the above cases where we need to introduce new
-    # variables and equations, don't add them when they already exist.
-    eq_var_matching = invview(var_eq_matching)
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
-    total_sub = Dict()
     for v in 1:length(var_to_diff)
         is_highest_discrete = begin
             var = fullvars[v]
             op = operation(var)
-            if (!is_only_discrete(state.structure) || op isa Shift) 
+            if (!is_discrete || op isa Shift) 
                 false
             elseif !haskey(lowest_shift, var)
                 false
@@ -450,7 +366,8 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
 
         solved = var_eq_matching[dv] isa Int
         solved && continue
-        # check if there's `D(x) = x_t` already
+
+        # Check if there's `D(x) = x_t` already
         local v_t, dummy_eq
         for eq in 𝑑neighbors(solvable_graph, dv)
             mi = get(linear_eqs, eq, 0)
@@ -468,6 +385,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                 @goto FOUND_DUMMY_EQ
             end
         end
+
         dx = fullvars[dv]
         # add `x_t`
         println()
@@ -499,7 +417,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         add_edge!(solvable_graph, dummy_eq, dv)
         @assert nsrcs(graph) == nsrcs(solvable_graph) == dummy_eq
         @label FOUND_DUMMY_EQ
-        is_only_discrete(state.structure) && begin
+        is_discrete && begin
             idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
             operation(dx) isa Shift && (total_sub[dx] = x_t)
             order == 1 && (total_sub[x_t] = fullvars[var_to_diff[dv]])
@@ -509,13 +427,51 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         eq_var_matching[dummy_eq] = dv
     end
     @show total_sub
+end
+
+"""
+Solve the solvable equations of the system and generate differential (or discrete)
+equations in terms of the selected states.
+
+Will reorder equations and unknowns to be:
+   [diffeqs; ...]
+   [diffvars; ...]
+such that the mass matrix is:
+   [I  0
+    0  0].
+
+Update the state to account for the new ordering and equations.
+"""
+function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, var_eq_matching, var_order)
+    @unpack fullvars, sys, structure = state 
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
+    eq_var_matching = invview(var_eq_matching)
+    diff_to_var = invview(var_to_diff)
+
+    if ModelingToolkit.has_iv(sys)
+        iv = get_iv(sys)
+        if is_only_discrete(structure)
+            D = Shift(iv, 1)
+            lower_name = lower_varname_withshift
+        else
+            D = Differential(iv)
+            lower_name = lower_varname_with_unit
+        end
+    else
+        iv = D = nothing
+        lower_name = lower_varname_with_unit
+    end
+
+    # if var is like D(x) or Shift(t, 1)(x)
+    isdervar = let diff_to_var = diff_to_var
+        var -> diff_to_var[var] !== nothing
+    end
+
+    ### Extract partition information
+    is_solvable = let solvable_graph = solvable_graph
+        (eq, iv) -> eq isa Int && iv isa Int && BipartiteEdge(eq, iv) in solvable_graph
+    end
 
-    # Will reorder equations and unknowns to be:
-    # [diffeqs; ...]
-    # [diffvars; ...]
-    # such that the mass matrix is:
-    # [I  0
-    #  0  0].
     diffeq_idxs = Int[]
     algeeq_idxs = Int[]
     diff_eqs = Equation[]
@@ -525,14 +481,11 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     solved_equations = Int[]
     solved_variables = Int[]
 
-    # Solve solvable equations
-    println()
-    println("SOLVING SOLVABLE EQUATIONS.")
     toporder = topological_sort(DiCMOBiGraph{false}(graph, var_eq_matching))
     eqs = Iterators.reverse(toporder)
-    idep = iv
-    @show eq_var_matching
+    idep = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
 
+    # Generate differential equations in terms of the new variables. 
     for ieq in eqs
         println()
         iv = eq_var_matching[ieq]
@@ -597,6 +550,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
             push!(algeeq_idxs, ieq)
         end
     end
+
     # TODO: BLT sorting
     neweqs = [diff_eqs; alge_eqs]
     inveqsperm = [diffeq_idxs; algeeq_idxs]
@@ -625,7 +579,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     graph = contract_variables(graph, var_eq_matching, varsperm, eqsperm,
         length(solved_variables), length(solved_variables_set))
 
-    # Update system
     new_var_to_diff = complete(DiffGraph(length(invvarsperm)))
     for (v, d) in enumerate(var_to_diff)
         v′ = varsperm[v]
@@ -641,6 +594,92 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         new_eq_to_diff[v′] = d′ > 0 ? d′ : nothing
     end
 
+    fullvars[invvarsperm], new_var_to_diff, new_eq_to_diff, neweqs, subeqs
+end
+
+# Documenting the differences to structural simplification for discrete systems:
+# In discrete systems the lowest-order term is x_k-i, instead of x(t). In order
+# to substitute dummy variables for x_k-1, x_k-2, ... instead you need to reverse 
+# the order. So for discrete systems `var_order` is defined a little differently.
+#
+# The orders will also be off by one. The reason this is is that the dynamics of
+# the system should be given in terms of Shift(t, 1)(x(t), x(t-1), ...). But 
+# having the observables be indexed by the next time step is not so nice. So we 
+# handle the shifts in the renaming, rather than explicitly.
+#
+# The substitution should look like the following: 
+#   x(t) -> Shift(t, 1)(x(t))
+#   x(k-1) -> x(t)
+#   x(k-2) -> x_{t-1}(t)
+#   x(k-3) -> x_{t-2}(t)
+#   and so on...
+#
+# In the implicit discrete case this shouldn't happen. The simplification should 
+# look like a NonlinearSystem.
+#
+# For discrete systems Shift(t, 2)(x(t)) is not equivalent to Shift(t, 1)(Shift(t,1)(x(t))
+# This is different from the continuous case where D(D(x)) can be substituted for 
+# by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
+# total_sub dict is updated at the time that the renamed variables are written,
+# inside the loop where new variables are generated.
+
+
+# Terminology and Definition:
+#
+# A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
+# characterize variables in `u(t)` into two classes: differential variables
+# (denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
+# variables are marked as `SelectedState` and they are differentiated in the
+# DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
+# appear in the system. Algebraic variables are variables that are not
+# differential variables.
+    
+import ModelingToolkit: Shift
+
+function tearing_reassemble(state::TearingState, var_eq_matching,
+        full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
+    @unpack fullvars, sys, structure = state
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
+    extra_vars = Int[]
+    if full_var_eq_matching !== nothing
+        for v in 𝑑vertices(state.structure.graph)
+            eq = full_var_eq_matching[v]
+            eq isa Int && continue
+            push!(extra_vars, v)
+        end
+    end
+    neweqs = collect(equations(state))
+    diff_to_var = invview(var_to_diff)
+    total_sub = Dict()
+    dummy_sub = Dict()
+    is_discrete = is_only_discrete(state.structure)
+
+    if is_discrete
+        idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
+        for (i,var) in enumerate(fullvars)
+            key = (operation(var) isa Shift) ? only(arguments(var)) : var
+            idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
+        end
+    end
+    var_order = let diff_to_var = diff_to_var
+        dv -> begin
+            order = 0
+            while (dv′ = diff_to_var[dv]) !== nothing
+                order += 1
+                dv = dv′
+            end
+            is_discrete && (order = -idx_to_lowest_shift[dv] - order)
+            order, dv
+        end
+    end
+
+    # Structural simplification 
+    substitute_dummy_derivatives!(state, neweqs, dummy_sub, var_eq_matching)
+    generate_derivative_variables!(state, neweqs, total_sub, var_eq_matching, var_order; 
+                                   is_discrete, mm)
+    new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs = solve_and_generate_equations!(state, neweqs, total_sub, var_eq_matching, var_order)
+
+    # Update system
     var_to_diff = new_var_to_diff
     eq_to_diff = new_eq_to_diff
     diff_to_var = invview(var_to_diff)
@@ -649,7 +688,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     @set! state.structure.graph = complete(graph)
     @set! state.structure.var_to_diff = var_to_diff
     @set! state.structure.eq_to_diff = eq_to_diff
-    @set! state.fullvars = fullvars = fullvars[invvarsperm]
+    @set! state.fullvars = fullvars = new_fullvars
     ispresent = let var_to_diff = var_to_diff, graph = graph
         i -> (!isempty(𝑑neighbors(graph, i)) ||
               (var_to_diff[i] !== nothing && !isempty(𝑑neighbors(graph, var_to_diff[i]))))
@@ -657,11 +696,13 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
 
     sys = state.sys
 
+    @show dummy_sub
     obs_sub = dummy_sub
     for eq in neweqs
         isdiffeq(eq) || continue
         obs_sub[eq.lhs] = eq.rhs
     end
+    @show obs_sub
     # TODO: compute the dependency correctly so that we don't have to do this
     obs = [fast_substitute(observed(sys), obs_sub); subeqs]
 
@@ -680,7 +721,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
 
     @set! sys.eqs = neweqs
     @set! sys.observed = obs
-
     @set! sys.substitutions = Substitutions(subeqs, deps)
 
     # Only makes sense for time-dependent

From 5a77f4347b1b9259cb6da56128146af254b827a2 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 6 Feb 2025 17:22:49 -0500
Subject: [PATCH 19/41] explicit case

---
 .../symbolics_tearing.jl                      | 180 ++++++++++--------
 src/structural_transformation/utils.jl        |   7 +-
 src/systems/systems.jl                        |   8 +-
 3 files changed, 104 insertions(+), 91 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index cf70914a51..76bba2a303 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -333,6 +333,32 @@ where `:=` denotes assignment.
 
 As a final note, in all the above cases where we need to introduce new
 variables and equations, don't add them when they already exist.
+
+###### DISCRETE SYSTEMS ####### 
+
+Documenting the differences to structural simplification for discrete systems:
+In discrete systems the lowest-order term is x_k-i, instead of x(t).
+
+The orders will also be off by one. The reason this is is that the dynamics of
+the system should be given in terms of Shift(t, 1)(x(t), x(t-1), ...). But 
+having the observables be indexed by the next time step is not so nice. So we 
+handle the shifts in the renaming, rather than explicitly.
+
+The substitution should look like the following: 
+  x(t) -> Shift(t, 1)(x(t))
+  Shift(t, -1)(x(t)) -> x(t)
+  Shift(t, -2)(x(t)) -> x_{t-1}(t)
+  Shift(t, -3)(x(t)) -> x_{t-2}(t)
+  and so on...
+
+In the implicit discrete case this shouldn't happen. The simplification should 
+look like a NonlinearSystem.
+
+For discrete systems Shift(t, 2)(x(t)) is not equivalent to Shift(t, 1)(Shift(t,1)(x(t))
+This is different from the continuous case where D(D(x)) can be substituted for 
+by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
+total_sub dict is updated at the time that the renamed variables are written,
+inside the loop where new variables are generated.
 """
 function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var_eq_matching, var_order;
         is_discrete = false, mm = nothing)
@@ -342,26 +368,41 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
     iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
-    lower_name = is_discrete ? lower_name_withshift : lower_name_with_unit
+    lower_name = is_discrete ? lower_varname_withshift : lower_varname_with_unit
+
+    # index v gets mapped to the lowest shift and the index of the unshifted variable
+    if is_discrete
+        idx_to_lowest_shift = Dict{Int, Tuple{Int, Int}}(var => (0,0) for var in 1:length(fullvars))
+        var_to_unshiftedidx = Dict{Any, Int}(var => findfirst(x -> isequal(x, var), fullvars) for var in keys(lowest_shift))
+
+        for (i,var) in enumerate(fullvars)
+            key = (operation(var) isa Shift) ? only(arguments(var)) : var
+            idx_to_lowest_shift[i] = (get(lowest_shift, key, 0), get(var_to_unshiftedidx, key, i))
+        end
+    end
 
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
+    # v is the index of the current variable, x = fullvars[v]
+    # dv is the index of the derivative dx = D(x), x_t is the substituted variable 
+    #
+    # For ODESystems: lv is the index of the lowest-order variable (x(t))
+    # For DiscreteSystems: 
+    # - lv is the index of the lowest-order variable (Shift(t, k)(x(t)))
+    # - uv is the index of the highest-order variable (x(t))
     for v in 1:length(var_to_diff)
-        is_highest_discrete = begin
-            var = fullvars[v]
-            op = operation(var)
-            if (!is_discrete || op isa Shift) 
-                false
-            elseif !haskey(lowest_shift, var)
-                false
-            else
-                low = lowest_shift[var]
-                idx = findfirst(x -> isequal(x, Shift(iv, low)(var)), fullvars)
-                true
+        dv = var_to_diff[v]
+        if is_discrete 
+            x = fullvars[v]
+            op = operation(x)
+            (low, uv) = idx_to_lowest_shift[v]
+
+            # If v is unshifted (i.e. x(t)), then substitute the lowest-shift variable
+            if !(op isa Shift) && (low != 0)
+                dv = findfirst(_x -> isequal(_x, Shift(iv, low)(x)), fullvars)
             end
         end
-        dv = is_highest_discrete ? idx : var_to_diff[v]
         dv isa Int || continue
 
         solved = var_eq_matching[dv] isa Int
@@ -388,13 +429,9 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
 
         dx = fullvars[dv]
         # add `x_t`
-        println()
-        @show order, lv = var_order(dv)
-        x_t = lower_name(fullvars[lv], iv, order)
-        @show fullvars[v]
-        @show fullvars[dv]
-        @show fullvars[lv]
-        @show dx, x_t
+        order, lv = var_order(dv)
+        x_t = is_discrete ? lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv]) :
+                            lower_name(fullvars[lv], iv, order)
         push!(fullvars, simplify_shifts(x_t))
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
@@ -402,11 +439,27 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
         # TODO: do we care about solvable_graph? We don't use them after
         # `dummy_derivative_graph`.
         add_vertex!(solvable_graph, DST)
-        # var_eq_matching is a bit odd.
-        # length(var_eq_matching) == length(invview(var_eq_matching))
         push!(var_eq_matching, unassigned)
         @assert v_t_idx == ndsts(graph) == ndsts(solvable_graph) == length(fullvars) ==
                 length(var_eq_matching)
+
+        # Add the substitutions to total_sub directly.  
+        is_discrete && begin
+            idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
+            @show dx
+            if operation(dx) isa Shift 
+                total_sub[dx] = x_t
+                for e in 𝑑neighbors(graph, dv)
+                    add_edge!(graph, e, v_t)
+                    rem_edge!(graph, e, dv)
+                end
+                @show graph
+                !(operation(x) isa Shift) && begin
+                    var_to_diff[v_t] = var_to_diff[dv]
+                    continue
+                end
+            end
+        end
         # add `D(x) - x_t ~ 0`
         push!(neweqs, 0 ~ x_t - dx)
         add_vertex!(graph, SRC)
@@ -417,16 +470,10 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
         add_edge!(solvable_graph, dummy_eq, dv)
         @assert nsrcs(graph) == nsrcs(solvable_graph) == dummy_eq
         @label FOUND_DUMMY_EQ
-        is_discrete && begin
-            idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
-            operation(dx) isa Shift && (total_sub[dx] = x_t)
-            order == 1 && (total_sub[x_t] = fullvars[var_to_diff[dv]])
-        end
         var_to_diff[v_t] = var_to_diff[dv]
         var_eq_matching[dv] = unassigned
         eq_var_matching[dummy_eq] = dv
     end
-    @show total_sub
 end
 
 """
@@ -442,7 +489,7 @@ such that the mass matrix is:
 
 Update the state to account for the new ordering and equations.
 """
-function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, var_eq_matching, var_order)
+function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, var_eq_matching, var_order; simplify = false)
     @unpack fullvars, sys, structure = state 
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
     eq_var_matching = invview(var_eq_matching)
@@ -467,7 +514,7 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
         var -> diff_to_var[var] !== nothing
     end
 
-    ### Extract partition information
+    # Extract partition information
     is_solvable = let solvable_graph = solvable_graph
         (eq, iv) -> eq isa Int && iv isa Int && BipartiteEdge(eq, iv) in solvable_graph
     end
@@ -485,9 +532,12 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
     eqs = Iterators.reverse(toporder)
     idep = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
 
-    # Generate differential equations in terms of the new variables. 
+    @show eq_var_matching
+    @show fullvars
+    @show neweqs
+
+    # Equation ieq is solved for the RHS of iv 
     for ieq in eqs
-        println()
         iv = eq_var_matching[ieq]
         if is_solvable(ieq, iv)
             # We don't solve differential equations, but we will need to try to
@@ -497,23 +547,14 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
                 isnothing(D) &&
                     error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(equations(sys)[ieq])")
                 order, lv = var_order(iv)
-                @show fullvars[iv]
-                @show (order, lv)
-                dx = D(simplify_shifts(lower_name(
-                    fullvars[lv], idep, order - 1)))
-                @show dx
+                dx = D(fullvars[lv])
                 eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
                     Symbolics.symbolic_linear_solve(neweqs[ieq],
                         fullvars[iv]),
                     total_sub; operator = ModelingToolkit.Shift))
-                @show total_sub
-                @show eq
                 for e in 𝑑neighbors(graph, iv)
-                    e == ieq && continue
-                    for v in 𝑠neighbors(graph, e)
-                        add_edge!(graph, e, v)
-                    end
                     rem_edge!(graph, e, iv)
+                    add_edge!(graph, e, lv)
                 end
                 push!(diff_eqs, eq)
                 total_sub[simplify_shifts(eq.lhs)] = eq.rhs
@@ -594,38 +635,10 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
         new_eq_to_diff[v′] = d′ > 0 ? d′ : nothing
     end
 
-    fullvars[invvarsperm], new_var_to_diff, new_eq_to_diff, neweqs, subeqs
+    fullvars[invvarsperm], new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph
 end
 
-# Documenting the differences to structural simplification for discrete systems:
-# In discrete systems the lowest-order term is x_k-i, instead of x(t). In order
-# to substitute dummy variables for x_k-1, x_k-2, ... instead you need to reverse 
-# the order. So for discrete systems `var_order` is defined a little differently.
-#
-# The orders will also be off by one. The reason this is is that the dynamics of
-# the system should be given in terms of Shift(t, 1)(x(t), x(t-1), ...). But 
-# having the observables be indexed by the next time step is not so nice. So we 
-# handle the shifts in the renaming, rather than explicitly.
-#
-# The substitution should look like the following: 
-#   x(t) -> Shift(t, 1)(x(t))
-#   x(k-1) -> x(t)
-#   x(k-2) -> x_{t-1}(t)
-#   x(k-3) -> x_{t-2}(t)
-#   and so on...
-#
-# In the implicit discrete case this shouldn't happen. The simplification should 
-# look like a NonlinearSystem.
-#
-# For discrete systems Shift(t, 2)(x(t)) is not equivalent to Shift(t, 1)(Shift(t,1)(x(t))
-# This is different from the continuous case where D(D(x)) can be substituted for 
-# by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
-# total_sub dict is updated at the time that the renamed variables are written,
-# inside the loop where new variables are generated.
-
-
 # Terminology and Definition:
-#
 # A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
 # characterize variables in `u(t)` into two classes: differential variables
 # (denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
@@ -654,13 +667,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     dummy_sub = Dict()
     is_discrete = is_only_discrete(state.structure)
 
-    if is_discrete
-        idx_to_lowest_shift = Dict{Int, Int}(var => 0 for var in 1:length(fullvars))
-        for (i,var) in enumerate(fullvars)
-            key = (operation(var) isa Shift) ? only(arguments(var)) : var
-            idx_to_lowest_shift[i] = get(lowest_shift, key, 0)
-        end
-    end
     var_order = let diff_to_var = diff_to_var
         dv -> begin
             order = 0
@@ -668,7 +674,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
                 order += 1
                 dv = dv′
             end
-            is_discrete && (order = -idx_to_lowest_shift[dv] - order)
             order, dv
         end
     end
@@ -677,7 +682,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     substitute_dummy_derivatives!(state, neweqs, dummy_sub, var_eq_matching)
     generate_derivative_variables!(state, neweqs, total_sub, var_eq_matching, var_order; 
                                    is_discrete, mm)
-    new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs = solve_and_generate_equations!(state, neweqs, total_sub, var_eq_matching, var_order)
+    new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph = solve_and_generate_equations!(state, neweqs, total_sub, var_eq_matching, var_order; simplify)
 
     # Update system
     var_to_diff = new_var_to_diff
@@ -693,16 +698,25 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         i -> (!isempty(𝑑neighbors(graph, i)) ||
               (var_to_diff[i] !== nothing && !isempty(𝑑neighbors(graph, var_to_diff[i]))))
     end
+    @show graph
+    println()
+    println("Shift test...")
+    @show neweqs
+    @show fullvars
+    @show 𝑑neighbors(graph, 5)
 
     sys = state.sys
-
-    @show dummy_sub
     obs_sub = dummy_sub
     for eq in neweqs
         isdiffeq(eq) || continue
         obs_sub[eq.lhs] = eq.rhs
     end
+    is_discrete && for eq in subeqs
+        obs_sub[eq.rhs] = eq.lhs
+    end
+
     @show obs_sub
+    @show observed(sys)
     # TODO: compute the dependency correctly so that we don't have to do this
     obs = [fast_substitute(observed(sys), obs_sub); subeqs]
 
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 707a57097b..8df2d4177b 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -449,10 +449,9 @@ end
 ### Misc
 ###
 
-function lower_varname_withshift(var, iv, order)
-    order == 0 && return var
-    #order == -1 && return Shift(iv, 1)(var)
-    ds = "$iv-$(order-1)"
+function lower_varname_withshift(var, iv, backshift; unshifted = nothing)
+    backshift == 0 && return unshifted 
+    ds = "$iv-$backshift"
     d_separator = 'ˍ'
 
     if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 9c8c272c5c..0a7e4264e4 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -41,10 +41,10 @@ function structural_simplify(
     end
     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.
-        """)
+        # error("""
+        #     Encountered algebraic equations when simplifying discrete system. This is \
+        #     not yet supported.
+        # """)
     end
     for pass in additional_passes
         newsys = pass(newsys)

From 91acf912e7651b1482a5e29552026774ed5fe5f0 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 6 Feb 2025 22:48:29 -0500
Subject: [PATCH 20/41] beginning implicit equation

---
 .../symbolics_tearing.jl                      | 76 +++++++++----------
 src/structural_transformation/utils.jl        |  1 +
 2 files changed, 37 insertions(+), 40 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 76bba2a303..fb1765964e 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -360,9 +360,8 @@ by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
 total_sub dict is updated at the time that the renamed variables are written,
 inside the loop where new variables are generated.
 """
-function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var_eq_matching, var_order;
+function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var_eq_matching;
         is_discrete = false, mm = nothing)
-
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
     eq_var_matching = invview(var_eq_matching)
@@ -386,13 +385,14 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
 
     # v is the index of the current variable, x = fullvars[v]
     # dv is the index of the derivative dx = D(x), x_t is the substituted variable 
-    #
     # For ODESystems: lv is the index of the lowest-order variable (x(t))
     # For DiscreteSystems: 
     # - lv is the index of the lowest-order variable (Shift(t, k)(x(t)))
     # - uv is the index of the highest-order variable (x(t))
     for v in 1:length(var_to_diff)
         dv = var_to_diff[v]
+        println()
+        @show (v, dv)
         if is_discrete 
             x = fullvars[v]
             op = operation(x)
@@ -405,6 +405,9 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
         end
         dv isa Int || continue
 
+        @show dv
+        @show var_eq_matching[dv]
+        @show fullvars
         solved = var_eq_matching[dv] isa Int
         solved && continue
 
@@ -429,9 +432,10 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
 
         dx = fullvars[dv]
         # add `x_t`
-        order, lv = var_order(dv)
+        order, lv = var_order(diff_to_var, dv)
         x_t = is_discrete ? lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv]) :
                             lower_name(fullvars[lv], iv, order)
+        @show dx, x_t
         push!(fullvars, simplify_shifts(x_t))
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
@@ -443,23 +447,23 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
         @assert v_t_idx == ndsts(graph) == ndsts(solvable_graph) == length(fullvars) ==
                 length(var_eq_matching)
 
-        # Add the substitutions to total_sub directly.  
+        # Add discrete substitutions to total_sub directly.  
         is_discrete && begin
             idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
-            @show dx
             if operation(dx) isa Shift 
                 total_sub[dx] = x_t
                 for e in 𝑑neighbors(graph, dv)
                     add_edge!(graph, e, v_t)
                     rem_edge!(graph, e, dv)
                 end
-                @show graph
+                # Do not add the lowest-order substitution as an equation, just substitute
                 !(operation(x) isa Shift) && begin
                     var_to_diff[v_t] = var_to_diff[dv]
                     continue
                 end
             end
         end
+
         # add `D(x) - x_t ~ 0`
         push!(neweqs, 0 ~ x_t - dx)
         add_vertex!(graph, SRC)
@@ -489,7 +493,7 @@ such that the mass matrix is:
 
 Update the state to account for the new ordering and equations.
 """
-function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, var_eq_matching, var_order; simplify = false)
+function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, var_eq_matching; simplify = false)
     @unpack fullvars, sys, structure = state 
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
     eq_var_matching = invview(var_eq_matching)
@@ -530,13 +534,11 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
 
     toporder = topological_sort(DiCMOBiGraph{false}(graph, var_eq_matching))
     eqs = Iterators.reverse(toporder)
-    idep = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
-
-    @show eq_var_matching
-    @show fullvars
-    @show neweqs
+    idep = iv
 
-    # Equation ieq is solved for the RHS of iv 
+    # Generate differential equations.
+    # fullvars[iv] is a differential variable of the form D^n(x), and neweqs[ieq]
+    # is solved to give the RHS.
     for ieq in eqs
         iv = eq_var_matching[ieq]
         if is_solvable(ieq, iv)
@@ -546,7 +548,7 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
             if isdervar(iv)
                 isnothing(D) &&
                     error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(equations(sys)[ieq])")
-                order, lv = var_order(iv)
+                order, lv = var_order(diff_to_var, iv)
                 dx = D(fullvars[lv])
                 eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
                     Symbolics.symbolic_linear_solve(neweqs[ieq],
@@ -634,8 +636,9 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
         d′ = eqsperm[d]
         new_eq_to_diff[v′] = d′ > 0 ? d′ : nothing
     end
+    new_fullvars = fullvars[invvarsperm]
 
-    fullvars[invvarsperm], new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph
+    new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph
 end
 
 # Terminology and Definition:
@@ -649,6 +652,16 @@ end
     
 import ModelingToolkit: Shift
 
+# Give the order of the variable indexed by dv
+function var_order(diff_to_var, dv) 
+    order = 0
+    while (dv′ = diff_to_var[dv]) !== nothing
+        order += 1
+        dv = dv′
+    end
+    order, dv
+end
+
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
     @unpack fullvars, sys, structure = state
@@ -667,22 +680,12 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     dummy_sub = Dict()
     is_discrete = is_only_discrete(state.structure)
 
-    var_order = let diff_to_var = diff_to_var
-        dv -> begin
-            order = 0
-            while (dv′ = diff_to_var[dv]) !== nothing
-                order += 1
-                dv = dv′
-            end
-            order, dv
-        end
-    end
-
     # Structural simplification 
     substitute_dummy_derivatives!(state, neweqs, dummy_sub, var_eq_matching)
-    generate_derivative_variables!(state, neweqs, total_sub, var_eq_matching, var_order; 
+    generate_derivative_variables!(state, neweqs, total_sub, var_eq_matching; 
                                    is_discrete, mm)
-    new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph = solve_and_generate_equations!(state, neweqs, total_sub, var_eq_matching, var_order; simplify)
+    new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph = 
+        solve_and_generate_equations!(state, neweqs, total_sub, var_eq_matching; simplify)
 
     # Update system
     var_to_diff = new_var_to_diff
@@ -698,25 +701,18 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         i -> (!isempty(𝑑neighbors(graph, i)) ||
               (var_to_diff[i] !== nothing && !isempty(𝑑neighbors(graph, var_to_diff[i]))))
     end
-    @show graph
-    println()
-    println("Shift test...")
-    @show neweqs
-    @show fullvars
+    @show ispresent.(collect(1:length(fullvars)))
     @show 𝑑neighbors(graph, 5)
+    @show var_to_diff[5]
 
+    @show neweqs
+    @show fullvars
     sys = state.sys
     obs_sub = dummy_sub
     for eq in neweqs
         isdiffeq(eq) || continue
         obs_sub[eq.lhs] = eq.rhs
     end
-    is_discrete && for eq in subeqs
-        obs_sub[eq.rhs] = eq.lhs
-    end
-
-    @show obs_sub
-    @show observed(sys)
     # TODO: compute the dependency correctly so that we don't have to do this
     obs = [fast_substitute(observed(sys), obs_sub); subeqs]
 
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 8df2d4177b..b2d67dc547 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -450,6 +450,7 @@ end
 ###
 
 function lower_varname_withshift(var, iv, backshift; unshifted = nothing)
+    backshift < 0 && return Shift(iv, -backshift)(var)
     backshift == 0 && return unshifted 
     ds = "$iv-$backshift"
     d_separator = 'ˍ'

From 56829f7c05e94a6048c55b482af8c1752e9cbad7 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 7 Feb 2025 02:26:13 -0500
Subject: [PATCH 21/41] up

---
 .../symbolics_tearing.jl                      | 112 ++++++++++--------
 src/structural_transformation/utils.jl        |  13 +-
 2 files changed, 73 insertions(+), 52 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index fb1765964e..020c85b05f 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -252,6 +252,7 @@ function substitute_dummy_derivatives!(ts::TearingState, neweqs, dummy_sub, var_
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
     diff_to_var = invview(var_to_diff)
+    iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
 
     for var in 1:length(fullvars)
         dv = var_to_diff[var]
@@ -337,16 +338,17 @@ variables and equations, don't add them when they already exist.
 ###### DISCRETE SYSTEMS ####### 
 
 Documenting the differences to structural simplification for discrete systems:
-In discrete systems the lowest-order term is x_k-i, instead of x(t).
+
+1. In discrete systems the lowest-order term is Shift(t, k)(x(t)), instead of x(t). We need to substitute the k-1 lowest order terms instead of the k-1 highest order terms.
 
 The orders will also be off by one. The reason this is is that the dynamics of
 the system should be given in terms of Shift(t, 1)(x(t), x(t-1), ...). But 
 having the observables be indexed by the next time step is not so nice. So we 
-handle the shifts in the renaming, rather than explicitly.
+handle the shifts in the renaming.
 
 The substitution should look like the following: 
   x(t) -> Shift(t, 1)(x(t))
-  Shift(t, -1)(x(t)) -> x(t)
+  Shift(t, -1)(x(t)) -> Shift(t, 0)(x(t))
   Shift(t, -2)(x(t)) -> x_{t-1}(t)
   Shift(t, -3)(x(t)) -> x_{t-2}(t)
   and so on...
@@ -354,14 +356,14 @@ The substitution should look like the following:
 In the implicit discrete case this shouldn't happen. The simplification should 
 look like a NonlinearSystem.
 
-For discrete systems Shift(t, 2)(x(t)) is not equivalent to Shift(t, 1)(Shift(t,1)(x(t))
+2. For discrete systems Shift(t, 2)(x(t)) cannot be substituted as Shift(t, 1)(Shift(t,1)(x(t)). 
 This is different from the continuous case where D(D(x)) can be substituted for 
 by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
-total_sub dict is updated at the time that the renamed variables are written,
+shift_sub dict is updated at the time that the renamed variables are written,
 inside the loop where new variables are generated.
 """
-function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var_eq_matching;
-        is_discrete = false, mm = nothing)
+function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching;
+        is_discrete = false, mm = nothing, shift_sub = nothing)
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
     eq_var_matching = invview(var_eq_matching)
@@ -391,23 +393,24 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
     # - uv is the index of the highest-order variable (x(t))
     for v in 1:length(var_to_diff)
         dv = var_to_diff[v]
-        println()
-        @show (v, dv)
+
         if is_discrete 
             x = fullvars[v]
             op = operation(x)
             (low, uv) = idx_to_lowest_shift[v]
 
             # If v is unshifted (i.e. x(t)), then substitute the lowest-shift variable
-            if !(op isa Shift) && (low != 0)
+            if !(op isa Shift)
                 dv = findfirst(_x -> isequal(_x, Shift(iv, low)(x)), fullvars)
             end
+            dx = fullvars[dv]
+            order, lv = var_order(diff_to_var, dv)
+            @show fullvars[uv]
+            x_t = lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv])
+            shift_sub[dx] = x_t
+            (var_eq_matching[dv] isa Int) ? continue : @goto DISCRETE_VARIABLE
         end
         dv isa Int || continue
-
-        @show dv
-        @show var_eq_matching[dv]
-        @show fullvars
         solved = var_eq_matching[dv] isa Int
         solved && continue
 
@@ -430,13 +433,13 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
             end
         end
 
-        dx = fullvars[dv]
         # add `x_t`
+        dx = fullvars[dv]
         order, lv = var_order(diff_to_var, dv)
-        x_t = is_discrete ? lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv]) :
-                            lower_name(fullvars[lv], iv, order)
-        @show dx, x_t
-        push!(fullvars, simplify_shifts(x_t))
+        x_t = lower_name(fullvars[lv], iv, order)
+        
+        @label DISCRETE_VARIABLE
+        push!(fullvars, x_t)
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
         add_vertex!(graph, DST)
@@ -444,23 +447,18 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
         # `dummy_derivative_graph`.
         add_vertex!(solvable_graph, DST)
         push!(var_eq_matching, unassigned)
-        @assert v_t_idx == ndsts(graph) == ndsts(solvable_graph) == length(fullvars) ==
-                length(var_eq_matching)
 
         # Add discrete substitutions to total_sub directly.  
         is_discrete && begin
             idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
-            if operation(dx) isa Shift 
-                total_sub[dx] = x_t
-                for e in 𝑑neighbors(graph, dv)
-                    add_edge!(graph, e, v_t)
-                    rem_edge!(graph, e, dv)
-                end
-                # Do not add the lowest-order substitution as an equation, just substitute
-                !(operation(x) isa Shift) && begin
-                    var_to_diff[v_t] = var_to_diff[dv]
-                    continue
-                end
+            for e in 𝑑neighbors(graph, dv)
+                add_edge!(graph, e, v_t)
+                rem_edge!(graph, e, dv)
+            end
+            # Do not add the lowest-order substitution as an equation, just substitute
+            !(operation(x) isa Shift) && begin
+                var_to_diff[v_t] = var_to_diff[dv]
+                continue
             end
         end
 
@@ -472,7 +470,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
         add_edge!(graph, dummy_eq, v_t)
         add_vertex!(solvable_graph, SRC)
         add_edge!(solvable_graph, dummy_eq, dv)
-        @assert nsrcs(graph) == nsrcs(solvable_graph) == dummy_eq
+
         @label FOUND_DUMMY_EQ
         var_to_diff[v_t] = var_to_diff[dv]
         var_eq_matching[dv] = unassigned
@@ -480,6 +478,14 @@ function generate_derivative_variables!(ts::TearingState, neweqs, total_sub, var
     end
 end
 
+function add_solvable_variable!() 
+    
+end
+
+function add_solvable_equation!() 
+    
+end
+
 """
 Solve the solvable equations of the system and generate differential (or discrete)
 equations in terms of the selected states.
@@ -492,21 +498,27 @@ such that the mass matrix is:
     0  0].
 
 Update the state to account for the new ordering and equations.
+
+####### DISCRETE CASE
+- only substitute Shift(t, -2) 
 """
-function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, var_eq_matching; simplify = false)
+function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false, shift_sub = Dict())
     @unpack fullvars, sys, structure = state 
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
+    dx_sub = Dict()
 
     if ModelingToolkit.has_iv(sys)
         iv = get_iv(sys)
         if is_only_discrete(structure)
             D = Shift(iv, 1)
             lower_name = lower_varname_withshift
+            total_sub = shift_sub
         else
             D = Differential(iv)
             lower_name = lower_varname_with_unit
+            total_sub = dx_sub
         end
     else
         iv = D = nothing
@@ -540,6 +552,7 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
     # fullvars[iv] is a differential variable of the form D^n(x), and neweqs[ieq]
     # is solved to give the RHS.
     for ieq in eqs
+        println()
         iv = eq_var_matching[ieq]
         if is_solvable(ieq, iv)
             # We don't solve differential equations, but we will need to try to
@@ -549,7 +562,9 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
                 isnothing(D) &&
                     error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(equations(sys)[ieq])")
                 order, lv = var_order(diff_to_var, iv)
-                dx = D(fullvars[lv])
+                @show fullvars[lv]
+                @show simplify_shifts(fullvars[lv])
+                dx = D(simplify_shifts(fullvars[lv]))
                 eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
                     Symbolics.symbolic_linear_solve(neweqs[ieq],
                         fullvars[iv]),
@@ -560,6 +575,9 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
                 end
                 push!(diff_eqs, eq)
                 total_sub[simplify_shifts(eq.lhs)] = eq.rhs
+                dx_sub[simplify_shifts(eq.lhs)] = eq.rhs
+                @show total_sub
+                @show eq
                 push!(diffeq_idxs, ieq)
                 push!(diff_vars, diff_to_var[iv])
                 continue
@@ -575,10 +593,10 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
                 @warn "Tearing: solving $eq for $var is singular!"
             else
                 rhs = -b / a
-                neweq = var ~ Symbolics.fixpoint_sub(
+                neweq = var ~ simplify_shifts(Symbolics.fixpoint_sub(
                     simplify ?
                     Symbolics.simplify(rhs) : rhs,
-                    total_sub; operator = ModelingToolkit.Shift)
+                    dx_sub; operator = ModelingToolkit.Shift))
                 push!(subeqs, neweq)
                 push!(solved_equations, ieq)
                 push!(solved_variables, iv)
@@ -589,7 +607,7 @@ function solve_and_generate_equations!(state::TearingState, neweqs, total_sub, v
             if !(eq.lhs isa Number && eq.lhs == 0)
                 rhs = eq.rhs - eq.lhs
             end
-            push!(alge_eqs, 0 ~ Symbolics.fixpoint_sub(rhs, total_sub))
+            push!(alge_eqs, 0 ~ simplify_shifts(Symbolics.fixpoint_sub(rhs, total_sub)))
             push!(algeeq_idxs, ieq)
         end
     end
@@ -676,16 +694,19 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     end
     neweqs = collect(equations(state))
     diff_to_var = invview(var_to_diff)
-    total_sub = Dict()
-    dummy_sub = Dict()
     is_discrete = is_only_discrete(state.structure)
 
+    shift_sub = Dict() 
+
     # Structural simplification 
+    dummy_sub = Dict()
     substitute_dummy_derivatives!(state, neweqs, dummy_sub, var_eq_matching)
-    generate_derivative_variables!(state, neweqs, total_sub, var_eq_matching; 
-                                   is_discrete, mm)
+
+    generate_derivative_variables!(state, neweqs, var_eq_matching; 
+                                   is_discrete, mm, shift_sub)
+
     new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph = 
-        solve_and_generate_equations!(state, neweqs, total_sub, var_eq_matching; simplify)
+        solve_and_generate_equations!(state, neweqs, var_eq_matching; simplify, shift_sub)
 
     # Update system
     var_to_diff = new_var_to_diff
@@ -701,12 +722,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         i -> (!isempty(𝑑neighbors(graph, i)) ||
               (var_to_diff[i] !== nothing && !isempty(𝑑neighbors(graph, var_to_diff[i]))))
     end
-    @show ispresent.(collect(1:length(fullvars)))
-    @show 𝑑neighbors(graph, 5)
-    @show var_to_diff[5]
 
-    @show neweqs
-    @show fullvars
     sys = state.sys
     obs_sub = dummy_sub
     for eq in neweqs
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index b2d67dc547..c202556906 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -449,10 +449,9 @@ end
 ### Misc
 ###
 
-function lower_varname_withshift(var, iv, backshift; unshifted = nothing)
-    backshift < 0 && return Shift(iv, -backshift)(var)
-    backshift == 0 && return unshifted 
-    ds = "$iv-$backshift"
+function lower_varname_withshift(var, iv, backshift; unshifted = nothing, allow_zero = true)
+    backshift <= 0 && return Shift(iv, -backshift)(unshifted, allow_zero)
+    ds = backshift > 0 ? "$iv-$backshift" : "$iv+$(-backshift)"
     d_separator = 'ˍ'
 
     if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
@@ -475,9 +474,15 @@ function isdoubleshift(var)
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
 end
 
+### Rules
+# 1. x(t) -> x(t)
+# 2. Shift(t, 0)(x(t)) -> x(t)
+# 3. Shift(t, 1)(x + z) -> Shift(t, 1)(x) + Shift(t, 1)(z)
+
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
     var isa Equation && return simplify_shifts(var.lhs) ~ simplify_shifts(var.rhs)
+    ((op = operation(var)) isa Shift) && op.steps == 0 && return simplify_shifts(arguments(var)[1])
     if isdoubleshift(var)
         op1 = operation(var)
         vv1 = arguments(var)[1]

From 1c578c3535fe0d7af1d82d08285d90fd244ce2af Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 7 Feb 2025 14:46:09 -0500
Subject: [PATCH 22/41] rename functions

---
 .../symbolics_tearing.jl                      | 37 +++++++-----
 src/structural_transformation/utils.jl        | 60 ++++++++++++++++++-
 test/structural_transformation/utils.jl       | 17 ++++++
 3 files changed, 98 insertions(+), 16 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 020c85b05f..778473907e 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -254,7 +254,9 @@ function substitute_dummy_derivatives!(ts::TearingState, neweqs, dummy_sub, var_
     diff_to_var = invview(var_to_diff)
     iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
 
+    @show neweqs
     for var in 1:length(fullvars)
+        #@show neweqs
         dv = var_to_diff[var]
         dv === nothing && continue
         if var_eq_matching[var] !== SelectedState()
@@ -286,9 +288,7 @@ function substitute_dummy_derivatives!(ts::TearingState, neweqs, dummy_sub, var_
     end
 end
 
-"""
-Generate new derivative variables for the system.
-
+#= 
 There are three cases where we want to generate new variables to convert
 the system into first order (semi-implicit) ODEs.
     
@@ -361,6 +361,16 @@ This is different from the continuous case where D(D(x)) can be substituted for
 by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
 shift_sub dict is updated at the time that the renamed variables are written,
 inside the loop where new variables are generated.
+=#
+"""
+Generate new derivative variables for the system.
+
+Effects on the state: 
+- fullvars: add the new derivative variables x_t
+- neweqs: add the identity equations for the new variables, D(x) ~ x_t
+- graph: update graph with the new equations and variables, and their connections
+- solvable_graph:
+- var_eq_matching: solvable equations
 """
 function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching;
         is_discrete = false, mm = nothing, shift_sub = nothing)
@@ -406,7 +416,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
             dx = fullvars[dv]
             order, lv = var_order(diff_to_var, dv)
             @show fullvars[uv]
-            x_t = lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv])
+            x_t = lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv], allow_zero = true)
             shift_sub[dx] = x_t
             (var_eq_matching[dv] isa Int) ? continue : @goto DISCRETE_VARIABLE
         end
@@ -439,7 +449,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         x_t = lower_name(fullvars[lv], iv, order)
         
         @label DISCRETE_VARIABLE
-        push!(fullvars, x_t)
+        push!(fullvars, simplify_shifts(x_t))
         v_t = length(fullvars)
         v_t_idx = add_vertex!(var_to_diff)
         add_vertex!(graph, DST)
@@ -448,7 +458,6 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         add_vertex!(solvable_graph, DST)
         push!(var_eq_matching, unassigned)
 
-        # Add discrete substitutions to total_sub directly.  
         is_discrete && begin
             idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
             for e in 𝑑neighbors(graph, dv)
@@ -463,7 +472,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         end
 
         # add `D(x) - x_t ~ 0`
-        push!(neweqs, 0 ~ x_t - dx)
+        push!(neweqs, 0 ~ dx - x_t)
         add_vertex!(graph, SRC)
         dummy_eq = length(neweqs)
         add_edge!(graph, dummy_eq, dv)
@@ -478,12 +487,11 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
     end
 end
 
-function add_solvable_variable!() 
+function add_solvable_variable!(state::TearingState)
     
 end
 
-function add_solvable_equation!() 
-    
+function add_solvable_equation!(s::SystemStructure, neweqs, eq)
 end
 
 """
@@ -500,7 +508,8 @@ such that the mass matrix is:
 Update the state to account for the new ordering and equations.
 
 ####### DISCRETE CASE
-- only substitute Shift(t, -2) 
+- Differential equations: substitute variables with everything shifted forward one timestep.
+- Algebraic and observable equations: substitute variables with everything shifted back one timestep.
 """
 function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false, shift_sub = Dict())
     @unpack fullvars, sys, structure = state 
@@ -562,8 +571,6 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
                 isnothing(D) &&
                     error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(equations(sys)[ieq])")
                 order, lv = var_order(diff_to_var, iv)
-                @show fullvars[lv]
-                @show simplify_shifts(fullvars[lv])
                 dx = D(simplify_shifts(fullvars[lv]))
                 eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
                     Symbolics.symbolic_linear_solve(neweqs[ieq],
@@ -576,8 +583,6 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
                 push!(diff_eqs, eq)
                 total_sub[simplify_shifts(eq.lhs)] = eq.rhs
                 dx_sub[simplify_shifts(eq.lhs)] = eq.rhs
-                @show total_sub
-                @show eq
                 push!(diffeq_idxs, ieq)
                 push!(diff_vars, diff_to_var[iv])
                 continue
@@ -611,6 +616,8 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
             push!(algeeq_idxs, ieq)
         end
     end
+    @show neweqs
+    @show subeqs
 
     # TODO: BLT sorting
     neweqs = [diff_eqs; alge_eqs]
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index c202556906..b947fa4269 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -271,6 +271,8 @@ end
 
 function find_solvables!(state::TearingState; kwargs...)
     @assert state.structure.solvable_graph === nothing
+    println("in find_solvables")
+    @show eqs
     eqs = equations(state)
     graph = state.structure.graph
     state.structure.solvable_graph = BipartiteGraph(nsrcs(graph), ndsts(graph))
@@ -278,6 +280,7 @@ function find_solvables!(state::TearingState; kwargs...)
     for ieq in 1:length(eqs)
         find_eq_solvables!(state, ieq, to_rm; kwargs...)
     end
+    @show eqs
     return nothing
 end
 
@@ -477,7 +480,7 @@ end
 ### Rules
 # 1. x(t) -> x(t)
 # 2. Shift(t, 0)(x(t)) -> x(t)
-# 3. Shift(t, 1)(x + z) -> Shift(t, 1)(x) + Shift(t, 1)(z)
+# 3. Shift(t, 3)(Shift(t, 2)(x(t)) -> Shift(t, 5)(x(t))
 
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
@@ -498,3 +501,58 @@ function simplify_shifts(var)
             unwrap(var).metadata)
     end
 end
+
+"""
+Power expand the shifts. Used for substitution.
+
+Shift(t, -3)(x(t)) -> Shift(t, -1)(Shift(t, -1)(Shift(t, -1)(x)))
+"""
+function expand_shifts(var)
+    ModelingToolkit.hasshift(var) || return var
+    var = ModelingToolkit.value(var)
+
+    var isa Equation && return expand_shifts(var.lhs) ~ expand_shifts(var.rhs)
+    op = operation(var)
+    s = sign(op.steps)
+    arg = only(arguments(var))
+
+    if ModelingToolkit.isvariable(arg) && (ModelingToolkit.getvariabletype(arg) === VARIABLE) && isequal(op.t, only(arguments(arg)))
+        out = arg
+        for i in 1:op.steps
+            out = Shift(op.t, s)(out)
+        end
+        return out
+    elseif iscall(arg)
+        return maketerm(typeof(var), operation(var), expand_shifts.(arguments(var)),
+            unwrap(var).metadata)
+    else
+        return arg
+    end
+end
+
+"""
+Shift(t, 1)(x + z) -> Shift(t, 1)(x) + Shift(t, 1)(z)
+"""
+function distribute_shift(var) 
+    ModelingToolkit.hasshift(var) || return var
+    var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
+    shift = operation(var)
+    expr = only(arguments(var))
+    _distribute_shift(expr, shift)
+end
+
+function _distribute_shift(expr, shift)
+    op = operation(expr)
+    args = arguments(expr)
+
+    if length(args) == 1 
+        if ModelingToolkit.isvariable(only(args)) && isequal(op.t, only(args))
+            return shift(only(args))
+        else
+            return only(args)
+        end
+    else iscall(op)
+        return maketerm(typeof(expr), operation(expr), _distribute_shift.(args, shift),
+            unwrap(var).metadata)
+    end
+end
diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl
index 621aa8b8a8..98d986ebd4 100644
--- a/test/structural_transformation/utils.jl
+++ b/test/structural_transformation/utils.jl
@@ -4,6 +4,7 @@ using Graphs
 using SparseArrays
 using UnPack
 using ModelingToolkit: t_nounits as t, D_nounits as D
+const ST = StructuralTransformations
 
 # Define some variables
 @parameters L g
@@ -161,3 +162,19 @@ end
     structural_simplify(sys; additional_passes = [pass])
     @test value[] == 1
 end
+
+@testset "Shift simplification" begin
+    @variables x(t) y(t) z(t)
+    @parameters a b c
+    
+    # Expand shifts
+    @test isequal(ST.expand_shifts(Shift(t, -3)(x)), Shift(t, -1)(Shift(t, -1)(Shift(t, -1)(x))))
+    expr = a * Shift(t, -2)(x) + Shift(t, 2)(y) + b
+    @test isequal(ST.expand_shifts(expr), 
+                    a * Shift(t, -1)(Shift(t, -1)(x)) + Shift(t, 1)(Shift(t, 1)(y)) + b)
+    @test isequal(ST.expand_shifts(Shift(t, 2)(Shift(t, 1)(a))), a)
+
+
+    # Distribute shifts
+
+end

From ed143a5d85d4b4660d478718fc9bcd9d79ec810c Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 10 Feb 2025 15:18:05 -0500
Subject: [PATCH 23/41] refactor: refactor tearing_assemble into functions

---
 .../symbolics_tearing.jl                      | 503 ++++++++----------
 src/structural_transformation/utils.jl        |  71 +--
 src/systems/systemstructure.jl                |   6 +-
 3 files changed, 239 insertions(+), 341 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 778473907e..b99e6ec21a 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -244,19 +244,17 @@ State selection may determine that some differential variables are
 algebraic variables in disguise. The derivative of such variables are
 called dummy derivatives.
 
-`SelectedState` information is no longer needed past here. State selection
-is done. All non-differentiated variables are algebraic variables, and all
-variables that appear differentiated are differential variables.
+`SelectedState` information is no longer needed after this function is called. 
+State selection is done. All non-differentiated variables are algebraic 
+variables, and all variables that appear differentiated are differential variables.
 """
-function substitute_dummy_derivatives!(ts::TearingState, neweqs, dummy_sub, var_eq_matching)
+function substitute_derivatives_algevars!(ts::TearingState, neweqs, dummy_sub, var_eq_matching)
     @unpack fullvars, sys, structure = ts
-    @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     diff_to_var = invview(var_to_diff)
     iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
 
-    @show neweqs
     for var in 1:length(fullvars)
-        #@show neweqs
         dv = var_to_diff[var]
         dv === nothing && continue
         if var_eq_matching[var] !== SelectedState()
@@ -339,199 +337,137 @@ variables and equations, don't add them when they already exist.
 
 Documenting the differences to structural simplification for discrete systems:
 
-1. In discrete systems the lowest-order term is Shift(t, k)(x(t)), instead of x(t). We need to substitute the k-1 lowest order terms instead of the k-1 highest order terms.
+In discrete systems the lowest-order term is Shift(t, k)(x(t)), instead of x(t). 
+We want to substitute the k-1 lowest order terms instead of the k-1 highest order terms.
 
-The orders will also be off by one. The reason this is is that the dynamics of
-the system should be given in terms of Shift(t, 1)(x(t), x(t-1), ...). But 
-having the observables be indexed by the next time step is not so nice. So we 
-handle the shifts in the renaming.
-
-The substitution should look like the following: 
-  x(t) -> Shift(t, 1)(x(t))
-  Shift(t, -1)(x(t)) -> Shift(t, 0)(x(t))
-  Shift(t, -2)(x(t)) -> x_{t-1}(t)
-  Shift(t, -3)(x(t)) -> x_{t-2}(t)
-  and so on...
-
-In the implicit discrete case this shouldn't happen. The simplification should 
-look like a NonlinearSystem.
-
-2. For discrete systems Shift(t, 2)(x(t)) cannot be substituted as Shift(t, 1)(Shift(t,1)(x(t)). 
-This is different from the continuous case where D(D(x)) can be substituted for 
-by iteratively substituting x_t ~ D(x), then x_tt ~ D(x_t). For this reason the
-shift_sub dict is updated at the time that the renamed variables are written,
-inside the loop where new variables are generated.
+In the system x(k) ~ x(k-1) + x(k-2), we want to lower
+Shift(t, -1)(x(t)) -> x\_{t-1}(t)
 =#
 """
 Generate new derivative variables for the system.
 
-Effects on the state: 
+Effects on the system structure: 
 - fullvars: add the new derivative variables x_t
 - neweqs: add the identity equations for the new variables, D(x) ~ x_t
 - graph: update graph with the new equations and variables, and their connections
 - solvable_graph:
-- var_eq_matching: solvable equations
+- var_eq_matching: match D(x) to the added identity equation
 """
 function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching;
-        is_discrete = false, mm = nothing, shift_sub = nothing)
+        is_discrete = false, mm = nothing)
     @unpack fullvars, sys, structure = ts
-    @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
     iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
     lower_name = is_discrete ? lower_varname_withshift : lower_varname_with_unit
 
-    # index v gets mapped to the lowest shift and the index of the unshifted variable
-    if is_discrete
-        idx_to_lowest_shift = Dict{Int, Tuple{Int, Int}}(var => (0,0) for var in 1:length(fullvars))
-        var_to_unshiftedidx = Dict{Any, Int}(var => findfirst(x -> isequal(x, var), fullvars) for var in keys(lowest_shift))
-
-        for (i,var) in enumerate(fullvars)
-            key = (operation(var) isa Shift) ? only(arguments(var)) : var
-            idx_to_lowest_shift[i] = (get(lowest_shift, key, 0), get(var_to_unshiftedidx, key, i))
-        end
-    end
-
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
-    # v is the index of the current variable, x = fullvars[v]
-    # dv is the index of the derivative dx = D(x), x_t is the substituted variable 
-    # For ODESystems: lv is the index of the lowest-order variable (x(t))
-    # For DiscreteSystems: 
-    # - lv is the index of the lowest-order variable (Shift(t, k)(x(t)))
-    # - uv is the index of the highest-order variable (x(t))
+    # Generate new derivative variables for all unsolved variables that have a derivative in the system 
     for v in 1:length(var_to_diff)
+        # Check if a derivative 1) exists and 2) is unsolved for
         dv = var_to_diff[v]
-
-        if is_discrete 
-            x = fullvars[v]
-            op = operation(x)
-            (low, uv) = idx_to_lowest_shift[v]
-
-            # If v is unshifted (i.e. x(t)), then substitute the lowest-shift variable
-            if !(op isa Shift)
-                dv = findfirst(_x -> isequal(_x, Shift(iv, low)(x)), fullvars)
-            end
-            dx = fullvars[dv]
-            order, lv = var_order(diff_to_var, dv)
-            @show fullvars[uv]
-            x_t = lower_name(fullvars[lv], iv, -low-order-1; unshifted = fullvars[uv], allow_zero = true)
-            shift_sub[dx] = x_t
-            (var_eq_matching[dv] isa Int) ? continue : @goto DISCRETE_VARIABLE
-        end
         dv isa Int || continue
         solved = var_eq_matching[dv] isa Int
         solved && continue
 
-        # Check if there's `D(x) = x_t` already
-        local v_t, dummy_eq
-        for eq in 𝑑neighbors(solvable_graph, dv)
-            mi = get(linear_eqs, eq, 0)
-            iszero(mi) && continue
-            row = @view mm[mi, :]
-            nzs = nonzeros(row)
-            rvs = SparseArrays.nonzeroinds(row)
-            # note that `v_t` must not be differentiated
-            if length(nzs) == 2 &&
-               (abs(nzs[1]) == 1 && nzs[1] == -nzs[2]) &&
-               (v_t = rvs[1] == dv ? rvs[2] : rvs[1];
-               diff_to_var[v_t] === nothing)
-                @assert dv in rvs
-                dummy_eq = eq
-                @goto FOUND_DUMMY_EQ
-            end
+        # If there's `D(x) = x_t` already, update mappings and continue without
+        # adding new equations/variables
+        dd = find_duplicate_dd(dv, lineareqs, mm)
+
+        if !isnothing(dd)
+            dummy_eq, v_t = dd
+            var_to_diff[v_t] = var_to_diff[dv]
+            var_eq_matching[dv] = unassigned
+            eq_var_matching[dummy_eq] = dv
+            continue
         end
 
-        # add `x_t`
         dx = fullvars[dv]
         order, lv = var_order(diff_to_var, dv)
-        x_t = lower_name(fullvars[lv], iv, order)
+        x_t = is_discrete ? lower_name(fullvars[lv], iv)
+                      : lower_name(fullvars[lv], iv, order)
         
-        @label DISCRETE_VARIABLE
-        push!(fullvars, simplify_shifts(x_t))
-        v_t = length(fullvars)
-        v_t_idx = add_vertex!(var_to_diff)
-        add_vertex!(graph, DST)
-        # TODO: do we care about solvable_graph? We don't use them after
-        # `dummy_derivative_graph`.
-        add_vertex!(solvable_graph, DST)
-        push!(var_eq_matching, unassigned)
-
-        is_discrete && begin
-            idx_to_lowest_shift[v_t] = idx_to_lowest_shift[dv]
-            for e in 𝑑neighbors(graph, dv)
-                add_edge!(graph, e, v_t)
-                rem_edge!(graph, e, dv)
-            end
-            # Do not add the lowest-order substitution as an equation, just substitute
-            !(operation(x) isa Shift) && begin
-                var_to_diff[v_t] = var_to_diff[dv]
-                continue
-            end
-        end
+        # Add `x_t` to the graph
+        add_dd_variable!(structure, x_t, dv)
+        # Add `D(x) - x_t ~ 0` to the graph
+        add_dd_equation!(structure, neweqs, 0 ~ dx - x_t, dv)
 
-        # add `D(x) - x_t ~ 0`
-        push!(neweqs, 0 ~ dx - x_t)
-        add_vertex!(graph, SRC)
-        dummy_eq = length(neweqs)
-        add_edge!(graph, dummy_eq, dv)
-        add_edge!(graph, dummy_eq, v_t)
-        add_vertex!(solvable_graph, SRC)
-        add_edge!(solvable_graph, dummy_eq, dv)
-
-        @label FOUND_DUMMY_EQ
-        var_to_diff[v_t] = var_to_diff[dv]
+        # Update matching
+        push!(var_eq_matching, unassigned)
         var_eq_matching[dv] = unassigned
         eq_var_matching[dummy_eq] = dv
     end
 end
 
-function add_solvable_variable!(state::TearingState)
-    
+"""
+Check if there's `D(x) = x_t` already. 
+"""
+function find_duplicate_dd(dv, lineareqs, mm)
+    for eq in 𝑑neighbors(solvable_graph, dv)
+        mi = get(linear_eqs, eq, 0)
+        iszero(mi) && continue
+        row = @view mm[mi, :]
+        nzs = nonzeros(row)
+        rvs = SparseArrays.nonzeroinds(row)
+        # note that `v_t` must not be differentiated
+        if length(nzs) == 2 &&
+           (abs(nzs[1]) == 1 && nzs[1] == -nzs[2]) &&
+           (v_t = rvs[1] == dv ? rvs[2] : rvs[1];
+           diff_to_var[v_t] === nothing)
+            @assert dv in rvs
+            return eq, v_t
+        end
+    end
+    return nothing 
+end
+
+function add_dd_variable!(s::SystemStructure, x_t, dv)
+    push!(s.fullvars, simplify_shifts(x_t))
+    v_t_idx = add_vertex!(s.var_to_diff)
+    @assert v_t_idx == ndsts(graph) == ndsts(solvable_graph) == length(fullvars) ==
+        length(var_eq_matching)
+    add_vertex!(s.graph, DST)
+    # TODO: do we care about solvable_graph? We don't use them after
+    # `dummy_derivative_graph`.
+    add_vertex!(s.solvable_graph, DST)
+    var_to_diff[v_t] = var_to_diff[dv]
 end
 
-function add_solvable_equation!(s::SystemStructure, neweqs, eq)
+# dv = index of D(x), v_t = index of x_t
+function add_dd_equation!(s::SystemStructure, neweqs, eq, dv)
+    push!(neweqs, eq)
+    add_vertex!(s.graph, SRC)
+    v_t = length(s.fullvars)
+    dummy_eq = length(neweqs)
+    add_edge!(s.graph, dummy_eq, dv)
+    add_edge!(s.graph, dummy_eq, v_t)
+    add_vertex!(s.solvable_graph, SRC)
+    add_edge!(s.solvable_graph, dummy_eq, dv)
 end
 
 """
 Solve the solvable equations of the system and generate differential (or discrete)
 equations in terms of the selected states.
-
-Will reorder equations and unknowns to be:
-   [diffeqs; ...]
-   [diffvars; ...]
-such that the mass matrix is:
-   [I  0
-    0  0].
-
-Update the state to account for the new ordering and equations.
-
-####### DISCRETE CASE
-- Differential equations: substitute variables with everything shifted forward one timestep.
-- Algebraic and observable equations: substitute variables with everything shifted back one timestep.
 """
-function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false, shift_sub = Dict())
+function generate_system_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false)
     @unpack fullvars, sys, structure = state 
-    @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
-    dx_sub = Dict()
+    total_sub = Dict()
 
     if ModelingToolkit.has_iv(sys)
         iv = get_iv(sys)
         if is_only_discrete(structure)
             D = Shift(iv, 1)
-            lower_name = lower_varname_withshift
-            total_sub = shift_sub
         else
             D = Differential(iv)
-            lower_name = lower_varname_with_unit
-            total_sub = dx_sub
         end
     else
         iv = D = nothing
-        lower_name = lower_varname_with_unit
     end
 
     # if var is like D(x) or Shift(t, 1)(x)
@@ -544,14 +480,14 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
         (eq, iv) -> eq isa Int && iv isa Int && BipartiteEdge(eq, iv) in solvable_graph
     end
 
-    diffeq_idxs = Int[]
-    algeeq_idxs = Int[]
     diff_eqs = Equation[]
-    alge_eqs = Equation[]
+    diffeq_idxs = Int[]
     diff_vars = Int[]
-    subeqs = Equation[]
-    solved_equations = Int[]
-    solved_variables = Int[]
+    alge_eqs = Equation[]
+    algeeq_idxs = Int[]
+    solved_eqs = Equation[]
+    solvedeq_idxs = Int[]
+    solved_vars = Int[]
 
     toporder = topological_sort(DiCMOBiGraph{false}(graph, var_eq_matching))
     eqs = Iterators.reverse(toporder)
@@ -561,100 +497,126 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
     # fullvars[iv] is a differential variable of the form D^n(x), and neweqs[ieq]
     # is solved to give the RHS.
     for ieq in eqs
-        println()
         iv = eq_var_matching[ieq]
         if is_solvable(ieq, iv)
-            # We don't solve differential equations, but we will need to try to
-            # convert it into the mass matrix form.
-            # We cannot solve the differential variable like D(x)
             if isdervar(iv)
                 isnothing(D) &&
                     error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(equations(sys)[ieq])")
-                order, lv = var_order(diff_to_var, iv)
-                dx = D(simplify_shifts(fullvars[lv]))
-                eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
-                    Symbolics.symbolic_linear_solve(neweqs[ieq],
-                        fullvars[iv]),
-                    total_sub; operator = ModelingToolkit.Shift))
-                for e in 𝑑neighbors(graph, iv)
-                    rem_edge!(graph, e, iv)
-                    add_edge!(graph, e, lv)
-                end
-                push!(diff_eqs, eq)
-                total_sub[simplify_shifts(eq.lhs)] = eq.rhs
-                dx_sub[simplify_shifts(eq.lhs)] = eq.rhs
-                push!(diffeq_idxs, ieq)
-                push!(diff_vars, diff_to_var[iv])
-                continue
-            end
-            eq = neweqs[ieq]
-            var = fullvars[iv]
-            residual = eq.lhs - eq.rhs
-            a, b, islinear = linear_expansion(residual, var)
-            @assert islinear
-            # 0 ~ a * var + b
-            # var ~ -b/a
-            if ModelingToolkit._iszero(a)
-                @warn "Tearing: solving $eq for $var is singular!"
+                add_differential_equation!(structure, iv, neweqs, ieq, 
+                                           diff_vars, diff_eqs, diffeq_idxs, total_sub)
             else
-                rhs = -b / a
-                neweq = var ~ simplify_shifts(Symbolics.fixpoint_sub(
-                    simplify ?
-                    Symbolics.simplify(rhs) : rhs,
-                    dx_sub; operator = ModelingToolkit.Shift))
-                push!(subeqs, neweq)
-                push!(solved_equations, ieq)
-                push!(solved_variables, iv)
+                add_solved_equation!(structure, iv, neweqs, ieq, 
+                                     solved_vars, solved_eqs, solvedeq_idxs, total_sub)
             end
         else
-            eq = neweqs[ieq]
-            rhs = eq.rhs
-            if !(eq.lhs isa Number && eq.lhs == 0)
-                rhs = eq.rhs - eq.lhs
-            end
-            push!(alge_eqs, 0 ~ simplify_shifts(Symbolics.fixpoint_sub(rhs, total_sub)))
-            push!(algeeq_idxs, ieq)
+            add_algebraic_equation!(structure, neweqs, ieq, 
+                                    alge_eqs, algeeq_idxs, total_sub)
         end
     end
-    @show neweqs
-    @show subeqs
 
-    # TODO: BLT sorting
+    # Generate new equations and orderings 
     neweqs = [diff_eqs; alge_eqs]
-    inveqsperm = [diffeq_idxs; algeeq_idxs]
-    eqsperm = zeros(Int, nsrcs(graph))
-    for (i, v) in enumerate(inveqsperm)
-        eqsperm[v] = i
-    end
+    eq_ordering = [diffeq_idxs; algeeq_idxs]
     diff_vars_set = BitSet(diff_vars)
     if length(diff_vars_set) != length(diff_vars)
         error("Tearing internal error: lowering DAE into semi-implicit ODE failed!")
     end
-    solved_variables_set = BitSet(solved_variables)
-    invvarsperm = [diff_vars;
+    solved_vars_set = BitSet(solved_vars)
+    var_ordering = [diff_vars;
                    setdiff!(setdiff(1:ndsts(graph), diff_vars_set),
-                       solved_variables_set)]
+                       solved_vars_set)]
+
+    return neweqs, solved_eqs, eq_ordering, var_ordering, length(solved_vars), length(solved_vars_set)
+end
+
+function add_differential_equation!(s::SystemStructure, iv, neweqs, ieq, diff_vars, diff_eqs, diffeqs_idxs, total_sub)
+    diff_to_var = invview(s.var_to_diff)
+
+    order, lv = var_order(diff_to_var, iv)
+    dx = D(simplify_shifts(fullvars[lv]))
+    eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
+        Symbolics.symbolic_linear_solve(neweqs[ieq],
+            fullvars[iv]),
+        total_sub; operator = ModelingToolkit.Shift))
+    for e in 𝑑neighbors(s.graph, iv)
+        e == ieq && continue
+        rem_edge!(s.graph, e, iv)
+    end
+
+    push!(diff_eqs, eq)
+    total_sub[simplify_shifts(eq.lhs)] = eq.rhs
+    push!(diffeq_idxs, ieq)
+    push!(diff_vars, diff_to_var[iv])
+end
+
+function add_algebraic_equation!(s::SystemStructure, neweqs, ieq, alge_eqs, algeeq_idxs, total_sub)
+    eq = neweqs[ieq]
+    rhs = eq.rhs
+    if !(eq.lhs isa Number && eq.lhs == 0)
+        rhs = eq.rhs - eq.lhs
+    end
+    push!(alge_eqs, 0 ~ simplify_shifts(Symbolics.fixpoint_sub(rhs, total_sub)))
+    push!(algeeq_idxs, ieq)
+end
+
+function add_solved_equation!(s::SystemStructure, iv, neweqs, ieq, solved_vars, solved_eqs, solvedeq_idxs, total_sub)
+    eq = neweqs[ieq]
+    var = fullvars[iv]
+    residual = eq.lhs - eq.rhs
+    a, b, islinear = linear_expansion(residual, var)
+    @assert islinear
+    # 0 ~ a * var + b
+    # var ~ -b/a
+    if ModelingToolkit._iszero(a)
+        @warn "Tearing: solving $eq for $var is singular!"
+    else
+        rhs = -b / a
+        neweq = var ~ simplify_shifts(Symbolics.fixpoint_sub(
+            simplify ?
+            Symbolics.simplify(rhs) : rhs,
+            total_sub; operator = ModelingToolkit.Shift))
+        push!(solved_eqs, neweq)
+        push!(solvedeq_idxs, ieq)
+        push!(solved_vars, iv)
+    end
+end
+
+"""
+Reorder the equations and unknowns to be:
+   [diffeqs; ...]
+   [diffvars; ...]
+such that the mass matrix is:
+   [I  0
+    0  0].
+
+Update the state to account for the new ordering and equations.
+"""
+# TODO: BLT sorting
+function reorder_vars!(s::SystemStructure, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
+
+    eqsperm = zeros(Int, nsrcs(graph))
+    for (i, v) in enumerate(eq_ordering)
+        eqsperm[v] = i
+    end
     varsperm = zeros(Int, ndsts(graph))
-    for (i, v) in enumerate(invvarsperm)
+    for (i, v) in enumerate(var_ordering)
         varsperm[v] = i
     end
 
-    deps = Vector{Int}[i == 1 ? Int[] : collect(1:(i - 1))
-                       for i in 1:length(solved_equations)]
-
     # Contract the vertices in the structure graph to make the structure match
     # the new reality of the system we've just created.
-    graph = contract_variables(graph, var_eq_matching, varsperm, eqsperm,
-        length(solved_variables), length(solved_variables_set))
+    new_graph = contract_variables(graph, var_eq_matching, varsperm, eqsperm,
+        nelim_eq, nelim_var)
 
-    new_var_to_diff = complete(DiffGraph(length(invvarsperm)))
+    new_var_to_diff = complete(DiffGraph(length(var_ordering)))
     for (v, d) in enumerate(var_to_diff)
         v′ = varsperm[v]
         (v′ > 0 && d !== nothing) || continue
         d′ = varsperm[d]
         new_var_to_diff[v′] = d′ > 0 ? d′ : nothing
     end
-    new_eq_to_diff = complete(DiffGraph(length(inveqsperm)))
+    new_eq_to_diff = complete(DiffGraph(length(eq_ordering)))
     for (v, d) in enumerate(eq_to_diff)
         v′ = eqsperm[v]
         (v′ > 0 && d !== nothing) || continue
@@ -663,7 +625,53 @@ function solve_and_generate_equations!(state::TearingState, neweqs, var_eq_match
     end
     new_fullvars = fullvars[invvarsperm]
 
-    new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph
+    # Update system structure 
+    @set! state.structure.graph = complete(new_graph)
+    @set! state.structure.var_to_diff = new_var_to_diff
+    @set! state.structure.eq_to_diff = new_eq_to_diff
+    @set! state.fullvars = new_fullvars
+end
+
+"""
+Set the system equations, unknowns, observables post-tearing.
+"""
+function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns)
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
+    diff_to_var = invview(var_to_diff)
+
+    ispresent = let var_to_diff = var_to_diff, graph = graph
+        i -> (!isempty(𝑑neighbors(graph, i)) ||
+              (var_to_diff[i] !== nothing && !isempty(𝑑neighbors(graph, var_to_diff[i]))))
+    end
+
+    sys = state.sys
+    obs_sub = dummy_sub
+    for eq in neweqs
+        isdiffeq(eq) || continue
+        obs_sub[eq.lhs] = eq.rhs
+    end
+    # TODO: compute the dependency correctly so that we don't have to do this
+    obs = [fast_substitute(observed(sys), obs_sub); solved_eqs]
+
+    unknowns = Any[v
+                   for (i, v) in enumerate(state.fullvars)
+                   if diff_to_var[i] === nothing && ispresent(i)]
+    unknowns = [unknowns; extra_unknowns]
+    @set! sys.unknowns = unknowns
+
+    obs, subeqs, deps = cse_and_array_hacks(
+        sys, obs, solved_eqs, unknowns, neweqs; cse = cse_hack, array = array_hack)
+
+    @set! sys.eqs = neweqs
+    @set! sys.observed = obs
+    @set! sys.substitutions = Substitutions(subeqs, deps)
+
+    # Only makes sense for time-dependent
+    # TODO: generalize to SDE
+    if sys isa ODESystem
+        @set! sys.schedule = Schedule(var_eq_matching, dummy_sub)
+    end
+    sys = schedule(sys)
 end
 
 # Terminology and Definition:
@@ -675,10 +683,8 @@ end
 # appear in the system. Algebraic variables are variables that are not
 # differential variables.
     
-import ModelingToolkit: Shift
-
 # Give the order of the variable indexed by dv
-function var_order(diff_to_var, dv) 
+function var_order(diff_to_var, dv)
     order = 0
     while (dv′ = diff_to_var[dv]) !== nothing
         order += 1
@@ -689,8 +695,7 @@ end
 
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
-    @unpack fullvars, sys, structure = state
-    @unpack solvable_graph, var_to_diff, eq_to_diff, graph, lowest_shift = structure
+
     extra_vars = Int[]
     if full_var_eq_matching !== nothing
         for v in 𝑑vertices(state.structure.graph)
@@ -699,69 +704,21 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
             push!(extra_vars, v)
         end
     end
+    extra_unknowns = fullvars[extra_vars]
     neweqs = collect(equations(state))
-    diff_to_var = invview(var_to_diff)
-    is_discrete = is_only_discrete(state.structure)
-
-    shift_sub = Dict() 
 
     # Structural simplification 
     dummy_sub = Dict()
-    substitute_dummy_derivatives!(state, neweqs, dummy_sub, var_eq_matching)
-
-    generate_derivative_variables!(state, neweqs, var_eq_matching; 
-                                   is_discrete, mm, shift_sub)
-
-    new_fullvars, new_var_to_diff, new_eq_to_diff, neweqs, subeqs, graph = 
-        solve_and_generate_equations!(state, neweqs, var_eq_matching; simplify, shift_sub)
-
-    # Update system
-    var_to_diff = new_var_to_diff
-    eq_to_diff = new_eq_to_diff
-    diff_to_var = invview(var_to_diff)
-
-    old_fullvars = fullvars
-    @set! state.structure.graph = complete(graph)
-    @set! state.structure.var_to_diff = var_to_diff
-    @set! state.structure.eq_to_diff = eq_to_diff
-    @set! state.fullvars = fullvars = new_fullvars
-    ispresent = let var_to_diff = var_to_diff, graph = graph
-        i -> (!isempty(𝑑neighbors(graph, i)) ||
-              (var_to_diff[i] !== nothing && !isempty(𝑑neighbors(graph, var_to_diff[i]))))
-    end
-
-    sys = state.sys
-    obs_sub = dummy_sub
-    for eq in neweqs
-        isdiffeq(eq) || continue
-        obs_sub[eq.lhs] = eq.rhs
-    end
-    # TODO: compute the dependency correctly so that we don't have to do this
-    obs = [fast_substitute(observed(sys), obs_sub); subeqs]
+    substitute_derivatives_algevars!(state, neweqs, dummy_sub, var_eq_matching)
 
-    unknowns = Any[v
-                   for (i, v) in enumerate(fullvars)
-                   if diff_to_var[i] === nothing && ispresent(i)]
-    if !isempty(extra_vars)
-        for v in extra_vars
-            push!(unknowns, old_fullvars[v])
-        end
-    end
-    @set! sys.unknowns = unknowns
+    generate_derivative_variables!(state, neweqs, var_eq_matching; mm)
 
-    obs, subeqs, deps = cse_and_array_hacks(
-        sys, obs, subeqs, unknowns, neweqs; cse = cse_hack, array = array_hack)
+    neweqs, solved_eqs, eq_ordering, var_ordering, nelim_eq, nelim_var = 
+        generate_system_equations!(state, neweqs, var_eq_matching; simplify)
 
-    @set! sys.eqs = neweqs
-    @set! sys.observed = obs
-    @set! sys.substitutions = Substitutions(subeqs, deps)
+    reorder_vars!(state.structure, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
 
-    # Only makes sense for time-dependent
-    # TODO: generalize to SDE
-    if sys isa ODESystem
-        @set! sys.schedule = Schedule(var_eq_matching, dummy_sub)
-    end
-    sys = schedule(sys)
+    sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns)
     @set! state.sys = sys
     @set! sys.tearing_state = state
     return invalidate_cache!(sys)
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index b947fa4269..77f0b62a1c 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -452,9 +452,13 @@ end
 ### Misc
 ###
 
-function lower_varname_withshift(var, iv, backshift; unshifted = nothing, allow_zero = true)
-    backshift <= 0 && return Shift(iv, -backshift)(unshifted, allow_zero)
-    ds = backshift > 0 ? "$iv-$backshift" : "$iv+$(-backshift)"
+function lower_varname_withshift(var, iv)
+    op = operation(var)
+    op isa Shift || return var
+    backshift = op.steps
+    backshift > 0 && return var
+
+    ds = "$iv-$(-backshift)"
     d_separator = 'ˍ'
 
     if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
@@ -477,15 +481,9 @@ function isdoubleshift(var)
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
 end
 
-### Rules
-# 1. x(t) -> x(t)
-# 2. Shift(t, 0)(x(t)) -> x(t)
-# 3. Shift(t, 3)(Shift(t, 2)(x(t)) -> Shift(t, 5)(x(t))
-
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
     var isa Equation && return simplify_shifts(var.lhs) ~ simplify_shifts(var.rhs)
-    ((op = operation(var)) isa Shift) && op.steps == 0 && return simplify_shifts(arguments(var)[1])
     if isdoubleshift(var)
         op1 = operation(var)
         vv1 = arguments(var)[1]
@@ -501,58 +499,3 @@ function simplify_shifts(var)
             unwrap(var).metadata)
     end
 end
-
-"""
-Power expand the shifts. Used for substitution.
-
-Shift(t, -3)(x(t)) -> Shift(t, -1)(Shift(t, -1)(Shift(t, -1)(x)))
-"""
-function expand_shifts(var)
-    ModelingToolkit.hasshift(var) || return var
-    var = ModelingToolkit.value(var)
-
-    var isa Equation && return expand_shifts(var.lhs) ~ expand_shifts(var.rhs)
-    op = operation(var)
-    s = sign(op.steps)
-    arg = only(arguments(var))
-
-    if ModelingToolkit.isvariable(arg) && (ModelingToolkit.getvariabletype(arg) === VARIABLE) && isequal(op.t, only(arguments(arg)))
-        out = arg
-        for i in 1:op.steps
-            out = Shift(op.t, s)(out)
-        end
-        return out
-    elseif iscall(arg)
-        return maketerm(typeof(var), operation(var), expand_shifts.(arguments(var)),
-            unwrap(var).metadata)
-    else
-        return arg
-    end
-end
-
-"""
-Shift(t, 1)(x + z) -> Shift(t, 1)(x) + Shift(t, 1)(z)
-"""
-function distribute_shift(var) 
-    ModelingToolkit.hasshift(var) || return var
-    var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
-    shift = operation(var)
-    expr = only(arguments(var))
-    _distribute_shift(expr, shift)
-end
-
-function _distribute_shift(expr, shift)
-    op = operation(expr)
-    args = arguments(expr)
-
-    if length(args) == 1 
-        if ModelingToolkit.isvariable(only(args)) && isequal(op.t, only(args))
-            return shift(only(args))
-        else
-            return only(args)
-        end
-    else iscall(op)
-        return maketerm(typeof(expr), operation(expr), _distribute_shift.(args, shift),
-            unwrap(var).metadata)
-    end
-end
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 89fcebf549..e5a227d9fb 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -154,7 +154,6 @@ Base.@kwdef mutable struct SystemStructure
     var_types::Union{Vector{VariableType}, Nothing}
     """Whether the system is discrete."""
     only_discrete::Bool
-    lowest_shift::Union{Dict, Nothing}
 end
 
 function Base.copy(structure::SystemStructure)
@@ -436,7 +435,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, lowest_shift),
+            complete(graph), nothing, var_types, sys isa DiscreteSystem),
         Any[])
     if sys isa DiscreteSystem
         ts = shift_discrete_system(ts, lowest_shift)
@@ -466,9 +465,8 @@ end
     Shift variable x by the largest shift s such that x(k-s) appears in the system of equations.
     The lowest-shift term will have.
 """
-function shift_discrete_system(ts::TearingState, lowest_shift)
+function shift_discrete_system(ts::TearingState)
     @unpack fullvars, sys = ts
-    return ts
     discvars = OrderedSet()
     eqs = equations(sys)
 

From 240ab215b2ec457186b889ca53160e9d23db02e8 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 10 Feb 2025 18:38:22 -0500
Subject: [PATCH 24/41] fix: properly rename variables inside
 generate_system_equations

---
 .../symbolics_tearing.jl                      | 174 ++++++++++--------
 src/structural_transformation/utils.jl        |  14 +-
 src/systems/systemstructure.jl                |  16 +-
 3 files changed, 119 insertions(+), 85 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index b99e6ec21a..74205fcf86 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -248,7 +248,7 @@ called dummy derivatives.
 State selection is done. All non-differentiated variables are algebraic 
 variables, and all variables that appear differentiated are differential variables.
 """
-function substitute_derivatives_algevars!(ts::TearingState, neweqs, dummy_sub, var_eq_matching)
+function substitute_derivatives_algevars!(ts::TearingState, neweqs, var_eq_matching, dummy_sub)
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     diff_to_var = invview(var_to_diff)
@@ -353,30 +353,32 @@ Effects on the system structure:
 - solvable_graph:
 - var_eq_matching: match D(x) to the added identity equation
 """
-function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching;
-        is_discrete = false, mm = nothing)
+function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching; mm = nothing)
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
     iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
-    lower_name = is_discrete ? lower_varname_withshift : lower_varname_with_unit
-
+    is_discrete = is_only_discrete(structure)
+    lower_varname = is_discrete ? lower_shift_varname : lower_varname_with_unit
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
-    # Generate new derivative variables for all unsolved variables that have a derivative in the system 
+    # For variable x, make dummy derivative x_t if the
+    # derivative is in the system
     for v in 1:length(var_to_diff)
-        # Check if a derivative 1) exists and 2) is unsolved for
         dv = var_to_diff[v]
+        # For discrete systems, directly substitute lowest-order variable 
+        if is_discrete && diff_to_var[v] == nothing
+            fullvars[v] = lower_varname(fullvars[v], iv)
+        end
         dv isa Int || continue
         solved = var_eq_matching[dv] isa Int
         solved && continue
 
         # If there's `D(x) = x_t` already, update mappings and continue without
         # adding new equations/variables
-        dd = find_duplicate_dd(dv, lineareqs, mm)
-
+        dd = find_duplicate_dd(dv, solvable_graph, linear_eqs, mm)
         if !isnothing(dd)
             dummy_eq, v_t = dd
             var_to_diff[v_t] = var_to_diff[dv]
@@ -386,26 +388,25 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         end
 
         dx = fullvars[dv]
-        order, lv = var_order(diff_to_var, dv)
-        x_t = is_discrete ? lower_name(fullvars[lv], iv)
-                      : lower_name(fullvars[lv], iv, order)
-        
+        order, lv = var_order(dv, diff_to_var)
+        x_t = is_discrete ? lower_varname(fullvars[dv], iv) : lower_varname(fullvars[lv], iv, order)
+
         # Add `x_t` to the graph
-        add_dd_variable!(structure, x_t, dv)
+        v_t = add_dd_variable!(structure, fullvars, x_t, dv)
         # Add `D(x) - x_t ~ 0` to the graph
-        add_dd_equation!(structure, neweqs, 0 ~ dx - x_t, dv)
+        dummy_eq = add_dd_equation!(structure, neweqs, 0 ~ dx - x_t, dv, v_t)
 
         # Update matching
         push!(var_eq_matching, unassigned)
         var_eq_matching[dv] = unassigned
-        eq_var_matching[dummy_eq] = dv
+        eq_var_matching[dummy_eq] = dv 
     end
 end
 
 """
-Check if there's `D(x) = x_t` already. 
+Check if there's `D(x) = x_t` already.
 """
-function find_duplicate_dd(dv, lineareqs, mm)
+function find_duplicate_dd(dv, solvable_graph, linear_eqs, mm)
     for eq in 𝑑neighbors(solvable_graph, dv)
         mi = get(linear_eqs, eq, 0)
         iszero(mi) && continue
@@ -424,28 +425,28 @@ function find_duplicate_dd(dv, lineareqs, mm)
     return nothing 
 end
 
-function add_dd_variable!(s::SystemStructure, x_t, dv)
-    push!(s.fullvars, simplify_shifts(x_t))
+function add_dd_variable!(s::SystemStructure, fullvars, x_t, dv)
+    push!(fullvars, simplify_shifts(x_t))
+    v_t = length(fullvars)
     v_t_idx = add_vertex!(s.var_to_diff)
-    @assert v_t_idx == ndsts(graph) == ndsts(solvable_graph) == length(fullvars) ==
-        length(var_eq_matching)
     add_vertex!(s.graph, DST)
     # TODO: do we care about solvable_graph? We don't use them after
     # `dummy_derivative_graph`.
     add_vertex!(s.solvable_graph, DST)
-    var_to_diff[v_t] = var_to_diff[dv]
+    s.var_to_diff[v_t] = s.var_to_diff[dv]
+    v_t
 end
 
 # dv = index of D(x), v_t = index of x_t
-function add_dd_equation!(s::SystemStructure, neweqs, eq, dv)
+function add_dd_equation!(s::SystemStructure, neweqs, eq, dv, v_t)
     push!(neweqs, eq)
     add_vertex!(s.graph, SRC)
-    v_t = length(s.fullvars)
     dummy_eq = length(neweqs)
     add_edge!(s.graph, dummy_eq, dv)
     add_edge!(s.graph, dummy_eq, v_t)
     add_vertex!(s.solvable_graph, SRC)
     add_edge!(s.solvable_graph, dummy_eq, dv)
+    dummy_eq
 end
 
 """
@@ -463,6 +464,10 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
         iv = get_iv(sys)
         if is_only_discrete(structure)
             D = Shift(iv, 1)
+            for v in fullvars
+                op = operation(v)
+                op isa Shift && (op.steps < 0) && (total_sub[v] = lower_shift_varname(v, iv))
+            end
         else
             D = Differential(iv)
         end
@@ -493,24 +498,40 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
     eqs = Iterators.reverse(toporder)
     idep = iv
 
-    # Generate differential equations.
-    # fullvars[iv] is a differential variable of the form D^n(x), and neweqs[ieq]
-    # is solved to give the RHS.
+    # Generate equations.
+    #   Solvable equations of differential variables D(x) become differential equations
+    #   Solvable equations of non-differential variables become observable equations
+    #   Non-solvable equations become algebraic equations.
     for ieq in eqs
         iv = eq_var_matching[ieq]
-        if is_solvable(ieq, iv)
-            if isdervar(iv)
-                isnothing(D) &&
-                    error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(equations(sys)[ieq])")
-                add_differential_equation!(structure, iv, neweqs, ieq, 
-                                           diff_vars, diff_eqs, diffeq_idxs, total_sub)
-            else
-                add_solved_equation!(structure, iv, neweqs, ieq, 
-                                     solved_vars, solved_eqs, solvedeq_idxs, total_sub)
+        var = fullvars[iv]
+        eq = neweqs[ieq]
+
+        if is_solvable(ieq, iv) && isdervar(iv)
+            isnothing(D) && throw(UnexpectedDifferentialError(equations(sys)[ieq]))
+            order, lv = var_order(iv, diff_to_var)
+            dx = D(simplify_shifts(fullvars[lv]))
+
+            neweq = make_differential_equation(var, dx, eq, total_sub)
+            for e in 𝑑neighbors(graph, iv)
+                e == ieq && continue
+                rem_edge!(graph, e, iv)
+            end
+
+            push!(diff_eqs, neweq)
+            push!(diffeq_idxs, ieq)
+            push!(diff_vars, diff_to_var[iv])
+        elseif is_solvable(ieq, iv)
+            neweq = make_solved_equation(var, eq, total_sub; simplify)
+            !isnothing(neweq) && begin
+                push!(solved_eqs, neweq)
+                push!(solvedeq_idxs, ieq)
+                push!(solved_vars, iv)
             end
         else
-            add_algebraic_equation!(structure, neweqs, ieq, 
-                                    alge_eqs, algeeq_idxs, total_sub)
+            neweq = make_algebraic_equation(var, eq, total_sub)
+            push!(alge_eqs, neweq)
+            push!(algeeq_idxs, ieq)
         end
     end
 
@@ -529,39 +550,29 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
     return neweqs, solved_eqs, eq_ordering, var_ordering, length(solved_vars), length(solved_vars_set)
 end
 
-function add_differential_equation!(s::SystemStructure, iv, neweqs, ieq, diff_vars, diff_eqs, diffeqs_idxs, total_sub)
-    diff_to_var = invview(s.var_to_diff)
+struct UnexpectedDifferentialError
+    eq::Equation
+end
 
-    order, lv = var_order(diff_to_var, iv)
-    dx = D(simplify_shifts(fullvars[lv]))
-    eq = dx ~ simplify_shifts(Symbolics.fixpoint_sub(
-        Symbolics.symbolic_linear_solve(neweqs[ieq],
-            fullvars[iv]),
-        total_sub; operator = ModelingToolkit.Shift))
-    for e in 𝑑neighbors(s.graph, iv)
-        e == ieq && continue
-        rem_edge!(s.graph, e, iv)
-    end
+function Base.showerror(io::IO, err::UnexpectedDifferentialError)
+    error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(err.eq)")
+end
 
-    push!(diff_eqs, eq)
-    total_sub[simplify_shifts(eq.lhs)] = eq.rhs
-    push!(diffeq_idxs, ieq)
-    push!(diff_vars, diff_to_var[iv])
+function make_differential_equation(var, dx, eq, total_sub)
+    dx ~ simplify_shifts(Symbolics.fixpoint_sub(
+        Symbolics.symbolic_linear_solve(eq, var),
+        total_sub; operator = ModelingToolkit.Shift))
 end
 
-function add_algebraic_equation!(s::SystemStructure, neweqs, ieq, alge_eqs, algeeq_idxs, total_sub)
-    eq = neweqs[ieq]
+function make_algebraic_equation(var, eq, total_sub)
     rhs = eq.rhs
     if !(eq.lhs isa Number && eq.lhs == 0)
         rhs = eq.rhs - eq.lhs
     end
-    push!(alge_eqs, 0 ~ simplify_shifts(Symbolics.fixpoint_sub(rhs, total_sub)))
-    push!(algeeq_idxs, ieq)
+    0 ~ simplify_shifts(Symbolics.fixpoint_sub(rhs, total_sub))
 end
 
-function add_solved_equation!(s::SystemStructure, iv, neweqs, ieq, solved_vars, solved_eqs, solvedeq_idxs, total_sub)
-    eq = neweqs[ieq]
-    var = fullvars[iv]
+function make_solved_equation(var, eq, total_sub; simplify = false)
     residual = eq.lhs - eq.rhs
     a, b, islinear = linear_expansion(residual, var)
     @assert islinear
@@ -569,15 +580,13 @@ function add_solved_equation!(s::SystemStructure, iv, neweqs, ieq, solved_vars,
     # var ~ -b/a
     if ModelingToolkit._iszero(a)
         @warn "Tearing: solving $eq for $var is singular!"
+        return nothing
     else
         rhs = -b / a
-        neweq = var ~ simplify_shifts(Symbolics.fixpoint_sub(
+        return var ~ simplify_shifts(Symbolics.fixpoint_sub(
             simplify ?
             Symbolics.simplify(rhs) : rhs,
             total_sub; operator = ModelingToolkit.Shift))
-        push!(solved_eqs, neweq)
-        push!(solvedeq_idxs, ieq)
-        push!(solved_vars, iv)
     end
 end
 
@@ -592,8 +601,8 @@ such that the mass matrix is:
 Update the state to account for the new ordering and equations.
 """
 # TODO: BLT sorting
-function reorder_vars!(s::SystemStructure, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
-    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
+function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
+    @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
 
     eqsperm = zeros(Int, nsrcs(graph))
     for (i, v) in enumerate(eq_ordering)
@@ -623,20 +632,26 @@ function reorder_vars!(s::SystemStructure, var_eq_matching, eq_ordering, var_ord
         d′ = eqsperm[d]
         new_eq_to_diff[v′] = d′ > 0 ? d′ : nothing
     end
-    new_fullvars = fullvars[invvarsperm]
+    new_fullvars = state.fullvars[var_ordering]
 
+    @show new_graph
+    @show new_var_to_diff
     # Update system structure 
     @set! state.structure.graph = complete(new_graph)
     @set! state.structure.var_to_diff = new_var_to_diff
     @set! state.structure.eq_to_diff = new_eq_to_diff
     @set! state.fullvars = new_fullvars
+    state
 end
 
 """
 Set the system equations, unknowns, observables post-tearing.
 """
-function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns)
+function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns; 
+        cse_hack = true, array_hack = true)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
+    @show graph
+    @show var_to_diff
     diff_to_var = invview(var_to_diff)
 
     ispresent = let var_to_diff = var_to_diff, graph = graph
@@ -656,7 +671,12 @@ function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dumm
     unknowns = Any[v
                    for (i, v) in enumerate(state.fullvars)
                    if diff_to_var[i] === nothing && ispresent(i)]
+    @show unknowns
+    @show state.fullvars
+    @show 𝑑neighbors(graph, 5)
+    @show neweqs
     unknowns = [unknowns; extra_unknowns]
+    @show unknowns
     @set! sys.unknowns = unknowns
 
     obs, subeqs, deps = cse_and_array_hacks(
@@ -684,7 +704,7 @@ end
 # differential variables.
     
 # Give the order of the variable indexed by dv
-function var_order(diff_to_var, dv)
+function var_order(dv, diff_to_var)
     order = 0
     while (dv′ = diff_to_var[dv]) !== nothing
         order += 1
@@ -695,7 +715,6 @@ end
 
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
-
     extra_vars = Int[]
     if full_var_eq_matching !== nothing
         for v in 𝑑vertices(state.structure.graph)
@@ -704,21 +723,22 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
             push!(extra_vars, v)
         end
     end
-    extra_unknowns = fullvars[extra_vars]
+    extra_unknowns = state.fullvars[extra_vars]
     neweqs = collect(equations(state))
+    dummy_sub = Dict()
 
     # Structural simplification 
-    dummy_sub = Dict()
-    substitute_derivatives_algevars!(state, neweqs, dummy_sub, var_eq_matching)
+    substitute_derivatives_algevars!(state, neweqs, var_eq_matching, dummy_sub)
 
     generate_derivative_variables!(state, neweqs, var_eq_matching; mm)
 
     neweqs, solved_eqs, eq_ordering, var_ordering, nelim_eq, nelim_var = 
         generate_system_equations!(state, neweqs, var_eq_matching; simplify)
 
-    reorder_vars!(state.structure, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
+    state = reorder_vars!(state, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
+
+    sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns; cse_hack, array_hack)
 
-    sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns)
     @set! state.sys = sys
     @set! sys.tearing_state = state
     return invalidate_cache!(sys)
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 77f0b62a1c..fec57db080 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -452,9 +452,10 @@ end
 ### Misc
 ###
 
-function lower_varname_withshift(var, iv)
+# 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 var
+    op isa Shift || return Shift(iv, 0)(var, true) # hack to prevent simplification of x(t) - x(t)
     backshift = op.steps
     backshift > 0 && return var
 
@@ -476,6 +477,14 @@ function lower_varname_withshift(var, iv)
     return ModelingToolkit._with_unit(identity, newvar, iv)
 end
 
+function lower_varname(var, iv, order; is_discrete = false)
+    if is_discrete
+        lower_shift_varname(var, iv)
+    else
+        lower_varname_with_unit(var, iv, order)
+    end
+end
+
 function isdoubleshift(var)
     return ModelingToolkit.isoperator(var, ModelingToolkit.Shift) &&
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
@@ -484,6 +493,7 @@ end
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
     var isa Equation && return simplify_shifts(var.lhs) ~ simplify_shifts(var.rhs)
+    (op = operation(var)) isa Shift && op.steps == 0 && return first(arguments(var))
     if isdoubleshift(var)
         op1 = operation(var)
         vv1 = arguments(var)[1]
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index e5a227d9fb..ae4d21f225 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -142,15 +142,16 @@ has_equations(::TransformationState) = true
 Base.@kwdef mutable struct SystemStructure
     """Maps the (index of) a variable to the (index of) the variable describing its derivative."""
     var_to_diff::DiffGraph
-    """Maps the (index of) a """
+    """Maps the (index of) an equation."""
     eq_to_diff::DiffGraph
     # Can be access as
     # `graph` to automatically look at the bipartite graph
     # or as `torn` to assert that tearing has run.
-    """Incidence graph of the system of equations. An edge from equation x to variable y exists if variable y appears in equation x."""
+    """Graph that maps equations to variables that appear in them."""
     graph::BipartiteGraph{Int, Nothing}
-    """."""
+    """Graph that connects equations to the variable they will be solved for during simplification."""
     solvable_graph::Union{BipartiteGraph{Int, Nothing}, Nothing}
+    """Variable types (brownian, variable, parameter) in the system."""
     var_types::Union{Vector{VariableType}, Nothing}
     """Whether the system is discrete."""
     only_discrete::Bool
@@ -200,7 +201,9 @@ function complete!(s::SystemStructure)
 end
 
 mutable struct TearingState{T <: AbstractSystem} <: AbstractTearingState{T}
+    """The system of equations."""
     sys::T
+    """The set of variables of the system."""
     fullvars::Vector
     structure::SystemStructure
     extra_eqs::Vector
@@ -438,7 +441,7 @@ function TearingState(sys; quick_cancel = false, check = true)
             complete(graph), nothing, var_types, sys isa DiscreteSystem),
         Any[])
     if sys isa DiscreteSystem
-        ts = shift_discrete_system(ts, lowest_shift)
+        ts = shift_discrete_system(ts)
     end
     return ts
 end
@@ -475,8 +478,9 @@ function shift_discrete_system(ts::TearingState)
     end
     iv = get_iv(sys)
 
-    discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, 1)(k))    for k in discvars 
-    if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold})) 
+    discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, 1)(k))    
+                   for k in discvars 
+                   if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold})) 
 
     for i in eachindex(fullvars)
         fullvars[i] = StructuralTransformations.simplify_shifts(fast_substitute(

From 22a39bd41b2123c48a7f88016f25470041271b69 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 10 Feb 2025 18:42:01 -0500
Subject: [PATCH 25/41] delete comments

---
 src/structural_transformation/utils.jl  |  3 ---
 src/systems/systems.jl                  |  8 ++++----
 test/structural_transformation/utils.jl | 16 ----------------
 3 files changed, 4 insertions(+), 23 deletions(-)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index fec57db080..757eb9a8db 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -271,8 +271,6 @@ end
 
 function find_solvables!(state::TearingState; kwargs...)
     @assert state.structure.solvable_graph === nothing
-    println("in find_solvables")
-    @show eqs
     eqs = equations(state)
     graph = state.structure.graph
     state.structure.solvable_graph = BipartiteGraph(nsrcs(graph), ndsts(graph))
@@ -280,7 +278,6 @@ function find_solvables!(state::TearingState; kwargs...)
     for ieq in 1:length(eqs)
         find_eq_solvables!(state, ieq, to_rm; kwargs...)
     end
-    @show eqs
     return nothing
 end
 
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 0a7e4264e4..9c8c272c5c 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -41,10 +41,10 @@ function structural_simplify(
     end
     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.
-        # """)
+        error("""
+            Encountered algebraic equations when simplifying discrete system. This is \
+            not yet supported.
+        """)
     end
     for pass in additional_passes
         newsys = pass(newsys)
diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl
index 98d986ebd4..67f6016bc4 100644
--- a/test/structural_transformation/utils.jl
+++ b/test/structural_transformation/utils.jl
@@ -162,19 +162,3 @@ end
     structural_simplify(sys; additional_passes = [pass])
     @test value[] == 1
 end
-
-@testset "Shift simplification" begin
-    @variables x(t) y(t) z(t)
-    @parameters a b c
-    
-    # Expand shifts
-    @test isequal(ST.expand_shifts(Shift(t, -3)(x)), Shift(t, -1)(Shift(t, -1)(Shift(t, -1)(x))))
-    expr = a * Shift(t, -2)(x) + Shift(t, 2)(y) + b
-    @test isequal(ST.expand_shifts(expr), 
-                    a * Shift(t, -1)(Shift(t, -1)(x)) + Shift(t, 1)(Shift(t, 1)(y)) + b)
-    @test isequal(ST.expand_shifts(Shift(t, 2)(Shift(t, 1)(a))), a)
-
-
-    # Distribute shifts
-
-end

From 9aaf04a512853d5c9336674455ade5d92487c2f5 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 10 Feb 2025 18:50:23 -0500
Subject: [PATCH 26/41] more cleanup

---
 docs/src/systems/DiscreteSystem.md            | 28 -------------------
 .../symbolics_tearing.jl                      |  9 ------
 src/systems/systemstructure.jl                | 11 ++------
 3 files changed, 3 insertions(+), 45 deletions(-)
 delete mode 100644 docs/src/systems/DiscreteSystem.md

diff --git a/docs/src/systems/DiscreteSystem.md b/docs/src/systems/DiscreteSystem.md
deleted file mode 100644
index b6a8061e50..0000000000
--- a/docs/src/systems/DiscreteSystem.md
+++ /dev/null
@@ -1,28 +0,0 @@
-# DiscreteSystem
-
-## System Constructors
-
-```@docs
-DiscreteSystem
-```
-
-## Composition and Accessor Functions
-
-  - `get_eqs(sys)` or `equations(sys)`: The equations that define the discrete system.
-  - `get_unknowns(sys)` or `unknowns(sys)`: The set of unknowns in the discrete system.
-  - `get_ps(sys)` or `parameters(sys)`: The parameters of the discrete system.
-  - `get_iv(sys)`: The independent variable of the discrete system.
-  - `discrete_events(sys)`: The set of discrete events in the discrete system.
-
-## Transformations
-
-```@docs; canonical=false
-structural_simplify
-```
-
-## Problem Constructors
-
-```@docs; canonical=false
-DiscreteProblem(sys::DiscreteSystem, u0map, tspan)
-DiscreteFunction(sys::DiscreteSystem, u0map, tspan)
-```
diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 74205fcf86..eed7bc3e29 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -634,8 +634,6 @@ function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_or
     end
     new_fullvars = state.fullvars[var_ordering]
 
-    @show new_graph
-    @show new_var_to_diff
     # Update system structure 
     @set! state.structure.graph = complete(new_graph)
     @set! state.structure.var_to_diff = new_var_to_diff
@@ -650,8 +648,6 @@ Set the system equations, unknowns, observables post-tearing.
 function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns; 
         cse_hack = true, array_hack = true)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
-    @show graph
-    @show var_to_diff
     diff_to_var = invview(var_to_diff)
 
     ispresent = let var_to_diff = var_to_diff, graph = graph
@@ -671,12 +667,7 @@ function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dumm
     unknowns = Any[v
                    for (i, v) in enumerate(state.fullvars)
                    if diff_to_var[i] === nothing && ispresent(i)]
-    @show unknowns
-    @show state.fullvars
-    @show 𝑑neighbors(graph, 5)
-    @show neweqs
     unknowns = [unknowns; extra_unknowns]
-    @show unknowns
     @set! sys.unknowns = unknowns
 
     obs, subeqs, deps = cse_and_array_hacks(
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index ae4d21f225..b98bb8c616 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -140,14 +140,14 @@ get_fullvars(ts::TransformationState) = ts.fullvars
 has_equations(::TransformationState) = true
 
 Base.@kwdef mutable struct SystemStructure
-    """Maps the (index of) a variable to the (index of) the variable describing its derivative."""
+    """Maps the index of variable x to the index of variable D(x)."""
     var_to_diff::DiffGraph
-    """Maps the (index of) an equation."""
+    """Maps the index of an equation."""
     eq_to_diff::DiffGraph
     # Can be access as
     # `graph` to automatically look at the bipartite graph
     # or as `torn` to assert that tearing has run.
-    """Graph that maps equations to variables that appear in them."""
+    """Graph that connects equations to variables that appear in them."""
     graph::BipartiteGraph{Int, Nothing}
     """Graph that connects equations to the variable they will be solved for during simplification."""
     solvable_graph::Union{BipartiteGraph{Int, Nothing}, Nothing}
@@ -464,15 +464,10 @@ function lower_order_var(dervar, t)
     diffvar
 end
 
-"""
-    Shift variable x by the largest shift s such that x(k-s) appears in the system of equations.
-    The lowest-shift term will have.
-"""
 function shift_discrete_system(ts::TearingState)
     @unpack fullvars, sys = ts
     discvars = OrderedSet()
     eqs = equations(sys)
-
     for eq in eqs
         vars!(discvars, eq; op = Union{Sample, Hold})
     end

From 66266e410c88a61a9c4e6af1c8c829433d5b4f76 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 10 Feb 2025 19:10:34 -0500
Subject: [PATCH 27/41] better comments

---
 src/structural_transformation/symbolics_tearing.jl | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index eed7bc3e29..5707ac8386 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -337,11 +337,19 @@ variables and equations, don't add them when they already exist.
 
 Documenting the differences to structural simplification for discrete systems:
 
-In discrete systems the lowest-order term is Shift(t, k)(x(t)), instead of x(t). 
-We want to substitute the k-1 lowest order terms instead of the k-1 highest order terms.
+In discrete systems everything gets shifted forward a timestep by `shift_discrete_system`
+in order to properly generate the difference equations. 
+
+In the system x(k) ~ x(k-1) + x(k-2), becomes Shift(t, 1)(x(t)) ~ x(t) + Shift(t, -1)(x(t))
+
+The lowest-order term is Shift(t, k)(x(t)), instead of x(t).
+As such we actually want dummy variables for the k-1 lowest order terms instead of the k-1 highest order terms.
 
-In the system x(k) ~ x(k-1) + x(k-2), we want to lower
 Shift(t, -1)(x(t)) -> x\_{t-1}(t)
+
+Since Shift(t, -1)(x) is not a derivative, it is directly substituted in `fullvars`. No equation or variable is added for it. 
+
+For ODESystems D(D(D(x))) in equations is recursively substituted as D(x) ~ x_t, D(x_t) ~ x_tt, etc. The analogue for discrete systems, Shift(t, 1)(Shift(t,1)(Shift(t,1)(Shift(t, -3)(x(t))))) does not actually appear. So `total_sub` in generate_system_equations` is directly initialized with all of the lowered variables `Shift(t, -3)(x) -> x_t-3(t)`, etc. 
 =#
 """
 Generate new derivative variables for the system.

From 1ab62463adb298386145f57933f4ccd4667d90d1 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 10 Feb 2025 19:30:56 -0500
Subject: [PATCH 28/41] fix unassigned indexing

---
 src/structural_transformation/symbolics_tearing.jl | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 5707ac8386..48f9141dbd 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -512,10 +512,10 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
     #   Non-solvable equations become algebraic equations.
     for ieq in eqs
         iv = eq_var_matching[ieq]
-        var = fullvars[iv]
         eq = neweqs[ieq]
 
         if is_solvable(ieq, iv) && isdervar(iv)
+            var = fullvars[iv]
             isnothing(D) && throw(UnexpectedDifferentialError(equations(sys)[ieq]))
             order, lv = var_order(iv, diff_to_var)
             dx = D(simplify_shifts(fullvars[lv]))
@@ -530,6 +530,7 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
             push!(diffeq_idxs, ieq)
             push!(diff_vars, diff_to_var[iv])
         elseif is_solvable(ieq, iv)
+            var = fullvars[iv]
             neweq = make_solved_equation(var, eq, total_sub; simplify)
             !isnothing(neweq) && begin
                 push!(solved_eqs, neweq)
@@ -537,7 +538,7 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
                 push!(solved_vars, iv)
             end
         else
-            neweq = make_algebraic_equation(var, eq, total_sub)
+            neweq = make_algebraic_equation(eq, total_sub)
             push!(alge_eqs, neweq)
             push!(algeeq_idxs, ieq)
         end
@@ -572,7 +573,7 @@ function make_differential_equation(var, dx, eq, total_sub)
         total_sub; operator = ModelingToolkit.Shift))
 end
 
-function make_algebraic_equation(var, eq, total_sub)
+function make_algebraic_equation(eq, total_sub)
     rhs = eq.rhs
     if !(eq.lhs isa Number && eq.lhs == 0)
         rhs = eq.rhs - eq.lhs

From 079901bd481b8ea4f505761bce6aa476602dd023 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 11 Feb 2025 15:13:17 -0500
Subject: [PATCH 29/41] move D, iv into tearing_reassemble

---
 .../symbolics_tearing.jl                      | 43 +++++++++----------
 1 file changed, 21 insertions(+), 22 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 48f9141dbd..45430595b6 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -248,11 +248,10 @@ called dummy derivatives.
 State selection is done. All non-differentiated variables are algebraic 
 variables, and all variables that appear differentiated are differential variables.
 """
-function substitute_derivatives_algevars!(ts::TearingState, neweqs, var_eq_matching, dummy_sub)
+function substitute_derivatives_algevars!(ts::TearingState, neweqs, var_eq_matching, dummy_sub; iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     diff_to_var = invview(var_to_diff)
-    iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
 
     for var in 1:length(fullvars)
         dv = var_to_diff[var]
@@ -361,12 +360,11 @@ Effects on the system structure:
 - solvable_graph:
 - var_eq_matching: match D(x) to the added identity equation
 """
-function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching; mm = nothing)
+function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching; mm = nothing, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
-    iv = ModelingToolkit.has_iv(sys) ? ModelingToolkit.get_iv(sys) : nothing
     is_discrete = is_only_discrete(structure)
     lower_varname = is_discrete ? lower_shift_varname : lower_varname_with_unit
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
@@ -461,27 +459,18 @@ end
 Solve the solvable equations of the system and generate differential (or discrete)
 equations in terms of the selected states.
 """
-function generate_system_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false)
+function generate_system_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = state 
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
     total_sub = Dict()
-
-    if ModelingToolkit.has_iv(sys)
-        iv = get_iv(sys)
-        if is_only_discrete(structure)
-            D = Shift(iv, 1)
-            for v in fullvars
-                op = operation(v)
-                op isa Shift && (op.steps < 0) && (total_sub[v] = lower_shift_varname(v, iv))
-            end
-        else
-            D = Differential(iv)
+    if is_only_discrete(structure)
+        for v in fullvars
+            op = operation(v)
+            op isa Shift && (op.steps < 0) && (total_sub[v] = lower_shift_varname(v, iv))
         end
-    else
-        iv = D = nothing
-    end
+   end
 
     # if var is like D(x) or Shift(t, 1)(x)
     isdervar = let diff_to_var = diff_to_var
@@ -727,13 +716,23 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     neweqs = collect(equations(state))
     dummy_sub = Dict()
 
+    if ModelingToolkit.has_iv(state.sys)
+        iv = get_iv(state.sys)
+        if !is_only_discrete(state.structure)
+            D = Differential(iv)
+        else
+            D = Shift(iv, 1)
+        end
+        iv = D = nothing
+    end
+
     # Structural simplification 
-    substitute_derivatives_algevars!(state, neweqs, var_eq_matching, dummy_sub)
+    substitute_derivatives_algevars!(state, neweqs, var_eq_matching, dummy_sub; iv, D)
 
-    generate_derivative_variables!(state, neweqs, var_eq_matching; mm)
+    generate_derivative_variables!(state, neweqs, var_eq_matching; mm, iv, D)
 
     neweqs, solved_eqs, eq_ordering, var_ordering, nelim_eq, nelim_var = 
-        generate_system_equations!(state, neweqs, var_eq_matching; simplify)
+        generate_system_equations!(state, neweqs, var_eq_matching; simplify, iv, D)
 
     state = reorder_vars!(state, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
 

From 9217c62941226cd8dc24a80fe3723810f2463ab5 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 11 Feb 2025 15:29:46 -0500
Subject: [PATCH 30/41] refactor iv

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

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 45430595b6..883166cfa8 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -723,6 +723,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         else
             D = Shift(iv, 1)
         end
+    else
         iv = D = nothing
     end
 

From 54910988509cfdf2f65ed8b8d24f565b3dab8584 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 11 Feb 2025 16:13:47 -0500
Subject: [PATCH 31/41] fix comment for eq_to_diff

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

diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index b98bb8c616..6fee78cfd6 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -142,7 +142,7 @@ has_equations(::TransformationState) = true
 Base.@kwdef mutable struct SystemStructure
     """Maps the index of variable x to the index of variable D(x)."""
     var_to_diff::DiffGraph
-    """Maps the index of an equation."""
+    """Maps the index of an algebraic equation to the index of the equation it is differentiated into."""
     eq_to_diff::DiffGraph
     # Can be access as
     # `graph` to automatically look at the bipartite graph

From 5878ad67ef6a78ad55034d8138a6fee741a9e7d6 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 11 Feb 2025 16:51:06 -0500
Subject: [PATCH 32/41] fix total_sub

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

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 883166cfa8..7cb9798f01 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -464,6 +464,7 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
+
     total_sub = Dict()
     if is_only_discrete(structure)
         for v in fullvars
@@ -515,6 +516,7 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
                 rem_edge!(graph, e, iv)
             end
 
+            total_sub[simplify_shifts(neweq.lhs)] = neweq.rhs
             push!(diff_eqs, neweq)
             push!(diffeq_idxs, ieq)
             push!(diff_vars, diff_to_var[iv])

From 397b279df53c45a6c923ee89a24125237d4ee44c Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 11 Feb 2025 23:08:07 -0500
Subject: [PATCH 33/41] add diff_to_var as argument for find_dumplicate_dd, fix
 substitution of lowest-order variable

---
 src/structural_transformation/symbolics_tearing.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 7cb9798f01..ec1ff45fd7 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -374,9 +374,9 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
     # derivative is in the system
     for v in 1:length(var_to_diff)
         dv = var_to_diff[v]
-        # For discrete systems, directly substitute lowest-order variable 
+        # For discrete systems, directly substitute lowest-order shift 
         if is_discrete && diff_to_var[v] == nothing
-            fullvars[v] = lower_varname(fullvars[v], iv)
+            operation(fullvars[v]) isa Shift && (fullvars[v] = lower_varname(fullvars[v], iv))
         end
         dv isa Int || continue
         solved = var_eq_matching[dv] isa Int
@@ -384,7 +384,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
 
         # If there's `D(x) = x_t` already, update mappings and continue without
         # adding new equations/variables
-        dd = find_duplicate_dd(dv, solvable_graph, linear_eqs, mm)
+        dd = find_duplicate_dd(dv, solvable_graph, diff_to_var, linear_eqs, mm)
         if !isnothing(dd)
             dummy_eq, v_t = dd
             var_to_diff[v_t] = var_to_diff[dv]
@@ -412,7 +412,7 @@ end
 """
 Check if there's `D(x) = x_t` already.
 """
-function find_duplicate_dd(dv, solvable_graph, linear_eqs, mm)
+function find_duplicate_dd(dv, solvable_graph, diff_to_var, linear_eqs, mm)
     for eq in 𝑑neighbors(solvable_graph, dv)
         mi = get(linear_eqs, eq, 0)
         iszero(mi) && continue

From bb42719158e876d21a9ddf12b934141632fd2823 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 13 Feb 2025 15:46:43 -0500
Subject: [PATCH 34/41] fix: fix initialization of DiscreteSystem with renamed
 variables

---
 .../StructuralTransformations.jl              |  3 +-
 .../symbolics_tearing.jl                      | 20 +++++++-----
 src/structural_transformation/utils.jl        | 32 ++++++++++---------
 .../discrete_system/discrete_system.jl        |  6 ++--
 src/utils.jl                                  |  2 ++
 src/variables.jl                              |  7 +++-
 test/runtests.jl                              |  3 ++
 7 files changed, 45 insertions(+), 28 deletions(-)

diff --git a/src/structural_transformation/StructuralTransformations.jl b/src/structural_transformation/StructuralTransformations.jl
index 1220d517cc..510352ad59 100644
--- a/src/structural_transformation/StructuralTransformations.jl
+++ b/src/structural_transformation/StructuralTransformations.jl
@@ -22,7 +22,7 @@ using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Di
                        get_postprocess_fbody, vars!,
                        IncrementalCycleTracker, add_edge_checked!, topological_sort,
                        invalidate_cache!, Substitutions, get_or_construct_tearing_state,
-                       filter_kwargs, lower_varname_with_unit, setio, SparseMatrixCLIL,
+                       filter_kwargs, lower_varname_with_unit, lower_shift_varname_with_unit, setio, SparseMatrixCLIL,
                        get_fullvars, has_equations, observed,
                        Schedule, schedule
 
@@ -63,6 +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
 
 include("utils.jl")
 include("pantelides.jl")
diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index ec1ff45fd7..ffbd6e3452 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -366,7 +366,6 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
     is_discrete = is_only_discrete(structure)
-    lower_varname = is_discrete ? lower_shift_varname : lower_varname_with_unit
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
@@ -375,9 +374,9 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
     for v in 1:length(var_to_diff)
         dv = var_to_diff[v]
         # For discrete systems, directly substitute lowest-order shift 
-        if is_discrete && diff_to_var[v] == nothing
-            operation(fullvars[v]) isa Shift && (fullvars[v] = lower_varname(fullvars[v], iv))
-        end
+        #if is_discrete && diff_to_var[v] == nothing
+        #    operation(fullvars[v]) isa Shift && (fullvars[v] = lower_shift_varname_with_unit(fullvars[v], iv))
+        #end
         dv isa Int || continue
         solved = var_eq_matching[dv] isa Int
         solved && continue
@@ -395,7 +394,8 @@ 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_varname(fullvars[dv], iv) : lower_varname(fullvars[lv], iv, order)
+        x_t = is_discrete ? lower_shift_varname_with_unit(fullvars[dv], iv) : 
+                            Symbolics.diff2term(fullvars[dv])
 
         # Add `x_t` to the graph
         v_t = add_dd_variable!(structure, fullvars, x_t, dv)
@@ -467,11 +467,15 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
 
     total_sub = Dict()
     if is_only_discrete(structure)
-        for v in fullvars
+        for (i, v) in enumerate(fullvars)
             op = operation(v)
-            op isa Shift && (op.steps < 0) && (total_sub[v] = lower_shift_varname(v, iv))
+            op isa Shift && (op.steps < 0) && begin
+                lowered = lower_shift_varname_with_unit(v, iv)
+                total_sub[v] = lowered
+                fullvars[i] = lowered
+            end
         end
-   end
+    end
 
     # if var is like D(x) or Shift(t, 1)(x)
     isdervar = let diff_to_var = diff_to_var
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 757eb9a8db..f64a8da132 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -453,16 +453,25 @@ end
 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)
-    backshift = op.steps
-    backshift > 0 && return var
+    if op.steps < 0
+        return shift2term(var)
+    else
+        return var
+    end
+end
 
-    ds = "$iv-$(-backshift)"
-    d_separator = 'ˍ'
+function shift2term(var) 
+    backshift = operation(var).steps
+    iv = operation(var).t
+    num = join(Char(0x2080 + d) for d in reverse!(digits(-backshift)))
+    ds = join([Char(0x209c), Char(0x208b), num])
+    #ds = "$iv-$(-backshift)"
+    #d_separator = 'ˍ'
 
     if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
         O = only(arguments(var))
         oldop = operation(O)
-        newname = Symbol(string(nameof(oldop)), d_separator, ds)
+        newname = Symbol(string(nameof(oldop)), ds)
     else
         O = var
         oldop = operation(var) 
@@ -470,16 +479,9 @@ function lower_shift_varname(var, iv)
         newname = Symbol(varname, d_separator, ds)
     end
     newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
-    setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
-    return ModelingToolkit._with_unit(identity, newvar, iv)
-end
-
-function lower_varname(var, iv, order; is_discrete = false)
-    if is_discrete
-        lower_shift_varname(var, iv)
-    else
-        lower_varname_with_unit(var, iv, order)
-    end
+    newvar = setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
+    newvar = setmetadata(newvar, ModelingToolkit.VariableUnshifted, O)
+    return newvar
 end
 
 function isdoubleshift(var)
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index cd6ff45b6c..9a9ac47853 100644
--- a/src/systems/discrete_system/discrete_system.jl
+++ b/src/systems/discrete_system/discrete_system.jl
@@ -275,10 +275,10 @@ function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
     end
     for var in unknowns(sys)
         op = operation(var)
-        op isa Shift || continue
         haskey(updated, var) && continue
-        root = first(arguments(var))
-        haskey(defs, root) || error("Initial condition for $var not provided.")
+        root = getunshifted(var)
+        isnothing(root) && continue
+        haskey(defs, root) || error("Initial condition for $root not provided.")
         updated[var] = defs[root]
     end
     return updated
diff --git a/src/utils.jl b/src/utils.jl
index cf49d9f445..c3a0f637a1 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1028,6 +1028,8 @@ end
 
 diff2term_with_unit(x, t) = _with_unit(diff2term, x, t)
 lower_varname_with_unit(var, iv, order) = _with_unit(lower_varname, var, iv, iv, order)
+shift2term_with_unit(x, t) = _with_unit(shift2term, x, t)
+lower_shift_varname_with_unit(var, iv) = _with_unit(lower_shift_varname, var, iv, iv)
 
 """
     $(TYPEDSIGNATURES)
diff --git a/src/variables.jl b/src/variables.jl
index 536119f107..a29a119607 100644
--- a/src/variables.jl
+++ b/src/variables.jl
@@ -6,6 +6,7 @@ struct VariableOutput end
 struct VariableIrreducible end
 struct VariableStatePriority end
 struct VariableMisc end
+struct VariableUnshifted end
 Symbolics.option_to_metadata_type(::Val{:unit}) = VariableUnit
 Symbolics.option_to_metadata_type(::Val{:connect}) = VariableConnectType
 Symbolics.option_to_metadata_type(::Val{:input}) = VariableInput
@@ -13,6 +14,7 @@ Symbolics.option_to_metadata_type(::Val{:output}) = VariableOutput
 Symbolics.option_to_metadata_type(::Val{:irreducible}) = VariableIrreducible
 Symbolics.option_to_metadata_type(::Val{:state_priority}) = VariableStatePriority
 Symbolics.option_to_metadata_type(::Val{:misc}) = VariableMisc
+Symbolics.option_to_metadata_type(::Val{:unshifted}) = VariableUnshifted
 
 """
     dump_variable_metadata(var)
@@ -133,7 +135,7 @@ function default_toterm(x)
     if iscall(x) && (op = operation(x)) isa Operator
         if !(op isa Differential)
             if op isa Shift && op.steps < 0
-                return x
+                return shift2term(x) 
             end
             x = normalize_to_differential(op)(arguments(x)...)
         end
@@ -600,3 +602,6 @@ getunit(x::Symbolic) = Symbolics.getmetadata(x, VariableUnit, nothing)
 Check if the variable `x` has a unit.
 """
 hasunit(x) = getunit(x) !== nothing
+
+getunshifted(x) = getunshifted(unwrap(x))
+getunshifted(x::Symbolic) = Symbolics.getmetadata(x, VariableUnshifted, nothing)
diff --git a/test/runtests.jl b/test/runtests.jl
index 9537b1b44e..e600305232 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -22,6 +22,8 @@ function activate_downstream_env()
     Pkg.instantiate()
 end
 
+@testset begin include("discrete_system.jl") end
+#=
 @time begin
     if GROUP == "All" || GROUP == "InterfaceI"
         @testset "InterfaceI" begin
@@ -136,3 +138,4 @@ end
         @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl")
     end
 end
+=#

From 09f31b23ca7338e3a646b26a3bc582bd58633a2a Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 13 Feb 2025 15:49:35 -0500
Subject: [PATCH 35/41] revert runtest

---
 test/runtests.jl | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index e600305232..9537b1b44e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -22,8 +22,6 @@ function activate_downstream_env()
     Pkg.instantiate()
 end
 
-@testset begin include("discrete_system.jl") end
-#=
 @time begin
     if GROUP == "All" || GROUP == "InterfaceI"
         @testset "InterfaceI" begin
@@ -138,4 +136,3 @@ end
         @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl")
     end
 end
-=#

From daa789856714d3926c265307a94bc6f3f7eecbbd Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 14 Feb 2025 06:25:44 -0500
Subject: [PATCH 36/41] frefactor use toterm instead of lower_varname

---
 src/structural_transformation/symbolics_tearing.jl | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index ffbd6e3452..47c7da371f 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -366,6 +366,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
     is_discrete = is_only_discrete(structure)
+    toterm = is_discrete ? shift2term_with_unit : diff2term_with_unit
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
@@ -373,10 +374,6 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
     # derivative is in the system
     for v in 1:length(var_to_diff)
         dv = var_to_diff[v]
-        # For discrete systems, directly substitute lowest-order shift 
-        #if is_discrete && diff_to_var[v] == nothing
-        #    operation(fullvars[v]) isa Shift && (fullvars[v] = lower_shift_varname_with_unit(fullvars[v], iv))
-        #end
         dv isa Int || continue
         solved = var_eq_matching[dv] isa Int
         solved && continue
@@ -394,8 +391,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])
+        x_t = toterm(fullvars[dv]) 
 
         # Add `x_t` to the graph
         v_t = add_dd_variable!(structure, fullvars, x_t, dv)
@@ -470,7 +466,7 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
         for (i, v) in enumerate(fullvars)
             op = operation(v)
             op isa Shift && (op.steps < 0) && begin
-                lowered = lower_shift_varname_with_unit(v, iv)
+                lowered = shift2term_with_unit(v)
                 total_sub[v] = lowered
                 fullvars[i] = lowered
             end

From ae19eeb84b3b2d3cf6ca6db1a658581b8f9ba714 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 14 Feb 2025 06:50:46 -0500
Subject: [PATCH 37/41] fix unit

---
 src/structural_transformation/StructuralTransformations.jl | 2 +-
 src/structural_transformation/symbolics_tearing.jl         | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/structural_transformation/StructuralTransformations.jl b/src/structural_transformation/StructuralTransformations.jl
index 510352ad59..7d2b9afa26 100644
--- a/src/structural_transformation/StructuralTransformations.jl
+++ b/src/structural_transformation/StructuralTransformations.jl
@@ -11,7 +11,7 @@ using SymbolicUtils: maketerm, iscall
 
 using ModelingToolkit
 using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Differential,
-                       unknowns, equations, vars, Symbolic, diff2term_with_unit, value,
+                       unknowns, equations, vars, Symbolic, diff2term_with_unit, shift2term_with_unit, value,
                        operation, arguments, Sym, Term, simplify, symbolic_linear_solve,
                        isdiffeq, isdifferential, isirreducible,
                        empty_substitutions, get_substitutions,
diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 47c7da371f..779f5931f5 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -366,7 +366,6 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
     is_discrete = is_only_discrete(structure)
-    toterm = is_discrete ? shift2term_with_unit : diff2term_with_unit
     linear_eqs = mm === nothing ? Dict{Int, Int}() :
                  Dict(reverse(en) for en in enumerate(mm.nzrows))
 
@@ -391,7 +390,8 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
 
         dx = fullvars[dv]
         order, lv = var_order(dv, diff_to_var)
-        x_t = toterm(fullvars[dv]) 
+        x_t = is_discrete ? lower_shift_varname_with_unit(fullvars[dv], iv)
+                          : lower_varname_with_unit(fullvars[lv], iv, order)
 
         # Add `x_t` to the graph
         v_t = add_dd_variable!(structure, fullvars, x_t, dv)
@@ -466,7 +466,7 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
         for (i, v) in enumerate(fullvars)
             op = operation(v)
             op isa Shift && (op.steps < 0) && begin
-                lowered = shift2term_with_unit(v)
+                lowered = lower_shift_varname_with_unit(v, iv)
                 total_sub[v] = lowered
                 fullvars[i] = lowered
             end

From 406d0a861d4904d10cf99ad11dbc118f38a0c718 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 14 Feb 2025 06:55:14 -0500
Subject: [PATCH 38/41] fix parse error

---
 src/structural_transformation/symbolics_tearing.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index 779f5931f5..a166bb064b 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -390,8 +390,8 @@ 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)
-                          : lower_varname_with_unit(fullvars[lv], iv, order)
+        x_t = is_discrete ? lower_shift_varname_with_unit(fullvars[dv], iv) :
+                          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 2a97325d0b4e9b70f3a18de8049d14559504880d Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 18 Feb 2025 15:19:38 -0500
Subject: [PATCH 39/41] feat: initialization of DiscreteSystem

---
 .../symbolics_tearing.jl                      | 114 +++++++++++++-----
 src/structural_transformation/utils.jl        |  55 ++++++---
 .../discrete_system/discrete_system.jl        |   8 +-
 src/variables.jl                              |  16 ++-
 test/discrete_system.jl                       |  68 ++++++++---
 5 files changed, 193 insertions(+), 68 deletions(-)

diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index a166bb064b..bbb6853a7c 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -341,14 +341,18 @@ in order to properly generate the difference equations.
 
 In the system x(k) ~ x(k-1) + x(k-2), becomes Shift(t, 1)(x(t)) ~ x(t) + Shift(t, -1)(x(t))
 
-The lowest-order term is Shift(t, k)(x(t)), instead of x(t).
-As such we actually want dummy variables for the k-1 lowest order terms instead of the k-1 highest order terms.
+The lowest-order term is Shift(t, k)(x(t)), instead of x(t). As such we actually want 
+dummy variables for the k-1 lowest order terms instead of the k-1 highest order terms.
 
 Shift(t, -1)(x(t)) -> x\_{t-1}(t)
 
-Since Shift(t, -1)(x) is not a derivative, it is directly substituted in `fullvars`. No equation or variable is added for it. 
+Since Shift(t, -1)(x) is not a derivative, it is directly substituted in `fullvars`. 
+No equation or variable is added for it. 
 
-For ODESystems D(D(D(x))) in equations is recursively substituted as D(x) ~ x_t, D(x_t) ~ x_tt, etc. The analogue for discrete systems, Shift(t, 1)(Shift(t,1)(Shift(t,1)(Shift(t, -3)(x(t))))) does not actually appear. So `total_sub` in generate_system_equations` is directly initialized with all of the lowered variables `Shift(t, -3)(x) -> x_t-3(t)`, etc. 
+For ODESystems D(D(D(x))) in equations is recursively substituted as D(x) ~ x_t, D(x_t) ~ x_tt, etc. 
+The analogue for discrete systems, Shift(t, 1)(Shift(t,1)(Shift(t,1)(Shift(t, -3)(x(t))))) 
+does not actually appear. So `total_sub` in generate_system_equations` is directly 
+initialized with all of the lowered variables `Shift(t, -3)(x) -> x_t-3(t)`, etc. 
 =#
 """
 Generate new derivative variables for the system.
@@ -358,7 +362,7 @@ Effects on the system structure:
 - neweqs: add the identity equations for the new variables, D(x) ~ x_t
 - graph: update graph with the new equations and variables, and their connections
 - solvable_graph:
-- var_eq_matching: match D(x) to the added identity equation
+- var_eq_matching: match D(x) to the added identity equation D(x) ~ x_t
 """
 function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching; mm = nothing, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = ts
@@ -406,7 +410,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
 end
 
 """
-Check if there's `D(x) = x_t` already.
+Check if there's `D(x) ~ x_t` already.
 """
 function find_duplicate_dd(dv, solvable_graph, diff_to_var, linear_eqs, mm)
     for eq in 𝑑neighbors(solvable_graph, dv)
@@ -427,6 +431,10 @@ function find_duplicate_dd(dv, solvable_graph, diff_to_var, linear_eqs, mm)
     return nothing 
 end
 
+"""
+Add a dummy derivative variable x_t corresponding to symbolic variable D(x) 
+which has index dv in `fullvars`. Return the new index of x_t.
+"""
 function add_dd_variable!(s::SystemStructure, fullvars, x_t, dv)
     push!(fullvars, simplify_shifts(x_t))
     v_t = length(fullvars)
@@ -439,7 +447,11 @@ function add_dd_variable!(s::SystemStructure, fullvars, x_t, dv)
     v_t
 end
 
-# dv = index of D(x), v_t = index of x_t
+"""
+Add the equation D(x) - x_t ~ 0 to `neweqs`. `dv` and `v_t` are the indices
+of the higher-order derivative variable and the newly-introduced dummy
+derivative variable. Return the index of the new equation in `neweqs`.
+"""
 function add_dd_equation!(s::SystemStructure, neweqs, eq, dv, v_t)
     push!(neweqs, eq)
     add_vertex!(s.graph, SRC)
@@ -452,8 +464,33 @@ function add_dd_equation!(s::SystemStructure, neweqs, eq, dv, v_t)
 end
 
 """
-Solve the solvable equations of the system and generate differential (or discrete)
-equations in terms of the selected states.
+Solve the equations in `neweqs` to obtain the final equations of the 
+system.
+
+For each equation of `neweqs`, do one of the following: 
+   1. If the equation is solvable for a differentiated variable D(x),
+      then solve for D(x), and add D(x) ~ sol as a differential equation
+      of the system.
+   2. If the equation is solvable for an un-differentiated variable x, 
+      solve for x and then add x ~ sol as a solved equation. These will
+      become observables.
+   3. If the equation is not solvable, add it as an algebraic equation.
+
+Solved equations are added to `total_sub`. Occurrences of differential
+or solved variables on the RHS of the final equations will get substituted.
+The topological sort of the equations ensures that variables are solved for
+before they appear in equations. 
+
+Reorder the equations and unknowns to be:
+   [diffeqs; ...]
+   [diffvars; ...]
+such that the mass matrix is:
+   [I  0
+    0  0].
+
+Order the new equations and variables such that the differential equations
+and variables come first. Return the new equations, the solved equations,
+the new orderings, and the number of solved variables and equations.
 """
 function generate_system_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = state 
@@ -550,6 +587,9 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
     return neweqs, solved_eqs, eq_ordering, var_ordering, length(solved_vars), length(solved_vars_set)
 end
 
+"""
+Occurs when a variable D(x) occurs in a non-differential system.
+"""
 struct UnexpectedDifferentialError
     eq::Equation
 end
@@ -558,12 +598,20 @@ function Base.showerror(io::IO, err::UnexpectedDifferentialError)
     error("Differential found in a non-differential system. Likely this is a bug in the construction of an initialization system. Please report this issue with a reproducible example. Offending equation: $(err.eq)")
 end
 
+"""
+Generate a first-order differential equation whose LHS is `dx`.
+
+`var` and `dx` represent the same variable, but `var` may be a higher-order differential and `dx` is always first-order. For example, if `var` is D(D(x)), then `dx` would be `D(x_t)`. Solve `eq` for `var`, substitute previously solved variables, and return the differential equation.
+"""
 function make_differential_equation(var, dx, eq, total_sub)
     dx ~ simplify_shifts(Symbolics.fixpoint_sub(
         Symbolics.symbolic_linear_solve(eq, var),
         total_sub; operator = ModelingToolkit.Shift))
 end
 
+"""
+Generate an algebraic equation. Substitute solved variables into `eq` and return the equation.
+"""
 function make_algebraic_equation(eq, total_sub)
     rhs = eq.rhs
     if !(eq.lhs isa Number && eq.lhs == 0)
@@ -572,6 +620,9 @@ function make_algebraic_equation(eq, total_sub)
     0 ~ simplify_shifts(Symbolics.fixpoint_sub(rhs, total_sub))
 end
 
+"""
+Solve equation `eq` for `var`, substitute previously solved variables, and return the solved equation.
+"""
 function make_solved_equation(var, eq, total_sub; simplify = false)
     residual = eq.lhs - eq.rhs
     a, b, islinear = linear_expansion(residual, var)
@@ -591,17 +642,13 @@ function make_solved_equation(var, eq, total_sub; simplify = false)
 end
 
 """
-Reorder the equations and unknowns to be:
-   [diffeqs; ...]
-   [diffvars; ...]
-such that the mass matrix is:
-   [I  0
-    0  0].
-
-Update the state to account for the new ordering and equations.
+Given the ordering returned by `generate_system_equations!`, update the 
+tearing state to account for the new order. Permute the variables and equations.
+Eliminate the solved variables and equations from the graph and permute the
+graph's vertices to account for the new variable/equation ordering.
 """
 # TODO: BLT sorting
-function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
+function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_ordering, nsolved_eq, nsolved_var)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
 
     eqsperm = zeros(Int, nsrcs(graph))
@@ -616,7 +663,7 @@ function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_or
     # Contract the vertices in the structure graph to make the structure match
     # the new reality of the system we've just created.
     new_graph = contract_variables(graph, var_eq_matching, varsperm, eqsperm,
-        nelim_eq, nelim_var)
+        nsolved_eq, nsolved_var)
 
     new_var_to_diff = complete(DiffGraph(length(var_ordering)))
     for (v, d) in enumerate(var_to_diff)
@@ -643,7 +690,7 @@ function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_or
 end
 
 """
-Set the system equations, unknowns, observables post-tearing.
+Update the system equations, unknowns, and observables after simplification.
 """
 function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns; 
         cse_hack = true, array_hack = true)
@@ -685,16 +732,10 @@ function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dumm
     sys = schedule(sys)
 end
 
-# Terminology and Definition:
-# A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
-# characterize variables in `u(t)` into two classes: differential variables
-# (denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
-# variables are marked as `SelectedState` and they are differentiated in the
-# DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
-# appear in the system. Algebraic variables are variables that are not
-# differential variables.
     
-# Give the order of the variable indexed by dv
+"""
+Give the order of the variable indexed by dv.
+"""
 function var_order(dv, diff_to_var)
     order = 0
     while (dv′ = diff_to_var[dv]) !== nothing
@@ -704,6 +745,21 @@ function var_order(dv, diff_to_var)
     order, dv
 end
 
+"""
+Main internal function for structural simplification for DAE systems and discrete systems.
+Generate dummy derivative variables, new equations in terms of variables, return updated
+system and tearing state.
+
+Terminology and Definition:
+
+A general DAE is in the form of `F(u'(t), u(t), p, t) == 0`. We can
+characterize variables in `u(t)` into two classes: differential variables
+(denoted `v(t)`) and algebraic variables (denoted `z(t)`). Differential
+variables are marked as `SelectedState` and they are differentiated in the
+DAE system, i.e. `v'(t)` are all the variables in `u'(t)` that actually
+appear in the system. Algebraic variables are variables that are not
+differential variables.
+"""
 function tearing_reassemble(state::TearingState, var_eq_matching,
         full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true)
     extra_vars = Int[]
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index f64a8da132..2753446a94 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -449,7 +449,12 @@ end
 ### Misc
 ###
 
-# For discrete variables. Turn Shift(t, k)(x(t)) into xₜ₋ₖ(t)
+"""
+Handle renaming variable names for discrete structural simplification. Three cases: 
+- positive shift: do nothing
+- zero shift: x(t) => Shift(t, 0)(x(t))
+- negative shift: rename the variable
+"""
 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)
@@ -460,30 +465,46 @@ function lower_shift_varname(var, iv)
     end
 end
 
-function shift2term(var) 
+"""
+Rename a Shift variable with negative shift, Shift(t, k)(x(t)) to xₜ₋ₖ(t).
+"""
+function shift2term(var)
     backshift = operation(var).steps
     iv = operation(var).t
-    num = join(Char(0x2080 + d) for d in reverse!(digits(-backshift)))
-    ds = join([Char(0x209c), Char(0x208b), num])
-    #ds = "$iv-$(-backshift)"
-    #d_separator = 'ˍ'
-
-    if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
-        O = only(arguments(var))
-        oldop = operation(O)
-        newname = Symbol(string(nameof(oldop)), ds)
-    else
-        O = var
-        oldop = operation(var) 
-        varname = split(string(nameof(oldop)), d_separator)[1]
-        newname = Symbol(varname, d_separator, ds)
-    end
+    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)
+
+    O = only(arguments(var))
+    oldop = operation(O)
+    newname = Symbol(string(nameof(oldop)), ds)
+
     newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
     newvar = setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
     newvar = setmetadata(newvar, ModelingToolkit.VariableUnshifted, O)
+    newvar = setmetadata(newvar, ModelingToolkit.VariableShift, backshift)
     return newvar
 end
 
+function term2shift(var)
+    var = Symbolics.unwrap(var)
+    name = Symbolics.getname(var)
+    O = only(arguments(var))
+    oldop = operation(O)
+    iv = only(arguments(x))
+    # Split on ₋
+    if occursin(Char(0x208b), name)
+        substrings = split(name, Char(0x208b))
+        shift = last(split(name, Char(0x208b)))
+        newname = join(substrings[1:end-1])[1:end-1]
+        newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
+        return Shift(iv, -shift)(newvar)
+    else
+        return var
+    end
+end
+
 function isdoubleshift(var)
     return ModelingToolkit.isoperator(var, ModelingToolkit.Shift) &&
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index 9a9ac47853..40bf632339 100644
--- a/src/systems/discrete_system/discrete_system.jl
+++ b/src/systems/discrete_system/discrete_system.jl
@@ -270,15 +270,19 @@ function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
         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)).")
+        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)(only(arguments(k)))).")
         end
+
         updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
     end
     for var in unknowns(sys)
         op = operation(var)
-        haskey(updated, var) && continue
         root = getunshifted(var)
+        shift = getshift(var)
         isnothing(root) && continue
-        haskey(defs, root) || error("Initial condition for $root not provided.")
+        (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
     return updated
diff --git a/src/variables.jl b/src/variables.jl
index a29a119607..4e13ad2c5d 100644
--- a/src/variables.jl
+++ b/src/variables.jl
@@ -6,7 +6,9 @@ struct VariableOutput end
 struct VariableIrreducible end
 struct VariableStatePriority end
 struct VariableMisc end
+# Metadata for renamed shift variables xₜ₋₁
 struct VariableUnshifted end
+struct VariableShift end
 Symbolics.option_to_metadata_type(::Val{:unit}) = VariableUnit
 Symbolics.option_to_metadata_type(::Val{:connect}) = VariableConnectType
 Symbolics.option_to_metadata_type(::Val{:input}) = VariableInput
@@ -15,6 +17,7 @@ Symbolics.option_to_metadata_type(::Val{:irreducible}) = VariableIrreducible
 Symbolics.option_to_metadata_type(::Val{:state_priority}) = VariableStatePriority
 Symbolics.option_to_metadata_type(::Val{:misc}) = VariableMisc
 Symbolics.option_to_metadata_type(::Val{:unshifted}) = VariableUnshifted
+Symbolics.option_to_metadata_type(::Val{:shift}) = VariableShift
 
 """
     dump_variable_metadata(var)
@@ -97,7 +100,7 @@ struct Stream <: AbstractConnectType end   # special stream connector
 
 Get the connect type of x. See also [`hasconnect`](@ref).
 """
-getconnect(x) = getconnect(unwrap(x))
+getconnect(x::Num) = getconnect(unwrap(x))
 getconnect(x::Symbolic) = Symbolics.getmetadata(x, VariableConnectType, nothing)
 """
     hasconnect(x)
@@ -264,7 +267,7 @@ end
 end
 
 struct IsHistory end
-ishistory(x) = ishistory(unwrap(x))
+ishistory(x::Num) = ishistory(unwrap(x))
 ishistory(x::Symbolic) = getmetadata(x, IsHistory, false)
 hist(x, t) = wrap(hist(unwrap(x), t))
 function hist(x::Symbolic, t)
@@ -575,7 +578,7 @@ end
 Fetch any miscellaneous data associated with symbolic variable `x`.
 See also [`hasmisc(x)`](@ref).
 """
-getmisc(x) = getmisc(unwrap(x))
+getmisc(x::Num) = getmisc(unwrap(x))
 getmisc(x::Symbolic) = Symbolics.getmetadata(x, VariableMisc, nothing)
 """
     hasmisc(x)
@@ -594,7 +597,7 @@ setmisc(x, miscdata) = setmetadata(x, VariableMisc, miscdata)
 
 Fetch the unit associated with variable `x`. This function is a metadata getter for an individual variable, while `get_unit` is used for unit inference on more complicated sdymbolic expressions.
 """
-getunit(x) = getunit(unwrap(x))
+getunit(x::Num) = getunit(unwrap(x))
 getunit(x::Symbolic) = Symbolics.getmetadata(x, VariableUnit, nothing)
 """
     hasunit(x)
@@ -603,5 +606,8 @@ Check if the variable `x` has a unit.
 """
 hasunit(x) = getunit(x) !== nothing
 
-getunshifted(x) = getunshifted(unwrap(x))
+getunshifted(x::Num) = getunshifted(unwrap(x))
 getunshifted(x::Symbolic) = Symbolics.getmetadata(x, VariableUnshifted, nothing)
+
+getshift(x::Num) = getshift(unwrap(x))
+getshift(x::Symbolic) = Symbolics.getmetadata(x, VariableShift, 0)
diff --git a/test/discrete_system.jl b/test/discrete_system.jl
index b63f66d4e4..fa7ba993e6 100644
--- a/test/discrete_system.jl
+++ b/test/discrete_system.jl
@@ -220,21 +220,6 @@ sol = solve(prob, FunctionMap())
 
 @test reduce(vcat, sol.u) == 1:11
 
-# test that default values apply to the entire history
-@variables x(t) = 1.0
-@mtkbuild de = DiscreteSystem([x ~ x(k - 1) + x(k - 2)], t)
-prob = DiscreteProblem(de, [], (0, 10))
-@test prob[x] == 2.0
-@test prob[x(k - 1)] == 1.0
-
-# must provide initial conditions for history
-@test_throws ErrorException DiscreteProblem(de, [x => 2.0], (0, 10))
-
-# initial values only affect _that timestep_, not the entire history
-prob = DiscreteProblem(de, [x(k - 1) => 2.0], (0, 10))
-@test prob[x] == 3.0
-@test prob[x(k - 1)] == 2.0
-
 # Issue#2585
 getdata(buffer, t) = buffer[mod1(Int(t), length(buffer))]
 @register_symbolic getdata(buffer::Vector, t)
@@ -272,6 +257,7 @@ k = ShiftIndex(t)
 @named sys = DiscreteSystem([x ~ x^2 + y^2, y ~ x(k - 1) + y(k - 1)], t)
 @test_throws ["algebraic equations", "not yet supported"] structural_simplify(sys)
 
+
 @testset "Passing `nothing` to `u0`" begin
     @variables x(t) = 1
     k = ShiftIndex()
@@ -279,3 +265,55 @@ k = ShiftIndex(t)
     prob = @test_nowarn DiscreteProblem(sys, nothing, (0.0, 1.0))
     @test_nowarn solve(prob, FunctionMap())
 end
+
+@testset "Initialization" begin
+    # test that default values apply to the entire history
+    @variables x(t) = 1.0
+    @mtkbuild de = DiscreteSystem([x ~ x(k - 1) + x(k - 2)], t)
+    prob = DiscreteProblem(de, [], (0, 10))
+    @test prob[x] == 2.0
+    @test prob[x(k - 1)] == 1.0
+    
+    # must provide initial conditions for history
+    @test_throws ErrorException DiscreteProblem(de, [x => 2.0], (0, 10))
+    @test_throws ErrorException DiscreteProblem(de, [x(k+1) => 2.], (0, 10))
+    
+    # initial values only affect _that timestep_, not the entire history
+    prob = DiscreteProblem(de, [x(k - 1) => 2.0], (0, 10))
+    @test prob[x] == 3.0
+    @test prob[x(k - 1)] == 2.0
+    @test prob[xₜ₋₁] == 2.0
+
+    # Test initial assignment with lowered variable
+    @variables xₜ₋₁(t)
+    prob = DiscreteProblem(de, [xₜ₋₁(k-1) => 4.0], (0, 10))
+    @test prob[x(k-1)] == prob[xₜ₋₁] == 1.0
+    @test prob[x] == 5.
+
+    # Test missing initial throws error
+    @variables x(t)
+    @mtkbuild de = DiscreteSystem([x ~ x(k-1) + x(k-2)*x(k-3)], t)
+    @test_throws ErrorException prob = DiscreteProblem(de, [x(k-3) => 2.], (0, 10))
+    @test_throws ErrorException prob = DiscreteProblem(de, [x(k-3) => 2., x(k-1) => 3.], (0, 10))
+
+    # Test non-assigned initials are given default value
+    @variables x(t) = 2.
+    prob = DiscreteProblem(de, [x(k-3) => 12.], (0, 10))
+    @test prob[x] == 26.0
+    @test prob[x(k-1)] == 2.0
+    @test prob[x(k-2)] == 2.0
+
+    # Elaborate test
+    eqs = [x ~ x(k-1) + z(k-2), 
+           z ~ x(k-2) * x(k-3) - z(k-1)^2]
+    @mtkbuild de = DiscreteSystem(eqs, t)
+    @variables xₜ₋₂(t) zₜ₋₁(t)
+    u0 = [x(k-1) => 3, 
+          xₜ₋₂(k-1) => 4, 
+          x(k-2) => 1, 
+          z(k-1) => 5, 
+          zₜ₋₁(k-1) => 12]
+    prob = DiscreteProblem(de, u0, (0, 10))
+    @test prob[x] == 15
+    @test prob[z] == -21
+end

From fc2a309f3d9f8e607329b9ee640672ca4f872eb3 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Wed, 19 Feb 2025 10:12:41 -0800
Subject: [PATCH 40/41] fix tests and shift2term

---
 src/structural_transformation/utils.jl        | 31 ++++++-------------
 .../discrete_system/discrete_system.jl        |  7 +++--
 test/discrete_system.jl                       | 13 ++++++--
 3 files changed, 25 insertions(+), 26 deletions(-)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 2753446a94..96f7f78f99 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -469,16 +469,21 @@ end
 Rename a Shift variable with negative shift, Shift(t, k)(x(t)) to xₜ₋ₖ(t).
 """
 function shift2term(var)
-    backshift = operation(var).steps
-    iv = operation(var).t
+    op = operation(var)
+    iv = op.t
+    arg = only(arguments(var))
+    is_lowered = !isnothing(ModelingToolkit.getunshifted(arg))
+
+    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)
 
-    O = only(arguments(var))
+    O = is_lowered ? ModelingToolkit.getunshifted(arg) : arg
     oldop = operation(O)
-    newname = Symbol(string(nameof(oldop)), ds)
+    newname = backshift != 0 ? Symbol(string(nameof(oldop)), ds) : Symbol(string(nameof(oldop)))
 
     newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
     newvar = setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
@@ -487,24 +492,6 @@ function shift2term(var)
     return newvar
 end
 
-function term2shift(var)
-    var = Symbolics.unwrap(var)
-    name = Symbolics.getname(var)
-    O = only(arguments(var))
-    oldop = operation(O)
-    iv = only(arguments(x))
-    # Split on ₋
-    if occursin(Char(0x208b), name)
-        substrings = split(name, Char(0x208b))
-        shift = last(split(name, Char(0x208b)))
-        newname = join(substrings[1:end-1])[1:end-1]
-        newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
-        return Shift(iv, -shift)(newvar)
-    else
-        return var
-    end
-end
-
 function isdoubleshift(var)
     return ModelingToolkit.isoperator(var, ModelingToolkit.Shift) &&
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index 40bf632339..3f8d9e85c6 100644
--- a/src/systems/discrete_system/discrete_system.jl
+++ b/src/systems/discrete_system/discrete_system.jl
@@ -269,12 +269,15 @@ function shift_u0map_forward(sys::DiscreteSystem, 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[Shift(iv, 1)(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)(only(arguments(k)))).")
+        else
+            updated[Shift(iv, op.steps + 1)(only(arguments(k)))] = v
         end
 
-        updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
     end
     for var in unknowns(sys)
         op = operation(var)
diff --git a/test/discrete_system.jl b/test/discrete_system.jl
index fa7ba993e6..756e5bca48 100644
--- a/test/discrete_system.jl
+++ b/test/discrete_system.jl
@@ -282,10 +282,10 @@ end
     prob = DiscreteProblem(de, [x(k - 1) => 2.0], (0, 10))
     @test prob[x] == 3.0
     @test prob[x(k - 1)] == 2.0
+    @variables xₜ₋₁(t)
     @test prob[xₜ₋₁] == 2.0
 
     # Test initial assignment with lowered variable
-    @variables xₜ₋₁(t)
     prob = DiscreteProblem(de, [xₜ₋₁(k-1) => 4.0], (0, 10))
     @test prob[x(k-1)] == prob[xₜ₋₁] == 1.0
     @test prob[x] == 5.
@@ -298,16 +298,17 @@ end
 
     # Test non-assigned initials are given default value
     @variables x(t) = 2.
+    @mtkbuild de = DiscreteSystem([x ~ x(k-1) + x(k-2)*x(k-3)], t)
     prob = DiscreteProblem(de, [x(k-3) => 12.], (0, 10))
     @test prob[x] == 26.0
     @test prob[x(k-1)] == 2.0
     @test prob[x(k-2)] == 2.0
 
     # Elaborate test
+    @variables xₜ₋₂(t) zₜ₋₁(t) z(t)
     eqs = [x ~ x(k-1) + z(k-2), 
            z ~ x(k-2) * x(k-3) - z(k-1)^2]
     @mtkbuild de = DiscreteSystem(eqs, t)
-    @variables xₜ₋₂(t) zₜ₋₁(t)
     u0 = [x(k-1) => 3, 
           xₜ₋₂(k-1) => 4, 
           x(k-2) => 1, 
@@ -316,4 +317,12 @@ end
     prob = DiscreteProblem(de, u0, (0, 10))
     @test prob[x] == 15
     @test prob[z] == -21
+
+    import ModelingToolkit: shift2term
+    # unknowns(de) = xₜ₋₁, x, zₜ₋₁, xₜ₋₂, z
+    vars = ModelingToolkit.value.(unknowns(de))
+    @test isequal(shift2term(Shift(t, 1)(vars[1])), vars[2])
+    @test isequal(shift2term(Shift(t, 1)(vars[4])), vars[1])
+    @test isequal(shift2term(Shift(t, -1)(vars[5])), vars[3])
+    @test isequal(shift2term(Shift(t, -2)(vars[2])), vars[4])
 end

From f29824e829041f01f60a32ada65c5249f91a9746 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 24 Feb 2025 09:25:21 -0500
Subject: [PATCH 41/41] Format

---
 .../StructuralTransformations.jl              |  6 ++-
 .../symbolics_tearing.jl                      | 54 +++++++++++--------
 src/structural_transformation/utils.jl        |  8 +--
 .../discrete_system/discrete_system.jl        |  6 +--
 src/systems/systemstructure.jl                |  6 +--
 src/variables.jl                              |  2 +-
 test/discrete_system.jl                       | 44 +++++++--------
 7 files changed, 69 insertions(+), 57 deletions(-)

diff --git a/src/structural_transformation/StructuralTransformations.jl b/src/structural_transformation/StructuralTransformations.jl
index 7d2b9afa26..f0124d7f4b 100644
--- a/src/structural_transformation/StructuralTransformations.jl
+++ b/src/structural_transformation/StructuralTransformations.jl
@@ -11,7 +11,8 @@ using SymbolicUtils: maketerm, iscall
 
 using ModelingToolkit
 using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Differential,
-                       unknowns, equations, vars, Symbolic, diff2term_with_unit, shift2term_with_unit, value,
+                       unknowns, equations, vars, Symbolic, diff2term_with_unit,
+                       shift2term_with_unit, value,
                        operation, arguments, Sym, Term, simplify, symbolic_linear_solve,
                        isdiffeq, isdifferential, isirreducible,
                        empty_substitutions, get_substitutions,
@@ -22,7 +23,8 @@ using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Di
                        get_postprocess_fbody, vars!,
                        IncrementalCycleTracker, add_edge_checked!, topological_sort,
                        invalidate_cache!, Substitutions, get_or_construct_tearing_state,
-                       filter_kwargs, lower_varname_with_unit, lower_shift_varname_with_unit, setio, SparseMatrixCLIL,
+                       filter_kwargs, lower_varname_with_unit,
+                       lower_shift_varname_with_unit, setio, SparseMatrixCLIL,
                        get_fullvars, has_equations, observed,
                        Schedule, schedule
 
diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl
index bbb6853a7c..ea594f2172 100644
--- a/src/structural_transformation/symbolics_tearing.jl
+++ b/src/structural_transformation/symbolics_tearing.jl
@@ -248,7 +248,8 @@ called dummy derivatives.
 State selection is done. All non-differentiated variables are algebraic 
 variables, and all variables that appear differentiated are differential variables.
 """
-function substitute_derivatives_algevars!(ts::TearingState, neweqs, var_eq_matching, dummy_sub; iv = nothing, D = nothing)
+function substitute_derivatives_algevars!(
+        ts::TearingState, neweqs, var_eq_matching, dummy_sub; iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     diff_to_var = invview(var_to_diff)
@@ -288,7 +289,7 @@ end
 #= 
 There are three cases where we want to generate new variables to convert
 the system into first order (semi-implicit) ODEs.
-    
+
 1. To first order:
 Whenever higher order differentiated variable like `D(D(D(x)))` appears,
 we introduce new variables `x_t`, `x_tt`, and `x_ttt` and new equations
@@ -364,7 +365,8 @@ Effects on the system structure:
 - solvable_graph:
 - var_eq_matching: match D(x) to the added identity equation D(x) ~ x_t
 """
-function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matching; mm = nothing, iv = nothing, D = nothing)
+function generate_derivative_variables!(
+        ts::TearingState, neweqs, var_eq_matching; mm = nothing, iv = nothing, D = nothing)
     @unpack fullvars, sys, structure = ts
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
@@ -395,7 +397,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) :
-                          lower_varname_with_unit(fullvars[lv], iv, order)
+              lower_varname_with_unit(fullvars[lv], iv, order)
 
         # Add `x_t` to the graph
         v_t = add_dd_variable!(structure, fullvars, x_t, dv)
@@ -405,7 +407,7 @@ function generate_derivative_variables!(ts::TearingState, neweqs, var_eq_matchin
         # Update matching
         push!(var_eq_matching, unassigned)
         var_eq_matching[dv] = unassigned
-        eq_var_matching[dummy_eq] = dv 
+        eq_var_matching[dummy_eq] = dv
     end
 end
 
@@ -428,7 +430,7 @@ function find_duplicate_dd(dv, solvable_graph, diff_to_var, linear_eqs, mm)
             return eq, v_t
         end
     end
-    return nothing 
+    return nothing
 end
 
 """
@@ -492,8 +494,9 @@ Order the new equations and variables such that the differential equations
 and variables come first. Return the new equations, the solved equations,
 the new orderings, and the number of solved variables and equations.
 """
-function generate_system_equations!(state::TearingState, neweqs, var_eq_matching; simplify = false, iv = nothing, D = nothing)
-    @unpack fullvars, sys, structure = state 
+function generate_system_equations!(state::TearingState, neweqs, var_eq_matching;
+        simplify = false, iv = nothing, D = nothing)
+    @unpack fullvars, sys, structure = state
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
     eq_var_matching = invview(var_eq_matching)
     diff_to_var = invview(var_to_diff)
@@ -502,11 +505,12 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
     if is_only_discrete(structure)
         for (i, v) in enumerate(fullvars)
             op = operation(v)
-            op isa Shift && (op.steps < 0) && begin
-                lowered = lower_shift_varname_with_unit(v, iv)
-                total_sub[v] = lowered
-                fullvars[i] = lowered
-            end
+            op isa Shift && (op.steps < 0) &&
+                begin
+                    lowered = lower_shift_varname_with_unit(v, iv)
+                    total_sub[v] = lowered
+                    fullvars[i] = lowered
+                end
         end
     end
 
@@ -581,10 +585,11 @@ function generate_system_equations!(state::TearingState, neweqs, var_eq_matching
     end
     solved_vars_set = BitSet(solved_vars)
     var_ordering = [diff_vars;
-                   setdiff!(setdiff(1:ndsts(graph), diff_vars_set),
-                       solved_vars_set)]
+                    setdiff!(setdiff(1:ndsts(graph), diff_vars_set),
+                        solved_vars_set)]
 
-    return neweqs, solved_eqs, eq_ordering, var_ordering, length(solved_vars), length(solved_vars_set)
+    return neweqs, solved_eqs, eq_ordering, var_ordering, length(solved_vars),
+    length(solved_vars_set)
 end
 
 """
@@ -648,7 +653,8 @@ Eliminate the solved variables and equations from the graph and permute the
 graph's vertices to account for the new variable/equation ordering.
 """
 # TODO: BLT sorting
-function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering, var_ordering, nsolved_eq, nsolved_var)
+function reorder_vars!(state::TearingState, var_eq_matching, eq_ordering,
+        var_ordering, nsolved_eq, nsolved_var)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
 
     eqsperm = zeros(Int, nsrcs(graph))
@@ -692,7 +698,8 @@ end
 """
 Update the system equations, unknowns, and observables after simplification.
 """
-function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns; 
+function update_simplified_system!(
+        state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns;
         cse_hack = true, array_hack = true)
     @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure
     diff_to_var = invview(var_to_diff)
@@ -732,7 +739,6 @@ function update_simplified_system!(state::TearingState, neweqs, solved_eqs, dumm
     sys = schedule(sys)
 end
 
-    
 """
 Give the order of the variable indexed by dv.
 """
@@ -790,12 +796,14 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
 
     generate_derivative_variables!(state, neweqs, var_eq_matching; mm, iv, D)
 
-    neweqs, solved_eqs, eq_ordering, var_ordering, nelim_eq, nelim_var = 
-        generate_system_equations!(state, neweqs, var_eq_matching; simplify, iv, D)
+    neweqs, solved_eqs, eq_ordering, var_ordering, nelim_eq, nelim_var = generate_system_equations!(
+        state, neweqs, var_eq_matching; simplify, iv, D)
 
-    state = reorder_vars!(state, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
+    state = reorder_vars!(
+        state, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var)
 
-    sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns; cse_hack, array_hack)
+    sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_eq_matching,
+        extra_unknowns; cse_hack, array_hack)
 
     @set! state.sys = sys
     @set! sys.tearing_state = state
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 96f7f78f99..f3cea9c7ba 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -477,15 +477,17 @@ 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]) 
+    ds = join([Char(0x209c), Char(0x208b), num])
     # Char(0x209c) = ₜ
     # Char(0x208b) = ₋ (subscripted minus)
 
     O = is_lowered ? ModelingToolkit.getunshifted(arg) : arg
     oldop = operation(O)
-    newname = backshift != 0 ? Symbol(string(nameof(oldop)), ds) : Symbol(string(nameof(oldop)))
+    newname = backshift != 0 ? Symbol(string(nameof(oldop)), ds) :
+              Symbol(string(nameof(oldop)))
 
-    newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname), Symbolics.children(O), Symbolics.metadata(O))
+    newvar = maketerm(typeof(O), Symbolics.rename(oldop, newname),
+        Symbolics.children(O), Symbolics.metadata(O))
     newvar = setmetadata(newvar, Symbolics.VariableSource, (:variables, newname))
     newvar = setmetadata(newvar, ModelingToolkit.VariableUnshifted, O)
     newvar = setmetadata(newvar, ModelingToolkit.VariableShift, backshift)
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index 3f8d9e85c6..e4949d812f 100644
--- a/src/systems/discrete_system/discrete_system.jl
+++ b/src/systems/discrete_system/discrete_system.jl
@@ -269,15 +269,15 @@ function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
     for k in collect(keys(u0map))
         v = u0map[k]
         if !((op = operation(k)) isa Shift)
-            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)).")
-            
+            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[Shift(iv, 1)(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)(only(arguments(k)))).")
         else
             updated[Shift(iv, op.steps + 1)(only(arguments(k)))] = v
         end
-
     end
     for var in unknowns(sys)
         op = operation(var)
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
index 6fee78cfd6..0643f32ec4 100644
--- a/src/systems/systemstructure.jl
+++ b/src/systems/systemstructure.jl
@@ -473,9 +473,9 @@ function shift_discrete_system(ts::TearingState)
     end
     iv = get_iv(sys)
 
-    discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, 1)(k))    
-                   for k in discvars 
-                   if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold})) 
+    discmap = Dict(k => StructuralTransformations.simplify_shifts(Shift(iv, 1)(k))
+    for k in discvars
+    if any(isequal(k), fullvars) && !isa(operation(k), Union{Sample, Hold}))
 
     for i in eachindex(fullvars)
         fullvars[i] = StructuralTransformations.simplify_shifts(fast_substitute(
diff --git a/src/variables.jl b/src/variables.jl
index 4e13ad2c5d..a7dde165f9 100644
--- a/src/variables.jl
+++ b/src/variables.jl
@@ -138,7 +138,7 @@ function default_toterm(x)
     if iscall(x) && (op = operation(x)) isa Operator
         if !(op isa Differential)
             if op isa Shift && op.steps < 0
-                return shift2term(x) 
+                return shift2term(x)
             end
             x = normalize_to_differential(op)(arguments(x)...)
         end
diff --git a/test/discrete_system.jl b/test/discrete_system.jl
index 756e5bca48..f232faaf81 100644
--- a/test/discrete_system.jl
+++ b/test/discrete_system.jl
@@ -257,7 +257,6 @@ k = ShiftIndex(t)
 @named sys = DiscreteSystem([x ~ x^2 + y^2, y ~ x(k - 1) + y(k - 1)], t)
 @test_throws ["algebraic equations", "not yet supported"] structural_simplify(sys)
 
-
 @testset "Passing `nothing` to `u0`" begin
     @variables x(t) = 1
     k = ShiftIndex()
@@ -273,11 +272,11 @@ end
     prob = DiscreteProblem(de, [], (0, 10))
     @test prob[x] == 2.0
     @test prob[x(k - 1)] == 1.0
-    
+
     # must provide initial conditions for history
     @test_throws ErrorException DiscreteProblem(de, [x => 2.0], (0, 10))
-    @test_throws ErrorException DiscreteProblem(de, [x(k+1) => 2.], (0, 10))
-    
+    @test_throws ErrorException DiscreteProblem(de, [x(k + 1) => 2.0], (0, 10))
+
     # initial values only affect _that timestep_, not the entire history
     prob = DiscreteProblem(de, [x(k - 1) => 2.0], (0, 10))
     @test prob[x] == 3.0
@@ -286,34 +285,35 @@ end
     @test prob[xₜ₋₁] == 2.0
 
     # Test initial assignment with lowered variable
-    prob = DiscreteProblem(de, [xₜ₋₁(k-1) => 4.0], (0, 10))
-    @test prob[x(k-1)] == prob[xₜ₋₁] == 1.0
-    @test prob[x] == 5.
+    prob = DiscreteProblem(de, [xₜ₋₁(k - 1) => 4.0], (0, 10))
+    @test prob[x(k - 1)] == prob[xₜ₋₁] == 1.0
+    @test prob[x] == 5.0
 
     # Test missing initial throws error
     @variables x(t)
-    @mtkbuild de = DiscreteSystem([x ~ x(k-1) + x(k-2)*x(k-3)], t)
-    @test_throws ErrorException prob = DiscreteProblem(de, [x(k-3) => 2.], (0, 10))
-    @test_throws ErrorException prob = DiscreteProblem(de, [x(k-3) => 2., x(k-1) => 3.], (0, 10))
+    @mtkbuild de = DiscreteSystem([x ~ x(k - 1) + x(k - 2) * x(k - 3)], t)
+    @test_throws ErrorException prob=DiscreteProblem(de, [x(k - 3) => 2.0], (0, 10))
+    @test_throws ErrorException prob=DiscreteProblem(
+        de, [x(k - 3) => 2.0, x(k - 1) => 3.0], (0, 10))
 
     # Test non-assigned initials are given default value
-    @variables x(t) = 2.
-    @mtkbuild de = DiscreteSystem([x ~ x(k-1) + x(k-2)*x(k-3)], t)
-    prob = DiscreteProblem(de, [x(k-3) => 12.], (0, 10))
+    @variables x(t) = 2.0
+    @mtkbuild de = DiscreteSystem([x ~ x(k - 1) + x(k - 2) * x(k - 3)], t)
+    prob = DiscreteProblem(de, [x(k - 3) => 12.0], (0, 10))
     @test prob[x] == 26.0
-    @test prob[x(k-1)] == 2.0
-    @test prob[x(k-2)] == 2.0
+    @test prob[x(k - 1)] == 2.0
+    @test prob[x(k - 2)] == 2.0
 
     # Elaborate test
     @variables xₜ₋₂(t) zₜ₋₁(t) z(t)
-    eqs = [x ~ x(k-1) + z(k-2), 
-           z ~ x(k-2) * x(k-3) - z(k-1)^2]
+    eqs = [x ~ x(k - 1) + z(k - 2),
+        z ~ x(k - 2) * x(k - 3) - z(k - 1)^2]
     @mtkbuild de = DiscreteSystem(eqs, t)
-    u0 = [x(k-1) => 3, 
-          xₜ₋₂(k-1) => 4, 
-          x(k-2) => 1, 
-          z(k-1) => 5, 
-          zₜ₋₁(k-1) => 12]
+    u0 = [x(k - 1) => 3,
+        xₜ₋₂(k - 1) => 4,
+        x(k - 2) => 1,
+        z(k - 1) => 5,
+        zₜ₋₁(k - 1) => 12]
     prob = DiscreteProblem(de, u0, (0, 10))
     @test prob[x] == 15
     @test prob[z] == -21