1# Copyright 2010-2024 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 abc
  45import collections
  46import dataclasses
  47import math
  48import typing
  49from typing import (
  50    Any,
  51    DefaultDict,
  52    Deque,
  53    Generic,
  54    Iterable,
  55    Iterator,
  56    Mapping,
  57    NamedTuple,
  58    NoReturn,
  59    Optional,
  60    Protocol,
  61    Tuple,
  62    Type,
  63    TypeVar,
  64    Union,
  65)
  66import weakref
  67
  68import immutabledict
  69
  70from ortools.math_opt import model_pb2
  71from ortools.math_opt import model_update_pb2
  72from ortools.math_opt.python import hash_model_storage
  73from ortools.math_opt.python import model_storage
  74
  75Storage = model_storage.ModelStorage
  76StorageClass = model_storage.ModelStorageImplClass
  77
  78LinearTypes = Union[int, float, "LinearBase"]
  79QuadraticTypes = Union[int, float, "LinearBase", "QuadraticBase"]
  80LinearTypesExceptVariable = Union[
  81    float, int, "LinearTerm", "LinearExpression", "LinearSum", "LinearProduct"
  82]
  83
  84_CHAINED_COMPARISON_MESSAGE = (
  85    "If you were trying to create a two-sided or "
  86    "ranged linear inequality of the form `lb <= "
  87    "expr <= ub`, try `(lb <= expr) <= ub` instead"
  88)
  89_EXPRESSION_COMP_EXPRESSION_MESSAGE = (
  90    "This error can occur when adding "
  91    "inequalities of the form `(a <= b) <= "
  92    "c` where (a, b, c) includes two or more"
  93    " non-constant linear expressions"
  94)
  95
  96
  97def _raise_binary_operator_type_error(
  98    operator: str,
  99    lhs: Type[Any],
 100    rhs: Type[Any],
 101    extra_message: Optional[str] = None,
 102) -> NoReturn:
 103    """Raises TypeError on unsupported operators."""
 104    message = (
 105        f"unsupported operand type(s) for {operator}: {lhs.__name__!r} and"
 106        f" {rhs.__name__!r}"
 107    )
 108    if extra_message is not None:
 109        message += "\n" + extra_message
 110    raise TypeError(message)
 111
 112
 113def _raise_ne_not_supported() -> NoReturn:
 114    raise TypeError("!= constraints are not supported")
 115
 116
 117class UpperBoundedLinearExpression:
 118    """An inequality of the form expression <= upper_bound.
 119
 120    Where:
 121     * expression is a linear expression, and
 122     * upper_bound is a float
 123    """
 124
 125    __slots__ = "_expression", "_upper_bound"
 126
 127    def __init__(self, expression: "LinearBase", upper_bound: float) -> None:
 128        """Operator overloading can be used instead: e.g. `x + y <= 2.0`."""
 129        self._expression: "LinearBase" = expression
 130        self._upper_bound: float = upper_bound
 131
 132    @property
 133    def expression(self) -> "LinearBase":
 134        return self._expression
 135
 136    @property
 137    def upper_bound(self) -> float:
 138        return self._upper_bound
 139
 140    def __ge__(self, lhs: float) -> "BoundedLinearExpression":
 141        if isinstance(lhs, (int, float)):
 142            return BoundedLinearExpression(lhs, self.expression, self.upper_bound)
 143        _raise_binary_operator_type_error(">=", type(self), type(lhs))
 144
 145    def __bool__(self) -> bool:
 146        raise TypeError(
 147            "__bool__ is unsupported for UpperBoundedLinearExpression"
 148            + "\n"
 149            + _CHAINED_COMPARISON_MESSAGE
 150        )
 151
 152    def __str__(self):
 153        return f"{self._expression!s} <= {self._upper_bound}"
 154
 155    def __repr__(self):
 156        return f"{self._expression!r} <= {self._upper_bound}"
 157
 158
 159class LowerBoundedLinearExpression:
 160    """An inequality of the form expression >= lower_bound.
 161
 162    Where:
 163     * expression is a linear expression, and
 164     * lower_bound is a float
 165    """
 166
 167    __slots__ = "_expression", "_lower_bound"
 168
 169    def __init__(self, expression: "LinearBase", lower_bound: float) -> None:
 170        """Operator overloading can be used instead: e.g. `x + y >= 2.0`."""
 171        self._expression: "LinearBase" = expression
 172        self._lower_bound: float = lower_bound
 173
 174    @property
 175    def expression(self) -> "LinearBase":
 176        return self._expression
 177
 178    @property
 179    def lower_bound(self) -> float:
 180        return self._lower_bound
 181
 182    def __le__(self, rhs: float) -> "BoundedLinearExpression":
 183        if isinstance(rhs, (int, float)):
 184            return BoundedLinearExpression(self.lower_bound, self.expression, rhs)
 185        _raise_binary_operator_type_error("<=", type(self), type(rhs))
 186
 187    def __bool__(self) -> bool:
 188        raise TypeError(
 189            "__bool__ is unsupported for LowerBoundedLinearExpression"
 190            + "\n"
 191            + _CHAINED_COMPARISON_MESSAGE
 192        )
 193
 194    def __str__(self):
 195        return f"{self._expression!s} >= {self._lower_bound}"
 196
 197    def __repr__(self):
 198        return f"{self._expression!r} >= {self._lower_bound}"
 199
 200
 201class BoundedLinearExpression:
 202    """An inequality of the form lower_bound <= expression <= upper_bound.
 203
 204    Where:
 205     * expression is a linear expression
 206     * lower_bound is a float, and
 207     * upper_bound is a float
 208
 209    Note: Because of limitations related to Python's handling of chained
 210    comparisons, bounded expressions cannot be directly created usign
 211    overloaded comparisons as in `lower_bound <= expression <= upper_bound`.
 212    One solution is to wrap one of the inequalities in parenthesis as in
 213    `(lower_bound <= expression) <= upper_bound`.
 214    """
 215
 216    __slots__ = "_expression", "_lower_bound", "_upper_bound"
 217
 218    def __init__(
 219        self, lower_bound: float, expression: "LinearBase", upper_bound: float
 220    ) -> None:
 221        self._expression: "LinearBase" = expression
 222        self._lower_bound: float = lower_bound
 223        self._upper_bound: float = upper_bound
 224
 225    @property
 226    def expression(self) -> "LinearBase":
 227        return self._expression
 228
 229    @property
 230    def lower_bound(self) -> float:
 231        return self._lower_bound
 232
 233    @property
 234    def upper_bound(self) -> float:
 235        return self._upper_bound
 236
 237    def __bool__(self) -> bool:
 238        raise TypeError(
 239            "__bool__ is unsupported for BoundedLinearExpression"
 240            + "\n"
 241            + _CHAINED_COMPARISON_MESSAGE
 242        )
 243
 244    def __str__(self):
 245        return f"{self._lower_bound} <= {self._expression!s} <= {self._upper_bound}"
 246
 247    def __repr__(self):
 248        return f"{self._lower_bound} <= {self._expression!r} <= {self._upper_bound}"
 249
 250
 251class VarEqVar:
 252    """The result of the equality comparison between two Variable.
 253
 254    We use an object here to delay the evaluation of equality so that we can use
 255    the operator== in two use-cases:
 256
 257      1. when the user want to test that two Variable values references the same
 258         variable. This is supported by having this object support implicit
 259         conversion to bool.
 260
 261      2. when the user want to use the equality to create a constraint of equality
 262         between two variables.
 263    """
 264
 265    __slots__ = "_first_variable", "_second_variable"
 266
 267    def __init__(
 268        self,
 269        first_variable: "Variable",
 270        second_variable: "Variable",
 271    ) -> None:
 272        self._first_variable: "Variable" = first_variable
 273        self._second_variable: "Variable" = second_variable
 274
 275    @property
 276    def first_variable(self) -> "Variable":
 277        return self._first_variable
 278
 279    @property
 280    def second_variable(self) -> "Variable":
 281        return self._second_variable
 282
 283    def __bool__(self) -> bool:
 284        return (
 285            self._first_variable.model is self._second_variable.model
 286            and self._first_variable.id == self._second_variable.id
 287        )
 288
 289    def __str__(self):
 290        return f"{self.first_variable!s} == {self._second_variable!s}"
 291
 292    def __repr__(self):
 293        return f"{self.first_variable!r} == {self._second_variable!r}"
 294
 295
 296BoundedLinearTypesList = (
 297    LowerBoundedLinearExpression,
 298    UpperBoundedLinearExpression,
 299    BoundedLinearExpression,
 300    VarEqVar,
 301)
 302BoundedLinearTypes = Union[BoundedLinearTypesList]
 303
 304
 305# TODO(b/231426528): consider using a frozen dataclass.
 306class QuadraticTermKey:
 307    """An id-ordered pair of variables used as a key for quadratic terms."""
 308
 309    __slots__ = "_first_var", "_second_var"
 310
 311    def __init__(self, a: "Variable", b: "Variable"):
 312        """Variables a and b will be ordered internally by their ids."""
 313        self._first_var: "Variable" = a
 314        self._second_var: "Variable" = b
 315        if self._first_var.id > self._second_var.id:
 316            self._first_var, self._second_var = self._second_var, self._first_var
 317
 318    @property
 319    def first_var(self) -> "Variable":
 320        return self._first_var
 321
 322    @property
 323    def second_var(self) -> "Variable":
 324        return self._second_var
 325
 326    def __eq__(self, other: "QuadraticTermKey") -> bool:
 327        return bool(
 328            self._first_var == other._first_var
 329            and self._second_var == other._second_var
 330        )
 331
 332    def __hash__(self) -> int:
 333        return hash((self._first_var, self._second_var))
 334
 335    def __str__(self):
 336        return f"{self._first_var!s} * {self._second_var!s}"
 337
 338    def __repr__(self):
 339        return f"QuadraticTermKey({self._first_var!r}, {self._second_var!r})"
 340
 341
 342@dataclasses.dataclass
 343class _ProcessedElements:
 344    """Auxiliary data class for LinearBase._flatten_once_and_add_to()."""
 345
 346    terms: DefaultDict["Variable", float] = dataclasses.field(
 347        default_factory=lambda: collections.defaultdict(float)
 348    )
 349    offset: float = 0.0
 350
 351
 352@dataclasses.dataclass
 353class _QuadraticProcessedElements(_ProcessedElements):
 354    """Auxiliary data class for QuadraticBase._quadratic_flatten_once_and_add_to()."""
 355
 356    quadratic_terms: DefaultDict["QuadraticTermKey", float] = dataclasses.field(
 357        default_factory=lambda: collections.defaultdict(float)
 358    )
 359
 360
 361class _ToProcessElements(Protocol):
 362    """Auxiliary to-process stack interface for LinearBase._flatten_once_and_add_to() and QuadraticBase._quadratic_flatten_once_and_add_to()."""
 363
 364    __slots__ = ()
 365
 366    def append(self, term: "LinearBase", scale: float) -> None:
 367        """Add a linear object and scale to the to-process stack."""
 368
 369
 370_T = TypeVar("_T", "LinearBase", Union["LinearBase", "QuadraticBase"])
 371
 372
 373class _ToProcessElementsImplementation(Generic[_T]):
 374    """Auxiliary data class for LinearBase._flatten_once_and_add_to()."""
 375
 376    __slots__ = ("_queue",)
 377
 378    def __init__(self, term: _T, scale: float) -> None:
 379        self._queue: Deque[Tuple[_T, float]] = collections.deque([(term, scale)])
 380
 381    def append(self, term: _T, scale: float) -> None:
 382        self._queue.append((term, scale))
 383
 384    def pop(self) -> Tuple[_T, float]:
 385        return self._queue.popleft()
 386
 387    def __bool__(self) -> bool:
 388        return bool(self._queue)
 389
 390
 391_LinearToProcessElements = _ToProcessElementsImplementation["LinearBase"]
 392_QuadraticToProcessElements = _ToProcessElementsImplementation[
 393    Union["LinearBase", "QuadraticBase"]
 394]
 395
 396
 397class LinearBase(metaclass=abc.ABCMeta):
 398    """Interface for types that can build linear expressions with +, -, * and /.
 399
 400    Classes derived from LinearBase (plus float and int scalars) are used to
 401    build expression trees describing a linear expression. Operations nodes of the
 402    expression tree include:
 403
 404      * LinearSum: describes a deferred sum of LinearTypes objects.
 405      * LinearProduct: describes a deferred product of a scalar and a
 406        LinearTypes object.
 407
 408    Leaf nodes of the expression tree include:
 409
 410      * float and int scalars.
 411      * Variable: a single variable.
 412      * LinearTerm: the product of a scalar and a Variable object.
 413      * LinearExpression: the sum of a scalar and LinearTerm objects.
 414
 415    LinearBase objects/expression-trees can be used directly to create
 416    constraints or objective functions. However, to facilitate their inspection,
 417    any LinearTypes object can be flattened to a LinearExpression
 418    through:
 419
 420     as_flat_linear_expression(value: LinearTypes) -> LinearExpression:
 421
 422    In addition, all LinearBase objects are immutable.
 423
 424    Performance notes:
 425
 426    Using an expression tree representation instead of an eager construction of
 427    LinearExpression objects reduces known inefficiencies associated with the
 428    use of operator overloading to construct linear expressions. In particular, we
 429    expect the runtime of as_flat_linear_expression() to be linear in the size of
 430    the expression tree. Additional performance can gained by using LinearSum(c)
 431    instead of sum(c) for a container c, as the latter creates len(c) LinearSum
 432    objects.
 433    """
 434
 435    __slots__ = ()
 436
 437    # TODO(b/216492143): explore requirements for this function so calculation of
 438    # coefficients and offsets follow expected associativity rules (so approximate
 439    # float calculations are as expected).
 440    # TODO(b/216492143): add more details of what subclasses need to do in
 441    # developers guide.
 442    @abc.abstractmethod
 443    def _flatten_once_and_add_to(
 444        self,
 445        scale: float,
 446        processed_elements: _ProcessedElements,
 447        target_stack: _ToProcessElements,
 448    ) -> None:
 449        """Flatten one level of tree if needed and add to targets.
 450
 451        Classes derived from LinearBase only need to implement this function
 452        to enable transformation to LinearExpression through
 453        as_flat_linear_expression().
 454
 455        Args:
 456          scale: multiply elements by this number when processing or adding to
 457            stack.
 458          processed_elements: where to add LinearTerms and scalars that can be
 459            processed immediately.
 460          target_stack: where to add LinearBase elements that are not scalars or
 461            LinearTerms (i.e. elements that need further flattening).
 462            Implementations should append() to this stack to avoid being recursive.
 463        """
 464
 465    def __eq__(
 466        self, rhs: LinearTypes
 467    ) -> (
 468        "BoundedLinearExpression"
 469    ):  # pytype: disable=signature-mismatch  # overriding-return-type-checks
 470        if isinstance(rhs, (int, float)):
 471            return BoundedLinearExpression(rhs, self, rhs)
 472        if not isinstance(rhs, LinearBase):
 473            _raise_binary_operator_type_error("==", type(self), type(rhs))
 474        return BoundedLinearExpression(0.0, self - rhs, 0.0)
 475
 476    def __ne__(
 477        self, rhs: LinearTypes
 478    ) -> (
 479        NoReturn
 480    ):  # pytype: disable=signature-mismatch  # overriding-return-type-checks
 481        _raise_ne_not_supported()
 482
 483    @typing.overload
 484    def __le__(self, rhs: float) -> "UpperBoundedLinearExpression": ...
 485
 486    @typing.overload
 487    def __le__(self, rhs: "LinearBase") -> "BoundedLinearExpression": ...
 488
 489    @typing.overload
 490    def __le__(self, rhs: "BoundedLinearExpression") -> NoReturn: ...
 491
 492    def __le__(self, rhs):
 493        if isinstance(rhs, (int, float)):
 494            return UpperBoundedLinearExpression(self, rhs)
 495        if isinstance(rhs, LinearBase):
 496            return BoundedLinearExpression(-math.inf, self - rhs, 0.0)
 497        if isinstance(rhs, BoundedLinearExpression):
 498            _raise_binary_operator_type_error(
 499                "<=", type(self), type(rhs), _EXPRESSION_COMP_EXPRESSION_MESSAGE
 500            )
 501        _raise_binary_operator_type_error("<=", type(self), type(rhs))
 502
 503    @typing.overload
 504    def __ge__(self, lhs: float) -> "LowerBoundedLinearExpression": ...
 505
 506    @typing.overload
 507    def __ge__(self, lhs: "LinearBase") -> "BoundedLinearExpression": ...
 508
 509    @typing.overload
 510    def __ge__(self, lhs: "BoundedLinearExpression") -> NoReturn: ...
 511
 512    def __ge__(self, lhs):
 513        if isinstance(lhs, (int, float)):
 514            return LowerBoundedLinearExpression(self, lhs)
 515        if isinstance(lhs, LinearBase):
 516            return BoundedLinearExpression(0.0, self - lhs, math.inf)
 517        if isinstance(lhs, BoundedLinearExpression):
 518            _raise_binary_operator_type_error(
 519                ">=", type(self), type(lhs), _EXPRESSION_COMP_EXPRESSION_MESSAGE
 520            )
 521        _raise_binary_operator_type_error(">=", type(self), type(lhs))
 522
 523    def __add__(self, expr: LinearTypes) -> "LinearSum":
 524        if not isinstance(expr, (int, float, LinearBase)):
 525            return NotImplemented
 526        return LinearSum((self, expr))
 527
 528    def __radd__(self, expr: LinearTypes) -> "LinearSum":
 529        if not isinstance(expr, (int, float, LinearBase)):
 530            return NotImplemented
 531        return LinearSum((expr, self))
 532
 533    def __sub__(self, expr: LinearTypes) -> "LinearSum":
 534        if not isinstance(expr, (int, float, LinearBase)):
 535            return NotImplemented
 536        return LinearSum((self, -expr))
 537
 538    def __rsub__(self, expr: LinearTypes) -> "LinearSum":
 539        if not isinstance(expr, (int, float, LinearBase, QuadraticBase)):
 540            return NotImplemented
 541        return LinearSum((expr, -self))
 542
 543    @typing.overload
 544    def __mul__(self, other: float) -> "LinearProduct": ...
 545
 546    @typing.overload
 547    def __mul__(self, other: "LinearBase") -> "LinearLinearProduct": ...
 548
 549    def __mul__(self, other):
 550        if not isinstance(other, (int, float, LinearBase)):
 551            return NotImplemented
 552        if isinstance(other, LinearBase):
 553            return LinearLinearProduct(self, other)
 554        return LinearProduct(other, self)
 555
 556    def __rmul__(self, constant: float) -> "LinearProduct":
 557        if not isinstance(constant, (int, float)):
 558            return NotImplemented
 559        return LinearProduct(constant, self)
 560
 561    # TODO(b/216492143): explore numerical consequences of 1.0 / constant below.
 562    def __truediv__(self, constant: float) -> "LinearProduct":
 563        if not isinstance(constant, (int, float)):
 564            return NotImplemented
 565        return LinearProduct(1.0 / constant, self)
 566
 567    def __neg__(self) -> "LinearProduct":
 568        return LinearProduct(-1.0, self)
 569
 570
 571class QuadraticBase(metaclass=abc.ABCMeta):
 572    """Interface for types that can build quadratic expressions with +, -, * and /.
 573
 574    Classes derived from QuadraticBase and LinearBase (plus float and int scalars)
 575    are used to build expression trees describing a quadratic expression.
 576    Operations nodes of the expression tree include:
 577
 578        * QuadraticSum: describes a deferred sum of QuadraticTypes objects.
 579        * QuadraticProduct: describes a deferred product of a scalar and a
 580          QuadraticTypes object.
 581        * LinearLinearProduct: describes a deferred product of two LinearTypes
 582          objects.
 583
 584      Leaf nodes of the expression tree include:
 585
 586        * float and int scalars.
 587        * Variable: a single variable.
 588        * LinearTerm: the product of a scalar and a Variable object.
 589        * LinearExpression: the sum of a scalar and LinearTerm objects.
 590        * QuadraticTerm: the product of a scalar and two Variable objects.
 591        * QuadraticExpression: the sum of a scalar, LinearTerm objects and
 592          QuadraticTerm objects.
 593
 594    QuadraticBase objects/expression-trees can be used directly to create
 595    objective functions. However, to facilitate their inspection, any
 596    QuadraticTypes object can be flattened to a QuadraticExpression
 597    through:
 598
 599     as_flat_quadratic_expression(value: QuadraticTypes) -> QuadraticExpression:
 600
 601    In addition, all QuadraticBase objects are immutable.
 602
 603    Performance notes:
 604
 605    Using an expression tree representation instead of an eager construction of
 606    QuadraticExpression objects reduces known inefficiencies associated with the
 607    use of operator overloading to construct quadratic expressions. In particular,
 608    we expect the runtime of as_flat_quadratic_expression() to be linear in the
 609    size of the expression tree. Additional performance can gained by using
 610    QuadraticSum(c) instead of sum(c) for a container c, as the latter creates
 611    len(c) QuadraticSum objects.
 612    """
 613
 614    __slots__ = ()
 615
 616    # TODO(b/216492143): explore requirements for this function so calculation of
 617    # coefficients and offsets follow expected associativity rules (so approximate
 618    # float calculations are as expected).
 619    # TODO(b/216492143): add more details of what subclasses need to do in
 620    # developers guide.
 621    @abc.abstractmethod
 622    def _quadratic_flatten_once_and_add_to(
 623        self,
 624        scale: float,
 625        processed_elements: _QuadraticProcessedElements,
 626        target_stack: _QuadraticToProcessElements,
 627    ) -> None:
 628        """Flatten one level of tree if needed and add to targets.
 629
 630        Classes derived from QuadraticBase only need to implement this function
 631        to enable transformation to QuadraticExpression through
 632        as_flat_quadratic_expression().
 633
 634        Args:
 635          scale: multiply elements by this number when processing or adding to
 636            stack.
 637          processed_elements: where to add linear terms, quadratic terms and scalars
 638            that can be processed immediately.
 639          target_stack: where to add LinearBase and QuadraticBase elements that are
 640            not scalars or linear terms or quadratic terms (i.e. elements that need
 641            further flattening). Implementations should append() to this stack to
 642            avoid being recursive.
 643        """
 644
 645    def __add__(self, expr: QuadraticTypes) -> "QuadraticSum":
 646        if not isinstance(expr, (int, float, LinearBase, QuadraticBase)):
 647            return NotImplemented
 648        return QuadraticSum((self, expr))
 649
 650    def __radd__(self, expr: QuadraticTypes) -> "QuadraticSum":
 651        if not isinstance(expr, (int, float, LinearBase, QuadraticBase)):
 652            return NotImplemented
 653        return QuadraticSum((expr, self))
 654
 655    def __sub__(self, expr: QuadraticTypes) -> "QuadraticSum":
 656        if not isinstance(expr, (int, float, LinearBase, QuadraticBase)):
 657            return NotImplemented
 658        return QuadraticSum((self, -expr))
 659
 660    def __rsub__(self, expr: QuadraticTypes) -> "QuadraticSum":
 661        if not isinstance(expr, (int, float, LinearBase, QuadraticBase)):
 662            return NotImplemented
 663        return QuadraticSum((expr, -self))
 664
 665    def __mul__(self, other: float) -> "QuadraticProduct":
 666        if not isinstance(other, (int, float)):
 667            return NotImplemented
 668        return QuadraticProduct(other, self)
 669
 670    def __rmul__(self, other: float) -> "QuadraticProduct":
 671        if not isinstance(other, (int, float)):
 672            return NotImplemented
 673        return QuadraticProduct(other, self)
 674
 675    # TODO(b/216492143): explore numerical consequences of 1.0 / constant below.
 676    def __truediv__(self, constant: float) -> "QuadraticProduct":
 677        if not isinstance(constant, (int, float)):
 678            return NotImplemented
 679        return QuadraticProduct(1.0 / constant, self)
 680
 681    def __neg__(self) -> "QuadraticProduct":
 682        return QuadraticProduct(-1.0, self)
 683
 684
 685class Variable(LinearBase):
 686    """A decision variable for an optimization model.
 687
 688    A decision variable takes a value from a domain, either the real numbers or
 689    the integers, and restricted to be in some interval [lb, ub] (where lb and ub
 690    can be infinite). The case of lb == ub is allowed, this means the variable
 691    must take a single value. The case of lb > ub is also allowed, this implies
 692    that the problem is infeasible.
 693
 694    A Variable is configured as follows:
 695      * lower_bound: a float property, lb above. Should not be NaN nor +inf.
 696      * upper_bound: a float property, ub above. Should not be NaN nor -inf.
 697      * integer: a bool property, if the domain is integer or continuous.
 698
 699    The name is optional, read only, and used only for debugging. Non-empty names
 700    should be distinct.
 701
 702    Every Variable is associated with a Model (defined below). Note that data
 703    describing the variable (e.g. lower_bound) is owned by Model.storage, this
 704    class is simply a reference to that data. Do not create a Variable directly,
 705    use Model.add_variable() instead.
 706    """
 707
 708    __slots__ = "__weakref__", "_model", "_id"
 709
 710    def __init__(self, model: "Model", vid: int) -> None:
 711        """Do not invoke directly, use Model.add_variable()."""
 712        self._model: "Model" = model
 713        self._id: int = vid
 714
 715    @property
 716    def lower_bound(self) -> float:
 717        return self.model.storage.get_variable_lb(self._id)
 718
 719    @lower_bound.setter
 720    def lower_bound(self, value: float) -> None:
 721        self.model.storage.set_variable_lb(self._id, value)
 722
 723    @property
 724    def upper_bound(self) -> float:
 725        return self.model.storage.get_variable_ub(self._id)
 726
 727    @upper_bound.setter
 728    def upper_bound(self, value: float) -> None:
 729        self.model.storage.set_variable_ub(self._id, value)
 730
 731    @property
 732    def integer(self) -> bool:
 733        return self.model.storage.get_variable_is_integer(self._id)
 734
 735    @integer.setter
 736    def integer(self, value: bool) -> None:
 737        self.model.storage.set_variable_is_integer(self._id, value)
 738
 739    @property
 740    def name(self) -> str:
 741        return self.model.storage.get_variable_name(self._id)
 742
 743    @property
 744    def id(self) -> int:
 745        return self._id
 746
 747    @property
 748    def model(self) -> "Model":
 749        return self._model
 750
 751    def __str__(self):
 752        """Returns the name, or a string containing the id if the name is empty."""
 753        return self.name if self.name else f"variable_{self.id}"
 754
 755    def __repr__(self):
 756        return f"<Variable id: {self.id}, name: {self.name!r}>"
 757
 758    @typing.overload
 759    def __eq__(self, rhs: "Variable") -> "VarEqVar": ...
 760
 761    @typing.overload
 762    def __eq__(self, rhs: LinearTypesExceptVariable) -> "BoundedLinearExpression": ...
 763
 764    def __eq__(self, rhs):
 765        if isinstance(rhs, Variable):
 766            return VarEqVar(self, rhs)
 767        return super().__eq__(rhs)
 768
 769    @typing.overload
 770    def __ne__(self, rhs: "Variable") -> bool: ...
 771
 772    @typing.overload
 773    def __ne__(self, rhs: LinearTypesExceptVariable) -> NoReturn: ...
 774
 775    def __ne__(self, rhs):
 776        if isinstance(rhs, Variable):
 777            return not self == rhs
 778        _raise_ne_not_supported()
 779
 780    def __hash__(self) -> int:
 781        return hash((self.model, self.id))
 782
 783    @typing.overload
 784    def __mul__(self, other: float) -> "LinearTerm": ...
 785
 786    @typing.overload
 787    def __mul__(self, other: Union["Variable", "LinearTerm"]) -> "QuadraticTerm": ...
 788
 789    @typing.overload
 790    def __mul__(self, other: "LinearBase") -> "LinearLinearProduct": ...
 791
 792    def __mul__(self, other):
 793        if not isinstance(other, (int, float, LinearBase)):
 794            return NotImplemented
 795        if isinstance(other, Variable):
 796            return QuadraticTerm(QuadraticTermKey(self, other), 1.0)
 797        if isinstance(other, LinearTerm):
 798            return QuadraticTerm(
 799                QuadraticTermKey(self, other.variable), other.coefficient
 800            )
 801        if isinstance(other, LinearBase):
 802            return LinearLinearProduct(self, other)  # pytype: disable=bad-return-type
 803        return LinearTerm(self, other)
 804
 805    def __rmul__(self, constant: float) -> "LinearTerm":
 806        if not isinstance(constant, (int, float)):
 807            return NotImplemented
 808        return LinearTerm(self, constant)
 809
 810    # TODO(b/216492143): explore numerical consequences of 1.0 / constant below.
 811    def __truediv__(self, constant: float) -> "LinearTerm":
 812        if not isinstance(constant, (int, float)):
 813            return NotImplemented
 814        return LinearTerm(self, 1.0 / constant)
 815
 816    def __neg__(self) -> "LinearTerm":
 817        return LinearTerm(self, -1.0)
 818
 819    def _flatten_once_and_add_to(
 820        self,
 821        scale: float,
 822        processed_elements: _ProcessedElements,
 823        target_stack: _ToProcessElements,
 824    ) -> None:
 825        processed_elements.terms[self] += scale
 826
 827
 828class LinearTerm(LinearBase):
 829    """The product of a scalar and a variable.
 830
 831    This class is immutable.
 832    """
 833
 834    __slots__ = "_variable", "_coefficient"
 835
 836    def __init__(self, variable: Variable, coefficient: float) -> None:
 837        self._variable: Variable = variable
 838        self._coefficient: float = coefficient
 839
 840    @property
 841    def variable(self) -> Variable:
 842        return self._variable
 843
 844    @property
 845    def coefficient(self) -> float:
 846        return self._coefficient
 847
 848    def _flatten_once_and_add_to(
 849        self,
 850        scale: float,
 851        processed_elements: _ProcessedElements,
 852        target_stack: _ToProcessElements,
 853    ) -> None:
 854        processed_elements.terms[self._variable] += self._coefficient * scale
 855
 856    @typing.overload
 857    def __mul__(self, other: float) -> "LinearTerm": ...
 858
 859    @typing.overload
 860    def __mul__(self, other: Union["Variable", "LinearTerm"]) -> "QuadraticTerm": ...
 861
 862    @typing.overload
 863    def __mul__(self, other: "LinearBase") -> "LinearLinearProduct": ...
 864
 865    def __mul__(self, other):
 866        if not isinstance(other, (int, float, LinearBase)):
 867            return NotImplemented
 868        if isinstance(other, Variable):
 869            return QuadraticTerm(
 870                QuadraticTermKey(self._variable, other), self._coefficient
 871            )
 872        if isinstance(other, LinearTerm):
 873            return QuadraticTerm(
 874                QuadraticTermKey(self.variable, other.variable),
 875                self._coefficient * other.coefficient,
 876            )
 877        if isinstance(other, LinearBase):
 878            return LinearLinearProduct(self, other)  # pytype: disable=bad-return-type
 879        return LinearTerm(self._variable, self._coefficient * other)
 880
 881    def __rmul__(self, constant: float) -> "LinearTerm":
 882        if not isinstance(constant, (int, float)):
 883            return NotImplemented
 884        return LinearTerm(self._variable, self._coefficient * constant)
 885
 886    def __truediv__(self, constant: float) -> "LinearTerm":
 887        if not isinstance(constant, (int, float)):
 888            return NotImplemented
 889        return LinearTerm(self._variable, self._coefficient / constant)
 890
 891    def __neg__(self) -> "LinearTerm":
 892        return LinearTerm(self._variable, -self._coefficient)
 893
 894    def __str__(self):
 895        return f"{self._coefficient} * {self._variable}"
 896
 897    def __repr__(self):
 898        return f"LinearTerm({self._variable!r}, {self._coefficient!r})"
 899
 900
 901class QuadraticTerm(QuadraticBase):
 902    """The product of a scalar and two variables.
 903
 904    This class is immutable.
 905    """
 906
 907    __slots__ = "_key", "_coefficient"
 908
 909    def __init__(self, key: QuadraticTermKey, coefficient: float) -> None:
 910        self._key: QuadraticTermKey = key
 911        self._coefficient: float = coefficient
 912
 913    @property
 914    def key(self) -> QuadraticTermKey:
 915        return self._key
 916
 917    @property
 918    def coefficient(self) -> float:
 919        return self._coefficient
 920
 921    def _quadratic_flatten_once_and_add_to(
 922        self,
 923        scale: float,
 924        processed_elements: _QuadraticProcessedElements,
 925        target_stack: _ToProcessElements,
 926    ) -> None:
 927        processed_elements.quadratic_terms[self._key] += self._coefficient * scale
 928
 929    def __mul__(self, constant: float) -> "QuadraticTerm":
 930        if not isinstance(constant, (int, float)):
 931            return NotImplemented
 932        return QuadraticTerm(self._key, self._coefficient * constant)
 933
 934    def __rmul__(self, constant: float) -> "QuadraticTerm":
 935        if not isinstance(constant, (int, float)):
 936            return NotImplemented
 937        return QuadraticTerm(self._key, self._coefficient * constant)
 938
 939    def __truediv__(self, constant: float) -> "QuadraticTerm":
 940        if not isinstance(constant, (int, float)):
 941            return NotImplemented
 942        return QuadraticTerm(self._key, self._coefficient / constant)
 943
 944    def __neg__(self) -> "QuadraticTerm":
 945        return QuadraticTerm(self._key, -self._coefficient)
 946
 947    def __str__(self):
 948        return f"{self._coefficient} * {self._key!s}"
 949
 950    def __repr__(self):
 951        return f"QuadraticTerm({self._key!r}, {self._coefficient})"
 952
 953
 954class LinearExpression(LinearBase):
 955    """For variables x, an expression: b + sum_{i in I} a_i * x_i.
 956
 957    This class is immutable.
 958    """
 959
 960    __slots__ = "__weakref__", "_terms", "_offset"
 961
 962    # TODO(b/216492143): consider initializing from a dictionary.
 963    def __init__(self, /, other: LinearTypes = 0) -> None:
 964        self._offset: float = 0.0
 965        if isinstance(other, (int, float)):
 966            self._offset = float(other)
 967            self._terms: Mapping[Variable, float] = immutabledict.immutabledict()
 968            return
 969
 970        to_process: _LinearToProcessElements = _LinearToProcessElements(other, 1.0)
 971        processed_elements = _ProcessedElements()
 972        while to_process:
 973            linear, coef = to_process.pop()
 974            linear._flatten_once_and_add_to(coef, processed_elements, to_process)
 975        # TODO(b/216492143): explore avoiding this copy.
 976        self._terms: Mapping[Variable, float] = immutabledict.immutabledict(
 977            processed_elements.terms
 978        )
 979        self._offset = processed_elements.offset
 980
 981    @property
 982    def terms(self) -> Mapping[Variable, float]:
 983        return self._terms
 984
 985    @property
 986    def offset(self) -> float:
 987        return self._offset
 988
 989    def evaluate(self, variable_values: Mapping[Variable, float]) -> float:
 990        """Returns the value of this expression for given variable values.
 991
 992        E.g. if this is 3 * x + 4 and variable_values = {x: 2.0}, then
 993        evaluate(variable_values) equals 10.0.
 994
 995        See also mathopt.evaluate_expression(), which works on any type in
 996        QuadraticTypes.
 997
 998        Args:
 999          variable_values: Must contain a value for every variable in expression.
1000
1001        Returns:
1002          The value of this expression when replacing variables by their value.
1003        """
1004        result = self._offset
1005        for var, coef in sorted(
1006            self._terms.items(), key=lambda var_coef_pair: var_coef_pair[0].id
1007        ):
1008            result += coef * variable_values[var]
1009        return result
1010
1011    def _flatten_once_and_add_to(
1012        self,
1013        scale: float,
1014        processed_elements: _ProcessedElements,
1015        target_stack: _ToProcessElements,
1016    ) -> None:
1017        for var, val in self._terms.items():
1018            processed_elements.terms[var] += val * scale
1019        processed_elements.offset += scale * self.offset
1020
1021    # TODO(b/216492143): change __str__ to match C++ implementation in
1022    # cl/421649402.
1023    def __str__(self):
1024        """Returns the name, or a string containing the id if the name is empty."""
1025        result = str(self.offset)
1026        sorted_keys = sorted(self._terms.keys(), key=str)
1027        for variable in sorted_keys:
1028            # TODO(b/216492143): consider how to better deal with `NaN` and try to
1029            # match C++ implementation in cl/421649402. See TODO for StrAndReprTest in
1030            # linear_expression_test.py.
1031            coefficient = self._terms[variable]
1032            if coefficient == 0.0:
1033                continue
1034            if coefficient > 0:
1035                result += " + "
1036            else:
1037                result += " - "
1038            result += str(abs(coefficient)) + " * " + str(variable)
1039        return result
1040
1041    def __repr__(self):
1042        result = f"LinearExpression({self.offset}, " + "{"
1043        result += ", ".join(
1044            f"{variable!r}: {coefficient}"
1045            for variable, coefficient in self._terms.items()
1046        )
1047        result += "})"
1048        return result
1049
1050
1051class QuadraticExpression(QuadraticBase):
1052    """For variables x, an expression: b + sum_{i in I} a_i * x_i + sum_{i,j in I, i<=j} a_i,j * x_i * x_j.
1053
1054    This class is immutable.
1055    """
1056
1057    __slots__ = "__weakref__", "_linear_terms", "_quadratic_terms", "_offset"
1058
1059    # TODO(b/216492143): consider initializing from a dictionary.
1060    def __init__(self, other: QuadraticTypes) -> None:
1061        self._offset: float = 0.0
1062        if isinstance(other, (int, float)):
1063            self._offset = float(other)
1064            self._linear_terms: Mapping[Variable, float] = immutabledict.immutabledict()
1065            self._quadratic_terms: Mapping[QuadraticTermKey, float] = (
1066                immutabledict.immutabledict()
1067            )
1068            return
1069
1070        to_process: _QuadraticToProcessElements = _QuadraticToProcessElements(
1071            other, 1.0
1072        )
1073        processed_elements = _QuadraticProcessedElements()
1074        while to_process:
1075            linear_or_quadratic, coef = to_process.pop()
1076            if isinstance(linear_or_quadratic, LinearBase):
1077                linear_or_quadratic._flatten_once_and_add_to(
1078                    coef, processed_elements, to_process
1079                )
1080            else:
1081                linear_or_quadratic._quadratic_flatten_once_and_add_to(
1082                    coef, processed_elements, to_process
1083                )
1084        # TODO(b/216492143): explore avoiding this copy.
1085        self._linear_terms: Mapping[Variable, float] = immutabledict.immutabledict(
1086            processed_elements.terms
1087        )
1088        self._quadratic_terms: Mapping[QuadraticTermKey, float] = (
1089            immutabledict.immutabledict(processed_elements.quadratic_terms)
1090        )
1091        self._offset = processed_elements.offset
1092
1093    @property
1094    def linear_terms(self) -> Mapping[Variable, float]:
1095        return self._linear_terms
1096
1097    @property
1098    def quadratic_terms(self) -> Mapping[QuadraticTermKey, float]:
1099        return self._quadratic_terms
1100
1101    @property
1102    def offset(self) -> float:
1103        return self._offset
1104
1105    def evaluate(self, variable_values: Mapping[Variable, float]) -> float:
1106        """Returns the value of this expression for given variable values.
1107
1108        E.g. if this is 3 * x * x + 4 and variable_values = {x: 2.0}, then
1109        evaluate(variable_values) equals 16.0.
1110
1111        See also mathopt.evaluate_expression(), which works on any type in
1112        QuadraticTypes.
1113
1114        Args:
1115          variable_values: Must contain a value for every variable in expression.
1116
1117        Returns:
1118          The value of this expression when replacing variables by their value.
1119        """
1120        result = self._offset
1121        for var, coef in sorted(
1122            self._linear_terms.items(),
1123            key=lambda var_coef_pair: var_coef_pair[0].id,
1124        ):
1125            result += coef * variable_values[var]
1126        for key, coef in sorted(
1127            self._quadratic_terms.items(),
1128            key=lambda quad_coef_pair: (
1129                quad_coef_pair[0].first_var.id,
1130                quad_coef_pair[0].second_var.id,
1131            ),
1132        ):
1133            result += (
1134                coef * variable_values[key.first_var] * variable_values[key.second_var]
1135            )
1136        return result
1137
1138    def _quadratic_flatten_once_and_add_to(
1139        self,
1140        scale: float,
1141        processed_elements: _QuadraticProcessedElements,
1142        target_stack: _QuadraticToProcessElements,
1143    ) -> None:
1144        for var, val in self._linear_terms.items():
1145            processed_elements.terms[var] += val * scale
1146        for key, val in self._quadratic_terms.items():
1147            processed_elements.quadratic_terms[key] += val * scale
1148        processed_elements.offset += scale * self.offset
1149
1150    # TODO(b/216492143): change __str__ to match C++ implementation in
1151    # cl/421649402.
1152    def __str__(self):
1153        result = str(self.offset)
1154        sorted_linear_keys = sorted(self._linear_terms.keys(), key=str)
1155        for variable in sorted_linear_keys:
1156            # TODO(b/216492143): consider how to better deal with `NaN` and try to
1157            # match C++ implementation in cl/421649402. See TODO for StrAndReprTest in
1158            # linear_expression_test.py.
1159            coefficient = self._linear_terms[variable]
1160            if coefficient == 0.0:
1161                continue
1162            if coefficient > 0:
1163                result += " + "
1164            else:
1165                result += " - "
1166            result += str(abs(coefficient)) + " * " + str(variable)
1167        sorted_quadratic_keys = sorted(self._quadratic_terms.keys(), key=str)
1168        for key in sorted_quadratic_keys:
1169            # TODO(b/216492143): consider how to better deal with `NaN` and try to
1170            # match C++ implementation in cl/421649402. See TODO for StrAndReprTest in
1171            # linear_expression_test.py.
1172            coefficient = self._quadratic_terms[key]
1173            if coefficient == 0.0:
1174                continue
1175            if coefficient > 0:
1176                result += " + "
1177            else:
1178                result += " - "
1179            result += str(abs(coefficient)) + " * " + str(key)
1180        return result
1181
1182    def __repr__(self):
1183        result = f"QuadraticExpression({self.offset}, " + "{"
1184        result += ", ".join(
1185            f"{variable!r}: {coefficient}"
1186            for variable, coefficient in self._linear_terms.items()
1187        )
1188        result += "}, {"
1189        result += ", ".join(
1190            f"{key!r}: {coefficient}"
1191            for key, coefficient in self._quadratic_terms.items()
1192        )
1193        result += "})"
1194        return result
1195
1196
1197class LinearConstraint:
1198    """A linear constraint for an optimization model.
1199
1200    A LinearConstraint adds the following restriction on feasible solutions to an
1201    optimization model:
1202      lb <= sum_{i in I} a_i * x_i <= ub
1203    where x_i are the decision variables of the problem. lb == ub is allowed, this
1204    models the equality constraint:
1205      sum_{i in I} a_i * x_i == b
1206    lb > ub is also allowed, but the optimization problem will be infeasible.
1207
1208    A LinearConstraint can be configured as follows:
1209      * lower_bound: a float property, lb above. Should not be NaN nor +inf.
1210      * upper_bound: a float property, ub above. Should not be NaN nor -inf.
1211      * set_coefficient() and get_coefficient(): get and set the a_i * x_i
1212        terms. The variable must be from the same model as this constraint, and
1213        the a_i must be finite and not NaN. The coefficient for any variable not
1214        set is 0.0, and setting a coefficient to 0.0 removes it from I above.
1215
1216    The name is optional, read only, and used only for debugging. Non-empty names
1217    should be distinct.
1218
1219    Every LinearConstraint is associated with a Model (defined below). Note that
1220    data describing the linear constraint (e.g. lower_bound) is owned by
1221    Model.storage, this class is simply a reference to that data. Do not create a
1222    LinearConstraint directly, use Model.add_linear_constraint() instead.
1223    """
1224
1225    __slots__ = "__weakref__", "_model", "_id"
1226
1227    def __init__(self, model: "Model", cid: int) -> None:
1228        """Do not invoke directly, use Model.add_linear_constraint()."""
1229        self._model: "Model" = model
1230        self._id: int = cid
1231
1232    @property
1233    def lower_bound(self) -> float:
1234        return self.model.storage.get_linear_constraint_lb(self._id)
1235
1236    @lower_bound.setter
1237    def lower_bound(self, value: float) -> None:
1238        self.model.storage.set_linear_constraint_lb(self._id, value)
1239
1240    @property
1241    def upper_bound(self) -> float:
1242        return self.model.storage.get_linear_constraint_ub(self._id)
1243
1244    @upper_bound.setter
1245    def upper_bound(self, value: float) -> None:
1246        self.model.storage.set_linear_constraint_ub(self._id, value)
1247
1248    @property
1249    def name(self) -> str:
1250        return self.model.storage.get_linear_constraint_name(self._id)
1251
1252    @property
1253    def id(self) -> int:
1254        return self._id
1255
1256    @property
1257    def model(self) -> "Model":
1258        return self._model
1259
1260    def set_coefficient(self, variable: Variable, coefficient: float) -> None:
1261        self.model.check_compatible(variable)
1262        self.model.storage.set_linear_constraint_coefficient(
1263            self._id, variable.id, coefficient
1264        )
1265
1266    def get_coefficient(self, variable: Variable) -> float:
1267        self.model.check_compatible(variable)
1268        return self.model.storage.get_linear_constraint_coefficient(
1269            self._id, variable.id
1270        )
1271
1272    def terms(self) -> Iterator[LinearTerm]:
1273        """Yields the variable/coefficient pairs with nonzero coefficient for this linear constraint."""
1274        for variable in self.model.row_nonzeros(self):
1275            yield LinearTerm(
1276                variable=variable,
1277                coefficient=self.model.storage.get_linear_constraint_coefficient(
1278                    self._id, variable.id
1279                ),
1280            )
1281
1282    def as_bounded_linear_expression(self) -> BoundedLinearExpression:
1283        """Returns the bounded expression from lower_bound, upper_bound and terms."""
1284        return BoundedLinearExpression(
1285            self.lower_bound, LinearSum(self.terms()), self.upper_bound
1286        )
1287
1288    def __str__(self):
1289        """Returns the name, or a string containing the id if the name is empty."""
1290        return self.name if self.name else f"linear_constraint_{self.id}"
1291
1292    def __repr__(self):
1293        return f"<LinearConstraint id: {self.id}, name: {self.name!r}>"
1294
1295
1296class Objective:
1297    """The objective for an optimization model.
1298
1299    An objective is either of the form:
1300      min o + sum_{i in I} c_i * x_i + sum_{i, j in I, i <= j} q_i,j * x_i * x_j
1301    or
1302      max o + sum_{i in I} c_i * x_i + sum_{(i, j) in Q} q_i,j * x_i * x_j
1303    where x_i are the decision variables of the problem and where all pairs (i, j)
1304    in Q satisfy i <= j. The values of o, c_i and q_i,j should be finite and not
1305    NaN.
1306
1307    The objective can be configured as follows:
1308      * offset: a float property, o above. Should be finite and not NaN.
1309      * is_maximize: a bool property, if the objective is to maximize or minimize.
1310      * set_linear_coefficient and get_linear_coefficient control the c_i * x_i
1311        terms. The variables must be from the same model as this objective, and
1312        the c_i must be finite and not NaN. The coefficient for any variable not
1313        set is 0.0, and setting a coefficient to 0.0 removes it from I above.
1314      * set_quadratic_coefficient and get_quadratic_coefficient control the
1315        q_i,j * x_i * x_j terms. The variables must be from the same model as this
1316        objective, and the q_i,j must be finite and not NaN. The coefficient for
1317        any pair of variables not set is 0.0, and setting a coefficient to 0.0
1318        removes the associated (i,j) from Q above.
1319
1320    Every Objective is associated with a Model (defined below). Note that data
1321    describing the objective (e.g. offset) is owned by Model.storage, this class
1322    is simply a reference to that data. Do not create an Objective directly,
1323    use Model.objective to access the objective instead.
1324
1325    The objective will be linear if only linear coefficients are set. This can be
1326    useful to avoid solve-time errors with solvers that do not accept quadratic
1327    objectives. To facilitate this linear objective guarantee we provide three
1328    functions to add to the objective:
1329      * add(), which accepts linear or quadratic expressions,
1330      * add_quadratic(), which also accepts linear or quadratic expressions and
1331        can be used to signal a quadratic objective is possible, and
1332      * add_linear(), which only accepts linear expressions and can be used to
1333        guarantee the objective remains linear.
1334    """
1335
1336    __slots__ = ("_model",)
1337
1338    def __init__(self, model: "Model") -> None:
1339        """Do no invoke directly, use Model.objective."""
1340        self._model: "Model" = model
1341
1342    @property
1343    def is_maximize(self) -> bool:
1344        return self.model.storage.get_is_maximize()
1345
1346    @is_maximize.setter
1347    def is_maximize(self, is_maximize: bool) -> None:
1348        self.model.storage.set_is_maximize(is_maximize)
1349
1350    @property
1351    def offset(self) -> float:
1352        return self.model.storage.get_objective_offset()
1353
1354    @offset.setter
1355    def offset(self, value: float) -> None:
1356        self.model.storage.set_objective_offset(value)
1357
1358    @property
1359    def model(self) -> "Model":
1360        return self._model
1361
1362    def set_linear_coefficient(self, variable: Variable, coef: float) -> None:
1363        self.model.check_compatible(variable)
1364        self.model.storage.set_linear_objective_coefficient(variable.id, coef)
1365
1366    def get_linear_coefficient(self, variable: Variable) -> float:
1367        self.model.check_compatible(variable)
1368        return self.model.storage.get_linear_objective_coefficient(variable.id)
1369
1370    def linear_terms(self) -> Iterator[LinearTerm]:
1371        """Yields variable coefficient pairs for variables with nonzero objective coefficient in undefined order."""
1372        yield from self.model.linear_objective_terms()
1373
1374    def add(self, objective: QuadraticTypes) -> None:
1375        """Adds the provided expression `objective` to the objective function.
1376
1377        To ensure the objective remains linear through type checks, use
1378        add_linear().
1379
1380        Args:
1381          objective: the expression to add to the objective function.
1382        """
1383        if isinstance(objective, (LinearBase, int, float)):
1384            self.add_linear(objective)
1385        elif isinstance(objective, QuadraticBase):
1386            self.add_quadratic(objective)
1387        else:
1388            raise TypeError(
1389                "unsupported type in objective argument for "
1390                f"Objective.add(): {type(objective).__name__!r}"
1391            )
1392
1393    def add_linear(self, objective: LinearTypes) -> None:
1394        """Adds the provided linear expression `objective` to the objective function."""
1395        if not isinstance(objective, (LinearBase, int, float)):
1396            raise TypeError(
1397                "unsupported type in objective argument for "
1398                f"Objective.add_linear(): {type(objective).__name__!r}"
1399            )
1400        objective_expr = as_flat_linear_expression(objective)
1401        self.offset += objective_expr.offset
1402        for var, coefficient in objective_expr.terms.items():
1403            self.set_linear_coefficient(
1404                var, self.get_linear_coefficient(var) + coefficient
1405            )
1406
1407    def add_quadratic(self, objective: QuadraticTypes) -> None:
1408        """Adds the provided quadratic expression `objective` to the objective function."""
1409        if not isinstance(objective, (QuadraticBase, LinearBase, int, float)):
1410            raise TypeError(
1411                "unsupported type in objective argument for "
1412                f"Objective.add(): {type(objective).__name__!r}"
1413            )
1414        objective_expr = as_flat_quadratic_expression(objective)
1415        self.offset += objective_expr.offset
1416        for var, coefficient in objective_expr.linear_terms.items():
1417            self.set_linear_coefficient(
1418                var, self.get_linear_coefficient(var) + coefficient
1419            )
1420        for key, coefficient in objective_expr.quadratic_terms.items():
1421            self.set_quadratic_coefficient(
1422                key.first_var,
1423                key.second_var,
1424                self.get_quadratic_coefficient(key.first_var, key.second_var)
1425                + coefficient,
1426            )
1427
1428    def set_quadratic_coefficient(
1429        self, first_variable: Variable, second_variable: Variable, coef: float
1430    ) -> None:
1431        self.model.check_compatible(first_variable)
1432        self.model.check_compatible(second_variable)
1433        self.model.storage.set_quadratic_objective_coefficient(
1434            first_variable.id, second_variable.id, coef
1435        )
1436
1437    def get_quadratic_coefficient(
1438        self, first_variable: Variable, second_variable: Variable
1439    ) -> float:
1440        self.model.check_compatible(first_variable)
1441        self.model.check_compatible(second_variable)
1442        return self.model.storage.get_quadratic_objective_coefficient(
1443            first_variable.id, second_variable.id
1444        )
1445
1446    def quadratic_terms(self) -> Iterator[QuadraticTerm]:
1447        """Yields quadratic terms with nonzero objective coefficient in undefined order."""
1448        yield from self.model.quadratic_objective_terms()
1449
1450    def as_linear_expression(self) -> LinearExpression:
1451        if any(self.quadratic_terms()):
1452            raise TypeError("Cannot get a quadratic objective as a linear expression")
1453        return as_flat_linear_expression(self.offset + LinearSum(self.linear_terms()))
1454
1455    def as_quadratic_expression(self) -> QuadraticExpression:
1456        return as_flat_quadratic_expression(
1457            self.offset
1458            + LinearSum(self.linear_terms())
1459            + QuadraticSum(self.quadratic_terms())
1460        )
1461
1462    def clear(self) -> None:
1463        """Clears objective coefficients and offset. Does not change direction."""
1464        self.model.storage.clear_objective()
1465
1466
1467class LinearConstraintMatrixEntry(NamedTuple):
1468    linear_constraint: LinearConstraint
1469    variable: Variable
1470    coefficient: float
1471
1472
1473class UpdateTracker:
1474    """Tracks updates to an optimization model from a ModelStorage.
1475
1476    Do not instantiate directly, instead create through
1477    ModelStorage.add_update_tracker().
1478
1479    Querying an UpdateTracker after calling Model.remove_update_tracker will
1480    result in a model_storage.UsedUpdateTrackerAfterRemovalError.
1481
1482    Example:
1483      mod = Model()
1484      x = mod.add_variable(0.0, 1.0, True, 'x')
1485      y = mod.add_variable(0.0, 1.0, True, 'y')
1486      tracker = mod.add_update_tracker()
1487      mod.set_variable_ub(x, 3.0)
1488      tracker.export_update()
1489        => "variable_updates: {upper_bounds: {ids: [0], values[3.0] }"
1490      mod.set_variable_ub(y, 2.0)
1491      tracker.export_update()
1492        => "variable_updates: {upper_bounds: {ids: [0, 1], values[3.0, 2.0] }"
1493      tracker.advance_checkpoint()
1494      tracker.export_update()
1495        => None
1496      mod.set_variable_ub(y, 4.0)
1497      tracker.export_update()
1498        => "variable_updates: {upper_bounds: {ids: [1], values[4.0] }"
1499      tracker.advance_checkpoint()
1500      mod.remove_update_tracker(tracker)
1501    """
1502
1503    def __init__(self, storage_update_tracker: model_storage.StorageUpdateTracker):
1504        """Do not invoke directly, use Model.add_update_tracker() instead."""
1505        self.storage_update_tracker = storage_update_tracker
1506
1507    def export_update(self) -> Optional[model_update_pb2.ModelUpdateProto]:
1508        """Returns changes to the model since last call to checkpoint/creation."""
1509        return self.storage_update_tracker.export_update()
1510
1511    def advance_checkpoint(self) -> None:
1512        """Track changes to the model only after this function call."""
1513        return self.storage_update_tracker.advance_checkpoint()
1514
1515
1516class Model:
1517    """An optimization model.
1518
1519    The objective function of the model can be linear or quadratic, and some
1520    solvers can only handle linear objective functions. For this reason Model has
1521    three versions of all objective setting functions:
1522      * A generic one (e.g. maximize()), which accepts linear or quadratic
1523        expressions,
1524      * a quadratic version (e.g. maximize_quadratic_objective()), which also
1525        accepts linear or quadratic expressions and can be used to signal a
1526        quadratic objective is possible, and
1527      * a linear version (e.g. maximize_linear_objective()), which only accepts
1528        linear expressions and can be used to avoid solve time errors for solvers
1529        that do not accept quadratic objectives.
1530
1531    Attributes:
1532      name: A description of the problem, can be empty.
1533      objective: A function to maximize or minimize.
1534      storage: Implementation detail, do not access directly.
1535      _variable_ids: Maps variable ids to Variable objects.
1536      _linear_constraint_ids: Maps linear constraint ids to LinearConstraint
1537        objects.
1538    """
1539
1540    __slots__ = "storage", "_variable_ids", "_linear_constraint_ids", "_objective"
1541
1542    def __init__(
1543        self,
1544        *,
1545        name: str = "",
1546        storage_class: StorageClass = hash_model_storage.HashModelStorage,
1547    ) -> None:
1548        self.storage: Storage = storage_class(name)
1549        # Weak references are used to eliminate reference cycles (so that Model will
1550        # be destroyed when we reach zero references, instead of relying on GC cycle
1551        # detection). Do not access a variable or constraint from these maps
1552        # directly, as they may no longer exist, use _get_or_make_variable() and
1553        # _get_or_make_linear_constraint() defined below instead.
1554        self._variable_ids: weakref.WeakValueDictionary[int, Variable] = (
1555            weakref.WeakValueDictionary()
1556        )
1557        self._linear_constraint_ids: weakref.WeakValueDictionary[
1558            int, LinearConstraint
1559        ] = weakref.WeakValueDictionary()
1560        self._objective: Objective = Objective(self)
1561
1562    @property
1563    def name(self) -> str:
1564        return self.storage.name
1565
1566    @property
1567    def objective(self) -> Objective:
1568        return self._objective
1569
1570    def add_variable(
1571        self,
1572        *,
1573        lb: float = -math.inf,
1574        ub: float = math.inf,
1575        is_integer: bool = False,
1576        name: str = "",
1577    ) -> Variable:
1578        """Adds a decision variable to the optimization model.
1579
1580        Args:
1581          lb: The new variable must take at least this value (a lower bound).
1582          ub: The new variable must be at most this value (an upper bound).
1583          is_integer: Indicates if the variable can only take integer values
1584            (otherwise, the variable can take any continuous value).
1585          name: For debugging purposes only, but nonempty names must be distinct.
1586
1587        Returns:
1588          A reference to the new decision variable.
1589        """
1590        variable_id = self.storage.add_variable(lb, ub, is_integer, name)
1591        result = Variable(self, variable_id)
1592        self._variable_ids[variable_id] = result
1593        return result
1594
1595    def add_integer_variable(
1596        self, *, lb: float = -math.inf, ub: float = math.inf, name: str = ""
1597    ) -> Variable:
1598        return self.add_variable(lb=lb, ub=ub, is_integer=True, name=name)
1599
1600    def add_binary_variable(self, *, name: str = "") -> Variable:
1601        return self.add_variable(lb=0.0, ub=1.0, is_integer=True, name=name)
1602
1603    def get_variable(self, var_id: int) -> Variable:
1604        """Returns the Variable for the id var_id, or raises KeyError."""
1605        if not self.storage.variable_exists(var_id):
1606            raise KeyError(f"variable does not exist with id {var_id}")
1607        return self._get_or_make_variable(var_id)
1608
1609    def delete_variable(self, variable: Variable) -> None:
1610        self.check_compatible(variable)
1611        self.storage.delete_variable(variable.id)
1612        del self._variable_ids[variable.id]
1613
1614    def variables(self) -> Iterator[Variable]:
1615        """Yields the variables in the order of creation."""
1616        for var_id in self.storage.get_variables():
1617            yield self._get_or_make_variable(var_id)
1618
1619    def maximize(self, objective: QuadraticTypes) -> None:
1620        """Sets the objective to maximize the provided expression `objective`."""
1621        self.set_objective(objective, is_maximize=True)
1622
1623    def maximize_linear_objective(self, objective: LinearTypes) -> None:
1624        """Sets the objective to maximize the provided linear expression `objective`."""
1625        self.set_linear_objective(objective, is_maximize=True)
1626
1627    def maximize_quadratic_objective(self, objective: QuadraticTypes) -> None:
1628        """Sets the objective to maximize the provided quadratic expression `objective`."""
1629        self.set_quadratic_objective(objective, is_maximize=True)
1630
1631    def minimize(self, objective: QuadraticTypes) -> None:
1632        """Sets the objective to minimize the provided expression `objective`."""
1633        self.set_objective(objective, is_maximize=False)
1634
1635    def minimize_linear_objective(self, objective: LinearTypes) -> None:
1636        """Sets the objective to minimize the provided linear expression `objective`."""
1637        self.set_linear_objective(objective, is_maximize=False)
1638
1639    def minimize_quadratic_objective(self, objective: QuadraticTypes) -> None:
1640        """Sets the objective to minimize the provided quadratic expression `objective`."""
1641        self.set_quadratic_objective(objective, is_maximize=False)
1642
1643    def set_objective(self, objective: QuadraticTypes, *, is_maximize: bool) -> None:
1644        """Sets the objective to optimize the provided expression `objective`."""
1645        if isinstance(objective, (LinearBase, int, float)):
1646            self.set_linear_objective(objective, is_maximize=is_maximize)
1647        elif isinstance(objective, QuadraticBase):
1648            self.set_quadratic_objective(objective, is_maximize=is_maximize)
1649        else:
1650            raise TypeError(
1651                "unsupported type in objective argument for "
1652                f"set_objective: {type(objective).__name__!r}"
1653            )
1654
1655    def set_linear_objective(
1656        self, objective: LinearTypes, *, is_maximize: bool
1657    ) -> None:
1658        """Sets the objective to optimize the provided linear expression `objective`."""
1659        if not isinstance(objective, (LinearBase, int, float)):
1660            raise TypeError(
1661                "unsupported type in objective argument for "
1662                f"set_linear_objective: {type(objective).__name__!r}"
1663            )
1664        self.storage.clear_objective()
1665        self._objective.is_maximize = is_maximize
1666        objective_expr = as_flat_linear_expression(objective)
1667        self._objective.offset = objective_expr.offset
1668        for var, coefficient in objective_expr.terms.items():
1669            self._objective.set_linear_coefficient(var, coefficient)
1670
1671    def set_quadratic_objective(
1672        self, objective: QuadraticTypes, *, is_maximize: bool
1673    ) -> None:
1674        """Sets the objective to optimize the provided linear expression `objective`."""
1675        if not isinstance(objective, (QuadraticBase, LinearBase, int, float)):
1676            raise TypeError(
1677                "unsupported type in objective argument for "
1678                f"set_quadratic_objective: {type(objective).__name__!r}"
1679            )
1680        self.storage.clear_objective()
1681        self._objective.is_maximize = is_maximize
1682        objective_expr = as_flat_quadratic_expression(objective)
1683        self._objective.offset = objective_expr.offset
1684        for var, coefficient in objective_expr.linear_terms.items():
1685            self._objective.set_linear_coefficient(var, coefficient)
1686        for key, coefficient in objective_expr.quadratic_terms.items():
1687            self._objective.set_quadratic_coefficient(
1688                key.first_var, key.second_var, coefficient
1689            )
1690
1691    # TODO(b/227214976): Update the note below and link to pytype bug number.
1692    # Note: bounded_expr's type includes bool only as a workaround to a pytype
1693    # issue. Passing a bool for bounded_expr will raise an error in runtime.
1694    def add_linear_constraint(
1695        self,
1696        bounded_expr: Optional[Union[bool, BoundedLinearTypes]] = None,
1697        *,
1698        lb: Optional[float] = None,
1699        ub: Optional[float] = None,
1700        expr: Optional[LinearTypes] = None,
1701        name: str = "",
1702    ) -> LinearConstraint:
1703        """Adds a linear constraint to the optimization model.
1704
1705        The simplest way to specify the constraint is by passing a one-sided or
1706        two-sided linear inequality as in:
1707          * add_linear_constraint(x + y + 1.0 <= 2.0),
1708          * add_linear_constraint(x + y >= 2.0), or
1709          * add_linear_constraint((1.0 <= x + y) <= 2.0).
1710
1711        Note the extra parenthesis for two-sided linear inequalities, which is
1712        required due to some language limitations (see
1713        https://peps.python.org/pep-0335/ and https://peps.python.org/pep-0535/).
1714        If the parenthesis are omitted, a TypeError will be raised explaining the
1715        issue (if this error was not raised the first inequality would have been
1716        silently ignored because of the noted language limitations).
1717
1718        The second way to specify the constraint is by setting lb, ub, and/o expr as
1719        in:
1720          * add_linear_constraint(expr=x + y + 1.0, ub=2.0),
1721          * add_linear_constraint(expr=x + y, lb=2.0),
1722          * add_linear_constraint(expr=x + y, lb=1.0, ub=2.0), or
1723          * add_linear_constraint(lb=1.0).
1724        Omitting lb is equivalent to setting it to -math.inf and omiting ub is
1725        equivalent to setting it to math.inf.
1726
1727        These two alternatives are exclusive and a combined call like:
1728          * add_linear_constraint(x + y <= 2.0, lb=1.0), or
1729          * add_linear_constraint(x + y <= 2.0, ub=math.inf)
1730        will raise a ValueError. A ValueError is also raised if expr's offset is
1731        infinite.
1732
1733        Args:
1734          bounded_expr: a linear inequality describing the constraint. Cannot be
1735            specified together with lb, ub, or expr.
1736          lb: The constraint's lower bound if bounded_expr is omitted (if both
1737            bounder_expr and lb are omitted, the lower bound is -math.inf).
1738          ub: The constraint's upper bound if bounded_expr is omitted (if both
1739            bounder_expr and ub are omitted, the upper bound is math.inf).
1740          expr: The constraint's linear expression if bounded_expr is omitted.
1741          name: For debugging purposes only, but nonempty names must be distinct.
1742
1743        Returns:
1744          A reference to the new linear constraint.
1745        """
1746        normalized_inequality = as_normalized_linear_inequality(
1747            bounded_expr, lb=lb, ub=ub, expr=expr
1748        )
1749        lin_con_id = self.storage.add_linear_constraint(
1750            normalized_inequality.lb, normalized_inequality.ub, name
1751        )
1752        result = LinearConstraint(self, lin_con_id)
1753        self._linear_constraint_ids[lin_con_id] = result
1754        for var, coefficient in normalized_inequality.coefficients.items():
1755            result.set_coefficient(var, coefficient)
1756        return result
1757
1758    def get_linear_constraint(self, lin_con_id: int) -> LinearConstraint:
1759        """Returns the LinearConstraint for the id lin_con_id."""
1760        if not self.storage.linear_constraint_exists(lin_con_id):
1761            raise KeyError(f"linear constraint does not exist with id {lin_con_id}")
1762        return self._get_or_make_linear_constraint(lin_con_id)
1763
1764    def delete_linear_constraint(self, linear_constraint: LinearConstraint) -> None:
1765        self.check_compatible(linear_constraint)
1766        self.storage.delete_linear_constraint(linear_constraint.id)
1767        del self._linear_constraint_ids[linear_constraint.id]
1768
1769    def linear_constraints(self) -> Iterator[LinearConstraint]:
1770        """Yields the linear constraints in the order of creation."""
1771        for lin_con_id in self.storage.get_linear_constraints():
1772            yield self._get_or_make_linear_constraint(lin_con_id)
1773
1774    def row_nonzeros(self, linear_constraint: LinearConstraint) -> Iterator[Variable]:
1775        """Yields the variables with nonzero coefficient for this linear constraint."""
1776        for var_id in self.storage.get_variables_for_linear_constraint(
1777            linear_constraint.id
1778        ):
1779            yield self._get_or_make_variable(var_id)
1780
1781    def column_nonzeros(self, variable: Variable) -> Iterator[LinearConstraint]:
1782        """Yields the linear constraints with nonzero coefficient for this variable."""
1783        for lin_con_id in self.storage.get_linear_constraints_with_variable(
1784            variable.id
1785        ):
1786            yield self._get_or_make_linear_constraint(lin_con_id)
1787
1788    def linear_objective_terms(self) -> Iterator[LinearTerm]:
1789        """Yields variable coefficient pairs for variables with nonzero objective coefficient in undefined order."""
1790        for term in self.storage.get_linear_objective_coefficients():
1791            yield LinearTerm(
1792                variable=self._get_or_make_variable(term.variable_id),
1793                coefficient=term.coefficient,
1794            )
1795
1796    def quadratic_objective_terms(self) -> Iterator[QuadraticTerm]:
1797        """Yields the quadratic terms with nonzero objective coefficient in undefined order."""
1798        for term in self.storage.get_quadratic_objective_coefficients():
1799            var1 = self._get_or_make_variable(term.id_key.id1)
1800            var2 = self._get_or_make_variable(term.id_key.id2)
1801            yield QuadraticTerm(
1802                key=QuadraticTermKey(var1, var2), coefficient=term.coefficient
1803            )
1804
1805    def linear_constraint_matrix_entries(
1806        self,
1807    ) -> Iterator[LinearConstraintMatrixEntry]:
1808        """Yields the nonzero elements of the linear constraint matrix in undefined order."""
1809        for entry in self.storage.get_linear_constraint_matrix_entries():
1810            yield LinearConstraintMatrixEntry(
1811                linear_constraint=self._get_or_make_linear_constraint(
1812                    entry.linear_constraint_id
1813                ),
1814                variable=self._get_or_make_variable(entry.variable_id),
1815                coefficient=entry.coefficient,
1816            )
1817
1818    def export_model(self) -> model_pb2.ModelProto:
1819        return self.storage.export_model()
1820
1821    def add_update_tracker(self) -> UpdateTracker:
1822        """Creates an UpdateTracker registered on this model to view changes."""
1823        return UpdateTracker(self.storage.add_update_tracker())
1824
1825    def remove_update_tracker(self, tracker: UpdateTracker):
1826        """Stops tracker from getting updates on changes to this model.
1827
1828        An error will be raised if tracker was not created by this Model or if
1829        tracker has been previously removed.
1830
1831        Using (via checkpoint or update) an UpdateTracker after it has been removed
1832        will result in an error.
1833
1834        Args:
1835          tracker: The UpdateTracker to unregister.
1836
1837        Raises:
1838          KeyError: The tracker was created by another model or was already removed.
1839        """
1840        self.storage.remove_update_tracker(tracker.storage_update_tracker)
1841
1842    def check_compatible(
1843        self, var_or_constraint: Union[Variable, LinearConstraint]
1844    ) -> None:
1845        """Raises a ValueError if the model of var_or_constraint is not self."""
1846        if var_or_constraint.model is not self:
1847            raise ValueError(
1848                f"{var_or_constraint} is from model {var_or_constraint.model} and"
1849                f" cannot be used with model {self}"
1850            )
1851
1852    def _get_or_make_variable(self, variable_id: int) -> Variable:
1853        result = self._variable_ids.get(variable_id)
1854        if result:
1855            return result
1856        result = Variable(self, variable_id)
1857        self._variable_ids[variable_id] = result
1858        return result
1859
1860    def _get_or_make_linear_constraint(
1861        self, linear_constraint_id: int
1862    ) -> LinearConstraint:
1863        result = self._linear_constraint_ids.get(linear_constraint_id)
1864        if result:
1865            return result
1866        result = LinearConstraint(self, linear_constraint_id)
1867        self._linear_constraint_ids[linear_constraint_id] = result
1868        return result
1869
1870
1871class LinearSum(LinearBase):
1872    # TODO(b/216492143): consider what details to move elsewhere and/or replace
1873    # by examples, and do complexity analysis.
1874    """A deferred sum of LinearBase objects.
1875
1876    LinearSum objects are automatically created when two linear objects are added
1877    and, as noted in the documentation for Linear, can reduce the inefficiencies.
1878    In particular, they are created when calling sum(iterable) when iterable is
1879    an Iterable[LinearTypes]. However, using LinearSum(iterable) instead
1880    can result in additional performance improvements:
1881
1882      * sum(iterable): creates a nested set of LinearSum objects (e.g.
1883        `sum([a, b, c])` is `LinearSum(0, LinearSum(a, LinearSum(b, c)))`).
1884      * LinearSum(iterable): creates a single LinearSum that saves a tuple with
1885        all the LinearTypes objects in iterable (e.g.
1886        `LinearSum([a, b, c])` does not create additional objects).
1887
1888    This class is immutable.
1889    """
1890
1891    __slots__ = "__weakref__", "_elements"
1892
1893    # Potentially unsafe use of Iterable argument is handled by immediate local
1894    # storage as tuple.
1895    def __init__(self, iterable: Iterable[LinearTypes]) -> None:
1896        """Creates a LinearSum object. A copy of iterable is saved as a tuple."""
1897
1898        self._elements = tuple(iterable)
1899        for item in self._elements:
1900            if not isinstance(item, (LinearBase, int, float)):
1901                raise TypeError(
1902                    "unsupported type in iterable argument for "
1903                    f"LinearSum: {type(item).__name__!r}"
1904                )
1905
1906    @property
1907    def elements(self) -> Tuple[LinearTypes, ...]:
1908        return self._elements
1909
1910    def _flatten_once_and_add_to(
1911        self,
1912        scale: float,
1913        processed_elements: _ProcessedElements,
1914        target_stack: _ToProcessElements,
1915    ) -> None:
1916        for term in self._elements:
1917            if isinstance(term, (int, float)):
1918                processed_elements.offset += scale * float(term)
1919            else:
1920                target_stack.append(term, scale)
1921
1922    def __str__(self):
1923        return str(as_flat_linear_expression(self))
1924
1925    def __repr__(self):
1926        result = "LinearSum(("
1927        result += ", ".join(repr(linear) for linear in self._elements)
1928        result += "))"
1929        return result
1930
1931
1932class QuadraticSum(QuadraticBase):
1933    # TODO(b/216492143): consider what details to move elsewhere and/or replace
1934    # by examples, and do complexity analysis.
1935    """A deferred sum of QuadraticTypes objects.
1936
1937    QuadraticSum objects are automatically created when a quadratic object is
1938    added to quadratic or linear objects and, as has performance optimizations
1939    similar to LinearSum.
1940
1941    This class is immutable.
1942    """
1943
1944    __slots__ = "__weakref__", "_elements"
1945
1946    # Potentially unsafe use of Iterable argument is handled by immediate local
1947    # storage as tuple.
1948    def __init__(self, iterable: Iterable[QuadraticTypes]) -> None:
1949        """Creates a QuadraticSum object. A copy of iterable is saved as a tuple."""
1950
1951        self._elements = tuple(iterable)
1952        for item in self._elements:
1953            if not isinstance(item, (LinearBase, QuadraticBase, int, float)):
1954                raise TypeError(
1955                    "unsupported type in iterable argument for "
1956                    f"QuadraticSum: {type(item).__name__!r}"
1957                )
1958
1959    @property
1960    def elements(self) -> Tuple[QuadraticTypes, ...]:
1961        return self._elements
1962
1963    def _quadratic_flatten_once_and_add_to(
1964        self,
1965        scale: float,
1966        processed_elements: _QuadraticProcessedElements,
1967        target_stack: _QuadraticToProcessElements,
1968    ) -> None:
1969        for term in self._elements:
1970            if isinstance(term, (int, float)):
1971                processed_elements.offset += scale * float(term)
1972            else:
1973                target_stack.append(term, scale)
1974
1975    def __str__(self):
1976        return str(as_flat_quadratic_expression(self))
1977
1978    def __repr__(self):
1979        result = "QuadraticSum(("
1980        result += ", ".join(repr(element) for element in self._elements)
1981        result += "))"
1982        return result
1983
1984
1985class LinearProduct(LinearBase):
1986    """A deferred multiplication computation for linear expressions.
1987
1988    This class is immutable.
1989    """
1990
1991    __slots__ = "_scalar", "_linear"
1992
1993    def __init__(self, scalar: float, linear: LinearBase) -> None:
1994        if not isinstance(scalar, (float, int)):
1995            raise TypeError(
1996                "unsupported type for scalar argument in "
1997                f"LinearProduct: {type(scalar).__name__!r}"
1998            )
1999        if not isinstance(linear, LinearBase):
2000            raise TypeError(
2001                "unsupported type for linear argument in "
2002                f"LinearProduct: {type(linear).__name__!r}"
2003            )
2004        self._scalar: float = float(scalar)
2005        self._linear: LinearBase = linear
2006
2007    @property
2008    def scalar(self) -> float:
2009        return self._scalar
2010
2011    @property
2012    def linear(self) -> LinearBase:
2013        return self._linear
2014
2015    def _flatten_once_and_add_to(
2016        self,
2017        scale: float,
2018        processed_elements: _ProcessedElements,
2019        target_stack: _ToProcessElements,
2020    ) -> None:
2021        target_stack.append(self._linear, self._scalar * scale)
2022
2023    def __str__(self):
2024        return str(as_flat_linear_expression(self))
2025
2026    def __repr__(self):
2027        result = f"LinearProduct({self._scalar!r}, "
2028        result += f"{self._linear!r})"
2029        return result
2030
2031
2032class QuadraticProduct(QuadraticBase):
2033    """A deferred multiplication computation for quadratic expressions.
2034
2035    This class is immutable.
2036    """
2037
2038    __slots__ = "_scalar", "_quadratic"
2039
2040    def __init__(self, scalar: float, quadratic: QuadraticBase) -> None:
2041        if not isinstance(scalar, (float, int)):
2042            raise TypeError(
2043                "unsupported type for scalar argument in "
2044                f"QuadraticProduct: {type(scalar).__name__!r}"
2045            )
2046        if not isinstance(quadratic, QuadraticBase):
2047            raise TypeError(
2048                "unsupported type for linear argument in "
2049                f"QuadraticProduct: {type(quadratic).__name__!r}"
2050            )
2051        self._scalar: float = float(scalar)
2052        self._quadratic: QuadraticBase = quadratic
2053
2054    @property
2055    def scalar(self) -> float:
2056        return self._scalar
2057
2058    @property
2059    def quadratic(self) -> QuadraticBase:
2060        return self._quadratic
2061
2062    def _quadratic_flatten_once_and_add_to(
2063        self,
2064        scale: float,
2065        processed_elements: _QuadraticProcessedElements,
2066        target_stack: _QuadraticToProcessElements,
2067    ) -> None:
2068        target_stack.append(self._quadratic, self._scalar * scale)
2069
2070    def __str__(self):
2071        return str(as_flat_quadratic_expression(self))
2072
2073    def __repr__(self):
2074        return f"QuadraticProduct({self._scalar}, {self._quadratic!r})"
2075
2076
2077class LinearLinearProduct(QuadraticBase):
2078    """A deferred multiplication of two linear expressions.
2079
2080    This class is immutable.
2081    """
2082
2083    __slots__ = "_first_linear", "_second_linear"
2084
2085    def __init__(self, first_linear: LinearBase, second_linear: LinearBase) -> None:
2086        if not isinstance(first_linear, LinearBase):
2087            raise TypeError(
2088                "unsupported type for first_linear argument in "
2089                f"LinearLinearProduct: {type(first_linear).__name__!r}"
2090            )
2091        if not isinstance(second_linear, LinearBase):
2092            raise TypeError(
2093                "unsupported type for second_linear argument in "
2094                f"LinearLinearProduct: {type(second_linear).__name__!r}"
2095            )
2096        self._first_linear: LinearBase = first_linear
2097        self._second_linear: LinearBase = second_linear
2098
2099    @property
2100    def first_linear(self) -> LinearBase:
2101        return self._first_linear
2102
2103    @property
2104    def second_linear(self) -> LinearBase:
2105        return self._second_linear
2106
2107    def _quadratic_flatten_once_and_add_to(
2108        self,
2109        scale: float,
2110        processed_elements: _QuadraticProcessedElements,
2111        target_stack: _QuadraticToProcessElements,
2112    ) -> None:
2113        # A recursion is avoided here because as_flat_linear_expression() must never
2114        # call _quadratic_flatten_once_and_add_to().
2115        first_expression = as_flat_linear_expression(self._first_linear)
2116        second_expression = as_flat_linear_expression(self._second_linear)
2117        processed_elements.offset += (
2118            first_expression.offset * second_expression.offset * scale
2119        )
2120        for first_var, first_val in first_expression.terms.items():
2121            processed_elements.terms[first_var] += (
2122                second_expression.offset * first_val * scale
2123            )
2124        for second_var, second_val in second_expression.terms.items():
2125            processed_elements.terms[second_var] += (
2126                first_expression.offset * second_val * scale
2127            )
2128
2129        for first_var, first_val in first_expression.terms.items():
2130            for second_var, second_val in second_expression.terms.items():
2131                processed_elements.quadratic_terms[
2132                    QuadraticTermKey(first_var, second_var)
2133                ] += (first_val * second_val * scale)
2134
2135    def __str__(self):
2136        return str(as_flat_quadratic_expression(self))
2137
2138    def __repr__(self):
2139        result = "LinearLinearProduct("
2140        result += f"{self._first_linear!r}, "
2141        result += f"{self._second_linear!r})"
2142        return result
2143
2144
2145def as_flat_linear_expression(value: LinearTypes) -> LinearExpression:
2146    """Converts floats, ints and Linear objects to a LinearExpression."""
2147    if isinstance(value, LinearExpression):
2148        return value
2149    return LinearExpression(value)
2150
2151
2152def as_flat_quadratic_expression(value: QuadraticTypes) -> QuadraticExpression:
2153    """Converts floats, ints, LinearBase and QuadraticBase objects to a QuadraticExpression."""
2154    if isinstance(value, QuadraticExpression):
2155        return value
2156    return QuadraticExpression(value)
2157
2158
2159@dataclasses.dataclass
2160class NormalizedLinearInequality:
2161    """Represents an inequality lb <= expr <= ub where expr's offset is zero.
2162
2163    The inequality is of the form:
2164        lb <= sum_{x in V} coefficients[x] * x <= ub
2165    where V is the set of keys of coefficients.
2166    """
2167
2168    lb: float
2169    ub: float
2170    coefficients: Mapping[Variable, float]
2171
2172    def __init__(self, *, lb: float, ub: float, expr: LinearTypes) -> None:
2173        """Raises a ValueError if expr's offset is infinite."""
2174        flat_expr = as_flat_linear_expression(expr)
2175        if math.isinf(flat_expr.offset):
2176            raise ValueError(
2177                "Trying to create a linear constraint whose expression has an"
2178                " infinite offset."
2179            )
2180        self.lb = lb - flat_expr.offset
2181        self.ub = ub - flat_expr.offset
2182        self.coefficients = flat_expr.terms
2183
2184
2185# TODO(b/227214976): Update the note below and link to pytype bug number.
2186# Note: bounded_expr's type includes bool only as a workaround to a pytype
2187# issue. Passing a bool for bounded_expr will raise an error in runtime.
2188def as_normalized_linear_inequality(
2189    bounded_expr: Optional[Union[bool, BoundedLinearTypes]] = None,
2190    *,
2191    lb: Optional[float] = None,
2192    ub: Optional[float] = None,
2193    expr: Optional[LinearTypes] = None,
2194) -> NormalizedLinearInequality:
2195    """Converts a linear constraint to a NormalizedLinearInequality.
2196
2197    The simplest way to specify the constraint is by passing a one-sided or
2198    two-sided linear inequality as in:
2199      * as_normalized_linear_inequality(x + y + 1.0 <= 2.0),
2200      * as_normalized_linear_inequality(x + y >= 2.0), or
2201      * as_normalized_linear_inequality((1.0 <= x + y) <= 2.0).
2202
2203    Note the extra parenthesis for two-sided linear inequalities, which is
2204    required due to some language limitations (see
2205    https://peps.python.org/pep-0335/ and https://peps.python.org/pep-0535/).
2206    If the parenthesis are omitted, a TypeError will be raised explaining the
2207    issue (if this error was not raised the first inequality would have been
2208    silently ignored because of the noted language limitations).
2209
2210    The second way to specify the constraint is by setting lb, ub, and/o expr as
2211    in:
2212      * as_normalized_linear_inequality(expr=x + y + 1.0, ub=2.0),
2213      * as_normalized_linear_inequality(expr=x + y, lb=2.0),
2214      * as_normalized_linear_inequality(expr=x + y, lb=1.0, ub=2.0), or
2215      * as_normalized_linear_inequality(lb=1.0).
2216    Omitting lb is equivalent to setting it to -math.inf and omiting ub is
2217    equivalent to setting it to math.inf.
2218
2219    These two alternatives are exclusive and a combined call like:
2220      * as_normalized_linear_inequality(x + y <= 2.0, lb=1.0), or
2221      * as_normalized_linear_inequality(x + y <= 2.0, ub=math.inf)
2222    will raise a ValueError. A ValueError is also raised if expr's offset is
2223    infinite.
2224
2225    Args:
2226      bounded_expr: a linear inequality describing the constraint. Cannot be
2227        specified together with lb, ub, or expr.
2228      lb: The constraint's lower bound if bounded_expr is omitted (if both
2229        bounder_expr and lb are omitted, the lower bound is -math.inf).
2230      ub: The constraint's upper bound if bounded_expr is omitted (if both
2231        bounder_expr and ub are omitted, the upper bound is math.inf).
2232      expr: The constraint's linear expression if bounded_expr is omitted.
2233
2234    Returns:
2235      A NormalizedLinearInequality representing the linear constraint.
2236    """
2237    if bounded_expr is None:
2238        if lb is None:
2239            lb = -math.inf
2240        if ub is None:
2241            ub = math.inf
2242        if expr is None:
2243            return NormalizedLinearInequality(lb=lb, ub=ub, expr=0)
2244        if not isinstance(expr, (LinearBase, int, float)):
2245            raise TypeError(
2246                f"unsupported type for expr argument: {type(expr).__name__!r}"
2247            )
2248        return NormalizedLinearInequality(lb=lb, ub=ub, expr=expr)
2249
2250    if isinstance(bounded_expr, bool):
2251        raise TypeError(
2252            "unsupported type for bounded_expr argument:"
2253            " bool. This error can occur when trying to add != constraints "
2254            "(which are not supported) or inequalities/equalities with constant "
2255            "left-hand-side and right-hand-side (which are redundant or make a "
2256            "model infeasible)."
2257        )
2258    if not isinstance(
2259        bounded_expr,
2260        (
2261            LowerBoundedLinearExpression,
2262            UpperBoundedLinearExpression,
2263            BoundedLinearExpression,
2264            VarEqVar,
2265        ),
2266    ):
2267        raise TypeError(
2268            f"unsupported type for bounded_expr: {type(bounded_expr).__name__!r}"
2269        )
2270    if lb is not None:
2271        raise ValueError("lb cannot be specified together with a linear inequality")
2272    if ub is not None:
2273        raise ValueError("ub cannot be specified together with a linear inequality")
2274    if expr is not None:
2275        raise ValueError("expr cannot be specified together with a linear inequality")
2276
2277    if isinstance(bounded_expr, VarEqVar):
2278        return NormalizedLinearInequality(
2279            lb=0.0,
2280            ub=0.0,
2281            expr=bounded_expr.first_variable - bounded_expr.second_variable,
2282        )
2283
2284    if isinstance(
2285        bounded_expr, (LowerBoundedLinearExpression, BoundedLinearExpression)
2286    ):
2287        lb = bounded_expr.lower_bound
2288    else:
2289        lb = -math.inf
2290    if isinstance(
2291        bounded_expr, (UpperBoundedLinearExpression, BoundedLinearExpression)
2292    ):
2293        ub = bounded_expr.upper_bound
2294    else:
2295        ub = math.inf
2296    return NormalizedLinearInequality(lb=lb, ub=ub, expr=bounded_expr.expression)