feat: shift all discrete systems by 1 to fix correctness issues

AayushSabharwal · AayushSabharwal · commit 71ebfdd00fe4 · 2024-03-21T16:48:53.000+05:30
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
@@ -433,6 +433,7 @@ end
 
 function simplify_shifts(var)
     ModelingToolkit.hasshift(var) || return var
+    var isa Equation && return simplify_shifts(var.lhs) ~ simplify_shifts(var.rhs)
     if isdoubleshift(var)
         op1 = operation(var)
         vv1 = arguments(var)[1]
diff --git a/src/systems/alias_elimination.jl b/src/systems/alias_elimination.jl
@@ -453,7 +453,7 @@ function observed2graph(eqs, unknowns)
         lhs_j === nothing &&
             throw(ArgumentError("The lhs $(eq.lhs) of $eq, doesn't appear in unknowns."))
         assigns[i] = lhs_j
-        vs = vars(eq.rhs)
+        vs = vars(eq.rhs; op = Symbolics.Operator)
         for v in vs
             j = get(v2j, v, nothing)
             j !== nothing && add_edge!(graph, i, j)
diff --git a/src/systems/clock_inference.jl b/src/systems/clock_inference.jl
@@ -133,7 +133,10 @@ function split_system(ci::ClockInference{S}) where {S}
     tss = similar(cid_to_eq, S)
     for (id, ieqs) in enumerate(cid_to_eq)
         ts_i = system_subset(ts, ieqs)
-        @set! ts_i.structure.only_discrete = id != continuous_id
+        if id != continuous_id
+            ts_i = shift_discrete_system(ts_i)
+            @set! ts_i.structure.only_discrete = true
+        end
         tss[id] = ts_i
     end
     return tss, inputs, continuous_id, id_to_clock
@@ -148,7 +151,7 @@ function generate_discrete_affect(
     end
     use_index_cache = has_index_cache(osys) && get_index_cache(osys) !== nothing
     out = Sym{Any}(:out)
-    appended_parameters = parameters(syss[continuous_id])
+    appended_parameters = full_parameters(syss[continuous_id])
     offset = length(appended_parameters)
     param_to_idx = if use_index_cache
         Dict{Any, ParameterIndex}(p => parameter_index(osys, p)
@@ -157,6 +160,7 @@ function generate_discrete_affect(
         Dict{Any, Int}(reverse(en) for en in enumerate(appended_parameters))
     end
     affect_funs = []
+    init_funs = []
     svs = []
     clocks = TimeDomain[]
     for (i, (sys, input)) in enumerate(zip(syss, inputs))
@@ -183,47 +187,38 @@ function generate_discrete_affect(
             if _v in fullvars
                 push!(needed_disc_to_cont_obs, _v)
                 push!(disc_to_cont_idxs, param_to_idx[v])
+                continue
             end
 
-            # In the above case, `_v` was in `observed(sys)`
-            # It may also be in `unknowns(sys)`, in which case it
-            # will be shifted back by one step
-            if istree(v) && (op = operation(v)) isa Shift
-                _v = arguments(_v)[1]
-                _v = Shift(op.t, op.steps - 1)(_v)
-            else
-                _v = Shift(get_iv(sys), -1)(_v)
-            end
+            # If the held quantity is calculated through observed
+            # it will be shifted forward by 1
+            _v = Shift(get_iv(sys), 1)(_v)
             if _v in fullvars
                 push!(needed_disc_to_cont_obs, _v)
                 push!(disc_to_cont_idxs, param_to_idx[v])
+                continue
             end
         end
-        append!(appended_parameters, input, unknowns(sys))
+        append!(appended_parameters, input)
         cont_to_disc_obs = build_explicit_observed_function(
             use_index_cache ? osys : syss[continuous_id],
             needed_cont_to_disc_obs,
             throw = false,
             expression = true,
             output_type = SVector)
-        @set! sys.ps = appended_parameters
         disc_to_cont_obs = build_explicit_observed_function(sys, needed_disc_to_cont_obs,
             throw = false,
             expression = true,
             output_type = SVector,
             op = Shift,
-            ps = reorder_parameters(osys, full_parameters(sys)))
+            ps = reorder_parameters(osys, appended_parameters))
         ni = length(input)
         ns = length(unknowns(sys))
         disc = Func(
             [
                 out,
                 DestructuredArgs(unknowns(osys)),
-                if use_index_cache
-                    DestructuredArgs.(reorder_parameters(osys, full_parameters(osys)))
-                else
-                    (DestructuredArgs(appended_parameters),)
-                end...,
+                DestructuredArgs.(reorder_parameters(osys, full_parameters(osys)))...,
                 get_iv(sys)
             ],
             [],
@@ -248,6 +243,36 @@ function generate_discrete_affect(
             end
         end
         empty_disc = isempty(disc_range)
+        disc_init = if use_index_cache
+            :(function (u, p, t)
+                c2d_obs = $cont_to_disc_obs
+                d2c_obs = $disc_to_cont_obs
+                result = c2d_obs(u, p..., t)
+                for (val, i) in zip(result, $cont_to_disc_idxs)
+                    $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+                end
+
+                disc_state = Tuple($(parameter_values)(p, i) for i in $disc_range)
+                result = d2c_obs(disc_state, p..., t)
+                for (val, i) in zip(result, $disc_to_cont_idxs)
+                    # prevent multiple updates to dependents
+                    _set_parameter_unchecked!(p, val, i; update_dependent = false)
+                end
+                discretes, repack, _ = $(SciMLStructures.canonicalize)(
+                    $(SciMLStructures.Discrete()), p)
+                repack(discretes) # to force recalculation of dependents
+            end)
+        else
+            :(function (u, p, t)
+                c2d_obs = $cont_to_disc_obs
+                d2c_obs = $disc_to_cont_obs
+                c2d_view = view(p, $cont_to_disc_idxs)
+                d2c_view = view(p, $disc_to_cont_idxs)
+                disc_unknowns = view(p, $disc_range)
+                copyto!(c2d_view, c2d_obs(u, p, t))
+                copyto!(d2c_view, d2c_obs(disc_unknowns, p, t))
+            end)
+        end
 
         # @show disc_to_cont_idxs
         # @show cont_to_disc_idxs
@@ -270,9 +295,6 @@ function generate_discrete_affect(
             # TODO: find a way to do this without allocating
             disc = $disc
 
-            push!(saved_values.t, t)
-            push!(saved_values.saveval, $save_vec)
-
             # Write continuous into to discrete: handles `Sample`
             # Write discrete into to continuous
             # Update discrete unknowns
@@ -322,6 +344,10 @@ function generate_discrete_affect(
                 :(copyto!(d2c_view, d2c_obs(disc_unknowns, p, t)))
             end
             )
+
+            push!(saved_values.t, t)
+            push!(saved_values.saveval, $save_vec)
+
             # @show "after d2c", p
             $(
                 if use_index_cache
@@ -335,15 +361,20 @@ function generate_discrete_affect(
         end)
         sv = SavedValues(Float64, Vector{Float64})
         push!(affect_funs, affect!)
+        push!(init_funs, disc_init)
         push!(svs, sv)
     end
     if eval_expression
         affects = map(affect_funs) do a
             drop_expr(@RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a))))
         end
+        inits = map(init_funs) do a
+            drop_expr(@RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a))))
+        end
     else
         affects = map(a -> toexpr(LiteralExpr(a)), affect_funs)
+        inits = map(a -> toexpr(LiteralExpr(a)), init_funs)
     end
     defaults = Dict{Any, Any}(v => 0.0 for v in Iterators.flatten(inputs))
-    return affects, clocks, svs, appended_parameters, defaults
+    return affects, inits, clocks, svs, appended_parameters, defaults
 end
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -1039,12 +1039,13 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
         t = tspan !== nothing ? tspan[1] : tspan,
         check_length, warn_initialize_determined, kwargs...)
     cbs = process_events(sys; callback, kwargs...)
