Add license to examples
[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 'Expression',
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, Expression):
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 Expression:
47 """
48 This class implements linear expressions.
49 """
50
51 def __new__(cls, coefficients=None, constant=0):
52 """
53 Create a new expression.
54 """
55 if isinstance(coefficients, str):
56 if constant != 0:
57 raise TypeError('too many arguments')
58 return Expression.fromstring(coefficients)
59 if coefficients is None:
60 return Rational(constant)
61 if isinstance(coefficients, Mapping):
62 coefficients = coefficients.items()
63 coefficients = list(coefficients)
64 for symbol, coefficient in coefficients:
65 if not isinstance(symbol, Symbol):
66 raise TypeError('symbols must be Symbol instances')
67 if not isinstance(coefficient, numbers.Rational):
68 raise TypeError('coefficients must be rational numbers')
69 if not isinstance(constant, numbers.Rational):
70 raise TypeError('constant must be a rational number')
71 if len(coefficients) == 0:
72 return Rational(constant)
73 if len(coefficients) == 1 and constant == 0:
74 symbol, coefficient = coefficients[0]
75 if coefficient == 1:
76 return symbol
77 coefficients = [(symbol, Fraction(coefficient))
78 for symbol, coefficient in coefficients if coefficient != 0]
79 coefficients.sort(key=lambda item: item[0].sortkey())
80 self = object().__new__(cls)
81 self._coefficients = OrderedDict(coefficients)
82 self._constant = Fraction(constant)
83 self._symbols = tuple(self._coefficients)
84 self._dimension = len(self._symbols)
85 return self
86
87 def coefficient(self, symbol):
88 """
89 Return the coefficient value of the given symbol.
90 """
91 if not isinstance(symbol, Symbol):
92 raise TypeError('symbol must be a Symbol instance')
93 return Rational(self._coefficients.get(symbol, 0))
94
95 __getitem__ = coefficient
96
97 def coefficients(self):
98 """
99 Return a list of the coefficients of an expression
100 """
101 for symbol, coefficient in self._coefficients.items():
102 yield symbol, Rational(coefficient)
103
104 @property
105 def constant(self):
106 """
107 Return the constant value of an expression.
108 """
109 return Rational(self._constant)
110
111 @property
112 def symbols(self):
113 """
114 Return a list of symbols in an expression.
115 """
116 return self._symbols
117
118 @property
119 def dimension(self):
120 """
121 Create and return a new linear expression from a string or a list of coefficients and a constant.
122 """
123 return self._dimension
124
125 def __hash__(self):
126 return hash((tuple(self._coefficients.items()), self._constant))
127
128 def isconstant(self):
129 """
130 Return true if an expression is a constant.
131 """
132 return False
133
134 def issymbol(self):
135 """
136 Return true if an expression is a symbol.
137 """
138 return False
139
140 def values(self):
141 """
142 Return the coefficient and constant values of an expression.
143 """
144 for coefficient in self._coefficients.values():
145 yield Rational(coefficient)
146 yield Rational(self._constant)
147
148 def __bool__(self):
149 return True
150
151 def __pos__(self):
152 return self
153
154 def __neg__(self):
155 return self * -1
156
157 @_polymorphic
158 def __add__(self, other):
159 """
160 Return the sum of two expressions.
161 """
162 coefficients = defaultdict(Fraction, self._coefficients)
163 for symbol, coefficient in other._coefficients.items():
164 coefficients[symbol] += coefficient
165 constant = self._constant + other._constant
166 return Expression(coefficients, constant)
167
168 __radd__ = __add__
169
170 @_polymorphic
171 def __sub__(self, other):
172 """
173 Return the difference between two expressions.
174 """
175 coefficients = defaultdict(Fraction, self._coefficients)
176 for symbol, coefficient in other._coefficients.items():
177 coefficients[symbol] -= coefficient
178 constant = self._constant - other._constant
179 return Expression(coefficients, constant)
180
181 @_polymorphic
182 def __rsub__(self, other):
183 return other - self
184
185 def __mul__(self, other):
186 """
187 Return the product of two expressions if other is a rational number.
188 """
189 if isinstance(other, numbers.Rational):
190 coefficients = ((symbol, coefficient * other)
191 for symbol, coefficient in self._coefficients.items())
192 constant = self._constant * other
193 return Expression(coefficients, constant)
194 return NotImplemented
195
196 __rmul__ = __mul__
197
198 def __truediv__(self, other):
199 if isinstance(other, numbers.Rational):
200 coefficients = ((symbol, coefficient / other)
201 for symbol, coefficient in self._coefficients.items())
202 constant = self._constant / other
203 return Expression(coefficients, constant)
204 return NotImplemented
205
206 @_polymorphic
207 def __eq__(self, other):
208 """
209 Test whether two expressions are equal
210 """
211 return isinstance(other, Expression) and \
212 self._coefficients == other._coefficients and \
213 self._constant == other._constant
214
215 def __le__(self, other):
216 from .polyhedra import Le
217 return Le(self, other)
218
219 def __lt__(self, other):
220 from .polyhedra import Lt
221 return Lt(self, other)
222
223 def __ge__(self, other):
224 from .polyhedra import Ge
225 return Ge(self, other)
226
227 def __gt__(self, other):
228 from .polyhedra import Gt
229 return Gt(self, other)
230
231 def scaleint(self):
232 """
233 Multiply an expression by a scalar to make all coefficients integer values.
234 """
235 lcm = functools.reduce(lambda a, b: a*b // gcd(a, b),
236 [value.denominator for value in self.values()])
237 return self * lcm
238
239 def subs(self, symbol, expression=None):
240 """
241 Subsitute symbol by expression in equations and return the resulting
242 expression.
243 """
244 if expression is None:
245 if isinstance(symbol, Mapping):
246 symbol = symbol.items()
247 substitutions = symbol
248 else:
249 substitutions = [(symbol, expression)]
250 result = self
251 for symbol, expression in substitutions:
252 if not isinstance(symbol, Symbol):
253 raise TypeError('symbols must be Symbol instances')
254 coefficients = [(othersymbol, coefficient)
255 for othersymbol, coefficient in result._coefficients.items()
256 if othersymbol != symbol]
257 coefficient = result._coefficients.get(symbol, 0)
258 constant = result._constant
259 result = Expression(coefficients, constant) + coefficient*expression
260 return result
261
262 @classmethod
263 def _fromast(cls, node):
264 if isinstance(node, ast.Module) and len(node.body) == 1:
265 return cls._fromast(node.body[0])
266 elif isinstance(node, ast.Expr):
267 return cls._fromast(node.value)
268 elif isinstance(node, ast.Name):
269 return Symbol(node.id)
270 elif isinstance(node, ast.Num):
271 return Rational(node.n)
272 elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
273 return -cls._fromast(node.operand)
274 elif isinstance(node, ast.BinOp):
275 left = cls._fromast(node.left)
276 right = cls._fromast(node.right)
277 if isinstance(node.op, ast.Add):
278 return left + right
279 elif isinstance(node.op, ast.Sub):
280 return left - right
281 elif isinstance(node.op, ast.Mult):
282 return left * right
283 elif isinstance(node.op, ast.Div):
284 return left / right
285 raise SyntaxError('invalid syntax')
286
287 _RE_NUM_VAR = re.compile(r'(\d+|\))\s*([^\W\d_]\w*|\()')
288
289 @classmethod
290 def fromstring(cls, string):
291 """
292 Create an expression from a string.
293 """
294 # add implicit multiplication operators, e.g. '5x' -> '5*x'
295 string = Expression._RE_NUM_VAR.sub(r'\1*\2', string)
296 tree = ast.parse(string, 'eval')
297 return cls._fromast(tree)
298
299 def __repr__(self):
300 string = ''
301 for i, (symbol, coefficient) in enumerate(self.coefficients()):
302 if coefficient == 1:
303 if i != 0:
304 string += ' + '
305 elif coefficient == -1:
306 string += '-' if i == 0 else ' - '
307 elif i == 0:
308 string += '{}*'.format(coefficient)
309 elif coefficient > 0:
310 string += ' + {}*'.format(coefficient)
311 else:
312 string += ' - {}*'.format(-coefficient)
313 string += '{}'.format(symbol)
314 constant = self.constant
315 if len(string) == 0:
316 string += '{}'.format(constant)
317 elif constant > 0:
318 string += ' + {}'.format(constant)
319 elif constant < 0:
320 string += ' - {}'.format(-constant)
321 return string
322
323 def _repr_latex_(self):
324 string = ''
325 for i, (symbol, coefficient) in enumerate(self.coefficients()):
326 if coefficient == 1:
327 if i != 0:
328 string += ' + '
329 elif coefficient == -1:
330 string += '-' if i == 0 else ' - '
331 elif i == 0:
332 string += '{}'.format(coefficient._repr_latex_().strip('$'))
333 elif coefficient > 0:
334 string += ' + {}'.format(coefficient._repr_latex_().strip('$'))
335 elif coefficient < 0:
336 string += ' - {}'.format((-coefficient)._repr_latex_().strip('$'))
337 string += '{}'.format(symbol._repr_latex_().strip('$'))
338 constant = self.constant
339 if len(string) == 0:
340 string += '{}'.format(constant._repr_latex_().strip('$'))
341 elif constant > 0:
342 string += ' + {}'.format(constant._repr_latex_().strip('$'))
343 elif constant < 0:
344 string += ' - {}'.format((-constant)._repr_latex_().strip('$'))
345 return '$${}$$'.format(string)
346
347 def _parenstr(self, always=False):
348 string = str(self)
349 if not always and (self.isconstant() or self.issymbol()):
350 return string
351 else:
352 return '({})'.format(string)
353
354 @classmethod
355 def fromsympy(cls, expr):
356 """
357 Convert sympy object to an expression.
358 """
359 import sympy
360 coefficients = []
361 constant = 0
362 for symbol, coefficient in expr.as_coefficients_dict().items():
363 coefficient = Fraction(coefficient.p, coefficient.q)
364 if symbol == sympy.S.One:
365 constant = coefficient
366 elif isinstance(symbol, sympy.Symbol):
367 symbol = Symbol(symbol.name)
368 coefficients.append((symbol, coefficient))
369 else:
370 raise ValueError('non-linear expression: {!r}'.format(expr))
371 return Expression(coefficients, constant)
372
373 def tosympy(self):
374 """
375 Return an expression as a sympy object.
376 """
377 import sympy
378 expr = 0
379 for symbol, coefficient in self.coefficients():
380 term = coefficient * sympy.Symbol(symbol.name)
381 expr += term
382 expr += self.constant
383 return expr
384
385
386 class Symbol(Expression):
387
388 def __new__(cls, name):
389 """
390 Create and return a symbol from a string.
391 """
392 if not isinstance(name, str):
393 raise TypeError('name must be a string')
394 self = object().__new__(cls)
395 self._name = name.strip()
396 self._coefficients = {self: Fraction(1)}
397 self._constant = Fraction(0)
398 self._symbols = (self,)
399 self._dimension = 1
400 return self
401
402 @property
403 def name(self):
404 return self._name
405
406 def __hash__(self):
407 return hash(self.sortkey())
408
409 def sortkey(self):
410 return self.name,
411
412 def issymbol(self):
413 return True
414
415 def __eq__(self, other):
416 return self.sortkey() == other.sortkey()
417
418 def asdummy(self):
419 """
420 Return a symbol as a Dummy Symbol.
421 """
422 return Dummy(self.name)
423
424 @classmethod
425 def _fromast(cls, node):
426 if isinstance(node, ast.Module) and len(node.body) == 1:
427 return cls._fromast(node.body[0])
428 elif isinstance(node, ast.Expr):
429 return cls._fromast(node.value)
430 elif isinstance(node, ast.Name):
431 return Symbol(node.id)
432 raise SyntaxError('invalid syntax')
433
434 def __repr__(self):
435 return self.name
436
437 def _repr_latex_(self):
438 return '$${}$$'.format(self.name)
439
440 @classmethod
441 def fromsympy(cls, expr):
442 import sympy
443 if isinstance(expr, sympy.Dummy):
444 return Dummy(expr.name)
445 elif isinstance(expr, sympy.Symbol):
446 return Symbol(expr.name)
447 else:
448 raise TypeError('expr must be a sympy.Symbol instance')
449
450
451 class Dummy(Symbol):
452 """
453 This class returns a dummy symbol to ensure that no variables are repeated in an expression
454 """
455 _count = 0
456
457 def __new__(cls, name=None):
458 """
459 Create and return a new dummy symbol.
460 """
461 if name is None:
462 name = 'Dummy_{}'.format(Dummy._count)
463 elif not isinstance(name, str):
464 raise TypeError('name must be a string')
465 self = object().__new__(cls)
466 self._index = Dummy._count
467 self._name = name.strip()
468 self._coefficients = {self: Fraction(1)}
469 self._constant = Fraction(0)
470 self._symbols = (self,)
471 self._dimension = 1
472 Dummy._count += 1
473 return self
474
475 def __hash__(self):
476 return hash(self.sortkey())
477
478 def sortkey(self):
479 return self._name, self._index
480
481 def __repr__(self):
482 return '_{}'.format(self.name)
483
484 def _repr_latex_(self):
485 return '$${}_{{{}}}$$'.format(self.name, self._index)
486
487
488 def symbols(names):
489 """
490 Transform strings into instances of the Symbol class
491 """
492 if isinstance(names, str):
493 names = names.replace(',', ' ').split()
494 return tuple(Symbol(name) for name in names)
495
496
497 class Rational(Expression, Fraction):
498 """
499 This class represents integers and rational numbers of any size.
500 """
501
502 def __new__(cls, numerator=0, denominator=None):
503 self = object().__new__(cls)
504 self._coefficients = {}
505 self._constant = Fraction(numerator, denominator)
506 self._symbols = ()
507 self._dimension = 0
508 self._numerator = self._constant.numerator
509 self._denominator = self._constant.denominator
510 return self
511
512 def __hash__(self):
513 return Fraction.__hash__(self)
514
515 @property
516 def constant(self):
517 """
518 Return rational as a constant.
519 """
520 return self
521
522 def isconstant(self):
523 """
524 Test whether a value is a constant.
525 """
526 return True
527
528 def __bool__(self):
529 return Fraction.__bool__(self)
530
531 def __repr__(self):
532 if self.denominator == 1:
533 return '{!r}'.format(self.numerator)
534 else:
535 return '{!r}/{!r}'.format(self.numerator, self.denominator)
536
537 def _repr_latex_(self):
538 if self.denominator == 1:
539 return '$${}$$'.format(self.numerator)
540 elif self.numerator < 0:
541 return '$$-\\frac{{{}}}{{{}}}$$'.format(-self.numerator,
542 self.denominator)
543 else:
544 return '$$\\frac{{{}}}{{{}}}$$'.format(self.numerator,
545 self.denominator)
546
547 @classmethod
548 def fromsympy(cls, expr):
549 """
550 Create a rational object from a sympy expression
551 """
552 import sympy
553 if isinstance(expr, sympy.Rational):
554 return Rational(expr.p, expr.q)
555 elif isinstance(expr, numbers.Rational):
556 return Rational(expr)
557 else:
558 raise TypeError('expr must be a sympy.Rational instance')