Add bridge attribute for choosing the Cholesky method

Author: blegat

src/Bridges/Constraint/bridges/QuadtoSOCBridge.jl

@@ -4,6 +4,29 @@
 # 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.
 
+"""
+    QuadtoSOCSquareRoot <: MOI.AbstractConstraintAttribute
+
+Constraint attribute to override the square root method used by
+[`QuadtoSOCBridge`](@ref) for a specific quadratic constraint.
+
+When set, this takes precedence over the default behavior of trying all
+available methods. When not set, [`MOI.get`](@ref) returns `nothing`.
+
+## Example
+
+```julia
+c = MOI.add_constraint(model, f, MOI.LessThan(1.0))
+MOI.set(
+    model,
+    MOI.Bridges.Constraint.QuadtoSOCSquareRoot(),
+    c,
+    MOI.Bridges.Constraint._LinearAlgebra(),
+)
+```
+"""
+struct QuadtoSOCSquareRoot <: MOI.AbstractConstraintAttribute end
+
 """
     QuadtoSOCBridge{T} <: Bridges.Constraint.AbstractBridge
 
@@ -46,15 +69,22 @@ Therefore, `QuadtoSOCBridge` implements the following reformulations:
 
 This bridge errors if `Q` is not positive definite.
 """
-struct QuadtoSOCBridge{T} <: AbstractBridge
-    soc::MOI.ConstraintIndex{
-        MOI.VectorAffineFunction{T},
-        MOI.RotatedSecondOrderCone,
+mutable struct QuadtoSOCBridge{T} <: AbstractBridge
+    soc::Union{
+        Nothing,
+        MOI.ConstraintIndex{
+            MOI.VectorAffineFunction{T},
+            MOI.RotatedSecondOrderCone,
+        },
     }
     dimension::Int  # dimension of the SOC constraint
     less_than::Bool # whether the constraint was ≤ or ≥
     set_constant::T # the constant that was on the set
     index_to_variable_map::Vector{MOI.VariableIndex}
+    # Stored for final_touch
+    func::MOI.ScalarQuadraticFunction{T}
+    set::Union{MOI.LessThan{T},MOI.GreaterThan{T}}
+    method::Union{Nothing,_AbstractExt}
 end
 
 const QuadtoSOC{T,OT<:MOI.ModelLike} =
@@ -157,13 +187,63 @@ function bridge_constraint(
     func::MOI.ScalarQuadraticFunction{T},
     set::Union{MOI.LessThan{T},MOI.GreaterThan{T}},
 ) where {T}
+    # Delay reformulation until `final_touch` so that the
+    # `QuadtoSOCSquareRoot` attribute can override the method first.
     less_than = set isa MOI.LessThan{T}
+    set_constant = MOI.constant(set)
+    MOI.throw_if_scalar_and_constant_not_zero(func, typeof(set))
+    return QuadtoSOCBridge{T}(
+        nothing,
+        0,
+        less_than,
+        set_constant,
+        MOI.VariableIndex[],
+        func,
+        set,
+        nothing,
+    )
+end
+
+MOI.supports(::MOI.ModelLike, ::QuadtoSOCSquareRoot, ::Type{<:QuadtoSOCBridge}) =
+    true
+
+function MOI.set(
+    ::MOI.ModelLike,
+    ::QuadtoSOCSquareRoot,
+    bridge::QuadtoSOCBridge,
+    value::Union{Nothing,_AbstractExt},
+)
+    bridge.method = value
+    return
+end
+
+function MOI.get(
+    ::MOI.ModelLike,
+    ::QuadtoSOCSquareRoot,
+    bridge::QuadtoSOCBridge,
+)
+    return bridge.method
+end
+
+MOI.Bridges.needs_final_touch(::QuadtoSOCBridge) = true
+
+function MOI.Bridges.final_touch(
+    bridge::QuadtoSOCBridge{T},
+    model::MOI.ModelLike,
+) where {T}
+    if bridge.soc !== nothing
+        return
+    end
+    func = bridge.func
+    set = bridge.set
+    less_than = bridge.less_than
     scale = less_than ? -1 : 1
     Q, index_to_variable_map =
         _matrix_from_quadratic_terms(func.quadratic_terms)
     if !less_than
         LinearAlgebra.rmul!(Q, -1)
     end
+    bridge.index_to_variable_map = index_to_variable_map
     # Construct the VectorAffineFunction. We're aiming for:
     #  |          1  |
     #  | -a^T x + ub | ∈ RotatedSecondOrderCone()
@@ -175,10 +255,15 @@ function bridge_constraint(
             MOI.ScalarAffineTerm(scale * term.coefficient, term.variable),
         ) for term in func.affine_terms
     ]
