Add chsh example

docs/src/tutorials/Noncommutative and Hermitian/chsh.jl

@@ -0,0 +1,116 @@
+import QuantumStuff: trace_monoid, Monoids
+using StarAlgebras
+
+M, A, C = trace_monoid(2, 2, A=:A, C=:C)
+
+RM = let M = M, A = A, C = C, level = 4
+    A_l, sizesA = Monoids.wlmetric_ball(A, radius=level)
+    C_l, sizesC = Monoids.wlmetric_ball(C, radius=level)
+
+    # starAlg(M, 1, half = unique!([a*c for a in A_l for c in C_l]))
+
+    @time words, sizes = Monoids.wlmetric_ball(
+        unique!([a * c for a in A_l for c in C_l]);
+        radius=2,
+    )
+    @info "Sizes of generated balls:" (A, C, combined) = (sizesA, sizesC, sizes)
+
+    b = @time StarAlgebras.FixedBasis(words, StarAlgebras.DiracMStructure(*), (UInt32(sizes[1]), UInt32(sizes[1])))
+    StarAlgebra(M, b)
+end
+
+A = RM.(A)
+C = RM.(C)
+chsh = A[1] * C[1] + A[1] * C[2] + A[2] * C[1] - A[2] * C[2]
+
+import StarAlgebras as SA
+struct Full{B} <: SA.ImplicitBasis{B,B} end
+Base.getindex(::Full{B}, b::B) where {B} = b
+import MultivariateBases as MB
+MB.implicit_basis(::SA.FixedBasis{B}) where {B} = Full{B}()
+MB.algebra(b::Full{B}) where {B} = SA.StarAlgebra(M, b)
+SA.mstructure(::Full) = SA.DiracMStructure(*)
+
+b = basis(chsh)
+import StarAlgebras as SA
+f = SA.AlgebraElement(
+    SA.SparseCoefficients(
+        [b[k] for (k, _) in SA.nonzero_pairs(coeffs(chsh))],
+        [v for (_, v) in SA.nonzero_pairs(coeffs(chsh))],
+    ),
+    SA.StarAlgebra(
+        parent(chsh).object,
+        Full{eltype(b)}()
+    ),
+)
+n = size(b.table, 1)
+gram_basis = @time StarAlgebras.FixedBasis(b.elts[1:n], StarAlgebras.DiracMStructure(*));
+one(f)
+SA.coeffs(f, b)
+using SumOfSquares
+function SumOfSquares._term_element(a, mono::Monoids.MonoidElement)
+    SA.AlgebraElement(
+        SA.SparseCoefficients((mono,), (a,)),
+        MB.algebra(Full{typeof(mono)}()),
+    )
+end
+
+cone = SumOfSquares.WeightedSOSCone{MOI.PositiveSemidefiniteConeTriangle}(
+    b,
+    [gram_basis],
+    [one(f)],
+)
+import SCS
+scs = optimizer_with_attributes(
+    SCS.Optimizer,
+    "eps_abs" => 1e-9,
+    "eps_rel" => 1e-9,
+    "max_iters" => 1000,
+)
+
+import Dualization
+#model = Model(Dualization.dual_optimizer(scs))
+model = Model(scs)
+SumOfSquares.Bridges.add_all_bridges(backend(model).optimizer, Float64)
+MOI.Bridges.remove_bridge(backend(model).optimizer, SumOfSquares.Bridges.Constraint.ImageBridge{Float64})
+@variable(model, λ)
+@objective(model, Min, λ)
+@constraint(model, SA.coeffs(λ * one(f) - f, b) in cone);
+optimize!(model)
+solution_summary(model)
+objective_value(model) ≈ 2√2
+function _add!(f, psd, model, F, S)
+    append!(psd, [
+        f(MOI.get(model, MOI.ConstraintSet(), ci))
+        for ci in MOI.get(model, MOI.ListOfConstraintIndices{F,S}())
+    ],
+    )
+end
+function summary(model)
+    free = MOI.get(model, MOI.NumberOfVariables())
+    psd = Int[]
+    F = MOI.VectorOfVariables
+    S = MOI.PositiveSemidefiniteConeTriangle
+    _add!(MOI.side_dimension, psd, model, F, S)
+    S = SCS.ScaledPSDCone
+    _add!(MOI.side_dimension, psd, model, F, S)
+    free -= sum(psd, init = 0) do d
+        div(d * (d + 1), 2)
+    end
+    F = MOI.VectorAffineFunction{Float64}
+    S = MOI.PositiveSemidefiniteConeTriangle
+    _add!(MOI.side_dimension, psd, model, F, S)
+    S = SCS.ScaledPSDCone
+    _add!(MOI.side_dimension, psd, model, F, S)
+    eq = Int[]
+    F = MOI.VectorAffineFunction{Float64}
+    S = MOI.Zeros
+    _add!(MOI.dimension, eq, model, F, S)
+    F = MOI.ScalarAffineFunction{Float64}
+    S = MOI.EqualTo{Float64}
+    _add!(MOI.dimension, eq, model, F, S)
+    return free, psd, sum(eq, init = 0)
+end
+summary(backend(model).optimizer.model)
+summary(backend(model).optimizer.model.optimizer.dual_problem.dual_model.model)
+print_active_bridges(model)

