1# Copyright 2010-2025 Google LLC
  2# Licensed under the Apache License, Version 2.0 (the "License");
  3# you may not use this file except in compliance with the License.
  4# You may obtain a copy of the License at
  5#
  6#     http://www.apache.org/licenses/LICENSE-2.0
  7#
  8# Unless required by applicable law or agreed to in writing, software
  9# distributed under the License is distributed on an "AS IS" BASIS,
 10# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 11# See the License for the specific language governing permissions and
 12# limitations under the License.
 13
 14"""A solver independent library for modeling optimization problems.
 15
 16Example use to model the optimization problem:
 17   max 2.0 * x + y
 18   s.t. x + y <= 1.5
 19            x in {0.0, 1.0}
 20            y in [0.0, 2.5]
 21
 22  model = mathopt.Model(name='my_model')
 23  x = model.add_binary_variable(name='x')
 24  y = model.add_variable(lb=0.0, ub=2.5, name='y')
 25  # We can directly use linear combinations of variables ...
 26  model.add_linear_constraint(x + y <= 1.5, name='c')
 27  # ... or build them incrementally.
 28  objective_expression = 0
 29  objective_expression += 2 * x
 30  objective_expression += y
 31  model.maximize(objective_expression)
 32
 33  # May raise a RuntimeError on invalid input or internal solver errors.
 34  result = mathopt.solve(model, mathopt.SolverType.GSCIP)
 35
 36  if result.termination.reason not in (mathopt.TerminationReason.OPTIMAL,
 37                                       mathopt.TerminationReason.FEASIBLE):
 38    raise RuntimeError(f'model failed to solve: {result.termination}')
 39
 40  print(f'Objective value: {result.objective_value()}')
 41  print(f'Value for variable x: {result.variable_values()[x]}')
 42"""
 43
 44import math
 45from typing import Iterator, Optional, Tuple, Union
 46
 47from ortools.math_opt import model_pb2
 48from ortools.math_opt import model_update_pb2
 49from ortools.math_opt.elemental.python import cpp_elemental
 50from ortools.math_opt.elemental.python import enums
 51from ortools.math_opt.python import from_model
 52from ortools.math_opt.python import indicator_constraints
 53from ortools.math_opt.python import linear_constraints as linear_constraints_mod
 54from ortools.math_opt.python import normalized_inequality
 55from ortools.math_opt.python import objectives
 56from ortools.math_opt.python import quadratic_constraints
 57from ortools.math_opt.python import variables as variables_mod
 58from ortools.math_opt.python.elemental import elemental
 59
 60
 61class UpdateTracker:
 62    """Tracks updates to an optimization model from a ModelStorage.
 63
 64    Do not instantiate directly, instead create through
 65    ModelStorage.add_update_tracker().
 66
 67    Querying an UpdateTracker after calling Model.remove_update_tracker will
 68    result in a model_storage.UsedUpdateTrackerAfterRemovalError.
 69
 70    Example:
 71      mod = Model()
 72      x = mod.add_variable(0.0, 1.0, True, 'x')
 73      y = mod.add_variable(0.0, 1.0, True, 'y')
 74      tracker = mod.add_update_tracker()
 75      mod.set_variable_ub(x, 3.0)
 76      tracker.export_update()
 77        => "variable_updates: {upper_bounds: {ids: [0], values[3.0] }"
 78      mod.set_variable_ub(y, 2.0)
 79      tracker.export_update()
 80        => "variable_updates: {upper_bounds: {ids: [0, 1], values[3.0, 2.0] }"
 81      tracker.advance_checkpoint()
 82      tracker.export_update()
 83        => None
 84      mod.set_variable_ub(y, 4.0)
 85      tracker.export_update()
 86        => "variable_updates: {upper_bounds: {ids: [1], values[4.0] }"
 87      tracker.advance_checkpoint()
 88      mod.remove_update_tracker(tracker)
 89    """
 90
 91    def __init__(
 92        self,
 93        diff_id: int,
 94        elem: elemental.Elemental,
 95    ):
 96        """Do not invoke directly, use Model.add_update_tracker() instead."""
 97        self._diff_id = diff_id
 98        self._elemental = elem
 99
