From 2bc3191e68360e0b57b21b53f900a5fe79243b7a Mon Sep 17 00:00:00 2001 From: Dimitri Justeau-Allaire Date: Fri, 20 Sep 2024 16:41:50 +0200 Subject: [PATCH] Add 2D matrix shape option in intvars and boolvars creation (https://github.com/chocoteam/pychoco/issues/18) --- pychoco/variables/variable_factory.py | 71 +++++++++++++++++++-------- tests/test_boolvar.py | 11 +++++ tests/test_intvars.py | 11 +++++ 3 files changed, 72 insertions(+), 21 deletions(-) diff --git a/pychoco/variables/variable_factory.py b/pychoco/variables/variable_factory.py index 988f548..3f43bd2 100644 --- a/pychoco/variables/variable_factory.py +++ b/pychoco/variables/variable_factory.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Union, List +from typing import Union, List, Tuple from pychoco import backend from pychoco._utils import make_int_array @@ -57,25 +57,40 @@ def intvar(self, lb: Union[int, List[int]], ub: Union[int, None] = None, name: U else backend.intvar_siib(self.handle, name, lb, ub, bounded_domain) return IntVar(var_handle, self) - def intvars(self, size: int, lb: Union[List[int], int], ub: Union[int, None] = None, name: Union[str, None] = None, bounded_domain: Union[bool, None] = None): + def intvars(self, size: Union[int, Tuple[int]], lb: Union[List[Union[int, List[int]]], int], ub: Union[int, None] = None, name: Union[str, None] = None, bounded_domain: Union[bool, None] = None): """ Creates a list of intvars. - :param size: Number of intvars. + :param size: Number of intvars. Either an int or a two-values tuple describing matrix dimensions (nrows, ncols) :param lb: Lower bound (integer). If lb is a list of ints, constant variables are created. :param ub: Upper bound (integer). If None: the variable is a constant equals to lb. :param bounded_domain: Force bounded (True) or enumerated domain (False). If None, Choco will automatically choose the best option. :param name: Prefix name of the intvars (automatically given if None). :return: A list of intvars. """ - names = [None for i in range(0, size)] - if name is not None: - names = ["{}_{}".format(name, i) for i in range(0, size)] - if isinstance(lb, list): - assert len(lb) == size - return [self.intvar(lb[i], None, names[i]) for i in range(0, size)] - else: - return [self.intvar(lb, ub, names[i], bounded_domain) for i in range(0, size)] + # Case 1D array + if isinstance(size, int): + names = [None for i in range(0, size)] + if name is not None: + names = ["{}_{}".format(name, i) for i in range(0, size)] + if isinstance(lb, list): + assert len(lb) == size + return [self.intvar(lb[i], None, names[i]) for i in range(0, size)] + else: + return [self.intvar(lb, ub, names[i], bounded_domain) for i in range(0, size)] + elif isinstance(size, tuple): + # Case 2D array + assert len(size) == 2, "Only 2D matrix of intvars are currently supported" + nrows = size[0] + ncols = size[1] + names = [[None for i in range(0, ncols)] for j in range(0, nrows)] + if name is not None: + names = [["{}_{},{}".format(name, j, i) for i in range(0, ncols)] for j in range(0, nrows)] + if isinstance(lb, list): + assert len(lb) == nrows and len(lb[0]) == ncols, "The value list has wrong dimensions" + return [[self.intvar(lb[j][i], None, names[j][i]) for i in range(0, ncols)] for j in range(0, nrows)] + else: + return [[self.intvar(lb, ub, names[j][i], bounded_domain) for i in range(0, ncols)] for j in range(0, nrows)] # Boolean variables # @@ -101,24 +116,38 @@ def boolvar(self, value: Union[bool, None] = None, name: Union[str, None] = None var_handle = backend.boolvar(self.handle) return BoolVar(var_handle, self) - def boolvars(self, size: int, value: Union[List[bool], bool, None] = None, name: Union[str, None] = None): + def boolvars(self, size: Union[int, Tuple[int]], value: Union[List[Union[bool, List[bool]]], None] = None, name: Union[str, None] = None): """ Creates a list of boolvars. - :param size: Number of boolvars. + :param size: Number of boolvars. Either an int or a two-values tuple describing matrix dimensions (nrows, ncols) :param value: If not None, a fixed value for the variables (which is thus a constant). This value is either the same for all variables (bool), or given as a list of bools. :param name: Prefix name of the variable (optional). :return: A list of boolvars. """ - names = [None for i in range(0, size)] - if name is not None: - names = ["{}_{}".format(name, i) for i in range(0, size)] - if isinstance(value, list): - assert len(value) == size - return [self.boolvar(value[i], names[i]) for i in range(0, size)] - else: - return [self.boolvar(value, names[i]) for i in range(0, size)] + if isinstance(size, int): + names = [None for i in range(0, size)] + if name is not None: + names = ["{}_{}".format(name, i) for i in range(0, size)] + if isinstance(value, list): + assert len(value) == size + return [self.boolvar(value[i], names[i]) for i in range(0, size)] + else: + return [self.boolvar(value, names[i]) for i in range(0, size)] + elif isinstance(size, tuple): + # Case 2D array + assert len(size) == 2, "Only 2D matrix of boolvars are currently supported" + nrows = size[0] + ncols = size[1] + names = [[None for i in range(0, ncols)] for j in range(0, nrows)] + if name is not None: + names = [["{}_{},{}".format(name, j, i) for i in range(0, ncols)] for j in range(0, nrows)] + if isinstance(value, list): + assert len(value) == nrows and len(value[0]) == ncols, "The value list has wrong dimensions" + return [[self.boolvar(value[j][i], name=names[j][i]) for i in range(0, ncols)] for j in range(0, nrows)] + else: + return [[self.boolvar(name=names[j][i]) for i in range(0, ncols)] for j in range(0, nrows)] # Task variables # diff --git a/tests/test_boolvar.py b/tests/test_boolvar.py index 4270303..0b09a26 100644 --- a/tests/test_boolvar.py +++ b/tests/test_boolvar.py @@ -68,6 +68,17 @@ def test_instantiate_array(self): s = True self.assertTrue(s) + def test_create_shape(self): + model = Model() + vars = model.boolvars((3, 4), name="var") + self.assertTrue(len(vars) == 3) + self.assertTrue(len(vars[0]) == 4) + self.assertTrue(vars[1][2].name == "var_1,2") + vals = [[False, False, True], + [True, True, False]] + others = model.boolvars((2, 3), vals) + self.assertEqual(others[1][1].get_value(), True) + def test_and(self): m = Model() a = m.boolvar() diff --git a/tests/test_intvars.py b/tests/test_intvars.py index ffcdedb..f668b9b 100644 --- a/tests/test_intvars.py +++ b/tests/test_intvars.py @@ -21,6 +21,17 @@ def test_create_intvar(self): self.assertEqual(b.get_lb(), 1) self.assertEqual(b.get_ub(), 2) + def test_create_shape(self): + model = Model() + vars = model.intvars((3, 4), 0, 10, name="var") + self.assertTrue(len(vars) == 3) + self.assertTrue(len(vars[0]) == 4) + self.assertTrue(vars[1][2].name == "var_1,2") + vals = [[0, 1, 2], + [3, 4, 5]] + others = model.intvars((2, 3), vals) + self.assertEqual(others[1][1].get_ub(), 4) + def test_created_enumerated(self): m = Model() a = m.intvar([0, 1, 4, 5], name="enum_a")