fix: fix clock inference accidentally merging all InferredDiscrete partitions

.github/workflows/CI.yml

@@ -13,7 +13,7 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    name: Test ${{ matrix.pkg }} - Julia v${{ matrix.version }}
     runs-on: ${{ matrix.os }}
     timeout-minutes: 60
     permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created
@@ -23,26 +23,38 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - 'nightly'
+          - '1'
         os:
           - ubuntu-latest
         arch:
           - x64
+        pkg:
+          - StateSelection
+          - ModelingToolkitTearing
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v1
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
-      - uses: julia-actions/julia-processcoverage@v1
-#      - uses: codecov/codecov-action@v4
-#        with:
-#          files: lcov.info
-#          token: ${{ secrets.CODECOV_TOKEN }}
-#          fail_ci_if_error: false
-#      - uses: julia-actions/julia-uploadcoveralls@v1
-#        env:
-#          COVERALLS_TOKEN: ${{ secrets.COVERALLS_TOKEN }}
+      - name: "Test ${{ matrix.pkg }}"
+        shell: julia --color=yes --check-bounds=yes --depwarn=yes {0}
+        env:
+          PKG_TO_TEST: ${{ matrix.pkg }}
+        run: |
+          using Pkg
+          const PKG = ENV["PKG_TO_TEST"]
+          if PKG == "StateSelection"
+            @info "Testing StateSelection"
+            Pkg.activate(".")
+            Pkg.test()
+          elseif PKG == "ModelingToolkitTearing"
+            @info "Testing ModelingToolkitTearing"
+            Pkg.activate("lib/ModelingToolkitTearing")
+            Pkg.develop(; path = ".")
+            Pkg.test()
+          else
+            @error "Unknown package" PKG
+            exit(1)
+          end

lib/ModelingToolkitTearing/Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkitTearing"
 uuid = "6bb917b9-1269-42b9-9f7c-b0dca72083ab"
-authors = ["Aayush Sabharwal <aayush.sabharwal@gmail.com>"]
 version = "1.2.6"
+authors = ["Aayush Sabharwal <aayush.sabharwal@gmail.com>"]
 
 [deps]
 BipartiteGraphs = "caf10ac8-0290-4205-88aa-f15908547e8d"
@@ -27,6 +27,7 @@ CommonSolve = "0.2"
 DocStringExtensions = "0.7, 0.8, 0.9"
 Graphs = "1"
 LinearAlgebra = "1"
+ModelingToolkit = "11"
 ModelingToolkitBase = "1.4"
 Moshi = "0.3"
 OffsetArrays = "1"
@@ -39,3 +40,10 @@ SymbolicIndexingInterface = "0.3"
 SymbolicUtils = "4.3"
 Symbolics = "7.1"
 julia = "1.10"
+
+[extras]
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test", "ModelingToolkit"]

lib/ModelingToolkitTearing/src/clock_inference/clock_inference.jl

@@ -3,7 +3,6 @@
     Equation(Int)
     InitEquation(Int)
     Clock(SciMLBase.AbstractClock)
-    AssertDiscrete
 end
 
 struct ClockInference{S <: StateSelection.TransformationState}
@@ -69,19 +68,21 @@ struct InferEquationClosure
     arg_hyperedge::Set{ClockVertex.Type}
     # mapping from `i` in `InferredDiscrete(i)` to the vertices in that inferred partition
     relative_hyperedges::Dict{Int, Set{ClockVertex.Type}}
+    # vertices that must be discrete
+    must_be_discrete::Set{ClockVertex.Type}
     var_to_idx::Dict{SymbolicT, Int}
     inference_graph::HyperGraph{ClockVertex.Type}
 end
 
 function InferEquationClosure(var_to_idx, inference_graph)
     InferEquationClosure(Set{SymbolicT}(), Set{SymbolicT}(), Set{ClockVertex.Type}(),
                          Set{ClockVertex.Type}(), Dict{Int, Set{ClockVertex.Type}}(),
-                         var_to_idx, inference_graph)
+                         Set{ClockVertex.Type}(), var_to_idx, inference_graph)
 end
 
 function (iec::InferEquationClosure)(ieq::Int, eq::Equation, is_initialization_equation::Bool)
     (; varsbuf, arg_varsbuf, hyperedge, arg_hyperedge, relative_hyperedges) = iec
