diff --git a/src/BurerMonteiro/model.jl b/src/BurerMonteiro/model.jl index 5139831..de4426e 100644 --- a/src/BurerMonteiro/model.jl +++ b/src/BurerMonteiro/model.jl @@ -165,7 +165,11 @@ function buffer_for_jtprod( dim::Dimensions, i::LRO.MatrixIndex, ) where {T} - row = view(model.A, i.value, :) + # Buffer selection runs only at setup; inspect every active component. + row = AbstractMatrix{T}[] + LRO._foreach_column(model.A, i.value) do values, slot + return append!(row, view(values, :, slot)) + end C = model.C[i.value] if any(A -> A isa _LowRank, row) || C isa _LowRank ncols = maximum(row; init = 0) do A @@ -191,6 +195,9 @@ end _buffer(_, ::AbstractMatrix, _) = nothing _buffer(buffer::AbstractVector, ::_RankOne, ::AbstractMatrix) = buffer +function _buffer(buffer::AbstractMatrix, ::_RankOne, ::AbstractMatrix) + return view(buffer, :, 1) +end function _buffer(buffer::AbstractMatrix, A::_LowRank, ::AbstractMatrix) # FIXME with this if-else, we return a small Union but the compiler # since to forget about this Union and allocates later @@ -211,11 +218,27 @@ function add_jtprod!( i::LRO.MatrixIndex, α = 2, ) + LRO._foreach_column(model.model.A, i.value) do values, slot + return _add_column_jtprod!( + values, + slot, + model.jtprod_buffer[i.value], + X, + y, + JtV, + α, + ) + end + return nothing +end + +function _add_column_jtprod!(values, slot, workspace, X, y, JtV, α) for j in eachindex(y) - A = LRO.jac(model.model, j, i) - buffer = _buffer(model.jtprod_buffer[i.value], A, X.factor) + A = values[j, slot] + buffer = _buffer(workspace, A, X.factor) LRO.buffered_mul!(JtV.factor, A, X.factor, α * y[j], true, buffer) end + return nothing end function jtprod!( diff --git a/src/LowRankOpt.jl b/src/LowRankOpt.jl index 07ae9ed..ea55cd7 100644 --- a/src/LowRankOpt.jl +++ b/src/LowRankOpt.jl @@ -19,6 +19,7 @@ include("Test/Test.jl") include("Bridges/Bridges.jl") include("solution.jl") +include("block_columns.jl") include("model.jl") include("buffer.jl") include("schur.jl") diff --git a/src/MOI_wrapper.jl b/src/MOI_wrapper.jl index 3f5e309..507e8bd 100644 --- a/src/MOI_wrapper.jl +++ b/src/MOI_wrapper.jl @@ -131,10 +131,14 @@ end const SOLVER_OPTIONS = ["solver", "sub_solver", "ranks", "square_scalars"] +# Options consumed while building the `Model` in `copy_to`; they must not be +# forwarded to the solver constructor nor to `solve!`. +const MODEL_OPTIONS = ["detect_rank_one"] + function MOI.optimize!(model::Optimizer) options = Dict{Symbol,Any}( - Symbol(key) => model.options[key] for - key in keys(model.options) if !(key in SOLVER_OPTIONS) + Symbol(key) => model.options[key] for key in keys(model.options) if + !(key in SOLVER_OPTIONS) && !(key in MODEL_OPTIONS) ) if model.silent options[:verbose] = 0 @@ -318,9 +322,16 @@ function MOI.copy_to( end b = dest.max_sense ? b0 : -b0 + A_con = _instantiate(A[:, 2:end]) + if get(dest.options, "detect_rank_one", false) + rank_one_A = detect_rank_one(A_con) + if !isnothing(rank_one_A) + A_con = rank_one_A + end + end dest.model = Model( _instantiate(A[:, 1]), - _instantiate(A[:, 2:end]), + A_con, b, convert( SparseArrays.SparseVector{T,Int64}, diff --git a/src/block_columns.jl b/src/block_columns.jl new file mode 100644 index 0000000..d9a6ae2 --- /dev/null +++ b/src/block_columns.jl @@ -0,0 +1,137 @@ +# Copyright (c) 2026: Benoît Legat and contributors +# +# Use of this source code is governed by an MIT-style license that can be found +# in the LICENSE.md file or at https://opensource.org/licenses/MIT. + +""" + BlockColumns(block_ids, values, nblocks) + +Store a component of the constraint matrices on selected PSD blocks. +`values[j, k]` is the matrix for constraint `j` on PSD block `block_ids[k]`. +The block identifiers must be strictly increasing. Other blocks have no +contribution. Missing blocks are skipped, without requiring `zero(M)` for +the matrix element type `M`. + +Constraints are rows and stored PSD blocks are columns, so each block's +constraint matrices are contiguous. This container stores whole columns; +it does not represent sparsity within a column. + +For example, with `n` equality constraints and three PSD blocks, +`BlockColumns([1, 3], values, 3)` stores an `n × 2` matrix of coefficients. +It contributes to blocks 1 and 3 and contributes nothing to block 2. +A model can sum this with another component by passing `(component1, +component2)` as its constraint data. +""" +struct BlockColumns{M<:AbstractMatrix} + block_ids::Vector{Int} + values::Matrix{M} + nblocks::Int + function BlockColumns( + block_ids, + values::Matrix{M}, + nblocks::Integer, + ) where {M<:AbstractMatrix} + ids = Int.(block_ids) + length(ids) == size(values, 2) || throw( + DimensionMismatch( + "one block identifier is required per stored column", + ), + ) + all( + i -> 1 <= ids[i] <= nblocks && (i == 1 || ids[i-1] < ids[i]), + eachindex(ids), + ) || throw( + ArgumentError( + "block identifiers must be strictly increasing and within 1:nblocks", + ), + ) + return new{M}(ids, values, Int(nblocks)) + end +end + +Base.size(A::BlockColumns) = (size(A.values, 1), A.nblocks) +Base.size(A::BlockColumns, d::Integer) = d == 2 ? A.nblocks : size(A.values, d) +function Base.:(==)(A::BlockColumns, B::BlockColumns) + return A.nblocks == B.nblocks && + A.block_ids == B.block_ids && + A.values == B.values +end + +# Legacy input has PSD blocks as rows and constraints as columns. +function BlockColumns(A::Matrix{<:AbstractMatrix}) + return BlockColumns(1:size(A, 1), permutedims(A), size(A, 1)) +end + +function _block_slot(A::BlockColumns, i::Integer) + k = searchsortedfirst(A.block_ids, i) + return k <= length(A.block_ids) && A.block_ids[k] == i ? k : 0 +end + +function _block_column(A::BlockColumns, i::Integer) + k = _block_slot(A, i) + return iszero(k) ? nothing : view(A.values, :, k) +end + +_as_component(A::BlockColumns) = A +_as_component(A::Matrix{<:AbstractMatrix}) = BlockColumns(A) + +@inline _foreach_column(f, ::Tuple{}, i) = nothing +@inline function _foreach_column(f, components::Tuple, i) + component = first(components) + slot = _block_slot(component, i) + if !iszero(slot) + f(component.values, slot) + end + _foreach_column(f, Base.tail(components), i) + return nothing +end + +@inline _sum_columns(f, ::Tuple{}, i, init) = init +@inline function _sum_columns(f, components::Tuple, i, init) + column = _block_column(first(components), i) + value = isnothing(column) ? init : init + f(column) + return _sum_columns(f, Base.tail(components), i, value) +end + +# Construction runs once. Split a heterogeneous legacy array into +# components, omitting blocks that contribute nothing to that component. +# Within a stored block, the entries belonging to other components are +# `Zeros`, so the element type is the small `Union{M,_ZerosMatrix{T}}` only +# when such entries exist. This avoids storing one zero matrix per entry. +const _ZerosMatrix{T} = + FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}} + +_matrix_iszero(A::AbstractMatrix) = iszero(A) +function _matrix_iszero(A::Factorization) + return all(iszero, A.factor) || all(iszero, A.scaling) +end + +function _typed_component(::Type{M}, A) where {M<:AbstractMatrix} + ids = Int[] + for i in axes(A, 1) + if any(j -> A[i, j] isa M && !_matrix_iszero(A[i, j]), axes(A, 2)) + push!(ids, i) + end + end + T = eltype(M) + E = + all(i -> all(j -> A[i, j] isa M, axes(A, 2)), ids) ? M : + Union{M,_ZerosMatrix{T}} + values = Matrix{E}(undef, size(A, 2), length(ids)) + for (k, i) in enumerate(ids) + d = LinearAlgebra.checksquare( + A[i, findfirst(j -> A[i, j] isa M, axes(A, 2))], + ) + for j in axes(A, 2) + values[j, k] = + A[i, j] isa M ? A[i, j] : FillArrays.Zeros{T}(Int(d), Int(d)) + end + end + return BlockColumns(ids, values, size(A, 1)) +end + +function _matrix_components(A::AbstractMatrix) + types = + unique(typeof(value) for value in A if !(value isa FillArrays.Zeros)) + return Tuple(_typed_component(M, A) for M in types) +end diff --git a/src/buffer.jl b/src/buffer.jl index 0b029de..11bc998 100644 --- a/src/buffer.jl +++ b/src/buffer.jl @@ -1,7 +1,7 @@ mutable struct BufferedModelForSchur{ T, C<:AbstractMatrix{T}, - A<:AbstractMatrix{T}, + A<:Tuple, JB, JTB, SB, @@ -118,23 +118,31 @@ Optimization Methods and Software 39.6 (2024): 1185-1215. semidefinite programming." Mathematical Programming 79 (1997): 235-253. """ function buffer_for_jprod(model::Model{T}, i::MatrixIndex) where {T} - nnz = sum(1:(model.meta.ncon); init = 0) do j - return _nnz(model.A[i.value, j]) + return _jacobian_buffer( + model.A, + i.value, + side_dimension(model, i), + model.meta.ncon, + T, + ) +end + +function _jacobian_buffer(components, i, d, ncon, ::Type{T}) where {T} + nnz = _sum_columns(components, i, 0) do column + return sum(_nnz, column; init = 0) end I = zeros(Int64, nnz) J = zeros(Int64, nnz) V = zeros(T, nnz) offset = 0 - for j in 1:(model.meta.ncon) - offset = _add_vec!(I, J, V, j, offset, model.A[i.value, j]) + # Components contribute to the same sparse columns; `sparse` adds + # duplicate entries, including overlaps between representations. + _foreach_column(components, i) do values, slot + for j in axes(values, 1) + offset = _add_vec!(I, J, V, j, offset, values[j, slot]) + end end - A = SparseArrays.sparse( - I, - J, - V, - side_dimension(model, i)^2, - model.meta.ncon, - ) + A = SparseArrays.sparse(I, J, V, d^2, ncon) return A end @@ -299,7 +307,9 @@ function buffer_for_jtprod( ncon = model.meta.ncon # If every `Aᵢⱼ` is zero then so is the product, and callers rely on # getting a `Zeros` back (`_sub` then aliases `C` instead of copying). - nnz_A = sum(j -> _nnz(model.A[mat_idx.value, j]), 1:ncon; init = 0) + nnz_A = _sum_columns(model.A, mat_idx.value, 0) do column + return sum(_nnz, column; init = 0) + end if iszero(nnz_A) return FillArrays.Zeros{T}(d, d) end @@ -309,12 +319,24 @@ function buffer_for_jtprod( max_nnz = min(density * d^2, (work_ratio * float(d)^3 - float(ncon) * d) / nnz_A) pattern = FillArrays.Zeros{T}(d, d) - for j in 1:ncon - pattern = _merge_sparsity(pattern, _abs(model.A[mat_idx.value, j])) - if _nnz(pattern) > max_nnz - return zeros(T, d, d) + dense = false + _foreach_column(model.A, mat_idx.value) do values, slot + for j in axes(values, 1) + A = values[j, slot] + if !(A isa Union{SparseArrays.SparseMatrixCSC,FillArrays.Zeros}) + dense = true + break + end + pattern = _merge_sparsity(pattern, _abs(A)) + if _nnz(pattern) > max_nnz + dense = true + break + end end end + if dense + return zeros(T, d, d) + end return pattern end @@ -397,12 +419,34 @@ function _jtprod!( i::MatrixIndex, ) _zero!(buffer) - for j in eachindex(y) - _add_mul!(buffer, model.model.A[i.value, j], y[j]) + _foreach_column(model.model.A, i.value) do values, slot + return _add_sparse_jtprod!(buffer, values, slot, y) end return buffer end +function _add_sparse_jtprod!(buffer, values, slot, y) + for j in eachindex(y) + _add_mul!(buffer, values[j, slot], y[j]) + end + return nothing +end + +# Setup fallback for arbitrary matrix representations. Specializations for +# sparse and rank-one matrices avoid visiting their implicit zeros. +function _add_vec!(I, J, V, j, offset, A::AbstractMatrix) + for col in axes(A, 2), row in axes(A, 1) + value = A[row, col] + if !iszero(value) + offset += 1 + I[offset] = row + (col - 1) * size(A, 1) + J[offset] = j + V[offset] = value + end + end + return offset +end + """ unsafe_jtprod(model::BufferedModelForSchur, y, i::MatrixIndex) diff --git a/src/factorization.jl b/src/factorization.jl index d587ff0..caf5751 100644 --- a/src/factorization.jl +++ b/src/factorization.jl @@ -155,6 +155,99 @@ function positive_semidefinite_factorization( return Factorization(factor, Ones{T}(Base.OneTo(size(factor, 2)))) end +# A uniform rank-one type lets detected coefficients share a concrete +# BlockColumns component without widening another component's element type. +const _SparseRankOne{T} = + Factorization{T,SparseArrays.SparseVector{T,Int64},Array{T,0}} + +function _sparse_rank_one(factor::SparseArrays.SparseVector{T}, s::T) where {T} + return Factorization( + SparseArrays.SparseVector{T,Int64}(factor), + s, + )::_SparseRankOne{T} +end + +""" + rank_one(A::AbstractMatrix{T}; atol = 0, rtol = √eps) + +If the symmetric matrix `A` is rank one (or zero), return an equivalent +[`Factorization`](@ref) `s * b * b'`; otherwise return `nothing`. The zero +matrix maps to a zero-scaling factorization so that an array mixing zeros and +rank-one matrices can still be converted uniformly (see [`detect_rank_one`](@ref)). + +No eigen/SVD is used: if `A = s b b'` then for any `p` with `A[p,p] ≠ 0`, +column `p` is `A[:,p] = (s b_p) b`, so `b ∝ A[:,p]` and +`A = A[:,p] A[:,p]' / A[p,p]`. We pick such a pivot `p` and verify the identity. +""" +function rank_one end + +function rank_one( + A::SparseArrays.SparseMatrixCSC{T}; + atol = zero(real(T)), + rtol = sqrt(eps(real(float(T)))), +) where {T} + n = LinearAlgebra.checksquare(A) + if iszero(A) + return _sparse_rank_one(SparseArrays.spzeros(T, n), zero(T)) + end + p = findfirst(j -> !iszero(A[j, j]), 1:n) + isnothing(p) && return nothing # nonzero but zero diagonal ⇒ not rank one + u = A[:, p] + s = inv(A[p, p]) + ui, uv = SparseArrays.findnz(u) + # A rank-one `s u u'` has exactly `length(ui)^2` nonzeros, all on `ui × ui`. + SparseArrays.nnz(A) == length(ui)^2 || return nothing + for a in eachindex(ui), b in eachindex(ui) + if !isapprox(A[ui[a], ui[b]], s * uv[a] * uv[b]; atol, rtol) + return nothing + end + end + return _sparse_rank_one(u, s) +end + +function rank_one(A::FillArrays.Zeros{T}; kws...) where {T} + n = LinearAlgebra.checksquare(A) + return _sparse_rank_one(SparseArrays.spzeros(T, n), zero(T)) +end + +# Already rank one (vector factor): just coerce to the uniform sparse type. +function rank_one( + A::Factorization{T,<:AbstractVector{T},<:AbstractArray{T,0}}; + kws..., +) where {T} + return _sparse_rank_one(SparseArrays.sparse(A.factor), only(A.scaling)) +end + +rank_one(::AbstractMatrix; kws...) = nothing + +""" + detect_rank_one(A::AbstractMatrix{<:AbstractMatrix}; kws...) + +Convert entries recognized by [`rank_one`](@ref), keeping other matrix +representations. Return a tuple of [`BlockColumns`](@ref) components grouped +by concrete matrix type. A PSD block missing from one component contributes +zero there; components may overlap and their contributions are summed. +Return `nothing` if no entries are convertible or the array is empty. +""" +function detect_rank_one( + A::AbstractMatrix{<:AbstractMatrix{T}}; + kws..., +) where {T} + isempty(A) && return nothing + changed = false + converted = map(A) do coefficient + # Keep `Zeros` so that `_matrix_components` stores them as such + coefficient isa FillArrays.Zeros && return coefficient + factor = rank_one(coefficient; kws...) + if isnothing(factor) + return coefficient + end + changed = true + return factor + end + return changed ? _matrix_components(converted) : nothing +end + struct AsymmetricFactorization{ T, F<:Union{AbstractVector{T},AbstractMatrix{T}}, @@ -302,6 +395,13 @@ function _abs2!!(a::AbstractArray) return a end +# Rank-one coefficients detected from ordinary matrices carry a 0-D Array +# scaling, unlike the unit FillArrays scalings used by PSD factorizations. +_lmul_diag!!(s::Array{<:Any,0}, value::Number) = s[] * value +_lmul_diag!!(s::Array{<:Any,0}, values) = LinearAlgebra.lmul!(s[], values) +_rmul_diag!!(value::Number, s::Array{<:Any,0}) = value * s[] +_rmul_diag!!(values, s::Array{<:Any,0}) = LinearAlgebra.rmul!(values, s[]) + _lmul_diag!!(::FillArrays.Ones, VtU) = VtU _rmul_diag!!(VtU, ::FillArrays.Ones) = VtU @@ -362,8 +462,11 @@ function _dot_mat_fact( res = zero(MA.promote_sum_mul(T, U)) for i in axes(left_factor(b), 2) res += - b.scaling[i] * - LinearAlgebra.dot(left_factor(b)[:, i], a, right_factor(b)[:, i]) + b.scaling[i] * LinearAlgebra.dot( + view(left_factor(b), :, i), + a, + view(right_factor(b), :, i), + ) end return res end diff --git a/src/model.jl b/src/model.jl index 5a7f330..465c0cd 100644 --- a/src/model.jl +++ b/src/model.jl @@ -59,14 +59,17 @@ The fields of the `struct` as related to the arrays of the above formulation as * The ``i``th PSD constraint is of size `msize[i] × msisze[i]` * The matrix ``C_i`` is given by `C[i]`. -* The matrix ``A_{i,j}`` is given by `A[i,j]`. +* `A` is a tuple of [`BlockColumns`](@ref) components. The matrix + ``A_{i,j}`` is the sum of their entries for PSD block `i` and constraint + `j`. A dense outer matrix in the original block-by-constraint layout is + accepted and split into homogeneous components. Tuple entries can also + use dense outer matrices in that layout. """ -mutable struct Model{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T}} <: - AbstractModel{T} +mutable struct Model{T,C<:AbstractMatrix{T},A<:Tuple} <: AbstractModel{T} meta::NLPModels.NLPModelMeta{T,Vector{T}} dim::Dimensions C::Vector{C} - A::Matrix{A} + A::A b::Vector{T} d_lin::SparseArrays.SparseVector{T,Int64} C_lin::SparseArrays.SparseMatrixCSC{T,Int64} @@ -74,15 +77,38 @@ mutable struct Model{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T}} <: function Model( C::Vector{CT}, - A::Matrix{AT}, + A::Tuple, b::Vector{T}, d_lin::SparseArrays.SparseVector{T,Int64}, C_lin::SparseArrays.SparseMatrixCSC{T,Int64}, msizes::Vector{Int64}, - ) where {T,CT<:AbstractMatrix{T},AT<:AbstractMatrix{T}} - model = new{T,CT,AT}() + ) where {T,CT<:AbstractMatrix{T}} + components = map(_as_component, A) + for component in components + size(component) == (length(b), length(C)) || throw( + DimensionMismatch( + "constraint components must have one row per constraint and one logical column per PSD block", + ), + ) + for (k, i) in enumerate(component.block_ids), + j in axes(component.values, 1) + + coefficient = component.values[j, k] + coefficient isa AbstractMatrix{T} || throw( + ArgumentError( + "constraint coefficients must have element type $T", + ), + ) + size(coefficient) == (msizes[i], msizes[i]) || throw( + DimensionMismatch( + "constraint matrix has the wrong size for PSD block $i", + ), + ) + end + end + model = new{T,CT,typeof(components)}() model.C = C - model.A = A + model.A = components model.b = b model.d_lin = d_lin model.C_lin = C_lin @@ -98,6 +124,11 @@ mutable struct Model{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T}} <: end end +# Preserve the existing constructor for one dense outer matrix of matrices. +function Model(C, A::Matrix{<:AbstractMatrix}, b, d_lin, C_lin, msizes) + return Model(C, _matrix_components(A), b, d_lin, C_lin, msizes) +end + function NLPModels.unconstrained(model::Model) return iszero(model.meta.ncon) end @@ -182,16 +213,41 @@ cons_constant(model::Model) = model.b # Should be only used with `norm` jac(model::Model, ::Type{ScalarIndex}) = model.C_lin function jac(model::Model, j::Integer, i::MatrixIndex) - return model.A[i.value, j] + return _coefficient_sum( + model.A, + j, + i.value, + side_dimension(model, i), + eltype(model.b), + ) +end + +function _coefficient_sum(::Tuple{}, j, i, d, ::Type{T}) where {T} + return FillArrays.Zeros{T}(d, d) +end +@inline function _coefficient_sum( + components::Tuple, + j, + i, + d, + ::Type{T}, +) where {T} + column = _block_column(first(components), i) + tail = _coefficient_sum(Base.tail(components), j, i, d, T) + if isnothing(column) + return tail + elseif tail isa FillArrays.Zeros + return column[j] + end + return column[j] + tail end + function jac(model::Model, j::Integer, ::Type{ScalarIndex}) return model.C_lin[j, :] end function norm_jac(model::Model{T}, i::MatrixIndex) where {T} - if isempty(model.A) - return zero(T) - end - return LinearAlgebra.norm(model.A[i.value, :]) + iszero(model.meta.ncon) && return zero(T) + return LinearAlgebra.norm((jac(model, j, i) for j in 1:model.meta.ncon)) end ####################### @@ -204,9 +260,40 @@ function add_jprod!( Jv::AbstractVector, i::MatrixIndex, ) - for j in 1:(model.meta.ncon) - Jv[j] += LinearAlgebra.dot(model.A[i.value, j], V) + _foreach_column(model.A, i.value) do values, slot + return _add_column_jprod!(values, slot, V, Jv) + end + return Jv +end + +function _coefficient_dot(A, V) + return LinearAlgebra.dot(A, V) +end +# `dot(Diagonal, Matrix)` copies the matrix diagonal on supported Julia +# versions. Contract the entries directly to keep the generic Schur path +# allocation-free for diagonal coefficients too. +function _coefficient_dot(A::LinearAlgebra.Diagonal, V::AbstractMatrix) + return sum( + eachindex(A.diag); + init = zero(promote_type(eltype(A), eltype(V))), + ) do k + return conj(A.diag[k]) * V[k, k] + end +end +function _coefficient_dot(::FillArrays.Zeros{T}, V) where {T} + return zero(promote_type(T, eltype(V))) +end +function _coefficient_dot(A::AbstractFactorization, V::AbstractMatrix) + return LinearAlgebra.dot(V, A) +end +function _coefficient_dot(A::AbstractFactorization, V::AbstractFactorization) + return LinearAlgebra.dot(A, V) +end +function _add_column_jprod!(values, slot, V, Jv) + for j in eachindex(Jv) + Jv[j] += _coefficient_dot(values[j, slot], V) end + return nothing end function jprod!( diff --git a/src/schur.jl b/src/schur.jl index 39529b8..a66a2e4 100644 --- a/src/schur.jl +++ b/src/schur.jl @@ -48,27 +48,6 @@ _nnz(A::SparseArrays.SparseMatrixCSC) = SparseArrays.nnz(A) # The `jprod!` buffer is guaranteed to be the first argument of the tuple. # This assumption is used by Loraine. -function buffer_for_schur_complement(model::Model{T}, κ) where {T} - n = model.meta.ncon - σ = zeros(Int64, n, num_matrices(model)) - last_dense = zeros(Int64, num_matrices(model)) - - for mat_idx in matrix_indices(model) - i = mat_idx.value - nzA = [_nnz(model.A[i, j]) for j in 1:n] - σ[:, i] = sortperm(nzA, rev = true) - sorted = nzA[σ[:, i]] - - # Last index for which nnz > κ - last_dense[i] = something(findlast(Base.Fix1(isless, κ), sorted), 0) - end - - AW = [zeros(T, dim, dim) # /!\ it's the same zero everywhere, might be an issue with BigFloat - for dim in model.msizes] - WAW = copy.(AW) - - return AW, WAW, σ, last_dense -end function add_schur_complement!( model::BufferedModelForSchur, @@ -82,8 +61,8 @@ function add_schur_complement!( return H end -# Adds the contribution `𝐀ᵢᵀ (W ⊗ W) 𝐀ᵢ` of the PSD block `mat_idx` to `H`, -# where `𝐀ᵢ` is the buffer documented in `buffer_for_jprod`. +# Adds one sparse component's contribution `𝐀ᵢᵀ (W ⊗ W) 𝐀ᵢ` to `H`. +# Its Jacobian cache follows the layout documented in `buffer_for_jprod`. # `H` is assembled column by column. For a constraint `i` in the dense # regime, that is, one of the `last_dense[ilmi]` first entries of the # permutation `σ` sorting the constraints by decreasing number of nonzeros, @@ -93,34 +72,33 @@ end # `add_sub_jprod!`. The constraints with at most one nonzero use the # dedicated low-nnz paths below, also following [FKN97]. # /!\ W needs to be symmetric -function add_schur_complement!( - model::BufferedModelForSchur, - mat_idx::MatrixIndex, - W::AbstractMatrix{T}, - H, -) where {T} - AW, WAW, σ, last_dense = model.schur_buffer - ilmi = mat_idx.value - n = model.meta.ncon +function _add_sparse_component!(workspace, slot, W, AW, WAW, H) + σ, last_dense = workspace.permutation, workspace.last_dense + column = view(workspace.component.values, :, slot) + n = length(column) for ii in axes(H, 1) - i = σ[ii, ilmi] - Ai = model.model.A[ilmi, i] + i = σ[ii, slot] + Ai = column[i] if _nnz(Ai) > 0 - if ii <= last_dense[ilmi] - LinearAlgebra.mul!(AW[ilmi], W, Ai) - LinearAlgebra.mul!(WAW[ilmi], AW[ilmi], W) - I = view(σ, ii:n, ilmi) - add_sub_jprod!(model, mat_idx, WAW[ilmi], view(H, I, i), I) + if ii <= last_dense[slot] + LinearAlgebra.mul!(AW, W, Ai) + LinearAlgebra.mul!(WAW, AW, W) + I = view(σ, ii:n, slot) + _add_jprod!( + WAW, + view(H, I, i), + view(workspace.jacobian[slot], :, I), + ) for jj in (ii+1):n - j = σ[jj, ilmi] + j = σ[jj, slot] H[i, j] = H[j, i] end else if _nnz(Ai) > 1 @inbounds for jj in ii:n - j = σ[jj, ilmi] - Aj = model.model.A[ilmi, j] + j = σ[jj, slot] + Aj = column[j] if !iszero(_nnz(Aj)) ttt = _dot(Ai, Aj, W) H[i, j] += ttt @@ -134,8 +112,8 @@ function add_schur_complement!( iiiiAi = jjjiAi = only(SparseArrays.rowvals(Ai)) vvvi = only(SparseArrays.nonzeros(Ai)) @inbounds for jj in ii:n - j = σ[jj, ilmi] - Ajjj = model.model.A[ilmi, j] + j = σ[jj, slot] + Ajjj = column[j] # As we sort the matrices in decreasing `nnz` order, # the rest of matrices is either zero or have only # one entry @@ -184,6 +162,16 @@ function add_schur_complement!( return H end +function _add_sparse_eval_column!(result, values, slot, Ay, W) + for k in eachindex(result) + A = values[k, slot] + if !iszero(_nnz(A)) + result[k] += _dot(A, Ay, W) + end + end + return nothing +end + # [HKS24, (5b)] # Returns the matrix equal to the sum, for each equation, of # ⟨A_i, WA_jW⟩ @@ -214,7 +202,7 @@ function _add_eval_schur_complement!( result, i::MatrixIndex, ) - AW, WAW = model.schur_buffer + AW, WAW, _, _ = model.schur_buffer LinearAlgebra.mul!(AW[i.value], W, Ay) LinearAlgebra.mul!(WAW[i.value], AW[i.value], W) return add_jprod!(model, WAW[i.value], result, i) @@ -231,11 +219,8 @@ function _add_eval_schur_complement!( result, i::MatrixIndex, ) - for k in eachindex(result) - A = model.model.A[i.value, k] - if !iszero(_nnz(A)) - result[k] += _dot(A, Ay, W) - end + _foreach_column(model.model.A, i.value) do values, slot + return _add_sparse_eval_column!(result, values, slot, Ay, W) end return result end @@ -245,22 +230,400 @@ end # ⟨A_i, WA(y)W⟩ function eval_schur_complement!(model::BufferedModelForSchur, W, y, result) fill!(result, zero(eltype(result))) + components = model.schur_buffer[3] for i in matrix_indices(model) - _add_eval_schur_complement!( - model, - W[i], - unsafe_jtprod(model, y, i), - result, - i, + if _has_general_component(components, i.value) + _add_eval_schur_complement!( + model, + W[i], + unsafe_jtprod(model, y, i), + result, + i, + ) + else + _prepare_components!(components, i.value, W[i]) + _add_component_matvec!(components, i.value, y, result) + end + end + C = model.model.C_lin + for k in axes(C, 2) + column = view(C, :, k) + weight = W[ScalarIndex][k] * LinearAlgebra.dot(column, y) + LinearAlgebra.axpy!(weight, column, result) + end + return result +end + +# Schur workspaces follow individual components, not the type of the whole +# model. A rank-one component can coexist with arbitrary other components. +const _RankOneFactorization{T} = + Factorization{T,<:AbstractVector{T},<:AbstractArray{T,0}} + +_nnz(A::AbstractMatrix) = count(!iszero, A) + +struct _GeneralSchurComponent{C,J} + component::C + jacobian::J + permutation::Matrix{Int} + last_dense::Vector{Int} + sparse::Vector{Bool} +end + +struct _RankOneSchurComponent{T,C} + component::C + U::Vector{Matrix{T}} + weights::Matrix{T} + UW::Vector{Matrix{T}} + gram::Matrix{T} +end + +function _schur_component(component::BlockColumns, model::Model{T}, κ) where {T} + n, p = size(component.values) + permutation = Matrix{Int}(undef, n, p) + last_dense = zeros(Int, p) + sparse = Vector{Bool}(undef, p) + jacobian = SparseArrays.SparseMatrixCSC{T,Int64}[] + for (slot, i) in enumerate(component.block_ids) + column = view(component.values, :, slot) + counts = _nnz.(column) + permutation[:, slot] = sortperm(counts; rev = true) + last_dense[slot] = count(>(κ), counts) + sparse[slot] = all( + A -> A isa Union{SparseArrays.SparseMatrixCSC,FillArrays.Zeros}, + column, ) + push!( + jacobian, + _jacobian_buffer((component,), i, model.msizes[i], n, T), + ) + end + return _GeneralSchurComponent( + component, + jacobian, + permutation, + last_dense, + sparse, + ) +end + +# `Zeros` entries get a zero row in `U` and a zero weight. +function _schur_component( + component::BlockColumns{<:Union{_RankOneFactorization{T},_ZerosMatrix{T}}}, + model::Model{T}, + _, +) where {T} + n, p = size(component.values) + U = Matrix{T}[] + UW = Matrix{T}[] + weights = zeros(T, n, p) + for (slot, i) in enumerate(component.block_ids) + Ui = zeros(T, n, model.msizes[i]) + for j in 1:n + A = component.values[j, slot] + if !(A isa FillArrays.Zeros) + Ui[j, :] = A.factor + weights[j, slot] = only(A.scaling) + end + end + push!(U, Ui) + push!(UW, similar(Ui)) + end + return _RankOneSchurComponent(component, U, weights, UW, zeros(T, n, n)) +end + +function buffer_for_schur_complement(model::Model{T}, κ) where {T} + AW = [zeros(T, d, d) for d in model.msizes] + WAW = copy.(AW) + components = map(A -> _schur_component(A, model, κ), model.A) + factor_work = map(matrix_indices(model)) do i + rank = _sum_columns(model.A, i.value, 0) do column + return maximum( + A -> A isa AbstractFactorization ? max_rank(A) : 0, + column; + init = 0, + ) + end + return zeros(T, model.msizes[i.value], rank) + end + # Loraine uses the first two per-block dense workspaces directly. + return AW, WAW, components, factor_work +end + +_prepare_schur!(::_GeneralSchurComponent, _, _) = nothing +function _prepare_schur!(workspace::_RankOneSchurComponent, i, W) + slot = _block_slot(workspace.component, i) + if !iszero(slot) + LinearAlgebra.mul!(workspace.UW[slot], workspace.U[slot], W) + end + return nothing +end + +@inline _prepare_components!(::Tuple{}, i, W) = nothing +@inline function _prepare_components!(components::Tuple, i, W) + _prepare_schur!(first(components), i, W) + _prepare_components!(Base.tail(components), i, W) + return nothing +end + +function _sandwich!(WAW, AW, W, A::AbstractMatrix, _) + LinearAlgebra.mul!(AW, W, A) + return LinearAlgebra.mul!(WAW, AW, W) +end +function _sandwich!( + WAW, + AW, + W, + A::AbstractFactorization{T,<:AbstractVector}, + work, +) where {T} + buffered_mul!(AW, A, W, true, false, view(work, :, 1)) + return LinearAlgebra.mul!(WAW, W, AW) +end +function _sandwich!( + WAW, + AW, + W, + A::AbstractFactorization{T,<:AbstractMatrix}, + work, +) where {T} + buffered_mul!(AW, A, W, true, false, view(work, :, 1:max_rank(A))) + return LinearAlgebra.mul!(WAW, W, AW) +end + +function _add_schur_pair!( + a::_GeneralSchurComponent, + b::_GeneralSchurComponent, + i, + W, + AW, + WAW, + work, + H, + same, +) + ia, ib = _block_slot(a.component, i), _block_slot(b.component, i) + (iszero(ia) || iszero(ib)) && return nothing + if same && a.sparse[ia] + _add_sparse_component!(a, ia, W, AW, WAW, H) + return nothing + end + ca, cb = view(a.component.values, :, ia), view(b.component.values, :, ib) + for j in eachindex(ca) + iszero(_nnz(ca[j])) && continue + _sandwich!(WAW, AW, W, ca[j], work) + for k in (same ? j : 1):length(cb) + value = _coefficient_dot(cb[k], WAW) + H[k, j] += value + if !same || k != j + H[j, k] += value + end + end end - # In the future, we'd want to allocate intermediate buffers for this operations - # to make it allocation-free. In the short-term, we just add this `if` - # to be allocation-free when there are no scalars to prevent a regerssion - # in https://github.com/kocvara/Loraine.jl/pull/29 - if !iszero(num_scalars(model)) - result .+= - model.model.C_lin * (W[ScalarIndex] .* (model.model.C_lin' * y)) + return nothing +end + +# For A_j = d_j b_j b_j', the pair contraction is +# d_j e_k (b_j' W c_k)^2. Distinct components contribute this matrix +# AND its transpose; omitting the transpose would lose cross terms. +function _add_schur_pair!( + a::_RankOneSchurComponent, + b::_RankOneSchurComponent, + i, + W, + AW, + WAW, + work, + H, + same, +) + ia, ib = _block_slot(a.component, i), _block_slot(b.component, i) + (iszero(ia) || iszero(ib)) && return nothing + LinearAlgebra.mul!(a.gram, a.UW[ia], b.U[ib]') + for k in axes(H, 2), j in axes(H, 1) + value = a.weights[j, ia] * b.weights[k, ib] * abs2(a.gram[j, k]) + H[j, k] += value + if !same + H[k, j] += value + end + end + return nothing +end + +_quadratic_form(u, A::AbstractMatrix) = LinearAlgebra.dot(u, A, u) +function _quadratic_form(u, ::FillArrays.Zeros{T}) where {T} + return zero(promote_type(eltype(u), T)) +end +function _quadratic_form( + u, + A::AbstractFactorization{T,<:AbstractVector}, +) where {T} + return only(A.scaling) * + LinearAlgebra.dot(u, left_factor(A)) * + LinearAlgebra.dot(right_factor(A), u) +end +function _quadratic_form( + u, + A::AbstractFactorization{T,<:AbstractMatrix}, +) where {T} + result = zero(T) + for k in eachindex(A.scaling) + result += + A.scaling[k] * + LinearAlgebra.dot(u, view(left_factor(A), :, k)) * + LinearAlgebra.dot(view(right_factor(A), :, k), u) end return result end + +function _add_schur_pair!( + a::_RankOneSchurComponent, + b::_GeneralSchurComponent, + i, + W, + AW, + WAW, + work, + H, + same, +) + ia, ib = _block_slot(a.component, i), _block_slot(b.component, i) + (iszero(ia) || iszero(ib)) && return nothing + column = view(b.component.values, :, ib) + for j in axes(H, 1) + u = view(a.UW[ia], j, :) + for k in eachindex(column) + value = a.weights[j, ia] * _quadratic_form(u, column[k]) + H[j, k] += value + H[k, j] += value + end + end + return nothing +end +function _add_schur_pair!( + a::_GeneralSchurComponent, + b::_RankOneSchurComponent, + i, + W, + AW, + WAW, + work, + H, + same, +) + return _add_schur_pair!(b, a, i, W, AW, WAW, work, H, same) +end + +@inline _add_cross_pairs!(a, ::Tuple{}, args...) = nothing +@inline function _add_cross_pairs!(a, others::Tuple, args...) + _add_schur_pair!(a, first(others), args..., false) + _add_cross_pairs!(a, Base.tail(others), args...) + return nothing +end +@inline _add_component_pairs!(::Tuple{}, args...) = nothing +@inline function _add_component_pairs!(components::Tuple, args...) + a, others = first(components), Base.tail(components) + _add_schur_pair!(a, a, args..., true) + _add_cross_pairs!(a, others, args...) + _add_component_pairs!(others, args...) + return nothing +end + +function add_schur_complement!( + model::BufferedModelForSchur, + i::MatrixIndex, + W::AbstractMatrix, + H, +) + AW, WAW, components, work = model.schur_buffer + _prepare_components!(components, i.value, W) + _add_component_pairs!( + components, + i.value, + W, + AW[i.value], + WAW[i.value], + work[i.value], + H, + ) + return H +end + +_has_general_component(::Tuple{}, i) = false +function _has_general_component(components::Tuple, i) + a = first(components) + return ( + a isa _GeneralSchurComponent && !iszero(_block_slot(a.component, i)) + ) || _has_general_component(Base.tail(components), i) +end + +function _add_rank_one_matvec!( + a::_RankOneSchurComponent, + b::_RankOneSchurComponent, + i, + y, + result, + same, +) + ia, ib = _block_slot(a.component, i), _block_slot(b.component, i) + (iszero(ia) || iszero(ib)) && return nothing + LinearAlgebra.mul!(a.gram, a.UW[ia], b.U[ib]') + for k in eachindex(y), j in eachindex(result) + value = a.weights[j, ia] * b.weights[k, ib] * abs2(a.gram[j, k]) + result[j] += value * y[k] + if !same + result[k] += value * y[j] + end + end + return nothing +end +_add_rank_one_matvec!(::_GeneralSchurComponent, _, args...) = nothing +function _add_rank_one_matvec!( + ::_RankOneSchurComponent, + ::_GeneralSchurComponent, + args..., +) + return nothing +end +@inline _add_cross_matvec!(a, ::Tuple{}, args...) = nothing +@inline function _add_cross_matvec!(a, others::Tuple, args...) + _add_rank_one_matvec!(a, first(others), args..., false) + _add_cross_matvec!(a, Base.tail(others), args...) + return nothing +end +@inline _add_component_matvec!(::Tuple{}, args...) = nothing +@inline function _add_component_matvec!(components::Tuple, args...) + a, others = first(components), Base.tail(components) + _add_rank_one_matvec!(a, a, args..., true) + _add_cross_matvec!(a, others, args...) + _add_component_matvec!(others, args...) + return nothing +end + +# Rank-one entries in the setup-time sparse Jacobian buffer. The component +# Schur kernels use their vector factors directly, while the combined +# Jacobian buffer supports products of sums without extra iteration storage. + +# Nonzero (index, value) pairs of a rank-one factor, for both sparse and dense +# factor vectors (dense factors are used e.g. by the SOS rank-one batch). +_rank_one_nz(b::SparseArrays.SparseVector) = SparseArrays.findnz(b) +_rank_one_nz(b::AbstractVector) = (eachindex(b), b) + +# `vec(s b b')` has `nnz(b)²` nonzeros at `support(b) × support(b)`. +function _nnz(F::_RankOneFactorization) + bi, _ = _rank_one_nz(F.factor) + return length(bi)^2 +end + +function _add_vec!(I, J, V, j, offset, F::_RankOneFactorization) + s = only(F.scaling) + d = length(F.factor) + bi, bv = _rank_one_nz(F.factor) + k = offset + for q in eachindex(bi), p in eachindex(bi) + k += 1 + I[k] = (bi[q] - 1) * d + bi[p] + J[k] = j + V[k] = s * bv[p] * bv[q] + end + return k +end diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl index 214d838..c544ee2 100644 --- a/test/BurerMonteiro.jl +++ b/test/BurerMonteiro.jl @@ -3,10 +3,13 @@ # Use of this source code is governed by an MIT-style license that can be found # in the LICENSE.md file or at https://opensource.org/licenses/MIT. +module TestBurerMonteiro + using Test using LinearAlgebra using SparseArrays using JuMP +import FillArrays import LowRankOpt as LRO using Dualization import SolverCore @@ -14,8 +17,7 @@ import Percival include("diff_check.jl") -@testset "Simple LP $opt" for opt in - [LRO.Optimizer, dual_optimizer(LRO.Optimizer)] +function _test_simple_lp(opt) model = Model(opt) @variable(model, x) @constraint(model, con_ref, 1 - x in Nonnegatives()) @@ -53,12 +55,18 @@ include("diff_check.jl") @test abs(sol[1]) < 1e-6 end diff_check(model) -end; + return +end + +function test_simple_lp() + @testset "Simple LP $opt" for opt in + [LRO.Optimizer, dual_optimizer(LRO.Optimizer)] + _test_simple_lp(opt) + end + return +end -@testset "Simple SDP $opt" for (is_dual, opt) in [ - (false, LRO.Optimizer), - (true, dual_optimizer(LRO.Optimizer)), -] +function _test_simple_sdp(is_dual, opt) model = Model(opt) @variable(model, x) @constraint(model, con_ref, Symmetric((1 - x) * ones(1, 1)) in PSDCone()) @@ -113,9 +121,20 @@ end; @test objective_value(model) ≈ 1 @test dual_objective_value(model) ≈ 1 diff_check(model) -end; + return +end -@testset "Trace" begin +function test_simple_sdp() + @testset "Simple SDP $opt" for (is_dual, opt) in [ + (false, LRO.Optimizer), + (true, dual_optimizer(LRO.Optimizer)), + ] + _test_simple_sdp(is_dual, opt) + end + return +end + +function test_trace() T = Float64 model = Model(LRO.Optimizer) cone = MOI.PositiveSemidefiniteConeTriangle(2) @@ -137,29 +156,28 @@ end; nlp = unsafe_backend(model).model @test nlp.C isa Vector{SparseMatrixCSC{T,Int64}} @test nlp.C[1] == [3 1; 1 4] - @test nlp.A isa Matrix{ + @test only(nlp.A).values isa Matrix{ Union{ - LRO.FillArrays.Zeros{ - T, - 2, - Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}, - }, LRO.Factorization{T,Matrix{T},Vector{T}}, + FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}}, }, } - @test nlp.A[1].factor == Matrix(I, 2, 2) - @test nlp.A[1].scaling == [1, 1] - @test nlp.A[2] === LRO.FillArrays.Zeros{T}(Int64(2), Int64(2)) -end; + @test only(nlp.A).block_ids == [1] + @test LRO.jac(nlp, 1, LRO.MatrixIndex(1)).factor == Matrix(I, 2, 2) + @test LRO.jac(nlp, 1, LRO.MatrixIndex(1)).scaling == [1, 1] + @test iszero(LRO.jac(nlp, 2, LRO.MatrixIndex(1))) + return +end -@testset "ResultCount" begin +function test_result_count() model = LRO.Optimizer() @test MOI.get(model, MOI.ResultCount()) == 0 @test MOI.get(model, MOI.PrimalStatus()) == MOI.NO_SOLUTION @test MOI.get(model, MOI.DualStatus()) == MOI.NO_SOLUTION + return end -@testset "MOI runtests" begin +function test_moi_runtests() model = LRO.Optimizer() MOI.set( model, @@ -173,9 +191,10 @@ end ) config = MOI.Test.Config() MOI.Test.runtests(model, config; include = ["Silent"]) -end; + return +end -@testset "No constraints" begin +function test_no_constraints() model = LRO.Model( [spzeros(1, 1)], [ones(1, 1) for _ in 1:1, _ in 1:0], @@ -188,13 +207,15 @@ end; @test LRO.norm_jac(model, LRO.MatrixIndex(1)) == 0 @test LRO.buffer_for_jtprod(model, LRO.MatrixIndex(1)) isa LRO.FillArrays.Zeros -end; + return +end -@testset "Fallback" begin +function test_fallback() @test LRO.Optimizer() isa LRO.Optimizer{Float64} + return end -@testset "Different factor ranks" begin +function test_different_factor_ranks() T = Float64 model = Model(LRO.Optimizer) cone = MOI.PositiveSemidefiniteConeTriangle(2) @@ -219,7 +240,22 @@ end MT = LRO.Factorization{T,Matrix{T},LRO.Ones{T}} @test nlp.C isa Vector{MT} @test nlp.C[] == factors[1] - @test nlp.A isa Matrix{MT} - @test nlp.A[1] == factors[2] - @test nlp.A[2] == factors[3] -end; + @test only(nlp.A).values isa Matrix{MT} + @test LRO.jac(nlp, 1, LRO.MatrixIndex(1)) == factors[2] + @test LRO.jac(nlp, 2, LRO.MatrixIndex(1)) == factors[3] + return +end + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$name", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end +end + +end + +TestBurerMonteiro.runtests() diff --git a/test/buffer.jl b/test/buffer.jl index 0f0d772..45bf894 100644 --- a/test/buffer.jl +++ b/test/buffer.jl @@ -28,25 +28,26 @@ function _test_zero_Ai(all_zero::Bool, matrix_in_objective::Bool) S = SparseArrays.SparseMatrixCSC{T,Int64} @test b.model.C isa Vector{matrix_in_objective ? S : Z} if all_zero - @test b.model.A isa Matrix{Z} + @test isempty(b.model.A) else - @test b.model.A isa Matrix{Union{Z,S}} + # Missing entries of a component use `Int` axes, unlike the `Int64` + # ones built by the MOI wrapper, which differ on 32-bit systems. + ZInt = FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}} + @test only(b.model.A).values isa Matrix{Union{S,ZInt}} end - @test b.model.A[1] isa Z - @test b.model.A[2] isa Z - if all_zero - @test b.model.A[3] isa Z - else - @test b.model.A[3] isa S - @test b.model.A[4] isa Z + i = LRO.MatrixIndex(1) + @test iszero(LRO.jac(b.model, 1, i)) + @test iszero(LRO.jac(b.model, 2, i)) + @test iszero(LRO.jac(b.model, 3, i)) == all_zero + if !all_zero + @test LRO.jac(b.model, 3, i) isa S + @test iszero(LRO.jac(b.model, 4, i)) end buf = LRO.BufferedModelForSchur(b.model, 1) - for A in b.model.A - if all_zero - @test buf.jtprod_buffer[] isa LRO.FillArrays.Zeros - else - @test buf.jtprod_buffer[] !== A - end + if all_zero + @test buf.jtprod_buffer[] isa LRO.FillArrays.Zeros + else + @test all(A -> buf.jtprod_buffer[] !== A, only(b.model.A).values) end for κ in 0:5 schur_test(model, κ) diff --git a/test/components.jl b/test/components.jl new file mode 100644 index 0000000..5d46759 --- /dev/null +++ b/test/components.jl @@ -0,0 +1,319 @@ +# Copyright (c) 2026: Benoît Legat and contributors +# +# Use of this source code is governed by an MIT-style license that can be found +# in the LICENSE.md file or at https://opensource.org/licenses/MIT. + +module TestComponents + +using Test, LinearAlgebra, SparseArrays, Random +import LowRankOpt as LRO +import NLPModels +import FillArrays + +function _rank_component(rng, ::Type{T}, ids, dims, n) where {T} + values = [ + LRO.Factorization(randn(rng, T, dims[i]), T((-1)^j * (j + i) / 5)) + for j in 1:n, i in ids + ] + return LRO.BlockColumns(ids, values, length(dims)) +end + +function _fixture(::Type{T}, kind) where {T} + rng = MersenneTwister(42) + dims, n = Int64[2, 3, 1], 3 + r1 = _rank_component(rng, T, [1, 2], dims, n) + r2 = _rank_component(rng, T, [1, 3], dims, n) + sparse = LRO.BlockColumns( + [1, 3], + [spdiagm(T.(1:dims[i]) .+ T(j)) for j in 1:n, i in (1, 3)], + 3, + ) + diagonal = LRO.BlockColumns( + [2], + [Diagonal(T.(1:3) .* T(j / 4)) for j in 1:n, _ in 1:1], + 3, + ) + lowrank = LRO.BlockColumns( + [2], + [ + LRO.Factorization(randn(rng, T, 3, 2), T[0.75, -0.5]) for + j in 1:n, _ in 1:1 + ], + 3, + ) + components = + kind == :rank_one ? (r1, r2) : (sparse, r1, r2, diagonal, lowrank) + # Build the reference directly from stored entries, independently of jac. + full = [zeros(T, d, d) for d in dims, _ in 1:n] + for component in components, + (k, i) in enumerate(component.block_ids), + j in 1:n + + full[i, j] .+= component.values[j, k] + end + C = [spdiagm(fill(T(i), d)) for (i, d) in enumerate(dims)] + b = T[1, -2, 3] + scalars = sparsevec(Int64[], T[], 0) + scalar_jacobian = spzeros(T, Int64, n, 0) + model = LRO.Model(C, components, b, scalars, scalar_jacobian, dims) + reference = LRO.Model(C, full, b, scalars, scalar_jacobian, dims) + W = LRO.ShapedSolution(T[], [ + let R = randn(rng, T, d, d) + R * R' + I + end for d in dims + ]) + x = LRO.ShapedSolution( + T[], + [Matrix(Symmetric(randn(rng, T, d, d))) for d in dims], + ) + y = randn(rng, T, n) + return model, reference, full, W, x, y +end + +function _iteration!(buffer, W, x, y, H, Hy, Jv, Jty) + fill!(H, zero(eltype(H))) + LRO.add_schur_complement!(buffer, W, LRO.MatrixIndex, H) + LRO.eval_schur_complement!(buffer, W, y, Hy) + NLPModels.jprod!(buffer, x, x, Jv) + NLPModels.jtprod!(buffer, x, y, Jty) + return nothing +end +function _allocations(args...) + _iteration!(args...) + return @allocated _iteration!(args...) +end + +@testset "Component sums and Schur cross terms" begin + for T in (Float32, Float64), kind in (:rank_one, :mixed), κ in (0, 1, 100) + model, reference, full, W, x, y = _fixture(T, kind) + buffer = LRO.BufferedModelForSchur(model, κ) + reference_buffer = LRO.BufferedModelForSchur(reference, κ) + H = zeros(T, 3, 3) + Hy = zeros(T, 3) + Jv = zeros(T, 3) + Jty = LRO.ShapedSolution(T[], similar.(x.matrices)) + for i in LRO.matrix_indices(model), j in 1:3 + @test LRO.jac(model, j, i) ≈ full[i.value, j] + end + for i in LRO.matrix_indices(model) + @test LRO.norm_jac(model, i) ≈ + sqrt(sum(sum(abs2, full[i.value, j]) for j in 1:3)) + @test LRO.unsafe_jtprod(buffer, y, i) ≈ + sum(y[j] * full[i.value, j] for j in 1:3) + @test LRO.unsafe_jtprod(buffer, sparsevec(y), i) ≈ + sum(y[j] * full[i.value, j] for j in 1:3) + end + expected_H = [ + sum( + dot(full[i.value, j], W[i] * full[i.value, k] * W[i]) for + i in LRO.matrix_indices(model) + ) for j in 1:3, k in 1:3 + ] + _iteration!(buffer, W, x, y, H, Hy, Jv, Jty) + @test H ≈ expected_H + @test Hy ≈ expected_H * y + expected_Jv = [ + sum(dot(full[i.value, j], x[i]) for i in LRO.matrix_indices(model)) for j in 1:3 + ] + @test Jv ≈ expected_Jv + unbuffered_Jv = similar(Jv) + NLPModels.jprod!(model, x, x, unbuffered_Jv) + @test unbuffered_Jv ≈ Jv + for i in LRO.matrix_indices(model) + @test Jty[i] ≈ sum(y[j] * full[i.value, j] for j in 1:3) + end + reference_H = similar(H) + fill!(reference_H, 0) + LRO.add_schur_complement!( + reference_buffer, + W, + LRO.MatrixIndex, + reference_H, + ) + @test H ≈ reference_H + @test _allocations(buffer, W, x, y, H, Hy, Jv, Jty) == 0 + end +end + +@testset "Missing component blocks" begin + T, n = Float64, 2 + component = LRO.BlockColumns( + [2], + reshape( + [LRO.Factorization([1.0], 2.0), LRO.Factorization([3.0], -1.0)], + n, + 1, + ), + 3, + ) + model = LRO.Model( + [spzeros(d, d) for d in (2, 1, 3)], + (component,), + zeros(n), + sparsevec(Int64[], T[], 0), + spzeros(T, Int64, n, 0), + Int64[2, 1, 3], + ) + @test size(component) == (2, 3) + @test isconcretetype(eltype(component.values)) + @test iszero(LRO.jac(model, 1, LRO.MatrixIndex(1))) + @test size(LRO.jac(model, 1, LRO.MatrixIndex(1))) == (2, 2) + @test iszero(LRO.jac(model, 2, LRO.MatrixIndex(3))) + @test LRO.jac(model, 1, LRO.MatrixIndex(2)) == fill(2.0, 1, 1) + @test_throws ArgumentError LRO.BlockColumns( + [2, 1], + repeat(component.values, 1, 2), + 3, + ) + @test_throws DimensionMismatch LRO.BlockColumns([1, 2], component.values, 3) + @test_throws DimensionMismatch LRO.Model( + model.C, + (LRO.BlockColumns([1], component.values, 2),), + model.b, + model.d_lin, + model.C_lin, + model.msizes, + ) +end + +@testset "Legacy matrix with Zeros entries" begin + T = Float64 + rng = MersenneTwister(7) + dims, n = Int64[3, 2, 2], 3 + Z(d) = FillArrays.Zeros{T}(d, d) + R(d) = LRO.Factorization(randn(rng, T, d), T(0.5)) + S(d) = sparse(Symmetric(sprandn(rng, T, d, d, 0.7) + I)) + # Block 1 mixes all three types, block 2 has no rank-one entry and + # block 3 has no sparse entry. + A = Matrix{Union{typeof(Z(1)),typeof(R(1)),typeof(S(1))}}(undef, 3, n) + A[1, :] = [R(3), S(3), Z(3)] + A[2, :] = [Z(2), S(2), S(2)] + A[3, :] = [R(2), Z(2), R(2)] + full = Matrix.(A) + C = [spdiagm(fill(T(i), d)) for (i, d) in enumerate(dims)] + b = T[1, -2, 3] + scalars = sparsevec(Int64[], T[], 0) + scalar_jacobian = spzeros(T, Int64, n, 0) + model = LRO.Model(C, A, b, scalars, scalar_jacobian, dims) + rank_one, sparse_component = model.A + @test rank_one.block_ids == [1, 3] + @test sparse_component.block_ids == [1, 2] + @test eltype(rank_one.values) == Union{typeof(R(1)),typeof(Z(1))} + @test eltype(sparse_component.values) == Union{typeof(S(1)),typeof(Z(1))} + W = LRO.ShapedSolution(T[], [ + let R = randn(rng, T, d, d) + R * R' + I + end for d in dims + ]) + x = LRO.ShapedSolution( + T[], + [Matrix(Symmetric(randn(rng, T, d, d))) for d in dims], + ) + y = randn(rng, T, n) + buffer = LRO.BufferedModelForSchur(model, 1) + H = zeros(T, n, n) + Hy = zeros(T, n) + Jv = zeros(T, n) + Jty = LRO.ShapedSolution(T[], similar.(x.matrices)) + _iteration!(buffer, W, x, y, H, Hy, Jv, Jty) + expected_H = [ + sum( + dot(full[i.value, j], W[i] * full[i.value, k] * W[i]) for + i in LRO.matrix_indices(model) + ) for j in 1:n, k in 1:n + ] + @test H ≈ expected_H + @test Hy ≈ expected_H * y + @test Jv ≈ [ + sum(dot(full[i.value, j], x[i]) for i in LRO.matrix_indices(model)) for + j in 1:n + ] + @test NLPModels.jprod!(model, x, x, similar(Jv)) ≈ Jv + for i in LRO.matrix_indices(model) + @test Jty[i] ≈ sum(y[j] * full[i.value, j] for j in 1:n) + for j in 1:n + @test LRO.jac(model, j, i) ≈ full[i.value, j] + end + end + @test _allocations(buffer, W, x, y, H, Hy, Jv, Jty) == 0 +end + +@testset "Vector factors outside the rank-one kernel" begin + T = Float64 + rng = MersenneTwister(11) + d, n = 3, 3 + # A symmetric `AsymmetricFactorization` with vector factors is not a + # `Factorization` so it uses the general Schur kernel, also when paired + # with a rank-one component. + asym = LRO.BlockColumns( + [1], + reshape( + [ + let u = randn(rng, T, d) + LRO.AsymmetricFactorization(u, u, fill(T(j))) + end for j in 1:n + ], + n, + 1, + ), + 1, + ) + @test size(asym, 1) == n + @test size(asym, 2) == 1 + rank_one = LRO.BlockColumns( + [1], + reshape( + [LRO.Factorization(randn(rng, T, d), T(-j)) for j in 1:n], + n, + 1, + ), + 1, + ) + # Tuple entries can also be in the legacy block-by-constraint layout. + legacy = reshape([sparse(T(j) * I, d, d) for j in 1:n], 1, n) + full = [ + Matrix(asym.values[j, 1]) + + Matrix(rank_one.values[j, 1]) + + Matrix(legacy[1, j]) for _ in 1:1, j in 1:n + ] + model = LRO.Model( + [spzeros(T, d, d)], + (asym, rank_one, legacy), + zeros(T, n), + sparsevec(Int64[], T[], 0), + spzeros(T, Int64, n, 0), + Int64[d], + ) + @test model.A[3] == LRO.BlockColumns(legacy) + buffer = LRO.BufferedModelForSchur(model, 0) + W = let R = randn(rng, T, d, d) + R * R' + I + end + y = randn(rng, T, n) + i = LRO.MatrixIndex(1) + H = zeros(T, n, n) + LRO.add_schur_complement!(buffer, i, W, H) + expected_H = [dot(full[j], W * full[k] * W) for j in 1:n, k in 1:n] + @test H ≈ expected_H + Hy = zeros(T, n) + LRO.eval_schur_complement!(buffer, LRO.ShapedSolution(T[], [W]), y, Hy) + @test Hy ≈ expected_H * y + @test LRO.unsafe_jtprod(buffer, y, i) ≈ sum(y[j] * full[j] for j in 1:n) +end + +@testset "Burer-Monteiro with sparse plus low-rank components" begin + model, reference, _, _, _, y = _fixture(Float64, :mixed) + ranks = Int64[1, 2, 1] + mixed = LRO.BurerMonteiro.Model{false}(model, ranks) + dense = LRO.BurerMonteiro.Model{false}(reference, ranks) + rng = MersenneTwister(3) + x, v = randn(rng, mixed.meta.nvar), randn(rng, mixed.meta.nvar) + @test NLPModels.obj(mixed, x) ≈ NLPModels.obj(dense, x) + @test NLPModels.cons(mixed, x) ≈ NLPModels.cons(dense, x) + @test NLPModels.grad(mixed, x) ≈ NLPModels.grad(dense, x) + @test NLPModels.jprod(mixed, x, v) ≈ NLPModels.jprod(dense, x, v) + @test NLPModels.jtprod(mixed, x, y) ≈ NLPModels.jtprod(dense, x, y) + @test NLPModels.hprod(mixed, x, y, v) ≈ NLPModels.hprod(dense, x, y, v) +end + +end diff --git a/test/diff_check.jl b/test/diff_check.jl index 7057a03..34236be 100644 --- a/test/diff_check.jl +++ b/test/diff_check.jl @@ -167,7 +167,7 @@ function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} Hy = similar(y) LRO.eval_schur_complement!(model, w, y, Hy) @test Hy ≈ H * y - if iszero(LRO.num_scalars(model)) && T == Float64 + if T == Float64 @test 0 == @allocated LRO.eval_schur_complement!(model, w, y, Hy) end # Reuse the workspaces for a different right-hand side, then assemble H again. @@ -190,7 +190,8 @@ function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} # may also alias `C` itself, which is sparse. @test ret isa Union{Matrix,SparseArrays.SparseMatrixCSC,FillArrays.Zeros} - Aty = sum(model.model.A[i.value, j] * y[j] for j in 1:(model.meta.ncon)) + Aty = + sum(LRO.jac(model.model, j, i) * y[j] for j in 1:(model.meta.ncon)) @test ret ≈ model.model.C[i.value] - Aty if Aty isa FillArrays.Zeros @test ret === model.model.C[i.value] diff --git a/test/factorization.jl b/test/factorization.jl index a108178..983c2c1 100644 --- a/test/factorization.jl +++ b/test/factorization.jl @@ -75,6 +75,16 @@ function test_dot() end end +# Rank-one factorizations with a 0-dimensional `Array` scaling, as built by +# `LRO.rank_one`. +function test_dot_rank_one_array_scaling() + u, v = [1.0, 2.0], [3.0, -1.0] + A = LRO.Factorization(u, 2.0) + B = LRO.Factorization(v, 3.0) + @test A.scaling isa Array{Float64,0} + @test dot(A, B) ≈ dot(Matrix(A), Matrix(B)) +end + function _test_mul(A, B, α, β) if B isa AbstractVector res = ones(length(B)) diff --git a/test/holy.jl b/test/holy.jl index 3374365..5da2679 100644 --- a/test/holy.jl +++ b/test/holy.jl @@ -3,14 +3,20 @@ # Use of this source code is governed by an MIT-style license that can be found # in the LICENSE.md file or at https://opensource.org/licenses/MIT. -include("diff_check.jl") -include(joinpath(dirname(@__DIR__), "examples", "holy_model.jl")) +module TestHoly -import Random +using Test +using SparseArrays +using JuMP using Dualization +import Random import Percival +import LowRankOpt as LRO + +include("diff_check.jl") +include(joinpath(dirname(@__DIR__), "examples", "holy_model.jl")) -function test_holy(; is_dual, low_rank, square_scalars, n = 10) +function _test_holy(; is_dual, low_rank, square_scalars, n = 10) opt = LRO.Optimizer if is_dual opt = dual_optimizer(opt) @@ -41,12 +47,29 @@ function test_holy(; is_dual, low_rank, square_scalars, n = 10) SparseMatrixCSC{T,Int64} end @test b.model.C isa Vector{F} - @test b.model.A isa Matrix{F} + @test only(b.model.A).values isa Matrix{F} end -@testset "Holy low-rank ? $low_rank" for low_rank in [false, true] - @testset "Square scalars ? $square_scalars" for square_scalars in - [false, true] - test_holy(; is_dual = true, low_rank, square_scalars) +function test_holy() + @testset "Holy low-rank ? $low_rank" for low_rank in [false, true] + @testset "Square scalars ? $square_scalars" for square_scalars in + [false, true] + _test_holy(; is_dual = true, low_rank, square_scalars) + end end -end; + return +end + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$name", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end +end + +end + +TestHoly.runtests() diff --git a/test/maxcut.jl b/test/maxcut.jl index 9c16c4a..0b5d820 100644 --- a/test/maxcut.jl +++ b/test/maxcut.jl @@ -3,14 +3,27 @@ # Use of this source code is governed by an MIT-style license that can be found # in the LICENSE.md file or at https://opensource.org/licenses/MIT. +module TestMaxcut + +using Test +using SparseArrays +using FillArrays +using JuMP +using Dualization +import Random import Percival +import NLPModels +import LowRankOpt as LRO include("diff_check.jl") include(joinpath(dirname(@__DIR__), "examples", "maxcut.jl")) -weights = [0 5 7 6; 5 0 0 1; 7 0 0 1; 6 1 1 0]; +weights = [0 5 7 6; 5 0 0 1; 7 0 0 1; 6 1 1 0] -function test_maxcut(; is_dual, sparse, vector) +function _test_maxcut(; is_dual, sparse, vector) + # The Burer-Monteiro starting point is `rand`-based (see + # `BurerMonteiro.meta`) + Random.seed!(0) opt = LRO.Optimizer if is_dual opt = dual_optimizer(opt) @@ -53,13 +66,14 @@ function test_maxcut(; is_dual, sparse, vector) end end D = vector ? LRO.One{Float64} : LRO.Ones{Float64} - @test lro_model.A isa Matrix{LowRankOpt.Factorization{Float64,F,D}} + @test only(lro_model.A).values isa + Matrix{LowRankOpt.Factorization{Float64,F,D}} else lro_model = unsafe_backend(model).model @test lro_model.C isa Vector{ FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}, } - @test lro_model.A isa Matrix{SparseMatrixCSC{T,Int64}} + @test only(lro_model.A).values isa Matrix{SparseMatrixCSC{T,Int64}} solver = unsafe_backend(model).solver LRO.BurerMonteiro.set_rank!(solver.model, LRO.MatrixIndex(1), 4) @test solver.model.dim.ranks == [4] @@ -82,15 +96,18 @@ function test_maxcut(; is_dual, sparse, vector) end end -@testset "Max-CUT dual ? $is_dual" for is_dual in [false, true] - @testset "Sparse ? $sparse" for sparse in [false, true] - @testset "Vector ? $vector" for vector in [false, true] - test_maxcut(; is_dual, sparse, vector) +function test_maxcut() + @testset "Max-CUT dual ? $is_dual" for is_dual in [false, true] + @testset "Sparse ? $sparse" for sparse in [false, true] + @testset "Vector ? $vector" for vector in [false, true] + _test_maxcut(; is_dual, sparse, vector) + end end end -end; + return +end -@testset "ConvexSolver $T" for T in [Float32, Float64] +function _test_convex_solver(T) model = maxcut(T.(weights), LRO.Optimizer{T}) set_attribute(model, "solver", ConvexSolver) b = unsafe_backend(model) @@ -143,4 +160,26 @@ end; schur_test(b.model, 0) schur_test(b.model, 1) schur_test(b.model, 2) -end; + return +end + +function test_convex_solver() + @testset "ConvexSolver $T" for T in [Float32, Float64] + _test_convex_solver(T) + end + return +end + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$name", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end +end + +end + +TestMaxcut.runtests() diff --git a/test/rank1_detection.jl b/test/rank1_detection.jl new file mode 100644 index 0000000..f6539fd --- /dev/null +++ b/test/rank1_detection.jl @@ -0,0 +1,227 @@ +# Copyright (c) 2024: Benoît Legat and contributors +# +# Use of this source code is governed by an MIT-style license that can be found +# in the LICENSE.md file or at https://opensource.org/licenses/MIT. +# +# Tests for `rank_one`/`detect_rank_one` and the `detect_rank_one` Optimizer +# option that converts sparse rank-one constraint matrices to rank-one +# `Factorization`s (so they hit the rank-1 Schur specialization in `schur.jl`). + +module TestRank1Detection + +using Test +using LinearAlgebra +using SparseArrays +import FillArrays +import MathOptInterface as MOI +using JuMP +import LowRankOpt as LRO + +include("diff_check.jl") + +function test_rank_one_unit_vector() + A = sparse([3], [3], [1.0], 4, 4) # e₃e₃' + f = LRO.rank_one(A) + @test f isa LRO.Factorization + @test Matrix(f) == Matrix(A) + @test only(f.scaling) == 1.0 + return +end + +function test_rank_one_general() + b = [1.0, 2.0, 0.0, 3.0] + M = sparse(2.0 .* (b * b')) + f = LRO.rank_one(M) + @test f isa LRO.Factorization + @test Matrix(f) ≈ Matrix(M) + return +end + +function test_rank_one_indefinite() + b = [1.0, 2.0, 0.0, 3.0] + M = sparse(-1.5 .* (b * b')) + f = LRO.rank_one(M) + @test f isa LRO.Factorization + @test Matrix(f) ≈ Matrix(M) + @test only(f.scaling) < 0 # negative scaling propagates + return +end + +function test_rank_one_zero() + @testset "$(typeof(Z))" for Z in + [spzeros(3, 3), FillArrays.Zeros{Float64}(3, 3)] + f = LRO.rank_one(Z) + @test f isa LRO.Factorization + @test iszero(Matrix(f)) + end + return +end + +function test_rank_one_rejects_higher_rank() + @test isnothing(LRO.rank_one(sparse([2.0 1.0; 1.0 2.0]))) # rank 2 + # Symmetric, nonzero, but zero diagonal ⇒ not rank one. + @test isnothing(LRO.rank_one(sparse([1, 2], [2, 1], [1.0, 1.0], 2, 2))) + return +end + +function test_rank_one_rejects_extra_nonzero() + # `e₁e₁'` plus a stray entry is not rank one. + A = sparse([1, 1], [1, 2], [1.0, 1.0], 2, 2) + @test isnothing(LRO.rank_one(A)) + return +end + +function test_rank_one_from_factorization() + g = LRO.positive_semidefinite_factorization([1.0, 2.0, 3.0]) + f = LRO.rank_one(g) + @test f isa + LRO.Factorization{Float64,<:SparseVector,<:AbstractArray{Float64,0}} + @test Matrix(f) ≈ Matrix(g) + return +end + +function test_detect_rank_one_all() + arr = [sparse([j], [j], [1.0], 4, 4) for i in 1:1, j in 1:4] + R = LRO.detect_rank_one(arr) + @test !isnothing(R) + @test isconcretetype(eltype(only(R).values)) + @test eltype(only(R).values) <: LRO.Factorization + @test all( + j -> Matrix(only(R).values[j, 1]) == Matrix(arr[1, j]), + axes(arr, 2), + ) + return +end + +function test_detect_rank_one_mixed() + arr = Matrix{SparseMatrixCSC{Float64,Int}}(undef, 1, 2) + arr[1, 1] = sparse([1], [1], [1.0], 2, 2) + arr[1, 2] = sparse([2.0 1.0; 1.0 2.0]) # rank 2 + components = LRO.detect_rank_one(arr) + @test length(components) == 2 + @test all( + j -> + sum(component.values[j, 1] for component in components) ≈ arr[1, j], + 1:2, + ) + return +end + +function test_detect_rank_one_with_zeros() + Z = FillArrays.Zeros{Float64}(2, 2) + S = sparse([1], [1], [1.0], 2, 2) + arr = Matrix{Union{typeof(Z),typeof(S)}}(undef, 1, 2) + arr[1, 1] = Z + arr[1, 2] = S + R = LRO.detect_rank_one(arr) + @test !isnothing(R) + @test iszero(Matrix(only(R).values[1, 1])) + @test Matrix(only(R).values[2, 1]) == Matrix(S) + return +end + +function test_rank_one_unsupported_type() + @test isnothing(LRO.rank_one(Diagonal([1.0, 2.0]))) + return +end + +# Detected rank-one matrices have sparse factors, they are used both by the +# rank-one Schur kernel and by the sparse Jacobian buffer. +function test_detected_schur_complement() + d, n = 4, 3 + us = [sparsevec([j, j + 1], [1.0, j], d) for j in 1:n] + arr = [sparse(u * u') for _ in 1:1, u in us] + C = [spzeros(d, d)] + b = zeros(n) + d_lin = sparsevec(Int64[], Float64[], 0) + C_lin = spzeros(Float64, Int64, n, 0) + model = LRO.Model(C, LRO.detect_rank_one(arr), b, d_lin, C_lin, Int64[d]) + @test eltype(only(model.A).values) <: LRO.Factorization + reference = LRO.Model(C, arr, b, d_lin, C_lin, Int64[d]) + buffer = LRO.BufferedModelForSchur(model, 0) + reference_buffer = LRO.BufferedModelForSchur(reference, 0) + W = Matrix(Symmetric(ones(d, d) + d * I)) + i = LRO.MatrixIndex(1) + H = zeros(n, n) + LRO.add_schur_complement!(buffer, i, W, H) + expected = zeros(n, n) + LRO.add_schur_complement!(reference_buffer, i, W, expected) + @test H ≈ expected + y = [1.0, -2.0, 3.0] + @test LRO.unsafe_jtprod(buffer, y, i) ≈ + LRO.unsafe_jtprod(reference_buffer, y, i) + return +end + +function test_detect_rank_one_empty() + @test isnothing( + LRO.detect_rank_one(Matrix{SparseMatrixCSC{Float64,Int}}(undef, 0, 0)), + ) + return +end + +# A tiny max-cut-like SDP in SDPA format whose data matrices `Fⱼ = eⱼeⱼ'` are +# rank one, written to a temporary file. +const _SDPA = """ +2 +1 +2 +1.0 1.0 +0 1 1 1 2.0 +0 1 2 2 2.0 +0 1 1 2 1.0 +1 1 1 1 1.0 +2 1 2 2 1.0 +""" + +function _build(detect) + path = tempname() * ".dat-s" + write(path, _SDPA) + src = MOI.FileFormats.SDPA.Model{Float64}() + MOI.read_from_file(src, path) + rm(path) + model = GenericModel{Float64}(LRO.Optimizer{Float64}) + LRO.Bridges.add_all_bridges(backend(model).optimizer, Float64) + MOI.copy_to(backend(model), src) + set_attribute(model, "solver", ConvexSolver) + set_attribute(model, "detect_rank_one", detect) + optimize!(model) # `ConvexSolver` is a no-op: this just triggers `copy_to` + return unsafe_backend(model).model +end + +function test_option_default_keeps_sparse() + model = _build(false) + @test eltype(only(model.A).values) <: SparseMatrixCSC + return +end + +function test_option_converts_to_rank_one() + sparse_model = _build(false) + rank_one_model = _build(true) + @test eltype(only(rank_one_model.A).values) <: LRO.Factorization + # `C` (the objective block) is not rank one, so it stays as is. + @test eltype(rank_one_model.C) <: SparseMatrixCSC + # The conversion is exact: same matrices, different representation. + @test size(only(rank_one_model.A)) == size(only(sparse_model.A)) + @test all( + j -> + Matrix(LRO.jac(rank_one_model, j, LRO.MatrixIndex(1))) == + Matrix(LRO.jac(sparse_model, j, LRO.MatrixIndex(1))), + 1:sparse_model.meta.ncon, + ) + return +end + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$name", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end +end + +end + +TestRank1Detection.runtests() diff --git a/test/schur_rank1.jl b/test/schur_rank1.jl new file mode 100644 index 0000000..e9c99b3 --- /dev/null +++ b/test/schur_rank1.jl @@ -0,0 +1,188 @@ +# Copyright (c) 2026: Benoît Legat and contributors +# +# Use of this source code is governed by an MIT-style license that can be found +# in the LICENSE.md file or at https://opensource.org/licenses/MIT. +# +# Verifies the rank-1 Schur-complement specialization in src/schur.jl: +# when constraint matrices are rank-1 `LRO.Factorization`s, the block +# contribution H[j,k] += d_j d_k (b_j' W b_k)² should match a brute-force +# evaluation of ⟨A_{i,j}, W A_{i,k} W⟩. + +module TestSchurRank1 + +using Test +using LinearAlgebra +using SparseArrays +import LowRankOpt as LRO +import Random + +# Build a tiny `LRO.Model` whose constraint matrices share the +# `view(U_i, j, :)` shape produced by the SOS bridge / sample-from-parent +# rank-1 batch (`LRO._rank_one_rowview_batch`). +function _rank1_model(::Type{T}, sizes::Vector{Int}, n::Int; seed = 0) where {T} + rng = Random.MersenneTwister(seed) + p = length(sizes) + Us = [randn(rng, T, n, m) for m in sizes] + weights = [randn(rng, T, n) for _ in 1:p] + AT = LRO.Factorization{ + T, + SubArray{T,1,Matrix{T},Tuple{Int,Base.Slice{Base.OneTo{Int}}},true}, + Array{T,0}, + } + A = Matrix{AT}(undef, p, n) + for i in 1:p, j in 1:n + A[i, j] = + LRO.Factorization(view(Us[i], j, :), reshape(T[weights[i][j]], ())) + end + CT = LRO.Factorization{T,Vector{T},Array{T,0}} + C = CT[ + LRO.Factorization(zeros(T, m), reshape(T[zero(T)], ())) for m in sizes + ] + b = randn(rng, T, n) + d_lin = sparsevec(Int64[], T[], 0) + C_lin = spzeros(T, Int64, n, 0) + return LRO.Model(C, A, b, d_lin, C_lin, Int64.(sizes)), Us, weights +end + +# Brute-force: H[j,k] = Σ_i ⟨A_{i,j}, W_i A_{i,k} W_i⟩, no rank-1 trick. +function _brute_force_H(model, W_blocks::Vector{<:AbstractMatrix}) + n = model.meta.ncon + T = eltype(model.b) + H = zeros(T, n, n) + for idx in LRO.matrix_indices(model) + i = idx.value + Wi = W_blocks[i] + for j in 1:n, k in 1:n + Aj = Matrix(LRO.jac(model, j, idx)) + Ak = Matrix(LRO.jac(model, k, idx)) + H[j, k] += tr(Aj * Wi * Ak * Wi) + end + end + return H +end + +function _random_sym_blocks(::Type{T}, sizes; seed = 1) where {T} + rng = Random.MersenneTwister(seed) + return map(sizes) do m + Wi = randn(rng, T, m, m) + return (Wi + Wi') / 2 + end +end + +# `LRO.add_schur_complement!` for one block matches the brute-force +# Schur contribution. +function test_single_block() + T = Float64 + sizes = [3] + n = 4 + model, _, _ = _rank1_model(T, sizes, n) + buf = LRO.BufferedModelForSchur(model, 1) + W = _random_sym_blocks(T, sizes) + H = zeros(T, n, n) + LRO.add_schur_complement!(buf, LRO.MatrixIndex(1), W[1], H) + @test H ≈ _brute_force_H(model, W) +end + +# Multi-block: each block contributes its own rank-1 trick; summed via +# per-block `add_schur_complement!` calls. +function test_multiple_blocks() + T = Float64 + sizes = [3, 2, 4] + n = 5 + model, _, _ = _rank1_model(T, sizes, n) + buf = LRO.BufferedModelForSchur(model, 1) + W = _random_sym_blocks(T, sizes) + H = zeros(T, n, n) + for idx in LRO.matrix_indices(model) + LRO.add_schur_complement!(buf, idx, W[idx.value], H) + end + @test H ≈ _brute_force_H(model, W) +end + +# Single constraint sanity (n = 1): H is the 1×1 matrix +# d² (b' W b)² for that constraint. +function test_single_constraint() + T = Float64 + sizes = [3] + n = 1 + model, Us, ws = _rank1_model(T, sizes, n) + buf = LRO.BufferedModelForSchur(model, 1) + W = _random_sym_blocks(T, sizes) + H = zeros(T, n, n) + LRO.add_schur_complement!(buf, LRO.MatrixIndex(1), W[1], H) + b = Us[1][1, :] + d = ws[1][1] + @test H[1, 1] ≈ d^2 * (b' * W[1] * b)^2 +end + +# Negative scalings should propagate through `d_j d_k` (the squaring is +# only of `b_j' W b_k`). +function test_negative_scalings() + T = Float64 + sizes = [2] + n = 2 + model, Us, ws = _rank1_model(T, sizes, n) + ws[1] .= [-1.5, 2.3] + AT = LRO.Factorization{ + T, + SubArray{T,1,Matrix{T},Tuple{Int,Base.Slice{Base.OneTo{Int}}},true}, + Array{T,0}, + } + A = Matrix{AT}(undef, 1, n) + for j in 1:n + A[1, j] = LRO.Factorization(view(Us[1], j, :), reshape(T[ws[1][j]], ())) + end + CT = LRO.Factorization{T,Vector{T},Array{T,0}} + C = CT[LRO.Factorization(zeros(T, 2), reshape(T[zero(T)], ()))] + model2 = LRO.Model(C, A, model.b, model.d_lin, model.C_lin, model.msizes) + buf = LRO.BufferedModelForSchur(model2, 1) + W = _random_sym_blocks(T, sizes; seed = 7) + H = zeros(T, n, n) + LRO.add_schur_complement!(buf, LRO.MatrixIndex(1), W[1], H) + @test H ≈ _brute_force_H(model2, W) + # off-diagonal sign comes from d_1 * d_2 = (-1.5) * 2.3 < 0. + @test sign(H[1, 2]) == sign(ws[1][1] * ws[1][2]) +end + +# Shim supporting `W[MatrixIndex(i)]` and `W[LRO.ScalarIndex]` indexing +# the way `LRO.VectorizedSolution` does on the solver side. +struct _WShim{T} + blocks::Vector{Matrix{T}} + scalars::Vector{T} +end +Base.getindex(W::_WShim, i::LRO.MatrixIndex) = W.blocks[i.value] +Base.getindex(W::_WShim, ::Type{LRO.ScalarIndex}) = W.scalars + +# `eval_schur_complement!` (rank-1 specialization) should agree with +# the explicit matrix-vector product `H * y`. +function test_eval_matches_matvec() + T = Float64 + sizes = [3, 2] + n = 4 + model, _, _ = _rank1_model(T, sizes, n) + buf = LRO.BufferedModelForSchur(model, 1) + W_blocks = _random_sym_blocks(T, sizes) + W = _WShim(W_blocks, T[]) + rng = Random.MersenneTwister(11) + y = randn(rng, T, n) + H = zeros(T, n, n) + LRO.add_schur_complement!(buf, W, LRO.MatrixIndex, H) + Hy_ref = H * y + Hy = zeros(T, n) + LRO.eval_schur_complement!(buf, W, y, Hy) + @test Hy ≈ Hy_ref +end + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$name", "test_") + @testset "$name" begin + getfield(@__MODULE__, name)() + end + end + end +end + +end + +TestSchurRank1.runtests()