feat: ImplicitDiscreteSystem

docs/pages.jl

@@ -42,7 +42,8 @@ pages = [
         "systems/NonlinearSystem.md",
         "systems/OptimizationSystem.md",
         "systems/PDESystem.md",
-        "systems/DiscreteSystem.md"],
+        "systems/DiscreteSystem.md",
+        "systems/ImplicitDiscreteSystem.md"],
     "comparison.md",
     "internals.md"
 ]

docs/src/systems/ImplicitDiscreteSystem.md

@@ -0,0 +1,28 @@
+# ImplicitDiscreteSystem
+
+## System Constructors
+
+```@docs
+ImplicitDiscreteSystem
+```
+
+## Composition and Accessor Functions
+
+  - `get_eqs(sys)` or `equations(sys)`: The equations that define the implicit discrete system.
+  - `get_unknowns(sys)` or `unknowns(sys)`: The set of unknowns in the implicit discrete system.
+  - `get_ps(sys)` or `parameters(sys)`: The parameters of the implicit discrete system.
+  - `get_iv(sys)`: The independent variable of the implicit discrete system
+  - `discrete_events(sys)`: The set of discrete events in the implicit discrete system.
+
+## Transformations
+
+```@docs; canonical=false
+structural_simplify
+```
+
+## Problem Constructors
+
+```@docs; canonical=false
+ImplicitDiscreteProblem(sys::ImplicitDiscreteSystem, u0map, tspan)
+ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...)
+```

src/ModelingToolkit.jl

@@ -124,6 +124,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
 
@@ -171,6 +172,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")
 
@@ -232,6 +234,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

src/discretedomain.jl

@@ -70,6 +70,19 @@ Base.literal_pow(f::typeof(^), D::Shift, ::Val{n}) where {n} = Shift(D.t, D.step
 
 hasshift(eq::Equation) = hasshift(eq.lhs) || hasshift(eq.rhs)
 
+"""
+    Next(x)
+
+An alias for Shift(t, 1)(x).
+"""
+Next(x) = Shift(t, 1)(x)
+"""
+    Prev(x)
+
+An alias for Shift(t, -1)(x).
+"""
+Prev(x) = Shift(t, -1)(x)
+
 """
     hasshift(O)
 

src/structural_transformation/StructuralTransformations.jl

@@ -65,7 +65,7 @@ export torn_system_jacobian_sparsity
 export full_equations
 export but_ordered_incidence, lowest_order_variable_mask, highest_order_variable_mask
 export computed_highest_diff_variables
-export shift2term, lower_shift_varname
+export shift2term, lower_shift_varname, simplify_shifts, distribute_shift
 
 include("utils.jl")
 include("pantelides.jl")

src/structural_transformation/symbolics_tearing.jl

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

src/structural_transformation/utils.jl

@@ -457,11 +457,10 @@ Handle renaming variable names for discrete structural simplification. Three cas
 """
 function lower_shift_varname(var, iv)
     op = operation(var)
-    op isa Shift || return Shift(iv, 0)(var, true) # hack to prevent simplification of x(t) - x(t)
-    if op.steps < 0
+    if op isa Shift
         return shift2term(var)
     else
-        return var
+        return Shift(iv, 0)(var, true)
     end
 end
 
@@ -476,10 +475,14 @@ function shift2term(var)
 
     backshift = is_lowered ? op.steps + ModelingToolkit.getshift(arg) : op.steps
 
-    num = join(Char(0x2080 + d) for d in reverse!(digits(-backshift))) # subscripted number, e.g. ₁
-    ds = join([Char(0x209c), Char(0x208b), num])
-    # Char(0x209c) = ₜ
     # Char(0x208b) = ₋ (subscripted minus)
+    # Char(0x208a) = ₊ (subscripted plus)
+    pm = backshift > 0 ? Char(0x208a) : Char(0x208b)
+    # subscripted number, e.g. ₁
+    num = join(Char(0x2080 + d) for d in reverse!(digits(abs(backshift))))
+    # Char(0x209c) = ₜ
+    # ds = ₜ₋₁
+    ds = join([Char(0x209c), pm, num])
 
     O = is_lowered ? ModelingToolkit.getunshifted(arg) : arg
     oldop = operation(O)
@@ -499,6 +502,9 @@ function isdoubleshift(var)
            ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
 end
 
+"""
+Simplify multiple shifts: Shift(t, k1)(Shift(t, k2)(x)) becomes Shift(t, k1+k2)(x).
+"""
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
     var isa Equation && return simplify_shifts(var.lhs) ~ simplify_shifts(var.rhs)
@@ -518,3 +524,43 @@ function simplify_shifts(var)
             unwrap(var).metadata)
     end
 end
+
+"""
+Distribute a shift applied to a whole expression or equation. 
+Shift(t, 1)(x + y) will become Shift(t, 1)(x) + Shift(t, 1)(y).
+Only shifts variables whose independent variable is the same t that appears in the Shift (i.e. constants, time-independent parameters, etc. do not get shifted).
+"""
+function distribute_shift(var) 
+    var = unwrap(var)
+    var isa Equation && return distribute_shift(var.lhs) ~ distribute_shift(var.rhs)
+
+    ModelingToolkit.hasshift(var) || return var
+    shift = operation(var)
+    shift isa Shift || return var
+
+    shift = operation(var)
+    expr = only(arguments(var))
+    if expr isa Equation
+        return distribute_shift(shift(expr.lhs)) ~ distribute_shift(shift(expr.rhs))
+    end
+    shiftexpr = _distribute_shift(expr, shift)
+    return simplify_shifts(shiftexpr)
+end
+
+function _distribute_shift(expr, shift)
+    if iscall(expr)
+        op = operation(expr)
+        args = arguments(expr)
+
+        if ModelingToolkit.isvariable(expr)
+            (length(args) == 1 && isequal(shift.t, only(args))) ? (return shift(expr)) : (return expr)
+        elseif op isa Shift
+            return shift(expr)
+        else
+            return maketerm(typeof(expr), operation(expr), Base.Fix2(_distribute_shift, shift).(args),
+                unwrap(expr).metadata)
+        end
+    else
+        return expr
+    end
+end

src/systems/discrete_system/discrete_system.jl

@@ -17,7 +17,7 @@ eqs = [x(k+1) ~ σ*(y-x),
 @named de = DiscreteSystem(eqs)
 ```
 """
-struct DiscreteSystem <: AbstractTimeDependentSystem
+struct DiscreteSystem <: AbstractDiscreteSystem
     """
     A tag for the system. If two systems have the same tag, then they are
     structurally identical.
@@ -237,6 +237,8 @@ function DiscreteSystem(eqs, iv; kwargs...)
         collect(allunknowns), collect(new_ps); kwargs...)
 end
 
+DiscreteSystem(eq::Equation, args...; kwargs...) = DiscreteSystem([eq], args...; kwargs...)
+
 function flatten(sys::DiscreteSystem, noeqs = false)
     systems = get_systems(sys)
     if isempty(systems)