Infer SampleTime

Author: YingboMa

src/ModelingToolkit.jl

@@ -269,7 +269,7 @@ export debug_system
 #export Continuous, 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
+export Sample, Hold, Shift, ShiftIndex, sampletime, SampleTime
 export Clock #, InferredDiscrete,
 
 end # module

src/clock.jl

@@ -117,8 +117,7 @@ end
 Clock(dt::Real) = Clock(nothing, dt)
 Clock() = Clock(nothing, nothing)
 
-sampletime() = InferredSampleTime()
-sampletime(c) = something(isdefined(c, :dt) ? c.dt : nothing, InferredSampleTime())
+sampletime(c) = isdefined(c, :dt) ? c.dt : nothing
 Base.hash(c::Clock, seed::UInt) = hash(c.dt, seed ⊻ 0x953d7a9a18874b90)
 function Base.:(==)(c1::Clock, c2::Clock)
     ((c1.t === nothing || c2.t === nothing) || isequal(c1.t, c2.t)) && c1.dt == c2.dt

src/discretedomain.jl

@@ -1,12 +1,9 @@
 using Symbolics: Operator, Num, Term, value, recursive_hasoperator
 
-struct InferredSampleTime <: Operator end
-function SymbolicUtils.promote_symtype(::Type{InferredSampleTime}, t...)
-    Real
-end
-function InferredSampleTime()
-    # Term{Real}(InferredSampleTime, Any[])
-    SymbolicUtils.term(InferredSampleTime, type = Real)
+struct SampleTime <: Operator end
+SymbolicUtils.promote_symtype(::Type{SampleTime}, t...) = Real
+function SampleTime()
+    SymbolicUtils.term(SampleTime, type = Real)
 end
 
 # Shift

src/systems/clock_inference.jl

@@ -21,6 +21,52 @@ function ClockInference(ts::TransformationState)
     ClockInference(ts, eq_domain, var_domain, inferred)
 end
 
+struct NotInferedTimeDomain end
+function error_sample_time(eq)
+    error("$eq\ncontains `SampleTime` but it is not an infered discrete equation.")
+end
+function substitute_sample_time(ci::ClockInference)
+    @unpack ts, eq_domain = ci
+    eqs = copy(equations(ts))
+    @assert length(eqs) == length(eq_domain)
+    for i in eachindex(eqs)
+        eq = eqs[i]
+        domain = eq_domain[i]
+        dt = sampletime(domain)
+        neweq = substitute_sample_time(eq, dt)
+        if neweq isa NotInferedTimeDomain
+            error_sample_time(eq)
+        end
+        eqs[i] = neweq
+    end
+    @set! ts.sys.eqs = eqs
+    @set! ci.ts = ts
+end
+
+function substitute_sample_time(eq::Equation, dt)
+    substitute_sample_time(eq.lhs, dt) ~ substitute_sample_time(eq.rhs, dt)
+end
+
+function substitute_sample_time(ex, dt)
+    istree(ex) || return ex
+    op = operation(ex)
+    args = arguments(ex)
+    if op == SampleTime
+        dt === nothing && return NotInferedTimeDomain()
+        return dt
+    else
+        new_args = similar(args)
+        for (i, arg) in enumerate(args)
+            ex_arg = substitute_sample_time(arg, dt)
+            if ex_arg isa NotInferedTimeDomain
+                return ex_arg
+            end
+            new_args[i] = ex_arg
+        end
+        similarterm(ex, op, new_args; metadata = metadata(ex))
+    end
+end
+
 function infer_clocks!(ci::ClockInference)
     @unpack ts, eq_domain, var_domain, inferred = ci
     @unpack var_to_diff, graph = ts.structure
@@ -66,6 +112,7 @@ function infer_clocks!(ci::ClockInference)
         end
     end
 
+    ci = substitute_sample_time(ci)
     return ci
 end
 

src/systems/systemstructure.jl

@@ -627,7 +627,7 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
         kwargs...)
     if state.sys isa ODESystem
         ci = ModelingToolkit.ClockInference(state)
-        ModelingToolkit.infer_clocks!(ci)
+        ci = ModelingToolkit.infer_clocks!(ci)
         time_domains = merge(Dict(state.fullvars .=> ci.var_domain),
             Dict(default_toterm.(state.fullvars) .=> ci.var_domain))
         tss, inputs, continuous_id, id_to_clock = ModelingToolkit.split_system(ci)
@@ -652,18 +652,6 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
                 append!(appended_parameters, inputs[i], unknowns(ss))
                 discrete_subsystems[i] = ss
             end
-            for i in eachindex(discrete_subsystems)
-                discsys = discrete_subsystems[i]
-                eqs = collect(discsys.eqs)
-                for eqi in eachindex(eqs)
-                    clock = id_to_clock[i]
-                    clock isa AbstractDiscrete || continue
-                    Ts = sampletime(clock)
-                    eqs[eqi] = substitute(eqs[eqi], InferredSampleTime() => Ts)
-                end
-                @set discsys.eqs = eqs
-                discrete_subsystems[i] = discsys
-            end
             @set! sys.discrete_subsystems = discrete_subsystems, inputs, continuous_id,
             id_to_clock
             @set! sys.ps = appended_parameters

test/clock.jl

@@ -347,7 +347,7 @@ k = ShiftIndex()
         y(t)
     end
     @equations begin
-        x(k) ~ x(k - 1) + ki * u(k) * sampletime() / dt
+        x(k) ~ x(k - 1) + ki * u(k) * SampleTime() / dt
         output.u(k) ~ y(k)
         input.u(k) ~ u(k)
         y(k) ~ x(k - 1) + kp * u(k)
@@ -374,7 +374,7 @@ end
 @mtkmodel ClosedLoop begin
     @components begin
         plant = FirstOrder(k = 0.3, T = 1)
-        sampler = Blocks.Sampler(; clock = d)
+        sampler = Sampler()
         holder = ZeroOrderHold()
         controller = DiscretePI(kp = 2, ki = 2)
         feedback = Feedback()
@@ -441,7 +441,7 @@ prob = ODEProblem(ssys,
     [model.plant.x => 0.0; model.controller.kp => 2.0; model.controller.ki => 2.0],
     (0.0, Tf))
 int = init(prob, Tsit5(); kwargshandle = KeywordArgSilent)
-@test int.ps[Hold(ssys.holder.input.u)] == 2 # constant output * kp issue https://github.com/SciML/ModelingToolkit.jl/issues/2356
+@test_broken int.ps[Hold(ssys.holder.input.u)] == 2 # constant output * kp issue https://github.com/SciML/ModelingToolkit.jl/issues/2356
 @test int.ps[ssys.controller.x] == 1 # c2d
 @test int.ps[Sample(d)(ssys.sampler.input.u)] == 0 # disc state
 sol = solve(prob,