Merge pull request #2138 from SciML/myb/clean_clock

Make clock inference more generic

Author: YingboMa

Project.toml

@@ -46,12 +46,6 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
-[weakdeps]
-DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
-
-[extensions]
-MTKDeepDiffsExt = "DeepDiffs"
-
 [compat]
 AbstractTrees = "0.3, 0.4"
 ArrayInterface = "6, 7"
@@ -90,6 +84,9 @@ UnPack = "0.1, 1.0"
 Unitful = "1.1"
 julia = "1.6"
 
+[extensions]
+MTKDeepDiffsExt = "DeepDiffs"
+
 [extras]
 AmplNLWriter = "7c4d4715-977e-5154-bfe0-e096adeac482"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
@@ -116,3 +113,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsMTK", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]
+
+[weakdeps]
+DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"

src/clock.jl

@@ -7,9 +7,6 @@ struct Inferred <: TimeDomain end
 struct InferredDiscrete <: AbstractDiscrete end
 struct Continuous <: TimeDomain end
 
-const UnknownDomain = Union{Nothing, Inferred, InferredDiscrete}
-const InferredDomain = Union{Inferred, InferredDiscrete}
-
 Symbolics.option_to_metadata_type(::Val{:timedomain}) = TimeDomain
 
 """
@@ -101,18 +98,27 @@ end
 abstract type AbstractClock <: AbstractDiscrete end
 
 """
-Clock <: AbstractClock
-Clock(t; dt)
+    Clock <: AbstractClock
+    Clock([t]; dt)
+
 The default periodic clock with independent variables `t` and tick interval `dt`.
 If `dt` is left unspecified, it will be inferred (if possible).
 """
 struct Clock <: AbstractClock
     "Independent variable"
-    t::Any
+    t::Union{Nothing, Symbolic}
     "Period"
-    dt::Any
-    Clock(t, dt = nothing) = new(value(t), dt)
+    dt::Union{Nothing, Float64}
+    Clock(t::Union{Num, Symbolic}, dt = nothing) = new(value(t), dt)
+    Clock(t::Nothing, dt = nothing) = new(t, dt)
 end
+Clock(dt::Real) = Clock(nothing, dt)
+Clock() = Clock(nothing, nothing)
 
 sampletime(c) = isdefined(c, :dt) ? c.dt : nothing
-Base.:(==)(c1::Clock, c2::Clock) = isequal(c1.t, c2.t) && c1.dt == c2.dt
+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
+end
+
+is_concrete_time_domain(x) = x isa Union{AbstractClock, Continuous}

src/discretedomain.jl

@@ -23,10 +23,11 @@ julia> Δ = Shift(t)
 """
 struct Shift <: Operator
     """Fixed Shift"""
-    t::Any
+    t::Union{Nothing, Symbolic}
     steps::Int
     Shift(t, steps = 1) = new(value(t), steps)
 end
+Shift(steps::Int) = new(nothing, steps)
 normalize_to_differential(s::Shift) = Differential(s.t)^s.steps
 function (D::Shift)(x, allow_zero = false)
     !allow_zero && D.steps == 0 && return x
@@ -38,7 +39,7 @@ function (D::Shift)(x::Num, allow_zero = false)
     if istree(vt)
         op = operation(vt)
         if op isa Shift
-            if isequal(D.t, op.t)
+            if D.t === nothing || isequal(D.t, op.t)
                 arg = arguments(vt)[1]
                 newsteps = D.steps + op.steps
                 return Num(newsteps == 0 ? arg : Shift(D.t, newsteps)(arg))
@@ -75,7 +76,7 @@ Represents a sample operator. A discrete-time signal is created by sampling a co
 
 # Constructors
 `Sample(clock::TimeDomain = InferredDiscrete())`
-`Sample(t, dt::Real)`
+`Sample([t], dt::Real)`
 
 `Sample(x::Num)`, with a single argument, is shorthand for `Sample()(x)`.
 

src/systems/clock_inference.jl

@@ -1,31 +1,30 @@
-struct ClockInference
-    ts::TearingState
+struct ClockInference{S}
+    ts::S
     eq_domain::Vector{TimeDomain}
     var_domain::Vector{TimeDomain}
     inferred::BitSet
 end
 
-function ClockInference(ts::TearingState)
-    @unpack fullvars, structure = ts
+function ClockInference(ts::TransformationState)
+    @unpack structure = ts
     @unpack graph = structure
     eq_domain = TimeDomain[Continuous() for _ in 1:nsrcs(graph)]
     var_domain = TimeDomain[Continuous() for _ in 1:ndsts(graph)]
     inferred = BitSet()
-    for (i, v) in enumerate(fullvars)
+    for (i, v) in enumerate(get_fullvars(ts))
         d = get_time_domain(v)
-        if d isa Union{AbstractClock, Continuous}
+        if is_concrete_time_domain(d)
             push!(inferred, i)
-            dd = d
-            var_domain[i] = dd
+            var_domain[i] = d
         end
     end
     ClockInference(ts, eq_domain, var_domain, inferred)
 end
 
 function infer_clocks!(ci::ClockInference)
     @unpack ts, eq_domain, var_domain, inferred = ci
-    @unpack fullvars = ts
     @unpack graph = ts.structure
+    fullvars = get_fullvars(ts)
     isempty(inferred) && return ci
     # TODO: add a graph type to do this lazily
     var_graph = SimpleGraph(ndsts(graph))
@@ -78,7 +77,7 @@ end
 
 function split_system(ci::ClockInference)
     @unpack ts, eq_domain, var_domain, inferred = ci
-    @unpack fullvars = ts
+    fullvars = get_fullvars(ts)
     @unpack graph, var_to_diff = ts.structure
     continuous_id = Ref(0)
     clock_to_id = Dict{TimeDomain, Int}()

src/systems/systemstructure.jl

@@ -31,6 +31,7 @@ export initialize_system_structure, find_linear_equations
 export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq, algeqs, is_only_discrete
 export dervars_range, diffvars_range, algvars_range
 export DiffGraph, complete!
+export get_fullvars
 
 struct DiffGraph <: Graphs.AbstractGraph{Int}
     primal_to_diff::Vector{Union{Int, Nothing}}
@@ -138,6 +139,8 @@ end
 abstract type TransformationState{T} end
 abstract type AbstractTearingState{T} <: TransformationState{T} end
 
+get_fullvars(ts::TransformationState) = ts.fullvars
+
 Base.@kwdef mutable struct SystemStructure
     # Maps the (index of) a variable to the (index of) the variable describing
     # its derivative.
@@ -201,6 +204,8 @@ mutable struct TearingState{T <: AbstractSystem} <: AbstractTearingState{T}
     extra_eqs::Vector
 end
 
+TransformationState(sys::AbstractSystem) = TearingState(sys)
+
 function Base.show(io::IO, state::TearingState)
     print(io, "TearingState of ", typeof(state.sys))
 end