Add LowRankBridge

Project.toml

@@ -8,6 +8,7 @@ CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LowRankOpt = "607ca3ad-272e-43c8-bcbe-fc71b56c935c"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 MultivariateBases = "be282fd4-ad43-11e9-1d11-8bd9d7e43378"
 MultivariateMoments = "f4abf1af-0426-5881-a0da-e2f168889b5e"
@@ -24,6 +25,7 @@ SymbolicWedderburn = "858aa9a9-4c7c-4c62-b466-2421203962a2"
 CliqueTrees = "1"
 DataStructures = "0.19"
 JuMP = "1.10"
+LowRankOpt = "0.2"
 MathOptInterface = "1.13"
 MultivariateBases = "0.3"
 MultivariateMoments = "0.5"

src/Bridges/Variable/Variable.jl

@@ -6,11 +6,13 @@ import StarAlgebras as SA
 import MultivariateBases as MB
 import MathOptInterface as MOI
 import MultivariatePolynomials as MP
+import LowRankOpt as LRO
 import SumOfSquares as SOS
 
 include("psd2x2.jl")
 include("scaled_diagonally_dominant.jl")
 include("copositive_inner.jl")
 include("kernel.jl")
+include("lowrank.jl")
 
 end

src/Bridges/Variable/kernel.jl

