Fix #1109 by flattening global args

cpmpy/expressions/globalfunctions.py

@@ -1057,6 +1057,8 @@ def __init__(self, arr: ListLike[ExprLike], val: ExprLike):
             raise TypeError(f"Count(arr, val) takes an array of expressions as first argument, not: {arr}")
         if is_any_list(val):
             raise TypeError(f"Count(arr, val) takes a numeric expression as second argument, not a list: {val}")
+        if isinstance(arr, np.ndarray):
+            arr = list(arr.flat)  # flatten multi-dimensional arrays
         super().__init__("count", (arr, val))
 
     def decompose(self) -> tuple[Expression, list[Expression]]:
@@ -1118,6 +1120,8 @@ def __init__(self, arr: ListLike[ExprLike], vals: ListLike[int|np.integer]):
             raise TypeError(f"Among takes as input two arrays, not: {arr} and {vals}")
         if any(isinstance(val, Expression) for val in vals):
             raise TypeError(f"Among takes a set of integer values as input, not {vals}")
+        if isinstance(arr, np.ndarray):
+            arr = list(arr.flat)  # flatten multi-dimensional arrays
         super().__init__("among", (arr, vals))
 
     def decompose(self) -> tuple[Expression, list[Expression]]:
@@ -1171,6 +1175,8 @@ def __init__(self, arr: ListLike[ExprLike]):
         """
         if not is_any_list(arr):
             raise ValueError(f"NValue(arr) takes an array as input, not: {arr}")
+        if isinstance(arr, np.ndarray):
+            arr = list(arr.flat)  # flatten multi-dimensional arrays
         super().__init__("nvalue", tuple(arr))
 
     def decompose(self) -> tuple[Expression, list[Expression]]:
@@ -1236,6 +1242,8 @@ def __init__(self, arr: ListLike[ExprLike], n: int|np.integer):
             raise ValueError("NValueExcept takes an array as input")
         if not is_num(n):
             raise ValueError(f"NValueExcept takes an integer as second argument, but got {n} of type {type(n)}")
+        if isinstance(arr, np.ndarray):
+            arr = list(arr.flat)  # flatten multi-dimensional arrays
         super().__init__("nvalue_except", (arr, n))
 
     def decompose(self) -> tuple[Expression, list[Expression]]:

tests/test_globalconstraints.py

@@ -1535,6 +1535,29 @@ def check_true():
         iv = cp.intvar(0, 10, shape=(3, 3))
         assert cp.Model([cp.NValueExcept(i, val) == 3 for i in iv.T]).solve()
 
+    @pytest.mark.parametrize("make_expr, expected, optimum", [
+        (lambda x: cp.Count(x, 1), 4, 6),
+        (lambda x: cp.Among(x, [1, 2]), 5, 6),
+        (lambda x: cp.NValue(x), 3, 4),
+        (lambda x: cp.NValueExcept(x, 0), 2, 3),
+    ], ids=["count", "among", "nvalue", "nvalue_except"])
+    def test_multidim_count_functions(self, make_expr, expected, optimum):
+        # https://github.com/CPMpy/cpmpy/issues/1109
+        x = cp.intvar(0, 3, shape=(2, 3), name="x")
+        expr = make_expr(x)
+        assert expr.get_bounds()[1] == 6
+
+        assert cp.Model(x == [[1, 1, 2], [1, 0, 1]]).solve()
+        assert expr.value() == expected
+        assert expr.value() == make_expr(x.flatten()).value()
+
+        # as constraint and as objective
+        assert cp.Model(expr == expected).solve()
+        assert expr.value() == expected
+        m = cp.Model(maximize=expr)
+        assert m.solve()
+        assert m.objective_value() == optimum
+
 
     @pytest.mark.skipif(not CPM_minizinc.supported(),
                         reason="Minizinc not installed")