Check symbol names
[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 Rational(self._coefficients.get(symbol, 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, Rational(coefficient)
136
137 @property
138 def constant(self):
139 """
140 The constant term of the expression.
141 """
142 return Rational(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 Rational(coefficient)
184 yield Rational(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 ValueError 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.Symbol):
425 symbol = Symbol(symbol.name)
426 coefficients.append((symbol, coefficient))
427 else:
428 raise ValueError('non-linear expression: {!r}'.format(expr))
429 return LinExpr(coefficients, constant)
430
431 def tosympy(self):
432 """
433 Convert the linear expression to a sympy expression.
434 """
435 import sympy
436 expr = 0
437 for symbol, coefficient in self.coefficients():
438 term = coefficient * sympy.Symbol(symbol.name)
439 expr += term
440 expr += self.constant
441 return expr
442
443
444 class Symbol(LinExpr):
445 """
446 Symbols are the basic components to build expressions and constraints.
447 They correspond to mathematical variables. Symbols are instances of
448 class LinExpr and inherit its functionalities.
449
450 Two instances of Symbol are equal if they have the same name.
451 """
452
453 def __new__(cls, name):
454 """
455 Return a symbol with the name string given in argument.
456 """
457 if not isinstance(name, str):
458 raise TypeError('name must be a string')
459 node = ast.parse(name)
460 try:
461 name = node.body[0].value.id
462 except (AttributeError, SyntaxError):
463 raise SyntaxError('invalid syntax')
464 self = object().__new__(cls)
465 self._name = name
466 self._coefficients = {self: Fraction(1)}
467 self._constant = Fraction(0)
468 self._symbols = (self,)
469 self._dimension = 1
470 return self
471
472 @property
473 def name(self):
474 """
475 The name of the symbol.
476 """
477 return self._name
478
479 def __hash__(self):
480 return hash(self.sortkey())
481
482 def sortkey(self):
483 """
484 Return a sorting key for the symbol. It is useful to sort a list of
485 symbols in a consistent order, as comparison functions are overridden
486 (see the documentation of class LinExpr).
487
488 >>> sort(symbols, key=Symbol.sortkey)
489 """
490 return self.name,
491
492 def issymbol(self):
493 return True
494
495 def __eq__(self, other):
496 return self.sortkey() == other.sortkey()
497
498 def asdummy(self):
499 """
500 Return a new Dummy symbol instance with the same name.
501 """
502 return Dummy(self.name)
503
504 def __repr__(self):
505 return self.name
506
507 def _repr_latex_(self):
508 return '$${}$$'.format(self.name)
509
510 @classmethod
511 def fromsympy(cls, expr):
512 import sympy
513 if isinstance(expr, sympy.Dummy):
514 return Dummy(expr.name)
515 elif isinstance(expr, sympy.Symbol):
516 return Symbol(expr.name)
517 else:
518 raise TypeError('expr must be a sympy.Symbol instance')
519
520
521 class Dummy(Symbol):
522 """
523 A variation of Symbol in which all symbols are unique and identified by
524 an internal count index. If a name is not supplied then a string value
525 of the count index will be used. This is useful when a unique, temporary
526 variable is needed and the name of the variable used in the expression
527 is not important.
528
529 Unlike Symbol, Dummy instances with the same name are not equal:
530
531 >>> x = Symbol('x')
532 >>> x1, x2 = Dummy('x'), Dummy('x')
533 >>> x == x1
534 False
535 >>> x1 == x2
536 False
537 >>> x1 == x1
538 True
539 """
540
541 _count = 0
542
543 def __new__(cls, name=None):
544 """
545 Return a fresh dummy symbol with the name string given in argument.
546 """
547 if name is None:
548 name = 'Dummy_{}'.format(Dummy._count)
549 elif not isinstance(name, str):
550 raise TypeError('name must be a string')
551 self = object().__new__(cls)
552 self._index = Dummy._count
553 self._name = name.strip()
554 self._coefficients = {self: Fraction(1)}
555 self._constant = Fraction(0)
556 self._symbols = (self,)
557 self._dimension = 1
558 Dummy._count += 1
559 return self
560
561 def __hash__(self):
562 return hash(self.sortkey())
563
564 def sortkey(self):
565 return self._name, self._index
566
567 def __repr__(self):
568 return '_{}'.format(self.name)
569
570 def _repr_latex_(self):
571 return '$${}_{{{}}}$$'.format(self.name, self._index)
572
573
574 def symbols(names):
575 """
576 This function returns a tuple of symbols whose names are taken from a comma
577 or whitespace delimited string, or a sequence of strings. It is useful to
578 define several symbols at once.
579
580 >>> x, y = symbols('x y')
581 >>> x, y = symbols('x, y')
582 >>> x, y = symbols(['x', 'y'])
583 """
584 if isinstance(names, str):
585 names = names.replace(',', ' ').split()
586 return tuple(Symbol(name) for name in names)
587
588
589 class Rational(LinExpr, Fraction):
590 """
591 A particular case of linear expressions are rational values, i.e. linear
592 expressions consisting only of a constant term, with no symbol. They are
593 implemented by the Rational class, that inherits from both LinExpr and
594 fractions.Fraction classes.
595 """
596
597 def __new__(cls, numerator=0, denominator=None):
598 self = object().__new__(cls)
599 self._coefficients = {}
600 self._constant = Fraction(numerator, denominator)
601 self._symbols = ()
602 self._dimension = 0
603 self._numerator = self._constant.numerator
604 self._denominator = self._constant.denominator
605 return self
606
607 def __hash__(self):
608 return Fraction.__hash__(self)
609
610 @property
611 def constant(self):
612 return self
613
614 def isconstant(self):
615 return True
616
617 def __bool__(self):
618 return Fraction.__bool__(self)
619
620 def __repr__(self):
621 if self.denominator == 1:
622 return '{!r}'.format(self.numerator)
623 else:
624 return '{!r}/{!r}'.format(self.numerator, self.denominator)
625
626 def _repr_latex_(self):
627 if self.denominator == 1:
628 return '$${}$$'.format(self.numerator)
629 elif self.numerator < 0:
630 return '$$-\\frac{{{}}}{{{}}}$$'.format(-self.numerator,
631 self.denominator)
632 else:
633 return '$$\\frac{{{}}}{{{}}}$$'.format(self.numerator,
634 self.denominator)
635
636 @classmethod
637 def fromsympy(cls, expr):
638 import sympy
639 if isinstance(expr, sympy.Rational):
640 return Rational(expr.p, expr.q)
641 elif isinstance(expr, numbers.Rational):
642 return Rational(expr)
643 else:
644 raise TypeError('expr must be a sympy.Rational instance')