Add low-rank Ai support for Schur complement

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!(

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")

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},

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

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)