@@ -37,9 +37,11 @@ end
 
 function MOI.Bridges.Variable.supports_constrained_variable(
     ::Type{<:KernelBridge},
-    ::Type{<:SOS.WeightedSOSCone},
-)
-    return true
+    ::Type{<:SOS.WeightedSOSCone{M,B}},
+) where {M,B}
+    # Could be made to work but doesn't work yet so it's best to use
+    # `LowRankBridge` which can then be bridged to classical PSD by MOI's bridge
+    return !(B <: MB.LagrangeBasis)
 end
 
 function MOI.Bridges.added_constrained_variable_types(

src/Bridges/Variable/lowrank.jl

@@ -0,0 +1,173 @@
+struct LowRankBridge{T,M} <: MOI.Bridges.Variable.AbstractBridge
+    affine::Vector{MOI.ScalarAffineFunction{T}}
+    variables::Vector{Vector{MOI.VariableIndex}}
+    constraints::Vector{MOI.ConstraintIndex{MOI.VectorOfVariables}}
+    set::SOS.WeightedSOSCone{M}
+end
+
+import LinearAlgebra
+
+function MOI.Bridges.Variable.bridge_constrained_variable(
+    ::Type{LowRankBridge{T,M}},
+    model::MOI.ModelLike,
+    set::SOS.WeightedSOSCone{M},
+) where {T,M}
+    variables = Vector{Vector{MOI.VariableIndex}}(undef, length(set.gram_bases))
+    constraints = Vector{MOI.ConstraintIndex{MOI.VectorOfVariables}}(
+        undef,
+        length(set.gram_bases),
+    )
+    for i in eachindex(set.gram_bases)
+        U = MB.transformation_to(set.gram_bases[i], set.basis)
+        weights = SA.coeffs(set.weights[i], set.basis)
+        variables[i], constraints[i] = MOI.add_constrained_variables(
+            model,
+            LRO.SetDotProducts{LRO.WITHOUT_SET}(
+                SOS.matrix_cone(M, length(set.gram_bases[i])),
+                [
+                    LRO.TriangleVectorization(
+                        LRO.Factorization(
+                            Matrix{T}(reshape(U[j, :], size(U, 2), 1)),
+                            T[weights[j]],
+                        ),
+                    ) for j in eachindex(set.basis)
+                ],
+            ),
+        )
+    end
+    return LowRankBridge{T,M}(
+        [
+            MOI.ScalarAffineFunction(
+                [
+                    MOI.ScalarAffineTerm(one(T), variables[i][j]) for
+                    i in eachindex(set.gram_bases)
+                ],
+                zero(T),
+            ) for j in eachindex(set.basis)
+        ],
+        variables,
+        constraints,
+        set,
+    )
+end
+
+function MOI.Bridges.Variable.supports_constrained_variable(
+    ::Type{<:LowRankBridge},
+    ::Type{<:SOS.WeightedSOSCone{M,B}},
+) where {M,B}
+    # Could be made to work for non-LagrangeBasis but it's not low rank in
+    # we we'll need a high bridge cost for them so that it's only UnsafeAddMul
+    # if the other bridges are removed
+    return B <: MB.LagrangeBasis
+end
+
+function MOI.Bridges.added_constrained_variable_types(
+    ::Type{LowRankBridge{T,M}},
+) where {T,M}
+    return Tuple{Type}[
+        (
+            LRO.SetDotProducts{
+                LRO.WITHOUT_SET,
+                S[1],
+                LRO.TriangleVectorization{
+                    T,
+                    LRO.Factorization{T,Matrix{T},Vector{T}},
+                },
+            },
+        ) for S in SOS.Bridges.Constraint.constrained_variable_types(M) if
+        S[1] == MOI.PositiveSemidefiniteConeTriangle # FIXME hack
+    ]
+end
+
+function MOI.Bridges.added_constraint_types(::Type{<:LowRankBridge})
+    return Tuple{Type,Type}[]
+end
+
+function MOI.Bridges.Variable.concrete_bridge_type(
+    ::Type{<:LowRankBridge{T}},
+    ::Type{<:SOS.WeightedSOSCone{M}},
+) where {T,M}
+    return LowRankBridge{T,M}
+end
+
+# Attributes, Bridge acting as a model
+function MOI.get(bridge::LowRankBridge, ::MOI.NumberOfVariables)
+    return sum(length, bridge.variables)
+end
+
+function MOI.get(bridge::LowRankBridge, ::MOI.ListOfVariableIndices)
+    return reduce(vcat, bridge.variables)
+end
+
+function MOI.get(
+    bridge::LowRankBridge,
+    ::MOI.NumberOfConstraints{MOI.VectorOfVariables,S},
+) where {S<:MOI.AbstractVectorSet}
+    return count(bridge.constraints) do ci
+        return ci isa MOI.ConstraintIndex{MOI.VectorOfVariables,S}
+    end
+end
+
+function MOI.get(
+    bridge::LowRankBridge,
+    ::MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S},
+) where {S}
+    return [
+        ci for ci in bridge.constraints if
+        ci isa MOI.ConstraintIndex{MOI.VectorOfVariables,S}
+    ]
+end
+
+# Indices
+function MOI.delete(model::MOI.ModelLike, bridge::LowRankBridge)
+    for vars in bridge.variables
+        MOI.delete(model, vars)
+    end
+    return
+end
+
+# Attributes, Bridge acting as a constraint
+
+function MOI.get(::MOI.ModelLike, ::MOI.ConstraintSet, bridge::LowRankBridge)
+    return bridge.set
+end
+
+function MOI.get(
+    model::MOI.ModelLike,
+    attr::MOI.ConstraintPrimal,
+    bridge::LowRankBridge,
+)
+    return [
+        MOI.get(
+            model,
+            MOI.VariablePrimal(attr.result_index),
+            bridge,
+            MOI.Bridges.IndexInVector(i),
+        ) for i in eachindex(bridge.affine)
+    ]
+end
+
+function MOI.get(
+    model::MOI.ModelLike,
+    attr::MOI.VariablePrimal,
+    bridge::LowRankBridge,
+    i::MOI.Bridges.IndexInVector,
+)
+    return MOI.Utilities.eval_variables(bridge.affine[i.value]) do vi
+        return MOI.get(model, MOI.VariablePrimal(attr.result_index), vi)
+    end
+end
+
+function MOI.Bridges.bridged_function(
+    bridge::LowRankBridge,
+    i::MOI.Bridges.IndexInVector,
+)
+    return bridge.affine[i.value]
+end
+
+function MOI.Bridges.Variable.unbridged_map(
+    ::LowRankBridge,
+    ::Vector{MOI.VariableIndex},
+)
+    return nothing
+end

src/variables.jl

@@ -9,10 +9,11 @@ function _bridge_coefficient_type(::Type{SOSPolynomialSet{D,B,C}}) where {D,B,C}
 end
 
 function PolyJuMP.bridges(S::Type{<:WeightedSOSCone})
-    return Tuple{Type,Type}[(
-        Bridges.Variable.KernelBridge,
-        _bridge_coefficient_type(S),
-    )]
+    T = _bridge_coefficient_type(S)
+    return Tuple{Type,Type}[
+        (Bridges.Variable.KernelBridge, T),
+        (Bridges.Variable.LowRankBridge, T),
+    ]
 end
 
 function PolyJuMP.bridges(::Type{<:PositiveSemidefinite2x2ConeTriangle})