-    (; var_to_idx, inference_graph) = iec
+    (; must_be_discrete, var_to_idx, inference_graph) = iec
     empty!(varsbuf)
     empty!(hyperedge)
     # get variables in equation
@@ -148,9 +149,8 @@ function (iec::InferEquationClosure)(ieq::Int, eq::Equation, is_initialization_e
                 InferredClock.InferredDiscrete(i) => begin
                     relative_edge = get!(Set{ClockVertex.Type}, relative_hyperedges, i)
                     union!(relative_edge, arg_hyperedge)
-                    # Ensure that this clock partition will be discrete. This is a separate
-                    # variant because I don't want to give `InferredDiscrete` too many meanings.
-                    push!(arg_hyperedge, ClockVertex.AssertDiscrete())
+                    # Ensure that this clock partition will be discrete.
+                    union!(must_be_discrete, arg_hyperedge)
                     add_edge!(inference_graph, arg_hyperedge)
                 end
             end
@@ -168,7 +168,7 @@ function (iec::InferEquationClosure)(ieq::Int, eq::Equation, is_initialization_e
                         union!(hyperedge, buffer)
                         delete!(relative_hyperedges, i)
                     end
-                    push!(hyperedge, ClockVertex.AssertDiscrete())
+                    union!(must_be_discrete, hyperedge)
                 end
             end
         else
@@ -222,11 +222,13 @@ function infer_clocks!(ci::ClockInference)
         infer_equation(ieq, eq, true)
     end
 
+    (; must_be_discrete) = infer_equation
+
     clock_partitions = Graphs.connected_components(inference_graph)
     for partition in clock_partitions
         clockidxs = findall(Base.Fix2(Moshi.Data.isa_variant, ClockVertex.Clock), partition)
         if isempty(clockidxs)
-            if any(isequal(ClockVertex.AssertDiscrete()), partition)
+            if any(in(must_be_discrete), partition)
                 throw(ExpectedDiscreteClockPartitionError(ts, partition, true))
             end
             push!(partition, ClockVertex.Clock(SciMLBase.ContinuousClock()))
@@ -246,7 +248,7 @@ function infer_clocks!(ci::ClockInference)
         clock = Moshi.Match.@match partition[only(clockidxs)] begin
             ClockVertex.Clock(clk) => clk
         end
-        if clock == SciMLBase.ContinuousClock() && any(isequal(ClockVertex.AssertDiscrete()), partition)
+        if clock == SciMLBase.ContinuousClock() && any(in(must_be_discrete), partition)
             throw(ExpectedDiscreteClockPartitionError(ts, partition, false))
         end
         for vert in partition
@@ -255,7 +257,6 @@ function infer_clocks!(ci::ClockInference)
                 ClockVertex.Equation(i) => (eq_domain[i] = clock)
                 ClockVertex.InitEquation(i) => (init_eq_domain[i] = clock)
                 ClockVertex.Clock(_) => nothing
-                ClockVertex.AssertDiscrete() => nothing
             end
         end
     end

lib/ModelingToolkitTearing/test/runtests.jl

@@ -0,0 +1,13 @@
+using Test
+using ModelingToolkitTearing
+using ModelingToolkit
+using ModelingToolkit: t_nounits as t, D_nounits as D
+
+@testset "`InferredDiscrete` validation" begin
+    k = ShiftIndex()
+    @variables x(t) y(t)
+    @named sys = System([x(k) ~ x(k-1) + x(k-2), D(y) ~ Hold(x) + t], t)
+    @test_throws ModelingToolkitTearing.ExpectedDiscreteClockPartitionError mtkcompile(sys)
+    @named sys = System([x(k) ~ x(k-1) + x(k-2), D(y) ~ x + t], t)
+    @test_throws ModelingToolkitTearing.ExpectedDiscreteClockPartitionError mtkcompile(sys)
+end