Pivot-solvability inline-linear row elimination with size-budgeted composition (#95)

lib/ModelingToolkitTearing/Project.toml

@@ -1,6 +1,6 @@
 name = "ModelingToolkitTearing"
 uuid = "6bb917b9-1269-42b9-9f7c-b0dca72083ab"
-version = "1.14.1"
+version = "1.14.2"
 authors = ["Aayush Sabharwal <aayush.sabharwal@gmail.com>"]
 
 [deps]
@@ -13,6 +13,7 @@ ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839"
 Moshi = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -34,6 +35,7 @@ ModelingToolkitBase = "1.37"
 Moshi = "0.3"
 OffsetArrays = "1"
 OrderedCollections = "1.8.1"
+Random = "1"
 SciMLBase = "2.108, 3"
 Setfield = "0.7, 0.8, 1"
 SparseArrays = "1"
@@ -47,7 +49,8 @@ julia = "1.10"
 [extras]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "ModelingToolkit", "ForwardDiff"]
+test = ["Test", "ModelingToolkit", "ForwardDiff", "SparseArrays"]

lib/ModelingToolkitTearing/src/ModelingToolkitTearing.jl

@@ -27,6 +27,7 @@ using SymbolicUtils: BSImpl, unwrap
 using SciMLBase: LinearProblem
 using SparseArrays: nonzeros
 import LinearAlgebra
+import Random
 import UUIDs: UUID, uuid4
 
 const TimeDomain = SciMLBase.AbstractClock

lib/ModelingToolkitTearing/src/stateselection_interface.jl

@@ -105,6 +105,13 @@ function StateSelection.linear_subsys_adjmat!(state::TearingState; kwargs...)
     return mm
 end
 
+# Structural zero check for symbolic CLIL values: `Base.iszero(::Num)` performs
+# a semantic (expansion-based) zero test that can OOM on large coefficient
+# expressions (e.g. multibody models), while explicit stored zeros produced by
+# duplicate-index summation are always structural `Const(0)`.
+StateSelection.CLIL.cheap_iszero(x::Num) = SU._iszero(Symbolics.unwrap(x))
+StateSelection.CLIL.cheap_iszero(x::SymbolicT) = SU._iszero(x)
+
 function maybe_zeros_descend(ex::SymbolicT)
     @match ex begin
         BSImpl.AddMul(; variant) => return variant === SU.AddMulVariant.MUL

lib/ModelingToolkitTearing/test/runtests.jl

@@ -11,6 +11,7 @@ import SymbolicUtils as SU
 using SymbolicUtils: unwrap
 using Setfield
 using ForwardDiff
+import SparseArrays
 
 @testset "`InferredDiscrete` validation" begin
     k = ShiftIndex()
@@ -171,6 +172,124 @@ end
     ) reassemble_alg = MTKTearing.DefaultReassembleAlgorithm(; inline_linear_sccs = true)
 end
 
