improve cp-sat python types

This commit is contained in:
Laurent Perron
2023-11-22 17:33:01 +01:00
parent dd765fbbc5
commit ddf88f35a9
16 changed files with 94 additions and 106 deletions

View File

@@ -726,7 +726,7 @@ def create_data_model() -> tuple[pd.DataFrame, pd.DataFrame]:
return capacity_df, tasks_df
def main():
def main() -> None:
"""Create the model and solves it."""
capacity_df, tasks_df = create_data_model()
@@ -834,7 +834,7 @@ def rank_tasks(
starts: list[cp_model.IntVar],
presences: list[cp_model.IntVar],
ranks: list[cp_model.IntVar],
):
) -> None:
"""This method adds constraints and variables to links tasks and ranks.
This method assumes that all starts are disjoint, meaning that all tasks have
@@ -852,7 +852,7 @@ def rank_tasks(
all_tasks = range(num_tasks)
# Creates precedence variables between pairs of intervals.
precedences = {}
precedences: dict[tuple[int, int], cp_model.IntVar] = {}
for i in all_tasks:
for j in all_tasks:
if i == j:
@@ -865,7 +865,10 @@ def rank_tasks(
# Treats optional intervals.
for i in range(num_tasks - 1):
for j in range(i + 1, num_tasks):
tmp_array = [precedences[(i, j)], precedences[(j, i)]]
tmp_array: list[cp_model.LiteralT] = [
precedences[(i, j)],
precedences[(j, i)],
]
if not cp_model.object_is_a_true_literal(presences[i]):
tmp_array.append(presences[i].negated())
# Makes sure that if i is not performed, all precedences are false.
@@ -898,7 +901,7 @@ def rank_tasks(
model.add(ranks[i] == sum(precedences[(j, i)] for j in all_tasks) - 1)
def ranking_sample_sat():
def ranking_sample_sat() -> None:
"""Ranks tasks in a NoOverlap constraint."""
model = cp_model.CpModel()

View File

@@ -101,7 +101,7 @@ parallelism. Therefore, the number of workers must be set to 1.
from ortools.sat.python import cp_model
def main():
def main() -> None:
"""Showcases assumptions."""
# Creates the model.
model = cp_model.CpModel()

View File

@@ -56,6 +56,7 @@ from typing import (
Dict,
Iterable,
List,
NoReturn,
Optional,
Sequence,
Tuple,
@@ -80,48 +81,30 @@ Domain = sorted_interval_list.Domain
# usual arithmetic operators + - * / and with constant numbers, which makes the
# python API very intuitive. See../ samples/*.py for examples.
INT_MIN = -9223372036854775808 # hardcoded to be platform independent.
INT_MAX = 9223372036854775807
INT32_MAX = 2147483647
INT32_MIN = -2147483648
INT_MIN = -(2**63) # hardcoded to be platform independent.
INT_MAX = 2**63 - 1
INT32_MIN = -(2**31)
INT32_MAX = 2**31 - 1
# CpSolver status (exported to avoid importing cp_model_cp2).
UNKNOWN: cp_model_pb2.CpSolverStatus = cp_model_pb2.UNKNOWN
MODEL_INVALID: cp_model_pb2.CpSolverStatus = cp_model_pb2.MODEL_INVALID
FEASIBLE: cp_model_pb2.CpSolverStatus = cp_model_pb2.FEASIBLE
INFEASIBLE: cp_model_pb2.CpSolverStatus = cp_model_pb2.INFEASIBLE
OPTIMAL: cp_model_pb2.CpSolverStatus = cp_model_pb2.OPTIMAL
UNKNOWN = cp_model_pb2.UNKNOWN
MODEL_INVALID = cp_model_pb2.MODEL_INVALID
FEASIBLE = cp_model_pb2.FEASIBLE
INFEASIBLE = cp_model_pb2.INFEASIBLE
OPTIMAL = cp_model_pb2.OPTIMAL
# Variable selection strategy
CHOOSE_FIRST: cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy = (
cp_model_pb2.DecisionStrategyProto.CHOOSE_FIRST
)
CHOOSE_LOWEST_MIN: (
cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy
) = cp_model_pb2.DecisionStrategyProto.CHOOSE_LOWEST_MIN
CHOOSE_HIGHEST_MAX: (
cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy
) = cp_model_pb2.DecisionStrategyProto.CHOOSE_HIGHEST_MAX
CHOOSE_MIN_DOMAIN_SIZE: (
cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy
) = cp_model_pb2.DecisionStrategyProto.CHOOSE_MIN_DOMAIN_SIZE
CHOOSE_MAX_DOMAIN_SIZE: (
cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy
) = cp_model_pb2.DecisionStrategyProto.CHOOSE_MAX_DOMAIN_SIZE
CHOOSE_FIRST = cp_model_pb2.DecisionStrategyProto.CHOOSE_FIRST
CHOOSE_LOWEST_MIN = cp_model_pb2.DecisionStrategyProto.CHOOSE_LOWEST_MIN
CHOOSE_HIGHEST_MAX = cp_model_pb2.DecisionStrategyProto.CHOOSE_HIGHEST_MAX
CHOOSE_MIN_DOMAIN_SIZE = cp_model_pb2.DecisionStrategyProto.CHOOSE_MIN_DOMAIN_SIZE
CHOOSE_MAX_DOMAIN_SIZE = cp_model_pb2.DecisionStrategyProto.CHOOSE_MAX_DOMAIN_SIZE
# Domain reduction strategy
SELECT_MIN_VALUE: cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy = (
cp_model_pb2.DecisionStrategyProto.SELECT_MIN_VALUE
)
SELECT_MAX_VALUE: cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy = (
cp_model_pb2.DecisionStrategyProto.SELECT_MAX_VALUE
)
SELECT_LOWER_HALF: (
cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy
) = cp_model_pb2.DecisionStrategyProto.SELECT_LOWER_HALF
SELECT_UPPER_HALF: (
cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy
) = cp_model_pb2.DecisionStrategyProto.SELECT_UPPER_HALF
SELECT_MIN_VALUE = cp_model_pb2.DecisionStrategyProto.SELECT_MIN_VALUE
SELECT_MAX_VALUE = cp_model_pb2.DecisionStrategyProto.SELECT_MAX_VALUE
SELECT_LOWER_HALF = cp_model_pb2.DecisionStrategyProto.SELECT_LOWER_HALF
SELECT_UPPER_HALF = cp_model_pb2.DecisionStrategyProto.SELECT_UPPER_HALF
# Search branching
AUTOMATIC_SEARCH = sat_parameters_pb2.SatParameters.AUTOMATIC_SEARCH
@@ -144,6 +127,7 @@ BoolVarT = Union["IntVar", "_NotBooleanVariable"]
VariableT = Union["IntVar", IntegralT]
LinearExprT = Union["LinearExpr", "IntVar", IntegralT]
ObjLinearExprT = Union["LinearExpr", "IntVar", NumberT]
BoundedLinearExprT = Union["BoundedLinearExpression", bool]
ArcT = Tuple[IntegralT, IntegralT, LiteralT]
_IndexOrSeries = Union[pd.Index, pd.Series]
@@ -405,34 +389,34 @@ class LinearExpr:
break
return coeffs, constant, is_integer
def __hash__(self):
def __hash__(self) -> int:
return object.__hash__(self)
def __abs__(self):
def __abs__(self) -> NoReturn:
raise NotImplementedError(
"calling abs() on a linear expression is not supported, "
"please use CpModel.add_abs_equality"
)
def __add__(self, arg):
def __add__(self, arg) -> LinearExprT:
if cmh.is_zero(arg):
return self
return _Sum(self, arg)
def __radd__(self, arg):
def __radd__(self, arg) -> LinearExprT:
if cmh.is_zero(arg):
return self
return _Sum(self, arg)
def __sub__(self, arg):
def __sub__(self, arg) -> LinearExprT:
if cmh.is_zero(arg):
return self
return _Sum(self, -arg)
def __rsub__(self, arg):
def __rsub__(self, arg) -> LinearExprT:
return _Sum(-self, arg)
def __mul__(self, arg):
def __mul__(self, arg) -> LinearExprT:
arg = cmh.assert_is_a_number(arg)
if cmh.is_one(arg):
return self
@@ -440,7 +424,7 @@ class LinearExpr:
return 0
return _ProductCst(self, arg)
def __rmul__(self, arg):
def __rmul__(self, arg) -> LinearExprT:
arg = cmh.assert_is_a_number(arg)
if cmh.is_one(arg):
return self
@@ -448,67 +432,67 @@ class LinearExpr:
return 0
return _ProductCst(self, arg)
def __div__(self, _):
def __div__(self, _) -> NoReturn:
raise NotImplementedError(
"calling / on a linear expression is not supported, "
"please use CpModel.add_division_equality"
)
def __truediv__(self, _):
def __truediv__(self, _) -> NoReturn:
raise NotImplementedError(
"calling // on a linear expression is not supported, "
"please use CpModel.add_division_equality"
)
def __mod__(self, _):
def __mod__(self, _) -> NoReturn:
raise NotImplementedError(
"calling %% on a linear expression is not supported, "
"please use CpModel.add_modulo_equality"
)
def __pow__(self, _):
def __pow__(self, _) -> NoReturn:
raise NotImplementedError(
"calling ** on a linear expression is not supported, "
"please use CpModel.add_multiplication_equality"
)
def __lshift__(self, _):
def __lshift__(self, _) -> NoReturn:
raise NotImplementedError(
"calling left shift on a linear expression is not supported"
)
def __rshift__(self, _):
def __rshift__(self, _) -> NoReturn:
raise NotImplementedError(
"calling right shift on a linear expression is not supported"
)
def __and__(self, _):
def __and__(self, _) -> NoReturn:
raise NotImplementedError(
"calling and on a linear expression is not supported, "
"please use CpModel.add_bool_and"
)
def __or__(self, _):
def __or__(self, _) -> NoReturn:
raise NotImplementedError(
"calling or on a linear expression is not supported, "
"please use CpModel.add_bool_or"
)
def __xor__(self, _):
def __xor__(self, _) -> NoReturn:
raise NotImplementedError(
"calling xor on a linear expression is not supported, "
"please use CpModel.add_bool_xor"
)
def __neg__(self):
def __neg__(self) -> LinearExprT:
return _ProductCst(self, -1)
def __bool__(self):
def __bool__(self) -> NoReturn:
raise NotImplementedError(
"Evaluating a LinearExpr instance as a Boolean is not implemented."
)
def __eq__(self, arg):
def __eq__(self, arg) -> BoundedLinearExprT:
if arg is None:
return False
if cmh.is_integral(arg):
@@ -517,21 +501,21 @@ class LinearExpr:
else:
return BoundedLinearExpression(self - arg, [0, 0])
def __ge__(self, arg):
def __ge__(self, arg) -> BoundedLinearExprT:
if cmh.is_integral(arg):
arg = cmh.assert_is_int64(arg)
return BoundedLinearExpression(self, [arg, INT_MAX])
else:
return BoundedLinearExpression(self - arg, [0, INT_MAX])
def __le__(self, arg):
def __le__(self, arg) -> BoundedLinearExprT:
if cmh.is_integral(arg):
arg = cmh.assert_is_int64(arg)
return BoundedLinearExpression(self, [INT_MIN, arg])
else:
return BoundedLinearExpression(self - arg, [INT_MIN, 0])
def __lt__(self, arg):
def __lt__(self, arg) -> BoundedLinearExprT:
if cmh.is_integral(arg):
arg = cmh.assert_is_int64(arg)
if arg == INT_MIN:
@@ -540,7 +524,7 @@ class LinearExpr:
else:
return BoundedLinearExpression(self - arg, [INT_MIN, -1])
def __gt__(self, arg):
def __gt__(self, arg) -> BoundedLinearExprT:
if cmh.is_integral(arg):
arg = cmh.assert_is_int64(arg)
if arg == INT_MAX:
@@ -549,7 +533,7 @@ class LinearExpr:
else:
return BoundedLinearExpression(self - arg, [1, INT_MAX])
def __ne__(self, arg):
def __ne__(self, arg) -> BoundedLinearExprT:
if arg is None:
return True
if cmh.is_integral(arg):
@@ -904,7 +888,7 @@ class _NotBooleanVariable(LinearExpr):
def name(self) -> str:
return "not(%s)" % str(self.__boolvar)
def __bool__(self) -> bool:
def __bool__(self) -> NoReturn:
raise NotImplementedError(
"Evaluating a literal as a Boolean value is not implemented."
)
@@ -964,8 +948,9 @@ class BoundedLinearExpression:
return self.__bounds
def __bool__(self) -> bool:
if isinstance(self.__expr, LinearExpr):
coeffs_map, constant = self.__expr.get_integer_var_value_map()
expr = self.__expr
if isinstance(expr, LinearExpr):
coeffs_map, constant = expr.get_integer_var_value_map()
all_coeffs = set(coeffs_map.values())
same_var = set([0])
eq_bounds = [0, 0]
@@ -3181,7 +3166,7 @@ class CpSolver:
"""Returns the indices of the infeasible assumptions."""
return self._solution.sufficient_assumptions_for_infeasibility
def status_name(self, status: Optional[cp_model_pb2.CpSolverStatus] = None) -> str:
def status_name(self, status: Optional[Any] = None) -> str:
"""Returns the name of the status returned by solve()."""
if status is None:
status = self._solution.status
@@ -3245,7 +3230,7 @@ class CpSolver:
def SolutionInfo(self) -> str:
return self.solution_info()
def StatusName(self, status: Optional[cp_model_pb2.CpSolverStatus] = None) -> str:
def StatusName(self, status: Optional[Any] = None) -> str:
return self.status_name(status)
def StopSearch(self) -> None:

View File

@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
# Data
# [START data]
costs = [

View File

@@ -24,7 +24,7 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
# Data
# [START data_model]
data_str = """

