pypol/linexprs.py

   1 import ast
   2 import functools
   3 import numbers
   4 import re
   5
   6 from collections import OrderedDict
   7 from fractions import Fraction, gcd
   8
   9
  10 __all__ = [
  11     'Expression',
  12     'Symbol', 'symbols',
  13     'Constant',
  14 ]
  15
  16
  17 def _polymorphic(func):
  18     @functools.wraps(func)
  19     def wrapper(left, right):
  20         if isinstance(right, Expression):
  21             return func(left, right)
  22         elif isinstance(right, numbers.Rational):
  23             right = Constant(right)
  24             return func(left, right)
  25         return NotImplemented
  26     return wrapper
  27
  28
  29 class Expression:
  30     """
  31     This class implements linear expressions.
  32     """
  33
  34     __slots__ = (
  35         '_coefficients',
  36         '_constant',
  37         '_symbols',
  38         '_dimension',
  39         '_hash',
  40     )
  41
  42     def __new__(cls, coefficients=None, constant=0):
  43         if isinstance(coefficients, str):
  44             if constant:
  45                 raise TypeError('too many arguments')
  46             return cls.fromstring(coefficients)
  47         if isinstance(coefficients, dict):
  48             coefficients = coefficients.items()
  49         if coefficients is None:
  50             return Constant(constant)
  51         coefficients = [(symbol, coefficient)
  52             for symbol, coefficient in coefficients if coefficient != 0]
  53         if len(coefficients) == 0:
  54             return Constant(constant)
  55         elif len(coefficients) == 1 and constant == 0:
  56             symbol, coefficient = coefficients[0]
  57             if coefficient == 1:
  58                 return Symbol(symbol)
  59         self = object().__new__(cls)
  60         self._coefficients = {}
  61         for symbol, coefficient in coefficients:
  62             if isinstance(symbol, Symbol):
  63                 symbol = symbol.name
  64             elif not isinstance(symbol, str):
  65                 raise TypeError('symbols must be strings or Symbol instances')
  66             if isinstance(coefficient, Constant):
  67                 coefficient = coefficient.constant
  68             if not isinstance(coefficient, numbers.Rational):
  69                 raise TypeError('coefficients must be rational numbers '
  70                     'or Constant instances')
  71             self._coefficients[symbol] = coefficient
  72         self._coefficients = OrderedDict(sorted(self._coefficients.items()))
  73         if isinstance(constant, Constant):
  74             constant = constant.constant
  75         if not isinstance(constant, numbers.Rational):
  76             raise TypeError('constant must be a rational number '
  77                 'or a Constant instance')
  78         self._constant = constant
  79         self._symbols = tuple(self._coefficients)
  80         self._dimension = len(self._symbols)
  81         self._hash = hash((tuple(self._coefficients.items()), self._constant))
  82         return self
  83
  84     def coefficient(self, symbol):
  85         if isinstance(symbol, Symbol):
  86             symbol = symbol.name
  87         elif not isinstance(symbol, str):
  88             raise TypeError('symbol must be a string or a Symbol instance')
  89         try:
  90             return self._coefficients[symbol]
  91         except KeyError:
  92             return 0
  93
  94     __getitem__ = coefficient
  95
  96     def coefficients(self):
  97         yield from self._coefficients.items()
  98
  99     @property
 100     def constant(self):
 101         return self._constant
 102
 103     @property
 104     def symbols(self):
 105         return self._symbols
 106
 107     @property
 108     def dimension(self):
 109         return self._dimension
 110
 111     def __hash__(self):
 112         return self._hash
 113
 114     def isconstant(self):
 115         return False
 116
 117     def issymbol(self):
 118         return False
 119
 120     def values(self):
 121         for symbol in self.symbols:
 122             yield self.coefficient(symbol)
 123         yield self.constant
 124
 125     def __bool__(self):
 126         return True
 127
 128     def __pos__(self):
 129         return self
 130
 131     def __neg__(self):
 132         return self * -1
 133
 134     @_polymorphic
 135     def __add__(self, other):
 136         coefficients = dict(self.coefficients())
 137         for symbol, coefficient in other.coefficients():
 138             if symbol in coefficients:
 139                 coefficients[symbol] += coefficient
 140             else:
 141                 coefficients[symbol] = coefficient
 142         constant = self.constant + other.constant
 143         return Expression(coefficients, constant)
 144
 145     __radd__ = __add__
 146
 147     @_polymorphic
 148     def __sub__(self, other):
 149         coefficients = dict(self.coefficients())
 150         for symbol, coefficient in other.coefficients():
 151             if symbol in coefficients:
 152                 coefficients[symbol] -= coefficient
 153             else:
 154                 coefficients[symbol] = -coefficient
 155         constant = self.constant - other.constant
 156         return Expression(coefficients, constant)
 157
 158     def __rsub__(self, other):
 159         return -(self - other)
 160
 161     @_polymorphic
 162     def __mul__(self, other):
 163         if other.isconstant():
 164             coefficients = dict(self.coefficients())
 165             for symbol in coefficients:
 166                 coefficients[symbol] *= other.constant
 167             constant = self.constant * other.constant
 168             return Expression(coefficients, constant)
 169         if isinstance(other, Expression) and not self.isconstant():
 170             raise ValueError('non-linear expression: '
 171                     '{} * {}'.format(self._parenstr(), other._parenstr()))
 172         return NotImplemented
 173
 174     __rmul__ = __mul__
 175
 176     @_polymorphic
 177     def __truediv__(self, other):
 178         if other.isconstant():
 179             coefficients = dict(self.coefficients())
 180             for symbol in coefficients:
 181                 coefficients[symbol] = \
 182                         Fraction(coefficients[symbol], other.constant)
 183             constant = Fraction(self.constant, other.constant)
 184             return Expression(coefficients, constant)
 185         if isinstance(other, Expression):
 186             raise ValueError('non-linear expression: '
 187                 '{} / {}'.format(self._parenstr(), other._parenstr()))
 188         return NotImplemented
 189
 190     def __rtruediv__(self, other):
 191         if isinstance(other, self):
 192             if self.isconstant():
 193                 constant = Fraction(other, self.constant)
 194                 return Expression(constant=constant)
 195             else:
 196                 raise ValueError('non-linear expression: '
 197                         '{} / {}'.format(other._parenstr(), self._parenstr()))
 198         return NotImplemented
 199
 200     @_polymorphic
 201     def __eq__(self, other):
 202         # "normal" equality
 203         # see http://docs.sympy.org/dev/tutorial/gotchas.html#equals-signs
 204         return isinstance(other, Expression) and \
 205                 self._coefficients == other._coefficients and \
 206                 self.constant == other.constant
 207
 208     @_polymorphic
 209     def __le__(self, other):
 210         from .polyhedra import Le
 211         return Le(self, other)
 212
 213     @_polymorphic
 214     def __lt__(self, other):
 215         from .polyhedra import Lt
 216         return Lt(self, other)
 217
 218     @_polymorphic
 219     def __ge__(self, other):
 220         from .polyhedra import Ge
 221         return Ge(self, other)
 222
 223     @_polymorphic
 224     def __gt__(self, other):
 225         from .polyhedra import Gt
 226         return Gt(self, other)
 227
 228     def _toint(self):
 229         lcm = functools.reduce(lambda a, b: a*b // gcd(a, b),
 230             [value.denominator for value in self.values()])
 231         return self * lcm
 232
 233     @classmethod
 234     def _fromast(cls, node):
 235         if isinstance(node, ast.Module) and len(node.body) == 1:
 236             return cls._fromast(node.body[0])
 237         elif isinstance(node, ast.Expr):
 238             return cls._fromast(node.value)
 239         elif isinstance(node, ast.Name):
 240             return Symbol(node.id)
 241         elif isinstance(node, ast.Num):
 242             return Constant(node.n)
 243         elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
 244             return -cls._fromast(node.operand)
 245         elif isinstance(node, ast.BinOp):
 246             left = cls._fromast(node.left)
 247             right = cls._fromast(node.right)
 248             if isinstance(node.op, ast.Add):
 249                 return left + right
 250             elif isinstance(node.op, ast.Sub):
 251                 return left - right
 252             elif isinstance(node.op, ast.Mult):
 253                 return left * right
 254             elif isinstance(node.op, ast.Div):
 255                 return left / right
 256         raise SyntaxError('invalid syntax')
 257
 258     _RE_NUM_VAR = re.compile(r'(\d+|\))\s*([^\W\d_]\w*|\()')
 259
 260     @classmethod
 261     def fromstring(cls, string):
 262         # add implicit multiplication operators, e.g. '5x' -> '5*x'
 263         string = cls._RE_NUM_VAR.sub(r'\1*\2', string)
 264         tree = ast.parse(string, 'eval')
 265         return cls._fromast(tree)
 266
 267     def __str__(self):
 268         string = ''
 269         i = 0
 270         for symbol in self.symbols:
 271             coefficient = self.coefficient(symbol)
 272             if coefficient == 1:
 273                 if i == 0:
 274                     string += symbol
 275                 else:
 276                     string += ' + {}'.format(symbol)
 277             elif coefficient == -1:
 278                 if i == 0:
 279                     string += '-{}'.format(symbol)
 280                 else:
 281                     string += ' - {}'.format(symbol)
 282             else:
 283                 if i == 0:
 284                     string += '{}*{}'.format(coefficient, symbol)
 285                 elif coefficient > 0:
 286                     string += ' + {}*{}'.format(coefficient, symbol)
 287                 else:
 288                     assert coefficient < 0
 289                     coefficient *= -1
 290                     string += ' - {}*{}'.format(coefficient, symbol)
 291             i += 1
 292         constant = self.constant
 293         if constant != 0 and i == 0:
 294             string += '{}'.format(constant)
 295         elif constant > 0:
 296             string += ' + {}'.format(constant)
 297         elif constant < 0:
 298             constant *= -1
 299             string += ' - {}'.format(constant)
 300         if string == '':
 301             string = '0'
 302         return string
 303
 304     def _parenstr(self, always=False):
 305         string = str(self)
 306         if not always and (self.isconstant() or self.issymbol()):
 307             return string
 308         else:
 309             return '({})'.format(string)
 310
 311     def __repr__(self):
 312         return '{}({!r})'.format(self.__class__.__name__, str(self))
 313
 314     @classmethod
 315     def fromsympy(cls, expr):
 316         import sympy
 317         coefficients = {}
 318         constant = 0
 319         for symbol, coefficient in expr.as_coefficients_dict().items():
 320             coefficient = Fraction(coefficient.p, coefficient.q)
 321             if symbol == sympy.S.One:
 322                 constant = coefficient
 323             elif isinstance(symbol, sympy.Symbol):
 324                 symbol = symbol.name
 325                 coefficients[symbol] = coefficient
 326             else:
 327                 raise ValueError('non-linear expression: {!r}'.format(expr))
 328         return cls(coefficients, constant)
 329
 330     def tosympy(self):
 331         import sympy
 332         expr = 0
 333         for symbol, coefficient in self.coefficients():
 334             term = coefficient * sympy.Symbol(symbol)
 335             expr += term
 336         expr += self.constant
 337         return expr
 338
 339
 340 class Symbol(Expression):
 341
 342     __slots__ = (
 343         '_name',
 344         '_hash',
 345     )
 346
 347     def __new__(cls, name):
 348         if isinstance(name, Symbol):
 349             name = name.name
 350         elif not isinstance(name, str):
 351             raise TypeError('name must be a string or a Symbol instance')
 352         name = name.strip()
 353         self = object().__new__(cls)
 354         self._name = name
 355         self._hash = hash(self._name)
 356         return self
 357
 358     @property
 359     def name(self):
 360         return self._name
 361
 362     def __hash__(self):
 363         return self._hash
 364
 365     def coefficient(self, symbol):
 366         if isinstance(symbol, Symbol):
 367             symbol = symbol.name
 368         elif not isinstance(symbol, str):
 369             raise TypeError('symbol must be a string or a Symbol instance')
 370         if symbol == self.name:
 371             return 1
 372         else:
 373             return 0
 374
 375     def coefficients(self):
 376         yield self.name, 1
 377
 378     @property
 379     def constant(self):
 380         return 0
 381
 382     @property
 383     def symbols(self):
 384         return self.name,
 385
 386     @property
 387     def dimension(self):
 388         return 1
 389
 390     def issymbol(self):
 391         return True
 392
 393     def __eq__(self, other):
 394         return isinstance(other, Symbol) and self.name == other.name
 395
 396     @classmethod
 397     def _fromast(cls, node):
 398         if isinstance(node, ast.Module) and len(node.body) == 1:
 399             return cls._fromast(node.body[0])
 400         elif isinstance(node, ast.Expr):
 401             return cls._fromast(node.value)
 402         elif isinstance(node, ast.Name):
 403             return Symbol(node.id)
 404         raise SyntaxError('invalid syntax')
 405
 406     def __repr__(self):
 407         return '{}({!r})'.format(self.__class__.__name__, self._name)
 408
 409     @classmethod
 410     def fromsympy(cls, expr):
 411         import sympy
 412         if isinstance(expr, sympy.Symbol):
 413             return cls(expr.name)
 414         else:
 415             raise TypeError('expr must be a sympy.Symbol instance')
 416
 417
 418 def symbols(names):
 419     if isinstance(names, str):
 420         names = names.replace(',', ' ').split()
 421     return (Symbol(name) for name in names)
 422
 423
 424 class Constant(Expression):
 425
 426     __slots__ = (
 427         '_constant',
 428         '_hash',
 429     )
 430
 431     def __new__(cls, numerator=0, denominator=None):
 432         self = object().__new__(cls)
 433         if denominator is None and isinstance(numerator, Constant):
 434             self._constant = numerator.constant
 435         else:
 436             self._constant = Fraction(numerator, denominator)
 437         self._hash = hash(self._constant)
 438         return self
 439
 440     def __hash__(self):
 441         return self._hash
 442
 443     def coefficient(self, symbol):
 444         if isinstance(symbol, Symbol):
 445             symbol = symbol.name
 446         elif not isinstance(symbol, str):
 447             raise TypeError('symbol must be a string or a Symbol instance')
 448         return 0
 449
 450     def coefficients(self):
 451         yield from []
 452
 453     @property
 454     def symbols(self):
 455         return ()
 456
 457     @property
 458     def dimension(self):
 459         return 0
 460
 461     def isconstant(self):
 462         return True
 463
 464     @_polymorphic
 465     def __eq__(self, other):
 466         return isinstance(other, Constant) and self.constant == other.constant
 467
 468     def __bool__(self):
 469         return self.constant != 0
 470
 471     @classmethod
 472     def fromstring(cls, string):
 473         if isinstance(string, str):
 474             return Constant(Fraction(string))
 475         else:
 476             raise TypeError('string must be a string instance')
 477
 478     def __repr__(self):
 479         if self.constant.denominator == 1:
 480             return '{}({!r})'.format(self.__class__.__name__,
 481                 self.constant.numerator)
 482         else:
 483             return '{}({!r}, {!r})'.format(self.__class__.__name__,
 484                 self.constant.numerator, self.constant.denominator)
 485
 486     @classmethod
 487     def fromsympy(cls, expr):
 488         import sympy
 489         if isinstance(expr, sympy.Rational):
 490             return cls(expr.p, expr.q)
 491         elif isinstance(expr, numbers.Rational):
 492             return cls(expr)
 493         else:
 494             raise TypeError('expr must be a sympy.Rational instance')