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 Rational(self._coefficients.get(symbol, 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         for symbol, coefficient in self._coefficients.items():
 135             yield symbol, Rational(coefficient)
 136
 137     @property
 138     def constant(self):
 139         """
 140         The constant term of the expression.
 141         """
 142         return Rational(self._constant)
 143
 144     @property
 145     def symbols(self):
 146         """
 147         The tuple of symbols present in the expression, sorted according to
 148         Symbol.sortkey().
 149         """
 150         return self._symbols
 151
 152     @property
 153     def dimension(self):
 154         """
 155         The dimension of the expression, i.e. the number of symbols present in
 156         it.
 157         """
 158         return self._dimension
 159
 160     def __hash__(self):
 161         return hash((tuple(self._coefficients.items()), self._constant))
 162
 163     def isconstant(self):
 164         """
 165         Return True if the expression only consists of a constant term. In this
 166         case, it is a Rational instance.
 167         """
 168         return False
 169
 170     def issymbol(self):
 171         """
 172         Return True if an expression only consists of a symbol with coefficient
 173         1. In this case, it is a Symbol instance.
 174         """
 175         return False
 176
 177     def values(self):
 178         """
 179         Iterate over the coefficient values in the expression, and the constant
 180         term.
 181         """
 182         for coefficient in self._coefficients.values():
 183             yield Rational(coefficient)
 184         yield Rational(self._constant)
 185
 186     def __bool__(self):
 187         return True
 188
 189     def __pos__(self):
 190         return self
 191
 192     def __neg__(self):
 193         return self * -1
 194
 195     @_polymorphic
 196     def __add__(self, other):
 197         """
 198         Return the sum of two linear expressions.
 199         """
 200         coefficients = defaultdict(Fraction, self._coefficients)
 201         for symbol, coefficient in other._coefficients.items():
 202             coefficients[symbol] += coefficient
 203         constant = self._constant + other._constant
 204         return LinExpr(coefficients, constant)
 205
 206     __radd__ = __add__
 207
 208     @_polymorphic
 209     def __sub__(self, other):
 210         """
 211         Return the difference between two linear expressions.
 212         """
 213         coefficients = defaultdict(Fraction, self._coefficients)
 214         for symbol, coefficient in other._coefficients.items():
 215             coefficients[symbol] -= coefficient
 216         constant = self._constant - other._constant
 217         return LinExpr(coefficients, constant)
 218
 219     @_polymorphic
 220     def __rsub__(self, other):
 221         return other - self
 222
 223     def __mul__(self, other):
 224         """
 225         Return the product of the linear expression by a rational.
 226         """
 227         if isinstance(other, numbers.Rational):
 228             coefficients = ((symbol, coefficient * other)
 229                 for symbol, coefficient in self._coefficients.items())
 230             constant = self._constant * other
 231             return LinExpr(coefficients, constant)
 232         return NotImplemented
 233
 234     __rmul__ = __mul__
 235
 236     def __truediv__(self, other):
 237         """
 238         Return the quotient of the linear expression by a rational.
 239         """
 240         if isinstance(other, numbers.Rational):
 241             coefficients = ((symbol, coefficient / other)
 242                 for symbol, coefficient in self._coefficients.items())
 243             constant = self._constant / other
 244             return LinExpr(coefficients, constant)
 245         return NotImplemented
 246
 247     @_polymorphic
 248     def __eq__(self, other):
 249         """
 250         Test whether two linear expressions are equal.
 251         """
 252         return isinstance(other, LinExpr) and \
 253             self._coefficients == other._coefficients and \
 254             self._constant == other._constant
 255
 256     def __le__(self, other):
 257         from .polyhedra import Le
 258         return Le(self, other)
 259
 260     def __lt__(self, other):
 261         from .polyhedra import Lt
 262         return Lt(self, other)
 263
 264     def __ge__(self, other):
 265         from .polyhedra import Ge
 266         return Ge(self, other)
 267
 268     def __gt__(self, other):
 269         from .polyhedra import Gt
 270         return Gt(self, other)
 271
 272     def scaleint(self):
 273         """
 274         Return the expression multiplied by its lowest common denominator to
 275         make all values integer.
 276         """
 277         lcm = functools.reduce(lambda a, b: a*b // gcd(a, b),
 278             [value.denominator for value in self.values()])
 279         return self * lcm
 280
 281     def subs(self, symbol, expression=None):
 282         """
 283         Substitute the given symbol by an expression and return the resulting
 284         expression. Raise TypeError if the resulting expression is not linear.
 285
 286         >>> x, y = symbols('x y')
 287         >>> e = x + 2*y + 1
 288         >>> e.subs(y, x - 1)
 289         3*x - 1
 290
 291         To perform multiple substitutions at once, pass a sequence or a
 292         dictionary of (old, new) pairs to subs.
 293
 294         >>> e.subs({x: y, y: x})
 295         2*x + y + 1
 296         """
 297         if expression is None:
 298             if isinstance(symbol, Mapping):
 299                 symbol = symbol.items()
 300             substitutions = symbol
 301         else:
 302             substitutions = [(symbol, expression)]
 303         result = self
 304         for symbol, expression in substitutions:
 305             if not isinstance(symbol, Symbol):
 306                 raise TypeError('symbols must be Symbol instances')
 307             coefficients = [(othersymbol, coefficient)
 308                 for othersymbol, coefficient in result._coefficients.items()
 309                 if othersymbol != symbol]
 310             coefficient = result._coefficients.get(symbol, 0)
 311             constant = result._constant
 312             result = LinExpr(coefficients, constant) + coefficient*expression
 313         return result
 314
 315     @classmethod
 316     def _fromast(cls, node):
 317         if isinstance(node, ast.Module) and len(node.body) == 1:
 318             return cls._fromast(node.body[0])
 319         elif isinstance(node, ast.Expr):
 320             return cls._fromast(node.value)
 321         elif isinstance(node, ast.Name):
 322             return Symbol(node.id)
 323         elif isinstance(node, ast.Num):
 324             return Rational(node.n)
 325         elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
 326             return -cls._fromast(node.operand)
 327         elif isinstance(node, ast.BinOp):
 328             left = cls._fromast(node.left)
 329             right = cls._fromast(node.right)
 330             if isinstance(node.op, ast.Add):
 331                 return left + right
 332             elif isinstance(node.op, ast.Sub):
 333                 return left - right
 334             elif isinstance(node.op, ast.Mult):
 335                 return left * right
 336             elif isinstance(node.op, ast.Div):
 337                 return left / right
 338         raise SyntaxError('invalid syntax')
 339
 340     _RE_NUM_VAR = re.compile(r'(\d+|\))\s*([^\W\d_]\w*|\()')
 341
 342     @classmethod
 343     def fromstring(cls, string):
 344         """
 345         Create an expression from a string. Raise SyntaxError if the string is
 346         not properly formatted.
 347         """
 348         # add implicit multiplication operators, e.g. '5x' -> '5*x'
 349         string = LinExpr._RE_NUM_VAR.sub(r'\1*\2', string)
 350         tree = ast.parse(string, 'eval')
 351         return cls._fromast(tree)
 352
 353     def __repr__(self):
 354         string = ''
 355         for i, (symbol, coefficient) in enumerate(self.coefficients()):
 356             if coefficient == 1:
 357                 if i != 0:
 358                     string += ' + '
 359             elif coefficient == -1:
 360                 string += '-' if i == 0 else ' - '
 361             elif i == 0:
 362                 string += '{}*'.format(coefficient)
 363             elif coefficient > 0:
 364                 string += ' + {}*'.format(coefficient)
 365             else:
 366                 string += ' - {}*'.format(-coefficient)
 367             string += '{}'.format(symbol)
 368         constant = self.constant
 369         if len(string) == 0:
 370             string += '{}'.format(constant)
 371         elif constant > 0:
 372             string += ' + {}'.format(constant)
 373         elif constant < 0:
 374             string += ' - {}'.format(-constant)
 375         return string
 376
 377     def _repr_latex_(self):
 378         string = ''
 379         for i, (symbol, coefficient) in enumerate(self.coefficients()):
 380             if coefficient == 1:
 381                 if i != 0:
 382                     string += ' + '
 383             elif coefficient == -1:
 384                 string += '-' if i == 0 else ' - '
 385             elif i == 0:
 386                 string += '{}'.format(coefficient._repr_latex_().strip('$'))
 387             elif coefficient > 0:
 388                 string += ' + {}'.format(coefficient._repr_latex_().strip('$'))
 389             elif coefficient < 0:
 390                 string += ' - {}'.format((-coefficient)._repr_latex_().strip('$'))
 391             string += '{}'.format(symbol._repr_latex_().strip('$'))
 392         constant = self.constant
 393         if len(string) == 0:
 394             string += '{}'.format(constant._repr_latex_().strip('$'))
 395         elif constant > 0:
 396             string += ' + {}'.format(constant._repr_latex_().strip('$'))
 397         elif constant < 0:
 398             string += ' - {}'.format((-constant)._repr_latex_().strip('$'))
 399         return '$${}$$'.format(string)
 400
 401     def _parenstr(self, always=False):
 402         string = str(self)
 403         if not always and (self.isconstant() or self.issymbol()):
 404             return string
 405         else:
 406             return '({})'.format(string)
 407
 408     @classmethod
 409     def fromsympy(cls, expr):
 410         """
 411         Create a linear expression from a sympy expression. Raise ValueError is
 412         the sympy expression is not linear.
 413         """
 414         import sympy
 415         coefficients = []
 416         constant = 0
 417         for symbol, coefficient in expr.as_coefficients_dict().items():
 418             coefficient = Fraction(coefficient.p, coefficient.q)
 419             if symbol == sympy.S.One:
 420                 constant = coefficient
 421             elif isinstance(symbol, sympy.Symbol):
 422                 symbol = Symbol(symbol.name)
 423                 coefficients.append((symbol, coefficient))
 424             else:
 425                 raise ValueError('non-linear expression: {!r}'.format(expr))
 426         return LinExpr(coefficients, constant)
 427
 428     def tosympy(self):
 429         """
 430         Convert the linear expression to a sympy expression.
 431         """
 432         import sympy
 433         expr = 0
 434         for symbol, coefficient in self.coefficients():
 435             term = coefficient * sympy.Symbol(symbol.name)
 436             expr += term
 437         expr += self.constant
 438         return expr
 439
 440
 441 class Symbol(LinExpr):
 442     """
 443     Symbols are the basic components to build expressions and constraints.
 444     They correspond to mathematical variables. Symbols are instances of
 445     class LinExpr and inherit its functionalities.
 446
 447     Two instances of Symbol are equal if they have the same name.
 448     """
 449
 450     def __new__(cls, name):
 451         """
 452         Return a symbol with the name string given in argument.
 453         """
 454         if not isinstance(name, str):
 455             raise TypeError('name must be a string')
 456         self = object().__new__(cls)
 457         self._name = name.strip()
 458         self._coefficients = {self: Fraction(1)}
 459         self._constant = Fraction(0)
 460         self._symbols = (self,)
 461         self._dimension = 1
 462         return self
 463
 464     @property
 465     def name(self):
 466         """
 467         The name of the symbol.
 468         """
 469         return self._name
 470
 471     def __hash__(self):
 472         return hash(self.sortkey())
 473
 474     def sortkey(self):
 475         """
 476         Return a sorting key for the symbol. It is useful to sort a list of
 477         symbols in a consistent order, as comparison functions are overridden
 478         (see the documentation of class LinExpr).
 479
 480         >>> sort(symbols, key=Symbol.sortkey)
 481         """
 482         return self.name,
 483
 484     def issymbol(self):
 485         return True
 486
 487     def __eq__(self, other):
 488         return self.sortkey() == other.sortkey()
 489
 490     def asdummy(self):
 491         """
 492         Return a new Dummy symbol instance with the same name.
 493         """
 494         return Dummy(self.name)
 495
 496     @classmethod
 497     def _fromast(cls, node):
 498         if isinstance(node, ast.Module) and len(node.body) == 1:
 499             return cls._fromast(node.body[0])
 500         elif isinstance(node, ast.Expr):
 501             return cls._fromast(node.value)
 502         elif isinstance(node, ast.Name):
 503             return Symbol(node.id)
 504         raise SyntaxError('invalid syntax')
 505
 506     def __repr__(self):
 507         return self.name
 508
 509     def _repr_latex_(self):
 510         return '$${}$$'.format(self.name)
 511
 512     @classmethod
 513     def fromsympy(cls, expr):
 514         import sympy
 515         if isinstance(expr, sympy.Dummy):
 516             return Dummy(expr.name)
 517         elif isinstance(expr, sympy.Symbol):
 518             return Symbol(expr.name)
 519         else:
 520             raise TypeError('expr must be a sympy.Symbol instance')
 521
 522
 523 class Dummy(Symbol):
 524     """
 525     A variation of Symbol in which all symbols are unique and identified by
 526     an internal count index. If a name is not supplied then a string value
 527     of the count index will be used. This is useful when a unique, temporary
 528     variable is needed and the name of the variable used in the expression
 529     is not important.
 530
 531     Unlike Symbol, Dummy instances with the same name are not equal:
 532
 533     >>> x = Symbol('x')
 534     >>> x1, x2 = Dummy('x'), Dummy('x')
 535     >>> x == x1
 536     False
 537     >>> x1 == x2
 538     False
 539     >>> x1 == x1
 540     True
 541     """
 542
 543     _count = 0
 544
 545     def __new__(cls, name=None):
 546         """
 547         Return a fresh dummy symbol with the name string given in argument.
 548         """
 549         if name is None:
 550             name = 'Dummy_{}'.format(Dummy._count)
 551         elif not isinstance(name, str):
 552             raise TypeError('name must be a string')
 553         self = object().__new__(cls)
 554         self._index = Dummy._count
 555         self._name = name.strip()
 556         self._coefficients = {self: Fraction(1)}
 557         self._constant = Fraction(0)
 558         self._symbols = (self,)
 559         self._dimension = 1
 560         Dummy._count += 1
 561         return self
 562
 563     def __hash__(self):
 564         return hash(self.sortkey())
 565
 566     def sortkey(self):
 567         return self._name, self._index
 568
 569     def __repr__(self):
 570         return '_{}'.format(self.name)
 571
 572     def _repr_latex_(self):
 573         return '$${}_{{{}}}$$'.format(self.name, self._index)
 574
 575
 576 def symbols(names):
 577     """
 578     This function returns a tuple of symbols whose names are taken from a comma
 579     or whitespace delimited string, or a sequence of strings. It is useful to
 580     define several symbols at once.
 581
 582     >>> x, y = symbols('x y')
 583     >>> x, y = symbols('x, y')
 584     >>> x, y = symbols(['x', 'y'])
 585     """
 586     if isinstance(names, str):
 587         names = names.replace(',', ' ').split()
 588     return tuple(Symbol(name) for name in names)
 589
 590
 591 class Rational(LinExpr, Fraction):
 592     """
 593     A particular case of linear expressions are rational values, i.e. linear
 594     expressions consisting only of a constant term, with no symbol. They are
 595     implemented by the Rational class, that inherits from both LinExpr and
 596     fractions.Fraction classes.
 597     """
 598
 599     def __new__(cls, numerator=0, denominator=None):
 600         self = object().__new__(cls)
 601         self._coefficients = {}
 602         self._constant = Fraction(numerator, denominator)
 603         self._symbols = ()
 604         self._dimension = 0
 605         self._numerator = self._constant.numerator
 606         self._denominator = self._constant.denominator
 607         return self
 608
 609     def __hash__(self):
 610         return Fraction.__hash__(self)
 611
 612     @property
 613     def constant(self):
 614         return self
 615
 616     def isconstant(self):
 617         return True
 618
 619     def __bool__(self):
 620         return Fraction.__bool__(self)
 621
 622     def __repr__(self):
 623         if self.denominator == 1:
 624             return '{!r}'.format(self.numerator)
 625         else:
 626             return '{!r}/{!r}'.format(self.numerator, self.denominator)
 627
 628     def _repr_latex_(self):
 629         if self.denominator == 1:
 630             return '$${}$$'.format(self.numerator)
 631         elif self.numerator < 0:
 632             return '$$-\\frac{{{}}}{{{}}}$$'.format(-self.numerator,
 633                 self.denominator)
 634         else:
 635             return '$$\\frac{{{}}}{{{}}}$$'.format(self.numerator,
 636                 self.denominator)
 637
 638     @classmethod
 639     def fromsympy(cls, expr):
 640         import sympy
 641         if isinstance(expr, sympy.Rational):
 642             return Rational(expr.p, expr.q)
 643         elif isinstance(expr, numbers.Rational):
 644             return Rational(expr)
 645         else:
 646             raise TypeError('expr must be a sympy.Rational instance')