b23eea8845db0a425efde0d1578445f5915b0dac
[linpy.git] / pypol / linexprs.py
1 import ast
2 import functools
3 import numbers
4 import re
5
6 from collections import OrderedDict, defaultdict, Mapping
7 from fractions import Fraction, gcd
8
9
10 __all__ = [
11 'Expression',
12 'Symbol', 'Dummy', 'symbols',
13 'Rational',
14 ]
15
16
17 def _polymorphic(func):
18 @functools.wraps(func)
19 def wrapper(left, right):
20 if isinstance(right, Expression):
21 return func(left, right)
22 elif isinstance(right, numbers.Rational):
23 right = Rational(right)
24 return func(left, right)
25 return NotImplemented
26 return wrapper
27
28
29 class Expression:
30 """
31 This class implements linear expressions.
32 """
33
34 __slots__ = (
35 '_coefficients',
36 '_constant',
37 '_symbols',
38 '_dimension',
39 )
40
41 def __new__(cls, coefficients=None, constant=0):
42 if isinstance(coefficients, str):
43 if constant:
44 raise TypeError('too many arguments')
45 return Expression.fromstring(coefficients)
46 if coefficients is None:
47 return Rational(constant)
48 if isinstance(coefficients, Mapping):
49 coefficients = coefficients.items()
50 for symbol, coefficient in coefficients:
51 if not isinstance(symbol, Symbol):
52 raise TypeError('symbols must be Symbol instances')
53 coefficients = [(symbol, coefficient)
54 for symbol, coefficient in coefficients if coefficient != 0]
55 if len(coefficients) == 0:
56 return Rational(constant)
57 if len(coefficients) == 1 and constant == 0:
58 symbol, coefficient = coefficients[0]
59 if coefficient == 1:
60 return symbol
61 self = object().__new__(cls)
62 self._coefficients = OrderedDict()
63 for symbol, coefficient in sorted(coefficients,
64 key=lambda item: item[0].sortkey()):
65 if isinstance(coefficient, Rational):
66 coefficient = coefficient.constant
67 if not isinstance(coefficient, numbers.Rational):
68 raise TypeError('coefficients must be Rational instances')
69 self._coefficients[symbol] = coefficient
70 if isinstance(constant, Rational):
71 constant = constant.constant
72 if not isinstance(constant, numbers.Rational):
73 raise TypeError('constant must be a Rational instance')
74 self._constant = constant
75 self._symbols = tuple(self._coefficients)
76 self._dimension = len(self._symbols)
77 return self
78
79 def coefficient(self, symbol):
80 if not isinstance(symbol, Symbol):
81 raise TypeError('symbol must be a Symbol instance')
82 try:
83 return self._coefficients[symbol]
84 except KeyError:
85 return 0
86
87 __getitem__ = coefficient
88
89 def coefficients(self):
90 yield from self._coefficients.items()
91
92 @property
93 def constant(self):
94 return self._constant
95
96 @property
97 def symbols(self):
98 return self._symbols
99
100 @property
101 def dimension(self):
102 return self._dimension
103
104 def __hash__(self):
105 return hash((tuple(self._coefficients.items()), self._constant))
106
107 def isconstant(self):
108 return False
109
110 def issymbol(self):
111 return False
112
113 def values(self):
114 yield from self._coefficients.values()
115 yield self.constant
116
117 def __bool__(self):
118 return True
119
120 def __pos__(self):
121 return self
122
123 def __neg__(self):
124 return self * -1
125
126 @_polymorphic
127 def __add__(self, other):
128 coefficients = defaultdict(Rational, self.coefficients())
129 for symbol, coefficient in other.coefficients():
130 coefficients[symbol] += coefficient
131 constant = self.constant + other.constant
132 return Expression(coefficients, constant)
133
134 __radd__ = __add__
135
136 @_polymorphic
137 def __sub__(self, other):
138 coefficients = defaultdict(Rational, self.coefficients())
139 for symbol, coefficient in other.coefficients():
140 coefficients[symbol] -= coefficient
141 constant = self.constant - other.constant
142 return Expression(coefficients, constant)
143
144 def __rsub__(self, other):
145 return -(self - other)
146
147 @_polymorphic
148 def __mul__(self, other):
149 if other.isconstant():
150 coefficients = dict(self.coefficients())
151 for symbol in coefficients:
152 coefficients[symbol] *= other.constant
153 constant = self.constant * other.constant
154 return Expression(coefficients, constant)
155 if isinstance(other, Expression) and not self.isconstant():
156 raise ValueError('non-linear expression: '
157 '{} * {}'.format(self._parenstr(), other._parenstr()))
158 return NotImplemented
159
160 __rmul__ = __mul__
161
162 @_polymorphic
163 def __truediv__(self, other):
164 if other.isconstant():
165 coefficients = dict(self.coefficients())
166 for symbol in coefficients:
167 coefficients[symbol] = Rational(coefficients[symbol], other.constant)
168 constant = Rational(self.constant, other.constant)
169 return Expression(coefficients, constant)
170 if isinstance(other, Expression):
171 raise ValueError('non-linear expression: '
172 '{} / {}'.format(self._parenstr(), other._parenstr()))
173 return NotImplemented
174
175 def __rtruediv__(self, other):
176 if isinstance(other, self):
177 if self.isconstant():
178 return Rational(other, self.constant)
179 else:
180 raise ValueError('non-linear expression: '
181 '{} / {}'.format(other._parenstr(), self._parenstr()))
182 return NotImplemented
183
184 @_polymorphic
185 def __eq__(self, other):
186 # "normal" equality
187 # see http://docs.sympy.org/dev/tutorial/gotchas.html#equals-signs
188 return isinstance(other, Expression) and \
189 self._coefficients == other._coefficients and \
190 self.constant == other.constant
191
192 @_polymorphic
193 def __le__(self, other):
194 from .polyhedra import Le
195 return Le(self, other)
196
197 @_polymorphic
198 def __lt__(self, other):
199 from .polyhedra import Lt
200 return Lt(self, other)
201
202 @_polymorphic
203 def __ge__(self, other):
204 from .polyhedra import Ge
205 return Ge(self, other)
206
207 @_polymorphic
208 def __gt__(self, other):
209 from .polyhedra import Gt
210 return Gt(self, other)
211
212 def scaleint(self):
213 lcm = functools.reduce(lambda a, b: a*b // gcd(a, b),
214 [value.denominator for value in self.values()])
215 return self * lcm
216
217 def subs(self, symbol, expression=None):
218 if expression is None:
219 if isinstance(symbol, Mapping):
220 symbol = symbol.items()
221 substitutions = symbol
222 else:
223 substitutions = [(symbol, expression)]
224 result = self
225 for symbol, expression in substitutions:
226 coefficients = [(othersymbol, coefficient)
227 for othersymbol, coefficient in result.coefficients()
228 if othersymbol != symbol]
229 coefficient = result.coefficient(symbol)
230 constant = result.constant
231 result = Expression(coefficients, constant) + coefficient*expression
232 return result
233
234 @classmethod
235 def _fromast(cls, node):
236 if isinstance(node, ast.Module) and len(node.body) == 1:
237 return cls._fromast(node.body[0])
238 elif isinstance(node, ast.Expr):
239 return cls._fromast(node.value)
240 elif isinstance(node, ast.Name):
241 return Symbol(node.id)
242 elif isinstance(node, ast.Num):
243 return Rational(node.n)
244 elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
245 return -cls._fromast(node.operand)
246 elif isinstance(node, ast.BinOp):
247 left = cls._fromast(node.left)
248 right = cls._fromast(node.right)
249 if isinstance(node.op, ast.Add):
250 return left + right
251 elif isinstance(node.op, ast.Sub):
252 return left - right
253 elif isinstance(node.op, ast.Mult):
254 return left * right
255 elif isinstance(node.op, ast.Div):
256 return left / right
257 raise SyntaxError('invalid syntax')
258
259 _RE_NUM_VAR = re.compile(r'(\d+|\))\s*([^\W\d_]\w*|\()')
260
261 @classmethod
262 def fromstring(cls, string):
263 # add implicit multiplication operators, e.g. '5x' -> '5*x'
264 string = Expression._RE_NUM_VAR.sub(r'\1*\2', string)
265 tree = ast.parse(string, 'eval')
266 return cls._fromast(tree)
267
268 def __repr__(self):
269 string = ''
270 for i, (symbol, coefficient) in enumerate(self.coefficients()):
271 if coefficient == 1:
272 string += '' if i == 0 else ' + '
273 string += '{!r}'.format(symbol)
274 elif coefficient == -1:
275 string += '-' if i == 0 else ' - '
276 string += '{!r}'.format(symbol)
277 else:
278 if i == 0:
279 string += '{}*{!r}'.format(coefficient, symbol)
280 elif coefficient > 0:
281 string += ' + {}*{!r}'.format(coefficient, symbol)
282 else:
283 string += ' - {}*{!r}'.format(-coefficient, symbol)
284 constant = self.constant
285 if len(string) == 0:
286 string += '{}'.format(constant)
287 elif constant > 0:
288 string += ' + {}'.format(constant)
289 elif constant < 0:
290 string += ' - {}'.format(-constant)
291 return string
292
293 def _parenstr(self, always=False):
294 string = str(self)
295 if not always and (self.isconstant() or self.issymbol()):
296 return string
297 else:
298 return '({})'.format(string)
299
300 @classmethod
301 def fromsympy(cls, expr):
302 import sympy
303 coefficients = []
304 constant = 0
305 for symbol, coefficient in expr.as_coefficients_dict().items():
306 coefficient = Fraction(coefficient.p, coefficient.q)
307 if symbol == sympy.S.One:
308 constant = coefficient
309 elif isinstance(symbol, sympy.Symbol):
310 symbol = Symbol(symbol.name)
311 coefficients.append((symbol, coefficient))
312 else:
313 raise ValueError('non-linear expression: {!r}'.format(expr))
314 return Expression(coefficients, constant)
315
316 def tosympy(self):
317 import sympy
318 expr = 0
319 for symbol, coefficient in self.coefficients():
320 term = coefficient * sympy.Symbol(symbol.name)
321 expr += term
322 expr += self.constant
323 return expr
324
325
326 class Symbol(Expression):
327
328 __slots__ = (
329 '_name',
330 )
331
332 def __new__(cls, name):
333 if not isinstance(name, str):
334 raise TypeError('name must be a string')
335 self = object().__new__(cls)
336 self._name = name.strip()
337 return self
338
339 @property
340 def name(self):
341 return self._name
342
343 def __hash__(self):
344 return hash(self.sortkey())
345
346 def coefficient(self, symbol):
347 if not isinstance(symbol, Symbol):
348 raise TypeError('symbol must be a Symbol instance')
349 if symbol == self:
350 return 1
351 else:
352 return 0
353
354 def coefficients(self):
355 yield self, 1
356
357 @property
358 def constant(self):
359 return 0
360
361 @property
362 def symbols(self):
363 return self,
364
365 @property
366 def dimension(self):
367 return 1
368
369 def sortkey(self):
370 return self.name,
371
372 def issymbol(self):
373 return True
374
375 def values(self):
376 yield 1
377
378 def __eq__(self, other):
379 return not isinstance(other, Dummy) and isinstance(other, Symbol) \
380 and self.name == other.name
381
382 def asdummy(self):
383 return Dummy(self.name)
384
385 @classmethod
386 def _fromast(cls, node):
387 if isinstance(node, ast.Module) and len(node.body) == 1:
388 return cls._fromast(node.body[0])
389 elif isinstance(node, ast.Expr):
390 return cls._fromast(node.value)
391 elif isinstance(node, ast.Name):
392 return Symbol(node.id)
393 raise SyntaxError('invalid syntax')
394
395 def __repr__(self):
396 return self.name
397
398 @classmethod
399 def fromsympy(cls, expr):
400 import sympy
401 if isinstance(expr, sympy.Symbol):
402 return cls(expr.name)
403 else:
404 raise TypeError('expr must be a sympy.Symbol instance')
405
406
407 class Dummy(Symbol):
408
409 __slots__ = (
410 '_name',
411 '_index',
412 )
413
414 _count = 0
415
416 def __new__(cls, name=None):
417 if name is None:
418 name = 'Dummy_{}'.format(Dummy._count)
419 self = object().__new__(cls)
420 self._name = name.strip()
421 self._index = Dummy._count
422 Dummy._count += 1
423 return self
424
425 def __hash__(self):
426 return hash(self.sortkey())
427
428 def sortkey(self):
429 return self._name, self._index
430
431 def __eq__(self, other):
432 return isinstance(other, Dummy) and self._index == other._index
433
434 def __repr__(self):
435 return '_{}'.format(self.name)
436
437
438 def symbols(names):
439 if isinstance(names, str):
440 names = names.replace(',', ' ').split()
441 return tuple(Symbol(name) for name in names)
442
443
444 class Rational(Expression):
445
446 __slots__ = (
447 '_constant',
448 )
449
450 def __new__(cls, numerator=0, denominator=None):
451 self = object().__new__(cls)
452 if denominator is None and isinstance(numerator, Rational):
453 self._constant = numerator.constant
454 else:
455 self._constant = Fraction(numerator, denominator)
456 return self
457
458 def __hash__(self):
459 return hash(self.constant)
460
461 def coefficient(self, symbol):
462 if not isinstance(symbol, Symbol):
463 raise TypeError('symbol must be a Symbol instance')
464 return 0
465
466 def coefficients(self):
467 yield from ()
468
469 @property
470 def symbols(self):
471 return ()
472
473 @property
474 def dimension(self):
475 return 0
476
477 def isconstant(self):
478 return True
479
480 def values(self):
481 yield self._constant
482
483 @_polymorphic
484 def __eq__(self, other):
485 return isinstance(other, Rational) and self.constant == other.constant
486
487 def __bool__(self):
488 return self.constant != 0
489
490 @classmethod
491 def fromstring(cls, string):
492 if not isinstance(string, str):
493 raise TypeError('string must be a string instance')
494 return Rational(Fraction(string))
495
496 @classmethod
497 def fromsympy(cls, expr):
498 import sympy
499 if isinstance(expr, sympy.Rational):
500 return Rational(expr.p, expr.q)
501 elif isinstance(expr, numbers.Rational):
502 return Rational(expr)
503 else:
504 raise TypeError('expr must be a sympy.Rational instance')