Mention that we use Z-polyhedra
[linpy.git] / 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 + 2*y + 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 + 2y + 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. Unlike methods
249 LinExpr.__lt__(), LinExpr.__le__(), LinExpr.__ge__(), LinExpr.__gt__(),
250 the result is a boolean value, not a polyhedron. To express that two
251 linear expressions are equal or not equal, use functions Eq() and Ne()
252 instead.
253 """
254 return self._coefficients == other._coefficients and \
255 self._constant == other._constant
256
257 @_polymorphic
258 def __lt__(self, other):
259 from .polyhedra import Polyhedron
260 return Polyhedron([], [other - self - 1])
261
262 @_polymorphic
263 def __le__(self, other):
264 from .polyhedra import Polyhedron
265 return Polyhedron([], [other - self])
266
267 @_polymorphic
268 def __ge__(self, other):
269 from .polyhedra import Polyhedron
270 return Polyhedron([], [self - other])
271
272 @_polymorphic
273 def __gt__(self, other):
274 from .polyhedra import Polyhedron
275 return Polyhedron([], [self - other - 1])
276
277 def scaleint(self):
278 """
279 Return the expression multiplied by its lowest common denominator to
280 make all values integer.
281 """
282 lcd = functools.reduce(lambda a, b: a*b // gcd(a, b),
283 [value.denominator for value in self.values()])
284 return self * lcd
285
286 def subs(self, symbol, expression=None):
287 """
288 Substitute the given symbol by an expression and return the resulting
289 expression. Raise TypeError if the resulting expression is not linear.
290
291 >>> x, y = symbols('x y')
292 >>> e = x + 2*y + 1
293 >>> e.subs(y, x - 1)
294 3*x - 1
295
296 To perform multiple substitutions at once, pass a sequence or a
297 dictionary of (old, new) pairs to subs.
298
299 >>> e.subs({x: y, y: x})
300 2*x + y + 1
301 """
302 if expression is None:
303 substitutions = dict(symbol)
304 else:
305 substitutions = {symbol: expression}
306 for symbol in substitutions:
307 if not isinstance(symbol, Symbol):
308 raise TypeError('symbols must be Symbol instances')
309 result = self._constant
310 for symbol, coefficient in self._coefficients.items():
311 expression = substitutions.get(symbol, symbol)
312 result += 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 expr = cls._fromast(tree)
352 if not isinstance(expr, cls):
353 raise SyntaxError('invalid syntax')
354 return expr
355
356 def __repr__(self):
357 string = ''
358 for i, (symbol, coefficient) in enumerate(self.coefficients()):
359 if coefficient == 1:
360 if i != 0:
361 string += ' + '
362 elif coefficient == -1:
363 string += '-' if i == 0 else ' - '
364 elif i == 0:
365 string += '{}*'.format(coefficient)
366 elif coefficient > 0:
367 string += ' + {}*'.format(coefficient)
368 else:
369 string += ' - {}*'.format(-coefficient)
370 string += '{}'.format(symbol)
371 constant = self.constant
372 if len(string) == 0:
373 string += '{}'.format(constant)
374 elif constant > 0:
375 string += ' + {}'.format(constant)
376 elif constant < 0:
377 string += ' - {}'.format(-constant)
378 return string
379
380 def _repr_latex_(self):
381 string = ''
382 for i, (symbol, coefficient) in enumerate(self.coefficients()):
383 if coefficient == 1:
384 if i != 0:
385 string += ' + '
386 elif coefficient == -1:
387 string += '-' if i == 0 else ' - '
388 elif i == 0:
389 string += '{}'.format(coefficient._repr_latex_().strip('$'))
390 elif coefficient > 0:
391 string += ' + {}'.format(coefficient._repr_latex_().strip('$'))
392 elif coefficient < 0:
393 string += ' - {}'.format((-coefficient)._repr_latex_().strip('$'))
394 string += '{}'.format(symbol._repr_latex_().strip('$'))
395 constant = self.constant
396 if len(string) == 0:
397 string += '{}'.format(constant._repr_latex_().strip('$'))
398 elif constant > 0:
399 string += ' + {}'.format(constant._repr_latex_().strip('$'))
400 elif constant < 0:
401 string += ' - {}'.format((-constant)._repr_latex_().strip('$'))
402 return '$${}$$'.format(string)
403
404 def _parenstr(self, always=False):
405 string = str(self)
406 if not always and (self.isconstant() or self.issymbol()):
407 return string
408 else:
409 return '({})'.format(string)
410
411 @classmethod
412 def fromsympy(cls, expr):
413 """
414 Create a linear expression from a SymPy expression. Raise TypeError is
415 the sympy expression is not linear.
416 """
417 import sympy
418 coefficients = []
419 constant = 0
420 for symbol, coefficient in expr.as_coefficients_dict().items():
421 coefficient = Fraction(coefficient.p, coefficient.q)
422 if symbol == sympy.S.One:
423 constant = coefficient
424 elif isinstance(symbol, sympy.Dummy):
425 # We cannot properly convert dummy symbols with respect to
426 # symbol equalities.
427 raise TypeError('cannot convert dummy symbols')
428 elif isinstance(symbol, sympy.Symbol):
429 symbol = Symbol(symbol.name)
430 coefficients.append((symbol, coefficient))
431 else:
432 raise TypeError('non-linear expression: {!r}'.format(expr))
433 expr = LinExpr(coefficients, constant)
434 if not isinstance(expr, cls):
435 raise TypeError('cannot convert to a {} instance'.format(cls.__name__))
436 return expr
437
438 def tosympy(self):
439 """
440 Convert the linear expression to a SymPy expression.
441 """
442 import sympy
443 expr = 0
444 for symbol, coefficient in self.coefficients():
445 term = coefficient * sympy.Symbol(symbol.name)
446 expr += term
447 expr += self.constant
448 return expr
449
450
451 class Symbol(LinExpr):
452 """
453 Symbols are the basic components to build expressions and constraints.
454 They correspond to mathematical variables. Symbols are instances of
455 class LinExpr and inherit its functionalities.
456
457 Two instances of Symbol are equal if they have the same name.
458 """
459
460 __slots__ = (
461 '_name',
462 '_constant',
463 '_symbols',
464 '_dimension',
465 )
466
467 def __new__(cls, name):
468 """
469 Return a symbol with the name string given in argument.
470 """
471 if not isinstance(name, str):
472 raise TypeError('name must be a string')
473 node = ast.parse(name)
474 try:
475 name = node.body[0].value.id
476 except (AttributeError, SyntaxError):
477 raise SyntaxError('invalid syntax')
478 self = object().__new__(cls)
479 self._name = name
480 self._constant = Fraction(0)
481 self._symbols = (self,)
482 self._dimension = 1
483 return self
484
485 @property
486 def _coefficients(self):
487 # This is not implemented as an attribute, because __hash__ is not
488 # callable in __new__ in class Dummy.
489 return {self: Fraction(1)}
490
491 @property
492 def name(self):
493 """
494 The name of the symbol.
495 """
496 return self._name
497
498 def __hash__(self):
499 return hash(self.sortkey())
500
501 def sortkey(self):
502 """
503 Return a sorting key for the symbol. It is useful to sort a list of
504 symbols in a consistent order, as comparison functions are overridden
505 (see the documentation of class LinExpr).
506
507 >>> sort(symbols, key=Symbol.sortkey)
508 """
509 return self.name,
510
511 def issymbol(self):
512 return True
513
514 def __eq__(self, other):
515 if isinstance(other, Symbol):
516 return self.sortkey() == other.sortkey()
517 return NotImplemented
518
519 def asdummy(self):
520 """
521 Return a new Dummy symbol instance with the same name.
522 """
523 return Dummy(self.name)
524
525 def __repr__(self):
526 return self.name
527
528 def _repr_latex_(self):
529 return '$${}$$'.format(self.name)
530
531
532 def symbols(names):
533 """
534 This function returns a tuple of symbols whose names are taken from a comma
535 or whitespace delimited string, or a sequence of strings. It is useful to
536 define several symbols at once.
537
538 >>> x, y = symbols('x y')
539 >>> x, y = symbols('x, y')
540 >>> x, y = symbols(['x', 'y'])
541 """
542 if isinstance(names, str):
543 names = names.replace(',', ' ').split()
544 return tuple(Symbol(name) for name in names)
545
546
547 class Dummy(Symbol):
548 """
549 A variation of Symbol in which all symbols are unique and identified by
550 an internal count index. If a name is not supplied then a string value
551 of the count index will be used. This is useful when a unique, temporary
552 variable is needed and the name of the variable used in the expression
553 is not important.
554
555 Unlike Symbol, Dummy instances with the same name are not equal:
556
557 >>> x = Symbol('x')
558 >>> x1, x2 = Dummy('x'), Dummy('x')
559 >>> x == x1
560 False
561 >>> x1 == x2
562 False
563 >>> x1 == x1
564 True
565 """
566
567 _count = 0
568
569 def __new__(cls, name=None):
570 """
571 Return a fresh dummy symbol with the name string given in argument.
572 """
573 if name is None:
574 name = 'Dummy_{}'.format(Dummy._count)
575 self = super().__new__(cls, name)
576 self._index = Dummy._count
577 Dummy._count += 1
578 return self
579
580 def __hash__(self):
581 return hash(self.sortkey())
582
583 def sortkey(self):
584 return self._name, self._index
585
586 def __repr__(self):
587 return '_{}'.format(self.name)
588
589 def _repr_latex_(self):
590 return '$${}_{{{}}}$$'.format(self.name, self._index)
591
592
593 class Rational(LinExpr, Fraction):
594 """
595 A particular case of linear expressions are rational values, i.e. linear
596 expressions consisting only of a constant term, with no symbol. They are
597 implemented by the Rational class, that inherits from both LinExpr and
598 fractions.Fraction classes.
599 """
600
601 __slots__ = (
602 '_coefficients',
603 '_constant',
604 '_symbols',
605 '_dimension',
606 ) + Fraction.__slots__
607
608 def __new__(cls, numerator=0, denominator=None):
609 self = object().__new__(cls)
610 self._coefficients = {}
611 self._constant = Fraction(numerator, denominator)
612 self._symbols = ()
613 self._dimension = 0
614 self._numerator = self._constant.numerator
615 self._denominator = self._constant.denominator
616 return self
617
618 def __hash__(self):
619 return Fraction.__hash__(self)
620
621 @property
622 def constant(self):
623 return self
624
625 def isconstant(self):
626 return True
627
628 def __bool__(self):
629 return Fraction.__bool__(self)
630
631 def __repr__(self):
632 if self.denominator == 1:
633 return '{!r}'.format(self.numerator)
634 else:
635 return '{!r}/{!r}'.format(self.numerator, self.denominator)
636
637 def _repr_latex_(self):
638 if self.denominator == 1:
639 return '$${}$$'.format(self.numerator)
640 elif self.numerator < 0:
641 return '$$-\\frac{{{}}}{{{}}}$$'.format(-self.numerator,
642 self.denominator)
643 else:
644 return '$$\\frac{{{}}}{{{}}}$$'.format(self.numerator,
645 self.denominator)