linpy/linexprs.py

   1 # Copyright 2014 MINES ParisTech
   2 #
   3 # This file is part of LinPy.
   4 #
   5 # LinPy is free software: you can redistribute it and/or modify
   6 # it under the terms of the GNU General Public License as published by
   7 # the Free Software Foundation, either version 3 of the License, or
   8 # (at your option) any later version.
   9 #
  10 # LinPy is distributed in the hope that it will be useful,
  11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 # GNU General Public License for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License
  16 # along with LinPy.  If not, see <http://www.gnu.org/licenses/>.
  17
  18 import ast
  19 import functools
  20 import numbers
  21 import re
  22
  23 from collections import OrderedDict, defaultdict, Mapping
  24 from fractions import Fraction, gcd
  25
  26
  27 __all__ = [
  28     'LinExpr',
  29     'Symbol', 'Dummy', 'symbols',
  30     'Rational',
  31 ]
  32
  33
  34 def _polymorphic(func):
  35     @functools.wraps(func)
  36     def wrapper(left, right):
  37         if isinstance(right, LinExpr):
  38             return func(left, right)
  39         elif isinstance(right, numbers.Rational):
  40             right = Rational(right)
  41             return func(left, right)
  42         return NotImplemented
  43     return wrapper
  44
  45
  46 class LinExpr:
  47     """
  48     A linear expression consists of a list of coefficient-variable pairs
  49     that capture the linear terms, plus a constant term. Linear expressions
  50     are used to build constraints. They are temporary objects that typically
  51     have short lifespans.
  52
  53     Linear expressions are generally built using overloaded operators. For
  54     example, if x is a Symbol, then x + 1 is an instance of LinExpr.
  55
  56     LinExpr instances are hashable, and should be treated as immutable.
  57     """
  58
  59     def __new__(cls, coefficients=None, constant=0):
  60         """
  61         Return a linear expression from a dictionary or a sequence, that maps
  62         symbols to their coefficients, and a constant term. The coefficients and
  63         the constant term must be rational numbers.
  64
  65         For example, the linear expression x + 2y + 1 can be constructed using
  66         one of the following instructions:
  67
  68         >>> x, y = symbols('x y')
  69         >>> LinExpr({x: 1, y: 2}, 1)
  70         >>> LinExpr([(x, 1), (y, 2)], 1)
  71
  72         However, it may be easier to use overloaded operators:
  73
  74         >>> x, y = symbols('x y')
  75         >>> x + 2*y + 1
  76
  77         Alternatively, linear expressions can be constructed from a string:
  78
  79         >>> LinExpr('x + 2*y + 1')
  80
  81         A linear expression with a single symbol of coefficient 1 and no
  82         constant term is automatically subclassed as a Symbol instance. A linear
  83         expression with no symbol, only a constant term, is automatically
  84         subclassed as a Rational instance.
  85         """
  86         if isinstance(coefficients, str):
  87             if constant != 0:
  88                 raise TypeError('too many arguments')
  89             return LinExpr.fromstring(coefficients)
  90         if coefficients is None:
  91             return Rational(constant)
  92         if isinstance(coefficients, Mapping):
  93             coefficients = coefficients.items()
  94         coefficients = list(coefficients)
  95         for symbol, coefficient in coefficients:
  96             if not isinstance(symbol, Symbol):
  97                 raise TypeError('symbols must be Symbol instances')
  98             if not isinstance(coefficient, numbers.Rational):
  99                 raise TypeError('coefficients must be rational numbers')
 100         if not isinstance(constant, numbers.Rational):
 101             raise TypeError('constant must be a rational number')
 102         if len(coefficients) == 0:
 103             return Rational(constant)
 104         if len(coefficients) == 1 and constant == 0:
 105             symbol, coefficient = coefficients[0]
 106             if coefficient == 1:
 107                 return symbol
 108         coefficients = [(symbol, Fraction(coefficient))
 109             for symbol, coefficient in coefficients if coefficient != 0]
 110         coefficients.sort(key=lambda item: item[0].sortkey())
 111         self = object().__new__(cls)
 112         self._coefficients = OrderedDict(coefficients)
 113         self._constant = Fraction(constant)
 114         self._symbols = tuple(self._coefficients)
 115         self._dimension = len(self._symbols)
 116         return self
 117
 118     def coefficient(self, symbol):
 119         """
 120         Return the coefficient value of the given symbol, or 0 if the symbol
 121         does not appear in the expression.
 122         """
 123         if not isinstance(symbol, Symbol):
 124             raise TypeError('symbol must be a Symbol instance')
 125         return self._coefficients.get(symbol, Fraction(0))
 126
 127     __getitem__ = coefficient
 128
 129     def coefficients(self):
 130         """
 131         Iterate over the pairs (symbol, value) of linear terms in the
 132         expression. The constant term is ignored.
 133         """
 134         yield from self._coefficients.items()
 135
 136     @property
 137     def constant(self):
 138         """
 139         The constant term of the expression.
 140         """
 141         return self._constant
 142
 143     @property
 144     def symbols(self):
 145         """
 146         The tuple of symbols present in the expression, sorted according to
 147         Symbol.sortkey().
 148         """
 149         return self._symbols
 150
 151     @property
 152     def dimension(self):
 153         """
 154         The dimension of the expression, i.e. the number of symbols present in
 155         it.
 156         """
 157         return self._dimension
 158
 159     def __hash__(self):
 160         return hash((tuple(self._coefficients.items()), self._constant))
 161
 162     def isconstant(self):
 163         """
 164         Return True if the expression only consists of a constant term. In this
 165         case, it is a Rational instance.
 166         """
 167         return False
 168
 169     def issymbol(self):
 170         """
 171         Return True if an expression only consists of a symbol with coefficient
 172         1. In this case, it is a Symbol instance.
 173         """
 174         return False
 175
 176     def values(self):
 177         """
 178         Iterate over the coefficient values in the expression, and the constant
 179         term.
 180         """
 181         yield from self._coefficients.values()
 182         yield self._constant
 183
 184     def __bool__(self):
 185         return True
 186
 187     def __pos__(self):
 188         return self
 189
 190     def __neg__(self):
 191         return self * -1
 192
 193     @_polymorphic
 194     def __add__(self, other):
 195         """
 196         Return the sum of two linear expressions.
 197         """
 198         coefficients = defaultdict(Fraction, self._coefficients)
 199         for symbol, coefficient in other._coefficients.items():
 200             coefficients[symbol] += coefficient
 201         constant = self._constant + other._constant
 202         return LinExpr(coefficients, constant)
 203
 204     __radd__ = __add__
 205
 206     @_polymorphic
 207     def __sub__(self, other):
 208         """
 209         Return the difference between two linear expressions.
 210         """
 211         coefficients = defaultdict(Fraction, self._coefficients)
 212         for symbol, coefficient in other._coefficients.items():
 213             coefficients[symbol] -= coefficient
 214         constant = self._constant - other._constant
 215         return LinExpr(coefficients, constant)
 216
 217     @_polymorphic
 218     def __rsub__(self, other):
 219         return other - self
 220
 221     def __mul__(self, other):
 222         """
 223         Return the product of the linear expression by a rational.
 224         """
 225         if isinstance(other, numbers.Rational):
 226             coefficients = ((symbol, coefficient * other)
 227                 for symbol, coefficient in self._coefficients.items())
 228             constant = self._constant * other
 229             return LinExpr(coefficients, constant)
 230         return NotImplemented
 231
 232     __rmul__ = __mul__
 233
 234     def __truediv__(self, other):
 235         """
 236         Return the quotient of the linear expression by a rational.
 237         """
 238         if isinstance(other, numbers.Rational):
 239             coefficients = ((symbol, coefficient / other)
 240                 for symbol, coefficient in self._coefficients.items())
 241             constant = self._constant / other
 242             return LinExpr(coefficients, constant)
 243         return NotImplemented
 244
 245     @_polymorphic
 246     def __eq__(self, other):
 247         """
 248         Test whether two linear expressions are equal.
 249         """
 250         return self._coefficients == other._coefficients and \
 251             self._constant == other._constant
 252
 253     @_polymorphic
 254     def __lt__(self, other):
 255         from .polyhedra import Polyhedron
 256         return Polyhedron([], [other - self - 1])
 257
 258     @_polymorphic
 259     def __le__(self, other):
 260         from .polyhedra import Polyhedron
 261         return Polyhedron([], [other - self])
 262
 263     @_polymorphic
 264     def __ge__(self, other):
 265         from .polyhedra import Polyhedron
 266         return Polyhedron([], [self - other])
 267
 268     @_polymorphic
 269     def __gt__(self, other):
 270         from .polyhedra import Polyhedron
 271         return Polyhedron([], [self - other - 1])
 272
 273     def scaleint(self):
 274         """
 275         Return the expression multiplied by its lowest common denominator to
 276         make all values integer.
 277         """
 278         lcd = functools.reduce(lambda a, b: a*b // gcd(a, b),
 279             [value.denominator for value in self.values()])
 280         return self * lcd
 281
 282     def subs(self, symbol, expression=None):
 283         """
 284         Substitute the given symbol by an expression and return the resulting
 285         expression. Raise TypeError if the resulting expression is not linear.
 286
 287         >>> x, y = symbols('x y')
 288         >>> e = x + 2*y + 1
 289         >>> e.subs(y, x - 1)
 290         3*x - 1
 291
 292         To perform multiple substitutions at once, pass a sequence or a
 293         dictionary of (old, new) pairs to subs.
 294
 295         >>> e.subs({x: y, y: x})
 296         2*x + y + 1
 297         """
 298         if expression is None:
 299             substitutions = dict(symbol)
 300         else:
 301             substitutions = {symbol: expression}
 302         for symbol in substitutions:
 303             if not isinstance(symbol, Symbol):
 304                 raise TypeError('symbols must be Symbol instances')
 305         result = self._constant
 306         for symbol, coefficient in self._coefficients.items():
 307             expression = substitutions.get(symbol, symbol)
 308             result += coefficient * expression
 309         return result
 310
 311     @classmethod
 312     def _fromast(cls, node):
 313         if isinstance(node, ast.Module) and len(node.body) == 1:
 314             return cls._fromast(node.body[0])
 315         elif isinstance(node, ast.Expr):
 316             return cls._fromast(node.value)
 317         elif isinstance(node, ast.Name):
 318             return Symbol(node.id)
 319         elif isinstance(node, ast.Num):
 320             return Rational(node.n)
 321         elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
 322             return -cls._fromast(node.operand)
 323         elif isinstance(node, ast.BinOp):
 324             left = cls._fromast(node.left)
 325             right = cls._fromast(node.right)
 326             if isinstance(node.op, ast.Add):
 327                 return left + right
 328             elif isinstance(node.op, ast.Sub):
 329                 return left - right
 330             elif isinstance(node.op, ast.Mult):
 331                 return left * right
 332             elif isinstance(node.op, ast.Div):
 333                 return left / right
 334         raise SyntaxError('invalid syntax')
 335
 336     _RE_NUM_VAR = re.compile(r'(\d+|\))\s*([^\W\d]\w*|\()')
 337
 338     @classmethod
 339     def fromstring(cls, string):
 340         """
 341         Create an expression from a string. Raise SyntaxError if the string is
 342         not properly formatted.
 343         """
 344         # Add implicit multiplication operators, e.g. '5x' -> '5*x'.
 345         string = LinExpr._RE_NUM_VAR.sub(r'\1*\2', string)
 346         tree = ast.parse(string, 'eval')
 347         expr = cls._fromast(tree)
 348         if not isinstance(expr, cls):
 349             raise SyntaxError('invalid syntax')
 350         return expr
 351
 352     def __repr__(self):
 353         string = ''
 354         for i, (symbol, coefficient) in enumerate(self.coefficients()):
 355             if coefficient == 1:
 356                 if i != 0:
 357                     string += ' + '
 358             elif coefficient == -1:
 359                 string += '-' if i == 0 else ' - '
 360             elif i == 0:
 361                 string += '{}*'.format(coefficient)
 362             elif coefficient > 0:
 363                 string += ' + {}*'.format(coefficient)
 364             else:
 365                 string += ' - {}*'.format(-coefficient)
 366             string += '{}'.format(symbol)
 367         constant = self.constant
 368         if len(string) == 0:
 369             string += '{}'.format(constant)
 370         elif constant > 0:
 371             string += ' + {}'.format(constant)
 372         elif constant < 0:
 373             string += ' - {}'.format(-constant)
 374         return string
 375
 376     def _repr_latex_(self):
 377         string = ''
 378         for i, (symbol, coefficient) in enumerate(self.coefficients()):
 379             if coefficient == 1:
 380                 if i != 0:
 381                     string += ' + '
 382             elif coefficient == -1:
 383                 string += '-' if i == 0 else ' - '
 384             elif i == 0:
 385                 string += '{}'.format(coefficient._repr_latex_().strip('$'))
 386             elif coefficient > 0:
 387                 string += ' + {}'.format(coefficient._repr_latex_().strip('$'))
 388             elif coefficient < 0:
 389                 string += ' - {}'.format((-coefficient)._repr_latex_().strip('$'))
 390             string += '{}'.format(symbol._repr_latex_().strip('$'))
 391         constant = self.constant
 392         if len(string) == 0:
 393             string += '{}'.format(constant._repr_latex_().strip('$'))
 394         elif constant > 0:
 395             string += ' + {}'.format(constant._repr_latex_().strip('$'))
 396         elif constant < 0:
 397             string += ' - {}'.format((-constant)._repr_latex_().strip('$'))
 398         return '$${}$$'.format(string)
 399
 400     def _parenstr(self, always=False):
 401         string = str(self)
 402         if not always and (self.isconstant() or self.issymbol()):
 403             return string
 404         else:
 405             return '({})'.format(string)
 406
 407     @classmethod
 408     def fromsympy(cls, expr):
 409         """
 410         Create a linear expression from a SymPy expression. Raise TypeError is
 411         the sympy expression is not linear.
 412         """
 413         import sympy
 414         coefficients = []
 415         constant = 0
 416         for symbol, coefficient in expr.as_coefficients_dict().items():
 417             coefficient = Fraction(coefficient.p, coefficient.q)
 418             if symbol == sympy.S.One:
 419                 constant = coefficient
 420             elif isinstance(symbol, sympy.Dummy):
 421                 # We cannot properly convert dummy symbols with respect to
 422                 # symbol equalities.
 423                 raise TypeError('cannot convert dummy symbols')
 424             elif isinstance(symbol, sympy.Symbol):
 425                 symbol = Symbol(symbol.name)
 426                 coefficients.append((symbol, coefficient))
 427             else:
 428                 raise TypeError('non-linear expression: {!r}'.format(expr))
 429         expr = LinExpr(coefficients, constant)
 430         if not isinstance(expr, cls):
 431             raise TypeError('cannot convert to a {} instance'.format(cls.__name__))
 432         return expr
 433
 434     def tosympy(self):
 435         """
 436         Convert the linear expression to a SymPy expression.
 437         """
 438         import sympy
 439         expr = 0
 440         for symbol, coefficient in self.coefficients():
 441             term = coefficient * sympy.Symbol(symbol.name)
 442             expr += term
 443         expr += self.constant
 444         return expr
 445
 446
 447 class Symbol(LinExpr):
 448     """
 449     Symbols are the basic components to build expressions and constraints.
 450     They correspond to mathematical variables. Symbols are instances of
 451     class LinExpr and inherit its functionalities.
 452
 453     Two instances of Symbol are equal if they have the same name.
 454     """
 455
 456     __slots__ = (
 457         '_name',
 458         '_constant',
 459         '_symbols',
 460         '_dimension',
 461     )
 462
 463     def __new__(cls, name):
 464         """
 465         Return a symbol with the name string given in argument.
 466         """
 467         if not isinstance(name, str):
 468             raise TypeError('name must be a string')
 469         node = ast.parse(name)
 470         try:
 471             name = node.body[0].value.id
 472         except (AttributeError, SyntaxError):
 473             raise SyntaxError('invalid syntax')
 474         self = object().__new__(cls)
 475         self._name = name
 476         self._constant = Fraction(0)
 477         self._symbols = (self,)
 478         self._dimension = 1
 479         return self
 480
 481     @property
 482     def _coefficients(self):
 483         # This is not implemented as an attribute, because __hash__ is not
 484         # callable in __new__ in class Dummy.
 485         return {self: Fraction(1)}
 486
 487     @property
 488     def name(self):
 489         """
 490         The name of the symbol.
 491         """
 492         return self._name
 493
 494     def __hash__(self):
 495         return hash(self.sortkey())
 496
 497     def sortkey(self):
 498         """
 499         Return a sorting key for the symbol. It is useful to sort a list of
 500         symbols in a consistent order, as comparison functions are overridden
 501         (see the documentation of class LinExpr).
 502
 503         >>> sort(symbols, key=Symbol.sortkey)
 504         """
 505         return self.name,
 506
 507     def issymbol(self):
 508         return True
 509
 510     def __eq__(self, other):
 511         if isinstance(other, Symbol):
 512             return self.sortkey() == other.sortkey()
 513         return NotImplemented
 514
 515     def asdummy(self):
 516         """
 517         Return a new Dummy symbol instance with the same name.
 518         """
 519         return Dummy(self.name)
 520
 521     def __repr__(self):
 522         return self.name
 523
 524     def _repr_latex_(self):
 525         return '$${}$$'.format(self.name)
 526
 527
 528 def symbols(names):
 529     """
 530     This function returns a tuple of symbols whose names are taken from a comma
 531     or whitespace delimited string, or a sequence of strings. It is useful to
 532     define several symbols at once.
 533
 534     >>> x, y = symbols('x y')
 535     >>> x, y = symbols('x, y')
 536     >>> x, y = symbols(['x', 'y'])
 537     """
 538     if isinstance(names, str):
 539         names = names.replace(',', ' ').split()
 540     return tuple(Symbol(name) for name in names)
 541
 542
 543 class Dummy(Symbol):
 544     """
 545     A variation of Symbol in which all symbols are unique and identified by
 546     an internal count index. If a name is not supplied then a string value
 547     of the count index will be used. This is useful when a unique, temporary
 548     variable is needed and the name of the variable used in the expression
 549     is not important.
 550
 551     Unlike Symbol, Dummy instances with the same name are not equal:
 552
 553     >>> x = Symbol('x')
 554     >>> x1, x2 = Dummy('x'), Dummy('x')
 555     >>> x == x1
 556     False
 557     >>> x1 == x2
 558     False
 559     >>> x1 == x1
 560     True
 561     """
 562
 563     _count = 0
 564
 565     def __new__(cls, name=None):
 566         """
 567         Return a fresh dummy symbol with the name string given in argument.
 568         """
 569         if name is None:
 570             name = 'Dummy_{}'.format(Dummy._count)
 571         self = super().__new__(cls, name)
 572         self._index = Dummy._count
 573         Dummy._count += 1
 574         return self
 575
 576     def __hash__(self):
 577         return hash(self.sortkey())
 578
 579     def sortkey(self):
 580         return self._name, self._index
 581
 582     def __repr__(self):
 583         return '_{}'.format(self.name)
 584
 585     def _repr_latex_(self):
 586         return '$${}_{{{}}}$$'.format(self.name, self._index)
 587
 588
 589 class Rational(LinExpr, Fraction):
 590     """
 591     A particular case of linear expressions are rational values, i.e. linear
 592     expressions consisting only of a constant term, with no symbol. They are
 593     implemented by the Rational class, that inherits from both LinExpr and
 594     fractions.Fraction classes.
 595     """
 596
 597     __slots__ = (
 598         '_coefficients',
 599         '_constant',
 600         '_symbols',
 601         '_dimension',
 602     ) + Fraction.__slots__
 603
 604     def __new__(cls, numerator=0, denominator=None):
 605         self = object().__new__(cls)
 606         self._coefficients = {}
 607         self._constant = Fraction(numerator, denominator)
 608         self._symbols = ()
 609         self._dimension = 0
 610         self._numerator = self._constant.numerator
 611         self._denominator = self._constant.denominator
 612         return self
 613
 614     def __hash__(self):
 615         return Fraction.__hash__(self)
 616
 617     @property
 618     def constant(self):
 619         return self
 620
 621     def isconstant(self):
 622         return True
 623
 624     def __bool__(self):
 625         return Fraction.__bool__(self)
 626
 627     def __repr__(self):
 628         if self.denominator == 1:
 629             return '{!r}'.format(self.numerator)
 630         else:
 631             return '{!r}/{!r}'.format(self.numerator, self.denominator)
 632
 633     def _repr_latex_(self):
 634         if self.denominator == 1:
 635             return '$${}$$'.format(self.numerator)
 636         elif self.numerator < 0:
 637             return '$$-\\frac{{{}}}{{{}}}$$'.format(-self.numerator,
 638                 self.denominator)
 639         else:
 640             return '$$\\frac{{{}}}{{{}}}$$'.format(self.numerator,
 641                 self.denominator)