+@testset "`__reduce_linear_system!` preserves the full-system residual" begin
+    SymT = Symbolics.SymbolicT
+    MVT = StateSelection.MatchedVarT
+
+    # Build the 4×4 SCC described in issue #98's plan. Variables x1..x4, equations e1..e4:
+    #   e1:  2*x2              = 0        # eliminates x2 (matched, const coeffs)
+    #   e2:  -3*x2 + x3        = 0        # eliminates x3 via x2 -> transitive chain
+    #   e3:  x1 + x2 + x3 + x4 = p        # RETAINED, references two eliminated vars, symbolic RHS
+    #   e4:  x1 + x4           = p        # RETAINED, makes the reduced block rank-deficient
+    # Matching: x2->e1, x3->e2 (eliminated); x1,x4 unassigned => e3,e4 retained.
+    @variables p
+    mkA() = StateSelection.CLIL.SparseMatrixCLIL{Num, Int}(
+        4, 4, collect(1:4),
+        [[2], [2, 3], [1, 2, 3, 4], [1, 4]],
+        [Num[2.0], Num[-3.0, 1.0], Num[1.0, 1.0, 1.0, 1.0], Num[1.0, 1.0]])
+    mkb() = SymT[unwrap(Num(0)), unwrap(Num(0)), unwrap(p), unwrap(p)]
+    vem = BipartiteGraphs.complete(
+        BipartiteGraphs.Matching{MVT}(Union{MVT, Int}[BipartiteGraphs.unassigned, 1, 2, BipartiteGraphs.unassigned]),
+        4)
+
+    Ar, br, em, vm = MTKTearing.__reduce_linear_system!(mkA(), mkb(), vem, collect(1:4), collect(1:4))
+
+    @test em == Bool[0, 0, 1, 1]
+    @test vm == Bool[1, 0, 0, 1]
+
+    # The reduction is exact: x2=0, x3=0, so both retained rows become `x1 + x4 = p`.
+    subs = Dict{Any, Float64}(unwrap(p) => 3.7)
+    ev(x) = MTKTearing._evalnum(x, subs)
+    @test ev.(collect(Ar)) == [1.0 1.0; 1.0 1.0]   # rank-deficient (rank 1), as expected
+    @test ev.(br) ≈ [3.7, 3.7]                      # consistent: b in range(A)
+
+    # Exercise the opt-in self-check code path end-to-end (snapshot + identity + rank report).
+    local res
+    withenv("MTKTEARING_CHECK_REDUCTION" => "1") do
+        res = MTKTearing.__reduce_linear_system!(mkA(), mkb(), vem, collect(1:4), collect(1:4))
+    end
+    @test res[3] == Bool[0, 0, 1, 1]
+end
+
+@testset "`_reduction_identity_ok` detects reduction errors" begin
+    SymT2 = Symbolics.SymbolicT
+    @variables p
+    # Full 2×2 system: e1: 2*x1 = p (eliminate x1), e2: x1 + x2 = 0 (retain x2).
+    # Correct reduction: x1 = p/2, so e2 becomes x2 = -p/2.
+    A0 = Num[2.0 0.0; 1.0 1.0]
+    b0 = SymT2[unwrap(p), unwrap(Num(0))]
+    aliases = Dict{Int, SparseArrays.SparseVector{Num, Int}}(1 => SparseArrays.spzeros(Num, 2))
+    constants = Dict{Int, SymT2}(1 => unwrap(p / 2))
+    eqs_mask = BitVector([false, true])
+    vars_mask = BitVector([false, true])
+    old_to_new_eq = [0, 1]
+
+    A_red = Num[1.0;;]
+    b_red_good = SymT2[unwrap(-p / 2)]
+    @test MTKTearing._reduction_identity_ok(
+        A0, b0, A_red, b_red_good, aliases, constants, eqs_mask, vars_mask, old_to_new_eq)
+
+    # A wrong RHS (off by a constant) must be caught.
+    b_red_bad = SymT2[unwrap(-p / 2 + 1)]
+    bad = @test_logs (:warn,) match_mode = :any MTKTearing._reduction_identity_ok(
+        A0, b0, A_red, b_red_bad, aliases, constants, eqs_mask, vars_mask, old_to_new_eq)
+    @test bad == false
+end
+
+@testset "`_group_inline_linear_families` merges coupled rank-deficible blocks" begin
+    @variables a(t) b(t) c(t) d(t)
+    @parameters p
+    # Equations indexed 1..4. Blocks 1,2,4 reference the `maybe_zeros` parameter `p`
+    # (their rank can drop); block 3 does not.
+    neweqs = [
+        p * a ~ 0,        # eq1
+        0 ~ b - p * a,    # eq2  (couples to block 1's variable below)
+        0 ~ c - a,        # eq3  (no `p`)
+        0 ~ d - p * c,    # eq4
+    ]
+    prepared = NTuple{2, Vector{Int}}[([1], [1]), ([2], [2]), ([3], [3]), ([4], [4])]
+    g = BipartiteGraph(4, 4)
+    add_edge!(g, BipartiteEdge(2, 1))  # block 2's equation references block 1's variable
+    # (block 4 is intentionally NOT coupled to block 3)
+    mz = Symbolics.SymbolicT[unwrap(p)]
+
+    # Coupled + rank-deficible blocks 1,2 merge; block 3 (not deficible) and block 4
+    # (not coupled to blocks 1,2) stay separate.
+    @test MTKTearing._group_inline_linear_families(prepared, mz, neweqs, g, false, true) ==
+          [[1, 2], [3], [4]]
+
+    # No `maybe_zeros` => no grouping at all (default one-block-per-SCC behaviour).
+    @test MTKTearing._group_inline_linear_families(prepared, Symbolics.SymbolicT[], neweqs, g, false, true) ==
+          [[1], [2], [3], [4]]
+    # Discrete or inline-linear disabled => singletons.
+    @test MTKTearing._group_inline_linear_families(prepared, mz, neweqs, g, true, true) ==
+          [[1], [2], [3], [4]]
+    @test MTKTearing._group_inline_linear_families(prepared, mz, neweqs, g, false, false) ==
+          [[1], [2], [3], [4]]
+
+    # Non-adjacent family members are connected through the dependency DAG. Blocks 1 and 3
+    # are rank-deficible; block 2 is not, but lies on the dependency path 1 → 2 → 3, so the
+    # path closure pulls it into the family.
+    neweqs2 = [
+        p * a ~ 0,        # block 1 (deficible)
+        0 ~ b - a,        # block 2 (not deficible, on the path between 1 and 3)
+        0 ~ c - p * b,    # block 3 (deficible)
+    ]
+    prepared2 = NTuple{2, Vector{Int}}[([1], [1]), ([2], [2]), ([3], [3])]
+    g2 = BipartiteGraph(3, 3)
+    add_edge!(g2, BipartiteEdge(2, 1))   # block 2's equation references block 1's variable
+    add_edge!(g2, BipartiteEdge(3, 2))   # block 3's equation references block 2's variable
+    @test MTKTearing._group_inline_linear_families(prepared2, mz, neweqs2, g2, false, true) ==
+          [[1, 2, 3]]
+
+    # An unrelated block between two family members is NOT pulled in (no path through it),
+    # and the family is emitted before it (contracted topological order, stable by index).
+    g3 = BipartiteGraph(3, 3)
+    add_edge!(g3, BipartiteEdge(3, 1))   # block 3's equation references block 1's variable
+    @test MTKTearing._group_inline_linear_families(prepared2, mz, neweqs2, g3, false, true) ==
+          [[1, 3], [2]]
+end
+
 @testset "`system_subset(::SystemStructure)` subsets `.state_priorities`" begin
     @variables x(t) y(t) [state_priority = 2] z(t) [state_priority = 5]
     @named sys = System([D(x) ~ x, D(y) ~ y, D(z) ~ z], t)