100    def export_update(
101        self, *, remove_names: bool = False
102    ) -> Optional[model_update_pb2.ModelUpdateProto]:
103        """Returns changes to the model since last call to checkpoint/creation."""
104        return self._elemental.export_model_update(
105            self._diff_id, remove_names=remove_names
106        )
107
108    def advance_checkpoint(self) -> None:
109        """Track changes to the model only after this function call."""
110        return self._elemental.advance_diff(self._diff_id)
111
112    @property
113    def diff_id(self) -> int:
114        return self._diff_id
115
116
117class Model:
118    """An optimization model.
119
120    The objective function of the model can be linear or quadratic, and some
121    solvers can only handle linear objective functions. For this reason Model has
122    three versions of all objective setting functions:
123      * A generic one (e.g. maximize()), which accepts linear or quadratic
124        expressions,
125      * a quadratic version (e.g. maximize_quadratic_objective()), which also
126        accepts linear or quadratic expressions and can be used to signal a
127        quadratic objective is possible, and
128      * a linear version (e.g. maximize_linear_objective()), which only accepts
129        linear expressions and can be used to avoid solve time errors for solvers
130        that do not accept quadratic objectives.
131
132    Attributes:
133      name: A description of the problem, can be empty.
134      objective: A function to maximize or minimize.
135      storage: Implementation detail, do not access directly.
136      _variable_ids: Maps variable ids to Variable objects.
137      _linear_constraint_ids: Maps linear constraint ids to LinearConstraint
138        objects.
139    """
140
141    __slots__ = ("_elemental",)
142
143    def __init__(
144        self,
145        *,
146        name: str = "",  # TODO(b/371236599): rename to model_name
147        primary_objective_name: str = "",
148    ) -> None:
149        self._elemental: elemental.Elemental = cpp_elemental.CppElemental(
150            model_name=name, primary_objective_name=primary_objective_name
151        )
152
153    @property
154    def name(self) -> str:
155        return self._elemental.model_name
156
157    ##############################################################################
158    # Variables
159    ##############################################################################
160
161    def add_variable(
162        self,
163        *,
164        lb: float = -math.inf,
165        ub: float = math.inf,
166        is_integer: bool = False,
167        name: str = "",
168    ) -> variables_mod.Variable:
169        """Adds a decision variable to the optimization model.
170
171        Args:
172          lb: The new variable must take at least this value (a lower bound).
173          ub: The new variable must be at most this value (an upper bound).
174          is_integer: Indicates if the variable can only take integer values
175            (otherwise, the variable can take any continuous value).
176          name: For debugging purposes only, but nonempty names must be distinct.
177
178        Returns:
179          A reference to the new decision variable.
180        """
181
182        variable_id = self._elemental.add_element(enums.ElementType.VARIABLE, name)
183        result = variables_mod.Variable(self._elemental, variable_id)
184        result.lower_bound = lb
185        result.upper_bound = ub
186        result.integer = is_integer
187        return result
188
189    def add_integer_variable(
190        self, *, lb: float = -math.inf, ub: float = math.inf, name: str = ""
191    ) -> variables_mod.Variable:
192        return self.add_variable(lb=lb, ub=ub, is_integer=True, name=name)
193
194    def add_binary_variable(self, *, name: str = "") -> variables_mod.Variable:
195        return self.add_variable(lb=0.0, ub=1.0, is_integer=True, name=name)
196
197    def get_variable(
198        self, var_id: int, *, validate: bool = True
199    ) -> variables_mod.Variable:
200        """Returns the Variable for the id var_id, or raises KeyError."""
201        if validate and not self._elemental.element_exists(
202            enums.ElementType.VARIABLE, var_id
203        ):
204            raise KeyError(f"Variable does not exist with id {var_id}.")
205        return variables_mod.Variable(self._elemental, var_id)
206
207    def has_variable(self, var_id: int) -> bool:
208        """Returns true if a Variable with this id is in the model."""
209        return self._elemental.element_exists(enums.ElementType.VARIABLE, var_id)
210
211    def get_num_variables(self) -> int:
212        """Returns the number of variables in the model."""
213        return self._elemental.get_num_elements(enums.ElementType.VARIABLE)
214
215    def get_next_variable_id(self) -> int:
216        """Returns the id of the next variable created in the model."""
217        return self._elemental.get_next_element_id(enums.ElementType.VARIABLE)
218
219    def ensure_next_variable_id_at_least(self, var_id: int) -> None:
220        """If the next variable id would be less than `var_id`, sets it to `var_id`."""
221        self._elemental.ensure_next_element_id_at_least(
222            enums.ElementType.VARIABLE, var_id
223        )
224
225    def delete_variable(self, var: variables_mod.Variable) -> None:
226        """Removes this variable from the model."""
227        self.check_compatible(var)
228        if not self._elemental.delete_element(enums.ElementType.VARIABLE, var.id):
229            raise ValueError(f"Variable with id {var.id} was not in the model.")
230
231    def variables(self) -> Iterator[variables_mod.Variable]:
232        """Yields the variables in the order of creation."""
233        var_ids = self._elemental.get_elements(enums.ElementType.VARIABLE)
234        var_ids.sort()
235        for var_id in var_ids:
236            yield variables_mod.Variable(self._elemental, int(var_id))
237
238    ##############################################################################
239    # Objective
240    ##############################################################################
241
242    @property
243    def objective(self) -> objectives.Objective:
244        return objectives.PrimaryObjective(self._elemental)
245
246    def maximize(self, obj: variables_mod.QuadraticTypes) -> None:
247        """Sets the objective to maximize the provided expression `obj`."""
248        self.set_objective(obj, is_maximize=True)
249
250    def maximize_linear_objective(self, obj: variables_mod.LinearTypes) -> None:
251        """Sets the objective to maximize the provided linear expression `obj`."""
252        self.set_linear_objective(obj, is_maximize=True)
253
254    def maximize_quadratic_objective(self, obj: variables_mod.QuadraticTypes) -> None:
255        """Sets the objective to maximize the provided quadratic expression `obj`."""
256        self.set_quadratic_objective(obj, is_maximize=True)
257
258    def minimize(self, obj: variables_mod.QuadraticTypes) -> None:
259        """Sets the objective to minimize the provided expression `obj`."""
260        self.set_objective(obj, is_maximize=False)
261
262    def minimize_linear_objective(self, obj: variables_mod.LinearTypes) -> None:
263        """Sets the objective to minimize the provided linear expression `obj`."""
264        self.set_linear_objective(obj, is_maximize=False)
265
266    def minimize_quadratic_objective(self, obj: variables_mod.QuadraticTypes) -> None:
267        """Sets the objective to minimize the provided quadratic expression `obj`."""
268        self.set_quadratic_objective(obj, is_maximize=False)
269
270    def set_objective(
271        self, obj: variables_mod.QuadraticTypes, *, is_maximize: bool
272    ) -> None:
273        """Sets the objective to optimize the provided expression `obj`."""
274        self.objective.set_to_expression(obj)
275        self.objective.is_maximize = is_maximize
276
277    def set_linear_objective(
278        self, obj: variables_mod.LinearTypes, *, is_maximize: bool
279    ) -> None:
280        """Sets the objective to optimize the provided linear expression `obj`."""
281        self.objective.set_to_linear_expression(obj)
282        self.objective.is_maximize = is_maximize
283
284    def set_quadratic_objective(
285        self, obj: variables_mod.QuadraticTypes, *, is_maximize: bool
286    ) -> None:
287        """Sets the objective to optimize the provided quadratic expression `obj`."""
288        self.objective.set_to_quadratic_expression(obj)
289        self.objective.is_maximize = is_maximize
290
291    def linear_objective_terms(self) -> Iterator[variables_mod.LinearTerm]:
292        """Yields variable coefficient pairs for variables with nonzero objective coefficient in undefined order."""
293        yield from self.objective.linear_terms()
294
295    def quadratic_objective_terms(self) -> Iterator[variables_mod.QuadraticTerm]:
296        """Yields the quadratic terms with nonzero objective coefficient in undefined order."""
297        yield from self.objective.quadratic_terms()
298
299    ##############################################################################
300    # Auxiliary Objectives
301    ##############################################################################
302
303    def add_auxiliary_objective(
304        self,
305        *,
306        priority: int,
307        name: str = "",
308        expr: Optional[variables_mod.LinearTypes] = None,
309        is_maximize: bool = False,
310    ) -> objectives.AuxiliaryObjective:
311        """Adds an additional objective to the model."""
312        obj_id = self._elemental.add_element(
313            enums.ElementType.AUXILIARY_OBJECTIVE, name
314        )
315        self._elemental.set_attr(
316            enums.IntAttr1.AUXILIARY_OBJECTIVE_PRIORITY, (obj_id,), priority
317        )
318        result = objectives.AuxiliaryObjective(self._elemental, obj_id)
319        if expr is not None:
320            result.set_to_linear_expression(expr)
321        result.is_maximize = is_maximize
322        return result
323
324    def add_maximization_objective(
325        self, expr: variables_mod.LinearTypes, *, priority: int, name: str = ""
326    ) -> objectives.AuxiliaryObjective:
327        """Adds an additional objective to the model that is maximizaition."""
328        result = self.add_auxiliary_objective(
329            priority=priority, name=name, expr=expr, is_maximize=True
330        )
331        return result
332
333    def add_minimization_objective(
334        self, expr: variables_mod.LinearTypes, *, priority: int, name: str = ""
335    ) -> objectives.AuxiliaryObjective:
336        """Adds an additional objective to the model that is minimizaition."""
337        result = self.add_auxiliary_objective(
338            priority=priority, name=name, expr=expr, is_maximize=False
339        )
340        return result
341
342    def delete_auxiliary_objective(self, obj: objectives.AuxiliaryObjective) -> None:
343        """Removes an auxiliary objective from the model."""
344        self.check_compatible(obj)
345        if not self._elemental.delete_element(
346            enums.ElementType.AUXILIARY_OBJECTIVE, obj.id
347        ):
348            raise ValueError(
349                f"Auxiliary objective with id {obj.id} is not in the model."
350            )
351
352    def has_auxiliary_objective(self, obj_id: int) -> bool:
353        """Returns true if the model has an auxiliary objective with id `obj_id`."""
354        return self._elemental.element_exists(
355            enums.ElementType.AUXILIARY_OBJECTIVE, obj_id
356        )
357
358    def next_auxiliary_objective_id(self) -> int:
359        """Returns the id of the next auxiliary objective added to the model."""
360        return self._elemental.get_next_element_id(
361            enums.ElementType.AUXILIARY_OBJECTIVE
362        )
363
364    def num_auxiliary_objectives(self) -> int:
365        """Returns the number of auxiliary objectives in this model."""
366        return self._elemental.get_num_elements(enums.ElementType.AUXILIARY_OBJECTIVE)
367
368    def ensure_next_auxiliary_objective_id_at_least(self, obj_id: int) -> None:
369        """If the next auxiliary objective id would be less than `obj_id`, sets it to `obj_id`."""
370        self._elemental.ensure_next_element_id_at_least(
371            enums.ElementType.AUXILIARY_OBJECTIVE, obj_id
372        )
373
374    def get_auxiliary_objective(
375        self, obj_id: int, *, validate: bool = True
376    ) -> objectives.AuxiliaryObjective:
377        """Returns the auxiliary objective with this id.
378
379        If there is no objective with this id, an exception is thrown if validate is
380        true, and an invalid AuxiliaryObjective is returned if validate is false
381        (later interactions with this object will cause unpredictable errors). Only
382        set validate=False if there is a known performance problem.
383
384        Args:
385          obj_id: The id of the auxiliary objective to look for.
386          validate: Set to false for more speed, but fails to raise an exception if
387            the objective is missing.
388
389        Raises:
390          KeyError: If `validate` is True and there is no objective with this id.
391        """
392        if validate and not self.has_auxiliary_objective(obj_id):
393            raise KeyError(f"Model has no auxiliary objective with id {obj_id}")
394        return objectives.AuxiliaryObjective(self._elemental, obj_id)
395
396    def auxiliary_objectives(self) -> Iterator[objectives.AuxiliaryObjective]:
397        """Returns the auxiliary objectives in the model in the order of creation."""
398        ids = self._elemental.get_elements(enums.ElementType.AUXILIARY_OBJECTIVE)
399        ids.sort()
400        for aux_obj_id in ids:
401            yield objectives.AuxiliaryObjective(self._elemental, int(aux_obj_id))
402
403    ##############################################################################
404    # Linear Constraints
405    ##############################################################################
406
407    # TODO(b/227214976): Update the note below and link to pytype bug number.
408    # Note: bounded_expr's type includes bool only as a workaround to a pytype
409    # issue. Passing a bool for bounded_expr will raise an error in runtime.
410    def add_linear_constraint(
411        self,
412        bounded_expr: Optional[Union[bool, variables_mod.BoundedLinearTypes]] = None,
413        *,
414        lb: Optional[float] = None,
415        ub: Optional[float] = None,
416        expr: Optional[variables_mod.LinearTypes] = None,
417        name: str = "",
418    ) -> linear_constraints_mod.LinearConstraint:
419        """Adds a linear constraint to the optimization model.
420
421        The simplest way to specify the constraint is by passing a one-sided or
422        two-sided linear inequality as in:
423          * add_linear_constraint(x + y + 1.0 <= 2.0),
424          * add_linear_constraint(x + y >= 2.0), or
425          * add_linear_constraint((1.0 <= x + y) <= 2.0).
426
427        Note the extra parenthesis for two-sided linear inequalities, which is
428        required due to some language limitations (see
429        https://peps.python.org/pep-0335/ and https://peps.python.org/pep-0535/).
430        If the parenthesis are omitted, a TypeError will be raised explaining the
431        issue (if this error was not raised the first inequality would have been
432        silently ignored because of the noted language limitations).
433
434        The second way to specify the constraint is by setting lb, ub, and/or expr
435        as in:
436          * add_linear_constraint(expr=x + y + 1.0, ub=2.0),
437          * add_linear_constraint(expr=x + y, lb=2.0),
438          * add_linear_constraint(expr=x + y, lb=1.0, ub=2.0), or
439          * add_linear_constraint(lb=1.0).
440        Omitting lb is equivalent to setting it to -math.inf and omiting ub is
441        equivalent to setting it to math.inf.
442
443        These two alternatives are exclusive and a combined call like:
444          * add_linear_constraint(x + y <= 2.0, lb=1.0), or
445          * add_linear_constraint(x + y <= 2.0, ub=math.inf)
446        will raise a ValueError. A ValueError is also raised if expr's offset is
447        infinite.
448
449        Args:
450          bounded_expr: a linear inequality describing the constraint. Cannot be
451            specified together with lb, ub, or expr.
452          lb: The constraint's lower bound if bounded_expr is omitted (if both
453            bounder_expr and lb are omitted, the lower bound is -math.inf).
454          ub: The constraint's upper bound if bounded_expr is omitted (if both
455            bounder_expr and ub are omitted, the upper bound is math.inf).
456          expr: The constraint's linear expression if bounded_expr is omitted.
457          name: For debugging purposes only, but nonempty names must be distinct.
458
459        Returns:
460          A reference to the new linear constraint.
461        """
462        norm_ineq = normalized_inequality.as_normalized_linear_inequality(
463            bounded_expr, lb=lb, ub=ub, expr=expr
464        )
465        lin_con_id = self._elemental.add_element(
466            enums.ElementType.LINEAR_CONSTRAINT, name
467        )
468
469        result = linear_constraints_mod.LinearConstraint(self._elemental, lin_con_id)
470        result.lower_bound = norm_ineq.lb
471        result.upper_bound = norm_ineq.ub
472        for var, coefficient in norm_ineq.coefficients.items():
473            result.set_coefficient(var, coefficient)
474        return result
475
476    def has_linear_constraint(self, con_id: int) -> bool:
477        """Returns true if a linear constraint with this id is in the model."""
478        return self._elemental.element_exists(
479            enums.ElementType.LINEAR_CONSTRAINT, con_id
480        )
481
482    def get_num_linear_constraints(self) -> int:
483        """Returns the number of linear constraints in the model."""
484        return self._elemental.get_num_elements(enums.ElementType.LINEAR_CONSTRAINT)
485
486    def get_next_linear_constraint_id(self) -> int:
487        """Returns the id of the next linear constraint created in the model."""
488        return self._elemental.get_next_element_id(enums.ElementType.LINEAR_CONSTRAINT)
489
490    def ensure_next_linear_constraint_id_at_least(self, con_id: int) -> None:
491        """If the next linear constraint id would be less than `con_id`, sets it to `con_id`."""
492        self._elemental.ensure_next_element_id_at_least(
493            enums.ElementType.LINEAR_CONSTRAINT, con_id
494        )
495
496    def get_linear_constraint(
497        self, con_id: int, *, validate: bool = True
498    ) -> linear_constraints_mod.LinearConstraint:
499        """Returns the LinearConstraint for the id con_id."""
500        if validate and not self._elemental.element_exists(
501            enums.ElementType.LINEAR_CONSTRAINT, con_id
502        ):
503            raise KeyError(f"Linear constraint does not exist with id {con_id}.")
504        return linear_constraints_mod.LinearConstraint(self._elemental, con_id)
505
506    def delete_linear_constraint(
507        self, lin_con: linear_constraints_mod.LinearConstraint
508    ) -> None:
509        self.check_compatible(lin_con)
510        if not self._elemental.delete_element(
511            enums.ElementType.LINEAR_CONSTRAINT, lin_con.id
512        ):
513            raise ValueError(
514                f"Linear constraint with id {lin_con.id} was not in the model."
515            )
516
517    def linear_constraints(
518        self,
519    ) -> Iterator[linear_constraints_mod.LinearConstraint]:
520        """Yields the linear constraints in the order of creation."""
521        lin_con_ids = self._elemental.get_elements(enums.ElementType.LINEAR_CONSTRAINT)
522        lin_con_ids.sort()
523        for lin_con_id in lin_con_ids:
524            yield linear_constraints_mod.LinearConstraint(
525                self._elemental, int(lin_con_id)
526            )
527
528    def row_nonzeros(
529        self, lin_con: linear_constraints_mod.LinearConstraint
530    ) -> Iterator[variables_mod.Variable]:
531        """Yields the variables with nonzero coefficient for this linear constraint."""
532        keys = self._elemental.slice_attr(
533            enums.DoubleAttr2.LINEAR_CONSTRAINT_COEFFICIENT, 0, lin_con.id
534        )
535        for var_id in keys[:, 1]:
536            yield variables_mod.Variable(self._elemental, int(var_id))
537
538    def column_nonzeros(
539        self, var: variables_mod.Variable
540    ) -> Iterator[linear_constraints_mod.LinearConstraint]:
541        """Yields the linear constraints with nonzero coefficient for this variable."""
542        keys = self._elemental.slice_attr(
543            enums.DoubleAttr2.LINEAR_CONSTRAINT_COEFFICIENT, 1, var.id
544        )
545        for lin_con_id in keys[:, 0]:
546            yield linear_constraints_mod.LinearConstraint(
547                self._elemental, int(lin_con_id)
548            )
549
550    def linear_constraint_matrix_entries(
551        self,
552    ) -> Iterator[linear_constraints_mod.LinearConstraintMatrixEntry]:
553        """Yields the nonzero elements of the linear constraint matrix in undefined order."""
554        keys = self._elemental.get_attr_non_defaults(
555            enums.DoubleAttr2.LINEAR_CONSTRAINT_COEFFICIENT
556        )
557        coefs = self._elemental.get_attrs(
558            enums.DoubleAttr2.LINEAR_CONSTRAINT_COEFFICIENT, keys
559        )
560        for i in range(len(keys)):
561            yield linear_constraints_mod.LinearConstraintMatrixEntry(
562                linear_constraint=linear_constraints_mod.LinearConstraint(
563                    self._elemental, int(keys[i, 0])
564                ),
565                variable=variables_mod.Variable(self._elemental, int(keys[i, 1])),
566                coefficient=float(coefs[i]),
567            )
568
569    ##############################################################################
570    # Quadratic Constraints
571    ##############################################################################
572
573    def add_quadratic_constraint(
574        self,
575        bounded_expr: Optional[
576            Union[
577                bool,
578                variables_mod.BoundedLinearTypes,
579                variables_mod.BoundedQuadraticTypes,
580            ]
581        ] = None,
582        *,
583        lb: Optional[float] = None,
584        ub: Optional[float] = None,
585        expr: Optional[variables_mod.QuadraticTypes] = None,
586        name: str = "",
587    ) -> quadratic_constraints.QuadraticConstraint:
588        """Adds a quadratic constraint to the optimization model.
589
590        The simplest way to specify the constraint is by passing a one-sided or
591        two-sided quadratic inequality as in:
592          * add_quadratic_constraint(x * x + y + 1.0 <= 2.0),
593          * add_quadratic_constraint(x * x + y >= 2.0), or
594          * add_quadratic_constraint((1.0 <= x * x + y) <= 2.0).
595
596        Note the extra parenthesis for two-sided linear inequalities, which is
597        required due to some language limitations (see add_linear_constraint for
598        details).
599
600        The second way to specify the constraint is by setting lb, ub, and/or expr
601        as in:
602          * add_quadratic_constraint(expr=x * x + y + 1.0, ub=2.0),
603          * add_quadratic_constraint(expr=x * x + y, lb=2.0),
604          * add_quadratic_constraint(expr=x * x + y, lb=1.0, ub=2.0), or
605          * add_quadratic_constraint(lb=1.0).
606        Omitting lb is equivalent to setting it to -math.inf and omiting ub is
607        equivalent to setting it to math.inf.
608
609        These two alternatives are exclusive and a combined call like:
610          * add_quadratic_constraint(x * x + y <= 2.0, lb=1.0), or
611          * add_quadratic_constraint(x * x+ y <= 2.0, ub=math.inf)
612        will raise a ValueError. A ValueError is also raised if expr's offset is
613        infinite.
614
615        Args:
616          bounded_expr: a quadratic inequality describing the constraint. Cannot be
617            specified together with lb, ub, or expr.
618          lb: The constraint's lower bound if bounded_expr is omitted (if both
619            bounder_expr and lb are omitted, the lower bound is -math.inf).
620          ub: The constraint's upper bound if bounded_expr is omitted (if both
621            bounder_expr and ub are omitted, the upper bound is math.inf).
622          expr: The constraint's quadratic expression if bounded_expr is omitted.
623          name: For debugging purposes only, but nonempty names must be distinct.
624
625        Returns:
626          A reference to the new quadratic constraint.
627        """
628        norm_quad = normalized_inequality.as_normalized_quadratic_inequality(
629            bounded_expr, lb=lb, ub=ub, expr=expr
630        )
631        quad_con_id = self._elemental.add_element(
632            enums.ElementType.QUADRATIC_CONSTRAINT, name
633        )
634        for var, coef in norm_quad.linear_coefficients.items():
635            self._elemental.set_attr(
636                enums.DoubleAttr2.QUADRATIC_CONSTRAINT_LINEAR_COEFFICIENT,
637                (quad_con_id, var.id),
638                coef,
639            )
640        for key, coef in norm_quad.quadratic_coefficients.items():
641            self._elemental.set_attr(
642                enums.SymmetricDoubleAttr3.QUADRATIC_CONSTRAINT_QUADRATIC_COEFFICIENT,
643                (quad_con_id, key.first_var.id, key.second_var.id),
644                coef,
645            )
646        if norm_quad.lb > -math.inf:
647            self._elemental.set_attr(
648                enums.DoubleAttr1.QUADRATIC_CONSTRAINT_LOWER_BOUND,
649                (quad_con_id,),
650                norm_quad.lb,
651            )
652        if norm_quad.ub < math.inf:
653            self._elemental.set_attr(
654                enums.DoubleAttr1.QUADRATIC_CONSTRAINT_UPPER_BOUND,
655                (quad_con_id,),
656                norm_quad.ub,
657            )
658        return quadratic_constraints.QuadraticConstraint(self._elemental, quad_con_id)
659
660    def has_quadratic_constraint(self, con_id: int) -> bool:
661        """Returns true if a quadratic constraint with this id is in the model."""
662        return self._elemental.element_exists(
663            enums.ElementType.QUADRATIC_CONSTRAINT, con_id
664        )
665
666    def get_num_quadratic_constraints(self) -> int:
667        """Returns the number of quadratic constraints in the model."""
668        return self._elemental.get_num_elements(enums.ElementType.QUADRATIC_CONSTRAINT)
669
670    def get_next_quadratic_constraint_id(self) -> int:
671        """Returns the id of the next quadratic constraint created in the model."""
672        return self._elemental.get_next_element_id(
673            enums.ElementType.QUADRATIC_CONSTRAINT
674        )
675
676    def ensure_next_quadratic_constraint_id_at_least(self, con_id: int) -> None:
677        """If the next quadratic constraint id would be less than `con_id`, sets it to `con_id`."""
678        self._elemental.ensure_next_element_id_at_least(
679            enums.ElementType.QUADRATIC_CONSTRAINT, con_id
680        )
681
682    def get_quadratic_constraint(
683        self, con_id: int, *, validate: bool = True
684    ) -> quadratic_constraints.QuadraticConstraint:
685        """Returns the constraint for the id, or raises KeyError if not in model."""
686        if validate and not self._elemental.element_exists(
687            enums.ElementType.QUADRATIC_CONSTRAINT, con_id
688        ):
689            raise KeyError(f"Quadratic constraint does not exist with id {con_id}.")
690        return quadratic_constraints.QuadraticConstraint(self._elemental, con_id)
691
692    def delete_quadratic_constraint(
693        self, quad_con: quadratic_constraints.QuadraticConstraint
694    ) -> None:
695        """Deletes the constraint with id, or raises ValueError if not in model."""
696        self.check_compatible(quad_con)
697        if not self._elemental.delete_element(
698            enums.ElementType.QUADRATIC_CONSTRAINT, quad_con.id
699        ):
700            raise ValueError(
701                f"Quadratic constraint with id {quad_con.id} was not in the model."
702            )
703
704    def get_quadratic_constraints(
705        self,
706    ) -> Iterator[quadratic_constraints.QuadraticConstraint]:
707        """Yields the quadratic constraints in the order of creation."""
708        quad_con_ids = self._elemental.get_elements(
709            enums.ElementType.QUADRATIC_CONSTRAINT
710        )
711        quad_con_ids.sort()
712        for quad_con_id in quad_con_ids:
713            yield quadratic_constraints.QuadraticConstraint(
714                self._elemental, int(quad_con_id)
715            )
716
717    def quadratic_constraint_linear_nonzeros(
718        self,
719    ) -> Iterator[
720        Tuple[
721            quadratic_constraints.QuadraticConstraint,
722            variables_mod.Variable,
723            float,
724        ]
725    ]:
726        """Yields the linear coefficients for all quadratic constraints in the model."""
727        keys = self._elemental.get_attr_non_defaults(
728            enums.DoubleAttr2.QUADRATIC_CONSTRAINT_LINEAR_COEFFICIENT
729        )
730        coefs = self._elemental.get_attrs(
731            enums.DoubleAttr2.QUADRATIC_CONSTRAINT_LINEAR_COEFFICIENT, keys
732        )
733        for i in range(len(keys)):
734            yield (
735                quadratic_constraints.QuadraticConstraint(
736                    self._elemental, int(keys[i, 0])
737                ),
738                variables_mod.Variable(self._elemental, int(keys[i, 1])),
739                float(coefs[i]),
740            )
741
742    def quadratic_constraint_quadratic_nonzeros(
743        self,
744    ) -> Iterator[
745        Tuple[
746            quadratic_constraints.QuadraticConstraint,
747            variables_mod.Variable,
748            variables_mod.Variable,
749            float,
750        ]
751    ]:
752        """Yields the quadratic coefficients for all quadratic constraints in the model."""
753        keys = self._elemental.get_attr_non_defaults(
754            enums.SymmetricDoubleAttr3.QUADRATIC_CONSTRAINT_QUADRATIC_COEFFICIENT
755        )
756        coefs = self._elemental.get_attrs(
757            enums.SymmetricDoubleAttr3.QUADRATIC_CONSTRAINT_QUADRATIC_COEFFICIENT,
758            keys,
759        )
760        for i in range(len(keys)):
761            yield (
762                quadratic_constraints.QuadraticConstraint(
763                    self._elemental, int(keys[i, 0])
764                ),
765                variables_mod.Variable(self._elemental, int(keys[i, 1])),
766                variables_mod.Variable(self._elemental, int(keys[i, 2])),
767                float(coefs[i]),
768            )
769
770    ##############################################################################
771    # Indicator Constraints
772    ##############################################################################
773
774    def add_indicator_constraint(
775        self,
776        *,
777        indicator: Optional[variables_mod.Variable] = None,
778        activate_on_zero: bool = False,
779        implied_constraint: Optional[
780            Union[bool, variables_mod.BoundedLinearTypes]
781        ] = None,
782        implied_lb: Optional[float] = None,
783        implied_ub: Optional[float] = None,
784        implied_expr: Optional[variables_mod.LinearTypes] = None,
785        name: str = "",
786    ) -> indicator_constraints.IndicatorConstraint:
787        """Adds an indicator constraint to the model.
788
789        If indicator is None or the variable equal to indicator is deleted from
790        the model, the model will be considered invalid at solve time (unless this
791        constraint is also deleted before solving). Likewise, the variable indicator
792        must be binary at solve time for the model to be valid.
793
794        If implied_constraint is set, you may not set implied_lb, implied_ub, or
795        implied_expr.
796
797        Args:
798          indicator: The variable whose value determines if implied_constraint must
799            be enforced.
800          activate_on_zero: If true, implied_constraint must hold when indicator is
801            zero, otherwise, the implied_constraint must hold when indicator is one.
802          implied_constraint: A linear constraint to conditionally enforce, if set.
803            If None, that information is instead passed via implied_lb, implied_ub,
804            and implied_expr.
805          implied_lb: The lower bound of the condtionally enforced linear constraint
806            (or -inf if None), used only when implied_constraint is None.
807          implied_ub: The upper bound of the condtionally enforced linear constraint
808            (or +inf if None), used only when implied_constraint is None.
809          implied_expr: The linear part of the condtionally enforced linear
810            constraint (or 0 if None), used only when implied_constraint is None. If
811            expr has a nonzero offset, it is subtracted from lb and ub.
812          name: For debugging purposes only, but nonempty names must be distinct.
813
814        Returns:
815          A reference to the new indicator constraint.
816        """
817        ind_con_id = self._elemental.add_element(
818            enums.ElementType.INDICATOR_CONSTRAINT, name
819        )
820        if indicator is not None:
821            self._elemental.set_attr(
822                enums.VariableAttr1.INDICATOR_CONSTRAINT_INDICATOR,
823                (ind_con_id,),
824                indicator.id,
825            )
826        self._elemental.set_attr(
827            enums.BoolAttr1.INDICATOR_CONSTRAINT_ACTIVATE_ON_ZERO,
828            (ind_con_id,),
829            activate_on_zero,
830        )
831        implied_inequality = normalized_inequality.as_normalized_linear_inequality(
832            implied_constraint, lb=implied_lb, ub=implied_ub, expr=implied_expr
833        )
834        self._elemental.set_attr(
835            enums.DoubleAttr1.INDICATOR_CONSTRAINT_LOWER_BOUND,
836            (ind_con_id,),
837            implied_inequality.lb,
838        )
839        self._elemental.set_attr(
840            enums.DoubleAttr1.INDICATOR_CONSTRAINT_UPPER_BOUND,
841            (ind_con_id,),
842            implied_inequality.ub,
843        )
844        for var, coef in implied_inequality.coefficients.items():
845            self._elemental.set_attr(
846                enums.DoubleAttr2.INDICATOR_CONSTRAINT_LINEAR_COEFFICIENT,
847                (ind_con_id, var.id),
848                coef,
849            )
850
851        return indicator_constraints.IndicatorConstraint(self._elemental, ind_con_id)
852
853    def has_indicator_constraint(self, con_id: int) -> bool:
854        """Returns true if an indicator constraint with this id is in the model."""
855        return self._elemental.element_exists(
856            enums.ElementType.INDICATOR_CONSTRAINT, con_id
857        )
858
859    def get_num_indicator_constraints(self) -> int:
860        """Returns the number of indicator constraints in the model."""
861        return self._elemental.get_num_elements(enums.ElementType.INDICATOR_CONSTRAINT)
862
863    def get_next_indicator_constraint_id(self) -> int:
864        """Returns the id of the next indicator constraint created in the model."""
865        return self._elemental.get_next_element_id(
866            enums.ElementType.INDICATOR_CONSTRAINT
867        )
868
869    def ensure_next_indicator_constraint_id_at_least(self, con_id: int) -> None:
870        """If the next indicator constraint id would be less than `con_id`, sets it to `con_id`."""
871        self._elemental.ensure_next_element_id_at_least(
872            enums.ElementType.INDICATOR_CONSTRAINT, con_id
873        )
874
875    def get_indicator_constraint(
876        self, con_id: int, *, validate: bool = True
877    ) -> indicator_constraints.IndicatorConstraint:
878        """Returns the IndicatorConstraint for the id con_id."""
879        if validate and not self._elemental.element_exists(
880            enums.ElementType.INDICATOR_CONSTRAINT, con_id
881        ):
882            raise KeyError(f"Indicator constraint does not exist with id {con_id}.")
883        return indicator_constraints.IndicatorConstraint(self._elemental, con_id)
884
885    def delete_indicator_constraint(
886        self, ind_con: indicator_constraints.IndicatorConstraint
887    ) -> None:
888        self.check_compatible(ind_con)
889        if not self._elemental.delete_element(
890            enums.ElementType.INDICATOR_CONSTRAINT, ind_con.id
891        ):
892            raise ValueError(
893                f"Indicator constraint with id {ind_con.id} was not in the model."
894            )
895
896    def get_indicator_constraints(
897        self,
898    ) -> Iterator[indicator_constraints.IndicatorConstraint]:
899        """Yields the indicator constraints in the order of creation."""
900        ind_con_ids = self._elemental.get_elements(
901            enums.ElementType.INDICATOR_CONSTRAINT
902        )
903        ind_con_ids.sort()
904        for ind_con_id in ind_con_ids:
905            yield indicator_constraints.IndicatorConstraint(
906                self._elemental, int(ind_con_id)
907            )
908
909    ##############################################################################
910    # Proto import/export
911    ##############################################################################
912
913    def export_model(self) -> model_pb2.ModelProto:
914        """Returns a protocol buffer equivalent to this model."""
915        return self._elemental.export_model(remove_names=False)
916
917    def add_update_tracker(self) -> UpdateTracker:
918        """Creates an UpdateTracker registered on this model to view changes."""
919        return UpdateTracker(self._elemental.add_diff(), self._elemental)
920
921    def remove_update_tracker(self, tracker: UpdateTracker):
922        """Stops tracker from getting updates on changes to this model.
923
924        An error will be raised if tracker was not created by this Model or if
925        tracker has been previously removed.
926
927        Using (via checkpoint or update) an UpdateTracker after it has been removed
928        will result in an error.
929
930        Args:
931          tracker: The UpdateTracker to unregister.
932
933        Raises:
934          KeyError: The tracker was created by another model or was already removed.
935        """
936        self._elemental.delete_diff(tracker.diff_id)
937
938    def check_compatible(self, e: from_model.FromModel) -> None:
939        """Raises a ValueError if the model of var_or_constraint is not self."""
940        if e.elemental is not self._elemental:
941            raise ValueError(
942                f"Expected element from model named: '{self._elemental.model_name}',"
943                f" but observed element {e} from model named:"
944                f" '{e.elemental.model_name}'."
945            )