fix: fix initialization cases

Author: vyudu

src/ModelingToolkit.jl

@@ -123,6 +123,7 @@ abstract type AbstractTimeIndependentSystem <: AbstractSystem end
 abstract type AbstractODESystem <: AbstractTimeDependentSystem end
 abstract type AbstractMultivariateSystem <: AbstractSystem end
 abstract type AbstractOptimizationSystem <: AbstractTimeIndependentSystem end
+abstract type AbstractDiscreteSystem <: AbstractTimeDependentSystem end
 
 function independent_variable end
 
@@ -165,6 +166,7 @@ include("systems/diffeqs/modelingtoolkitize.jl")
 include("systems/diffeqs/basic_transformations.jl")
 
 include("systems/discrete_system/discrete_system.jl")
+include("systems/discrete_system/implicit_discrete_system.jl")
 
 include("systems/jumps/jumpsystem.jl")
 
@@ -230,6 +232,7 @@ export DAEFunctionExpr, DAEProblemExpr
 export SDESystem, SDEFunction, SDEFunctionExpr, SDEProblemExpr
 export SystemStructure
 export DiscreteSystem, DiscreteProblem, DiscreteFunction, DiscreteFunctionExpr
+export ImplicitDiscreteSystem, ImplicitDiscreteProblem, ImplicitDiscreteFunction, ImplicitDiscreteFunctionExpr
 export JumpSystem
 export ODEProblem, SDEProblem
 export NonlinearFunction, NonlinearFunctionExpr

src/structural_transformation/StructuralTransformations.jl

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

src/structural_transformation/utils.jl

@@ -452,11 +452,10 @@ end
 # For discrete variables. Turn Shift(t, k)(x(t)) into xₜ₋ₖ(t)
 function lower_shift_varname(var, iv)
     op = operation(var)
-    op isa Shift || return Shift(iv, 0)(var, true) # hack to prevent simplification of x(t) - x(t)
-    if op.steps < 0
+    if op isa Shift && op.steps < 0
         return shift2term(var)
     else
-        return var
+        return Shift(iv, 0)(var, true)
     end
 end
 

src/systems/discrete_system/discrete_system.jl

