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