-    affects = []
+    inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
+        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
-                PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
+                PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt;
+                    final_affect = true)
             elseif clock isa SolverStepClock
                 affect = DiscreteSaveAffect(affect, sv)
                 DiscreteCallback(Returns(true), affect,
@@ -1062,12 +1063,6 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
         else
             cbs = CallbackSet(cbs, discrete_cbs...)
         end
-        # initialize by running affects
-        dummy_saveval = (; t = [], saveval = [])
-        for affect! in affects
-            affect!(
-                (; u = u0, p = p, t = tspan !== nothing ? tspan[1] : tspan), dummy_saveval)
-        end
     else
         svs = nothing
     end
@@ -1081,7 +1076,14 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
     if svs !== nothing
         kwargs1 = merge(kwargs1, (disc_saved_values = svs,))
     end
-    ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...)
+
+    prob = ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...)
+    if !isempty(inits)
+        for init in inits
+            init(prob.u0, prob.p, tspan[1])
+        end
+    end
+    prob
 end
 get_callback(prob::ODEProblem) = prob.kwargs[:callback]
 
@@ -1150,12 +1152,12 @@ function DiffEqBase.DDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     h(p::MTKParameters, t) = h_oop(p..., t)
     u0 = h(p, tspan[1])
     cbs = process_events(sys; callback, kwargs...)
-    inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
+        affects, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
-                PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
+                PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt;
+                    final_affect = true, initial_affect = true)
             else
                 error("$clock is not a supported clock type.")
             end
@@ -1181,13 +1183,7 @@ function DiffEqBase.DDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     if svs !== nothing
         kwargs1 = merge(kwargs1, (disc_saved_values = svs,))
     end
-    prob = DDEProblem{iip}(f, u0, h, tspan, p; kwargs1..., kwargs...)
-    if !isempty(inits)
-        for init in inits
-            init(prob.p, tspan[1])
-        end
-    end
-    prob
+    DDEProblem{iip}(f, u0, h, tspan, p; kwargs1..., kwargs...)
 end
 
 function DiffEqBase.SDDEProblem(sys::AbstractODESystem, args...; kwargs...)
