Merge pull request Pyomo#3058 from emma58/all-diff

Adding `all_different` and `count_if` to the logical expression system

pyomo/contrib/cp/repn/docplex_writer.py

@@ -64,7 +64,7 @@
     IndexedBooleanVar,
 )
 from pyomo.core.base.expression import ScalarExpression, _GeneralExpressionData
-from pyomo.core.base.param import IndexedParam, ScalarParam
+from pyomo.core.base.param import IndexedParam, ScalarParam, _ParamData
 from pyomo.core.base.var import ScalarVar, _GeneralVarData, IndexedVar
 import pyomo.core.expr as EXPR
 from pyomo.core.expr.visitor import StreamBasedExpressionVisitor, identify_variables
@@ -805,6 +805,14 @@ def _handle_at_least_node(visitor, node, *args):
     )
 
 
+def _handle_all_diff_node(visitor, node, *args):
+    return (_GENERAL, cp.all_diff(_get_int_valued_expr(arg) for arg in args))
+
+
+def _handle_count_if_node(visitor, node, *args):
+    return (_GENERAL, cp.count((_get_bool_valued_expr(arg) for arg in args), 1))
+
+
 ## CallExpression handllers
 
 
@@ -932,6 +940,8 @@ class LogicalToDoCplex(StreamBasedExpressionVisitor):
         EXPR.ExactlyExpression: _handle_exactly_node,
         EXPR.AtMostExpression: _handle_at_most_node,
         EXPR.AtLeastExpression: _handle_at_least_node,
+        EXPR.AllDifferentExpression: _handle_all_diff_node,
+        EXPR.CountIfExpression: _handle_count_if_node,
         EXPR.EqualityExpression: _handle_equality_node,
         EXPR.NotEqualExpression: _handle_not_equal_node,
         EXPR.InequalityExpression: _handle_inequality_node,
@@ -960,6 +970,7 @@ class LogicalToDoCplex(StreamBasedExpressionVisitor):
         ScalarExpression: _before_named_expression,
         IndexedParam: _before_indexed_param,  # Because of indirection
         ScalarParam: _before_param,
+        _ParamData: _before_param,
     }
 
     def __init__(self, cpx_model, symbolic_solver_labels=False):

pyomo/contrib/cp/tests/test_docplex_walker.py

@@ -21,7 +21,14 @@
 
 from pyomo.core.base.range import NumericRange
 from pyomo.core.expr.numeric_expr import MinExpression, MaxExpression
-from pyomo.core.expr.logical_expr import equivalent, exactly, atleast, atmost
+from pyomo.core.expr.logical_expr import (
+    equivalent,
+    exactly,
+    atleast,
+    atmost,
+    all_different,
+    count_if,
+)
 from pyomo.core.expr.relational_expr import NotEqualExpression
 
 from pyomo.environ import (
@@ -401,6 +408,38 @@ def test_atmost_expression(self):
             expr[1].equals(cp.less_or_equal(cp.count([a[i] == 4 for i in m.I], 1), 3))
         )
 
+    def test_all_diff_expression(self):
+        m = self.get_model()
+        m.a.domain = Integers
+        m.a.bounds = (11, 20)
+        m.c = LogicalConstraint(expr=all_different(m.a))
+
+        visitor = self.get_visitor()
+        expr = visitor.walk_expression((m.c.body, m.c, 0))
+
+        a = {}
+        for i in m.I:
+            self.assertIn(id(m.a[i]), visitor.var_map)
+            a[i] = visitor.var_map[id(m.a[i])]
+
+        self.assertTrue(expr[1].equals(cp.all_diff(a[i] for i in m.I)))
+
+    def test_count_if_expression(self):
+        m = self.get_model()
+        m.a.domain = Integers
+        m.a.bounds = (11, 20)
+        m.c = Constraint(expr=count_if(m.a[i] == i for i in m.I) == 5)
+
+        visitor = self.get_visitor()
+        expr = visitor.walk_expression((m.c.expr, m.c, 0))
+
+        a = {}
+        for i in m.I:
+            self.assertIn(id(m.a[i]), visitor.var_map)
+            a[i] = visitor.var_map[id(m.a[i])]
+
+        self.assertTrue(expr[1].equals(cp.count((a[i] == i for i in m.I), 1) == 5))
+
     def test_interval_var_is_present(self):
         m = self.get_model()
         m.a.domain = Integers

pyomo/contrib/cp/tests/test_docplex_writer.py

