Use property decorators for symbols method
[linpy.git] / polyp.py
1
2 import functools
3 import numbers
4 import re
5 import subprocess
6
7 from fractions import Fraction, gcd
8
9
10 __all__ = [
11 'Expression',
12 'Constant', 'Symbol', 'symbols',
13 'Eq', 'Le', 'Lt', 'Ge', 'Gt',
14 'Polyhedron',
15 'empty', 'universe'
16 ]
17
18
19 def _strsymbol(symbol):
20 if isinstance(symbol, str):
21 return symbol
22 elif isinstance(symbol, Expression) and symbol.issymbol():
23 return str(symbol)
24 else:
25 raise TypeError('symbol must be a string or a symbol')
26
27 def _strsymbols(symbols):
28 if isinstance(symbols, str):
29 return symbols.replace(',', ' ').split()
30 else:
31 return [_strsymbol(symbol) for symbol in symbols]
32
33
34 class _ISCC:
35
36 command =['iscc']
37 debug = False
38
39 @classmethod
40 def eval(cls, input):
41 if not input.endswith(';'):
42 input += ';'
43 proc = subprocess.Popen(cls.command,
44 stdin=subprocess.PIPE, stdout=subprocess.PIPE,
45 universal_newlines=True)
46 output, error = proc.communicate(input=input)
47 output = output.strip()
48 if cls.debug:
49 print('iscc: {!r} -> {!r}'.format(input, output))
50 return output
51
52 @classmethod
53 def symbols(cls, symbols):
54 symbols = _strsymbols(symbols)
55 return '[{}]'.format(', '.join(symbols))
56
57 @classmethod
58 def constraints(cls, equalities, inequalities):
59 strings = []
60 for equality in equalities:
61 strings.append('{} = 0'.format(equality))
62 for inequality in inequalities:
63 strings.append('{} <= 0'.format(inequality))
64 return ' and '.join(strings)
65
66 @classmethod
67 def set(cls, symbols, equalities, inequalities):
68 return '{{ {} : {} }}'.format(cls.symbols(symbols),
69 cls.constraints(equalities, inequalities))
70
71 @classmethod
72 def map(cls, lsymbols, rsymbols, equalities, inequalities):
73 return '{{ {} -> {} : {} }}'.format(
74 cls.symbols(lsymbols), cls.symbols(rsymbols),
75 cls.constraints(equalities, inequalities))
76
77 _iscc = _ISCC()
78
79
80 def _polymorphic_method(func):
81 @functools.wraps(func)
82 def wrapper(a, b):
83 if isinstance(b, Expression):
84 return func(a, b)
85 if isinstance(b, numbers.Rational):
86 b = Constant(b)
87 return func(a, b)
88 return NotImplemented
89 return wrapper
90
91 def _polymorphic_operator(func):
92 # A polymorphic operator should call a polymorphic method, hence we just
93 # have to test the left operand.
94 @functools.wraps(func)
95 def wrapper(a, b):
96 if isinstance(a, numbers.Rational):
97 a = Constant(a)
98 return func(a, b)
99 elif isinstance(a, Expression):
100 return func(a, b)
101 raise TypeError('arguments must be linear expressions')
102 return wrapper
103
104
105 class Expression:
106 """
107 This class implements linear expressions.
108 """
109
110 def __new__(cls, coefficients=None, constant=0):
111 self = super().__new__(cls)
112 self._coefficients = {}
113 if isinstance(coefficients, dict):
114 coefficients = coefficients.items()
115 if coefficients is not None:
116 for symbol, coefficient in coefficients:
117 symbol = _strsymbol(symbol)
118 if not isinstance(coefficient, numbers.Rational):
119 raise TypeError('coefficients must be rational numbers')
120 if coefficient != 0:
121 self._coefficients[symbol] = coefficient
122 if not isinstance(constant, numbers.Rational):
123 raise TypeError('constant must be a rational number')
124 self._constant = constant
125 return self
126
127 @classmethod
128 def _fromiscc(cls, symbols, string):
129 symbols = _strsymbols(symbols)
130 string = re.sub(r'(\d+)\s*([a-zA-Z_]\w*)',
131 lambda m: '{}*{}'.format(m.group(1), m.group(2)),
132 string)
133 context = {}
134 for symbol in symbols:
135 context[symbol] = Symbol(symbol)
136 return eval(string, context)
137
138 @property
139 def symbols(self):
140 return tuple(sorted(self._coefficients))
141
142 @property
143 def dimension(self):
144 return len(list(self.symbols))
145
146 def coefficient(self, symbol):
147 symbol = _strsymbol(symbol)
148 try:
149 return self._coefficients[symbol]
150 except KeyError:
151 return 0
152
153 __getitem__ = coefficient
154
155 def coefficients(self):
156 for symbol in self.symbols:
157 yield symbol, self.coefficient(symbol)
158
159 @property
160 def constant(self):
161 return self._constant
162
163 def isconstant(self):
164 return len(self._coefficients) == 0
165
166 def values(self):
167 for symbol in self.symbols:
168 yield self.coefficient(symbol)
169 yield self.constant
170
171 def symbol(self):
172 if not self.issymbol():
173 raise ValueError('not a symbol: {}'.format(self))
174 for symbol in self.symbols:
175 return symbol
176
177 def issymbol(self):
178 return len(self._coefficients) == 1 and self._constant == 0
179
180 def __bool__(self):
181 return (not self.isconstant()) or bool(self.constant)
182
183 def __pos__(self):
184 return self
185
186 def __neg__(self):
187 return self * -1
188
189 @_polymorphic_method
190 def __add__(self, other):
191 coefficients = dict(self.coefficients())
192 for symbol, coefficient in other.coefficients():
193 if symbol in coefficients:
194 coefficients[symbol] += coefficient
195 else:
196 coefficients[symbol] = coefficient
197 constant = self.constant + other.constant
198 return Expression(coefficients, constant)
199
200 __radd__ = __add__
201
202 @_polymorphic_method
203 def __sub__(self, other):
204 coefficients = dict(self.coefficients())
205 for symbol, coefficient in other.coefficients():
206 if symbol in coefficients:
207 coefficients[symbol] -= coefficient
208 else:
209 coefficients[symbol] = -coefficient
210 constant = self.constant - other.constant
211 return Expression(coefficients, constant)
212
213 def __rsub__(self, other):
214 return -(self - other)
215
216 @_polymorphic_method
217 def __mul__(self, other):
218 if other.isconstant():
219 coefficients = dict(self.coefficients())
220 for symbol in coefficients:
221 coefficients[symbol] *= other.constant
222 constant = self.constant * other.constant
223 return Expression(coefficients, constant)
224 if isinstance(other, Expression) and not self.isconstant():
225 raise ValueError('non-linear expression: '
226 '{} * {}'.format(self._prepr(), other._prepr()))
227 return NotImplemented
228
229 __rmul__ = __mul__
230
231 def __repr__(self):
232 string = ''
233 symbols = self.symbols
234 i = 0
235 for symbol in symbols:
236 coefficient = self[symbol]
237 if coefficient == 1:
238 if i == 0:
239 string += symbol
240 else:
241 string += ' + {}'.format(symbol)
242 elif coefficient == -1:
243 if i == 0:
244 string += '-{}'.format(symbol)
245 else:
246 string += ' - {}'.format(symbol)
247 else:
248 if i == 0:
249 string += '{}*{}'.format(coefficient, symbol)
250 elif coefficient > 0:
251 string += ' + {}*{}'.format(coefficient, symbol)
252 else:
253 assert coefficient < 0
254 coefficient *= -1
255 string += ' - {}*{}'.format(coefficient, symbol)
256 i += 1
257 constant = self.constant
258 if constant != 0 and i == 0:
259 string += '{}'.format(constant)
260 elif constant > 0:
261 string += ' + {}'.format(constant)
262 elif constant < 0:
263 constant *= -1
264 string += ' - {}'.format(constant)
265 if string == '':
266 string = '0'
267 return string
268
269 def _prepr(self, always=False):
270 string = str(self)
271 if not always and (self.isconstant() or self.issymbol()):
272 return string
273 else:
274 return '({})'.format(string)
275
276 @_polymorphic_method
277 def __eq__(self, other):
278 # "normal" equality
279 # see http://docs.sympy.org/dev/tutorial/gotchas.html#equals-signs
280 return isinstance(other, Expression) and \
281 self._coefficients == other._coefficients and \
282 self.constant == other.constant
283
284 def __hash__(self):
285 return hash((tuple(self._coefficients), self._constant))
286
287 @_polymorphic_method
288 def _eq(self, other):
289 return Polyhedron(equalities=[self - other])
290
291 @_polymorphic_method
292 def __le__(self, other):
293 return Polyhedron(inequalities=[self - other])
294
295 @_polymorphic_method
296 def __lt__(self, other):
297 return Polyhedron(inequalities=[self - other + 1])
298
299 @_polymorphic_method
300 def __ge__(self, other):
301 return Polyhedron(inequalities=[other - self])
302
303 @_polymorphic_method
304 def __gt__(self, other):
305 return Polyhedron(inequalities=[other - self + 1])
306
307
308 def Constant(numerator=0, denominator=None):
309 if denominator is None and isinstance(numerator, numbers.Rational):
310 return Expression(constant=numerator)
311 else:
312 return Expression(constant=Fraction(numerator, denominator))
313
314 def Symbol(name):
315 name = _strsymbol(name)
316 return Expression(coefficients={name: 1})
317
318 def symbols(names):
319 names = _strsymbols(names)
320 return (Symbol(name) for name in names)
321
322
323 @_polymorphic_operator
324 def Eq(a, b):
325 return a._eq(b)
326
327 @_polymorphic_operator
328 def Le(a, b):
329 return a <= b
330
331 @_polymorphic_operator
332 def Lt(a, b):
333 return a < b
334
335 @_polymorphic_operator
336 def Ge(a, b):
337 return a >= b
338
339 @_polymorphic_operator
340 def Gt(a, b):
341 return a > b
342
343
344 class Polyhedron:
345 """
346 This class implements polyhedrons.
347 """
348
349 def __new__(cls, equalities=None, inequalities=None):
350 if equalities is None:
351 equalities = []
352 if inequalities is None:
353 inequalities = []
354 symbols = set()
355 for equality in equalities:
356 symbols.update(equality.symbols)
357 for inequality in inequalities:
358 symbols.update(inequality.symbols)
359 symbols = list(symbols)
360 string = _iscc.set(symbols, equalities, inequalities)
361 string = _iscc.eval(string)
362 return cls._fromiscc(symbols, string)
363
364 def _toiscc(self, symbols):
365 return _iscc.set(symbols, self.equalities, self.inequalities)
366
367 @classmethod
368 def _fromiscc(cls, symbols, string):
369 if re.match(r'^\s*\{\s*\}\s*$', string):
370 return empty
371 self = super().__new__(cls)
372 self._symbols = sorted(_strsymbols(symbols))
373 self._equalities = []
374 self._inequalities = []
375 string = re.sub(r'^\s*\{\s*(.*?)\s*\}\s*$', lambda m: m.group(1), string)
376 if ':' not in string:
377 string = string + ':'
378 vstr, cstr = re.split(r'\s*:\s*', string)
379 assert vstr != ''
380 vstr = re.sub(r'^\s*\[\s*(.*?)\s*\]\s*$', lambda m: m.group(1), vstr)
381 toks = list(filter(None, re.split(r'\s*,\s*', vstr)))
382 assert len(toks) == len(symbols)
383 for i in range(len(symbols)):
384 symbol = symbols[i]
385 if toks[i] != symbol:
386 expr = Expression._fromiscc(symbols, toks[i])
387 self._equalities.append(Symbol(symbol) - expr)
388 if cstr != '':
389 cstrs = re.split(r'\s*and\s*', cstr)
390 for cstr in cstrs:
391 lhs, op, rhs = re.split(r'\s*(<=|>=|<|>|=)\s*', cstr)
392 lhs = Expression._fromiscc(symbols, lhs)
393 rhs = Expression._fromiscc(symbols, rhs)
394 if op == '=':
395 self._equalities.append(lhs - rhs)
396 elif op == '<=':
397 self._inequalities.append(lhs - rhs)
398 elif op == '>=':
399 self._inequalities.append(rhs - lhs)
400 elif op == '<':
401 self._inequalities.append(lhs - rhs + 1)
402 elif op == '>':
403 self._inequalities.append(rhs - lhs + 1)
404 return self
405
406 @property
407 def equalities(self):
408 yield from self._equalities
409
410 @property
411 def inequalities(self):
412 yield from self._inequalities
413
414 def constraints(self):
415 yield from self.equalities
416 yield from self.inequalities
417
418 @property
419 def symbols(self):
420 return tuple(self._symbols)
421
422 @property
423 def dimension(self):
424 return len(self.symbols)
425
426 def __bool__(self):
427 # return false if the polyhedron is empty, true otherwise
428 raise not self.isempty()
429
430 def _symbolunion(self, *others):
431 symbols = set(self.symbols)
432 for other in others:
433 symbols.update(other.symbols)
434 return sorted(symbols)
435
436 def __eq__(self, other):
437 symbols = self._symbolunion(other)
438 string = '{} = {}'.format(self._toiscc(symbols), other._toiscc(symbols))
439 string = _iscc.eval(string)
440 return string == 'True'
441
442 def isempty(self):
443 return self == empty
444
445 def isuniverse(self):
446 return self == universe
447
448 def isdisjoint(self, other):
449 # return true if the polyhedron has no elements in common with other
450 return (self & other).isempty()
451
452 def issubset(self, other):
453 symbols = self._symbolunion(other)
454 string = '{} <= {}'.format(self._toiscc(symbols), other._toiscc(symbols))
455 string = _iscc.eval(string)
456 return string == 'True'
457
458 def __le__(self, other):
459 return self.issubset(other)
460
461 def __lt__(self, other):
462 symbols = self._symbolunion(other)
463 string = '{} < {}'.format(self._toiscc(symbols), other._toiscc(symbols))
464 string = _iscc.eval(string)
465 return string == 'True'
466
467 def issuperset(self, other):
468 # test whether every element in other is in the polyhedron
469 symbols = self._symbolunion(other)
470 string = '{} >= {}'.format(self._toiscc(symbols), other._toiscc(symbols))
471 string = _iscc.eval(string)
472 return string == 'True'
473
474 def __ge__(self, other):
475 return self.issuperset(other)
476
477 def __gt__(self, other):
478 symbols = self._symbolunion(other)
479 string = '{} > {}'.format(self._toiscc(symbols), other._toiscc(symbols))
480 string = _iscc.eval(string)
481 return string == 'True'
482
483 def union(self, *others):
484 # return a new polyhedron with elements from the polyhedron and all
485 # others (convex union)
486 symbols = self._symbolunion(*others)
487 strings = [self._toiscc(symbols)]
488 for other in others:
489 strings.append(other._toiscc(symbols))
490 string = ' + '.join(strings)
491 string = _iscc.eval('poly ({})'.format(string))
492 return Polyhedron._fromiscc(symbols, string)
493
494 def __or__(self, other):
495 return self.union(other)
496
497 def intersection(self, *others):
498 symbols = self._symbolunion(*others)
499 strings = [self._toiscc(symbols)]
500 for other in others:
501 strings.append(other._toiscc(symbols))
502 string = ' * '.join(strings)
503 string = _iscc.eval('poly ({})'.format(string))
504 return Polyhedron._fromiscc(symbols, string)
505
506 def __and__(self, other):
507 return self.intersection(other)
508
509 def difference(self, *others):
510 # return a new polyhedron with elements in the polyhedron that are not
511 # in the others
512 symbols = self._symbolunion(*others)
513 strings = [self._toiscc(symbols)]
514 for other in others:
515 strings.append(other._toiscc(symbols))
516 string = ' - '.join(strings)
517 string = _iscc.eval('poly ({})'.format(string))
518 return Polyhedron._fromiscc(symbols, string)
519
520 def __sub__(self, other):
521 return self.difference(other)
522
523 def projection(self, symbols):
524 symbols = _strsymbols(symbols)
525 string = _iscc.map(symbols, self.symbols,
526 self.equalities, self.inequalities)
527 string = _iscc.eval('poly (dom ({}))'.format(string))
528 return Polyhedron._fromiscc(symbols, string)
529
530 def __getitem__(self, symbol):
531 return self.projection([symbol])
532
533 def __repr__(self):
534 constraints = []
535 for constraint in self.equalities:
536 constraints.append('{} == 0'.format(constraint))
537 for constraint in self.inequalities:
538 constraints.append('{} <= 0'.format(constraint))
539 if len(constraints) == 0:
540 return 'universe'
541 elif len(constraints) == 1:
542 string = constraints[0]
543 return 'empty' if string == '1 == 0' else string
544 else:
545 strings = ['({})'.format(constraint) for constraint in constraints]
546 return ' & '.join(strings)
547
548 empty = Eq(1, 0)
549
550 universe = Polyhedron()
551
552
553 if __name__ == '__main__':
554 x, y, z = symbols('x y z')
555 sq0 = (0 <= x) & (x <= 1) & (0 <= y) & (y <= 1)
556 sq1 = (1 <= x) & (x <= 2) & (1 <= y) & (y <= 2)
557 print('sq0 ∪ sq1 =', sq0 | sq1)
558 print('sq0 ∩ sq1 =', sq0 & sq1)
559 print('sq0[x] =', sq0[x])