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