Inline-linear-SCC: assert b-free invariant + opt-in reduction self-check (#98)

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/reassemble.jl

@@ -540,6 +540,37 @@ end
 
 const INLINE_LINEAR_SCC_OP = (\)
 
+"""
+    $TYPEDSIGNATURES
+
+Assert the invariant that the inline-linear-SCC right-hand side `b` is free of all SCC
+variables `vars`. A correct construction expands every SCC variable out of `b` into the
+coefficient matrix `A`; a live SCC variable left buried inside some `b[eqidx]` indicates the
+structural incidence graph desynced from the `total_sub`-substituted residual (issue #98),
+which would otherwise be smuggled onto the RHS by `__reduce_linear_system!`. Throws an error
+identifying the offending variable(s) if the invariant is violated.
+
+Only invoked when the self-check is enabled (`MTKTEARING_CHECK_REDUCTION`); the construction
+is expected to maintain this invariant unconditionally.
+"""
+function _assert_b_free_of_scc_vars(b::AbstractVector, vars)
+    scc_atoms = Set{SymbolicT}()
+    for var in vars
+        Symbolics.get_variables!(scc_atoms, var)
+    end
+    isempty(scc_atoms) && return nothing
+    bsyms = Set{SymbolicT}()
+    for eqidx in eachindex(b)
+        Symbolics.get_variables!(empty!(bsyms), b[eqidx])
+        buried = intersect(bsyms, scc_atoms)
+        isempty(buried) && continue
+        error("Inline-linear-SCC construction left SCC variable(s) $(collect(buried)) \
+               buried in the right-hand side `b[$eqidx]`; the structural incidence graph \
+               desynced from the substituted residual (issue #98).")
+    end
+    return nothing
+end
+
 """
     $TYPEDSIGNATURES
 
@@ -606,6 +637,15 @@ function get_linear_scc_linsol(state::TearingState, alg_eqs::Vector{Int},
         end
     end
 
+    # The `has_edge` gate above relies on the structural incidence `graph`, which is mutated
+    # during reassembly and could in principle desync from the `total_sub`-substituted
+    # residual, leaving a live SCC variable buried inside some `b[eqidx]`. Such a variable
+    # would be smuggled onto the RHS of retained equations by `__reduce_linear_system!`
+    # (which inspects only the coefficient row, never `b`), producing an inconsistent runtime
+    # `A \ b` (issue #98). The construction is expected to keep `b` free of all SCC
+    # variables; when the self-check is enabled, assert it so any regression surfaces loudly.
+    _inline_scc_check_enabled() && _assert_b_free_of_scc_vars(b, vars)
+
     # `-` is important! `b` is on the other side of the equality.
     for i in eachindex(b)
         b[i] = -b[i]
@@ -706,8 +746,135 @@ function get_linear_scc_linsol(state::TearingState, alg_eqs::Vector{Int},
     return (INLINE_LINEAR_SCC_OP(A, b), eqs_mask, vars_mask)
 end
 
+# ---------------------------------------------------------------------------
+# Opt-in self-checks for the inline-linear-SCC pass (issue #98).
+#
+# These are disabled unless the `MTKTEARING_CHECK_REDUCTION` environment variable
+# is set to a non-empty value. When off, the only added cost is a single `get(ENV,
+# ...)` per SCC and the snapshots below are skipped, so production pays nothing.
+# ---------------------------------------------------------------------------
+
+"""
+    $TYPEDSIGNATURES
+
+Whether the inline-linear-SCC self-checks are enabled. Controlled by the
+`MTKTEARING_CHECK_REDUCTION` environment variable (any non-empty value enables it).
+"""
+_inline_scc_check_enabled() = !isempty(get(ENV, "MTKTEARING_CHECK_REDUCTION", ""))
+
+# Numerically evaluate a symbolic expression under a substitution of *all* its free
+# symbols to numbers. Deliberately avoids `iszero`/`simplify`/`expand`, which can OOM
+# on large multibody coefficient expressions (see StateSelection.jl#95).
+_evalnum(x, subs::AbstractDict) =
+    Float64(Symbolics.value(Symbolics.substitute(unwrap(x), subs; fold = Val(true))))
+
+_free_syms_into!(s::AbstractSet, x) =
+    (union!(s, Symbolics.get_variables(unwrap(x))); s)
+
+# Build a deterministic random substitution for all free symbols appearing in the
+# given symbolic containers, plus an RNG seeded from the symbol *names* (so a reported
+# failure reproduces across runs). Returns `(subs, rng)`.
+function _deterministic_subs(containers...)
+    syms = Set{Any}()
+    for c in containers, x in c
+        _free_syms_into!(syms, x)
+    end
+    symvec = sort!(collect(syms); by = string)
+    seed = foldl((h, s) -> hash(string(s), h), symvec; init = UInt(0x5eed))
+    rng = Random.MersenneTwister(seed % typemax(UInt) + one(UInt))
+    subs = Dict{Any, Float64}(s => randn(rng) for s in symvec)
+    return subs, rng
+end
+
+"""
+    $TYPEDSIGNATURES
+
+Verify that the reduced linear system `A_red x_ret = b_red` produced by
+`__reduce_linear_system!` is the *exact* substitution-projection of the full system
+`A0 x = b0`, i.e. for every retained equation the reduced residual equals the full
+residual after replacing each eliminated variable `v` by `aliases[v]·x_ret +
+constants[v]`. Returns `true` iff the identity holds at a deterministic pseudo-random
+probe point. Pure and positional: it does not need to know which symbols are the SCC
+variables (those are represented by column position, not as free symbols).
+"""
+function _reduction_identity_ok(
+        A0::AbstractMatrix, b0::AbstractVector, A_red::AbstractMatrix, b_red::AbstractVector,
+        aliases::AbstractDict, constants::AbstractDict, eqs_mask::BitVector,
+        vars_mask::BitVector, old_to_new_eq::Vector{Int}; rtol::Float64 = 1e-7)
+    N = length(b0)
+    aliasvals = (SparseArrays.nonzeros(sv) for sv in values(aliases))
+    subs, rng = _deterministic_subs(A0, b0, A_red, b_red, values(constants), aliasvals...)
+
+    # New index of each retained variable (mirrors `cumsum(vars_mask)` in the caller).
+    old_to_new_var = zeros(Int, N)
+    let c = 0
+        for j in 1:N
+            vars_mask[j] || continue
+            old_to_new_var[j] = (c += 1)
+        end
+    end
+    nret = count(vars_mask)
+    xret = randn(rng, nret)
+
+    # Full-length solution: retained vars from `xret`, eliminated vars from their alias.
+    xfull = Vector{Float64}(undef, N)
+    for j in 1:N
+        vars_mask[j] && (xfull[j] = xret[old_to_new_var[j]])
+    end
+    for j in 1:N
+        vars_mask[j] && continue
+        acc = _evalnum(constants[j], subs)
+        I, V = SparseArrays.findnz(aliases[j])
+        for (k, coeff) in zip(I, V)
+            # After Phase-2 flattening, `k` references only retained variables.
+            acc += _evalnum(coeff, subs) * xfull[k]
+        end
+        xfull[j] = acc
+    end
+
+    nr = length(b_red)
+    ok = true
+    for i in 1:N
+        eqs_mask[i] || continue
+        ired = old_to_new_eq[i]
+        rfull = sum(_evalnum(A0[i, j], subs) * xfull[j] for j in 1:N; init = 0.0) - _evalnum(b0[i], subs)
+        rred = sum(_evalnum(A_red[ired, j], subs) * xret[j] for j in 1:nr; init = 0.0) - _evalnum(b_red[ired], subs)
+        if abs(rfull - rred) > rtol * (1 + abs(rfull))
+            ok = false
+            @warn "Inline-linear-SCC reduction identity violated" full_row=i reduced_row=ired residual_full=rfull residual_reduced=rred mismatch=abs(rfull - rred)
+        end
+    end
+    return ok
+end
+
+# Report the rank/consistency of the full and reduced blocks at a deterministic probe
+# point. Distinguishes "full block already bad (upstream construction)" from "reduction
+# broke it". Logs via `@info`; only called when the self-check is enabled.
+function _reduction_rank_report(A0::AbstractMatrix, b0::AbstractVector,
+        A_red::AbstractMatrix, b_red::AbstractVector, alg_vars::Vector{Int})
+    subs, _ = _deterministic_subs(A0, b0, A_red, b_red)
+    N = length(b0)
+    A0n = [_evalnum(A0[i, j], subs) for i in 1:N, j in 1:N]
+    b0n = [_evalnum(b0[i], subs) for i in 1:N]
+    nr = length(b_red)
+    Arn = [_evalnum(A_red[i, j], subs) for i in 1:nr, j in 1:nr]
+    brn = [_evalnum(b_red[i], subs) for i in 1:nr]
+    # Rank-tolerant residual: min-norm least-squares via the pseudoinverse, so a
+    # (legitimately) rank-deficient block does not throw and we can still tell whether
+    # `b` lies in the range of `A` (relres ≈ 0 consistent, relres ≈ 1 inconsistent).
+    relres(A, b) = isempty(b) ? 0.0 : LinearAlgebra.norm(A * (LinearAlgebra.pinv(A) * b) - b) / max(LinearAlgebra.norm(b), eps())
+    consistent(A, b) = LinearAlgebra.rank(A) == LinearAlgebra.rank(hcat(A, b))
+    @info "Inline-linear-SCC block diagnostics" alg_vars full_n=N full_rank=LinearAlgebra.rank(A0n) full_consistent=consistent(A0n, b0n) full_relres=relres(A0n, b0n) reduced_n=nr reduced_rank=(nr == 0 ? 0 : LinearAlgebra.rank(Arn)) reduced_consistent=(nr == 0 ? true : consistent(Arn, brn)) reduced_relres=relres(Arn, brn)
+    return nothing
+end
+
 function __reduce_linear_system!(A::StateSelection.CLIL.SparseMatrixCLIL{Num, Int}, b::Vector{SymbolicT}, var_eq_matching::StateSelection.VarEqMatchingT, alg_eqs::Vector{Int}, alg_vars::Vector{Int})
     N = length(b)
+    # Snapshot the pre-reduction system for the opt-in self-check. `A`'s rows are
+    # mutated in place below and `b` is reassigned, so these must be copies.
+    _check = _inline_scc_check_enabled()
+    A0_check = _check ? collect(A)::Matrix{Num} : nothing
+    b0_check = _check ? copy(b) : nothing
     # Identify rows (equations) not worth involving in the linear solve.
     #
     # The current heuristic is to find all rows with constant coefficients
@@ -817,6 +984,13 @@ function __reduce_linear_system!(A::StateSelection.CLIL.SparseMatrixCLIL{Num, In
     old_to_new_eq[.!eqs_mask] .= 0
     A = StateSelection.get_new_mm(aliases, old_to_new_eq, old_to_new_var, A)
 
+    if _check
+        A_red_check = collect(A)::Matrix{Num}
+        _reduction_identity_ok(A0_check, b0_check, A_red_check, b, aliases, constants,
+                               eqs_mask, vars_mask, old_to_new_eq)
+        _reduction_rank_report(A0_check, b0_check, A_red_check, b, alg_vars)
+    end
+
     return A, b, eqs_mask, vars_mask
 end
 

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,90 @@ 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 "`_assert_b_free_of_scc_vars` guards the b-free invariant (#98)" begin
+    SymT = Symbolics.SymbolicT
+    @variables x(t) y(t)
+    vars = [unwrap(x), unwrap(y)]
+
+    # A correct construction expands every SCC variable out of `b`, so the invariant holds
+    # and the assertion passes (returns `nothing`).
+    @test MTKTearing._assert_b_free_of_scc_vars(
+        SymT[unwrap(Num(5)), unwrap(Num(2))], vars) === nothing
+
+    # The #98 broken state — a live SCC variable left buried in `b` because the structural
+    # `has_edge` gate desynced from the substituted residual — is exactly what, if it
+    # reached `__reduce_linear_system!`, would be smuggled onto the RHS. The assertion must
+    # catch it.
+    @test_throws ErrorException MTKTearing._assert_b_free_of_scc_vars(
+        SymT[unwrap(3y + 5)], vars)
+    @test_throws ErrorException MTKTearing._assert_b_free_of_scc_vars(
+        SymT[unwrap(Num(5)), unwrap(2x + 1)], vars)
+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)