-    sqrt_ret = _compute_sparse_sqrt(LinearAlgebra.Symmetric(Q))
+    Q_sym = LinearAlgebra.Symmetric(Q)
+    sqrt_ret = if bridge.method !== nothing
+        _compute_sparse_sqrt(bridge.method, Q_sym)
+    else
+        _compute_sparse_sqrt(Q_sym)
+    end
     if sqrt_ret === nothing
         msg = _get_sqrt_error_message(is_defined(_CliqueTrees()))
-        return throw(MOI.UnsupportedConstraint{typeof(func),typeof(set)}(msg))
+        throw(MOI.UnsupportedConstraint{typeof(func),typeof(set)}(msg))
     end
     for (i, j, v) in zip(sqrt_ret[1], sqrt_ret[2], sqrt_ret[3])
         push!(
@@ -190,19 +275,13 @@ function bridge_constraint(
         )
     end
     # This is the [1, ub, 0] vector...
-    set_constant = MOI.constant(set)
-    MOI.throw_if_scalar_and_constant_not_zero(func, typeof(set))
+    set_constant = bridge.set_constant
     vector_constant = vcat(one(T), -scale * set_constant, zeros(T, size(Q, 1)))
     f = MOI.VectorAffineFunction(vector_terms, vector_constant)
-    dimension = MOI.output_dimension(f)
-    soc = MOI.add_constraint(model, f, MOI.RotatedSecondOrderCone(dimension))
-    return QuadtoSOCBridge(
-        soc,
-        dimension,
-        less_than,
-        set_constant,
-        index_to_variable_map,
-    )
+    bridge.dimension = MOI.output_dimension(f)
+    bridge.soc =
+        MOI.add_constraint(model, f, MOI.RotatedSecondOrderCone(bridge.dimension))
+    return
 end
 
 function _matrix_from_quadratic_terms(
@@ -267,13 +346,13 @@ end
 
 # Attributes, Bridge acting as a model
 function MOI.get(
-    ::QuadtoSOCBridge{T},
+    bridge::QuadtoSOCBridge{T},
     ::MOI.NumberOfConstraints{
         MOI.VectorAffineFunction{T},
         MOI.RotatedSecondOrderCone,
     },
 )::Int64 where {T}
-    return 1
+    return bridge.soc === nothing ? 0 : 1
 end
 
 function MOI.get(
@@ -283,12 +362,21 @@ function MOI.get(
         MOI.RotatedSecondOrderCone,
     },
 ) where {T}
+    if bridge.soc === nothing
+        return MOI.ConstraintIndex{
+            MOI.VectorAffineFunction{T},
+            MOI.RotatedSecondOrderCone,
+        }[]
+    end
     return [bridge.soc]
 end
 
 # References
 function MOI.delete(model::MOI.ModelLike, bridge::QuadtoSOCBridge)
-    MOI.delete(model, bridge.soc)
+    if bridge.soc !== nothing
+        MOI.delete(model, bridge.soc)
+        bridge.soc = nothing
+    end
     return
 end
 
@@ -485,6 +573,9 @@ function MOI.get(
     attr::MOI.ConstraintFunction,
     b::QuadtoSOCBridge{T},
 ) where {T}
+    if b.soc === nothing
+        return copy(b.func)
+    end
     f = MOI.get(model, attr, b.soc)
     fs = MOI.Utilities.eachscalar(f)
     q = zero(MOI.ScalarQuadraticFunction{T})

test/Bridges/Constraint/test_QuadtoSOCBridge.jl

@@ -31,13 +31,19 @@ function test_error_for_nonconvex_quadratic_constraints()
     model = MOI.Bridges.Constraint.QuadtoSOC{Float64}(inner)
     x = MOI.add_variable(model)
     F = MOI.ScalarQuadraticFunction{Float64}
+    # Error is now thrown at final_touch, not add_constraint
+    MOI.add_constraint(model, 1.0 * x * x, MOI.GreaterThan(0.0))
     @test_throws(
         MOI.UnsupportedConstraint{F,MOI.GreaterThan{Float64}},
-        MOI.add_constraint(model, 1.0 * x * x, MOI.GreaterThan(0.0))
+        MOI.Bridges.final_touch(model),
     )
+    MOI.empty!(model)
+    MOI.add_variable(model)
+    x = MOI.get(model, MOI.ListOfVariableIndices())[1]
+    MOI.add_constraint(model, -1.0 * x * x, MOI.LessThan(0.0))
     @test_throws(
         MOI.UnsupportedConstraint{F,MOI.LessThan{Float64}},
-        MOI.add_constraint(model, -1.0 * x * x, MOI.LessThan(0.0))
+        MOI.Bridges.final_touch(model),
     )
     return
 end
@@ -65,6 +71,7 @@ function test_quadratic_constraints_with_2_variables()
         ),
     )
     MOI.Test.test_constraint_qcp_duplicate_off_diagonal(bridged_mock, config)
