Merge pull request #90 from JuliaComputing/as/operator-arg-clocks

feat: track clocks of clock operator arguments in clock inference

Author: AayushSabharwal

lib/ModelingToolkitTearing/src/clock_inference/clock_inference.jl

@@ -1,9 +1,17 @@
 @data ClockVertex begin
+    # A variable identified by its index in `fullvars`
     Variable(Int)
+    # An equation identified by its index in `equations(state)`
     Equation(Int)
+    # An initialization equation identified by its index in `initialization_equations(state.sys)`
     InitEquation(Int)
+    # A concrete, known clock
     Clock(SciMLBase.AbstractClock)
+    # An unnamed, unnknown clock identified by some UUID stored in `ShiftIndex`
     InferredClock(UUID)
+    # An arbitrary expression for which we want to preserve clocking information
+    Expression(SymbolicT)
+    # Compatibility for ModelingToolkit's `DiscreteSystem` semantics.
     IntegerSequence
 end
 
@@ -21,6 +29,10 @@ struct ClockInference{S <: StateSelection.TransformationState}
     inference_graph::HyperGraph{ClockVertex.Type}
     """The set of variables with concrete domains."""
     inferred::BitSet
+    """A mapping from every argument of every clock operator to the clock it is on. More \
+    generally extensible to any expression for which it is deemed necessary to preserve \
+    clock information."""
+    expression_clocks::Dict{SymbolicT, TimeDomain}
 end
 
 function ClockInference(ts::StateSelection.TransformationState)
@@ -73,7 +85,7 @@ function ClockInference(ts::StateSelection.TransformationState)
         add_edge!(inference_graph, (varvert, dvert))
     end
     ClockInference{typeof(ts)}(ts, eq_domain, init_eq_domain, var_domain, inference_graph,
-                               inferred)
+                               inferred, Dict{SymbolicT, TimeDomain}())
 end
 
 struct NotInferredTimeDomain end
@@ -262,6 +274,7 @@ function (ivc::InferVariableClosure)(var::SymbolicT)
             nested_ivc(v)
         end
 
+        push!(arg_hyperedge, ClockVertex.Expression(arg))
         # NOTE: Ensure all branches here add `arg_hyperedge` to the inference graph. It
         # is the parent hyperedge for `nested_ivc`, so this closure is responsible for
         # adding it to the graph.
@@ -422,6 +435,7 @@ function infer_clocks!(ci::ClockInference)
                 ClockVertex.Clock(_) => nothing
                 ClockVertex.InferredClock(_) => nothing
                 ClockVertex.IntegerSequence() => nothing
+                ClockVertex.Expression(expr) => (ci.expression_clocks[expr] = clock)
             end
         end
     end

lib/ModelingToolkitTearing/test/runtests.jl

@@ -265,8 +265,10 @@ MTKTearing.is_timevarying_operator(::Op3) = true
 
 @testset "Clock inference handles nested clock operators correctly" begin
     @variables x(t) y(t) z(t) w(t)
-    k1 = ShiftIndex(Clock(0.1))
-    k2 = ShiftIndex(Clock(0.2))
+    clk1 = Clock(0.1)
+    clk2 = Clock(0.2)
+    k1 = ShiftIndex(clk1)
+    k2 = ShiftIndex(clk2)
 
     # From `Op1`, we should infer that `clock_of(z) == clock_of(y)`
     # From `Op2`, we should infer that `clock_of(w) == clock_of(Op3(x))`
@@ -281,6 +283,20 @@ MTKTearing.is_timevarying_operator(::Op3) = true
     @test issetequal(ts.fullvars, [x, x(k2-1), y, z, w, Op3(x), Op2(Op3(x), w), Op1(Op2(Op3(x), w) + x, 3sin(y))])
     ci = MTKTearing.ClockInference(ts)
     MTKTearing.infer_clocks!(ci)
+    idx = findfirst(isequal(w), ts.fullvars)
+    @test ci.var_domain[idx] == clk2
+    idx = findfirst(isequal(Op3(x)), ts.fullvars)
+    @test ci.var_domain[idx] == clk2
+    idx = findfirst(isequal(Op2(Op3(x), w)), ts.fullvars)
+    @test ci.var_domain[idx] == clk2
+    idx = findfirst(isequal(Op1(Op2(Op3(x), w) + x, 3sin(y))), ts.fullvars)
+    @test ci.var_domain[idx] == clk1
+
+    @test ci.expression_clocks[x] == clk2
+    @test ci.expression_clocks[Op3(x)] == clk2
+    @test ci.expression_clocks[w] == clk2
+    @test ci.expression_clocks[Op2(Op3(x), w) + x] == clk2
+    @test ci.expression_clocks[3sin(y)] == clk1
 
     # Check that this errors if `+ x` is not present.
     eqs = [