Add support for rank-1 constraints

Author: blegat

src/MOI_wrapper.jl

@@ -14,6 +14,10 @@ mutable struct Optimizer <: MOI.AbstractOptimizer
     # * `-d` means a diagonal block with diagonal of length `d`
     # * `d` means a symmetric `d x d` block
     blockdims::Vector{Int}
+    # MOI variable index -> rank 1 matrix it corresponds to
+    rank_one::Vector{Union{Nothing,MOI.LowRankMatrix{Cdouble}}}
+    # To avoid it being free'd
+    cached_ind::Vector{Vector{Cint}}
     varmap::Vector{Tuple{Int, Int, Int}} # Variable Index vi -> blk, i, j
     # If `blockdims[i] < 0`, `blk[i]` is the offset in `lpdvars`.
     # That is the **sum of length** of diagonal block before
@@ -33,15 +37,18 @@ mutable struct Optimizer <: MOI.AbstractOptimizer
 
     y::Vector{Cdouble}
 
+    solve_time::Float64
     silent::Bool
     options::Dict{Symbol,Any}
     function Optimizer()
         optimizer = new(
             C_NULL, C_NULL, 0.0, 1, Cdouble[], Int[],
+            Union{Nothing,MOI.LowRankMatrix{Cdouble}}[],
+            Vector{Cint}[],
             Tuple{Int, Int, Int}[], Int[], C_NULL, 0,
             Int[], Int[], Cdouble[],
             Vector{Int}[], Vector{Cdouble}[],
-            Cdouble[], false, Dict{Symbol, Any}(),
+            Cdouble[], NaN, false, Dict{Symbol, Any}(),
         )
         finalizer(_free, optimizer)
         return optimizer
@@ -64,6 +71,8 @@ function MOI.empty!(optimizer::Optimizer)
     optimizer.objective_sign = 1
     empty!(optimizer.b)
     empty!(optimizer.blockdims)
+    empty!(optimizer.rank_one)
+    empty!(optimizer.cached_ind)
     empty!(optimizer.varmap)
     empty!(optimizer.blk)
     optimizer.nlpdrows = 0
@@ -73,6 +82,8 @@ function MOI.empty!(optimizer::Optimizer)
     empty!(optimizer.sdpdinds)
     empty!(optimizer.sdpdcoefs)
     empty!(optimizer.y)
+    optimizer.solve_time = NaN
+    return
 end
 
 function MOI.is_empty(optimizer::Optimizer)