+    MOI.Bridges.final_touch(bridged_mock)
     ci = first(
         MOI.get(
             mock,
@@ -164,6 +171,7 @@ function test_fill_reducing_permutation()
     Q = Float64[2 1 1; 1 2 0; 1 0 2]
     f = 0.5 * x' * Q * x
     MOI.add_constraint(bridge, f, MOI.LessThan(2.0))
+    MOI.Bridges.final_touch(bridge)
     indices = MOI.get(
         model,
         MOI.ListOfConstraintIndices{
@@ -331,6 +339,7 @@ function test_semidefinite_cholesky_fail()
     x = MOI.add_variables(model, 2)
     f = 0.5 * x[1] * x[1] + 1.0 * x[1] * x[2] + 0.5 * x[2] * x[2]
     c = MOI.add_constraint(model, f, MOI.LessThan(1.0))
+    MOI.Bridges.final_touch(model)
     F, S = MOI.VectorAffineFunction{Float64}, MOI.RotatedSecondOrderCone
     ci = only(MOI.get(inner, MOI.ListOfConstraintIndices{F,S}()))
     g = MOI.get(inner, MOI.ConstraintFunction(), ci)
@@ -435,6 +444,7 @@ function test_clique_trees_semidefinite_cholesky_fail()
     x = MOI.add_variables(model, 2)
     f = 0.5 * x[1] * x[1] + 1.0 * x[1] * x[2] + 0.5 * x[2] * x[2]
     c = MOI.add_constraint(model, f, MOI.LessThan(1.0))
+    MOI.Bridges.final_touch(model)
     F, S = MOI.VectorAffineFunction{Float64}, MOI.RotatedSecondOrderCone
     ci = only(MOI.get(inner, MOI.ListOfConstraintIndices{F,S}()))
     g = MOI.get(inner, MOI.ConstraintFunction(), ci)
@@ -454,6 +464,7 @@ function test_clique_trees_early_zero_pivot()
     # Q = [1 1 0; 1 1 0; 0 0 1]
     f = sum(0.5 * x[i] * x[i] for i in 1:3) + 1.0 * x[1] * x[2]
     c = MOI.add_constraint(model, f, MOI.LessThan(1.0))
+    MOI.Bridges.final_touch(model)
     F, S = MOI.VectorAffineFunction{Float64}, MOI.RotatedSecondOrderCone
     ci = only(MOI.get(inner, MOI.ListOfConstraintIndices{F,S}()))
     g = MOI.get(inner, MOI.ConstraintFunction(), ci)
@@ -482,6 +493,68 @@ function test_is_defined_default_fallback()
     return
 end
 
+function test_quad_to_soc_square_root_attribute()
+    inner = MOI.Utilities.Model{Float64}()
+    model = MOI.Bridges.Constraint.QuadtoSOC{Float64}(inner)
+    x = MOI.add_variables(model, 2)
+    f = 1.0 * x[1] * x[1] + 1.0 * x[2] * x[2]
+    c = MOI.add_constraint(model, f, MOI.LessThan(1.0))
+    attr = MOI.Bridges.Constraint.QuadtoSOCSquareRoot()
+    F = MOI.ScalarQuadraticFunction{Float64}
+    S = MOI.LessThan{Float64}
+    @test MOI.supports(model, attr, MOI.ConstraintIndex{F,S})
+    # Default is nothing
+    @test MOI.get(model, attr, c) === nothing
+    # Set to _LinearAlgebra
+    la = MOI.Bridges.Constraint._LinearAlgebra()
+    MOI.set(model, attr, c, la)
+    @test MOI.get(model, attr, c) === la
+    # final_touch uses the specified method
+    MOI.Bridges.final_touch(model)
+    F2, S2 = MOI.VectorAffineFunction{Float64}, MOI.RotatedSecondOrderCone
+    ci = only(MOI.get(inner, MOI.ListOfConstraintIndices{F2,S2}()))
+    g = MOI.get(inner, MOI.ConstraintFunction(), ci)
+    @test MOI.output_dimension(g) == 4  # [1, rhs, Ux...]
+    return
+end
+
+function test_quad_to_soc_square_root_attribute_clique_trees()
+    inner = MOI.Utilities.Model{Float64}()
+    model = MOI.Bridges.Constraint.QuadtoSOC{Float64}(inner)
+    x = MOI.add_variables(model, 2)
+    f = 1.0 * x[1] * x[1] + 1.0 * x[2] * x[2]
+    c = MOI.add_constraint(model, f, MOI.LessThan(1.0))
+    attr = MOI.Bridges.Constraint.QuadtoSOCSquareRoot()
+    ct = MOI.Bridges.Constraint._CliqueTrees()
+    MOI.set(model, attr, c, ct)
+    @test MOI.get(model, attr, c) === ct
+    MOI.Bridges.final_touch(model)
+    F, S = MOI.VectorAffineFunction{Float64}, MOI.RotatedSecondOrderCone
+    ci = only(MOI.get(inner, MOI.ListOfConstraintIndices{F,S}()))
+    g = MOI.get(inner, MOI.ConstraintFunction(), ci)
+    @test MOI.output_dimension(g) == 4
+    return
+end
+
+function test_quad_to_soc_square_root_attribute_reset()
+    inner = MOI.Utilities.Model{Float64}()
+    model = MOI.Bridges.Constraint.QuadtoSOC{Float64}(inner)
+    x = MOI.add_variables(model, 2)
+    f = 1.0 * x[1] * x[1] + 1.0 * x[2] * x[2]
+    c = MOI.add_constraint(model, f, MOI.LessThan(1.0))
+    attr = MOI.Bridges.Constraint.QuadtoSOCSquareRoot()
+    la = MOI.Bridges.Constraint._LinearAlgebra()
+    MOI.set(model, attr, c, la)
+    # Reset to nothing (default behavior)
+    MOI.set(model, attr, c, nothing)
+    @test MOI.get(model, attr, c) === nothing
+    MOI.Bridges.final_touch(model)
+    F, S = MOI.VectorAffineFunction{Float64}, MOI.RotatedSecondOrderCone
+    ci = only(MOI.get(inner, MOI.ListOfConstraintIndices{F,S}()))
+    @test ci isa MOI.ConstraintIndex
+    return
+end
+
 end  # module
 
 TestConstraintQuadToSOC.runtests()