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