Refine asymmetric starts with conjugate gradients

docs/src/index.md

@@ -156,17 +156,27 @@ julia> round.(extrema(P2 ./ P1); digits=3)
 
 [`cover`](@ref) is exactly covariant when every row and column has the same
 nonzero pattern, including dense matrices without zeros. On irregular sparse
-support it may be only approximately covariant:
+support its geometric-mean start is only approximately covariant; the
+conjugate-gradient refinement (`cgiter`, default 4) moves it toward a
+covariant least-squares fit, reaching it here but not necessarily on larger
+problems:
 
 ```jldoctest covariance
 julia> e = [2.0, 0.3, 5.0]; E = Diagonal(e);
 
-julia> r1, c1 = cover(A); r2, c2 = cover(D * A * E);
+julia> function covariance_spread(cgiter)
+           r1, c1 = cover(A; cgiter)
+           r2, c2 = cover(D * A * E; cgiter)
+           Q1 = (d .* r1) * (e .* c1)'
+           Q2 = r2 * c2'
+           return round.(extrema(Q2 ./ Q1); digits=3)
+       end;
 
-julia> Q1 = (d .* r1) * (e .* c1)'; Q2 = r2 * c2';
-
-julia> round.(extrema(Q2 ./ Q1); digits=3)
+julia> covariance_spread(0)
 (0.872, 1.205)
+
+julia> covariance_spread(4)
+(1.0, 1.0)
 ```
 
 Use [`symcover_min`](@ref) or [`cover_min`](@ref) when exact covariance is

src/dense_heuristic.jl

@@ -253,7 +253,7 @@ end
 # `cover!` over a dense log-magnitude grid, up to but not including the balance
 # convention.
 function _cover_dense!(a::AbstractVector, b::AbstractVector, A::AbstractMatrix,
-                       ::Type{T}, maxiter::Int) where {T}
+                       ::Type{T}, maxiter::Int, cgiter::Int) where {T}
     or = first(axes(A, 1)) - 1
     oc = first(axes(A, 2)) - 1
     m, n = size(A)
