From 50595d5ad0d55c78205a690bce5b6a35501a3ef6 Mon Sep 17 00:00:00 2001 From: Tim Holy Date: Sun, 25 Oct 2026 03:15:32 -0500 Subject: [PATCH] Max-balance dense blocks with an adjacency-matrix tree MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_maxbalance` holds a block with more than `n²/8` arcs in a `DenseParametricTree`: two `n×n` weight matrices with parallel arcs merged into the heaviest, and an indexed heap over nodes keyed by the best candidate arc entering each, rescanned lazily when a key may have fallen. A node whose depth changes then costs one pass over the live nodes instead of a heap update for each incident arc on a heap of `n²` entries. Sparser blocks keep the arc-list `ParametricTree`; the keyword `dense` selects the representation explicitly. The initial tree and the summation of the contraction chains are shared. On dense random lognormal matrices, `cover_transversal(PowerMean{2}(), A; maxbalance=true)` takes 0.8 s at n = 1600 and 20 s at n = 6400, and the shifts agree with the arc-list path to roundoff. Assisted-by: Claude Opus 5.5 and Claude Fable 5.1 --- docs/src/objectives.md | 5 +- src/balanced.jl | 8 +- src/maxbalance.jl | 320 +++++++++++++++++++++++++++++++++++++---- test/maxbalance.jl | 46 ++++-- 4 files changed, 335 insertions(+), 44 deletions(-) diff --git a/docs/src/objectives.md b/docs/src/objectives.md index 9d30aba..766226b 100644 --- a/docs/src/objectives.md +++ b/docs/src/objectives.md @@ -88,8 +88,9 @@ The method used to minimize the objective depends on the penalty: - `PowerMean{p}` with a target ratio, for [`cover_transversal`](@ref) (native), uses an active-set Newton method that keeps every iterate a hard cover tight on the transversal. With `maxbalance=true`, a parametric shortest-path - algorithm max-balances each block first, and the Newton method acts on one - scale per block. + algorithm max-balances each block first, at a cost comparable to that of + finding the transversal for dense and for typical sparse matrices, and the + Newton method acts on one scale per block. - `AbsLog{1}` is a linear program solved with JuMP and HiGHS. - `AbsLinear` is not convex; several starts are each refined to a local minimum with JuMP and Ipopt, and the best is returned. diff --git a/src/balanced.jl b/src/balanced.jl index 489087c..5f33beb 100644 --- a/src/balanced.jl +++ b/src/balanced.jl @@ -107,8 +107,12 @@ shortest-path algorithm of Young, Tarjan & Orlin (Networks 21:205–221, 1991). For a block with `n` rows and `m` entries, its cost is at most `O(nm + n² log n)`. The cost is roughly linear in `m` when the block's shortest-path tree stays shallow, as it does for most sparse matrices, and -about `O(nm)` for a banded block, whose tree is a path. The active-set Newton -method then runs with one scale per block. +about `O(nm)` for a banded block, whose tree is a path. A block with more than +`n²/8` entries is held in two dense `n×n` arrays, so that a row moving in the +tree costs `O(n)` operations without a heap update for each of its entries; on +dense random matrices the cost then grows about as `n^2.4` and is comparable +to that of finding the transversal. The active-set Newton method then runs with +one scale per block. `tol` (default `1e-9`) is the stopping tolerance: the groups are stationary when the derivative of the objective with respect to each group's scale is at most diff --git a/src/maxbalance.jl b/src/maxbalance.jl index c14fad8..f046763 100644 --- a/src/maxbalance.jl +++ b/src/maxbalance.jl @@ -141,21 +141,54 @@ function ParametricTree(nv::Int, src::Vector{Int}, dst::Vector{Int}, w::Vector{T zeros(Int, 2ne), zeros(Int, 2ne), IndexedHeap{T}(ne)) end -function _child_append!(S::ParametricTree, p::Int, c::Int) +# The parametric shortest-path tree of a dense digraph, with its arcs held in +# matrices. Parallel arcs are merged into one of largest weight: they share the +# coefficient of `λ` in their reduced lengths, so that one becomes tight first. +# Between live nodes `p ≠ q` there is then at most one arc `p → q`, of weight +# `win[p, q] == wout[q, p]`, or `-Inf` if there is none; column `q` of `win` +# holds the arcs entering `q`, and column `p` of `wout` the arcs leaving `p`. +# +# - `live` lists the live nodes, and `islive` marks them. +# - `a`, `b`, and the lists of children are as for `ParametricTree`; +# `parent[v]` is the tail of the tree arc entering `v`, or 0 at the root. +# - The candidate arcs are the non-tree arcs `p → q` with `a[p] + 1 > a[q]`, +# keyed by the `λ` at which they become tight. `heap` holds each live node `q` +# that a candidate arc enters, keyed by an upper bound on the keys of those +# arcs; unless `dirty[q]`, the bound is the key of the arc from `best[q]`. +struct DenseParametricTree{T} + win::Matrix{T} + wout::Matrix{T} + live::Vector{Int} + islive::BitVector + a::Vector{Int} + b::Vector{T} + parent::Vector{Int} + fchild::Vector{Int} + lchild::Vector{Int} + nsib::Vector{Int} + psib::Vector{Int} + heap::IndexedHeap{T} + best::Vector{Int} + dirty::BitVector +end + +const ParametricTrees = Union{ParametricTree,DenseParametricTree} + +function _child_append!(S::ParametricTrees, p::Int, c::Int) S.psib[c] = S.lchild[p] S.nsib[c] = 0 S.lchild[p] == 0 ? (S.fchild[p] = c) : (S.nsib[S.lchild[p]] = c) S.lchild[p] = c return nothing end -function _child_remove!(S::ParametricTree, p::Int, c::Int) +function _child_remove!(S::ParametricTrees, p::Int, c::Int) pr, nx = S.psib[c], S.nsib[c] pr == 0 ? (S.fchild[p] = nx) : (S.nsib[pr] = nx) nx == 0 ? (S.lchild[p] = pr) : (S.psib[nx] = pr) return nothing end # Moves the children of `z` to the end of the children of `p`. -function _child_concat!(S::ParametricTree, p::Int, z::Int) +function _child_concat!(S::ParametricTrees, p::Int, z::Int) S.fchild[z] == 0 && return nothing if S.lchild[p] == 0 S.fchild[p] = S.fchild[z] @@ -235,7 +268,7 @@ function _shift_subtree!(S::ParametricTree, sub::Vector{Int}, Δa::Int, Δb, sta end # Appends the descendants of `v` to `sub`, with `stack` as workspace. -function _subtree!(sub::Vector{Int}, S::ParametricTree, v::Int, stack::Vector{Int}) +function _subtree!(sub::Vector{Int}, S::ParametricTrees, v::Int, stack::Vector{Int}) empty!(stack) c = S.fchild[v] while c != 0 @@ -277,35 +310,20 @@ end # largest weight. Each time an arc becomes tight, the work is proportional to # the size of the subtrees whose depths change plus the number of arcs incident # to them: `O(nm + n² log n)` in all, and far less when the tree stays shallow. -function _maxbalance(nv::Int, src::Vector{Int}, dst::Vector{Int}, w::Vector{T}) where {T} +# +# On a dense digraph, with `dense`, the tree is held by `DenseParametricTree`: +# each node whose depth changes then costs a pass over the live nodes rather +# than a heap update for each of its arcs. +function _maxbalance(nv::Int, src::Vector{Int}, dst::Vector{Int}, w::Vector{T}; + dense::Bool=8 * length(w) > nv^2) where {T} x = zeros(T, nv) nv == 1 && return x + dense && return _maxbalance_dense(nv, src, dst, w) ne = length(w) S = ParametricTree(nv, src, dst, w) a, b, tarc, heap = S.a, S.b, S.tarc, S.heap csrc, cdst, cw = S.src, S.dst, S.w - fill!(a, -1) - ptr, adj = _out_arcs(nv, src, e -> S.alive[e]) - queue = Int[1] - a[1] = 0 - i = 1 - while i <= length(queue) - v = queue[i] - i += 1 - for s in ptr[v]:ptr[v+1]-1 - e = adj[s] - q = dst[e] - if a[q] == -1 - a[q] = a[v] + 1 - push!(queue, q) - end - if a[q] == a[v] + 1 && (tarc[q] == 0 || b[v] + w[e] > b[q]) - b[q] = b[v] + w[e] - tarc[q] = e - end - end - end - length(queue) == nv || error("internal error: the digraph is not strongly connected") + _initial_tree!(a, b, tarc, src, dst, w) for q in 2:nv _child_append!(S, src[tarc[q]], q) end @@ -414,9 +432,47 @@ function _maxbalance(nv::Int, src::Vector{Int}, dst::Vector{Int}, w::Vector{T}) nlive -= length(path) end end - # Sum `uΔ` along each vertex's chain of contractions, compressing the - # chains as they are summed. - for z in 1:nv + return _chain_sums!(x, uf, uΔ, stack) +end + +# The shortest-path tree for `λ → ∞` on the strongly connected digraph on +# `1:length(a)` with arcs `src[e] → dst[e]` of weight `w[e]`, rooted at vertex +# 1: fewest arcs, then largest weight. Sets the depths `a`, the path weights +# `b`, and the tree arcs `tarc`, with `tarc[1] == 0`. +function _initial_tree!(a::Vector{Int}, b::Vector, tarc::Vector{Int}, src::Vector{Int}, dst::Vector{Int}, w::Vector) + nv = length(a) + fill!(a, -1) + fill!(b, zero(eltype(b))) + fill!(tarc, 0) + ptr, adj = _out_arcs(nv, src, e -> src[e] != dst[e]) + queue = Int[1] + a[1] = 0 + i = 1 + while i <= length(queue) + v = queue[i] + i += 1 + for s in ptr[v]:ptr[v+1]-1 + e = adj[s] + q = dst[e] + if a[q] == -1 + a[q] = a[v] + 1 + push!(queue, q) + end + if a[q] == a[v] + 1 && (tarc[q] == 0 || b[v] + w[e] > b[q]) + b[q] = b[v] + w[e] + tarc[q] = e + end + end + end + length(queue) == nv || error("internal error: the digraph is not strongly connected") + return nothing +end + +# The shift `x` of each vertex `z`, as the sum of `uΔ` along its chain of +# contractions `z`, `uf[z]`, `uf[uf[z]]`, …, with `x[1] == 0`; the chains are +# compressed as they are summed, with `stack` as workspace. +function _chain_sums!(x::Vector{T}, uf::Vector{Int}, uΔ::Vector{T}, stack::Vector{Int}) where {T} + for z in eachindex(uf) empty!(stack) y = z while uf[y] != y @@ -437,6 +493,210 @@ function _maxbalance(nv::Int, src::Vector{Int}, dst::Vector{Int}, w::Vector{T}) return x end +# The `λ` at which the arc `p → q` of weight `wpq` between live nodes becomes +# tight, as for `_key!`, or `-Inf` if it is not a candidate arc. +function _dense_key(S::DenseParametricTree{T}, p::Int, q::Int, wpq::T) where {T} + c = S.a[p] + 1 - S.a[q] + (c <= 0 || S.parent[q] == p) && return -T(Inf) + return (S.b[p] + wpq - S.b[q]) / c +end + +# Sets `best[q]` to `p` and the key of `q` to `k`, or removes `q` from the heap +# if `p == 0`. +function _dense_set_best!(S::DenseParametricTree, q::Int, p::Int, k) + S.best[q] = p + S.dirty[q] = false + p == 0 ? _heap_remove!(S.heap, q) : _heap_set!(S.heap, q, k) + return nothing +end + +# Sets `best[q]` to the tail of the candidate arc entering `q` with the +# largest key. +function _dense_best!(S::DenseParametricTree{T}, q::Int) where {T} + f, kf = 0, -T(Inf) + for p in S.live + k = _dense_key(S, p, q, S.win[p, q]) + if k > kf + f, kf = p, k + end + end + _dense_set_best!(S, q, f, kf) + return nothing +end + +# Updates `q` after the key of the arc `p → q` of weight `wpq` has changed: the +# arc from `p` becomes `best[q]` if its key exceeds the key of `q`, and `q` +# becomes dirty if `best[q] == p` and the key fell. +function _dense_offer!(S::DenseParametricTree{T}, p::Int, q::Int, wpq::T) where {T} + k = _dense_key(S, p, q, wpq) + h = S.heap + if k > -T(Inf) && (h.pos[q] == 0 || k > h.key[q]) + S.best[q] = p + S.dirty[q] = false + _heap_set!(h, q, k) + elseif S.best[q] == p && k != h.key[q] + S.dirty[q] = true + end + return nothing +end + +# Adds `Δa` to the depth and `Δb` to the path weight of each node `y` of the +# subtree listed in `sub`, marks it with `stamp[y] = st`, and updates `best` at +# the subtree and at the heads of the arcs leaving it. Arcs within the subtree +# keep their keys, since both ends shift alike. +function _dense_shift!(S::DenseParametricTree{T}, sub::Vector{Int}, Δa::Int, Δb, stamp::Vector{Int}, + st::Int) where {T} + for y in sub + S.a[y] += Δa + S.b[y] += Δb + stamp[y] = st + end + for y in sub + f, kf = 0, -T(Inf) + for p in S.live + k = _dense_key(S, p, y, S.win[p, y]) + if k > kf + f, kf = p, k + end + wyp = S.wout[p, y] + stamp[p] == st || wyp == -T(Inf) || _dense_offer!(S, y, p, wyp) + end + _dense_set_best!(S, y, f, kf) + end + return nothing +end + +# `_maxbalance` with the tree held by `DenseParametricTree`. +function _maxbalance_dense(nv::Int, src::Vector{Int}, dst::Vector{Int}, w::Vector{T}) where {T} + all(isfinite, w) || error("internal error: the weights must be finite") + tarc = zeros(Int, nv) + S = DenseParametricTree{T}(fill(-T(Inf), nv, nv), fill(-T(Inf), nv, nv), collect(1:nv), trues(nv), + zeros(Int, nv), zeros(T, nv), zeros(Int, nv), zeros(Int, nv), zeros(Int, nv), + zeros(Int, nv), zeros(Int, nv), IndexedHeap{T}(nv), zeros(Int, nv), falses(nv)) + a, b, parent, win, wout, heap = S.a, S.b, S.parent, S.win, S.wout, S.heap + for e in eachindex(w) + p, q = src[e], dst[e] + if p != q && w[e] > win[p, q] + win[p, q] = wout[q, p] = w[e] + end + end + # The tree arc chosen among parallel arcs has the largest weight, so + # `b[q] == b[parent[q]] + win[parent[q], q]`. + _initial_tree!(a, b, tarc, src, dst, w) + for q in 2:nv + parent[q] = src[tarc[q]] + _child_append!(S, parent[q], q) + end + for q in 1:nv + _dense_best!(S, q) + end + uf = collect(1:nv) + uΔ = zeros(T, nv) + stamp = zeros(Int, nv) + st = 0 + sub = Int[] + stack = Int[] + path = Int[] + while length(S.live) > 1 + isempty(heap.heap) && error("internal error: no arc can become tight on a strongly connected digraph") + v = heap.heap[1] + while S.dirty[v] + _dense_best!(S, v) + isempty(heap.heap) && error("internal error: no arc can become tight on a strongly connected digraph") + v = heap.heap[1] + end + u = S.best[v] + λ = heap.key[v] + # Decide whether `u` is a descendant of `v` as in `_maxbalance`. + cycle = false + z = u + empty!(sub) + empty!(stack) + push!(stack, v) + while true + if z != 0 + z == v && (cycle = true; break) + z = parent[z] + end + isempty(stack) && break + y = pop!(stack) + y == v || push!(sub, y) + y == u && (cycle = true; break) + c = S.fchild[y] + while c != 0 + push!(stack, c) + c = S.nsib[c] + end + end + if !cycle + # The arc `u → v` replaces the tree arc into `v`: the subtree of `v` + # moves under `u` and becomes deeper. + Δa = a[u] + 1 - a[v] + Δb = b[u] + win[u, v] - b[v] + _child_remove!(S, parent[v], v) + _child_append!(S, u, v) + parent[v] = u + push!(sub, v) + _dense_shift!(S, sub, Δa, Δb, stamp, st += 1) + else + # Contract the cycle `v → ⋯ → u → v` into `v`. + empty!(path) + z = u + while z != v + push!(path, z) + z = parent[z] + end + for z in path + uf[z] = v + uΔ[z] = (a[z] - a[v]) * λ - (b[z] - b[v]) + _child_remove!(S, parent[z], z) + S.islive[z] = false + end + filter!(y -> S.islive[y], S.live) + for z in path + # The arcs between `z` and the live nodes join `v`, merged with + # the parallel arcs there; a tree arc keeps its weight. + δ = uΔ[z] + for p in S.live + p == v && continue + wi = win[p, z] + if wi != -T(Inf) && parent[v] != p + wv = wi + δ + wv > win[p, v] && (win[p, v] = wout[v, p] = wv) + end + wo = wout[p, z] + if wo != -T(Inf) && parent[p] != v + wv = wo - δ + (parent[p] == z || wv > wout[p, v]) && (wout[p, v] = win[v, p] = wv) + end + end + # The children of `z` move under `v`, which is shallower, and + # their subtrees rise by `a[z] - a[v]`. + c = S.fchild[z] + while c != 0 + parent[c] = v + c = S.nsib[c] + end + Δa = a[v] - a[z] + _dense_shift!(S, _subtree!(empty!(sub), S, z, stack), Δa, Δa * λ, stamp, st += 1) + _child_concat!(S, v, z) + end + for z in path + _dense_set_best!(S, z, 0, λ) + end + for q in S.live + q == v && continue + f = S.best[q] + f == 0 || S.islive[f] || (S.dirty[q] = true) + wvq = wout[q, v] + wvq == -T(Inf) || _dense_offer!(S, v, q, wvq) + end + _dense_best!(S, v) + end + end + return _chain_sums!(zeros(T, nv), uf, uΔ, stack) +end + # The scales `y` of the difference system `sys` shifted by the max-balancing # shift of each strongly connected component, `scc` and `nscc` as returned by # `_strong_components`. The shift of each component is fixed by making its diff --git a/test/maxbalance.jl b/test/maxbalance.jl index 53c4569..2d655c0 100644 --- a/test/maxbalance.jl +++ b/test/maxbalance.jl @@ -70,7 +70,8 @@ end @testset "_maxbalance" begin # Random strongly connected digraphs with parallel arcs, self loops, and - # (on even trials) many tied weights. + # (on even trials) many tied weights, with the arcs held in lists and in + # matrices. rng = StableRNG(3) for trial in 1:200 nv = rand(rng, 2:9) @@ -79,9 +80,32 @@ end src = vcat(perm, rand(rng, 1:nv, extra)) dst = vcat(perm[mod1.(2:nv+1, nv)], rand(rng, 1:nv, extra)) w = isodd(trial) ? randn(rng, length(src)) : Float64.(rand(rng, -2:2, length(src))) - x = _maxbalance(nv, src, dst, w) - @test x[1] == 0 - @test maxbalanced(nv, src, dst, w, x) + for dense in (false, true) + x = _maxbalance(nv, src, dst, w; dense) + @test x[1] == 0 + @test maxbalanced(nv, src, dst, w, x) + end + end + # Larger digraphs of either density: the two representations agree. + for (nv, extra) in ((60, 120), (60, 3000), (200, 30000)) + perm = randperm(rng, nv) + src = vcat(perm, rand(rng, 1:nv, extra)) + dst = vcat(perm[mod1.(2:nv+1, nv)], rand(rng, 1:nv, extra)) + for w in (randn(rng, length(src)), Float64.(rand(rng, -2:2, length(src)))) + @test _maxbalance(nv, src, dst, w; dense=true) ≈ _maxbalance(nv, src, dst, w; dense=false) atol = 1e-12 + end + end + # Other weight types. + nv = 30 + perm = randperm(rng, nv) + src = vcat(perm, rand(rng, 1:nv, 200)) + dst = vcat(perm[mod1.(2:nv+1, nv)], rand(rng, 1:nv, 200)) + w = randn(rng, length(src)) + x = _maxbalance(nv, src, dst, w) + for dense in (false, true) + xB = _maxbalance(nv, src, dst, big.(w); dense) + @test eltype(xB) == BigFloat + @test xB ≈ x atol = 1e-12 end # A banded digraph, whose shortest-path tree is deep and moves as a whole at # every contraction. It has too many vertices to check every partition, so @@ -90,12 +114,14 @@ end src = vcat(1:nv-1, 2:nv, 1:nv-2) dst = vcat(2:nv, 1:nv-1, 3:nv) w = randn(rng, length(src)) - x = _maxbalance(nv, src, dst, w) - wx = [w[e] - x[src[e]] + x[dst[e]] for e in eachindex(w)] - for v in 1:nv - out = maximum(wx[e] for e in eachindex(w) if src[e] == v) - inn = maximum(wx[e] for e in eachindex(w) if dst[e] == v) - @test out ≈ inn atol = 1e-9 + for dense in (false, true) + x = _maxbalance(nv, src, dst, w; dense) + wx = [w[e] - x[src[e]] + x[dst[e]] for e in eachindex(w)] + for v in 1:nv + out = maximum(wx[e] for e in eachindex(w) if src[e] == v) + inn = maximum(wx[e] for e in eachindex(w) if dst[e] == v) + @test out ≈ inn atol = 1e-9 + end end end