@@ -15,10 +15,13 @@
 from pyomo.contrib.cp import IntervalVar, Pulse, Step, AlwaysIn
 from pyomo.contrib.cp.repn.docplex_writer import LogicalToDoCplex
 from pyomo.environ import (
+    all_different,
+    count_if,
     ConcreteModel,
     Set,
     Var,
     Integers,
+    Param,
     LogicalConstraint,
     implies,
     value,
@@ -254,3 +257,138 @@ def x_bounds(m, i):
         self.assertEqual(results.problem.sense, minimize)
         self.assertEqual(results.problem.lower_bound, 6)
         self.assertEqual(results.problem.upper_bound, 6)
+
+    def test_matching_problem(self):
+        m = ConcreteModel()
+
+        m.People = Set(initialize=['P1', 'P2', 'P3', 'P4', 'P5', 'P6', 'P7'])
+        m.Languages = Set(initialize=['English', 'Spanish', 'Hindi', 'Swedish'])
+        # People have integer names because we don't have categorical vars yet.
+        m.Names = Set(initialize=range(len(m.People)))
+
+        m.Observed = Param(
+            m.Names,
+            m.Names,
+            m.Languages,
+            initialize={
+                (0, 1, 'English'): 1,
+                (1, 0, 'English'): 1,
+                (0, 2, 'English'): 1,
+                (2, 0, 'English'): 1,
+                (0, 3, 'English'): 1,
+                (3, 0, 'English'): 1,
+                (0, 4, 'English'): 1,
+                (4, 0, 'English'): 1,
+                (0, 5, 'English'): 1,
+                (5, 0, 'English'): 1,
+                (0, 6, 'English'): 1,
+                (6, 0, 'English'): 1,
+                (1, 2, 'Spanish'): 1,
+                (2, 1, 'Spanish'): 1,
+                (1, 5, 'Hindi'): 1,
+                (5, 1, 'Hindi'): 1,
+                (1, 6, 'Hindi'): 1,
+                (6, 1, 'Hindi'): 1,
+                (2, 3, 'Swedish'): 1,
+                (3, 2, 'Swedish'): 1,
+                (3, 4, 'English'): 1,
+                (4, 3, 'English'): 1,
+            },
+            default=0,
+            mutable=True,
+        )  # TODO: shouldn't need to
+        # be mutable, but waiting
+        # on #3045
+
+        m.Expected = Param(
+            m.People,
+            m.People,
+            m.Languages,
+            initialize={
+                ('P1', 'P2', 'English'): 1,
+                ('P2', 'P1', 'English'): 1,
+                ('P1', 'P3', 'English'): 1,
+                ('P3', 'P1', 'English'): 1,
+                ('P1', 'P4', 'English'): 1,
+                ('P4', 'P1', 'English'): 1,
+                ('P1', 'P5', 'English'): 1,
+                ('P5', 'P1', 'English'): 1,
+                ('P1', 'P6', 'English'): 1,
+                ('P6', 'P1', 'English'): 1,
+                ('P1', 'P7', 'English'): 1,
+                ('P7', 'P1', 'English'): 1,
+                ('P2', 'P3', 'Spanish'): 1,
+                ('P3', 'P2', 'Spanish'): 1,
+                ('P2', 'P6', 'Hindi'): 1,
+                ('P6', 'P2', 'Hindi'): 1,
+                ('P2', 'P7', 'Hindi'): 1,
+                ('P7', 'P2', 'Hindi'): 1,
+                ('P3', 'P4', 'Swedish'): 1,
+                ('P4', 'P3', 'Swedish'): 1,
+                ('P4', 'P5', 'English'): 1,
+                ('P5', 'P4', 'English'): 1,
+            },
+            default=0,
+            mutable=True,
+        )  # TODO: shouldn't need to be mutable, but
+        # waiting on #3045
+
+        m.person_name = Var(m.People, bounds=(0, max(m.Names)), domain=Integers)
+
+        m.one_to_one = LogicalConstraint(
+            expr=all_different(m.person_name[person] for person in m.People)
+        )
+
+        m.obj = Objective(
+            expr=count_if(
+                m.Observed[m.person_name[p1], m.person_name[p2], l]
+                == m.Expected[p1, p2, l]
+                for p1 in m.People
+                for p2 in m.People
+                for l in m.Languages
+            ),
+            sense=maximize,
+        )
+
+        results = SolverFactory('cp_optimizer').solve(m)
+
+        # we can get one of two perfect matches:
+        perfect = 7 * 7 * 4
+        self.assertEqual(results.problem.lower_bound, perfect)
+        self.assertEqual(results.problem.upper_bound, perfect)
+        self.assertEqual(
+            results.solver.termination_condition, TerminationCondition.optimal
+        )
+        self.assertEqual(value(m.obj), perfect)
+        m.person_name.pprint()
+        self.assertEqual(value(m.person_name['P1']), 0)
+        self.assertEqual(value(m.person_name['P2']), 1)
+        self.assertEqual(value(m.person_name['P3']), 2)
+        self.assertEqual(value(m.person_name['P4']), 3)
+        self.assertEqual(value(m.person_name['P5']), 4)
+        # We can't distinguish P6 and P7, so they could each have either of
+        # names 5 and 6
+        self.assertTrue(
+            value(m.person_name['P6']) == 5 or value(m.person_name['P6']) == 6
+        )
+        self.assertTrue(
+            value(m.person_name['P7']) == 5 or value(m.person_name['P7']) == 6
+        )
+
+        m.person_name['P6'].fix(5)
+        m.person_name['P7'].fix(6)
+
+        results = SolverFactory('cp_optimizer').solve(m)
+        self.assertEqual(
+            results.solver.termination_condition, TerminationCondition.optimal
+        )
+        self.assertEqual(value(m.obj), perfect)
+
+        m.person_name['P6'].fix(6)
+        m.person_name['P7'].fix(5)
+
+        results = SolverFactory('cp_optimizer').solve(m)
+        self.assertEqual(
+            results.solver.termination_condition, TerminationCondition.optimal
+        )
+        self.assertEqual(value(m.obj), perfect)

pyomo/core/__init__.py

@@ -33,6 +33,8 @@
     exactly,
     atleast,
     atmost,
+    all_different,
+    count_if,
     implies,
     lnot,
     xor,

pyomo/core/expr/__init__.py

@@ -70,6 +70,8 @@
     ExactlyExpression,
     AtMostExpression,
     AtLeastExpression,
+    AllDifferentExpression,
+    CountIfExpression,
     #
     land,
     lnot,
@@ -79,6 +81,8 @@
     exactly,
     atleast,
     atmost,
+    all_different,
+    count_if,
     implies,
 )
 from .numeric_expr import (