Merge pull request #3386 from vyudu/implicit_discrete_system

feat: `ImplicitDiscreteSystem`

Author: ChrisRackauckas

docs/pages.jl

@@ -43,7 +43,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/DiscreteSystem.md

@@ -26,3 +26,9 @@ structural_simplify
 DiscreteProblem(sys::DiscreteSystem, u0map, tspan)
 DiscreteFunction(sys::DiscreteSystem, args...)
 ```
+
+## Discrete Domain
+
+```@docs; canonical=false
+Shift
+```

docs/src/systems/ImplicitDiscreteSystem.md

@@ -0,0 +1,34 @@
+# 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...)
+```
+
+## Discrete Domain
+
+```@docs; canonical=false
+Shift
+```

docs/src/tutorials/SampledData.md

@@ -25,8 +25,10 @@ The operators [`Sample`](@ref) and [`Hold`](@ref) are thus providing the interfa
 The [`ShiftIndex`](@ref) operator is used to refer to past and future values of discrete-time variables. The example below illustrates its use, implementing the discrete-time system
 
 ```math
-x(k+1) = 0.5x(k) + u(k)
-y(k) = x(k)
+\begin{align}
+    x(k+1) &= 0.5x(k) + u(k) \\ 
+    y(k) &= x(k)
+\end{align}
 ```
 
 ```@example clocks
@@ -187,3 +189,10 @@ connections = [r ~ sin(t)          # reference signal
 
 @named cl = ODESystem(connections, t, systems = [f, c, p])
 ```
+
+```@docs
+Sample
+Hold
+ShiftIndex
+Clock
+```

src/ModelingToolkit.jl

@@ -127,6 +127,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
 
@@ -174,6 +175,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")
 
@@ -235,6 +237,8 @@ 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
@@ -298,7 +302,7 @@ export debug_system
 #export ContinuousClock, Discrete, sampletime, input_timedomain, output_timedomain
 #export has_discrete_domain, has_continuous_domain
 #export is_discrete_domain, is_continuous_domain, is_hybrid_domain
-export Sample, Hold, Shift, ShiftIndex, sampletime, SampleTime
+export Sample, Hold, Shift, ShiftIndex, sampletime, SampleTime, Next, Prev
 export Clock, SolverStepClock, TimeDomain
 
 export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunables

src/doc

@@ -0,0 +1 @@
+

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,45 @@ 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)