refactor: enforce non-positive shifts, update documentation

Author: AayushSabharwal

docs/src/tutorials/discrete_system.md

@@ -15,18 +15,30 @@ end
 @constants h = 1
 @variables S(t) I(t) R(t)
 k = ShiftIndex(t)
-infection = rate_to_proportion(β * c * I / (S * h + I + R), δt * h) * S
-recovery = rate_to_proportion(γ * h, δt) * I
+infection = rate_to_proportion(β * c * I(k-1) / (S(k-1) * h + I(k-1) + R(k-1)), δt * h) * S(k-1)
+recovery = rate_to_proportion(γ * h, δt) * I(k-1)
 
 # Equations
-eqs = [S(k + 1) ~ S - infection * h,
-    I(k + 1) ~ I + infection - recovery,
-    R(k + 1) ~ R + recovery]
+eqs = [S(k) ~ S(k-1) - infection * h,
+    I(k) ~ I(k-1) + infection - recovery,
+    R(k) ~ R(k-1) + recovery]
 @mtkbuild sys = DiscreteSystem(eqs, t)
 
-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]
 tspan = (0.0, 100.0)
 prob = DiscreteProblem(sys, u0, tspan, p)
 sol = solve(prob, FunctionMap())
 ```
+
+All shifts must be negative. If default values are provided, they are treated as the value
+for the variable at the previous timestep. For example, consider the following system to
+generate the Fibonacci series:
+
+```@example discrete
+@variables x(t) = 1.0
+@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.

src/discretedomain.jl

@@ -38,7 +38,9 @@ function (D::Shift)(x::Num, allow_zero = false)
     vt = value(x)
     if istree(vt)
         op = operation(vt)
-        if op isa Shift
+        if op isa Sample
+            error("Cannot shift a `Sample`. Create a variable to represent the sampled value and shift that instead")
+        elseif op isa Shift
             if D.t === nothing || isequal(D.t, op.t)
                 arg = arguments(vt)[1]
                 newsteps = D.steps + op.steps

src/structural_transformation/utils.jl

@@ -449,21 +449,4 @@ function simplify_shifts(var)
         end
     end
     return Postwalk(PassThrough(r))(var)
-    while ModelingToolkit.isoperator(var, ModelingToolkit.Shift) &&
-        ModelingToolkit.isoperator(arguments(var)[1], ModelingToolkit.Shift)
-        op1 = operation(var)
-        vv1 = arguments(var)[1]
-        op2 = operation(vv1)
-        vv2 = arguments(vv1)[1]
-        s1 = op1.steps
-        s2 = op2.steps
-        t1 = op1.t
-        t2 = op2.t
-        if t1 === nothing
-            var = Shift(t2, s1 + s2)(vv2)
-        else
-            var = Shift(t1, s1 + s2)(vv2)
-        end
-    end
-    return var
 end

src/systems/clock_inference.jl

@@ -180,9 +180,23 @@ function generate_discrete_affect(
             disc_to_cont_idxs = Int[]
         end
         for v in inputs[continuous_id]
-            vv = arguments(v)[1]
-            if vv in fullvars
-                push!(needed_disc_to_cont_obs, vv)
+            _v = arguments(v)[1]
+            if _v in fullvars
+                push!(needed_disc_to_cont_obs, _v)
+                push!(disc_to_cont_idxs, param_to_idx[v])
+            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 _v in fullvars
+                push!(needed_disc_to_cont_obs, _v)
                 push!(disc_to_cont_idxs, param_to_idx[v])
             end
         end
@@ -198,6 +212,7 @@ function generate_discrete_affect(
             throw = false,
             expression = true,
             output_type = SVector,
+            op = Shift,
             ps = reorder_parameters(osys, full_parameters(sys)))
         ni = length(input)
         ns = length(unknowns(sys))
@@ -213,7 +228,7 @@ function generate_discrete_affect(
                 get_iv(sys)
             ],
             [],
-            let_block)
+            let_block) |> toexpr
         if use_index_cache
             cont_to_disc_idxs = [parameter_index(osys, sym) for sym in input]
             disc_range = [parameter_index(osys, sym) for sym in unknowns(sys)]

src/systems/diffeqs/odesystem.jl

@@ -378,14 +378,15 @@ function build_explicit_observed_function(sys, ts;
         checkbounds = true,
         drop_expr = drop_expr,
         ps = full_parameters(sys),
+        op = Differential,
         throw = true)
     if (isscalar = !(ts isa AbstractVector))
         ts = [ts]
     end
     ts = unwrap.(Symbolics.scalarize(ts))
 
     vars = Set()
-    foreach(Base.Fix1(vars!, vars), ts)
+    foreach(v -> vars!(vars, v; op), ts)
     ivs = independent_variables(sys)
     dep_vars = scalarize(setdiff(vars, ivs))
 

src/systems/systemstructure.jl