src/Bridges/Bridges.jl

@@ -6,6 +6,11 @@ import SumOfSquares as SOS
 include("Variable/Variable.jl")
 include("Constraint/Constraint.jl")
 
+function add_all_bridges(model, ::Type{T}) where {T}
+    Variable.add_all_bridges(model, T)
+    Constraint.add_all_bridges(model, T)
+end
+
 function MOI.get(
     model::MOI.ModelLike,
     attr::Union{

src/Bridges/Constraint/Constraint.jl

@@ -26,4 +26,13 @@ include("image.jl")
 include("sos_polynomial.jl")
 include("sos_polynomial_in_semialgebraic_set.jl")
 
+function add_all_bridges(model, ::Type{T}) where {T}
+    MOI.Bridges.add_bridge(model, EmptyBridge{T})
+    MOI.Bridges.add_bridge(model, PositiveSemidefinite2x2Bridge{T})
+    MOI.Bridges.add_bridge(model, DiagonallyDominantBridge{T})
+    MOI.Bridges.add_bridge(model, ImageBridge{T})
+    MOI.Bridges.add_bridge(model, SOSPolynomialBridge{T})
+    MOI.Bridges.add_bridge(model, SOSPolynomialInSemialgebraicSetBridge{T})
+end
+
 end

src/Bridges/Constraint/image.jl

@@ -81,17 +81,18 @@ function MOI.Bridges.Constraint.bridge_constraint(
     @assert MOI.output_dimension(g) == length(set.basis)
     scalars = MOI.Utilities.scalarize(g)
     k = 0
-    found = Dict{eltype(set.basis.monomials),Int}()
+    found = Dict{eltype(set.basis),Int}()
     first = Union{Nothing,Int}[nothing for _ in eachindex(scalars)]
     variables = MOI.VariableIndex[]
     constraints = MOI.ConstraintIndex{F}[]
     for (gram_basis, weight) in zip(set.gram_bases, set.weights)
+        # TODO don't ignore weight
         cone = SOS.matrix_cone(M, length(gram_basis))
         f = MOI.Utilities.zero_with_output_dimension(F, MOI.dimension(cone))
-        for j in eachindex(gram_basis.monomials)
+        for j in eachindex(gram_basis)
             for i in 1:j
                 k += 1
-                mono = gram_basis.monomials[i] * gram_basis.monomials[j]
+                mono = SA.star(gram_basis[i]) * gram_basis[j]
                 is_diag = i == j
                 if haskey(found, mono)
                     var = MOI.add_variable(model)
@@ -119,8 +120,9 @@ function MOI.Bridges.Constraint.bridge_constraint(
                     MOI.Utilities.operate_output_index!(-, T, k, f, var)
                 else
                     found[mono] = k
-                    t = MB.monomial_index(set.basis, mono)
-                    if !isnothing(t)
+                    @show mono, set.basis[mono]
+                    if mono in set.basis
+                        t = set.basis[mono]
                         first[t] = k
                         if is_diag
                             MOI.Utilities.operate_output_index!(
@@ -139,6 +141,8 @@ function MOI.Bridges.Constraint.bridge_constraint(
                                 inv(T(2)) * scalars[t],
                             )
                         end
+                    else
+                        @warn("$mono not in basis")
                     end
                 end
             end
@@ -167,14 +171,8 @@ end
 function MOI.supports_constraint(
     ::Type{ImageBridge{T}},
     ::Type{<:MOI.AbstractVectorFunction},
-    ::Type{
-        <:SOS.WeightedSOSCone{
-            M,
-            <:MB.SubBasis{MB.Monomial},
-            <:MB.SubBasis{MB.Monomial},
-        },
-    },
-) where {T,M}
+    ::Type{<:SOS.WeightedSOSCone},
+) where {T}
     return true
 end
 

src/Bridges/Variable/Variable.jl

@@ -13,4 +13,13 @@ include("scaled_diagonally_dominant.jl")
 include("copositive_inner.jl")
 include("kernel.jl")
 
+function add_all_bridges(model, ::Type{T}) where {T}
+    MOI.Bridges.add_bridge(model, PositiveSemidefinite2x2Bridge{T})
+    MOI.Bridges.add_bridge(model, ScaledDiagonallyDominantBridge{T})
+    MOI.Bridges.add_bridge(model, CopositiveInnerBridge{T})
+    MOI.Bridges.add_bridge(model, KernelBridge{T})
+    MOI.Bridges.add_bridge(model, LowRankBridge{T})
+    return
+end
+
 end

src/Bridges/Variable/lagrange.jl

@@ -0,0 +1,178 @@
+struct LagrangeBridge{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
+
+function MOI.Bridges.Variable.bridge_constrained_variable(
+    ::Type{KernelBridge{T,M}},
+    model::MOI.ModelLike,
+    set::SOS.WeightedSOSCone{M},
+) where {T,M}
+    variables = Vector{MOI.VariableIndex}[]
+    constraints = MOI.ConstraintIndex{MOI.VectorOfVariables}[]
+    acc = zero(
+        MOI.ScalarAffineFunction{T},
+        MB.algebra(MB.implicit_basis(set.basis)),
+    )
+    for (gram_basis, weight) in zip(set.gram_bases, set.weights)
+        gram, vars, con = SOS.add_gram_matrix(model, M, gram_basis, T)
+        push!(variables, vars)
+        push!(constraints, con)
+        MA.operate!(SA.UnsafeAddMul(*), acc, gram, weight)
+    end
+    MA.operate!(SA.canonical, SA.coeffs(acc))
+    return KernelBridge{T,M}(
+        SA.coeffs(acc, set.basis),
+        variables,
+        constraints,
+        set,
+    )
+end
+
+function MOI.Bridges.Variable.supports_constrained_variable(
+    ::Type{<:KernelBridge},
+    ::Type{<:SOS.WeightedSOSCone},
+)
+    return true
+end
+
+function MOI.Bridges.added_constrained_variable_types(
+    ::Type{KernelBridge{T,M}},
+) where {T,M}
+    return SOS.Bridges.Constraint.constrained_variable_types(M)
+end
+
+function MOI.Bridges.added_constraint_types(::Type{<:KernelBridge})
+    return Tuple{Type,Type}[]
+end
+
+function MOI.Bridges.Variable.concrete_bridge_type(
+    ::Type{<:KernelBridge{T}},
+    ::Type{<:SOS.WeightedSOSCone{M}},
+) where {T,M}
+    return KernelBridge{T,M}
+end
+
+# Attributes, Bridge acting as a model
+function MOI.get(bridge::KernelBridge, ::MOI.NumberOfVariables)
+    return sum(length, bridge.variables)
+end
+
+function MOI.get(bridge::KernelBridge, ::MOI.ListOfVariableIndices)
+    return reduce(vcat, bridge.variables)
+end
+
+function MOI.get(
+    bridge::KernelBridge,
+    ::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::KernelBridge,
+    ::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::KernelBridge)
+    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::KernelBridge)
+    return bridge.set
+end
+
+function MOI.get(
+    model::MOI.ModelLike,
+    attr::MOI.ConstraintPrimal,
+    bridge::KernelBridge,
+)
+    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::KernelBridge,
+    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.get(
+    model::MOI.ModelLike,
+    attr::SOS.GramMatrixAttribute,
+    bridge::KernelBridge{T,M},
+) where {T,M}
+    SOS.check_multiplier_index_bounds(attr, eachindex(bridge.constraints))
+    return SOS.build_gram_matrix(
+        convert(
+            Vector{T},
+            MOI.get(
+                model,
+                MOI.VariablePrimal(attr.result_index),
+                bridge.variables[attr.multiplier_index],
+            ),
+        ),
+        bridge.set.gram_bases[attr.multiplier_index],
+        M,
+        T,
+    )
+end
+
+function MOI.get(
+    model::MOI.ModelLike,
+    attr::SOS.MomentMatrixAttribute,
+    bridge::KernelBridge{T,M},
+) where {T,M}
+    SOS.check_multiplier_index_bounds(attr, eachindex(bridge.constraints))
+    return SOS.build_moment_matrix(
+        convert(
+            Vector{T},
+            MOI.get(
+                model,
+                MOI.ConstraintDual(attr.result_index),
+                bridge.constraints[attr.multiplier_index],
+            ),
+        ),
+        bridge.set.gram_bases[attr.multiplier_index],
+    )
+end
+
+function MOI.Bridges.bridged_function(
+    bridge::KernelBridge,
+    i::MOI.Bridges.IndexInVector,
+)
+    return bridge.affine[i.value]
+end
+
+function MOI.Bridges.Variable.unbridged_map(
+    ::KernelBridge{T},
+    ::Vector{MOI.VariableIndex},
+) where {T}
+    return nothing
+end