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