99ad4ec3e162e36f5d6a5e855a19a47c2dca945a
[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 for symbol, coefficient in self._coefficients.items():
135 yield symbol, coefficient
136
137 @property
138 def constant(self):
139 """
140 The constant term of the expression.
141 """
142 return self._constant
143
144 @property
145 def symbols(self):
146 """
147 The tuple of symbols present in the expression, sorted according to
148 Symbol.sortkey().
149 """
150 return self._symbols
151
152 @property
153 def dimension(self):
154 """
155 The dimension of the expression, i.e. the number of symbols present in
156 it.
157 """
158 return self._dimension
159
160 def __hash__(self):
161 return hash((tuple(self._coefficients.items()), self._constant))
162
163 def isconstant(self):
164 """
165 Return True if the expression only consists of a constant term. In this
166 case, it is a Rational instance.
167 """
168 return False
169
170 def issymbol(self):
171 """
172 Return True if an expression only consists of a symbol with coefficient
173 1. In this case, it is a Symbol instance.
174 """
175 return False
176
177 def values(self):
178 """
179 Iterate over the coefficient values in the expression, and the constant
180 term.
181 """
182 for coefficient in self._coefficients.values():
183 yield coefficient
184 yield self._constant
185
186 def __bool__(self):
187 return True
188
189 def __pos__(self):
190 return self
191
192 def __neg__(self):
193 return self * -1
194
195 @_polymorphic
196 def __add__(self, other):
197 """
198 Return the sum of two linear expressions.
199 """
200 coefficients = defaultdict(Fraction, self._coefficients)
201 for symbol, coefficient in other._coefficients.items():
202 coefficients[symbol] += coefficient
203 constant = self._constant + other._constant
204 return LinExpr(coefficients, constant)
205
206 __radd__ = __add__
207
208 @_polymorphic
209 def __sub__(self, other):
210 """
211 Return the difference between two linear expressions.
212 """
213 coefficients = defaultdict(Fraction, self._coefficients)
214 for symbol, coefficient in other._coefficients.items():
215 coefficients[symbol] -= coefficient
216 constant = self._constant - other._constant
217 return LinExpr(coefficients, constant)
218
219 @_polymorphic
220 def __rsub__(self, other):
221 return other - self
222
223 def __mul__(self, other):
224 """
225 Return the product of the linear expression by a rational.
226 """
227 if isinstance(other, numbers.Rational):
228 coefficients = ((symbol, coefficient * other)
229 for symbol, coefficient in self._coefficients.items())
230 constant = self._constant * other
231 return LinExpr(coefficients, constant)
232 return NotImplemented
233
234 __rmul__ = __mul__
235
236 def __truediv__(self, other):
237 """
238 Return the quotient of the linear expression by a rational.
239 """
240 if isinstance(other, numbers.Rational):
241 coefficients = ((symbol, coefficient / other)
242 for symbol, coefficient in self._coefficients.items())
243 constant = self._constant / other
244 return LinExpr(coefficients, constant)
245 return NotImplemented
246
247 @_polymorphic
248 def __eq__(self, other):
249 """
250 Test whether two linear expressions are equal.
251 """
252 return isinstance(other, LinExpr) and \
253 self._coefficients == other._coefficients and \
254 self._constant == other._constant
255
256 def __le__(self, other):
257 from .polyhedra import Le
258 return Le(self, other)
259
260 def __lt__(self, other):
261 from .polyhedra import Lt
262 return Lt(self, other)
263
264 def __ge__(self, other):
265 from .polyhedra import Ge
266 return Ge(self, other)
267
268 def __gt__(self, other):
269 from .polyhedra import Gt
270 return Gt(self, other)
271
272 def scaleint(self):
273 """
274 Return the expression multiplied by its lowest common denominator to
275 make all values integer.
276 """
277 lcm = functools.reduce(lambda a, b: a*b // gcd(a, b),
278 [value.denominator for value in self.values()])
279 return self * lcm
280
281 def subs(self, symbol, expression=None):
282 """
283 Substitute the given symbol by an expression and return the resulting
284 expression. Raise TypeError if the resulting expression is not linear.
285
286 >>> x, y = symbols('x y')
287 >>> e = x + 2*y + 1
288 >>> e.subs(y, x - 1)
289 3*x - 1
290
291 To perform multiple substitutions at once, pass a sequence or a
292 dictionary of (old, new) pairs to subs.
293
294 >>> e.subs({x: y, y: x})
295 2*x + y + 1
296 """
297 if expression is None:
298 if isinstance(symbol, Mapping):
299 symbol = symbol.items()
300 substitutions = symbol
301 else:
302 substitutions = [(symbol, expression)]
303 result = self
304 for symbol, expression in substitutions:
305 if not isinstance(symbol, Symbol):
306 raise TypeError('symbols must be Symbol instances')
307 coefficients = [(othersymbol, coefficient)
308 for othersymbol, coefficient in result._coefficients.items()
309 if othersymbol != symbol]
310 coefficient = result._coefficients.get(symbol, 0)
311 constant = result._constant
312 result = LinExpr(coefficients, constant) + 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
426 raise TypeError('cannot convert dummy symbols')
427 elif isinstance(symbol, sympy.Symbol):
428 symbol = Symbol(symbol.name)
429 coefficients.append((symbol, coefficient))
430 else:
431 raise TypeError('non-linear expression: {!r}'.format(expr))
432 expr = LinExpr(coefficients, constant)
433 if not isinstance(expr, cls):
434 raise TypeError('cannot convert to a {} instance'.format(cls.__name__))
435 return expr
436
437 def tosympy(self):
438 """
439 Convert the linear expression to a sympy expression.
440 """
441 import sympy
442 expr = 0
443 for symbol, coefficient in self.coefficients():
444 term = coefficient * sympy.Symbol(symbol.name)
445 expr += term
446 expr += self.constant
447 return expr
448
449
450 class Symbol(LinExpr):
451 """
452 Symbols are the basic components to build expressions and constraints.
453 They correspond to mathematical variables. Symbols are instances of
454 class LinExpr and inherit its functionalities.
455
456 Two instances of Symbol are equal if they have the same name.
457 """
458
459 def __new__(cls, name):
460 """
461 Return a symbol with the name string given in argument.
462 """
463 if not isinstance(name, str):
464 raise TypeError('name must be a string')
465 node = ast.parse(name)
466 try:
467 name = node.body[0].value.id
468 except (AttributeError, SyntaxError):
469 raise SyntaxError('invalid syntax')
470 self = object().__new__(cls)
471 self._name = name
472 self._coefficients = {self: Fraction(1)}
473 self._constant = Fraction(0)
474 self._symbols = (self,)
475 self._dimension = 1
476 return self
477
478 @property
479 def name(self):
480 """
481 The name of the symbol.
482 """
483 return self._name
484
485 def __hash__(self):
486 return hash(self.sortkey())
487
488 def sortkey(self):
489 """
490 Return a sorting key for the symbol. It is useful to sort a list of
491 symbols in a consistent order, as comparison functions are overridden
492 (see the documentation of class LinExpr).
493
494 >>> sort(symbols, key=Symbol.sortkey)
495 """
496 return self.name,
497
498 def issymbol(self):
499 return True
500
501 def __eq__(self, other):
502 return self.sortkey() == other.sortkey()
503
504 def asdummy(self):
505 """
506 Return a new Dummy symbol instance with the same name.
507 """
508 return Dummy(self.name)
509
510 def __repr__(self):
511 return self.name
512
513 def _repr_latex_(self):
514 return '$${}$$'.format(self.name)
515
516
517 def symbols(names):
518 """
519 This function returns a tuple of symbols whose names are taken from a comma
520 or whitespace delimited string, or a sequence of strings. It is useful to
521 define several symbols at once.
522
523 >>> x, y = symbols('x y')
524 >>> x, y = symbols('x, y')
525 >>> x, y = symbols(['x', 'y'])
526 """
527 if isinstance(names, str):
528 names = names.replace(',', ' ').split()
529 return tuple(Symbol(name) for name in names)
530
531
532 class Dummy(Symbol):
533 """
534 A variation of Symbol in which all symbols are unique and identified by
535 an internal count index. If a name is not supplied then a string value
536 of the count index will be used. This is useful when a unique, temporary
537 variable is needed and the name of the variable used in the expression
538 is not important.
539
540 Unlike Symbol, Dummy instances with the same name are not equal:
541
542 >>> x = Symbol('x')
543 >>> x1, x2 = Dummy('x'), Dummy('x')
544 >>> x == x1
545 False
546 >>> x1 == x2
547 False
548 >>> x1 == x1
549 True
550 """
551
552 _count = 0
553
554 def __new__(cls, name=None):
555 """
556 Return a fresh dummy symbol with the name string given in argument.
557 """
558 if name is None:
559 name = 'Dummy_{}'.format(Dummy._count)
560 elif not isinstance(name, str):
561 raise TypeError('name must be a string')
562 self = object().__new__(cls)
563 self._index = Dummy._count
564 self._name = name.strip()
565 self._coefficients = {self: Fraction(1)}
566 self._constant = Fraction(0)
567 self._symbols = (self,)
568 self._dimension = 1
569 Dummy._count += 1
570 return self
571
572 def __hash__(self):
573 return hash(self.sortkey())
574
575 def sortkey(self):
576 return self._name, self._index
577
578 def __repr__(self):
579 return '_{}'.format(self.name)
580
581 def _repr_latex_(self):
582 return '$${}_{{{}}}$$'.format(self.name, self._index)
583
584
585 class Rational(LinExpr, Fraction):
586 """
587 A particular case of linear expressions are rational values, i.e. linear
588 expressions consisting only of a constant term, with no symbol. They are
589 implemented by the Rational class, that inherits from both LinExpr and
590 fractions.Fraction classes.
591 """
592
593 def __new__(cls, numerator=0, denominator=None):
594 self = object().__new__(cls)
595 self._coefficients = {}
596 self._constant = Fraction(numerator, denominator)
597 self._symbols = ()
598 self._dimension = 0
599 self._numerator = self._constant.numerator
600 self._denominator = self._constant.denominator
601 return self
602
603 def __hash__(self):
604 return Fraction.__hash__(self)
605
606 @property
607 def constant(self):
608 return self
609
610 def isconstant(self):
611 return True
612
613 def __bool__(self):
614 return Fraction.__bool__(self)
615
616 def __repr__(self):
617 if self.denominator == 1:
618 return '{!r}'.format(self.numerator)
619 else:
620 return '{!r}/{!r}'.format(self.numerator, self.denominator)
621
622 def _repr_latex_(self):
623 if self.denominator == 1:
624 return '$${}$$'.format(self.numerator)
625 elif self.numerator < 0:
626 return '$$-\\frac{{{}}}{{{}}}$$'.format(-self.numerator,
627 self.denominator)
628 else:
629 return '$$\\frac{{{}}}{{{}}}$$'.format(self.numerator,
630 self.denominator)