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