Add CountAtLeast set

Author: odow

docs/src/reference/standard_form.md

@@ -101,6 +101,7 @@ Complements
 ```@docs
 AllDifferent
 Among
+CountAtLeast
 CountDistinct
 ```
 

src/Test/test_basic_constraint.jl

@@ -112,6 +112,7 @@ _set(::Type{MOI.Complements}) = MOI.Complements(2)
 _set(::Type{MOI.AllDifferent}) = MOI.AllDifferent(3)
 _set(::Type{MOI.CountDistinct}) = MOI.CountDistinct(4)
 _set(::Type{MOI.Among}) = MOI.Among(4, Set([3, 4]))
+_set(::Type{MOI.CountAtLeast}) = MOI.CountAtLeast(1, [2, 2], Set([3]))
 
 function _set(
     ::Type{MOI.Indicator{MOI.ACTIVATE_ON_ONE,MOI.LessThan{T}}},
@@ -278,6 +279,7 @@ for s in [
     :AllDifferent,
     :CountDistinct,
     :Among,
+    :CountAtLeast,
 ]
     S = getfield(MOI, s)
     functions = if S <: MOI.AbstractScalarSet

src/Test/test_cpsat.jl

@@ -122,3 +122,53 @@ function setup_test(
     )
     return
 end
+
+"""
+    test_cpsat_CountAtLeast(model::MOI.ModelLike, config::Config)
+
+Add a VectorOfVariables-in-CountAtLeast constraint.
+"""
+function test_cpsat_CountAtLeast(
+    model::MOI.ModelLike,
+    config::Config{T},
+) where {T}
+    @requires MOI.supports_constraint(
+        model,
+        MOI.VectorOfVariables,
+        MOI.CountAtLeast,
+    )
+    @requires MOI.supports_add_constrained_variable(model, MOI.Integer)
+    @requires _supports(config, MOI.optimize!)
+    x, _ = MOI.add_constrained_variable(model, MOI.Integer())
+    y, _ = MOI.add_constrained_variable(model, MOI.Integer())
+    z, _ = MOI.add_constrained_variable(model, MOI.Integer())
+    variables = [x, y, y, z]
+    partitions = [2, 2]
+    set = Set([3])
+    MOI.add_constraint(
+        model,
+        MOI.VectorOfVariables(variables),
+        MOI.CountAtLeast(1, partitions, set),
+    )
+    MOI.optimize!(model)
+    x_val = round.(Int, MOI.get.(model, MOI.VariablePrimal(), [x, y, z]))
+    @test x_val[1] == 3 || x_val[2] == 3
+    @test x_val[2] == 3 || x_val[3] == 3
+    return
+end
+
+function setup_test(
+    ::typeof(test_cpsat_CountAtLeast),
+    model::MOIU.MockOptimizer,
+    ::Config{T},
+) where {T}
+    MOIU.set_mock_optimize!(
+        model,
+        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(
+            mock,
+            MOI.OPTIMAL,
+            (MOI.FEASIBLE_POINT, T[0, 3, 0]),
+        ),
+    )
+    return
+end

src/Utilities/model.jl

@@ -790,6 +790,7 @@ const LessThanIndicatorZero{T} =
         MOI.AllDifferent,
         MOI.CountDistinct,
         MOI.Among,
+        MOI.CountAtLeast,
     ),
     (
         MOI.PowerCone,

src/sets.jl

@@ -1216,6 +1216,52 @@ end
 
 Base.:(==)(x::Among, y::Among) = x.dimension == y.dimension && x.set == y.set
 
+"""
+    CountAtLeast(n::Int, d::Vector{Int}, set::Set{Int})
+
+The set ``\\{x \\in \\mathbb{Z}^{d_1 + d_2 + \\ldots d_N}\\}``, where `x` is
+partitioned into `N` subsets (``\\{x_1,  \\ldots, x_{d_1}\\}``,
+``\\{x_{d_1, 1},  \\ldots, x_{d_1 + d_2}\\}`` and so on), and at least ``n``
+elements of each subset take one of the values in `set`.
+
+## Also known as
+
+This constraint is called `at_least` in MiniZinc.
+
+## Example
+
+```julia
+model = Utilities.Model{Float64}()
+a, _ = add_constrained_variable(model, Integer())
+b, _ = add_constrained_variable(model, Integer())
+c, _ = add_constrained_variable(model, Integer())
+# To ensure that `3` appears at least once in each of the subsets {a, b}, {b, c}
+x, d, set = [a, b, b, c], [2, 2], [3]
+add_constraint(model, VectorOfVariables(x), CountAtLeast(1, d, Set(set)))
+```
+"""
+struct CountAtLeast <: AbstractVectorSet
+    n::Int
+    partitions::Vector{Int}
+    set::Set{Int}
+    function CountAtLeast(
+        n::Base.Integer,
+        partitions::Vector{Int},
+        set::Set{Int},
+    )
+        if any(p <= 0 for p in partitions)
+            throw(DimensionMismatch("Invalid partition dimension."))
+        end
+        return new(n, partitions, set)
+    end
+end
+
+dimension(s::CountAtLeast) = sum(s.partitions)
+
+function Base.:(==)(x::CountAtLeast, y::CountAtLeast)
+    return x.n == y.n && x.partitions == y.partitions && x.set == y.set
+end
+
 # isbits types, nothing to copy
 function Base.copy(
     set::Union{
@@ -1253,13 +1299,18 @@ function Base.copy(
         AllDifferent,
         CountDistinct,
         Among,
+        CountAtLeast,
     },
 )
     return set
 end
 Base.copy(set::S) where {S<:Union{SOS1,SOS2}} = S(copy(set.weights))
 Base.copy(set::Among) = Among(set.dimension, copy(set.set))
 
+function Base.copy(set::CountAtLeast)
+    return CountAtLeast(set.n, copy(set.partitions), copy(set.set))
+end
+
 """
     supports_dimension_update(S::Type{<:MOI.AbstractVectorSet})
 

test/sets.jl

@@ -156,6 +156,7 @@ function test_sets_DimensionMismatch()
     @test_throws DimensionMismatch MOI.Complements(-3)
     @test_throws DimensionMismatch MOI.Complements(3)
     @test_throws DimensionMismatch MOI.Among(0, Set([1, 2]))
+    @test_throws DimensionMismatch MOI.CountAtLeast(1, [-1, 2], Set([1, 2]))
     return
 end