src/math/sparsematrixclil.jl

@@ -91,14 +91,26 @@ zero!(a::SparseVector) = (empty!(a.nzind); empty!(a.nzval))
 zero!(a::CLILVector) = zero!(a.vec)
 SparseArrays.dropzeros!(a::CLILVector) = SparseArrays.dropzeros!(a.vec)
 
+"""
+    cheap_iszero(x)
+
+Structural zero check used by [`SparseArrays.dropzeros!`](@ref) on
+`SparseMatrixCLIL`. Defaults to `Base.iszero`. Downstream packages whose CLIL
+value type is symbolic should overload this with a cheap *structural* check:
+`Base.iszero` on e.g. `Symbolics.Num` performs a semantic (expansion-based)
+zero test that can be arbitrarily expensive on large expressions, while
+explicit stored zeros are always structural zeros.
+"""
+cheap_iszero(x) = iszero(x)
+
 # Remove explicitly-stored zeros from each row, in place.
 function SparseArrays.dropzeros!(S::SparseMatrixCLIL)
     for r in eachindex(S.row_vals)
         cols = S.row_cols[r]
         vals = S.row_vals[r]
         j = 0
         for k in eachindex(vals)
-            iszero(vals[k]) && continue
+            cheap_iszero(vals[k]) && continue
             j += 1
             cols[j] = cols[k]
             vals[j] = vals[k]

src/partial_state_selection.jl