@@ -350,6 +350,9 @@ function TearingState(sys; quick_cancel = false, check = true)
     for var in fullvars
         if ModelingToolkit.isoperator(var, ModelingToolkit.Shift)
             steps = operation(var).steps
+            if steps > 0
+                error("Only non-positive shifts allowed. Found $var with a shift of $steps")
+            end
             v = arguments(var)[1]
             lowest_shift[v] = min(get(lowest_shift, v, 0), steps)
         end
@@ -369,7 +372,7 @@ function TearingState(sys; quick_cancel = false, check = true)
             if lshift < 0
                 defs = ModelingToolkit.get_defaults(sys)
                 if (_val = get(defs, var, nothing)) !== nothing
-                    defs[Shift(tt, lshift)(v)] = _val
+                    defs[Shift(tt, -1)(v)] = _val
                 end
             end
         else
@@ -387,14 +390,13 @@ function TearingState(sys; quick_cancel = false, check = true)
             end
         end
     end
-
     # sort `fullvars` such that the mass matrix is as diagonal as possible.
     dervaridxs = collect(dervaridxs)
     sorted_fullvars = OrderedSet(fullvars[dervaridxs])
     var_to_old_var = Dict(zip(fullvars, fullvars))
     for dervaridx in dervaridxs
         dervar = fullvars[dervaridx]
-        diffvar = var_to_old_var[lower_order_var(dervar)]
+        diffvar = var_to_old_var[lower_order_var(dervar, iv)]
         if !(diffvar in sorted_fullvars)
             push!(sorted_fullvars, diffvar)
         end
@@ -416,7 +418,7 @@ function TearingState(sys; quick_cancel = false, check = true)
     var_to_diff = DiffGraph(nvars, true)
     for dervaridx in dervaridxs
         dervar = fullvars[dervaridx]
-        diffvar = lower_order_var(dervar)
+        diffvar = lower_order_var(dervar, iv)
         diffvaridx = var2idx[diffvar]
         push!(diffvars, diffvar)
         var_to_diff[diffvaridx] = dervaridx
@@ -438,7 +440,7 @@ function TearingState(sys; quick_cancel = false, check = true)
         Any[])
 end
 
-function lower_order_var(dervar)
+function lower_order_var(dervar, t)
     if isdifferential(dervar)
         diffvar = arguments(dervar)[1]
     elseif ModelingToolkit.isoperator(dervar, ModelingToolkit.Shift)
@@ -451,8 +453,7 @@ function lower_order_var(dervar)
             diffvar = vv
         end
     else
-        iv = only(arguments(dervar))
-        return Shift(iv, -1)(dervar)
+        return Shift(t, -1)(dervar)
     end
     diffvar
 end

test/clock.jl

@@ -89,40 +89,31 @@ d = Clock(t, dt)
 
 @info "Testing shift normalization"
 dt = 0.1
-@variables x(t) y(t) u(t) yd(t) ud(t) r(t) z(t)
+@variables x(t) y(t) u(t) yd(t) ud(t)
 @parameters kp
 d = Clock(t, dt)
 k = ShiftIndex(d)
 
 eqs = [yd ~ Sample(t, dt)(y)
-       ud ~ kp * (r - yd) + z(k)
-       r ~ 1.0
+       ud ~ kp * yd + ud(k - 2)
 
        # plant (time continuous part)
        u ~ Hold(ud)
        D(x) ~ -x + u
-       y ~ x
-       z(k + 2) ~ z(k) + yd
-       #=
-       z(k + 2) ~ z(k) + yd
-       =>
-       z′(k + 1) ~ z(k) + yd
-       z(k + 1)  ~ z′(k)
-       =#
-       ]
+       y ~ x]
 @named sys = ODESystem(eqs, t)
 ss = structural_simplify(sys);
 
 Tf = 1.0
 prob = ODEProblem(ss, [x => 0.0, y => 0.0], (0.0, Tf),
-    [kp => 1.0; z => 3.0; z(k + 1) => 2.0])
-@test sort(vcat(prob.p...)) == [0, 1.0, 2.0, 3.0, 4.0] # yd, kp, z(k+1), z(k), ud
+    [kp => 1.0; ud(k - 1) => 2.0; ud(k - 2) => 2.0])
+@test sort(vcat(prob.p...)) == [0, 1.0, 2.0, 2.0, 2.0] # yd, Hold(ud), kp, ud(k - 1)
 sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
 
 ss_nosplit = structural_simplify(sys; split = false)
 prob_nosplit = ODEProblem(ss_nosplit, [x => 0.0, y => 0.0], (0.0, Tf),
-    [kp => 1.0; z => 3.0; z(k + 1) => 2.0])
-@test sort(prob_nosplit.p) == [0, 1.0, 2.0, 3.0, 4.0] # yd, kp, z(k+1), z(k), ud
+    [kp => 1.0; ud(k - 1) => 2.0; ud(k - 2) => 2.0])
+@test sort(prob_nosplit.p) == [0, 1.0, 2.0, 2.0, 2.0] # yd, Hold(ud), kp, ud(k - 1)
 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.
