From 8b44ceddd5907db2625e963e185ba505aff198b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Tue, 29 Jul 2025 16:10:55 +0200 Subject: [PATCH 01/15] Fix Schur with zero matrices --- src/model.jl | 35 +++++++++++++++++++++++++++-------- src/schur.jl | 15 +++++++++------ 2 files changed, 36 insertions(+), 14 deletions(-) diff --git a/src/model.jl b/src/model.jl index 6a0eb63..f119429 100644 --- a/src/model.jl +++ b/src/model.jl @@ -241,15 +241,26 @@ function buffer_for_jtprod(model::Model) return map(Base.Fix1(buffer_for_jtprod, model), matrix_indices(model)) end +_merge_sparsity(A::SparseArrays.SparseMatrixCSC, B::SparseArrays.SparseMatrixCSC) = abs.(A) + abs.(B) +_merge_sparsity(::FillArrays.Zeros, B::SparseArrays.SparseMatrixCSC) = B + function buffer_for_jtprod(model::Model, mat_idx::MatrixIndex) if iszero(model.meta.ncon) return end # FIXME: at some point, switch to dense - return sum(abs.(model.A[mat_idx.value, j]) for j in 1:model.meta.ncon) + return reduce(_merge_sparsity, model.A[mat_idx.value, j] for j in 1:model.meta.ncon) end # Computes `A .+= B * α` +function _add_mul!( + A::SparseArrays.SparseMatrixCSC, + ::FillArrays.Zeros, + _, +) + return A +end + function _add_mul!( A::SparseArrays.SparseMatrixCSC, B::SparseArrays.SparseMatrixCSC, @@ -323,6 +334,19 @@ function NLPModels.cons!( return cx end +function _add_vec!(_, _, _, _, offset, ::FillArrays.Zeros) + return offset +end + +function _add_vec!(I, J, V, j, offset, A::SparseArrays.SparseMatrixCSC) + Ai, Av = SparseArrays.findnz(A[:]) + K = offset .+ eachindex(Ai) + I[K] = Ai + J[K] .= j + V[K] = Av + return offset + length(Ai) +end + # `SparseMatrixCSC` is stored with an offset by column. # This means that getting view `view(A, :, I)` can be handles efficently, # these give `SparseMatrixCSCView` (if `I` is a `UnitRange`) and @@ -332,19 +356,14 @@ end # of `A` for constraint indices and the rows of `A` for matrix indices. function buffer_for_jprod(model::Model{T}, i::MatrixIndex) where {T} nnz = sum(1:model.meta.ncon; init = 0) do j - return SparseArrays.nnz(model.A[i.value, j]) + return _nnz(model.A[i.value, j]) end I = zeros(Int64, nnz) J = zeros(Int64, nnz) V = zeros(T, nnz) offset = 0 for j in 1:model.meta.ncon - Ai, Av = SparseArrays.findnz(model.A[i.value, j][:]) - K = offset .+ eachindex(Ai) - I[K] = Ai - J[K] .= j - V[K] = Av - offset += length(Ai) + offset = _add_vec!(I, J, V, j, offset, model.A[i.value, j]) end A = SparseArrays.sparse( I, diff --git a/src/schur.jl b/src/schur.jl index 8760ca6..ef0ca96 100644 --- a/src/schur.jl +++ b/src/schur.jl @@ -42,6 +42,9 @@ function _dot( return result end +_nnz(::FillArrays.Zeros) = 0 +_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} @@ -51,7 +54,7 @@ function buffer_for_schur_complement(model::Model{T}, κ) where {T} for mat_idx in matrix_indices(model) i = mat_idx.value - nzA = [SparseArrays.nnz(model.A[i, j]) for j in 1:n] + nzA = [_nnz(model.A[i, j]) for j in 1:n] σ[:, i] = sortperm(nzA, rev = true) sorted = nzA[σ[:, i]] @@ -89,7 +92,7 @@ function add_schur_complement!( for ii in axes(H, 1) i = σ[ii, ilmi] Ai = model.A[ilmi, i] - if SparseArrays.nnz(Ai) > 0 + if _nnz(Ai) > 0 if ii <= last_dense[ilmi] LinearAlgebra.mul!(AW[ilmi], W, Ai) LinearAlgebra.mul!(WAW[ilmi], AW[ilmi], W) @@ -100,11 +103,11 @@ function add_schur_complement!( H[i, j] = H[j, i] end else - if SparseArrays.nnz(Ai) > 1 + if _nnz(Ai) > 1 @inbounds for jj in ii:n j = σ[jj, ilmi] Aj = model.A[ilmi, j] - if !iszero(SparseArrays.nnz(Aj)) + if !iszero(_nnz(Aj)) ttt = _dot(Ai, Aj, W) H[i, j] += ttt if i != j @@ -112,7 +115,7 @@ function add_schur_complement!( end end end - elseif SparseArrays.nnz(Ai) == 1 + elseif _nnz(Ai) == 1 # A is symmetric iiiiAi = jjjiAi = only(SparseArrays.rowvals(Ai)) vvvi = only(SparseArrays.nonzeros(Ai)) @@ -122,7 +125,7 @@ function add_schur_complement!( # As we sort the matrices in decreasing `nnz` order, # the rest of matrices is either zero or have only # one entry - if !iszero(SparseArrays.nnz(Ajjj)) + if !iszero(_nnz(Ajjj)) iiijAj = jjjjAj = only(SparseArrays.rowvals(Ajjj)) vvvj = only(SparseArrays.nonzeros(Ajjj)) ttt = From 944c2d67e99881454d806776feea5df1e954f566 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 30 Jul 2025 10:59:33 +0200 Subject: [PATCH 02/15] Add buffered model --- src/LowRankOpt.jl | 1 + src/MOI_wrapper.jl | 17 +++- src/buffer.jl | 203 +++++++++++++++++++++++++++++++++++++++++++++ src/model.jl | 149 ++------------------------------- src/schur.jl | 41 ++++----- test/maxcut.jl | 8 +- 6 files changed, 247 insertions(+), 172 deletions(-) create mode 100644 src/buffer.jl diff --git a/src/LowRankOpt.jl b/src/LowRankOpt.jl index e9da0a7..1ccb06a 100644 --- a/src/LowRankOpt.jl +++ b/src/LowRankOpt.jl @@ -18,6 +18,7 @@ include("Test/Test.jl") include("Bridges/Bridges.jl") include("model.jl") +include("buffer.jl") include("schur.jl") include("MOI_wrapper.jl") include("BurerMonteiro.jl") diff --git a/src/MOI_wrapper.jl b/src/MOI_wrapper.jl index 049b334..b5c99d8 100644 --- a/src/MOI_wrapper.jl +++ b/src/MOI_wrapper.jl @@ -335,7 +335,19 @@ function MOI.copy_to( key in SOLVER_OPTIONS && key != "solver" ) dest.solver = dest.options["solver"](dest.model; options...) - return MOI.Utilities.identity_index_map(src) + index_map = MOI.Utilities.identity_index_map(src) + vis_src = MOI.get(src, MOI.ListOfVariableIndices()) + # Just to throw a nice error saying we don't support such attributes + # Filter them out since we already took care of them + no_obj = MOI.Utilities.ModelFilter(attr -> !(attr isa MOI.ObjectiveFunction || attr isa MOI.ObjectiveSense), src) + # We error for the rest. In the future, we should also take care of starting values + MOI.Utilities.pass_attributes(dest, no_obj, index_map) + MOI.Utilities.pass_attributes(dest, src, index_map, vis_src) + for (F, S) in constraint_types + cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F,S}()) + MOI.Utilities.pass_attributes(dest, src, index_map, cis_src) + end + return index_map end function MOI.copy_to(dest::Optimizer{T}, src::MOI.ModelLike) where {T} @@ -462,9 +474,6 @@ end struct Solution <: MOI.AbstractModelAttribute end MOI.is_set_by_optimize(::Solution) = true -function MOI.get(solver::SolverCore.AbstractOptimizationSolver, ::Solution) - return VectorizedSolution(solver.stats.solution, solver.model.dim) -end MOI.get(optimizer::Optimizer, attr::Solution) = MOI.get(optimizer.solver, attr) function MOI.get( diff --git a/src/buffer.jl b/src/buffer.jl new file mode 100644 index 0000000..98109d5 --- /dev/null +++ b/src/buffer.jl @@ -0,0 +1,203 @@ +mutable struct BufferedModelForSchur{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T},JB,JTB,SB} <: NLPModels.AbstractNLPModel{T,Vector{T}} + model::Model{T,C,A} + meta::NLPModels.NLPModelMeta{T,Vector{T}} + jprod_buffer::JB + jtprod_buffer::JTB + schur_buffer::SB +end + +function BufferedModelForSchur(model, datasparsity) + return BufferedModelForSchur( + model, + model.meta, + buffer_for_jprod(model), + buffer_for_jtprod(model), + buffer_for_schur_complement(model, datasparsity), + ) +end + +num_scalars(model::BufferedModelForSchur) = num_scalars(model.model) +num_matrices(model::BufferedModelForSchur) = num_matrices(model.model) +matrix_indices(model::BufferedModelForSchur) = matrix_indices(model.model) +side_dimension(model::BufferedModelForSchur, i) = side_dimension(model.model, i) +cons_constant(model::BufferedModelForSchur) = cons_constant(model.model) + +NLPModels.grad(model::BufferedModelForSchur, ::Type{ScalarIndex}) = NLPModels.grad(model.model, ScalarIndex) +NLPModels.grad(model::BufferedModelForSchur, i::MatrixIndex) = NLPModels.grad(model.model, i) +NLPModels.jac(model::BufferedModelForSchur, j::Integer, ::Type{ScalarIndex}) = NLPModels.jac(model.model, j, ScalarIndex) +norm_jac(model::BufferedModelForSchur, i::MatrixIndex) = norm_jac(model.model, i) + +errors(model::BufferedModelForSchur, x; kws...) = errors(model.model, x; kws...) + +####################### +###### Objective ###### +####################### + +function NLPModels.obj(model::BufferedModelForSchur, x::AbstractVector) + return NLPModels.obj(model.model, x) +end + +function dual_obj(model::BufferedModelForSchur, y::AbstractVector) + return dual_obj(model.model, y) +end + +####################### +###### J product ###### +####################### + +function buffer_for_jprod(model::Model{T}) where {T} + return SparseArrays.SparseMatrixCSC{T,Int64}[ + buffer_for_jprod(model, i) for i in matrix_indices(model) + ] +end + +function _add_vec!(_, _, _, _, offset, ::FillArrays.Zeros) + return offset +end + +function _add_vec!(I, J, V, j, offset, A::SparseArrays.SparseMatrixCSC) + Ai, Av = SparseArrays.findnz(A[:]) + K = offset .+ eachindex(Ai) + I[K] = Ai + J[K] .= j + V[K] = Av + return offset + length(Ai) +end + +# `SparseMatrixCSC` is stored with an offset by column. +# This means that getting view `view(A, :, I)` can be handles efficently, +# these give `SparseMatrixCSCView` (if `I` is a `UnitRange`) and +# `SparseMatrixCSCColumnSubset` otherwise. +# In `schur.jl`, we therefore get a `SparseMatrixCSCColumnSubset`. +# Since we want to use subsets of constraint indices, we use the columns +# of `A` for constraint indices and the rows of `A` for matrix indices. +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]) + 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]) + end + A = SparseArrays.sparse( + I, + J, + V, + side_dimension(model, i)^2, + model.meta.ncon, + ) + return A +end + +_vec(x::AbstractVector) = x +_vec(x::AbstractArray) = UnsafeArrays.uview(x, :) +_vec(x::Base.ReshapedArray) = _vec(parent(x)) + +function _add_jprod!(V, Jv::AbstractArray{T}, A) where {T} + return LinearAlgebra.mul!(Jv, A', _vec(V), true, true) +end + +function add_sub_jprod!( + model::BufferedModelForSchur, + i::MatrixIndex, + V::AbstractMatrix, + Jv::AbstractVector, + I, +) + # `view(cache, I)` would be terribly slow, only the number of elements of `I` matter here + A = model.jprod_buffer[i.value] + return _add_jprod!(V, Jv, view(A, :, I)) +end + +function add_jprod!( + model::BufferedModelForSchur, + V::AbstractMatrix, + Jv::AbstractVector, + i::MatrixIndex, +) + return _add_jprod!(V, Jv, model.jprod_buffer[i.value]) +end + +function NLPModels.jprod!(model::BufferedModelForSchur, x::AbstractVector, v::AbstractVector, Jv::AbstractVector) + return NLPModels.jprod!(model.model, x, v, Jv, model.jprod_buffer) +end + +function NLPModels.cons!( + model::BufferedModelForSchur, + x::AbstractVector, + cx::AbstractVector, +) + return NLPModels.cons!(model.model, x, cx, model.jprod_buffer) +end + +######################## +###### Jᵀ product ###### +######################## + +function jtprod!(model::BufferedModelForSchur, y::AbstractVector, vJ::AbstractVector, ::Type{ScalarIndex}) + jtprod!(model.model, y, vJ, ScalarIndex) +end + +function buffer_for_jtprod(model::Model) + if iszero(num_matrices(model)) + return + end + return map(Base.Fix1(buffer_for_jtprod, model), matrix_indices(model)) +end + +_merge_sparsity(A::SparseArrays.SparseMatrixCSC, B::SparseArrays.SparseMatrixCSC) = abs.(A) + abs.(B) +_merge_sparsity(::FillArrays.Zeros, B::SparseArrays.SparseMatrixCSC) = B +_merge_sparsity(A::SparseArrays.SparseMatrixCSC, ::FillArrays.Zeros) = A +_merge_sparsity(A::FillArrays.Zeros, ::FillArrays.Zeros) = A + +function buffer_for_jtprod(model::Model{T}, mat_idx::MatrixIndex) where {T} + if iszero(model.meta.ncon) + d = side_dimension(model, mat_idx) + return FillArrays.Zeros{T}(d, d) + end + # FIXME: at some point, switch to dense + return reduce(_merge_sparsity, model.A[mat_idx.value, j] for j in 1:model.meta.ncon) +end + +function NLPModels.jtprod!( + model::BufferedModelForSchur, + _::AbstractVector, + y::AbstractVector, + vJ::AbstractVector, +) + jtprod!(model, y, vJ[ScalarIndex], ScalarIndex) + for mat_idx in matrix_indices(model) + i = mat_idx.value + vJ[mat_idx] .= jtprod!(model, y, buffer[i], mat_idx) + end +end + +_zero!(A::FillArrays.Zeros) = A +_zero!(A::SparseArrays.SparseMatrixCSC) = fill!(SparseArrays.nonzeros(A), 0.0) + +function jtprod!(model::Model, y, buffer, mat_idx::MatrixIndex) + _zero!(buffer) + for j in eachindex(y) + _add_mul!(buffer, model.A[mat_idx.value, j], y[j]) + end + return buffer +end + +function jtprod!(model::BufferedModelForSchur, y, mat_idx::MatrixIndex) + return jtprod!(model.model, y, model.jtprod_buffer[mat_idx.value], mat_idx) +end + +function dual_cons!(model::BufferedModelForSchur, y::AbstractVector, res, ::Type{ScalarIndex}) + return dual_cons!(model.model, y, res, ScalarIndex) +end + +function dual_cons!( + model::BufferedModelForSchur, + y::AbstractVector, + i::MatrixIndex, +) + return model.model.C[i.value] - jtprod!(model, y, i) +end diff --git a/src/model.jl b/src/model.jl index f119429..a1c6f60 100644 --- a/src/model.jl +++ b/src/model.jl @@ -230,26 +230,8 @@ end dual_obj(model::Model, y::AbstractVector) = LinearAlgebra.dot(model.b, y) -function jtprod(model::Model, ::Type{ScalarIndex}, y::AbstractVector) - return model.C_lin' * y -end - -function buffer_for_jtprod(model::Model) - if iszero(num_matrices(model)) - return - end - return map(Base.Fix1(buffer_for_jtprod, model), matrix_indices(model)) -end - -_merge_sparsity(A::SparseArrays.SparseMatrixCSC, B::SparseArrays.SparseMatrixCSC) = abs.(A) + abs.(B) -_merge_sparsity(::FillArrays.Zeros, B::SparseArrays.SparseMatrixCSC) = B - -function buffer_for_jtprod(model::Model, mat_idx::MatrixIndex) - if iszero(model.meta.ncon) - return - end - # FIXME: at some point, switch to dense - return reduce(_merge_sparsity, model.A[mat_idx.value, j] for j in 1:model.meta.ncon) +function jtprod!(model::Model, y::AbstractVector, vJ::AbstractVector, ::Type{ScalarIndex}) + return LinearAlgebra.mul!(vJ, model.C_lin', y) end # Computes `A .+= B * α` @@ -280,42 +262,9 @@ function _add_mul!( end end -_zero!(A::SparseArrays.SparseMatrixCSC) = fill!(SparseArrays.nonzeros(A), 0.0) - -function jtprod!(buffer, model::Model, mat_idx::MatrixIndex, y) - _zero!(buffer) - for j in eachindex(y) - _add_mul!(buffer, model.A[mat_idx.value, j], y[j]) - end - return buffer -end - -function dual_cons(model::Model, ::Type{ScalarIndex}, y::AbstractVector) - return model.d_lin - jtprod(model, ScalarIndex, y) -end - -function dual_cons!( - buffer, - model::Model, - mat_idx::MatrixIndex, - y::AbstractVector, -) - i = mat_idx.value - return model.C[i] - jtprod!(buffer[i], model, mat_idx, y) -end - -function NLPModels.jtprod!( - model::Model, - _::AbstractVector, - y::AbstractVector, - vJ::AbstractVector, - buffer, -) - vJ[ScalarIndex] .= jtprod(model, ScalarIndex, y) - for mat_idx in matrix_indices(model) - i = mat_idx.value - vJ[mat_idx] .= jtprod!(buffer[i], model, mat_idx, y) - end +function dual_cons!(model::Model, y::AbstractVector, res, ::Type{ScalarIndex}) + copyto!(res, model.d_lin) + return LinearAlgebra.mul!(res, model.C_lin', y, -1, true) end NLPModels.grad(model::Model, ::Type{ScalarIndex}) = model.d_lin @@ -334,95 +283,11 @@ function NLPModels.cons!( return cx end -function _add_vec!(_, _, _, _, offset, ::FillArrays.Zeros) - return offset -end - -function _add_vec!(I, J, V, j, offset, A::SparseArrays.SparseMatrixCSC) - Ai, Av = SparseArrays.findnz(A[:]) - K = offset .+ eachindex(Ai) - I[K] = Ai - J[K] .= j - V[K] = Av - return offset + length(Ai) -end - -# `SparseMatrixCSC` is stored with an offset by column. -# This means that getting view `view(A, :, I)` can be handles efficently, -# these give `SparseMatrixCSCView` (if `I` is a `UnitRange`) and -# `SparseMatrixCSCColumnSubset` otherwise. -# In `schur.jl`, we therefore get a `SparseMatrixCSCColumnSubset`. -# Since we want to use subsets of constraint indices, we use the columns -# of `A` for constraint indices and the rows of `A` for matrix indices. -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]) - 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]) - end - A = SparseArrays.sparse( - I, - J, - V, - side_dimension(model, i)^2, - model.meta.ncon, - ) - return A -end - -# We define a new type so that we can define a custom `getindex` -struct JProdBuffer{T} - A::Vector{SparseArrays.SparseMatrixCSC{T,Int64}} -end - -function buffer_for_jprod(model::Model{T}) where {T} - return JProdBuffer([ - buffer_for_jprod(model, i) for i in matrix_indices(model) - ],) -end - -Base.getindex(buf::JProdBuffer, i::MatrixIndex) = buf.A[i.value] - -_vec(x::AbstractVector) = x -_vec(x::AbstractArray) = UnsafeArrays.uview(x, :) -_vec(x::Base.ReshapedArray) = _vec(parent(x)) - -function _add_jprod!(V, Jv::AbstractArray{T}, A) where {T} - return LinearAlgebra.mul!(Jv, A', _vec(V), true, true) -end - -function add_sub_jprod!( - _::Model, - _::MatrixIndex, - V::AbstractMatrix, - Jv::AbstractVector, - I, - A, -) - # `view(cache, I)` would be terribly slow, only the number of elements of `I` matter here - return _add_jprod!(V, Jv, view(A, :, I)) -end - -function add_jprod!( - ::Model, - ::MatrixIndex, - V::AbstractMatrix, - Jv::AbstractVector, - buffer, -) - return _add_jprod!(V, Jv, buffer) -end - function add_jprod!( model::Model, - i::MatrixIndex, V::AbstractMatrix, Jv::AbstractVector, + i::MatrixIndex, ) for j in 1:model.meta.ncon Jv[j] += LinearAlgebra.dot(model.A[i.value, j], V) @@ -438,7 +303,7 @@ function NLPModels.jprod!( ) where {N} LinearAlgebra.mul!(Jv, model.C_lin, v[ScalarIndex]) for i in matrix_indices(model) - add_jprod!(model, i, v[i], Jv, getindex.(args, i)...) + add_jprod!(model, v[i], Jv, i) end return Jv end diff --git a/src/schur.jl b/src/schur.jl index ef0ca96..83caef8 100644 --- a/src/schur.jl +++ b/src/schur.jl @@ -66,38 +66,36 @@ function buffer_for_schur_complement(model::Model{T}, κ) where {T} for dim in model.msizes] WAW = copy.(AW) - return buffer_for_jprod(model), AW, WAW, σ, last_dense + return AW, WAW, σ, last_dense end -function add_schur_complement!(model::Model, W, ::Type{MatrixIndex}, H, buffer) +function add_schur_complement!(model::BufferedModelForSchur, W, ::Type{MatrixIndex}, H) for i in matrix_indices(model) - add_schur_complement!(model, i, W[i], H, buffer) + add_schur_complement!(model, i, W[i], H) end return H end # /!\ W needs to be symmetric function add_schur_complement!( - model::Model, + model::BufferedModelForSchur, mat_idx::MatrixIndex, W::AbstractMatrix{T}, H, - buffer, ) where {T} - jprod_buffer, AW, WAW, σ, last_dense = buffer - buf = jprod_buffer[mat_idx] + AW, WAW, σ, last_dense = model.schur_buffer ilmi = mat_idx.value n = model.meta.ncon for ii in axes(H, 1) i = σ[ii, ilmi] - Ai = model.A[ilmi, i] + Ai = model.model.A[ilmi, 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, buf) + add_sub_jprod!(model, mat_idx, WAW[ilmi], view(H, I, i), I) for jj in (ii+1):n j = σ[jj, ilmi] H[i, j] = H[j, i] @@ -106,7 +104,7 @@ function add_schur_complement!( if _nnz(Ai) > 1 @inbounds for jj in ii:n j = σ[jj, ilmi] - Aj = model.A[ilmi, j] + Aj = model.model.A[ilmi, j] if !iszero(_nnz(Aj)) ttt = _dot(Ai, Aj, W) H[i, j] += ttt @@ -121,7 +119,7 @@ function add_schur_complement!( vvvi = only(SparseArrays.nonzeros(Ai)) @inbounds for jj in ii:n j = σ[jj, ilmi] - Ajjj = model.A[ilmi, j] + Ajjj = model.model.A[ilmi, j] # As we sort the matrices in decreasing `nnz` order, # the rest of matrices is either zero or have only # one entry @@ -146,18 +144,18 @@ function add_schur_complement!( return H end -function add_schur_complement!(model::Model, w, ::Type{ScalarIndex}, H) - H .+= model.C_lin * SparseArrays.spdiagm(w) * model.C_lin' +function add_schur_complement!(model::BufferedModelForSchur, w, ::Type{ScalarIndex}, H) + H .+= model.model.C_lin * SparseArrays.spdiagm(w) * model.model.C_lin' return H end # [HKS24, (5b)] # Returns the matrix equal to the sum, for each equation, of # ⟨A_i, WA_jW⟩ -function schur_complement!(model::Model, W::AbstractVector, H, buffer) +function schur_complement!(model::BufferedModelForSchur, W::AbstractVector, H) fill!(H, zero(eltype(H))) if num_matrices(model) > 0 - add_schur_complement!(model, W, MatrixIndex, H, buffer) + add_schur_complement!(model, W, MatrixIndex, H) end if num_scalars(model) > 0 add_schur_complement!(model, W[ScalarIndex], ScalarIndex, H) @@ -169,23 +167,20 @@ end # Returns the matrix equal to the sum, for each equation, of # ⟨A_i, WA(y)W⟩ function eval_schur_complement!( - result, - model::Model, + model::BufferedModelForSchur, W, y, - jprod_buffer, - jtprod_buffer, + result, ) fill!(result, zero(eltype(result))) for i in matrix_indices(model) add_jprod!( model, - i, - W[i] * jtprod!(jtprod_buffer[i.value], model, i, y) * W[i], + W[i] * jtprod!(model, y, i) * W[i], result, - jprod_buffer[i], + i, ) end - result .+= model.C_lin * (W[ScalarIndex] .* (model.C_lin' * y)) + result .+= model.model.C_lin * (W[ScalarIndex] .* (model.model.C_lin' * y)) return result end diff --git a/test/maxcut.jl b/test/maxcut.jl index 09befc5..77443eb 100644 --- a/test/maxcut.jl +++ b/test/maxcut.jl @@ -122,17 +122,19 @@ function schur_test(model::LRO.Model{T}, w, κ) where {T} H = zeros(n, n) H = LRO.schur_complement!(model, w, H, schur_buffer) Hy = similar(y) - LRO.eval_schur_complement!(Hy, model, w, y, schur_buffer[1], jtprod_buffer) + LRO.eval_schur_complement!(model, w, y, Hy) @test Hy ≈ H * y for i in LRO.matrix_indices(model) Wi = @inferred w[i] _alloc_schur_complement(model, i, Wi, H, schur_buffer) end for i in LRO.matrix_indices(model) - ret = LRO.dual_cons!(jtprod_buffer, model, i, y) + ret = LRO.dual_cons!(model, y, i) @test ret isa SparseMatrixCSC end - @test LRO.dual_cons(model, LRO.ScalarIndex, y) isa SparseArrays.SparseVector + dcons = ones(model.dim.num_scalars) + LRO.dual_cons(model, y, dcons, LRO.ScalarIndex) + @test dcons ≈ model.d_lin - model.C_lin' * y end function schur_test(model::LRO.Model{T}, κ) where {T} From 0d06e624f83d0b8f530d4adcd83e487b5e6797a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 30 Jul 2025 13:37:02 +0200 Subject: [PATCH 03/15] Fixes --- src/LowRankOpt.jl | 1 + src/buffer.jl | 78 +++++++++++------ src/model.jl | 192 +++++++++++------------------------------- src/solution.jl | 90 ++++++++++++++++++++ test/BurerMonteiro.jl | 2 +- test/maxcut.jl | 31 +++---- 6 files changed, 206 insertions(+), 188 deletions(-) create mode 100644 src/solution.jl diff --git a/src/LowRankOpt.jl b/src/LowRankOpt.jl index 1ccb06a..cf56cd8 100644 --- a/src/LowRankOpt.jl +++ b/src/LowRankOpt.jl @@ -17,6 +17,7 @@ include("distance_to_set.jl") include("Test/Test.jl") include("Bridges/Bridges.jl") +include("solution.jl") include("model.jl") include("buffer.jl") include("schur.jl") diff --git a/src/buffer.jl b/src/buffer.jl index 98109d5..4730734 100644 --- a/src/buffer.jl +++ b/src/buffer.jl @@ -1,4 +1,4 @@ -mutable struct BufferedModelForSchur{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T},JB,JTB,SB} <: NLPModels.AbstractNLPModel{T,Vector{T}} +mutable struct BufferedModelForSchur{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T},JB,JTB,SB} <: AbstractModel{T} model::Model{T,C,A} meta::NLPModels.NLPModelMeta{T,Vector{T}} jprod_buffer::JB @@ -20,14 +20,6 @@ num_scalars(model::BufferedModelForSchur) = num_scalars(model.model) num_matrices(model::BufferedModelForSchur) = num_matrices(model.model) matrix_indices(model::BufferedModelForSchur) = matrix_indices(model.model) side_dimension(model::BufferedModelForSchur, i) = side_dimension(model.model, i) -cons_constant(model::BufferedModelForSchur) = cons_constant(model.model) - -NLPModels.grad(model::BufferedModelForSchur, ::Type{ScalarIndex}) = NLPModels.grad(model.model, ScalarIndex) -NLPModels.grad(model::BufferedModelForSchur, i::MatrixIndex) = NLPModels.grad(model.model, i) -NLPModels.jac(model::BufferedModelForSchur, j::Integer, ::Type{ScalarIndex}) = NLPModels.jac(model.model, j, ScalarIndex) -norm_jac(model::BufferedModelForSchur, i::MatrixIndex) = norm_jac(model.model, i) - -errors(model::BufferedModelForSchur, x; kws...) = errors(model.model, x; kws...) ####################### ###### Objective ###### @@ -41,14 +33,25 @@ function dual_obj(model::BufferedModelForSchur, y::AbstractVector) return dual_obj(model.model, y) end +NLPModels.grad(model::BufferedModelForSchur, ::Type{ScalarIndex}) = NLPModels.grad(model.model, ScalarIndex) +NLPModels.grad(model::BufferedModelForSchur, i::MatrixIndex) = NLPModels.grad(model.model, i) + +######################### +###### Constraints ###### +######################### + +cons_constant(model::BufferedModelForSchur) = cons_constant(model.model) +NLPModels.jac(model::BufferedModelForSchur, j::Integer, ::Type{ScalarIndex}) = NLPModels.jac(model.model, j, ScalarIndex) +norm_jac(model::BufferedModelForSchur, i::MatrixIndex) = norm_jac(model.model, i) + +errors(model::BufferedModelForSchur, x; kws...) = errors(model.model, x; kws...) + ####################### ###### J product ###### ####################### -function buffer_for_jprod(model::Model{T}) where {T} - return SparseArrays.SparseMatrixCSC{T,Int64}[ - buffer_for_jprod(model, i) for i in matrix_indices(model) - ] +function jprod!(model, x, v, Jv, ::Type{ScalarIndex}) + return jprod!(model.model, x, v, Jv, ScalarIndex) end function _add_vec!(_, _, _, _, offset, ::FillArrays.Zeros) @@ -92,6 +95,12 @@ function buffer_for_jprod(model::Model{T}, i::MatrixIndex) where {T} return A end +function buffer_for_jprod(model::Model{T}) where {T} + return SparseArrays.SparseMatrixCSC{T,Int64}[ + buffer_for_jprod(model, i) for i in matrix_indices(model) + ] +end + _vec(x::AbstractVector) = x _vec(x::AbstractArray) = UnsafeArrays.uview(x, :) _vec(x::Base.ReshapedArray) = _vec(parent(x)) @@ -121,18 +130,6 @@ function add_jprod!( return _add_jprod!(V, Jv, model.jprod_buffer[i.value]) end -function NLPModels.jprod!(model::BufferedModelForSchur, x::AbstractVector, v::AbstractVector, Jv::AbstractVector) - return NLPModels.jprod!(model.model, x, v, Jv, model.jprod_buffer) -end - -function NLPModels.cons!( - model::BufferedModelForSchur, - x::AbstractVector, - cx::AbstractVector, -) - return NLPModels.cons!(model.model, x, cx, model.jprod_buffer) -end - ######################## ###### Jᵀ product ###### ######################## @@ -170,14 +167,41 @@ function NLPModels.jtprod!( ) jtprod!(model, y, vJ[ScalarIndex], ScalarIndex) for mat_idx in matrix_indices(model) - i = mat_idx.value - vJ[mat_idx] .= jtprod!(model, y, buffer[i], mat_idx) + vJ[mat_idx] .= jtprod!(model, y, mat_idx) end end _zero!(A::FillArrays.Zeros) = A _zero!(A::SparseArrays.SparseMatrixCSC) = fill!(SparseArrays.nonzeros(A), 0.0) +# Computes `A .+= B * α` +function _add_mul!( + A::SparseArrays.SparseMatrixCSC, + ::FillArrays.Zeros, + _, +) + return A +end + +function _add_mul!( + A::SparseArrays.SparseMatrixCSC, + B::SparseArrays.SparseMatrixCSC, + α, +) + for col in axes(A, 2) + range_A = SparseArrays.nzrange(A, col) + it_A = iterate(range_A) + for k in SparseArrays.nzrange(B, col) + row_B = SparseArrays.rowvals(B)[k] + while SparseArrays.rowvals(A)[it_A[1]] < row_B + it_A = iterate(range_A, it_A[2]) + end + @assert row_B == SparseArrays.rowvals(A)[it_A[1]] + SparseArrays.nonzeros(A)[it_A[1]] += SparseArrays.nonzeros(B)[k] * α + end + end +end + function jtprod!(model::Model, y, buffer, mat_idx::MatrixIndex) _zero!(buffer) for j in eachindex(y) diff --git a/src/model.jl b/src/model.jl index a1c6f60..1c975e5 100644 --- a/src/model.jl +++ b/src/model.jl @@ -5,97 +5,31 @@ import MathOptInterface as MOI import NLPModels import UnsafeArrays -struct Dimensions - num_scalars::Int64 - side_dimensions::Vector{Int64} - offsets::Vector{Int64} -end - -num_matrices(d::Dimensions) = length(d.side_dimensions) -Base.length(d::Dimensions) = d.offsets[end] - -struct ScalarIndex - value::Int64 -end - -struct MatrixIndex - value::Int64 -end -Base.broadcastable(i::MatrixIndex) = Ref(i) +abstract type AbstractModel{T} <: NLPModels.AbstractNLPModel{T,Vector{T}} end -abstract type AbstractSolution{T} <: AbstractVector{T} end - -function Base.getindex( - s::AbstractSolution, - i::Union{Type{ScalarIndex},MatrixIndex}, +function NLPModels.cons!( + model::AbstractModel, + x::AbstractVector, + cx::AbstractVector, ) - return view(s, i) -end - -struct VectorizedSolution{T} <: AbstractSolution{T} - x::Vector{T} - dim::Dimensions -end - -function LinearAlgebra.dot(x::VectorizedSolution, z::VectorizedSolution) - return LinearAlgebra.dot(x.x, z.x) -end - -num_matrices(x::VectorizedSolution) = num_matrices(x.dim) - -Base.similar(s::VectorizedSolution) = VectorizedSolution(similar(s.x), s.dim) - -Base.size(s::VectorizedSolution) = (length(s.dim),) - -function Base.to_index(s::VectorizedSolution, ::Type{ScalarIndex}) - return Base.OneTo(s.dim.num_scalars) -end - -function Base.to_index(s::VectorizedSolution, mi::MatrixIndex) - i = mi.value - return (1+s.dim.offsets[i]):s.dim.offsets[i+1] -end - -function Base.view(s::VectorizedSolution, ::Type{ScalarIndex}) - return view(s.x, Base.to_index(s, ScalarIndex)) -end - -# `s[i] .= ...` calleds `copyto!(view(s, i), Broadcasted(...))` -function Base.view(s::VectorizedSolution, i::MatrixIndex) - v = view(s.x, Base.to_index(s, i)) - dim = s.dim.side_dimensions[i.value] - X = reshape(v, dim, dim) - return X -end - -Base.setindex!(s::VectorizedSolution, v, i::Integer) = setindex!(s.x, v, i) -Base.getindex(s::VectorizedSolution, i::Integer) = getindex(s.x, i) - -struct ShapedSolution{T,MT<:AbstractMatrix{T}} <: AbstractSolution{T} - scalars::Vector{T} - matrices::Vector{MT} -end - -function Base.size(s::ShapedSolution) - return (length(s.scalars) + sum(length, s.matrices, init = 0),) -end -num_matrices(s::ShapedSolution) = length(s.matrices) - -function LinearAlgebra.norm2(s::ShapedSolution{T}) where {T} - # `LinearAlgebra.generic_norm2` starts by computing the ∞ norm and do a rescaling, we don't do that here - return √(LinearAlgebra.dot(s, s)) + NLPModels.jprod!(model, x, x, cx) + cx .-= cons_constant(model) + return cx end -function LinearAlgebra.dot(a::ShapedSolution{T}, b::ShapedSolution{T}) where {T} - return LinearAlgebra.dot(a.scalars, b.scalars) + - sum(eachindex(a.matrices); init = zero(T)) do i - return LinearAlgebra.dot(a.matrices[i], b.matrices[i]) +function NLPModels.jprod!( + model::AbstractModel, + x::AbstractVector, + v::AbstractVector, + Jv::AbstractVector, +) + jprod!(model, x, v[ScalarIndex], Jv, ScalarIndex) + for i in matrix_indices(model) + add_jprod!(model, v[i], Jv, i) end + return Jv end -Base.view(s::ShapedSolution, ::Type{ScalarIndex}) = s.scalars -Base.view(s::ShapedSolution, i::MatrixIndex) = s.matrices[i.value] - """ Model @@ -126,8 +60,7 @@ The fields of the `struct` as related to the arrays of the above formulation as * The matrix ``C_i`` is given by `C[i]`. * The matrix ``A_{i,j}`` is given by `A[i,j]`. """ -mutable struct Model{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T}} <: - NLPModels.AbstractNLPModel{T,Vector{T}} +mutable struct Model{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T}} <: AbstractModel{T} meta::NLPModels.NLPModelMeta{T,Vector{T}} dim::Dimensions C::Vector{C} @@ -184,20 +117,9 @@ end side_dimension(model::Model, i::MatrixIndex) = model.msizes[i.value] -# Should be only used with `norm` -NLPModels.jac(model::Model, ::Type{ScalarIndex}) = model.C_lin -function NLPModels.jac(model::Model, j::Integer, i::MatrixIndex) - return model.A[i.value, j] -end -function NLPModels.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, :]) -end +####################### +###### Objective ###### +####################### function NLPModels.obj(model::Model, X::AbstractMatrix, i::MatrixIndex) return LinearAlgebra.dot(model.C[i.value], X) @@ -234,34 +156,6 @@ function jtprod!(model::Model, y::AbstractVector, vJ::AbstractVector, ::Type{Sca return LinearAlgebra.mul!(vJ, model.C_lin', y) end -# Computes `A .+= B * α` -function _add_mul!( - A::SparseArrays.SparseMatrixCSC, - ::FillArrays.Zeros, - _, -) - return A -end - -function _add_mul!( - A::SparseArrays.SparseMatrixCSC, - B::SparseArrays.SparseMatrixCSC, - α, -) - for col in axes(A, 2) - range_A = SparseArrays.nzrange(A, col) - it_A = iterate(range_A) - for k in SparseArrays.nzrange(B, col) - row_B = SparseArrays.rowvals(B)[k] - while SparseArrays.rowvals(A)[it_A[1]] < row_B - it_A = iterate(range_A, it_A[2]) - end - @assert row_B == SparseArrays.rowvals(A)[it_A[1]] - SparseArrays.nonzeros(A)[it_A[1]] += SparseArrays.nonzeros(B)[k] * α - end - end -end - function dual_cons!(model::Model, y::AbstractVector, res, ::Type{ScalarIndex}) copyto!(res, model.d_lin) return LinearAlgebra.mul!(res, model.C_lin', y, -1, true) @@ -270,19 +164,31 @@ end NLPModels.grad(model::Model, ::Type{ScalarIndex}) = model.d_lin NLPModels.grad(model::Model, i::MatrixIndex) = model.C[i.value] +######################### +###### Constraints ###### +######################### + cons_constant(model::Model) = model.b -function NLPModels.cons!( - model::Model, - x::AbstractVector, - cx::AbstractVector, - args::Vararg{Any,N}, -) where {N} - NLPModels.jprod!(model, x, x, cx, args...) - cx .-= model.b - return cx +# Should be only used with `norm` +NLPModels.jac(model::Model, ::Type{ScalarIndex}) = model.C_lin +function NLPModels.jac(model::Model, j::Integer, i::MatrixIndex) + return model.A[i.value, j] +end +function NLPModels.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, :]) end +####################### +###### J product ###### +####################### + function add_jprod!( model::Model, V::AbstractMatrix, @@ -294,16 +200,12 @@ function add_jprod!( end end -function NLPModels.jprod!( +function jprod!( model::Model, _::AbstractVector, v::AbstractVector, Jv::AbstractVector, - args::Vararg{Any,N}, # Optional buffer -) where {N} - LinearAlgebra.mul!(Jv, model.C_lin, v[ScalarIndex]) - for i in matrix_indices(model) - add_jprod!(model, v[i], Jv, i) - end - return Jv + ::Type{ScalarIndex}, +) + return LinearAlgebra.mul!(Jv, model.C_lin, v) end diff --git a/src/solution.jl b/src/solution.jl new file mode 100644 index 0000000..3009403 --- /dev/null +++ b/src/solution.jl @@ -0,0 +1,90 @@ +struct Dimensions + num_scalars::Int64 + side_dimensions::Vector{Int64} + offsets::Vector{Int64} +end + +num_matrices(d::Dimensions) = length(d.side_dimensions) +Base.length(d::Dimensions) = d.offsets[end] + +struct ScalarIndex + value::Int64 +end + +struct MatrixIndex + value::Int64 +end +Base.broadcastable(i::MatrixIndex) = Ref(i) + +abstract type AbstractSolution{T} <: AbstractVector{T} end + +function Base.getindex( + s::AbstractSolution, + i::Union{Type{ScalarIndex},MatrixIndex}, +) + return view(s, i) +end + +struct VectorizedSolution{T} <: AbstractSolution{T} + x::Vector{T} + dim::Dimensions +end + +function LinearAlgebra.dot(x::VectorizedSolution, z::VectorizedSolution) + return LinearAlgebra.dot(x.x, z.x) +end + +num_matrices(x::VectorizedSolution) = num_matrices(x.dim) + +Base.similar(s::VectorizedSolution) = VectorizedSolution(similar(s.x), s.dim) + +Base.size(s::VectorizedSolution) = (length(s.dim),) + +function Base.to_index(s::VectorizedSolution, ::Type{ScalarIndex}) + return Base.OneTo(s.dim.num_scalars) +end + +function Base.to_index(s::VectorizedSolution, mi::MatrixIndex) + i = mi.value + return (1+s.dim.offsets[i]):s.dim.offsets[i+1] +end + +function Base.view(s::VectorizedSolution, ::Type{ScalarIndex}) + return view(s.x, Base.to_index(s, ScalarIndex)) +end + +# `s[i] .= ...` calleds `copyto!(view(s, i), Broadcasted(...))` +function Base.view(s::VectorizedSolution, i::MatrixIndex) + v = view(s.x, Base.to_index(s, i)) + dim = s.dim.side_dimensions[i.value] + X = reshape(v, dim, dim) + return X +end + +Base.setindex!(s::VectorizedSolution, v, i::Integer) = setindex!(s.x, v, i) +Base.getindex(s::VectorizedSolution, i::Integer) = getindex(s.x, i) + +struct ShapedSolution{T,MT<:AbstractMatrix{T}} <: AbstractSolution{T} + scalars::Vector{T} + matrices::Vector{MT} +end + +function Base.size(s::ShapedSolution) + return (length(s.scalars) + sum(length, s.matrices, init = 0),) +end +num_matrices(s::ShapedSolution) = length(s.matrices) + +function LinearAlgebra.norm2(s::ShapedSolution{T}) where {T} + # `LinearAlgebra.generic_norm2` starts by computing the ∞ norm and do a rescaling, we don't do that here + return √(LinearAlgebra.dot(s, s)) +end + +function LinearAlgebra.dot(a::ShapedSolution{T}, b::ShapedSolution{T}) where {T} + return LinearAlgebra.dot(a.scalars, b.scalars) + + sum(eachindex(a.matrices); init = zero(T)) do i + return LinearAlgebra.dot(a.matrices[i], b.matrices[i]) + end +end + +Base.view(s::ShapedSolution, ::Type{ScalarIndex}) = s.scalars +Base.view(s::ShapedSolution, i::MatrixIndex) = s.matrices[i.value] diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl index 458344c..caa678a 100644 --- a/test/BurerMonteiro.jl +++ b/test/BurerMonteiro.jl @@ -182,7 +182,7 @@ end; ) @test model.meta.ncon == 0 @test LRO.norm_jac(model, LRO.MatrixIndex(1)) == 0 - @test isnothing(LRO.buffer_for_jtprod(model, LRO.MatrixIndex(1))) + @test LRO.buffer_for_jtprod(model, LRO.MatrixIndex(1)) isa LRO.FillArrays.Zeros end; @testset "Fallback" begin diff --git a/test/maxcut.jl b/test/maxcut.jl index 77443eb..a9953de 100644 --- a/test/maxcut.jl +++ b/test/maxcut.jl @@ -94,47 +94,48 @@ function ConvexSolver(model::LRO.Model) return ConvexSolver(model, stats) end +function LRO.MOI.get(solver::ConvexSolver, ::LRO.Solution) + return LRO.VectorizedSolution(solver.stats.solution, solver.model.dim) +end + function SolverCore.solve!(::ConvexSolver, ::LRO.Model) return end -function _alloc_schur_complement(model, i, Wi, H, schur_buffer) +function _alloc_schur_complement(model, i, Wi, H) if VERSION < v"1.11" return end - LRO.add_schur_complement!(model, i, Wi, H, schur_buffer) - @test 0 == - @allocated LRO.add_schur_complement!(model, i, Wi, H, schur_buffer) + LRO.add_schur_complement!(model, i, Wi, H) + @test 0 == @allocated LRO.add_schur_complement!(model, i, Wi, H) end -function schur_test(model::LRO.Model{T}, w, κ) where {T} - schur_buffer = LRO.buffer_for_schur_complement(model, κ) - jtprod_buffer = LRO.buffer_for_jtprod(model) +function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} n = model.meta.ncon y = rand(T, n) Jv = similar(y) vJ = similar(w) - NLPModels.jprod!(model, w, w, Jv, schur_buffer[1]) - NLPModels.jtprod!(model, w, y, vJ, jtprod_buffer) + NLPModels.jprod!(model, w, w, Jv) + NLPModels.jtprod!(model, w, y, vJ) @test dot(Jv, y) ≈ dot(vJ, w) H = zeros(n, n) - H = LRO.schur_complement!(model, w, H, schur_buffer) + H = LRO.schur_complement!(model, w, H) Hy = similar(y) LRO.eval_schur_complement!(model, w, y, Hy) @test Hy ≈ H * y for i in LRO.matrix_indices(model) Wi = @inferred w[i] - _alloc_schur_complement(model, i, Wi, H, schur_buffer) + _alloc_schur_complement(model, i, Wi, H) end for i in LRO.matrix_indices(model) ret = LRO.dual_cons!(model, y, i) @test ret isa SparseMatrixCSC end - dcons = ones(model.dim.num_scalars) - LRO.dual_cons(model, y, dcons, LRO.ScalarIndex) - @test dcons ≈ model.d_lin - model.C_lin' * y + dcons = ones(LRO.num_scalars(model)) + LRO.dual_cons!(model, y, dcons, LRO.ScalarIndex) + @test dcons ≈ model.model.d_lin - model.model.C_lin' * y end function schur_test(model::LRO.Model{T}, κ) where {T} @@ -143,7 +144,7 @@ function schur_test(model::LRO.Model{T}, κ) where {T} for i in LRO.matrix_indices(model) W[i] .= W[i] .+ W[i]' end - return schur_test(model, W, κ) + return schur_test(LRO.BufferedModelForSchur(model, κ), W) end @testset "ConvexSolver $T" for T in [Float32, Float64] From da74a3ecc83703859334ee41c7f7d2a806ca5103 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 30 Jul 2025 14:19:31 +0200 Subject: [PATCH 04/15] Fix --- src/buffer.jl | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/buffer.jl b/src/buffer.jl index 4730734..76f69d8 100644 --- a/src/buffer.jl +++ b/src/buffer.jl @@ -145,18 +145,24 @@ function buffer_for_jtprod(model::Model) return map(Base.Fix1(buffer_for_jtprod, model), matrix_indices(model)) end -_merge_sparsity(A::SparseArrays.SparseMatrixCSC, B::SparseArrays.SparseMatrixCSC) = abs.(A) + abs.(B) +_merge_sparsity(A::SparseArrays.SparseMatrixCSC, B::SparseArrays.SparseMatrixCSC) = A + B _merge_sparsity(::FillArrays.Zeros, B::SparseArrays.SparseMatrixCSC) = B _merge_sparsity(A::SparseArrays.SparseMatrixCSC, ::FillArrays.Zeros) = A _merge_sparsity(A::FillArrays.Zeros, ::FillArrays.Zeros) = A +_abs(A::SparseArrays.SparseMatrixCSC) = abs.(A) +_abs(A::FillArrays.Zeros) = A + function buffer_for_jtprod(model::Model{T}, mat_idx::MatrixIndex) where {T} if iszero(model.meta.ncon) d = side_dimension(model, mat_idx) return FillArrays.Zeros{T}(d, d) end # FIXME: at some point, switch to dense - return reduce(_merge_sparsity, model.A[mat_idx.value, j] for j in 1:model.meta.ncon) + # /!\ If there is only one nonzero matrix and we didn't have `_abs`, + # we would return an alias of that only matrix so that `_abs` has the + # non-obvious role of avoid this as well as avoiding cancellations. + return reduce(_merge_sparsity, _abs(model.A[mat_idx.value, j]) for j in 1:model.meta.ncon) end function NLPModels.jtprod!( From e2a1d65b96493c01029c1f9ab96fc2625f6b0700 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 30 Jul 2025 14:21:44 +0200 Subject: [PATCH 05/15] Fix format --- src/MOI_wrapper.jl | 7 ++++- src/buffer.jl | 63 ++++++++++++++++++++++++++++++++----------- src/model.jl | 10 +++++-- src/schur.jl | 28 +++++++++---------- test/BurerMonteiro.jl | 3 ++- 5 files changed, 78 insertions(+), 33 deletions(-) diff --git a/src/MOI_wrapper.jl b/src/MOI_wrapper.jl index b5c99d8..3f34c24 100644 --- a/src/MOI_wrapper.jl +++ b/src/MOI_wrapper.jl @@ -339,7 +339,12 @@ function MOI.copy_to( vis_src = MOI.get(src, MOI.ListOfVariableIndices()) # Just to throw a nice error saying we don't support such attributes # Filter them out since we already took care of them - no_obj = MOI.Utilities.ModelFilter(attr -> !(attr isa MOI.ObjectiveFunction || attr isa MOI.ObjectiveSense), src) + no_obj = MOI.Utilities.ModelFilter( + attr -> !( + attr isa MOI.ObjectiveFunction || attr isa MOI.ObjectiveSense + ), + src, + ) # We error for the rest. In the future, we should also take care of starting values MOI.Utilities.pass_attributes(dest, no_obj, index_map) MOI.Utilities.pass_attributes(dest, src, index_map, vis_src) diff --git a/src/buffer.jl b/src/buffer.jl index 76f69d8..8eb5f99 100644 --- a/src/buffer.jl +++ b/src/buffer.jl @@ -1,4 +1,11 @@ -mutable struct BufferedModelForSchur{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T},JB,JTB,SB} <: AbstractModel{T} +mutable struct BufferedModelForSchur{ + T, + C<:AbstractMatrix{T}, + A<:AbstractMatrix{T}, + JB, + JTB, + SB, +} <: AbstractModel{T} model::Model{T,C,A} meta::NLPModels.NLPModelMeta{T,Vector{T}} jprod_buffer::JB @@ -33,16 +40,28 @@ function dual_obj(model::BufferedModelForSchur, y::AbstractVector) return dual_obj(model.model, y) end -NLPModels.grad(model::BufferedModelForSchur, ::Type{ScalarIndex}) = NLPModels.grad(model.model, ScalarIndex) -NLPModels.grad(model::BufferedModelForSchur, i::MatrixIndex) = NLPModels.grad(model.model, i) +function NLPModels.grad(model::BufferedModelForSchur, ::Type{ScalarIndex}) + return NLPModels.grad(model.model, ScalarIndex) +end +function NLPModels.grad(model::BufferedModelForSchur, i::MatrixIndex) + return NLPModels.grad(model.model, i) +end ######################### ###### Constraints ###### ######################### cons_constant(model::BufferedModelForSchur) = cons_constant(model.model) -NLPModels.jac(model::BufferedModelForSchur, j::Integer, ::Type{ScalarIndex}) = NLPModels.jac(model.model, j, ScalarIndex) -norm_jac(model::BufferedModelForSchur, i::MatrixIndex) = norm_jac(model.model, i) +function NLPModels.jac( + model::BufferedModelForSchur, + j::Integer, + ::Type{ScalarIndex}, +) + return NLPModels.jac(model.model, j, ScalarIndex) +end +function norm_jac(model::BufferedModelForSchur, i::MatrixIndex) + return norm_jac(model.model, i) +end errors(model::BufferedModelForSchur, x; kws...) = errors(model.model, x; kws...) @@ -134,8 +153,13 @@ end ###### Jᵀ product ###### ######################## -function jtprod!(model::BufferedModelForSchur, y::AbstractVector, vJ::AbstractVector, ::Type{ScalarIndex}) - jtprod!(model.model, y, vJ, ScalarIndex) +function jtprod!( + model::BufferedModelForSchur, + y::AbstractVector, + vJ::AbstractVector, + ::Type{ScalarIndex}, +) + return jtprod!(model.model, y, vJ, ScalarIndex) end function buffer_for_jtprod(model::Model) @@ -145,7 +169,12 @@ function buffer_for_jtprod(model::Model) return map(Base.Fix1(buffer_for_jtprod, model), matrix_indices(model)) end -_merge_sparsity(A::SparseArrays.SparseMatrixCSC, B::SparseArrays.SparseMatrixCSC) = A + B +function _merge_sparsity( + A::SparseArrays.SparseMatrixCSC, + B::SparseArrays.SparseMatrixCSC, +) + return A + B +end _merge_sparsity(::FillArrays.Zeros, B::SparseArrays.SparseMatrixCSC) = B _merge_sparsity(A::SparseArrays.SparseMatrixCSC, ::FillArrays.Zeros) = A _merge_sparsity(A::FillArrays.Zeros, ::FillArrays.Zeros) = A @@ -162,7 +191,10 @@ function buffer_for_jtprod(model::Model{T}, mat_idx::MatrixIndex) where {T} # /!\ If there is only one nonzero matrix and we didn't have `_abs`, # we would return an alias of that only matrix so that `_abs` has the # non-obvious role of avoid this as well as avoiding cancellations. - return reduce(_merge_sparsity, _abs(model.A[mat_idx.value, j]) for j in 1:model.meta.ncon) + return reduce( + _merge_sparsity, + _abs(model.A[mat_idx.value, j]) for j in 1:model.meta.ncon + ) end function NLPModels.jtprod!( @@ -181,11 +213,7 @@ _zero!(A::FillArrays.Zeros) = A _zero!(A::SparseArrays.SparseMatrixCSC) = fill!(SparseArrays.nonzeros(A), 0.0) # Computes `A .+= B * α` -function _add_mul!( - A::SparseArrays.SparseMatrixCSC, - ::FillArrays.Zeros, - _, -) +function _add_mul!(A::SparseArrays.SparseMatrixCSC, ::FillArrays.Zeros, _) return A end @@ -220,7 +248,12 @@ function jtprod!(model::BufferedModelForSchur, y, mat_idx::MatrixIndex) return jtprod!(model.model, y, model.jtprod_buffer[mat_idx.value], mat_idx) end -function dual_cons!(model::BufferedModelForSchur, y::AbstractVector, res, ::Type{ScalarIndex}) +function dual_cons!( + model::BufferedModelForSchur, + y::AbstractVector, + res, + ::Type{ScalarIndex}, +) return dual_cons!(model.model, y, res, ScalarIndex) end diff --git a/src/model.jl b/src/model.jl index 1c975e5..6150ea7 100644 --- a/src/model.jl +++ b/src/model.jl @@ -60,7 +60,8 @@ The fields of the `struct` as related to the arrays of the above formulation as * The matrix ``C_i`` is given by `C[i]`. * The matrix ``A_{i,j}`` is given by `A[i,j]`. """ -mutable struct Model{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T}} <: AbstractModel{T} +mutable struct Model{T,C<:AbstractMatrix{T},A<:AbstractMatrix{T}} <: + AbstractModel{T} meta::NLPModels.NLPModelMeta{T,Vector{T}} dim::Dimensions C::Vector{C} @@ -152,7 +153,12 @@ end dual_obj(model::Model, y::AbstractVector) = LinearAlgebra.dot(model.b, y) -function jtprod!(model::Model, y::AbstractVector, vJ::AbstractVector, ::Type{ScalarIndex}) +function jtprod!( + model::Model, + y::AbstractVector, + vJ::AbstractVector, + ::Type{ScalarIndex}, +) return LinearAlgebra.mul!(vJ, model.C_lin', y) end diff --git a/src/schur.jl b/src/schur.jl index 83caef8..68e4990 100644 --- a/src/schur.jl +++ b/src/schur.jl @@ -69,7 +69,12 @@ function buffer_for_schur_complement(model::Model{T}, κ) where {T} return AW, WAW, σ, last_dense end -function add_schur_complement!(model::BufferedModelForSchur, W, ::Type{MatrixIndex}, H) +function add_schur_complement!( + model::BufferedModelForSchur, + W, + ::Type{MatrixIndex}, + H, +) for i in matrix_indices(model) add_schur_complement!(model, i, W[i], H) end @@ -144,7 +149,12 @@ function add_schur_complement!( return H end -function add_schur_complement!(model::BufferedModelForSchur, w, ::Type{ScalarIndex}, H) +function add_schur_complement!( + model::BufferedModelForSchur, + w, + ::Type{ScalarIndex}, + H, +) H .+= model.model.C_lin * SparseArrays.spdiagm(w) * model.model.C_lin' return H end @@ -166,20 +176,10 @@ end # [HKS24, (5b)] # Returns the matrix equal to the sum, for each equation, of # ⟨A_i, WA(y)W⟩ -function eval_schur_complement!( - model::BufferedModelForSchur, - W, - y, - result, -) +function eval_schur_complement!(model::BufferedModelForSchur, W, y, result) fill!(result, zero(eltype(result))) for i in matrix_indices(model) - add_jprod!( - model, - W[i] * jtprod!(model, y, i) * W[i], - result, - i, - ) + add_jprod!(model, W[i] * jtprod!(model, y, i) * W[i], result, i) end result .+= model.model.C_lin * (W[ScalarIndex] .* (model.model.C_lin' * y)) return result diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl index caa678a..e38f242 100644 --- a/test/BurerMonteiro.jl +++ b/test/BurerMonteiro.jl @@ -182,7 +182,8 @@ end; ) @test model.meta.ncon == 0 @test LRO.norm_jac(model, LRO.MatrixIndex(1)) == 0 - @test LRO.buffer_for_jtprod(model, LRO.MatrixIndex(1)) isa LRO.FillArrays.Zeros + @test LRO.buffer_for_jtprod(model, LRO.MatrixIndex(1)) isa + LRO.FillArrays.Zeros end; @testset "Fallback" begin From be7b84d6f74641ce19cb409a730ccca1c8552d04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 30 Jul 2025 14:46:32 +0200 Subject: [PATCH 06/15] Add tests --- test/buffer.jl | 46 +++++++++++++++++++++++++++++ test/diff_check.jl | 73 +++++++++++++++++++++++++++++++++++++++++++--- test/maxcut.jl | 63 --------------------------------------- 3 files changed, 115 insertions(+), 67 deletions(-) create mode 100644 test/buffer.jl diff --git a/test/buffer.jl b/test/buffer.jl new file mode 100644 index 0000000..65ad9ab --- /dev/null +++ b/test/buffer.jl @@ -0,0 +1,46 @@ +module TestBuffer + +import FillArrays, SparseArrays +using JuMP, Dualization +include("diff_check.jl") + +# Test with zero Ai matrices +function test_zero_Ai() + model = Model(dual_optimizer(LRO.Optimizer)) + @variable(model, x[1:2] in MOI.Nonnegatives(2)) + @variable(model, X[1:2, 1:2] in PSDCone()) + @constraint(model, sum(x) == 1) + @constraint(model, 2sum(x) == 2) + @constraint(model, sum(X) == 2) + @constraint(model, x[1] - x[2] == 1) + @objective(model, Max, x[1]) + set_attribute(model, "solver", ConvexSolver) + optimize!(model) + b = _backend(model) + T = Float64 + Z = FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}} + S = SparseArrays.SparseMatrixCSC{T,Int} + @test b.model.A isa Matrix{Union{Z,S}} + @test b.model.A[1] isa Z + @test b.model.A[2] isa Z + @test b.model.A[3] isa S + @test b.model.A[4] isa Z + buf = LRO.BufferedModelForSchur(b.model, 1) + for A in b.model.A + @test buf.jtprod_buffer[] !== A + end +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 + +TestBuffer.runtests() diff --git a/test/diff_check.jl b/test/diff_check.jl index 94d5567..0aaa9c4 100644 --- a/test/diff_check.jl +++ b/test/diff_check.jl @@ -4,8 +4,10 @@ # in the LICENSE.md file or at https://opensource.org/licenses/MIT. using Test +import SolverCore +import LowRankOpt as LRO -function test_vecprod(f, len, J; tol = 1e-6) +function _test_vecprod(f, len, J; tol = 1e-6) v = ones(len) @test f(v) ≈ J * v rtol = tol atol = tol v = -ones(len) @@ -33,7 +35,7 @@ function jac_check(model, x; kws...) f(x) = NLPModels.cons(model, x) J = FiniteDiff.finite_difference_jacobian(f, x) @testset "jprod" begin - test_vecprod( + _test_vecprod( v -> NLPModels.jprod(model, x, v), model.meta.nvar, J; @@ -41,7 +43,7 @@ function jac_check(model, x; kws...) ) end @testset "jtprod" begin - test_vecprod( + _test_vecprod( v -> NLPModels.jtprod(model, x, v), model.meta.ncon, J'; @@ -56,7 +58,7 @@ function hess_check(model, x; kws...) f(x) = obj_weight * NLPModels.obj(model, x) - dot(y, NLPModels.cons(model, x)) J = FiniteDiff.finite_difference_hessian(f, x) - return test_vecprod( + return _test_vecprod( v -> NLPModels.hprod(model, x, y, v; obj_weight), model.meta.nvar, J; @@ -89,3 +91,66 @@ function diff_check(model) hess_check(bm, x) end end + +struct ConvexSolver{T} <: SolverCore.AbstractOptimizationSolver + model::LRO.Model{T} + stats::SolverCore.GenericExecutionStats{T,Vector{T},Vector{T},Any} +end + +function ConvexSolver(model::LRO.Model) + stats = SolverCore.GenericExecutionStats(model) + return ConvexSolver(model, stats) +end + +function LRO.MOI.get(solver::ConvexSolver, ::LRO.Solution) + return LRO.VectorizedSolution(solver.stats.solution, solver.model.dim) +end + +function SolverCore.solve!(::ConvexSolver, ::LRO.Model) + return +end + +function _alloc_schur_complement(model, i, Wi, H) + if VERSION < v"1.11" + return + end + LRO.add_schur_complement!(model, i, Wi, H) + @test 0 == @allocated LRO.add_schur_complement!(model, i, Wi, H) +end + +function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} + n = model.meta.ncon + y = rand(T, n) + + Jv = similar(y) + vJ = similar(w) + NLPModels.jprod!(model, w, w, Jv) + NLPModels.jtprod!(model, w, y, vJ) + @test dot(Jv, y) ≈ dot(vJ, w) + + H = zeros(n, n) + H = LRO.schur_complement!(model, w, H) + Hy = similar(y) + LRO.eval_schur_complement!(model, w, y, Hy) + @test Hy ≈ H * y + for i in LRO.matrix_indices(model) + Wi = @inferred w[i] + _alloc_schur_complement(model, i, Wi, H) + end + for i in LRO.matrix_indices(model) + ret = LRO.dual_cons!(model, y, i) + @test ret isa SparseMatrixCSC + end + dcons = ones(LRO.num_scalars(model)) + LRO.dual_cons!(model, y, dcons, LRO.ScalarIndex) + @test dcons ≈ model.model.d_lin - model.model.C_lin' * y +end + +function schur_test(model::LRO.Model{T}, κ) where {T} + w = rand(T, model.meta.nvar) + W = LRO.VectorizedSolution(w, model.dim) + for i in LRO.matrix_indices(model) + W[i] .= W[i] .+ W[i]' + end + return schur_test(LRO.BufferedModelForSchur(model, κ), W) +end diff --git a/test/maxcut.jl b/test/maxcut.jl index a9953de..f6d11f6 100644 --- a/test/maxcut.jl +++ b/test/maxcut.jl @@ -84,69 +84,6 @@ end end end; -struct ConvexSolver{T} <: SolverCore.AbstractOptimizationSolver - model::LRO.Model{T} - stats::SolverCore.GenericExecutionStats{T,Vector{T},Vector{T},Any} -end - -function ConvexSolver(model::LRO.Model) - stats = SolverCore.GenericExecutionStats(model) - return ConvexSolver(model, stats) -end - -function LRO.MOI.get(solver::ConvexSolver, ::LRO.Solution) - return LRO.VectorizedSolution(solver.stats.solution, solver.model.dim) -end - -function SolverCore.solve!(::ConvexSolver, ::LRO.Model) - return -end - -function _alloc_schur_complement(model, i, Wi, H) - if VERSION < v"1.11" - return - end - LRO.add_schur_complement!(model, i, Wi, H) - @test 0 == @allocated LRO.add_schur_complement!(model, i, Wi, H) -end - -function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} - n = model.meta.ncon - y = rand(T, n) - - Jv = similar(y) - vJ = similar(w) - NLPModels.jprod!(model, w, w, Jv) - NLPModels.jtprod!(model, w, y, vJ) - @test dot(Jv, y) ≈ dot(vJ, w) - - H = zeros(n, n) - H = LRO.schur_complement!(model, w, H) - Hy = similar(y) - LRO.eval_schur_complement!(model, w, y, Hy) - @test Hy ≈ H * y - for i in LRO.matrix_indices(model) - Wi = @inferred w[i] - _alloc_schur_complement(model, i, Wi, H) - end - for i in LRO.matrix_indices(model) - ret = LRO.dual_cons!(model, y, i) - @test ret isa SparseMatrixCSC - end - dcons = ones(LRO.num_scalars(model)) - LRO.dual_cons!(model, y, dcons, LRO.ScalarIndex) - @test dcons ≈ model.model.d_lin - model.model.C_lin' * y -end - -function schur_test(model::LRO.Model{T}, κ) where {T} - w = rand(T, model.meta.nvar) - W = LRO.VectorizedSolution(w, model.dim) - for i in LRO.matrix_indices(model) - W[i] .= W[i] .+ W[i]' - end - return schur_test(LRO.BufferedModelForSchur(model, κ), W) -end - @testset "ConvexSolver $T" for T in [Float32, Float64] model = maxcut(T.(weights), LRO.Optimizer{T}) set_attribute(model, "solver", ConvexSolver) From f2ba081e318f3d0bb979ee398dbe909dc80cd4b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 30 Jul 2025 16:56:46 +0200 Subject: [PATCH 07/15] Add tests --- test/buffer.jl | 3 +++ test/diff_check.jl | 38 ++++++++++++++++++++++++++------------ test/maxcut.jl | 2 ++ 3 files changed, 31 insertions(+), 12 deletions(-) diff --git a/test/buffer.jl b/test/buffer.jl index 65ad9ab..1c321a6 100644 --- a/test/buffer.jl +++ b/test/buffer.jl @@ -29,6 +29,9 @@ function test_zero_Ai() for A in b.model.A @test buf.jtprod_buffer[] !== A end + for κ in 0:5 + schur_test(model, κ) + end end function runtests() diff --git a/test/diff_check.jl b/test/diff_check.jl index 0aaa9c4..afe43cf 100644 --- a/test/diff_check.jl +++ b/test/diff_check.jl @@ -4,7 +4,9 @@ # in the LICENSE.md file or at https://opensource.org/licenses/MIT. using Test +using LinearAlgebra import SolverCore +using Dualization import LowRankOpt as LRO function _test_vecprod(f, len, J; tol = 1e-6) @@ -76,22 +78,25 @@ function _backend(model) return b end -function diff_check(model) - b = _backend(model) - bm = b.solver.model - x = rand(bm.meta.nvar) +function diff_check(model::NLPModels.AbstractNLPModel) + x = rand(model.meta.nvar) @testset "Gradient" begin - grad_check(bm, x) - @test isempty(NLPModelsTest.gradient_check(bm; x)) + grad_check(model, x) + @test isempty(NLPModelsTest.gradient_check(model; x)) end @testset "Jacobian" begin - jac_check(bm, x) + jac_check(model, x) end @testset "Hessian" begin - hess_check(bm, x) + hess_check(model, x) end end +function diff_check(model::JuMP.AbstractModel) + b = _backend(model) + diff_check(b.solver.model) +end + struct ConvexSolver{T} <: SolverCore.AbstractOptimizationSolver model::LRO.Model{T} stats::SolverCore.GenericExecutionStats{T,Vector{T},Vector{T},Any} @@ -139,18 +144,27 @@ function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} end for i in LRO.matrix_indices(model) ret = LRO.dual_cons!(model, y, i) - @test ret isa SparseMatrixCSC + @test ret isa SparseArrays.SparseMatrixCSC end dcons = ones(LRO.num_scalars(model)) LRO.dual_cons!(model, y, dcons, LRO.ScalarIndex) @test dcons ≈ model.model.d_lin - model.model.C_lin' * y end -function schur_test(model::LRO.Model{T}, κ) where {T} +function schur_test(model::LRO.BufferedModelForSchur{T}) where {T} w = rand(T, model.meta.nvar) - W = LRO.VectorizedSolution(w, model.dim) + W = LRO.VectorizedSolution(w, model.model.dim) for i in LRO.matrix_indices(model) W[i] .= W[i] .+ W[i]' end - return schur_test(LRO.BufferedModelForSchur(model, κ), W) + schur_test(model, W) +end + +function schur_test(model::LRO.Model, κ) + return schur_test(LRO.BufferedModelForSchur(model, κ)) +end + +function schur_test(model::JuMP.AbstractModel, κ) + b = _backend(model) + schur_test(b.solver.model, κ) end diff --git a/test/maxcut.jl b/test/maxcut.jl index f6d11f6..356bac4 100644 --- a/test/maxcut.jl +++ b/test/maxcut.jl @@ -3,6 +3,8 @@ # 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. +import Percival + include("diff_check.jl") include(joinpath(dirname(@__DIR__), "examples", "maxcut.jl")) From c1026a6be4a7c7326aba5ed9b8b0af2e6739dce9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 31 Jul 2025 08:11:50 +0200 Subject: [PATCH 08/15] Fix coverage --- src/buffer.jl | 14 ++++++++++---- src/errors.jl | 2 +- src/model.jl | 1 + test/buffer.jl | 33 ++++++++++++++++++++++++++------- test/diff_check.jl | 27 ++++++++++++++++++++++++--- 5 files changed, 62 insertions(+), 15 deletions(-) diff --git a/src/buffer.jl b/src/buffer.jl index 8eb5f99..16e4247 100644 --- a/src/buffer.jl +++ b/src/buffer.jl @@ -43,6 +43,7 @@ end function NLPModels.grad(model::BufferedModelForSchur, ::Type{ScalarIndex}) return NLPModels.grad(model.model, ScalarIndex) end + function NLPModels.grad(model::BufferedModelForSchur, i::MatrixIndex) return NLPModels.grad(model.model, i) end @@ -63,8 +64,6 @@ function norm_jac(model::BufferedModelForSchur, i::MatrixIndex) return norm_jac(model.model, i) end -errors(model::BufferedModelForSchur, x; kws...) = errors(model.model, x; kws...) - ####################### ###### J product ###### ####################### @@ -121,6 +120,7 @@ function buffer_for_jprod(model::Model{T}) where {T} end _vec(x::AbstractVector) = x +_vec(x::FillArrays.Zeros{T}) where {T} = FillArrays.Zeros{T}(length(x)) _vec(x::AbstractArray) = UnsafeArrays.uview(x, :) _vec(x::Base.ReshapedArray) = _vec(parent(x)) @@ -209,10 +209,12 @@ function NLPModels.jtprod!( end end -_zero!(A::FillArrays.Zeros) = A +_zero!(A::FillArrays.Zeros) = (@show @__LINE__; A) _zero!(A::SparseArrays.SparseMatrixCSC) = fill!(SparseArrays.nonzeros(A), 0.0) # Computes `A .+= B * α` +function _add_mul!(::FillArrays.Zeros, ::FillArrays.Zeros, _) end + function _add_mul!(A::SparseArrays.SparseMatrixCSC, ::FillArrays.Zeros, _) return A end @@ -257,10 +259,14 @@ function dual_cons!( return dual_cons!(model.model, y, res, ScalarIndex) end +# Note that we can't use `-` because of https://github.com/JuliaArrays/FillArrays.jl/issues/412 +_sub(A::AbstractArray, ::FillArrays.Zeros) = copy(A) +_sub(A::AbstractArray, B::AbstractArray) = A - B + function dual_cons!( model::BufferedModelForSchur, y::AbstractVector, i::MatrixIndex, ) - return model.model.C[i.value] - jtprod!(model, y, i) + return _sub(model.model.C[i.value], jtprod!(model, y, i)) end diff --git a/src/errors.jl b/src/errors.jl index 8f7349a..15f8ee3 100644 --- a/src/errors.jl +++ b/src/errors.jl @@ -4,7 +4,7 @@ Return [the 6 standard DIMACS errors](https://plato.asu.edu/dimacs/node3.html). """ function errors( - model::Model, + model::AbstractModel, x; y = nothing, primal_err = NLPModels.cons(model, x), diff --git a/src/model.jl b/src/model.jl index 6150ea7..e2a803c 100644 --- a/src/model.jl +++ b/src/model.jl @@ -6,6 +6,7 @@ import NLPModels import UnsafeArrays abstract type AbstractModel{T} <: NLPModels.AbstractNLPModel{T,Vector{T}} end +Base.broadcastable(model::AbstractModel) = Ref(model) function NLPModels.cons!( model::AbstractModel, diff --git a/test/buffer.jl b/test/buffer.jl index 1c321a6..d759e64 100644 --- a/test/buffer.jl +++ b/test/buffer.jl @@ -5,13 +5,15 @@ using JuMP, Dualization include("diff_check.jl") # Test with zero Ai matrices -function test_zero_Ai() +function _test_zero_Ai(all_zero::Bool) model = Model(dual_optimizer(LRO.Optimizer)) @variable(model, x[1:2] in MOI.Nonnegatives(2)) @variable(model, X[1:2, 1:2] in PSDCone()) @constraint(model, sum(x) == 1) @constraint(model, 2sum(x) == 2) - @constraint(model, sum(X) == 2) + if !all_zero + @constraint(model, sum(X) == 2) + end @constraint(model, x[1] - x[2] == 1) @objective(model, Max, x[1]) set_attribute(model, "solver", ConvexSolver) @@ -19,21 +21,38 @@ function test_zero_Ai() b = _backend(model) T = Float64 Z = FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}} - S = SparseArrays.SparseMatrixCSC{T,Int} - @test b.model.A isa Matrix{Union{Z,S}} + if all_zero + @test b.model.A isa Matrix{Z} + else + S = SparseArrays.SparseMatrixCSC{T,Int} + @test b.model.A isa Matrix{Union{Z,S}} + end @test b.model.A[1] isa Z @test b.model.A[2] isa Z - @test b.model.A[3] isa S - @test b.model.A[4] 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 + end buf = LRO.BufferedModelForSchur(b.model, 1) for A in b.model.A - @test buf.jtprod_buffer[] !== A + if all_zero + @test buf.jtprod_buffer[] isa LRO.FillArrays.Zeros + else + @test buf.jtprod_buffer[] !== A + end end for κ in 0:5 schur_test(model, κ) end end +function test_zero_Ai() + _test_zero_Ai(false) + return _test_zero_Ai(true) +end + function runtests() for name in names(@__MODULE__; all = true) if startswith("$name", "test_") diff --git a/test/diff_check.jl b/test/diff_check.jl index afe43cf..784179c 100644 --- a/test/diff_check.jl +++ b/test/diff_check.jl @@ -94,7 +94,7 @@ end function diff_check(model::JuMP.AbstractModel) b = _backend(model) - diff_check(b.solver.model) + return diff_check(b.solver.model) end struct ConvexSolver{T} <: SolverCore.AbstractOptimizationSolver @@ -127,6 +127,22 @@ function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} n = model.meta.ncon y = rand(T, n) + idx = LRO.matrix_indices(model) + @test NLPModels.obj(model, w) ≈ + dot(NLPModels.grad(model, LRO.ScalarIndex), w[LRO.ScalarIndex]) + + dot(NLPModels.grad.(model, idx), getindex.(Ref(w), idx)) + @test LRO.dual_obj(model, y) ≈ dot(LRO.cons_constant(model), y) + + @test NLPModels.jac(model, 1, LRO.ScalarIndex) == + NLPModels.jac(model.model, 1, LRO.ScalarIndex) + @test LRO.norm_jac.(model, idx) == LRO.norm_jac.(model.model, idx) + @test LRO.side_dimension.(model, idx) == + LRO.side_dimension.(model.model, idx) + if !isempty(idx) + @test LRO.side_dimension.([model], idx[1]) == + [LRO.side_dimension(model, idx[1])] + end + Jv = similar(y) vJ = similar(w) NLPModels.jprod!(model, w, w, Jv) @@ -143,8 +159,13 @@ function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} _alloc_schur_complement(model, i, Wi, H) end for i in LRO.matrix_indices(model) + @test model.jtprod_buffer[i.value] isa + Union{FillArrays.Zeros,SparseArrays.SparseMatrixCSC} + ret = LRO.jtprod!(model, y, i) + @test ret === model.jtprod_buffer[i.value] ret = LRO.dual_cons!(model, y, i) @test ret isa SparseArrays.SparseMatrixCSC + @test ret !== model.model.C[i.value] end dcons = ones(LRO.num_scalars(model)) LRO.dual_cons!(model, y, dcons, LRO.ScalarIndex) @@ -157,7 +178,7 @@ function schur_test(model::LRO.BufferedModelForSchur{T}) where {T} for i in LRO.matrix_indices(model) W[i] .= W[i] .+ W[i]' end - schur_test(model, W) + return schur_test(model, W) end function schur_test(model::LRO.Model, κ) @@ -166,5 +187,5 @@ end function schur_test(model::JuMP.AbstractModel, κ) b = _backend(model) - schur_test(b.solver.model, κ) + return schur_test(b.solver.model, κ) end From eff533dbef0618557196d3b4e57333d32dee746c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 31 Jul 2025 08:14:59 +0200 Subject: [PATCH 09/15] Fix --- src/buffer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/buffer.jl b/src/buffer.jl index 16e4247..ada0f2b 100644 --- a/src/buffer.jl +++ b/src/buffer.jl @@ -209,7 +209,7 @@ function NLPModels.jtprod!( end end -_zero!(A::FillArrays.Zeros) = (@show @__LINE__; A) +_zero!(A::FillArrays.Zeros) = A _zero!(A::SparseArrays.SparseMatrixCSC) = fill!(SparseArrays.nonzeros(A), 0.0) # Computes `A .+= B * α` From 88e587f7d9d4a50992fecd6ab9b6ef2716003bc0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 31 Jul 2025 08:20:39 +0200 Subject: [PATCH 10/15] Fix --- src/MOI_wrapper.jl | 4 +++- src/buffer.jl | 2 ++ test/buffer.jl | 1 + test/diff_check.jl | 7 +++++-- 4 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/MOI_wrapper.jl b/src/MOI_wrapper.jl index 3f34c24..3ac4f57 100644 --- a/src/MOI_wrapper.jl +++ b/src/MOI_wrapper.jl @@ -268,7 +268,9 @@ function _add_constraints( A[lmi_id, k] = _MatrixBuilder{T}(d) end for row in eachindex(func.constants) - _add!(A[lmi_id, 1], row, func.constants[row], set) + if !iszero(func.constants[row]) + _add!(A[lmi_id, 1], row, func.constants[row], set) + end end for term in func.terms scalar = term.scalar_term diff --git a/src/buffer.jl b/src/buffer.jl index ada0f2b..45cbf25 100644 --- a/src/buffer.jl +++ b/src/buffer.jl @@ -260,7 +260,9 @@ function dual_cons!( end # Note that we can't use `-` because of https://github.com/JuliaArrays/FillArrays.jl/issues/412 +_sub(A::FillArrays.Zeros, ::FillArrays.Zeros) = A _sub(A::AbstractArray, ::FillArrays.Zeros) = copy(A) +_sub(::FillArrays.Zeros, B::AbstractArray) = -B _sub(A::AbstractArray, B::AbstractArray) = A - B function dual_cons!( diff --git a/test/buffer.jl b/test/buffer.jl index d759e64..147ab73 100644 --- a/test/buffer.jl +++ b/test/buffer.jl @@ -21,6 +21,7 @@ function _test_zero_Ai(all_zero::Bool) b = _backend(model) T = Float64 Z = FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}} + @test b.model.C isa Vector{Z} if all_zero @test b.model.A isa Matrix{Z} else diff --git a/test/diff_check.jl b/test/diff_check.jl index 784179c..ac80a9c 100644 --- a/test/diff_check.jl +++ b/test/diff_check.jl @@ -164,8 +164,11 @@ function schur_test(model::LRO.BufferedModelForSchur{T}, w) where {T} ret = LRO.jtprod!(model, y, i) @test ret === model.jtprod_buffer[i.value] ret = LRO.dual_cons!(model, y, i) - @test ret isa SparseArrays.SparseMatrixCSC - @test ret !== model.model.C[i.value] + if ret isa SparseArrays.SparseMatrixCSC + @test ret !== model.model.C[i.value] + else + @test ret isa FillArrays.Zeros + end end dcons = ones(LRO.num_scalars(model)) LRO.dual_cons!(model, y, dcons, LRO.ScalarIndex) From 32ed009c3cfbacb287c11196b6dceb954c428023 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 31 Jul 2025 09:20:55 +0200 Subject: [PATCH 11/15] Fix --- test/diff_check.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/diff_check.jl b/test/diff_check.jl index ac80a9c..026a462 100644 --- a/test/diff_check.jl +++ b/test/diff_check.jl @@ -5,6 +5,7 @@ using Test using LinearAlgebra +using FillArrays import SolverCore using Dualization import LowRankOpt as LRO From de1f27be3867063d3b32db54eb837c106734756c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 31 Jul 2025 09:32:51 +0200 Subject: [PATCH 12/15] Fix --- test/maxcut.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/maxcut.jl b/test/maxcut.jl index 356bac4..773aa26 100644 --- a/test/maxcut.jl +++ b/test/maxcut.jl @@ -54,7 +54,7 @@ function test_maxcut(; is_dual, sparse, vector) @test lro_model.A isa Matrix{LowRankOpt.Factorization{Float64,F,D}} else lro_model = unsafe_backend(model).model - @test lro_model.C isa Vector{SparseMatrixCSC{T,Int64}} + @test lro_model.C isa Vector{FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}}} @test lro_model.A isa Matrix{SparseMatrixCSC{T,Int64}} solver = unsafe_backend(model).solver LRO.BurerMonteiro.set_rank!(solver.model, LRO.MatrixIndex(1), 4) From 3ae247f6836b483eddd1cb656e81dbb3dd67c92b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 31 Jul 2025 09:33:55 +0200 Subject: [PATCH 13/15] Fix format --- test/maxcut.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/maxcut.jl b/test/maxcut.jl index 773aa26..b11a619 100644 --- a/test/maxcut.jl +++ b/test/maxcut.jl @@ -54,7 +54,9 @@ function test_maxcut(; is_dual, sparse, vector) @test lro_model.A 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{Int},Base.OneTo{Int}}}} + @test lro_model.C isa Vector{ + FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}}, + } @test lro_model.A isa Matrix{SparseMatrixCSC{T,Int64}} solver = unsafe_backend(model).solver LRO.BurerMonteiro.set_rank!(solver.model, LRO.MatrixIndex(1), 4) From a1b168b24b7b484c01ccc5e54d4c41916614512f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 31 Jul 2025 10:31:02 +0200 Subject: [PATCH 14/15] Add coverage --- src/buffer.jl | 3 +++ test/buffer.jl | 20 ++++++++++++++------ 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/src/buffer.jl b/src/buffer.jl index 45cbf25..90d5678 100644 --- a/src/buffer.jl +++ b/src/buffer.jl @@ -259,6 +259,9 @@ function dual_cons!( return dual_cons!(model.model, y, res, ScalarIndex) end +# TODO If we rename `dual_cons!` to `unsafe_dual_cons`, +# we can remove the `copy` and remplace `-B` with a mutation + # Note that we can't use `-` because of https://github.com/JuliaArrays/FillArrays.jl/issues/412 _sub(A::FillArrays.Zeros, ::FillArrays.Zeros) = A _sub(A::AbstractArray, ::FillArrays.Zeros) = copy(A) diff --git a/test/buffer.jl b/test/buffer.jl index 147ab73..677602c 100644 --- a/test/buffer.jl +++ b/test/buffer.jl @@ -5,7 +5,7 @@ using JuMP, Dualization include("diff_check.jl") # Test with zero Ai matrices -function _test_zero_Ai(all_zero::Bool) +function _test_zero_Ai(all_zero::Bool, matrix_in_objective::Bool) model = Model(dual_optimizer(LRO.Optimizer)) @variable(model, x[1:2] in MOI.Nonnegatives(2)) @variable(model, X[1:2, 1:2] in PSDCone()) @@ -15,17 +15,21 @@ function _test_zero_Ai(all_zero::Bool) @constraint(model, sum(X) == 2) end @constraint(model, x[1] - x[2] == 1) - @objective(model, Max, x[1]) + if matrix_in_objective + @objective(model, Max, x[1] + X[1, 2] - X[1, 1]) + else + @objective(model, Max, x[1]) + end set_attribute(model, "solver", ConvexSolver) optimize!(model) b = _backend(model) T = Float64 Z = FillArrays.Zeros{T,2,Tuple{Base.OneTo{Int},Base.OneTo{Int}}} - @test b.model.C isa Vector{Z} + S = SparseArrays.SparseMatrixCSC{T,Int} + @test b.model.C isa Vector{matrix_in_objective ? S : Z} if all_zero @test b.model.A isa Matrix{Z} else - S = SparseArrays.SparseMatrixCSC{T,Int} @test b.model.A isa Matrix{Union{Z,S}} end @test b.model.A[1] isa Z @@ -50,8 +54,12 @@ function _test_zero_Ai(all_zero::Bool) end function test_zero_Ai() - _test_zero_Ai(false) - return _test_zero_Ai(true) + @testset "all_zero=$all_zero" for all_zero in [false, true] + @testset "matrix_in_objective=$matrix_in_objective" for matrix_in_objective in [false, true] + _test_zero_Ai(all_zero, matrix_in_objective) + end + end + return end function runtests() From 563332bf7a3a812f42c9ada16f8718cfcb1c8174 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 31 Jul 2025 10:31:09 +0200 Subject: [PATCH 15/15] Fix format --- test/buffer.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/buffer.jl b/test/buffer.jl index 677602c..068b61e 100644 --- a/test/buffer.jl +++ b/test/buffer.jl @@ -55,7 +55,8 @@ end function test_zero_Ai() @testset "all_zero=$all_zero" for all_zero in [false, true] - @testset "matrix_in_objective=$matrix_in_objective" for matrix_in_objective in [false, true] + @testset "matrix_in_objective=$matrix_in_objective" for matrix_in_objective in + [false, true] _test_zero_Ai(all_zero, matrix_in_objective) end end