@@ -275,6 +275,21 @@ function _cover_dense!(a::AbstractVector, b::AbstractVector, A::AbstractMatrix,
         β[jp] = _uncon_scale(β[jp], nb[jp], halfmu)
         lβ[jp] = log(β[jp])
     end
+    if cgiter > 0
+        function foreach_grid(f)
+            for jp in 1:n, ip in 1:m
+                l = L[ip, jp]
+                isfinite(l) && f(ip, jp, l)   # -Inf marks a zero entry
+            end
+        end
+        _cg_refine!(lα, lβ, foreach_grid, cgiter)
+        for ip in 1:m
+            iszero(na[ip]) || (α[ip] = max(exp(lα[ip]), floatmin(T)); lα[ip] = log(α[ip]))
+        end
+        for jp in 1:n
+            iszero(nb[jp]) || (β[jp] = max(exp(lβ[jp]), floatmin(T)); lβ[jp] = log(β[jp]))
+        end
+    end
 
     nviol = 0
     zmax = zero(T)

src/heuristic_covers.jl

@@ -90,13 +90,16 @@ function _symcover!(a::AbstractVector, A::AbstractMatrix; maxiter::Int=3)
 end
 
 """
-    a, b = cover(ϕ, A; maxiter=3)
-    a, b = cover(A; maxiter=3)
+    a, b = cover(ϕ, A; maxiter=3, cgiter=4)
+    a, b = cover(A; maxiter=3, cgiter=4)
 
 Given a matrix `A`, return vectors `a` and `b` such that
 `a[i] * b[j] >= abs(A[i, j])` for all `i`, `j`. The method initializes from row
-and column geometric means, covers the most-violated entries first, then applies
-`maxiter` tightening iterations.
+and column geometric means, refines that start with up to `cgiter` conjugate-gradient
+iterations, covers the most-violated entries first, then applies `maxiter`
+tightening iterations.
+
+`cgiter=0` disables refinement.
 
 The factors use the per-component balance convention described by
 [`cover_min`](@ref).
@@ -119,6 +122,13 @@ julia> a * b'
  2.16541  2.03444  3.0
  6.0      5.63709  8.31251
 ```
+
+# Extended help
+
+Conjugate-gradient refinement fits `log(a[i]*b[j])` to `log(abs(A[i, j]))`
+in least squares over the nonzero entries. The exact fit is scale-covariant;
+a finite number of iterations need not achieve this. Each iteration makes
+one pass over the support, stopping early when the residual is small.
 """
 cover(ϕ::AbstractCoverPenalty, A::AbstractMatrix; kwargs...) = cover(A; kwargs...)
 
@@ -160,13 +170,15 @@ function cover!(a::AbstractVector, b::AbstractVector, A::AbstractMatrix; kwargs.
     return _cover!(a, b, A; kwargs...)
 end
 
-function _cover!(a::AbstractVector, b::AbstractVector, A::AbstractMatrix; maxiter::Int=3)
+function _cover!(a::AbstractVector, b::AbstractVector, A::AbstractMatrix; maxiter::Int=3, cgiter::Int=4)
+    cgiter >= 0 || throw(ArgumentError("cgiter must be nonnegative, got $cgiter"))
     T = float(promote_type(eltype(a), eltype(b)))
     if _use_dense_grid(A, T)
-        _cover_dense!(a, b, A, T, maxiter)
+        _cover_dense!(a, b, A, T, maxiter, cgiter)
     else
         sup = flat_support(A, T)
         unconstrained_min!(AbsLog{2}(), a, b, sup)
+        cg_refine_start!(a, b, sup, cgiter)
         boost_feasible!(a, b, sup)
         tighten_cover!(a, b, sup; maxiter)
         # Apply the package's balance convention, then restore coverage lost to rounding.
@@ -480,6 +492,103 @@ function unconstrained_min!(::AbsLog{2}, a::AbstractVector, b::AbstractVector, s
     return nza, nzb
 end
 
+# Refine the log scales by CG on the normal equations for
+# ∑_{ij ∈ support} (lα[i] + lβ[j] - log|A_ij|)².
+# `foreach_entries(f)` calls `f(i, j, log|A_ij|)` over the support.
+# Updates preserve the initial gauge and leave unsupported log scales at -Inf.
+function _cg_refine!(lα::AbstractVector{T}, lβ::AbstractVector{T}, foreach_entries::F, maxiter::Int) where {T,F}
+    maxiter > 0 || return lα, lβ
+    axa, axb = eachindex(lα), eachindex(lβ)
+    na = zeros(Int, axa)
+    nb = zeros(Int, axb)
+    fa = fill!(similar(lα, T), zero(T))
+    fb = fill!(similar(lβ, T), zero(T))
+    ra = fill!(similar(lα, T), zero(T))
+    rb = fill!(similar(lβ, T), zero(T))
+    # Accumulate support counts, right-hand sides, and off-diagonal residuals.
+    foreach_entries() do i, j, lv
+        na[i] += 1
+        nb[j] += 1
+        fa[i] += lv
+        fb[j] += lv
+        ra[i] += lv - lβ[j]
+        rb[j] += lv - lα[i]
+    end
+    rr = zero(T)     # squared residual norm
+    rref = zero(T)   # squared right-hand-side norm
+    for i in axa
+        if iszero(na[i])
+            ra[i] = zero(T)
+        else
+            ra[i] -= na[i] * lα[i]
+            rr += ra[i]^2
+            rref += fa[i]^2
+        end
+    end
+    for j in axb
+        if iszero(nb[j])
+            rb[j] = zero(T)
+        else
+            rb[j] -= nb[j] * lβ[j]
+            rr += rb[j]^2
+            rref += fb[j]^2
+        end
+    end
+    tol = eps(T) * rref
+    pa = copy(ra)
+    pb = copy(rb)
+    Apa = similar(ra)
+    Apb = similar(rb)
+    for _ in 1:maxiter
+        rr <= tol && break
+        for i in axa
+            Apa[i] = na[i] * pa[i]
+        end
+        for j in axb
+            Apb[j] = nb[j] * pb[j]
+        end
+        foreach_entries() do i, j, _
+            Apa[i] += pb[j]
+            Apb[j] += pa[i]
+        end
+        pAp = LinearAlgebra.dot(pa, Apa) + LinearAlgebra.dot(pb, Apb)
+        pAp > 0 || break
+        γ = rr / pAp
+        lα .+= γ .* pa
+        lβ .+= γ .* pb
+        ra .-= γ .* Apa
+        rb .-= γ .* Apb
+        rr2 = LinearAlgebra.dot(ra, ra) + LinearAlgebra.dot(rb, rb)
+        δ = rr2 / rr
+        pa .= ra .+ δ .* pa
+        pb .= rb .+ δ .* pb
+        rr = rr2
+    end
+    return lα, lβ
+end
+
+# Apply log-space refinement to `a`, `b`, preserving zeros and clamping underflow.
+function cg_refine_start!(a::AbstractVector, b::AbstractVector, sup::FlatSupport, maxiter::Int)
+    maxiter > 0 || return a, b
+    T = float(promote_type(eltype(a), eltype(b)))
+    lα = map(x -> log(T(x)), a)
+    lβ = map(x -> log(T(x)), b)
+    is, js, lv = sup.is, sup.js, sup.lv
+    function foreach_entries(f)
+        for k in eachindex(is, js, lv)
+            f(is[k], js[k], lv[k])
+        end
+    end
+    _cg_refine!(lα, lβ, foreach_entries, maxiter)
+    for i in eachindex(a)
+        iszero(a[i]) || (a[i] = max(exp(lα[i]), floatmin(T)))
+    end
+    for j in eachindex(b)
+        iszero(b[j]) || (b[j] = max(exp(lβ[j]), floatmin(T)))
+    end
+    return a, b
+end
+
 # Feasible cover starting from the diagonal alone, resolved by
 # `boost_feasible_seq!`. Unlike `boost_feasible!`, a zero entry of `a` going
 # into that call means "not yet resolved", not "permanently unsupported" —

test/heuristic_covers.jl

@@ -284,3 +284,58 @@ end
     α = (n * I + ones(n, n)) \ vec(sum(L, dims=2))
     @test start(A) ≈ exp.(α) rtol=1e-10
 end
+
+@testset "cover: conjugate-gradient refinement of the start" begin
+    # Centered log-product deviation under row and column scaling.
+    function covdev(coverfn, A, dr, dc)
+        a, b = coverfn(A)
+        aD, bD = coverfn(dr .* A .* dc')
+        d = log.((aD .* bD') ./ ((dr .* a) .* (dc .* b)'))[A .!= 0]
+        return maximum(abs, d .- sum(d) / length(d))
+    end
+    rng = StableRNG(20260905)
+    # Refinement reduces scaling dependence on banded support.
+    n = 12
+    A = zeros(n, n)
+    for i in 1:n, j in max(1, i - 1):min(n, i + 1)
+        A[i, j] = exp(2 * randn(rng))
+    end
+    dr, dc = exp.(4 .* randn(rng, n)), exp.(4 .* randn(rng, n))
+    dev0 = covdev(A -> cover(A; cgiter=0), A, dr, dc)
+    dev4 = covdev(A -> cover(A), A, dr, dc)
+    devx = covdev(A -> cover(A; cgiter=4n), A, dr, dc)
+    @test dev0 > 0.1
+    @test dev4 < dev0 / 3
+    @test devx < 1e-6
+    @test covaries(A -> cover(A; cgiter=4n), A, dr, dc; rtol=1e-6)
+    for k in (0, 4, 4n)
+        a, b = cover(A; cgiter=k)
+        @test iscover(a, b, A; rtol=8eps())
+    end
+    @test_throws ArgumentError cover(A; cgiter=-1)
+
+    # The dense-grid and flattened-support paths refine identically.
+    m = 2 * MatrixCovers.DENSE_GRID_MIN
+    B = zeros(m, m)
+    for i in 1:m, j in max(1, i - 2):min(m, i + 1)
+        B[i, j] = exp(2 * randn(rng)) * (rand(rng) < 0.7)
+    end
+    for i in 1:m
+        B[i, i] = exp(2 * randn(rng))
+    end
+    for k in (0, 4, 40)
+        ad, bd = cover(B; cgiter=k)
+        as, bs = cover(sparse(B); cgiter=k)
+        @test ad .* bd' ≈ as .* bs' rtol = 1e-10
+    end
+
+    # Fully populated support needs no refinement.
+    C = exp.(2 .* randn(rng, 9, 7))
+    @test cover(C; cgiter=4) == cover(C; cgiter=0)
+    @test cover(sparse(C); cgiter=4) == cover(sparse(C); cgiter=0)
+
+    # Empty rows and columns keep zero scales through the refinement.
+    Z = [0.0 2.0 0.0; 0.0 0.0 0.0; 1.0 0.0 0.0]
+    a, b = cover(Z)
+    @test a[2] == 0 && b[3] == 0 && iscover(a, b, Z; rtol=8eps())
+end