fix: implement proper initialization convention for discrete variables

Author: AayushSabharwal

docs/src/tutorials/discrete_system.md

@@ -42,5 +42,10 @@ the Fibonacci series:
 @mtkbuild sys = DiscreteSystem([x ~ x(k - 1) + x(k - 2)], t)
 ```
 
-Note that the default value is treated as the initial value of `x(k-1)`. The value for
-`x(k-2)` must be provided during problem construction.
+The "default value" here should be interpreted as the value of `x` at all past timesteps.
+For example, here `x(k-1)` and `x(k-2)` will be `1.0`, and the inital value of `x(k)` will
+thus be `2.0`. During problem construction, the _past_ value of a variable should be
+provided. For example, providing `[x => 1.0]` while constructing this problem will error.
+Provide `[x(k-1) => 1.0]` instead. Note that values provided during problem construction
+_do not_ apply to the entire history. Hence, if `[x(k-1) => 2.0]` is provided, the value of
+`x(k-2)` will still be `1.0`.

src/systems/diffeqs/abstractodesystem.jl

@@ -864,6 +864,36 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     # since they will be checked in the initialization problem's construction
     # TODO: make check for if a DAE cheaper than calculating the mass matrix a second time!
     ci = infer_clocks!(ClockInference(TearingState(sys)))
+
+    if eltype(parammap) <: Pair
+        parammap = Dict(unwrap(k) => v for (k, v) in todict(parammap))
+    elseif parammap isa AbstractArray
+        parammap = Dict(unwrap.(parameters(sys)) .=> parammap)
+    end
+    clockedparammap = Dict()
+    defs = ModelingToolkit.get_defaults(sys)
+    for v in ps
+        is_discrete_domain(v) || continue
+        op = operation(v)
+        if !isa(op, Symbolics.Operator) && parammap != SciMLBase.NullParameters() &&
+           haskey(parammap, v)
+            error("Initial conditions for discrete variables must be for the past state of the unknown. Instead of providing the condition for $v, provide the condition for $(Shift(iv, -1)(v)).")
+        end
+        shiftedv = StructuralTransformations.simplify_shifts(Shift(iv, -1)(v))
+        if parammap != SciMLBase.NullParameters() &&
+           (val = get(parammap, shiftedv, nothing)) !== nothing
+            clockedparammap[v] = val
+        elseif op isa Shift
+            root = arguments(v)[1]
+            haskey(defs, root) || error("Initial condition for $v not provided.")
+            clockedparammap[v] = defs[root]
+        end
+    end
+    parammap = if parammap == SciMLBase.NullParameters()
+        clockedparammap
+    else
+        merge(parammap, clockedparammap)
+    end
     # TODO: make it work with clocks
     # ModelingToolkit.get_tearing_state(sys) !== nothing => Requires structural_simplify first
     if sys isa ODESystem && (implicit_dae || !isempty(missingvars)) &&

src/systems/discrete_system/discrete_system.jl

@@ -171,8 +171,8 @@ function DiscreteSystem(eqs, iv; kwargs...)
     ps = OrderedSet()
     iv = value(iv)
     for eq in eqs
-        collect_vars!(allunknowns, ps, eq.lhs, iv)
-        collect_vars!(allunknowns, ps, eq.rhs, iv)
+        collect_vars!(allunknowns, ps, eq.lhs, iv; op = Shift)
+        collect_vars!(allunknowns, ps, eq.rhs, iv; op = Shift)
         if istree(eq.lhs) && operation(eq.lhs) isa Shift
             isequal(iv, operation(eq.lhs).t) ||
                 throw(ArgumentError("A DiscreteSystem can only have one independent variable."))
@@ -205,21 +205,38 @@ function generate_function(
 end
 
 function process_DiscreteProblem(constructor, sys::DiscreteSystem, u0map, parammap;
-        version = nothing,
         linenumbers = true, parallel = SerialForm(),
         eval_expression = true,
         use_union = false,
         tofloat = !use_union,
         kwargs...)
+    iv = get_iv(sys)
     eqs = equations(sys)
     dvs = unknowns(sys)
     ps = parameters(sys)
 
+    trueu0map = Dict()
+    for (k, v) in u0map
+        k = unwrap(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
+        trueu0map[Shift(iv, op.steps + 1)(arguments(k)[1])] = v
+    end
+    defs = ModelingToolkit.get_defaults(sys)
+    for var in dvs
+        if (op = operation(var)) isa Shift && !haskey(trueu0map, var)
+            root = arguments(var)[1]
+            haskey(defs, root) || error("Initial condition for $var not provided.")
+            trueu0map[var] = defs[root]
+        end
+    end
+    @show trueu0map u0map
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        u0, defs = get_u0(sys, u0map, parammap)
-        p = MTKParameters(sys, parammap)
+        u0, defs = get_u0(sys, trueu0map, parammap)
+        p = MTKParameters(sys, parammap, trueu0map)
     else
-        u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union)
+        u0, p, defs = get_u0_p(sys, trueu0map, parammap; tofloat, use_union)
     end
 
     check_eqs_u0(eqs, dvs, u0; kwargs...)

src/utils.jl

@@ -243,11 +243,14 @@ function collect_var_to_name!(vars, xs)
         x = unwrap(x)
         if hasmetadata(x, Symbolics.GetindexParent)
             xarr = getmetadata(x, Symbolics.GetindexParent)
+            hasname(xarr) || continue
             vars[Symbolics.getname(xarr)] = xarr
         else
             if istree(x) && operation(x) === getindex
                 x = arguments(x)[1]
             end
+            x = unwrap(x)
+            hasname(x) || continue
             vars[Symbolics.getname(unwrap(x))] = x
         end
     end
@@ -434,11 +437,11 @@ function find_derivatives!(vars, expr, f)
     return vars
 end
 
-function collect_vars!(unknowns, parameters, expr, iv)
+function collect_vars!(unknowns, parameters, expr, iv; op = Differential)
     if issym(expr)
         collect_var!(unknowns, parameters, expr, iv)
     else
-        for var in vars(expr)
+        for var in vars(expr; op)
             if istree(var) && operation(var) isa Differential
                 var, _ = var_from_nested_derivative(var)
             end

test/clock.jl

@@ -109,21 +109,21 @@ ss = structural_simplify(sys);
 
 Tf = 1.0
 prob = ODEProblem(ss, [x => 0.1], (0.0, Tf),
-    [kp => 1.0; ud => 2.1; ud(k - 1) => 2.0])
+    [kp => 1.0; ud(k - 1) => 2.1; ud(k - 2) => 2.0])
 # create integrator so callback is evaluated at t=0 and we can test correct param values
 int = init(prob, Tsit5(); kwargshandle = KeywordArgSilent)
 @test sort(vcat(int.p...)) == [0.1, 1.0, 2.1, 2.1, 2.1] # yd, kp, ud(k-1), ud, Hold(ud)
 prob = ODEProblem(ss, [x => 0.1], (0.0, Tf),
-    [kp => 1.0; ud => 2.1; ud(k - 1) => 2.0]) # recreate problem to empty saved values
+    [kp => 1.0; ud(k - 1) => 2.1; ud(k - 2) => 2.0]) # recreate problem to empty saved values
 sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
 
 ss_nosplit = structural_simplify(sys; split = false)
 prob_nosplit = ODEProblem(ss_nosplit, [x => 0.1], (0.0, Tf),
-    [kp => 1.0; ud => 2.1; ud(k - 1) => 2.0])
+    [kp => 1.0; ud(k - 1) => 2.1; ud(k - 2) => 2.0])
 int = init(prob_nosplit, Tsit5(); kwargshandle = KeywordArgSilent)
 @test sort(int.p) == [0.1, 1.0, 2.1, 2.1, 2.1] # yd, kp, ud(k-1), ud, Hold(ud)
 prob_nosplit = ODEProblem(ss_nosplit, [x => 0.1], (0.0, Tf),
-    [kp => 1.0; ud => 2.1; ud(k - 1) => 2.0]) # recreate problem to empty saved values
+    [kp => 1.0; ud(k - 1) => 2.1; ud(k - 2) => 2.0]) # recreate problem to empty saved values
 sol_nosplit = solve(prob_nosplit, Tsit5(), kwargshandle = KeywordArgSilent)
 # For all inputs in parameters, just initialize them to 0.0, and then set them
 # in the callback.
@@ -516,3 +516,18 @@ sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
 @test sol.prob.kwargs[:disc_saved_values][1].t == sol.t[1:2:end] # Test that the discrete-tiem system executed at every step of the continuous solver. The solver saves each time step twice, one state value before discrete affect and one after.
 @test_nowarn ModelingToolkit.build_explicit_observed_function(
     model, model.counter.ud)(sol.u[1], prob.p..., sol.t[1])
+
+@variables x(t)=1.0 y(t)=1.0
+eqs = [D(y) ~ Hold(x)
+       x ~ x(k - 1) + x(k - 2)]
+@mtkbuild sys = ODESystem(eqs, t)
+prob = ODEProblem(sys, [], (0.0, 10.0))
+int = init(prob, Tsit5(); kwargshandle = KeywordArgSilent)
+@test int.ps[x] == 2.0
+@test int.ps[x(k - 1)] == 1.0
+
+@test_throws ErrorException ODEProblem(sys, [], (0.0, 10.0), [x => 2.0])
+prob = ODEProblem(sys, [], (0.0, 10.0), [x(k - 1) => 2.0])
+int = init(prob, Tsit5(); kwargshandle = KeywordArgSilent)
+@test int.ps[x] == 3.0
+@test int.ps[x(k - 1)] == 2.0

test/discrete_system.jl

@@ -51,7 +51,7 @@ for df in [
 end
 
 # Problem
-u0 = [S => 990.0, I => 10.0, R => 0.0]
+u0 = [S(k - 1) => 990.0, I(k - 1) => 10.0, R(k - 1) => 0.0]
 p = [β => 0.05, c => 10.0, γ => 0.25, δt => 0.1, nsteps => 400]
 tspan = (0.0, ModelingToolkit.value(substitute(nsteps, p))) # value function (from Symbolics) is used to convert a Num to Float64
 prob_map = DiscreteProblem(syss, u0, tspan, p)
@@ -215,7 +215,22 @@ eqs = [u ~ 1
        x ~ x(k - 1) + u
        y ~ x + u]
 @mtkbuild de = DiscreteSystem(eqs, t)
-prob = DiscreteProblem(de, [x => 0.0], (0, 10))
+prob = DiscreteProblem(de, [x(k - 1) => 0.0], (0, 10))
 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