@@ -17,7 +17,7 @@ eqs = [x(k+1) ~ σ*(y-x),
 @named de = DiscreteSystem(eqs)
 ```
 """
-struct DiscreteSystem <: AbstractTimeDependentSystem
+struct DiscreteSystem <: AbstractDiscreteSystem
     """
     A tag for the system. If two systems have the same tag, then they are
     structurally identical.
@@ -237,6 +237,8 @@ function DiscreteSystem(eqs, iv; kwargs...)
         collect(allunknowns), collect(new_ps); kwargs...)
 end
 
+DiscreteSystem(eq::Equation, args...; kwargs...) = DiscreteSystem([eq], args...; kwargs...)
+
 function flatten(sys::DiscreteSystem, noeqs = false)
     systems = get_systems(sys)
     if isempty(systems)
@@ -271,14 +273,16 @@ function shift_u0map_forward(sys::DiscreteSystem, u0map, defs)
         if !((op = operation(k)) isa Shift)
             error("Initial conditions must be for the past state of the unknowns. Instead of providing the condition for $k, provide the condition for $(Shift(iv, -1)(k)).")
         end
-        updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
+        k_next = Shift(iv, op.steps + 1)(arguments(k)[1])
+        operation(k_next) isa Shift ? updated[shift2term(k_next)] = v :
+                                      updated[k_next] = v
     end
     for var in unknowns(sys)
         op = operation(var)
         haskey(updated, var) && continue
         root = getunshifted(var)
         isnothing(root) && continue
-        haskey(defs, root) || error("Initial condition for $root not provided.")
+        haskey(defs, root) || error("Initial condition for $var not provided.")
         updated[var] = defs[root]
     end
     return updated

src/systems/discrete_system/implicit_discrete_system.jl

@@ -239,6 +239,8 @@ function ImplicitDiscreteSystem(eqs, iv; kwargs...)
         collect(allunknowns), collect(new_ps); kwargs...)
 end
 
+ImplicitDiscreteSystem(eq::Equation, args...; kwargs...) = ImplicitDiscreteSystem([eq], args...; kwargs...)
+
 function flatten(sys::ImplicitDiscreteSystem, noeqs = false)
     systems = get_systems(sys)
     if isempty(systems)
@@ -261,10 +263,14 @@ end
 
 function generate_function(
         sys::ImplicitDiscreteSystem, dvs = unknowns(sys), ps = parameters(sys); wrap_code = identity, kwargs...)
-    exprs = [eq.lhs - eq.rhs for eq in equations(sys)]
-    u = dvs
-    u_next = map(Shift(iv, 1), u)
-    generate_custom_function(sys, exprs, u_next, u, ps..., get_iv(sys); kwargs...)
+    iv = get_iv(sys)
+    exprs = map(equations(sys)) do eq
+        _iszero(eq.lhs) ? eq.rhs : (simplify_shifts(Shift(iv, -1)(eq.rhs)) - simplify_shifts(Shift(iv, -1)(eq.lhs)))
+    end
+
+    u_next = dvs
+    u = map(Shift(iv, -1), u_next)
+    build_function_wrapper(sys, exprs, u_next, u, ps..., iv; p_start = 3, kwargs...)
 end
 
 function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
@@ -275,13 +281,13 @@ function shift_u0map_forward(sys::ImplicitDiscreteSystem, u0map, defs)
         if !((op = operation(k)) isa Shift)
             error("Initial conditions must be for the past state of the unknowns. Instead of providing the condition for $k, provide the condition for $(Shift(iv, -1)(k)).")
         end
-        updated[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
+        updated[shift2term(k)] = v
     end
     for var in unknowns(sys)
         op = operation(var)
-        op isa Shift || continue
         haskey(updated, var) && continue
-        root = first(arguments(var))
+        root = getunshifted(var)
+        isnothing(root) && continue
         haskey(defs, root) || error("Initial condition for $var not provided.")
         updated[var] = defs[root]
     end
@@ -301,7 +307,7 @@ function SciMLBase.ImplicitDiscreteProblem(
         kwargs...
 )
     if !iscomplete(sys)
-        error("A completed `ImplicitDiscreteSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `ImplicitDiscreteProblem`")
+        error("A completed `ImplicitDiscreteSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `ImplicitDiscreteProblem`.")
     end
     dvs = unknowns(sys)
     ps = parameters(sys)
@@ -312,8 +318,7 @@ function SciMLBase.ImplicitDiscreteProblem(
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
     f, u0, p = process_SciMLProblem(
         ImplicitDiscreteFunction, sys, u0map, parammap; eval_expression, eval_module)
-    u0 = f(u0, p, tspan[1])
-    NonlinearProblem(f, u0, tspan, p; kwargs...)
+    ImplicitDiscreteProblem(f, u0, tspan, p; kwargs...)
 end
 
 function SciMLBase.ImplicitDiscreteFunction(sys::ImplicitDiscreteSystem, args...; kwargs...)

src/systems/systems.jl

@@ -42,8 +42,8 @@ function structural_simplify(
     if newsys isa DiscreteSystem &&
        any(eq -> symbolic_type(eq.lhs) == NotSymbolic(), equations(newsys))
         error("""
-            Encountered algebraic equations when simplifying discrete system. This is \
-            not yet supported.
+            Encountered algebraic equations when simplifying discrete system. Please construct \
+            an ImplicitDiscreteSystem instead.
         """)
     end
     for pass in additional_passes

src/systems/systemstructure.jl

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

test/implicit_discrete_system.jl

@@ -0,0 +1,30 @@
+using ModelingToolkit, Test
+using ModelingToolkit: t_nounits as t
+
+k = ShiftIndex(t)
+@variables x(t) = 1
+@mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+tspan = (0, 10)
+
+# Shift(t, -1)(x(t)) - x_{t-1}(t)
+# -3 - x(t) + x(t)*x_{t-1}
+f = ImplicitDiscreteFunction(sys)
+u_next = [3., 1.5]
+@test f(u_next, [2.,3.], [], t) ≈ [0., 0.]
+u_next = [0., 0.]
+@test f(u_next, [2.,3.], [], t) ≈ [3., -3.]
+
+resid = rand(2)
+f(resid, u_next, [2.,3.], [], t)
+@test resid ≈ [3., -3.]
+
+# Initialization cases.
+prob = ImplicitDiscreteProblem(sys, [x(k-1) => 3.], tspan)
+@test prob.u0 == [3., 1.]
+prob = ImplicitDiscreteProblem(sys, [], tspan)
+@test prob.u0 == [1., 1.]
+@variables x(t)
+@mtkbuild sys = ImplicitDiscreteSystem([x(k) ~ x(k)*x(k-1) - 3], t)
+@test_throws ErrorException prob = ImplicitDiscreteProblem(sys, [], tspan)
+
+# Test solvers

test/runtests.jl

@@ -80,6 +80,7 @@ end
             @safetestset "Variable Metadata Test" include("test_variable_metadata.jl")
             @safetestset "OptimizationSystem Test" include("optimizationsystem.jl")
             @safetestset "Discrete System" include("discrete_system.jl")
+            @safetestset "Implicit Discrete System" include("implicit_discrete_system.jl")
             @safetestset "SteadyStateSystem Test" include("steadystatesystems.jl")
             @safetestset "SDESystem Test" include("sdesystem.jl")
             @safetestset "DDESystem Test" include("dde.jl")