@@ -193,8 +204,20 @@ function MOI.supports(
 end
 
 MOI.supports_add_constrained_variables(::Optimizer, ::Type{MOI.Reals}) = false
-const SupportedSets = Union{MOI.Nonnegatives, MOI.PositiveSemidefiniteConeTriangle}
+
+const _SetWithDotProd = MOI.SetWithDotProducts{
+    MOI.PositiveSemidefiniteConeTriangle,
+    MOI.TriangleVectorization{Cdouble,MOI.LowRankMatrix{Cdouble}},
+}
+
+const SupportedSets = Union{
+    MOI.Nonnegatives,
+    MOI.PositiveSemidefiniteConeTriangle,
+    _SetWithDotProd,
+}
+
 MOI.supports_add_constrained_variables(::Optimizer, ::Type{<:SupportedSets}) = true
+
 function MOI.supports_constraint(
     ::Optimizer, ::Type{MOI.ScalarAffineFunction{Cdouble}},
     ::Type{MOI.EqualTo{Cdouble}})
@@ -206,6 +229,7 @@ function new_block(optimizer::Optimizer, set::MOI.Nonnegatives)
     blk = length(optimizer.blockdims)
     for i in 1:MOI.dimension(set)
         push!(optimizer.varmap, (blk, i, i))
+        push!(optimizer.rank_one, nothing)
     end
 end
 
@@ -215,13 +239,26 @@ function new_block(optimizer::Optimizer, set::MOI.PositiveSemidefiniteConeTriang
     for j in 1:set.side_dimension
         for i in 1:j
             push!(optimizer.varmap, (blk, i, j))
+            push!(optimizer.rank_one, nothing)
         end
     end
 end
 
+function new_block(model::Optimizer, set::_SetWithDotProd)
+    println("______________________Low-Rank")
+    blk = length(model.blockdims) + 1
+    for i in eachindex(set.vectors)
+        push!(model.varmap, (blk, 0, 0))
+        push!(model.rank_one, set.vectors[i].matrix)
+    end
+    new_block(model, set.set)
+end
+
 function _add_constrained_variables(optimizer::Optimizer, set::SupportedSets)
     offset = length(optimizer.varmap)
     new_block(optimizer, set)
+    @assert length(optimizer.varmap) == offset + MOI.dimension(set)
+    @assert length(optimizer.rank_one) == offset + MOI.dimension(set)
     ci = MOI.ConstraintIndex{MOI.VectorOfVariables, typeof(set)}(offset + 1)
     return [MOI.VariableIndex(i) for i in offset .+ (1:MOI.dimension(set))], ci
 end
@@ -265,32 +302,51 @@ function constrain_variables_on_creation(
     end
 end
 
-function _setcoefficient!(m::Optimizer, coef, constr::Integer, blk::Integer, i::Integer, j::Integer)
-    if m.blockdims[blk] < 0
-        @assert i == j
-        push!(m.lpdvars, constr + 1)
-        push!(m.lpdrows, m.blk[blk] + i - 1) # -1 because indexing starts at 0 in DSDP
-        push!(m.lpcoefs, coef)
-    else
-        sdp = m.blk[blk]
-        push!(m.sdpdinds[end][sdp], i + (j - 1) * m.blockdims[blk] - 1)
-        if i != j
-            coef /= 2
+function _setcoefficient!(dest::Optimizer, coef, constr::Integer, vi::MOI.VariableIndex)
+    blk, i, j = varmap(dest, vi)
+    rank_one = dest.rank_one[vi.value]
+    if isnothing(rank_one)
+        if dest.blockdims[blk] < 0
+            @assert i == j
+            push!(dest.lpdvars, constr + 1)
+            push!(dest.lpdrows, dest.blk[blk] + i - 1) # -1 because indexing starts at 0 in DSDP
+            push!(dest.lpcoefs, coef)
+        else
+            sdp = dest.blk[blk]
+            push!(dest.sdpdinds[end][sdp], i + (j - 1) * dest.blockdims[blk] - 1)
+            if i != j
+                coef /= 2
+            end
+            push!(dest.sdpdcoefs[end][sdp], coef)
         end
-        push!(m.sdpdcoefs[end][sdp], coef)
+    else
+        d = Cint(dest.blockdims[blk])
+        push!(dest.cached_ind, collect(Cint(0):(d - 1)))
+        # We use `Add` and not `Set` because I think (if I interpret the name correctly) that would allow mixing with sparse matrices for the same block and constraint
+        DSDP.SDPCone.SetARankOneMat(
+            dest.sdpcone,
+            dest.blk[blk] - 1,
+            constr,
+            d,
+            coef * rank_one.diagonal[],
+            0,
+            last(dest.cached_ind),
+            collect(eachcol(rank_one.factor)[]),
+            d,
+        )
     end
 end
 
 # Loads objective coefficient α * vi
-function load_objective_term!(optimizer::Optimizer, index_map, α, vi::MOI.VariableIndex)
-    blk, i, j = varmap(optimizer, vi)
+function load_objective_term!(optimizer::Optimizer, α, vi::MOI.VariableIndex)
     coef = optimizer.objective_sign * α
-    _setcoefficient!(optimizer, coef, 0, blk, i, j)
+    _setcoefficient!(optimizer, coef, 0, vi)
 end
 
 function _set_A_matrices(m::Optimizer, i)
     for (blk, blkdim) in zip(m.blk, m.blockdims)
-        if blkdim > 0
+        if blkdim > 0 && !isempty(m.sdpdcoefs[end][blk])
+            @show (blk - 1, i, blkdim, 1.0, 0, m.sdpdinds[end][blk], m.sdpdcoefs[end][blk], length(m.sdpdcoefs[end][blk]))
             SDPCone.SetASparseVecMat(m.sdpcone, blk - 1, i, blkdim, 1.0, 0, m.sdpdinds[end][blk], m.sdpdcoefs[end][blk], length(m.sdpdcoefs[end][blk]))
         end
     end
@@ -321,6 +377,12 @@ function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike)
         index_map,
         MOI.PositiveSemidefiniteConeTriangle,
     )
+    constrain_variables_on_creation(
+        dest,
+        src,
+        index_map,
+        _SetWithDotProd,
+    )
     vis_src = MOI.get(src, MOI.ListOfVariableIndices())
     if length(vis_src) < length(index_map.var_map)
         _error(
@@ -392,8 +454,7 @@ function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike)
         _new_A_matrix(dest)
         for t in func.terms
             if !iszero(t.coefficient)
-                blk, i, j = varmap(dest, index_map[t.variable])
-                _setcoefficient!(dest, t.coefficient, k, blk, i, j)
+                _setcoefficient!(dest, t.coefficient, k, index_map[t.variable])
             end
         end
         _set_A_matrices(dest, k)
@@ -428,7 +489,6 @@ function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike)
             if !iszero(term.coefficient)
                 load_objective_term!(
                     dest,
-                    index_map,
                     term.coefficient,
                     index_map[term.variable],
                 )
