refactor: store mm in TearingState

Project.toml

@@ -87,7 +87,7 @@ LinearSolve = "3.66"
 Logging = "1"
 ModelingToolkitBase = "1.48"
 ModelingToolkitStandardLibrary = "2.20"
-ModelingToolkitTearing = "1.16"
+ModelingToolkitTearing = "1.17.2"
 Moshi = "0.3.6"
 NonlinearSolve = "4.3"
 OffsetArrays = "1.3.0"
@@ -116,7 +116,7 @@ Serialization = "1"
 Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "2.10.0"
 SparseArrays = "1"
-StateSelection = "1.9.1"
+StateSelection = "1.10.1"
 StaticArrays = "1.9.14"
 StochasticDiffEq = "6.82.0, 7"
 SymbolicIndexingInterface = "0.3.46"

src/structural_transformation/symbolics_tearing.jl

@@ -16,12 +16,12 @@ new residual equations after tearing. End users are encouraged to call [`mtkcomp
 instead, which calls this function internally.
 """
 function tearing(
-        sys::AbstractSystem, state = TearingState(sys); mm = nothing,
+        sys::AbstractSystem, state = TearingState(sys);
         reassemble_alg::ReassembleAlgorithm = DefaultReassembleAlgorithm(),
         fully_determined = true, kwargs...
     )
     tearing_result, extras = tearing(state; kwargs...)
-    return invalidate_cache!(reassemble_alg(state, tearing_result, mm; fully_determined, kwargs...))
+    return invalidate_cache!(reassemble_alg(state, tearing_result, state.mm; fully_determined, kwargs...))
 end
 
 function safe_isinteger(@nospecialize(x::Number))
@@ -57,7 +57,7 @@ the system is balanced.
 function dummy_derivative(
         sys, state = TearingState(sys);
         reassemble_alg::ReassembleAlgorithm = DefaultReassembleAlgorithm(),
-        mm = nothing, fully_determined = true, kwargs...
+        fully_determined = true, kwargs...
     )
     jac = let state = state
         (eqs, vars) -> begin
@@ -81,5 +81,5 @@ function dummy_derivative(
     tearing_result, extras = StateSelection.dummy_derivative_graph!(
         state, jac; state_priority, kwargs...
     )
-    return reassemble_alg(state, tearing_result, mm; fully_determined, kwargs...)
+    return reassemble_alg(state, tearing_result, state.mm; fully_determined, kwargs...)
 end

src/systems/alias_elimination.jl

@@ -110,9 +110,10 @@ equation, while physics variables appear in many dynamics equations); if that is
 tied, one of the tied variables is chosen arbitrarily.
 """
 function pick_alias_target(
-        fullvars::Vector{SymbolicT}, group_vars::Vector{Int}, state_priorities, irreducibles::AtomicSetT,
-        graph = nothing
+        state::TearingState, group_vars::Vector{Int}, irreducibles::AtomicSetT
     )
+    (; fullvars, structure) = state
+    (; graph, canonical_ranks, state_priorities) = structure
     irr_idx = findfirst(
         Base.Fix1(contains_possibly_indexed_element, irreducibles) ∘ Base.Fix1(getindex, fullvars),
         group_vars
@@ -121,13 +122,14 @@ function pick_alias_target(
     max_priority = maximum(Base.Fix1(getindex, state_priorities), group_vars)
     candidates = filter(v -> state_priorities[v] == max_priority, group_vars)
     if length(candidates) > 1 && max_priority > 0
-        tied_names = getindex.(Ref(fullvars), candidates)
-        @warn "Multiple variables in an alias group share the highest state_priority \
-($max_priority); choosing alias target by equation count. Tied variables: $tied_names"
-        if graph !== nothing
-            max_degree = maximum(v -> length(𝑑neighbors(graph, v)), candidates)
-            candidates = filter(v -> length(𝑑neighbors(graph, v)) == max_degree, candidates)
+        if max_priority >= 100
+            tied_names = getindex.(Ref(fullvars), candidates)
+            @warn "Multiple variables in an alias group share the highest state_priority \
+            ($max_priority); choosing alias target by equation count. Tied variables: $tied_names"
         end
+        max_degree = maximum(v -> length(𝑑neighbors(graph, v)), candidates)
+        filter!(v -> length(𝑑neighbors(graph, v)) == max_degree, candidates)
+        sort!(candidates; by = Base.Fix1(getindex, canonical_ranks))
     end
     return candidates[1]
 end
@@ -253,7 +255,7 @@ function find_perfect_aliases!(
         end
         # For consistent groups pick a target (survives as unknown) and rebase
         # parities relative to it via `target_p`.
-        target = pick_alias_target(fullvars, group_vars, state_priorities, irreducibles, graph)
+        target = pick_alias_target(state, group_vars, irreducibles)
         group_target[root] = target
         target_p = parity[target]
         for v in group_vars

src/systems/systemstructure.jl

@@ -216,44 +216,44 @@ function _mtkcompile!(
     old_to_new_eq, old_to_new_var, aliases = eliminate_perfect_aliases!(state)
     sys = state.sys
     mm = StateSelection.get_new_mm(aliases, old_to_new_eq, old_to_new_var, mm)
+    state.mm = mm
+    @assert mm.nparentrows == nsrcs(state.structure.graph) && mm.ncols == ndsts(state.structure.graph) lazy"""
+    Invalid `mm`. Got `nparentrows, ncols` = ($(mm.nparentrows), $(mm.ncols)).
+    Expected ($(nsrcs(state.structure.graph)), $(ndsts(state.structure.graph))).
+    """
     if check_consistency
         fully_determined = StateSelection.check_consistency(
             state, orig_inputs; nothrow = fully_determined === nothing
         )
     end
-    # This phrasing avoids making the `kwcall` dynamic dispatch due to the type of a
-    # keyword (`mm`) being non-concrete
-    if mm isa CLIL.SparseMatrixCLIL{BigInt, Int}
-        sys = _mtkcompile_worker!(state, sys, mm; fully_determined, dummy_derivative, kwargs...)
-    else
-        sys = _mtkcompile_worker!(state, sys, mm; fully_determined, dummy_derivative, kwargs...)
-    end
+    sys = _mtkcompile_worker!(state, sys; fully_determined, dummy_derivative, kwargs...)
     fullunknowns = [observables(sys); unknowns(sys)]
     @set! sys.observed = MTKBase.topsort_equations(sys, observed(sys), fullunknowns)
 
     return MTKBase.invalidate_cache!(sys)
 end
 
 function _mtkcompile_worker!(
-        state::TearingState, sys::System, mm::CLIL.SparseMatrixCLIL{T, Int};
+        state::TearingState, sys::System;
         fully_determined::Bool, dummy_derivative::Bool,
         kwargs...
-    ) where {T}
+    )
     if fully_determined && dummy_derivative
         sys = ModelingToolkit.dummy_derivative(
-            sys, state; mm, kwargs...
+            sys, state; kwargs...
         )
     elseif fully_determined
         var_eq_matching = StateSelection.pantelides!(state; finalize = false, kwargs...)
         sys = pantelides_reassemble(state, var_eq_matching)
         state = TearingState(sys)
-        sys, mm::CLIL.SparseMatrixCLIL{T, Int} = ModelingToolkit.alias_elimination!(state; fully_determined, kwargs...)
+        sys, mm = ModelingToolkit.alias_elimination!(state; fully_determined, kwargs...)
+        state.mm = mm
         sys = ModelingToolkit.dummy_derivative(
-            sys, state; mm, fully_determined, kwargs...
+            sys, state; fully_determined, kwargs...
         )
     else
         sys = ModelingToolkit.tearing(
-            sys, state; mm, fully_determined, kwargs...
+            sys, state; fully_determined, kwargs...
         )
     end
     return sys

test/reduction.jl

@@ -377,7 +377,7 @@ end
     # so the equation-count heuristic picks x as target — but the priority tie
     # (both at 5) must still emit a warning regardless of how the tie is broken.
     @variables x(t) y(t)
-    @named sys = System([D(x) ~ -x, x ~ y], t; state_priorities = [x => 5, y => 5])
+    @named sys = System([D(x) ~ -x, x ~ y], t; state_priorities = [x => 105, y => 105])
     state = TearingState(sys)
     @test_logs (:warn, r"state_priority") match_mode = :any ModelingToolkit.eliminate_perfect_aliases!(state)
     # Exactly one of the two variables is eliminated as observed.