@@ -187,6 +187,86 @@ struct DummyDerivativeSummary
     state_priority::Vector{Vector{Float64}}
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Merge the SCCs of the Pantelides matching into the blocks on which dummy-derivative
+selection must operate.
+
+The Mattsson–Söderlind dummy-derivative selection problem is posed on the subproblem of
+differentiated equations and their highest-derivative candidate variables. The SCCs of
+the full matching can be strictly finer than the blocks of that subproblem: a
+differentiated equation matched to a candidate in one SCC may be incident to candidate
+variables in other SCCs. A common example is a twice-differentiated connection alias
+`0 ~ D(D(x)) - D(D(y))` matched to `D(D(x))`, with `D(D(y))` belonging to a
+kinematic-loop SCC: `D(D(x))` then sits in a singleton SCC where it is demoted
+unconditionally, and a high `state_priority` on `x` cannot prevent it even though
+demoting `D(D(y))` instead would be structurally valid (see issue #101). Selecting per
+merged block restores the full selection freedom of the subproblem, and the
+priority-sorted greedy selection inside `dummy_derivative_graph!` then maximizes the
+total priority of the kept states.
+
+Returns the merged list of variable blocks; SCCs without coupling are returned
+unchanged (in particular the result is `===` the input when nothing merges).
+"""
+function merge_dummy_derivative_blocks(
+        structure::SystemStructure, var_eq_matching, var_sccs::Vector{Vector{Int}})
+    (; eq_to_diff, var_to_diff, graph) = structure
+    diff_to_eq = invview(eq_to_diff)
+    diff_to_var = invview(var_to_diff)
+
+    # SCC index of every variable that is a dummy-derivative candidate of its SCC,
+    # mirroring the candidate filter in `dummy_derivative_graph!`.
+    scc_of_candidate = zeros(Int, ndsts(graph))
+    for (i, vars) in enumerate(var_sccs), var in vars
+        var_eq_matching[var] isa Int || continue
+        (diff_to_var[var] !== nothing && is_present(structure, var)) || continue
+        scc_of_candidate[var] = i
+    end
+
+    # Union-find over SCC indices, merging along differentiated equations that are
+    # incident to candidate variables outside the SCC they are matched in.
+    parent = collect(1:length(var_sccs))
+    function root(i::Int)
+        while parent[i] != i
+            parent[i] = parent[parent[i]]
+            i = parent[i]
+        end
+        i
+    end
+    merged_any = false
+    for (i, vars) in enumerate(var_sccs), var in vars
+        eq = var_eq_matching[var]
+        eq isa Int || continue
+        diff_to_eq[eq] === nothing && continue
+        for var2 in 𝑠neighbors(graph, eq)
+            j = scc_of_candidate[var2]
+            (j == 0 || j == i) && continue
+            ri = root(i)
+            rj = root(j)
+            ri == rj && continue
+            # union by min keeps roots at the first SCC of each block, which
+            # preserves the original SCC order in the output
+            parent[max(ri, rj)] = min(ri, rj)
+            merged_any = true
+        end
+    end
+    merged_any || return var_sccs
+
+    buckets = Dict{Int, Vector{Int}}()
+    order = Int[]
+    for (i, vars) in enumerate(var_sccs)
+        r = root(i)
+        b = get!(buckets, r) do
+            push!(order, r)
+            Int[]
+        end
+        append!(b, vars)
+    end
+    sort!(order)
+    return [buckets[r] for r in order]
+end
+
 """
     $TYPEDSIGNATURES
 
@@ -227,6 +307,7 @@ function dummy_derivative_graph!(
     end
 
     var_sccs = find_var_sccs(graph, var_eq_matching)
+    var_sccs = merge_dummy_derivative_blocks(structure, var_eq_matching, var_sccs)
     var_perm = Int[]
     var_dummy_scc = Vector{Int}[]
     var_state_priority = Vector{Float64}[]
@@ -281,6 +362,9 @@ function dummy_derivative_graph!(
                 sortperm!(var_perm, sp)
                 permute!(vars, var_perm)
                 permute!(sp, var_perm)
+                # keep the Jacobian columns aligned with the permuted variable
+                # order; `col_order` below indexes into `vars` (#102)
+                J === nothing || (J = J[:, var_perm])
                 push!(var_dummy_scc, copy(vars))
                 push!(var_state_priority, sp)
             end

src/utils.jl

@@ -345,7 +345,10 @@ function get_new_mm(
             # entry: a prior cancellation may have `pop!`ed the matching entry.
             if !isempty(final_row_cols) && col == final_row_cols[end]
                 final_row_vals[end] += new_row_val_i[indices[i]]
-                if iszero(final_row_vals[end])
+                # Syntactic zero test: semantic `iszero` on symbolic coefficients
+                # can OOM via polynomial expansion (#95); an uncancelled exact zero
+                # is pruned later by CLIL `dropzeros!`.
+                if CLIL.cheap_iszero(final_row_vals[end])
                     pop!(final_row_cols)
                     pop!(final_row_vals)
                 end