@@ -449,7 +509,9 @@ function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike)
 end
 
 function MOI.optimize!(m::Optimizer)
+    start_time = time()
     Solve(m.dsdp)
+    m.solve_time = time() - start_time
     # Calling `ComputeX` not right after `Solve` seems to sometime cause segfaults or weird Heisenbug's
     # let's call it directly what `DSDP/examples/readsdpa.c` does
     ComputeX(m.dsdp)
@@ -459,6 +521,8 @@ function MOI.optimize!(m::Optimizer)
     return
 end
 
+MOI.get(optimizer::Optimizer, ::MOI.SolveTimeSec) = optimizer.solve_time
+
 function MOI.get(m::Optimizer, ::MOI.RawStatusString)
     if m.dsdp == C_NULL
         return "`optimize!` not called"

src/sdpcone.jl

@@ -53,8 +53,8 @@ function SetADenseVecMat(sdpcone::SDPConeT, arg2::Integer, arg3::Integer, arg4::
     @dsdp_ccall SDPConeSetADenseVecMat (SDPConeT, Cint, Cint, Cint, Cdouble, Ptr{Cdouble}, Cint) sdpcone arg2 arg3 arg4 arg5 arg6 arg7
 end
 
-function SetARankOneMat(sdpcone::SDPConeT, arg2::Integer, arg3::Integer, arg4::Integer, arg5::Cdouble, arg6::Integer, arg7::Vector{Cint}, arg8::Vector{Cdouble}, arg9::Integer)
-    @dsdp_ccall SDPConeSetARankOneMat (SDPConeT, Cint, Cint, Cint, Cdouble, Cint, Ptr{Cint}, Ptr{Cdouble}, Cint) sdpcone arg2 arg3 arg4 arg5 arg6 arg7 arg8 arg9
+function SetARankOneMat(sdpcone::SDPConeT, blockj::Integer, vari::Integer, n::Integer, alpha::Cdouble, ishift::Integer, ind::Vector{Cint}, val::Vector{Cdouble}, nnz::Integer)
+    @dsdp_ccall SDPConeSetARankOneMat (SDPConeT, Cint, Cint, Cint, Cdouble, Cint, Ptr{Cint}, Ptr{Cdouble}, Cint) sdpcone blockj vari n alpha ishift ind val nnz
 end
 
 function SetConstantMat(sdpcone::SDPConeT, arg2::Integer, arg3::Integer, arg4::Integer, arg5::Cdouble)
@@ -93,8 +93,8 @@ function AddIdentity(sdpcone::SDPConeT, arg2::Integer, arg3::Integer, arg4::Inte
     @dsdp_ccall SDPConeAddIdentity (SDPConeT, Cint, Cint, Cint, Cdouble) sdpcone arg2 arg3 arg4 arg5
 end
 
-function AddARankOneMat(sdpcone::SDPConeT, arg2::Integer, arg3::Integer, arg4::Integer, arg5::Cdouble, arg6::Integer, arg7::Vector{Cint}, arg8::Vector{Cdouble}, arg9::Integer)
-    @dsdp_ccall SDPConeAddARankOneMat (SDPConeT, Cint, Cint, Cint, Cdouble, Cint, Ptr{Cint}, Ptr{Cdouble}, Cint) sdpcone arg2 arg3 arg4 arg5 arg6 arg7 arg8 arg9
+function AddARankOneMat(sdpcone::SDPConeT, blockj::Integer, vari::Integer, n::Integer, alpha::Cdouble, ishift::Integer, ind::Vector{Cint}, val::Vector{Cdouble}, nnz::Integer)
+    @dsdp_ccall SDPConeAddARankOneMat (SDPConeT, Cint, Cint, Cint, Cdouble, Cint, Ptr{Cint}, Ptr{Cdouble}, Cint) sdpcone blockj vari n alpha ishift ind val nnz
 end
 
 function AddSparseVecMat(sdpcone::SDPConeT, arg2::Integer, arg3::Integer, arg4::Integer, arg5::Integer, arg6::Vector{Cint}, arg7::Vector{Cdouble}, arg8::Integer)