View File

@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
# Data
# [START data]
costs = [

View File

@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
# Data
# [START data]
costs = [

View File

@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
"""Showcases assumptions."""
# Creates the model.
# [START model]

View File

@@ -61,7 +61,7 @@ def create_data_model() -> tuple[pd.DataFrame, pd.DataFrame]:
# [END data_model]
def main():
def main() -> None:
# [START data]
items, bins = create_data_model()
# [END data]

View File

@@ -46,7 +46,7 @@ class VarArraySolutionPrinter(cp_model.CpSolverSolutionCallback):
# [END solution_printer]
def main():
def main() -> None:
"""solve the CP+IS+FUN==TRUE cryptarithm."""
# Constraint programming engine
# [START model]

View File

@@ -20,7 +20,7 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
"""Minimal jobshop problem."""
# Data.
# [START data]

View File

@@ -19,18 +19,18 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
# [START data]
data = {}
data["weights"] = [48, 30, 42, 36, 36, 48, 42, 42, 36, 24, 30, 30, 42, 36, 36]
data["values"] = [10, 30, 25, 50, 35, 30, 15, 40, 30, 35, 45, 10, 20, 30, 25]
assert len(data["weights"]) == len(data["values"])
data["num_items"] = len(data["weights"])
data["all_items"] = range(data["num_items"])
num_items = len(data["weights"])
all_items = range(num_items)
data["bin_capacities"] = [100, 100, 100, 100, 100]
data["num_bins"] = len(data["bin_capacities"])
data["all_bins"] = range(data["num_bins"])
num_bins = len(data["bin_capacities"])
all_bins = range(num_bins)
# [END data]
# [START model]
@@ -41,21 +41,21 @@ def main():
# [START variables]
# x[i, b] = 1 if item i is packed in bin b.
x = {}
for i in data["all_items"]:
for b in data["all_bins"]:
for i in all_items:
for b in all_bins:
x[i, b] = model.new_bool_var(f"x_{i}_{b}")
# [END variables]
# Constraints.
# [START constraints]
# Each item is assigned to at most one bin.
for i in data["all_items"]:
model.add_at_most_one(x[i, b] for b in data["all_bins"])
for i in all_items:
model.add_at_most_one(x[i, b] for b in all_bins)
# The amount packed in each bin cannot exceed its capacity.
for b in data["all_bins"]:
for b in all_bins:
model.add(
sum(x[i, b] * data["weights"][i] for i in data["all_items"])
sum(x[i, b] * data["weights"][i] for i in all_items)
<= data["bin_capacities"][b]
)
# [END constraints]
@@ -64,8 +64,8 @@ def main():
# [START objective]
# maximize total value of packed items.
objective = []
for i in data["all_items"]:
for b in data["all_bins"]:
for i in all_items:
for b in all_bins:
objective.append(cp_model.LinearExpr.term(x[i, b], data["values"][i]))
model.maximize(cp_model.LinearExpr.sum(objective))
# [END objective]
@@ -79,11 +79,11 @@ def main():
if status == cp_model.OPTIMAL:
print(f"Total packed value: {solver.objective_value}")
total_weight = 0
for b in data["all_bins"]:
for b in all_bins:
print(f"Bin {b}")
bin_weight = 0
bin_value = 0
for i in data["all_items"]:
for i in all_items:
if solver.value(x[i, b]) > 0:
print(
f"Item:{i} weight:{data['weights'][i]} value:{data['values'][i]}"

View File

@@ -57,7 +57,7 @@ class NQueenSolutionPrinter(cp_model.CpSolverSolutionCallback):
# [END solution_printer]
def main(board_size):
def main(board_size: int) -> None:
# Creates the solver.
# [START model]
model = cp_model.CpModel()

View File

@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
# Data.
# [START data]
num_nurses = 4

View File

@@ -20,9 +20,9 @@ from ortools.sat.python import cp_model
def rank_tasks(
model: cp_model.CpModel,
starts: list[cp_model.IntVar],
presences: list[cp_model.IntVar],
presences: list[cp_model.BoolVarT],
ranks: list[cp_model.IntVar],
):
) -> None:
"""This method adds constraints and variables to links tasks and ranks.
This method assumes that all starts are disjoint, meaning that all tasks have
@@ -32,7 +32,7 @@ def rank_tasks(
Args:
model: The CpModel to add the constraints to.
starts: The array of starts variables of all tasks.
presences: The array of presence variables of all tasks.
presences: The array of presence variables or constants of all tasks.
ranks: The array of rank variables of all tasks.
"""
@@ -40,7 +40,7 @@ def rank_tasks(
all_tasks = range(num_tasks)
# Creates precedence variables between pairs of intervals.
precedences = {}
precedences: dict[tuple[int, int], cp_model.BoolVarT] = {}
for i in all_tasks:
for j in all_tasks:
if i == j:
@@ -53,7 +53,7 @@ def rank_tasks(
# Treats optional intervals.
for i in range(num_tasks - 1):
for j in range(i + 1, num_tasks):
tmp_array: list[cp_model.LiteralT] = [
tmp_array: list[cp_model.BoolVarT] = [
precedences[(i, j)],
precedences[(j, i)],
]
@@ -89,7 +89,7 @@ def rank_tasks(
model.add(ranks[i] == sum(precedences[(j, i)] for j in all_tasks) - 1)
def ranking_sample_sat():
def ranking_sample_sat() -> None:
"""Ranks tasks in a NoOverlap constraint."""
model = cp_model.CpModel()
@@ -100,7 +100,7 @@ def ranking_sample_sat():
starts = []
ends = []
intervals = []
presences = []
presences: list[cp_model.BoolVarT] = []
ranks = []
# Creates intervals, half of them are optional.
@@ -110,7 +110,7 @@ def ranking_sample_sat():
end = model.new_int_var(0, horizon, f"end[{t}]")
if t < num_tasks // 2:
interval = model.new_interval_var(start, duration, end, f"interval[{t}]")
presence = True
presence = model.new_constant(1)
else:
presence = model.new_bool_var(f"presence[{t}]")
interval = model.new_optional_interval_var(

View File

@@ -19,7 +19,7 @@ from ortools.sat.python import cp_model
# [END import]
def main():
def main() -> None:
# This program tries to find an optimal assignment of nurses to shifts
# (3 shifts per day, for 7 days), subject to some constraints (see below).
# Each nurse can request to be assigned to specific shifts.