@@ -1212,12 +1208,12 @@ function DiffEqBase.SDDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     h(p, t) = h_oop(p, t)
     u0 = h(p, tspan[1])
     cbs = process_events(sys; callback, kwargs...)
-    inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
+        affects, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
-                PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
+                PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt;
+                    final_affect = true, initial_affect = true)
             else
                 error("$clock is not a supported clock type.")
             end
@@ -1254,15 +1250,9 @@ function DiffEqBase.SDDEProblem{iip}(sys::AbstractODESystem, u0map = [],
     else
         noise_rate_prototype = zeros(eltype(u0), size(noiseeqs))
     end
-    prob = SDDEProblem{iip}(f, f.g, u0, h, tspan, p;
+    SDDEProblem{iip}(f, f.g, u0, h, tspan, p;
         noise_rate_prototype =
         noise_rate_prototype, kwargs1..., kwargs...)
-    if !isempty(inits)
-        for init in inits
-            init(prob.p, tspan[1])
-        end
-    end
-    prob
 end
 
 """
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
@@ -453,13 +453,16 @@ function build_explicit_observed_function(sys, ts;
     if inputs !== nothing
         ps = setdiff(ps, inputs) # Inputs have been converted to parameters by io_preprocessing, remove those from the parameter list
     end
-    if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        ps = DestructuredArgs.(reorder_parameters(get_index_cache(sys), ps))
-    elseif ps isa Tuple
+    if ps isa Tuple
         ps = DestructuredArgs.(ps, inbounds = !checkbounds)
+    elseif has_index_cache(sys) && get_index_cache(sys) !== nothing
+        ps = DestructuredArgs.(reorder_parameters(get_index_cache(sys), ps))
     else
         ps = (DestructuredArgs(ps, inbounds = !checkbounds),)
     end
+    if isempty(ps)
+        ps = (DestructuredArgs([]),)
+    end
     dvs = DestructuredArgs(unknowns(sys), inbounds = !checkbounds)
     if inputs === nothing
         args = [dvs, ps..., ivs...]
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
@@ -139,8 +139,8 @@ function DiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
     iv′ = value(iv)
     dvs′ = value.(dvs)
     ps′ = value.(ps)
-    if !all(hasshift, eqs)
-        error("All equations in a `DiscreteSystem` must be difference equations")
+    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(
diff --git a/src/systems/systemstructure.jl b/src/systems/systemstructure.jl
@@ -369,12 +369,6 @@ function TearingState(sys; quick_cancel = false, check = true)
             steps = 0
             tt = iv
             v = var
-            if lshift < 0
-                defs = ModelingToolkit.get_defaults(sys)
-                if (_val = get(defs, var, nothing)) !== nothing
-                    defs[Shift(tt, -1)(v)] = _val
-                end
-            end
         else
             continue
         end
@@ -434,10 +428,14 @@ function TearingState(sys; quick_cancel = false, check = true)
 
     eq_to_diff = DiffGraph(nsrcs(graph))
 
-    return TearingState(sys, fullvars,
+    ts = TearingState(sys, fullvars,
         SystemStructure(complete(var_to_diff), complete(eq_to_diff),
             complete(graph), nothing, var_types, sys isa DiscreteSystem),
         Any[])
+    if sys isa DiscreteSystem
+        ts = shift_discrete_system(ts)
+    end
+    return ts
 end
 
 function lower_order_var(dervar, t)
@@ -458,6 +456,30 @@ function lower_order_var(dervar, t)
     diffvar
 end
 
+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
+    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}))
+    for i in eachindex(fullvars)
+        fullvars[i] = StructuralTransformations.simplify_shifts(fast_substitute(
+            fullvars[i], discmap; operator = Union{Sample, Hold}))
+    end
+    for i in eachindex(eqs)
+        eqs[i] = StructuralTransformations.simplify_shifts(fast_substitute(
+            eqs[i], discmap; operator = Union{Sample, Hold}))
+    end
+    @set! ts.sys.eqs = eqs
+    @set! ts.fullvars = fullvars
+    return ts
+end
+
 using .BipartiteGraphs: Label, BipartiteAdjacencyList
 struct SystemStructurePrintMatrix <:
        AbstractMatrix{Union{Label, BipartiteAdjacencyList}}
diff --git a/src/variables.jl b/src/variables.jl
@@ -195,7 +195,8 @@ function _varmap_to_vars(varmap::Dict, varlist; defaults = Dict(), check = false
     values = Dict()
     for var in varlist
         var = unwrap(var)
-        val = unwrap(fixpoint_sub(fixpoint_sub(var, varmap), defaults))
+        val = unwrap(fixpoint_sub(fixpoint_sub(var, varmap; operator = Symbolics.Operator),
+            defaults; operator = Symbolics.Operator))
         if symbolic_type(val) === NotSymbolic()
             values[var] = val
         end
diff --git a/test/clock.jl b/test/clock.jl
diff --git a/test/discrete_system.jl b/test/discrete_system.jl