@@ -134,30 +125,23 @@ function foo!(du, u, p, t)
     du[1] = -x + ud
 end
 function affect!(integrator, saved_values)
-    z_t, z = integrator.p[3], integrator.p[4]
     yd = integrator.u[1]
     kp = integrator.p[1]
-    r = 1.0
+    ud = integrator.p[2]
+    udd = integrator.p[3]
 
     push!(saved_values.t, integrator.t)
-    push!(saved_values.saveval, [z_t, z])
-
-    # Update the discrete state
-    z_t, z = z + yd, z_t
-    # @show z_t, z
-    integrator.p[3] = z_t
-    integrator.p[4] = z
+    push!(saved_values.saveval, [ud, udd])
 
-    ud = kp * (r - yd) + z
-    integrator.p[2] = ud
+    integrator.p[2] = kp * yd + udd
+    integrator.p[3] = ud
 
     nothing
 end
 saved_values = SavedValues(Float64, Vector{Float64})
 cb = PeriodicCallback(Base.Fix2(affect!, saved_values), 0.1)
-#                                           kp   ud   z_t  z
-prob = ODEProblem(foo!, [0.0], (0.0, Tf), [1.0, 4.0, 2.0, 3.0], callback = cb)
-#                               ud initializes to kp * (r - yd) + z = 1 * (1 - 0) + 3 = 4
+#                                           kp   ud
+prob = ODEProblem(foo!, [0.0], (0.0, Tf), [1.0, 2.0, 2.0], callback = cb)
 sol2 = solve(prob, Tsit5())
 @test sol.u == sol2.u
 @test sol_nosplit.u == sol2.u
@@ -217,7 +201,7 @@ end
 function filt(; name)
     @variables x(t)=0 u(t)=0 y(t)=0
     a = 1 / exp(dt)
-    eqs = [x(k + 1) ~ a * x + (1 - a) * u(k)
+    eqs = [x ~ a * x(k - 1) + (1 - a) * u(k - 1)
            y ~ x]
     ODESystem(eqs, t, name = name)
 end
@@ -487,9 +471,11 @@ k = ShiftIndex(c)
     @variables begin
         count(t) = 0
         u(t) = 0
+        ud(t) = 0
     end
     @equations begin
-        count(k + 1) ~ Sample(c)(u)
+        ud ~ Sample(c)(u)
+        count ~ ud(k - 1)
     end
 end
 

test/discrete_system.jl

@@ -7,6 +7,11 @@ using ModelingToolkit, Test
 using ModelingToolkit: t_nounits as t
 using ModelingToolkit: get_metadata, MTKParameters
 
+# Make sure positive shifts error
+@variables x(t)
+k = ShiftIndex(t)
+@test_throws ErrorException @mtkbuild sys = DiscreteSystem([x(k + 1) ~ x + x(k - 1)], t)
+
 @inline function rate_to_proportion(r, t)
     1 - exp(-r * t)
 end;
@@ -15,7 +20,6 @@ end;
 @parameters c nsteps δt β γ
 @constants h = 1
 @variables S(t) I(t) R(t)
-k = ShiftIndex(t)
 infection = rate_to_proportion(
     β * c * I(k - 1) / (S(k - 1) * h + I(k - 1) + R(k - 1)), δt * h) * S(k - 1)
 recovery = rate_to_proportion(γ * h, δt) * I(k - 1)

test/parameter_dependencies.jl

@@ -61,12 +61,12 @@ end
            u ~ Hold(ud)
            D(x) ~ -x + u
            y ~ x
-           z(k + 2) ~ z(k) + yd]
+           z(k) ~ z(k - 2) + yd(k - 2)]
     @mtkbuild sys = ODESystem(eqs, t; parameter_dependencies = [kq => 2kp])
 
     Tf = 1.0
     prob = ODEProblem(sys, [x => 0.0, y => 0.0], (0.0, Tf),
-        [kp => 1.0; z => 3.0; z(k + 1) => 2.0])
+        [kp => 1.0; z(k - 1) => 3.0; yd(k - 1) => 0.0])
     @test_nowarn solve(prob, Tsit5(); kwargshandle = KeywordArgSilent)
 
     @mtkbuild sys = ODESystem(eqs, t; parameter_dependencies = [kq => 2kp],
@@ -77,7 +77,7 @@ end
     @test prob.ps[kq] == 2.0
     @test_nowarn solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
     prob = ODEProblem(sys, [x => 0.0, y => 0.0], (0.0, Tf),
-        [kp => 1.0; z => 3.0; z(k + 1) => 2.0])
+        [kp => 1.0; z(k - 1) => 3.0; yd(k - 1) => 0.0])
     integ = init(prob, Tsit5(), kwargshandle = KeywordArgSilent)
     @test integ.ps[kp] == 1.0
     @test integ.ps[kq] == 2.0