Add CountGreaterThan set (#1829)

Author: odow

docs/src/reference/standard_form.md

@@ -103,6 +103,7 @@ AllDifferent
 Among
 CountAtLeast
 CountDistinct
+CountGreaterThan
 ```
 
 ## Matrix sets

src/Test/test_basic_constraint.jl

@@ -113,6 +113,7 @@ _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]))
+_set(::Type{MOI.CountGreaterThan}) = MOI.CountGreaterThan(5)
 
 function _set(
     ::Type{MOI.Indicator{MOI.ACTIVATE_ON_ONE,MOI.LessThan{T}}},
@@ -280,6 +281,7 @@ for s in [
     :CountDistinct,
     :Among,
     :CountAtLeast,
+    :CountGreaterThan,
 ]
     S = getfield(MOI, s)
     functions = if S <: MOI.AbstractScalarSet

src/Test/test_cpsat.jl

@@ -172,3 +172,51 @@ function setup_test(
     )
     return
 end
+
+"""
+    test_cpsat_CountGreaterThan(model::MOI.ModelLike, config::Config)
+
+Add a VectorOfVariables-in-CountGreaterThan constraint.
+"""
+function test_cpsat_CountGreaterThan(
+    model::MOI.ModelLike,
+    config::Config{T},
+) where {T}
+    @requires MOI.supports_constraint(
+        model,
+        MOI.VectorOfVariables,
+        MOI.CountGreaterThan,
+    )
+    @requires MOI.supports_add_constrained_variable(model, MOI.Integer)
+    @requires _supports(config, MOI.optimize!)
+    c, _ = MOI.add_constrained_variable(model, MOI.Integer())
+    y, _ = MOI.add_constrained_variable(model, MOI.Integer())
+    x = [MOI.add_constrained_variable(model, MOI.Integer())[1] for _ in 1:3]
+    MOI.add_constraint(
+        model,
+        MOI.VectorOfVariables([c; y; x]),
+        MOI.CountGreaterThan(5),
+    )
+    MOI.optimize!(model)
+    c_val = round(Int, MOI.get(model, MOI.VariablePrimal(), c))
+    y_val = round(Int, MOI.get(model, MOI.VariablePrimal(), y))
+    x_val = round.(Int, MOI.get.(model, MOI.VariablePrimal(), x))
+    @test c_val > sum(x_val[i] == y_val for i in 1:length(x))
+    return
+end
+
+function setup_test(
+    ::typeof(test_cpsat_CountGreaterThan),
+    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[2, 4, 0, 4, 0]),
+        ),
+    )
+    return
+end

src/Utilities/model.jl

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

src/sets.jl

@@ -1262,6 +1262,40 @@ function Base.:(==)(x::CountAtLeast, y::CountAtLeast)
     return x.n == y.n && x.partitions == y.partitions && x.set == y.set
 end
 
+"""
+    CountGreaterThan(dimension::Int)
+
+The set ``\\{(c, y, x) \\in \\mathbb{Z}^{1+1+d}\\}`` such that `c` is strictly
+greater than the number of occurances of `y` in `x`.
+
+## Also known as
+
+This constraint is called `count_gt` in MiniZinc.
+
+## Example
+
+```julia
+model = Utilities.Model{Float64}()
+c, _ = add_constrained_variable(model, Integer())
+y, _ = add_constrained_variable(model, Integer())
+x = [add_constrained_variable(model, Integer())[1] for _ in 1:3]
+add_constraint(model, VectorOfVariables([c; y; x]), CountGreaterThan(5))
+```
+"""
+struct CountGreaterThan <: AbstractVectorSet
+    dimension::Int
+    function CountGreaterThan(dimension::Base.Integer)
+        if dimension < 2
+            throw(
+                DimensionMismatch(
+                    "Dimension of CountGreaterThan must be >= 2.",
+                ),
+            )
+        end
+        return new(dimension)
+    end
+end
+
 # isbits types, nothing to copy
 function Base.copy(
     set::Union{
@@ -1300,6 +1334,7 @@ function Base.copy(
         CountDistinct,
         Among,
         CountAtLeast,
+        CountGreaterThan,
     },
 )
     return set

test/sets.jl

@@ -141,6 +141,7 @@ function test_sets_DimensionMismatch()
         (MOI.RootDetConeSquare, 0),
         (MOI.AllDifferent, 0),
         (MOI.CountDistinct, 1),
+        (MOI.CountGreaterThan, 2),
     )
         @test_throws DimensionMismatch S(min_dimension - 1)
         @test S(min_dimension) isa S